





'''simple docstring'''






import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 					UpperCAmelCase	(     a__				,						unittest.TestCase  ):
      _lowercase:     List[str]			=  AudioLDMPipeline
      _lowercase:     str			=  TEXT_TO_AUDIO_PARAMS
      _lowercase:     str			=  TEXT_TO_AUDIO_BATCH_PARAMS
      _lowercase:     Any			=  frozenset(
          [
              '''num_inference_steps''',
              '''num_waveforms_per_prompt''',
              '''generator''',
              '''latents''',
              '''output_type''',
              '''return_dict''',
              '''callback''',
              '''callback_steps''',
          ]  )




      def 							lowercase__       (     self     :					Tuple    )							-> Any:
            torch.manual_seed(0    )
            _lowerCAmelCase					=				UNetaDConditionModel(
                block_out_channels=(32, 64)     ,     layers_per_block=2     ,     sample_size=32     ,     in_channels=4     ,     out_channels=4     ,     down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")     ,     up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")     ,     cross_attention_dim=(32, 64)     ,     class_embed_type="""simple_projection"""     ,     projection_class_embeddings_input_dim=32     ,     class_embeddings_concat=__a     ,     )
            _lowerCAmelCase					=				DDIMScheduler(
                beta_start=0.0_00_85     ,     beta_end=0.0_12     ,     beta_schedule="""scaled_linear"""     ,     clip_sample=__a     ,     set_alpha_to_one=__a     ,     )
            torch.manual_seed(0    )
            _lowerCAmelCase					=				AutoencoderKL(
                block_out_channels=[32, 64]     ,     in_channels=1     ,     out_channels=1     ,     down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]     ,     up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]     ,     latent_channels=4     ,     )
            torch.manual_seed(0    )
            _lowerCAmelCase					=				ClapTextConfig(
                bos_token_id=0     ,     eos_token_id=2     ,     hidden_size=32     ,     intermediate_size=37     ,     layer_norm_eps=1E-0_5     ,     num_attention_heads=4     ,     num_hidden_layers=5     ,     pad_token_id=1     ,     vocab_size=10_00     ,     projection_dim=32     ,     )
            _lowerCAmelCase					=				ClapTextModelWithProjection(__a    )
            _lowerCAmelCase					=				RobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-random-roberta"""     ,     model_max_length=77    )

            _lowerCAmelCase					=				SpeechTaHifiGanConfig(
                model_in_dim=8     ,     sampling_rate=1_60_00     ,     upsample_initial_channel=16     ,     upsample_rates=[2, 2]     ,     upsample_kernel_sizes=[4, 4]     ,     resblock_kernel_sizes=[3, 7]     ,     resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]     ,     normalize_before=__a     ,     )

            _lowerCAmelCase					=				SpeechTaHifiGan(__a    )

            _lowerCAmelCase					=				{
                """unet""": unet,
                """scheduler""": scheduler,
                """vae""": vae,
                """text_encoder""": text_encoder,
                """tokenizer""": tokenizer,
                """vocoder""": vocoder,
            }
            return components




      def 							lowercase__       (     self     :					Tuple     ,     __snake_case     :					Tuple     ,     __snake_case     :					Optional[Any]=0    )							-> List[Any]:
            if str(__a    ).startswith("""mps"""    ):
                  _lowerCAmelCase					=				torch.manual_seed(__a    )
            else:
                  _lowerCAmelCase					=				torch.Generator(device=__a    ).manual_seed(__a    )
            _lowerCAmelCase					=				{
                """prompt""": """A hammer hitting a wooden surface""",
                """generator""": generator,
                """num_inference_steps""": 2,
                """guidance_scale""": 6.0,
            }
            return inputs




      def 							lowercase__       (     self     :					int    )							-> Union[str, Any]:
            _lowerCAmelCase					=				"""cpu"""  # ensure determinism for the device-dependent torch.Generator

            _lowerCAmelCase					=				self.get_dummy_components()
            _lowerCAmelCase					=				AudioLDMPipeline(**__a    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            audioldm_pipe.set_progress_bar_config(disable=__a    )

            _lowerCAmelCase					=				self.get_dummy_inputs(__a    )
            _lowerCAmelCase					=				audioldm_pipe(**__a    )
            _lowerCAmelCase					=				output.audios[0]

            assert audio.ndim == 1
            assert len(__a    ) == 2_56

            _lowerCAmelCase					=				audio[:10]
            _lowerCAmelCase					=				np.array(
                [-0.00_50, 0.00_50, -0.00_60, 0.00_33, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_33]    )

            assert np.abs(audio_slice - expected_slice    ).max() < 1E-2




      def 							lowercase__       (     self     :					Union[str, Any]    )							-> Union[str, Any]:
            _lowerCAmelCase					=				self.get_dummy_components()
            _lowerCAmelCase					=				AudioLDMPipeline(**__a    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            audioldm_pipe.set_progress_bar_config(disable=__a    )

            _lowerCAmelCase					=				self.get_dummy_inputs(__a    )
            _lowerCAmelCase					=				3 * [inputs["""prompt"""]]

            # forward
            _lowerCAmelCase					=				audioldm_pipe(**__a    )
            _lowerCAmelCase					=				output.audios[0]

            _lowerCAmelCase					=				self.get_dummy_inputs(__a    )
            _lowerCAmelCase					=				3 * [inputs.pop("""prompt"""    )]

            _lowerCAmelCase					=				audioldm_pipe.tokenizer(
                __a     ,     padding="""max_length"""     ,     max_length=audioldm_pipe.tokenizer.model_max_length     ,     truncation=__a     ,     return_tensors="""pt"""     ,     )
            _lowerCAmelCase					=				text_inputs["""input_ids"""].to(__a    )

            _lowerCAmelCase					=				audioldm_pipe.text_encoder(
                __a     ,     )
            _lowerCAmelCase					=				prompt_embeds.text_embeds
            # additional L_2 normalization over each hidden-state
            _lowerCAmelCase					=				F.normalize(__a     ,     dim=-1    )

            _lowerCAmelCase					=				prompt_embeds

            # forward
            _lowerCAmelCase					=				audioldm_pipe(**__a    )
            _lowerCAmelCase					=				output.audios[0]

            assert np.abs(audio_a - audio_a    ).max() < 1E-2




      def 							lowercase__       (     self     :					Tuple    )							-> Dict:
            _lowerCAmelCase					=				self.get_dummy_components()
            _lowerCAmelCase					=				AudioLDMPipeline(**__a    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            audioldm_pipe.set_progress_bar_config(disable=__a    )

            _lowerCAmelCase					=				self.get_dummy_inputs(__a    )
            _lowerCAmelCase					=				3 * ["""this is a negative prompt"""]
            _lowerCAmelCase					=				negative_prompt
            _lowerCAmelCase					=				3 * [inputs["""prompt"""]]

            # forward
            _lowerCAmelCase					=				audioldm_pipe(**__a    )
            _lowerCAmelCase					=				output.audios[0]

            _lowerCAmelCase					=				self.get_dummy_inputs(__a    )
            _lowerCAmelCase					=				3 * [inputs.pop("""prompt"""    )]

            _lowerCAmelCase					=				[]
            for p in [prompt, negative_prompt]:
                  _lowerCAmelCase					=				audioldm_pipe.tokenizer(
                      __a     ,     padding="""max_length"""     ,     max_length=audioldm_pipe.tokenizer.model_max_length     ,     truncation=__a     ,     return_tensors="""pt"""     ,     )
                  _lowerCAmelCase					=				text_inputs["""input_ids"""].to(__a    )

                  _lowerCAmelCase					=				audioldm_pipe.text_encoder(
                      __a     ,     )
                  _lowerCAmelCase					=				text_embeds.text_embeds
                  # additional L_2 normalization over each hidden-state
                  _lowerCAmelCase					=				F.normalize(__a     ,     dim=-1    )

                  embeds.append(__a    )

            _lowerCAmelCase					=				embeds

            # forward
            _lowerCAmelCase					=				audioldm_pipe(**__a    )
            _lowerCAmelCase					=				output.audios[0]

            assert np.abs(audio_a - audio_a    ).max() < 1E-2




      def 							lowercase__       (     self     :					Optional[Any]    )							-> str:
            _lowerCAmelCase					=				"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            _lowerCAmelCase					=				self.get_dummy_components()
            _lowerCAmelCase					=				PNDMScheduler(skip_prk_steps=__a    )
            _lowerCAmelCase					=				AudioLDMPipeline(**__a    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            audioldm_pipe.set_progress_bar_config(disable=__a    )

            _lowerCAmelCase					=				self.get_dummy_inputs(__a    )
            _lowerCAmelCase					=				"""egg cracking"""
            _lowerCAmelCase					=				audioldm_pipe(**__a     ,     negative_prompt=__a    )
            _lowerCAmelCase					=				output.audios[0]

            assert audio.ndim == 1
            assert len(__a    ) == 2_56

            _lowerCAmelCase					=				audio[:10]
            _lowerCAmelCase					=				np.array(
                [-0.00_51, 0.00_50, -0.00_60, 0.00_34, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_32]    )

            assert np.abs(audio_slice - expected_slice    ).max() < 1E-2




      def 							lowercase__       (     self     :					Union[str, Any]    )							-> Dict:
            _lowerCAmelCase					=				"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            _lowerCAmelCase					=				self.get_dummy_components()
            _lowerCAmelCase					=				PNDMScheduler(skip_prk_steps=__a    )
            _lowerCAmelCase					=				AudioLDMPipeline(**__a    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            audioldm_pipe.set_progress_bar_config(disable=__a    )

            _lowerCAmelCase					=				"""A hammer hitting a wooden surface"""

            # test num_waveforms_per_prompt=1 (default)
            _lowerCAmelCase					=				audioldm_pipe(__a     ,     num_inference_steps=2    ).audios

            assert audios.shape == (1, 2_56)

            # test num_waveforms_per_prompt=1 (default) for batch of prompts
            _lowerCAmelCase					=				2
            _lowerCAmelCase					=				audioldm_pipe([prompt] * batch_size     ,     num_inference_steps=2    ).audios

            assert audios.shape == (batch_size, 2_56)

            # test num_waveforms_per_prompt for single prompt
            _lowerCAmelCase					=				2
            _lowerCAmelCase					=				audioldm_pipe(__a     ,     num_inference_steps=2     ,     num_waveforms_per_prompt=__a    ).audios

            assert audios.shape == (num_waveforms_per_prompt, 2_56)

            # test num_waveforms_per_prompt for batch of prompts
            _lowerCAmelCase					=				2
            _lowerCAmelCase					=				audioldm_pipe(
                [prompt] * batch_size     ,     num_inference_steps=2     ,     num_waveforms_per_prompt=__a    ).audios

            assert audios.shape == (batch_size * num_waveforms_per_prompt, 2_56)




      def 							lowercase__       (     self     :					Tuple    )							-> Optional[Any]:
            _lowerCAmelCase					=				"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            _lowerCAmelCase					=				self.get_dummy_components()
            _lowerCAmelCase					=				AudioLDMPipeline(**__a    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            audioldm_pipe.set_progress_bar_config(disable=__a    )
            _lowerCAmelCase					=				audioldm_pipe.vocoder.config.sampling_rate

            _lowerCAmelCase					=				self.get_dummy_inputs(__a    )
            _lowerCAmelCase					=				audioldm_pipe(audio_length_in_s=0.0_16     ,     **__a    )
            _lowerCAmelCase					=				output.audios[0]

            assert audio.ndim == 1
            assert len(__a    ) / vocoder_sampling_rate == 0.0_16

            _lowerCAmelCase					=				audioldm_pipe(audio_length_in_s=0.0_32     ,     **__a    )
            _lowerCAmelCase					=				output.audios[0]

            assert audio.ndim == 1
            assert len(__a    ) / vocoder_sampling_rate == 0.0_32




      def 							lowercase__       (     self     :					Dict    )							-> str:
            _lowerCAmelCase					=				self.get_dummy_components()
            _lowerCAmelCase					=				AudioLDMPipeline(**__a    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            audioldm_pipe.set_progress_bar_config(disable=__a    )

            _lowerCAmelCase					=				["""hey"""]

            _lowerCAmelCase					=				audioldm_pipe(__a     ,     num_inference_steps=1    )
            _lowerCAmelCase					=				output.audios.shape
            assert audio_shape == (1, 2_56)

            _lowerCAmelCase					=				audioldm_pipe.vocoder.config
            config.model_in_dim *= 2
            _lowerCAmelCase					=				SpeechTaHifiGan(__a    ).to(__a    )
            _lowerCAmelCase					=				audioldm_pipe(__a     ,     num_inference_steps=1    )
            _lowerCAmelCase					=				output.audios.shape
            # waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
            assert audio_shape == (1, 2_56)




      def 							lowercase__       (     self     :					Optional[int]    )							-> Optional[Any]:
            self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__a    )




      def 							lowercase__       (     self     :					str    )							-> Dict:
            self._test_inference_batch_single_identical(test_mean_pixel_difference=__a    )




      @unittest.skipIf(
          torch_device != """cuda""" or not is_xformers_available()     ,     reason="""XFormers attention is only available with CUDA and `xformers` installed"""     ,     )
      def 							lowercase__       (     self     :					Dict    )							-> Dict:
            self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__a    )




@slow
class 					UpperCAmelCase	(     unittest.TestCase  ):




      def 							lowercase__       (     self     :					List[Any]    )							-> Optional[Any]:
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      def 							lowercase__       (     self     :					List[str]     ,     __snake_case     :					Dict     ,     __snake_case     :					Optional[Any]="cpu"     ,     __snake_case     :					Tuple=torch.floataa     ,     __snake_case     :					Optional[int]=0    )							-> Dict:
            _lowerCAmelCase					=				torch.Generator(device=__a    ).manual_seed(__a    )
            _lowerCAmelCase					=				np.random.RandomState(__a    ).standard_normal((1, 8, 1_28, 16)    )
            _lowerCAmelCase					=				torch.from_numpy(__a    ).to(device=__a     ,     dtype=__a    )
            _lowerCAmelCase					=				{
                """prompt""": """A hammer hitting a wooden surface""",
                """latents""": latents,
                """generator""": generator,
                """num_inference_steps""": 3,
                """guidance_scale""": 2.5,
            }
            return inputs




      def 							lowercase__       (     self     :					Optional[int]    )							-> Tuple:
            _lowerCAmelCase					=				AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            audioldm_pipe.set_progress_bar_config(disable=__a    )

            _lowerCAmelCase					=				self.get_inputs(__a    )
            _lowerCAmelCase					=				25
            _lowerCAmelCase					=				audioldm_pipe(**__a    ).audios[0]

            assert audio.ndim == 1
            assert len(__a    ) == 8_19_20

            _lowerCAmelCase					=				audio[7_72_30:7_72_40]
            _lowerCAmelCase					=				np.array(
                [-0.48_84, -0.46_07, 0.00_23, 0.50_07, 0.58_96, 0.51_51, 0.38_13, -0.02_08, -0.36_87, -0.43_15]    )
            _lowerCAmelCase					=				np.abs(expected_slice - audio_slice    ).max()
            assert max_diff < 1E-2




      def 							lowercase__       (     self     :					Any    )							-> str:
            _lowerCAmelCase					=				AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""    )
            _lowerCAmelCase					=				LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config    )
            _lowerCAmelCase					=				audioldm_pipe.to(__a    )
            audioldm_pipe.set_progress_bar_config(disable=__a    )

            _lowerCAmelCase					=				self.get_inputs(__a    )
            _lowerCAmelCase					=				audioldm_pipe(**__a    ).audios[0]

            assert audio.ndim == 1
            assert len(__a    ) == 8_19_20

            _lowerCAmelCase					=				audio[2_77_80:2_77_90]
            _lowerCAmelCase					=				np.array([-0.21_31, -0.08_73, -0.01_24, -0.01_89, 0.05_69, 0.13_73, 0.18_83, 0.28_86, 0.32_97, 0.22_12]    )
            _lowerCAmelCase					=				np.abs(expected_slice - audio_slice    ).max()
            assert max_diff < 3E-2










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    __lowerCamelCase   (      A__						: int					)     ->    List[str]:
   lowerCamelCase_						:   int		= 384
   if "tiny" in model_name:
      lowerCamelCase_						:   str		= [3, 3, 9, 3]
      lowerCamelCase_						:   List[str]		= [96, 192, 384, 768]
   if "small" in model_name:
      lowerCamelCase_						:   Any		= [3, 3, 27, 3]
      lowerCamelCase_						:   List[Any]		= [96, 192, 384, 768]
   if "base" in model_name:
      lowerCamelCase_						:   Optional[Any]		= [3, 3, 27, 3]
      lowerCamelCase_						:   int		= [128, 256, 512, 1024]
      lowerCamelCase_						:   int		= 512
   if "large" in model_name:
      lowerCamelCase_						:   Union[str, Any]		= [3, 3, 27, 3]
      lowerCamelCase_						:   Dict		= [192, 384, 768, 1536]
      lowerCamelCase_						:   Union[str, Any]		= 768
   if "xlarge" in model_name:
      lowerCamelCase_						:   Union[str, Any]		= [3, 3, 27, 3]
      lowerCamelCase_						:   Optional[int]		= [256, 512, 1024, 2048]
      lowerCamelCase_						:   List[str]		= 1024

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

   lowerCamelCase_						:   Dict		= ConvNextConfig(
       depths=A__						,						hidden_sizes=A__						,						out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]					)
   lowerCamelCase_						:   str		= UperNetConfig(
       backbone_config=A__						,						auxiliary_in_channels=A__						,						num_labels=A__						,						idalabel=A__						,						labelaid=A__						,						)

   return config


def    __lowerCamelCase   (      A__						: Optional[Any]					)     ->    Optional[Any]:
   lowerCamelCase_						:   Tuple		= []

   # 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    __lowerCamelCase   (      A__						: Optional[Any]						,						A__						: List[Any]						,						A__						: int					)     ->    Union[str, Any]:
   lowerCamelCase_						:   Dict		= dct.pop(A__					)
   lowerCamelCase_						:   int		= val


def    __lowerCamelCase   (      A__						: Tuple						,						A__						: str						,						A__						: Any					)     ->    Optional[int]:
   lowerCamelCase_						:   int		= {
       """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""",
   }
   lowerCamelCase_						:   List[str]		= model_name_to_url[model_name]
   lowerCamelCase_						:   Dict		= torch.hub.load_state_dict_from_url(A__						,						map_location="""cpu"""					)["""state_dict"""]

   lowerCamelCase_						:   List[Any]		= get_upernet_config(A__					)
   lowerCamelCase_						:   Any		= UperNetForSemanticSegmentation(A__					)
   model.eval()

   # replace "bn" => "batch_norm"
   for key in state_dict.copy().keys():
      lowerCamelCase_						:   Tuple		= state_dict.pop(A__					)
      if "bn" in key:
         lowerCamelCase_						:   Dict		= key.replace("""bn"""						,						"""batch_norm"""					)
      lowerCamelCase_						:   Tuple		= val

   # rename keys
   lowerCamelCase_						:   List[Any]		= create_rename_keys(A__					)
   for src, dest in rename_keys:
      rename_key(A__						,						A__						,						A__					)

   model.load_state_dict(A__					)

   # verify on image
   lowerCamelCase_						:   Any		= """https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"""
   lowerCamelCase_						:   Optional[int]		= Image.open(requests.get(A__						,						stream=A__					).raw					).convert("""RGB"""					)

   lowerCamelCase_						:   List[str]		= SegformerImageProcessor()
   lowerCamelCase_						:   Union[str, Any]		= processor(A__						,						return_tensors="""pt"""					).pixel_values

   with torch.no_grad():
      lowerCamelCase_						:   List[str]		= model(A__					)

   if model_name == "upernet-convnext-tiny":
      lowerCamelCase_						:   Dict		= torch.tensor(
          [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]]					)
   elif model_name == "upernet-convnext-small":
      lowerCamelCase_						:   Optional[Any]		= torch.tensor(
          [[-8.8236, -8.8236, -8.6771], [-8.8236, -8.8236, -8.6771], [-8.7638, -8.7638, -8.6240]]					)
   elif model_name == "upernet-convnext-base":
      lowerCamelCase_						:   str		= torch.tensor(
          [[-8.8558, -8.8558, -8.6905], [-8.8558, -8.8558, -8.6905], [-8.7669, -8.7669, -8.6021]]					)
   elif model_name == "upernet-convnext-large":
      lowerCamelCase_						:   List[str]		= torch.tensor(
          [[-8.6660, -8.6660, -8.6210], [-8.6660, -8.6660, -8.6210], [-8.6310, -8.6310, -8.5964]]					)
   elif model_name == "upernet-convnext-xlarge":
      lowerCamelCase_						:   Any		= torch.tensor(
          [[-8.4980, -8.4980, -8.3977], [-8.4980, -8.4980, -8.3977], [-8.4379, -8.4379, -8.3412]]					)
   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__":
      snake_case__				:	Dict					       =					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.'
      )

      snake_case__				:	List[Any]					       =					parser.parse_args()
      convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





