


"""simple docstring"""







import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class 						_A  (  _a   ,unittest.TestCase  ):







    """simple docstring"""




    UpperCAmelCase	:							List[Any]     =	KandinskyVaaImgaImgPipeline
    UpperCAmelCase	:							Dict     =	["""image_embeds""", """negative_image_embeds""", """image"""]
    UpperCAmelCase	:							Dict     =	[
        """image_embeds""",
        """negative_image_embeds""",
        """image""",
    ]
    UpperCAmelCase	:							Dict     =	[
        """generator""",
        """height""",
        """width""",
        """strength""",
        """guidance_scale""",
        """num_inference_steps""",
        """return_dict""",
        """guidance_scale""",
        """num_images_per_prompt""",
        """output_type""",
        """return_dict""",
    ]
    UpperCAmelCase	:							List[Any]     =	False






    @property
    def      __snake_case							( self    :   Tuple):
      return 32






    @property
    def      __snake_case							( self    :   Union[str, Any]):
      return 32






    @property
    def      __snake_case							( self    :   List[Any]):
      return self.time_input_dim






    @property
    def      __snake_case							( self    :   List[str]):
      return self.time_input_dim * 4






    @property
    def      __snake_case							( self    :   List[str]):
      return 100






    @property
    def      __snake_case							( self    :   Any):
      torch.manual_seed(0)

      a						:		Optional[int]											=      {
          "in_channels": 4,
          # Out channels is double in channels because predicts mean and variance
          "out_channels": 8,
          "addition_embed_type": "image",
          "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
          "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
          "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
          "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
          "layers_per_block": 1,
          "encoder_hid_dim": self.text_embedder_hidden_size,
          "encoder_hid_dim_type": "image_proj",
          "cross_attention_dim": self.cross_attention_dim,
          "attention_head_dim": 4,
          "resnet_time_scale_shift": "scale_shift",
          "class_embed_type": None,
      }

      a						:		str											=      UNetaDConditionModel(**__UpperCAmelCase)
      return model






    @property
    def      __snake_case							( self    :   Union[str, Any]):
      return {
          "block_out_channels": [32, 64],
          "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
          "in_channels": 3,
          "latent_channels": 4,
          "layers_per_block": 1,
          "norm_num_groups": 8,
          "norm_type": "spatial",
          "num_vq_embeddings": 12,
          "out_channels": 3,
          "up_block_types": [
              "AttnUpDecoderBlock2D",
              "UpDecoderBlock2D",
          ],
          "vq_embed_dim": 4,
      }






    @property
    def      __snake_case							( self    :   Dict):
      torch.manual_seed(0)
      a						:		Optional[int]											=      VQModel(**self.dummy_movq_kwargs)
      return model






    def      __snake_case							( self    :   Union[str, Any]):
      a						:		Union[str, Any]											=      self.dummy_unet
      a						:		int											=      self.dummy_movq

      a						:		Dict											=      {
          "num_train_timesteps": 1000,
          "beta_schedule": "linear",
          "beta_start": 0.00_085,
          "beta_end": 0.012,
          "clip_sample": False,
          "set_alpha_to_one": False,
          "steps_offset": 0,
          "prediction_type": "epsilon",
          "thresholding": False,
      }

      a						:		List[str]											=      DDIMScheduler(**__UpperCAmelCase)

      a						:		int											=      {
          "unet": unet,
          "scheduler": scheduler,
          "movq": movq,
      }

      return components






    def      __snake_case							( self    :   Tuple    ,			__UpperCAmelCase    :   Any    ,			__UpperCAmelCase    :   List[str]=0):
      a						:		Union[str, Any]											=      floats_tensor((1, self.text_embedder_hidden_size)    ,			rng=random.Random(__UpperCAmelCase)).to(__UpperCAmelCase)
      a						:		List[str]											=      floats_tensor((1, self.text_embedder_hidden_size)    ,			rng=random.Random(seed + 1)).to(
          __UpperCAmelCase)
      # create init_image
      a						:		Union[str, Any]											=      floats_tensor((1, 3, 64, 64)    ,			rng=random.Random(__UpperCAmelCase)).to(__UpperCAmelCase)
      a						:		Tuple											=      image.cpu().permute(0    ,			2    ,			3    ,			1)[0]
      a						:		Tuple											=      Image.fromarray(np.uinta(__UpperCAmelCase)).convert("RGB").resize((256, 256))

      if str(__UpperCAmelCase).startswith("mps"):
        a						:		Dict											=      torch.manual_seed(__UpperCAmelCase)
      else:
        a						:		Any											=      torch.Generator(device=__UpperCAmelCase).manual_seed(__UpperCAmelCase)
      a						:		Dict											=      {
          "image": init_image,
          "image_embeds": image_embeds,
          "negative_image_embeds": negative_image_embeds,
          "generator": generator,
          "height": 64,
          "width": 64,
          "num_inference_steps": 10,
          "guidance_scale": 7.0,
          "strength": 0.2,
          "output_type": "np",
      }
      return inputs






    def      __snake_case							( self    :   Tuple):
      a						:		int											=      "cpu"

      a						:		Union[str, Any]											=      self.get_dummy_components()

      a						:		List[str]											=      self.pipeline_class(**__UpperCAmelCase)
      a						:		Optional[Any]											=      pipe.to(__UpperCAmelCase)

      pipe.set_progress_bar_config(disable=__UpperCAmelCase)

      a						:		Optional[int]											=      pipe(**self.get_dummy_inputs(__UpperCAmelCase))
      a						:		List[str]											=      output.images

      a						:		List[Any]											=      pipe(
          **self.get_dummy_inputs(__UpperCAmelCase)    ,			return_dict=__UpperCAmelCase    ,			)[0]

      a						:		Optional[int]											=      image[0, -3:, -3:, -1]
      a						:		str											=      image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)

      a						:		List[str]											=      np.array(
          [0.6_199_778, 0.63_984_406, 0.46_145_785, 0.62_944_984, 0.5_622_215, 0.47_306_132, 0.47_441_456, 0.4_607_606, 0.48_719_263])
      assert (
          np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
      ), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
      assert (
          np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2
      ), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''






