






"""simple docstring"""







import argparse
import json

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

from transformers import SegformerImageProcessor, SwinConfig, UperNetConfig, UperNetForSemanticSegmentation






def       _lowerCAmelCase       (      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase	    =			384
       UpperCAmelCase	    =			7
       if "tiny" in model_name:
              UpperCAmelCase	    =			96
              UpperCAmelCase	    =			(2, 2, 6, 2)
              UpperCAmelCase	    =			(3, 6, 12, 24)
       elif "small" in model_name:
              UpperCAmelCase	    =			96
              UpperCAmelCase	    =			(2, 2, 18, 2)
              UpperCAmelCase	    =			(3, 6, 12, 24)
       elif "base" in model_name:
              UpperCAmelCase	    =			128
              UpperCAmelCase	    =			(2, 2, 18, 2)
              UpperCAmelCase	    =			(4, 8, 16, 32)
              UpperCAmelCase	    =			12
              UpperCAmelCase	    =			512
       elif "large" in model_name:
              UpperCAmelCase	    =			192
              UpperCAmelCase	    =			(2, 2, 18, 2)
              UpperCAmelCase	    =			(6, 12, 24, 48)
              UpperCAmelCase	    =			12
              UpperCAmelCase	    =			768

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

       UpperCAmelCase	    =			SwinConfig(
           embed_dim=SCREAMING_SNAKE_CASE_	,      depths=SCREAMING_SNAKE_CASE_	,      num_heads=SCREAMING_SNAKE_CASE_	,      window_size=SCREAMING_SNAKE_CASE_	,      out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]	,      )
       UpperCAmelCase	    =			UperNetConfig(
           backbone_config=SCREAMING_SNAKE_CASE_	,      auxiliary_in_channels=SCREAMING_SNAKE_CASE_	,      num_labels=SCREAMING_SNAKE_CASE_	,      idalabel=SCREAMING_SNAKE_CASE_	,      labelaid=SCREAMING_SNAKE_CASE_	,      )

       return config






