

'''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       UpperCAmelCase__    (      UpperCAmelCase__				)			->			Union[str, Any]:
       A_						=      3_84
       if "tiny" in model_name:
              A_						=      [3, 3, 9, 3]
              A_						=      [96, 1_92, 3_84, 7_68]
       if "small" in model_name:
              A_						=      [3, 3, 27, 3]
              A_						=      [96, 1_92, 3_84, 7_68]
       if "base" in model_name:
              A_						=      [3, 3, 27, 3]
              A_						=      [1_28, 2_56, 5_12, 10_24]
              A_						=      5_12
       if "large" in model_name:
              A_						=      [3, 3, 27, 3]
              A_						=      [1_92, 3_84, 7_68, 15_36]
              A_						=      7_68
       if "xlarge" in model_name:
              A_						=      [3, 3, 27, 3]
              A_						=      [2_56, 5_12, 10_24, 20_48]
              A_						=      10_24

       # set label information
       A_						=      1_50
       A_						=      "huggingface/label-files"
       A_						=      "ade20k-id2label.json"
       A_						=      json.load(open(hf_hub_download(_A,       _A,       repo_type="""dataset"""				),       """r"""				)				)
       A_						=      {int(_A				): v for k, v in idalabel.items()}
       A_						=      {v: k for k, v in idalabel.items()}

       A_						=      ConvNextConfig(
           depths=_A,       hidden_sizes=_A,       out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]				)
       A_						=      UperNetConfig(
           backbone_config=_A,       auxiliary_in_channels=_A,       num_labels=_A,       idalabel=_A,       labelaid=_A,       )

       return config
def       UpperCAmelCase__    (      UpperCAmelCase__				)			->			Optional[Any]:
       A_						=      []

       # 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       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__				)			->			Dict:
       A_						=      dct.pop(_A				)
       A_						=      val
def       UpperCAmelCase__    (      UpperCAmelCase__,       UpperCAmelCase__,       UpperCAmelCase__				)			->			Any:
       A_						=      {
           "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_						=      model_name_to_url[model_name]
       A_						=      torch.hub.load_state_dict_from_url(_A,       map_location="""cpu"""				)["state_dict"]

       A_						=      get_upernet_config(_A				)
       A_						=      UperNetForSemanticSegmentation(_A				)
       model.eval()

       # replace "bn" => "batch_norm"
       for key in state_dict.copy().keys():
              A_						=      state_dict.pop(_A				)
              if "bn" in key:
                     A_						=      key.replace("""bn""",       """batch_norm"""				)
              A_						=      val

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

       model.load_state_dict(_A				)

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

       A_						=      SegformerImageProcessor()
       A_						=      processor(_A,       return_tensors="""pt"""				).pixel_values

       with torch.no_grad():
              A_						=      model(_A				)

       if model_name == "upernet-convnext-tiny":
              A_						=      torch.tensor(
                  [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]				)
       elif model_name == "upernet-convnext-small":
              A_						=      torch.tensor(
                  [[-8.8_236, -8.8_236, -8.6_771], [-8.8_236, -8.8_236, -8.6_771], [-8.7_638, -8.7_638, -8.6_240]]				)
       elif model_name == "upernet-convnext-base":
              A_						=      torch.tensor(
                  [[-8.8_558, -8.8_558, -8.6_905], [-8.8_558, -8.8_558, -8.6_905], [-8.7_669, -8.7_669, -8.6_021]]				)
       elif model_name == "upernet-convnext-large":
              A_						=      torch.tensor(
                  [[-8.6_660, -8.6_660, -8.6_210], [-8.6_660, -8.6_660, -8.6_210], [-8.6_310, -8.6_310, -8.5_964]]				)
       elif model_name == "upernet-convnext-xlarge":
              A_						=      torch.tensor(
                  [[-8.4_980, -8.4_980, -8.3_977], [-8.4_980, -8.4_980, -8.3_977], [-8.4_379, -8.4_379, -8.3_412]]				)
       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__":
      __lowerCamelCase		       =    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.'''
      )

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



'''simple docstring'''







import unittest