@slow
@require_torch_gpu
class 						_A  (  unittest.TestCase  ):







    """simple docstring"""




    def      __snake_case							( self    :   str):
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()






    def      __snake_case							( self    :   List[str]):
      a						:		Optional[int]											=      load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
          "/kandinskyv22/kandinskyv22_img2img_frog.npy")

      a						:		Union[str, Any]											=      load_image(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png")
      a						:		Tuple											=      "A red cartoon frog, 4k"

      a						:		List[Any]											=      KandinskyVaaPriorPipeline.from_pretrained(
          "kandinsky-community/kandinsky-2-2-prior"    ,			torch_dtype=torch.floataa)
      pipe_prior.to(__UpperCAmelCase)

      a						:		Dict											=      KandinskyVaaImgaImgPipeline.from_pretrained(
          "kandinsky-community/kandinsky-2-2-decoder"    ,			torch_dtype=torch.floataa)
      a						:		Any											=      pipeline.to(__UpperCAmelCase)

      pipeline.set_progress_bar_config(disable=__UpperCAmelCase)

      a						:		str											=      torch.Generator(device="cpu").manual_seed(0)
      a							,		a						:		Tuple											=      pipe_prior(
          __UpperCAmelCase    ,			generator=__UpperCAmelCase    ,			num_inference_steps=5    ,			negative_prompt=""    ,			).to_tuple()

      a						:		Any											=      pipeline(
          image=__UpperCAmelCase    ,			image_embeds=__UpperCAmelCase    ,			negative_image_embeds=__UpperCAmelCase    ,			generator=__UpperCAmelCase    ,			num_inference_steps=100    ,			height=768    ,			width=768    ,			strength=0.2    ,			output_type="np"    ,			)

      a						:		Optional[Any]											=      output.images[0]

      assert image.shape == (768, 768, 3)

      assert_mean_pixel_difference(__UpperCAmelCase    ,			__UpperCAmelCase)










"""simple docstring"""







import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextConfig, SegformerImageProcessor, UperNetConfig, UperNetForSemanticSegmentation


def 					lowercase      (    A_	)->				Optional[Any]:
  '''simple docstring'''







  a						:		Optional[Any]											=      384
  if "tiny" in model_name:
    a						:		List[str]											=      [3, 3, 9, 3]
    a						:		Optional[Any]											=      [96, 192, 384, 768]
  if "small" in model_name:
    a						:		Tuple											=      [3, 3, 27, 3]
    a						:		str											=      [96, 192, 384, 768]
  if "base" in model_name:
    a						:		Union[str, Any]											=      [3, 3, 27, 3]
    a						:		Dict											=      [128, 256, 512, 1_024]
    a						:		Any											=      512
  if "large" in model_name:
    a						:		Optional[Any]											=      [3, 3, 27, 3]
    a						:		str											=      [192, 384, 768, 1_536]
    a						:		Dict											=      768
  if "xlarge" in model_name:
    a						:		str											=      [3, 3, 27, 3]
    a						:		List[Any]											=      [256, 512, 1_024, 2_048]
    a						:		List[Any]											=      1_024

  # set label information
  a						:		int											=      150
  a						:		str											=      "huggingface/label-files"
  a						:		Tuple											=      "ade20k-id2label.json"
  a						:		Dict											=      json.load(open(hf_hub_download(A_  ,       A_  ,       repo_type="dataset"	)  ,       "r"	)	)
  a						:		int											=      {int(A_	): v for k, v in idalabel.items()}
  a						:		List[Any]											=      {v: k for k, v in idalabel.items()}

  a						:		Optional[int]											=      ConvNextConfig(
      depths=A_  ,       hidden_sizes=A_  ,       out_features=["stage1", "stage2", "stage3", "stage4"]	)
  a						:		Tuple											=      UperNetConfig(
      backbone_config=A_  ,       auxiliary_in_channels=A_  ,       num_labels=A_  ,       idalabel=A_  ,       labelaid=A_  ,       )

  return config


def 					lowercase      (    A_	)->				Tuple:
  '''simple docstring'''







  a						:		int											=      []

  # fmt: off
  # stem
  rename_keys.append(("backbone.downsample_layers.0.0.weight", "backbone.embeddings.patch_embeddings.weight")	)
  rename_keys.append(("backbone.downsample_layers.0.0.bias", "backbone.embeddings.patch_embeddings.bias")	)
  rename_keys.append(("backbone.downsample_layers.0.1.weight", "backbone.embeddings.layernorm.weight")	)
  rename_keys.append(("backbone.downsample_layers.0.1.bias", "backbone.embeddings.layernorm.bias")	)
  # stages
  for i in range(len(config.backbone_config.depths	)	):
    for j in range(config.backbone_config.depths[i]	):
      rename_keys.append((F'''backbone.stages.{i}.{j}.gamma''', F'''backbone.encoder.stages.{i}.layers.{j}.layer_scale_parameter''')	)
      rename_keys.append((F'''backbone.stages.{i}.{j}.depthwise_conv.weight''', F'''backbone.encoder.stages.{i}.layers.{j}.dwconv.weight''')	)
      rename_keys.append((F'''backbone.stages.{i}.{j}.depthwise_conv.bias''', F'''backbone.encoder.stages.{i}.layers.{j}.dwconv.bias''')	)
      rename_keys.append((F'''backbone.stages.{i}.{j}.norm.weight''', F'''backbone.encoder.stages.{i}.layers.{j}.layernorm.weight''')	)
      rename_keys.append((F'''backbone.stages.{i}.{j}.norm.bias''', F'''backbone.encoder.stages.{i}.layers.{j}.layernorm.bias''')	)
      rename_keys.append((F'''backbone.stages.{i}.{j}.pointwise_conv1.weight''', F'''backbone.encoder.stages.{i}.layers.{j}.pwconv1.weight''')	)
      rename_keys.append((F'''backbone.stages.{i}.{j}.pointwise_conv1.bias''', F'''backbone.encoder.stages.{i}.layers.{j}.pwconv1.bias''')	)
      rename_keys.append((F'''backbone.stages.{i}.{j}.pointwise_conv2.weight''', F'''backbone.encoder.stages.{i}.layers.{j}.pwconv2.weight''')	)
      rename_keys.append((F'''backbone.stages.{i}.{j}.pointwise_conv2.bias''', F'''backbone.encoder.stages.{i}.layers.{j}.pwconv2.bias''')	)
    if i > 0:
      rename_keys.append((F'''backbone.downsample_layers.{i}.0.weight''', F'''backbone.encoder.stages.{i}.downsampling_layer.0.weight''')	)
      rename_keys.append((F'''backbone.downsample_layers.{i}.0.bias''', F'''backbone.encoder.stages.{i}.downsampling_layer.0.bias''')	)
      rename_keys.append((F'''backbone.downsample_layers.{i}.1.weight''', F'''backbone.encoder.stages.{i}.downsampling_layer.1.weight''')	)
      rename_keys.append((F'''backbone.downsample_layers.{i}.1.bias''', F'''backbone.encoder.stages.{i}.downsampling_layer.1.bias''')	)

    rename_keys.append((F'''backbone.norm{i}.weight''', F'''backbone.hidden_states_norms.stage{i+1}.weight''')	)
    rename_keys.append((F'''backbone.norm{i}.bias''', F'''backbone.hidden_states_norms.stage{i+1}.bias''')	)

  # decode head
  rename_keys.extend(
      [
          ("decode_head.conv_seg.weight", "decode_head.classifier.weight"),
          ("decode_head.conv_seg.bias", "decode_head.classifier.bias"),
          ("auxiliary_head.conv_seg.weight", "auxiliary_head.classifier.weight"),
          ("auxiliary_head.conv_seg.bias", "auxiliary_head.classifier.bias"),
      ]	)
  # fmt: on

  return rename_keys


def 					lowercase      (    A_  ,       A_  ,       A_	)->				str:
  '''simple docstring'''







  a						:		str											=      dct.pop(A_	)
  a						:		str											=      val


def 					lowercase      (    A_  ,       A_  ,       A_	)->				int:
  '''simple docstring'''







  a						:		Any											=      {
      "upernet-convnext-tiny": "https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_tiny_fp16_512x512_160k_ade20k/upernet_convnext_tiny_fp16_512x512_160k_ade20k_20220227_124553-cad485de.pth",
      "upernet-convnext-small": "https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_small_fp16_512x512_160k_ade20k/upernet_convnext_small_fp16_512x512_160k_ade20k_20220227_131208-1b1e394f.pth",
      "upernet-convnext-base": "https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_base_fp16_512x512_160k_ade20k/upernet_convnext_base_fp16_512x512_160k_ade20k_20220227_181227-02a24fc6.pth",
      "upernet-convnext-large": "https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_large_fp16_640x640_160k_ade20k/upernet_convnext_large_fp16_640x640_160k_ade20k_20220226_040532-e57aa54d.pth",
      "upernet-convnext-xlarge": "https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_xlarge_fp16_640x640_160k_ade20k/upernet_convnext_xlarge_fp16_640x640_160k_ade20k_20220226_080344-95fc38c2.pth",
  }
  a						:		Tuple											=      model_name_to_url[model_name]
  a						:		int											=      torch.hub.load_state_dict_from_url(A_  ,       map_location="cpu"	)["state_dict"]

  a						:		Optional[Any]											=      get_upernet_config(A_	)
  a						:		int											=      UperNetForSemanticSegmentation(A_	)
  model.eval()

  # replace "bn" => "batch_norm"
  for key in state_dict.copy().keys():
    a						:		Tuple											=      state_dict.pop(A_	)
    if "bn" in key:
      a						:		Optional[int]											=      key.replace("bn"  ,       "batch_norm"	)
    a						:		Any											=      val

  # rename keys
  a						:		Any											=      create_rename_keys(A_	)
  for src, dest in rename_keys:
    rename_key(A_  ,       A_  ,       A_	)

  model.load_state_dict(A_	)

  # verify on image
  a						:		Dict											=      "https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"
  a						:		Any											=      Image.open(requests.get(A_  ,       stream=A_	).raw	).convert("RGB"	)

  a						:		Union[str, Any]											=      SegformerImageProcessor()
  a						:		Dict											=      processor(A_  ,       return_tensors="pt"	).pixel_values

  with torch.no_grad():
    a						:		Any											=      model(A_	)

  if model_name == "upernet-convnext-tiny":
    a						:		List[str]											=      torch.tensor(
        [[-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.7_7_4_6, -8.7_7_4_6, -8.6_1_3_0]]	)
  elif model_name == "upernet-convnext-small":
    a						:		int											=      torch.tensor(
        [[-8.8_2_3_6, -8.8_2_3_6, -8.6_7_7_1], [-8.8_2_3_6, -8.8_2_3_6, -8.6_7_7_1], [-8.7_6_3_8, -8.7_6_3_8, -8.6_2_4_0]]	)
  elif model_name == "upernet-convnext-base":
    a						:		Union[str, Any]											=      torch.tensor(
        [[-8.8_5_5_8, -8.8_5_5_8, -8.6_9_0_5], [-8.8_5_5_8, -8.8_5_5_8, -8.6_9_0_5], [-8.7_6_6_9, -8.7_6_6_9, -8.6_0_2_1]]	)
  elif model_name == "upernet-convnext-large":
    a						:		Union[str, Any]											=      torch.tensor(
        [[-8.6_6_6_0, -8.6_6_6_0, -8.6_2_1_0], [-8.6_6_6_0, -8.6_6_6_0, -8.6_2_1_0], [-8.6_3_1_0, -8.6_3_1_0, -8.5_9_6_4]]	)
  elif model_name == "upernet-convnext-xlarge":
    a						:		str											=      torch.tensor(
        [[-8.4_9_8_0, -8.4_9_8_0, -8.3_9_7_7], [-8.4_9_8_0, -8.4_9_8_0, -8.3_9_7_7], [-8.4_3_7_9, -8.4_3_7_9, -8.3_4_1_2]]	)
  print("Logits:"  ,       outputs.logits[0, 0, :3, :3]	)
  assert torch.allclose(outputs.logits[0, 0, :3, :3]  ,       A_  ,       atol=1e-4	)
  print("Looks ok!"	)

  if pytorch_dump_folder_path is not None:
    print(F'''Saving model {model_name} to {pytorch_dump_folder_path}'''	)
    model.save_pretrained(A_	)
    print(F'''Saving processor to {pytorch_dump_folder_path}'''	)
    processor.save_pretrained(A_	)

  if push_to_hub:
    print(F'''Pushing model and processor for {model_name} to hub'''	)
    model.push_to_hub(F'''openmmlab/{model_name}'''	)
    processor.push_to_hub(F'''openmmlab/{model_name}'''	)


if __name__ == "__main__":
 __lowercase    =				argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     """--model_name""",
     default="""upernet-convnext-tiny""",
     type=str,
     choices=[f'''upernet-convnext-{size}''' for size in ["""tiny""", """small""", """base""", """large""", """xlarge"""]],
     help="""Name of the ConvNext UperNet model you'd like to convert.""",
 )
 parser.add_argument(
     """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
 )
 parser.add_argument(
     """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
 )

 __lowercase    =				parser.parse_args()
 convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