def       _lowerCAmelCase       (      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase	    =			[]

       # fmt: off
       # stem
       rename_keys.append(("""backbone.patch_embed.projection.weight""", """backbone.embeddings.patch_embeddings.projection.weight""")		)
       rename_keys.append(("""backbone.patch_embed.projection.bias""", """backbone.embeddings.patch_embeddings.projection.bias""")		)
       rename_keys.append(("""backbone.patch_embed.norm.weight""", """backbone.embeddings.norm.weight""")		)
       rename_keys.append(("""backbone.patch_embed.norm.bias""", """backbone.embeddings.norm.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}.blocks.{j}.norm1.weight''', F'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.norm1.bias''', F'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_bias_table''', F'''backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_index''', F'''backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.weight''', F'''backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.bias''', F'''backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.norm2.weight''', F'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.norm2.bias''', F'''backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.weight''', F'''backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.bias''', F'''backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.ffn.layers.1.weight''', F'''backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight''')		)
                     rename_keys.append((F'''backbone.stages.{i}.blocks.{j}.ffn.layers.1.bias''', F'''backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias''')		)

              if i < 3:
                     rename_keys.append((F'''backbone.stages.{i}.downsample.reduction.weight''', F'''backbone.encoder.layers.{i}.downsample.reduction.weight''')		)
                     rename_keys.append((F'''backbone.stages.{i}.downsample.norm.weight''', F'''backbone.encoder.layers.{i}.downsample.norm.weight''')		)
                     rename_keys.append((F'''backbone.stages.{i}.downsample.norm.bias''', F'''backbone.encoder.layers.{i}.downsample.norm.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       (      lowerCAmelCase	,      lowerCAmelCase	,      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase	    =			dct.pop(SCREAMING_SNAKE_CASE_		)
       UpperCAmelCase	    =			val






def       _lowerCAmelCase       (      lowerCAmelCase	,      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase	    =			[int(backbone_config.embed_dim * 2**i		) for i in range(len(backbone_config.depths		)		)]
       for i in range(len(backbone_config.depths		)		):
              UpperCAmelCase	    =			num_features[i]
              for j in range(backbone_config.depths[i]		):
                     # fmt: off
                     # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
                     UpperCAmelCase	    =			state_dict.pop(F'''backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight'''		)
                     UpperCAmelCase	    =			state_dict.pop(F'''backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.bias'''		)
                     # next, add query, keys and values (in that order) to the state dict
                     UpperCAmelCase	    =			in_proj_weight[:dim, :]
                     UpperCAmelCase	    =			in_proj_bias[: dim]
                     UpperCAmelCase	    =			in_proj_weight[
                         dim : dim * 2, :
                     ]
                     UpperCAmelCase	    =			in_proj_bias[
                         dim : dim * 2
                     ]
                     UpperCAmelCase	    =			in_proj_weight[
                         -dim :, :
                     ]
                     UpperCAmelCase	    =			in_proj_bias[-dim :]
                     # fmt: on






def       _lowerCAmelCase       (      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase    ,       UpperCAmelCase	    =			x.shape
       UpperCAmelCase	    =			x.reshape(SCREAMING_SNAKE_CASE_	,      4	,      in_channel // 4		)
       UpperCAmelCase	    =			x[:, [0, 2, 1, 3], :].transpose(1	,      2		).reshape(SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_		)
       return x






def       _lowerCAmelCase       (      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase    ,       UpperCAmelCase	    =			x.shape
       UpperCAmelCase	    =			x.reshape(SCREAMING_SNAKE_CASE_	,      in_channel // 4	,      4		)
       UpperCAmelCase	    =			x[:, :, [0, 2, 1, 3]].transpose(1	,      2		).reshape(SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_		)

       return x






def       _lowerCAmelCase       (      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase	    =			x.shape[0]
       UpperCAmelCase	    =			x.reshape(4	,      in_channel // 4		)
       UpperCAmelCase	    =			x[[0, 2, 1, 3], :].transpose(0	,      1		).reshape(SCREAMING_SNAKE_CASE_		)
       return x






def       _lowerCAmelCase       (      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase	    =			x.shape[0]
       UpperCAmelCase	    =			x.reshape(in_channel // 4	,      4		)
       UpperCAmelCase	    =			x[:, [0, 2, 1, 3]].transpose(0	,      1		).reshape(SCREAMING_SNAKE_CASE_		)
       return x






def       _lowerCAmelCase       (      lowerCAmelCase	,      lowerCAmelCase	,      lowerCAmelCase		):







       '''simple docstring'''
       UpperCAmelCase	    =			{
           """upernet-swin-tiny""": """https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210531_112542-e380ad3e.pth""",
           """upernet-swin-small""": """https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210526_192015-ee2fff1c.pth""",
           """upernet-swin-base""": """https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K_20210531_125459-429057bf.pth""",
           """upernet-swin-large""": """https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k_20220318_091743-9ba68901.pth""",
       }
       UpperCAmelCase	    =			model_name_to_url[model_name]
       UpperCAmelCase	    =			torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_	,      map_location="""cpu"""	,      file_name=SCREAMING_SNAKE_CASE_		)[
           """state_dict"""
       ]

       for name, param in state_dict.items():
              print(SCREAMING_SNAKE_CASE_	,      param.shape		)

       UpperCAmelCase	    =			get_upernet_config(SCREAMING_SNAKE_CASE_		)
       UpperCAmelCase	    =			UperNetForSemanticSegmentation(SCREAMING_SNAKE_CASE_		)
       model.eval()

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

       # rename keys
       UpperCAmelCase	    =			create_rename_keys(SCREAMING_SNAKE_CASE_		)
       for src, dest in rename_keys:
              rename_key(SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_	,      SCREAMING_SNAKE_CASE_		)
       read_in_q_k_v(SCREAMING_SNAKE_CASE_	,      config.backbone_config		)

       # fix downsample parameters
       for key, value in state_dict.items():
              if "downsample" in key:
                     if "reduction" in key:
                            UpperCAmelCase	    =			reverse_correct_unfold_reduction_order(SCREAMING_SNAKE_CASE_		)
                     if "norm" in key:
                            UpperCAmelCase	    =			reverse_correct_unfold_norm_order(SCREAMING_SNAKE_CASE_		)

       model.load_state_dict(SCREAMING_SNAKE_CASE_		)

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

       UpperCAmelCase	    =			SegformerImageProcessor()
       UpperCAmelCase	    =			processor(SCREAMING_SNAKE_CASE_	,      return_tensors="""pt"""		).pixel_values

       with torch.no_grad():
              UpperCAmelCase	    =			model(SCREAMING_SNAKE_CASE_		)
              UpperCAmelCase	    =			outputs.logits

       print(logits.shape		)
       print("""First values of logits:"""	,      logits[0, 0, :3, :3]		)
       # assert values
       if model_name == "upernet-swin-tiny":
              UpperCAmelCase	    =			torch.tensor(
                  [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]]		)
       elif model_name == "upernet-swin-small":
              UpperCAmelCase	    =			torch.tensor(
                  [[-7.19_21, -7.19_21, -6.95_32], [-7.19_21, -7.19_21, -6.95_32], [-7.09_08, -7.09_08, -6.85_34]]		)
       elif model_name == "upernet-swin-base":
              UpperCAmelCase	    =			torch.tensor(
                  [[-6.58_51, -6.58_51, -6.43_30], [-6.58_51, -6.58_51, -6.43_30], [-6.47_63, -6.47_63, -6.32_54]]		)
       elif model_name == "upernet-swin-large":
              UpperCAmelCase	    =			torch.tensor(
                  [[-7.52_97, -7.52_97, -7.38_02], [-7.52_97, -7.52_97, -7.38_02], [-7.40_44, -7.40_44, -7.25_86]]		)
       print("""Logits:"""	,      outputs.logits[0, 0, :3, :3]		)
       assert torch.allclose(outputs.logits[0, 0, :3, :3]	,      SCREAMING_SNAKE_CASE_	,      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(SCREAMING_SNAKE_CASE_		)
              print(F'''Saving processor to {pytorch_dump_folder_path}'''		)
              processor.save_pretrained(SCREAMING_SNAKE_CASE_		)

       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_	:    Optional[Any]            =      argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--model_name''',
     default='''upernet-swin-tiny''',
     type=str,
     choices=[F'upernet-swin-{size}' for size in ['''tiny''', '''small''', '''base''', '''large''']],
     help='''Name of the Swin + 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_	:    Union[str, Any]            =      parser.parse_args()
 convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








"""simple docstring"""







import math
from typing import Optional

import numpy as np

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase_	:    List[Any]            =      logging.get_logger(__name__)

lowerCAmelCase_	:    Optional[Any]            =      {
    '''facebook/encodec_24khz''': '''https://huggingface.co/facebook/encodec_24khz/resolve/main/config.json''',
    '''facebook/encodec_48khz''': '''https://huggingface.co/facebook/encodec_48khz/resolve/main/config.json''',
}





class   UpperCamelCase_ (       a_  ):
  _A       :					Dict				=	'encodec'



  def __init__(      self ,		snake_case__=[1.5, 3.0, 6.0, 12.0, 24.0] ,		snake_case__=2_40_00 ,		snake_case__=1 ,		snake_case__=False ,		snake_case__=None ,		snake_case__=None ,		snake_case__=1_28 ,		snake_case__=32 ,		snake_case__=1 ,		snake_case__=[8, 5, 4, 2] ,		snake_case__="weight_norm" ,		snake_case__=7 ,		snake_case__=7 ,		snake_case__=3 ,		snake_case__=2 ,		snake_case__=True ,		snake_case__="reflect" ,		snake_case__=2 ,		snake_case__=2 ,		snake_case__=1.0 ,		snake_case__=10_24 ,		snake_case__=None ,		snake_case__=True ,		**snake_case__ ,		)			->			List[Any]:



         """simple docstring"""





         UpperCAmelCase	    =			target_bandwidths
         UpperCAmelCase	    =			sampling_rate
         UpperCAmelCase	    =			audio_channels
         UpperCAmelCase	    =			normalize
         UpperCAmelCase	    =			chunk_length_s
         UpperCAmelCase	    =			overlap
         UpperCAmelCase	    =			hidden_size
         UpperCAmelCase	    =			num_filters
         UpperCAmelCase	    =			num_residual_layers
         UpperCAmelCase	    =			upsampling_ratios
         UpperCAmelCase	    =			norm_type
         UpperCAmelCase	    =			kernel_size
         UpperCAmelCase	    =			last_kernel_size
         UpperCAmelCase	    =			residual_kernel_size
         UpperCAmelCase	    =			dilation_growth_rate
         UpperCAmelCase	    =			use_causal_conv
         UpperCAmelCase	    =			pad_mode
         UpperCAmelCase	    =			compress
         UpperCAmelCase	    =			num_lstm_layers
         UpperCAmelCase	    =			trim_right_ratio
         UpperCAmelCase	    =			codebook_size
         UpperCAmelCase	    =			codebook_dim if codebook_dim is not None else hidden_size
         UpperCAmelCase	    =			use_conv_shortcut

         if self.norm_type not in ["weight_norm", "time_group_norm"]:
                raise ValueError(
                    f'''self.norm_type must be one of `"weight_norm"`, `"time_group_norm"`), got {self.norm_type}'''				)

         super().__init__(**snake_case__				)



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



         """simple docstring"""





         if self.chunk_length_s is None:
                return None
         else:
                return int(self.chunk_length_s * self.sampling_rate				)



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



         """simple docstring"""





         if self.chunk_length_s is None or self.overlap is None:
                return None
         else:
                return max(1 ,		int((1.0 - self.overlap) * self.chunk_length				)				)



  @property
  def 		UpperCamelCase_				(      self				)			->			int:



         """simple docstring"""





         UpperCAmelCase	    =			np.prod(self.upsampling_ratios				)
         return math.ceil(self.sampling_rate / hop_length				)



  @property
  def 		UpperCamelCase_				(      self				)			->			int:



         """simple docstring"""





         return int(10_00 * self.target_bandwidths[-1] // (self.frame_rate * 10)				)