import numpy as np
import torch
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class 					A__      (      _snake_case						,						unittest.TestCase			):
    lowercase						=		KandinskyVaaPriorPipeline
    lowercase						=		["prompt"]
    lowercase						=		["prompt", "negative_prompt"]
    lowercase						=		[
        "num_images_per_prompt",
        "generator",
        "num_inference_steps",
        "latents",
        "negative_prompt",
        "guidance_scale",
        "output_type",
        "return_dict",
    ]
    lowercase						=		False



    @property
    def 					snake_case_			(					self       )						->      Any:





           '''simple docstring'''
           return 32



    @property
    def 					snake_case_			(					self       )						->      Optional[int]:





           '''simple docstring'''
           return 32



    @property
    def 					snake_case_			(					self       )						->      Optional[Any]:





           '''simple docstring'''
           return self.time_input_dim



    @property
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           return self.time_input_dim * 4



    @property
    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           return 100



    @property
    def 					snake_case_			(					self       )						->      Optional[Any]:





           '''simple docstring'''
           A_						=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )
           return tokenizer



    @property
    def 					snake_case_			(					self       )						->      Union[str, Any]:





           '''simple docstring'''
           torch.manual_seed(0       )
           A_						=      CLIPTextConfig(
               bos_token_id=0		,				eos_token_id=2		,				hidden_size=self.text_embedder_hidden_size		,				projection_dim=self.text_embedder_hidden_size		,				intermediate_size=37		,				layer_norm_eps=1e-0_5		,				num_attention_heads=4		,				num_hidden_layers=5		,				pad_token_id=1		,				vocab_size=1000		,				)
           return CLIPTextModelWithProjection(UpperCamelCase__       )



    @property
    def 					snake_case_			(					self       )						->      Tuple:





           '''simple docstring'''
           torch.manual_seed(0       )

           A_						=      {
               """num_attention_heads""": 2,
               """attention_head_dim""": 12,
               """embedding_dim""": self.text_embedder_hidden_size,
               """num_layers""": 1,
           }

           A_						=      PriorTransformer(**UpperCamelCase__       )
           # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
           A_						=      nn.Parameter(torch.ones(model.clip_std.shape       )       )
           return model



    @property
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           torch.manual_seed(0       )
           A_						=      CLIPVisionConfig(
               hidden_size=self.text_embedder_hidden_size		,				image_size=224		,				projection_dim=self.text_embedder_hidden_size		,				intermediate_size=37		,				num_attention_heads=4		,				num_channels=3		,				num_hidden_layers=5		,				patch_size=14		,				)

           A_						=      CLIPVisionModelWithProjection(UpperCamelCase__       )
           return model



    @property
    def 					snake_case_			(					self       )						->      str:





           '''simple docstring'''
           A_						=      CLIPImageProcessor(
               crop_size=224		,				do_center_crop=UpperCamelCase__		,				do_normalize=UpperCamelCase__		,				do_resize=UpperCamelCase__		,				image_mean=[0.48145466, 0.4578275, 0.40821073]		,				image_std=[0.26862954, 0.26130258, 0.27577711]		,				resample=3		,				size=224		,				)

           return image_processor



    def 					snake_case_			(					self       )						->      Union[str, Any]:





           '''simple docstring'''
           A_						=      self.dummy_prior
           A_						=      self.dummy_image_encoder
           A_						=      self.dummy_text_encoder
           A_						=      self.dummy_tokenizer
           A_						=      self.dummy_image_processor

           A_						=      UnCLIPScheduler(
               variance_type="""fixed_small_log"""		,				prediction_type="""sample"""		,				num_train_timesteps=1000		,				clip_sample=UpperCamelCase__		,				clip_sample_range=10.0		,				)

           A_						=      {
               """prior""": prior,
               """image_encoder""": image_encoder,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
               """scheduler""": scheduler,
               """image_processor""": image_processor,
           }

           return components



    def 					snake_case_			(					self		,				UpperCamelCase__		,				UpperCamelCase__=0       )						->      Optional[int]:





           '''simple docstring'''
           if str(UpperCamelCase__       ).startswith("""mps"""       ):
                  A_						=      torch.manual_seed(UpperCamelCase__       )
           else:
                  A_						=      torch.Generator(device=UpperCamelCase__       ).manual_seed(UpperCamelCase__       )
           A_						=      {
               """prompt""": """horse""",
               """generator""": generator,
               """guidance_scale""": 4.0,
               """num_inference_steps""": 2,
               """output_type""": """np""",
           }
           return inputs



    def 					snake_case_			(					self       )						->      List[str]:





           '''simple docstring'''
           A_						=      """cpu"""

           A_						=      self.get_dummy_components()

           A_						=      self.pipeline_class(**UpperCamelCase__       )
           A_						=      pipe.to(UpperCamelCase__       )

           pipe.set_progress_bar_config(disable=UpperCamelCase__       )

           A_						=      pipe(**self.get_dummy_inputs(UpperCamelCase__       )       )
           A_						=      output.image_embeds

           A_						=      pipe(
               **self.get_dummy_inputs(UpperCamelCase__       )		,				return_dict=UpperCamelCase__		,				)[0]

           A_						=      image[0, -10:]
           A_						=      image_from_tuple[0, -10:]

           assert image.shape == (1, 32)

           A_						=      np.array(
               [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]       )

           assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2
           assert np.abs(image_from_tuple_slice.flatten() - expected_slice       ).max() < 1e-2



    @skip_mps
    def 					snake_case_			(					self       )						->      int:





           '''simple docstring'''
           A_						=      torch_device == """cpu"""
           A_						=      True
           A_						=      False

           self._test_inference_batch_single_identical(
               test_max_difference=UpperCamelCase__		,				relax_max_difference=UpperCamelCase__		,				test_mean_pixel_difference=UpperCamelCase__		,				)







    @skip_mps
    def 					snake_case_			(					self       )						->      Union[str, Any]:





           '''simple docstring'''
           A_						=      torch_device == """cpu"""
           A_						=      False

           self._test_attention_slicing_forward_pass(
               test_max_difference=UpperCamelCase__		,				test_mean_pixel_difference=UpperCamelCase__		,				)
