



'''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__ (		_A       ):
  a		:      Optional[int]       =  384
  a		:      Dict       =  7
  if "tiny" in model_name:
    a		:      Any       =  96
    a		:      int       =  (2, 2, 6, 2)
    a		:      Union[str, Any]       =  (3, 6, 12, 24)
  elif "small" in model_name:
    a		:      str       =  96
    a		:      Any       =  (2, 2, 18, 2)
    a		:      Any       =  (3, 6, 12, 24)
  elif "base" in model_name:
    a		:      Optional[int]       =  128
    a		:      int       =  (2, 2, 18, 2)
    a		:      Any       =  (4, 8, 16, 32)
    a		:      str       =  12
    a		:      List[Any]       =  512
  elif "large" in model_name:
    a		:      int       =  192
    a		:      Optional[Any]       =  (2, 2, 18, 2)
    a		:      List[Any]       =  (6, 12, 24, 48)
    a		:      List[Any]       =  12
    a		:      List[Any]       =  768

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

  a		:      Optional[int]       =  SwinConfig(
      embed_dim=_A				,    depths=_A				,    num_heads=_A				,    window_size=_A				,    out_features=['stage1', 'stage2', 'stage3', 'stage4']				,    )
  a		:      int       =  UperNetConfig(
      backbone_config=_A				,    auxiliary_in_channels=_A				,    num_labels=_A				,    idalabel=_A				,    labelaid=_A				,    )

  return config
def    lowerCamelCase__ (		_A       ):
  a		:      List[Any]       =  []

  # 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__ (		_A				,    _A				,    _A       ):
  a		:      Optional[Any]       =  dct.pop(_A       )
  a		:      Any       =  val
def    lowerCamelCase__ (		_A				,    _A       ):
  a		:      str       =  [int(backbone_config.embed_dim * 2**i       ) for i in range(len(backbone_config.depths       )       )]
  for i in range(len(backbone_config.depths       )       ):
    a		:      List[Any]       =  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)
      a		:      Optional[int]       =  state_dict.pop(f"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight"""       )
      a		:      Any       =  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
      a		:      Dict       =  in_proj_weight[:dim, :]
      a		:      Optional[Any]       =  in_proj_bias[: dim]
      a		:      int       =  in_proj_weight[
          dim : dim * 2, :
      ]
      a		:      Union[str, Any]       =  in_proj_bias[
          dim : dim * 2
      ]
      a		:      List[str]       =  in_proj_weight[
          -dim :, :
      ]
      a		:      Tuple       =  in_proj_bias[-dim :]
      # fmt: on
def    lowerCamelCase__ (		_A       ):
  a     ,			a		:      str       =  x.shape
  a		:      List[Any]       =  x.reshape(_A				,    4				,    in_channel // 4       )
  a		:      Tuple       =  x[:, [0, 2, 1, 3], :].transpose(1				,    2       ).reshape(_A				,    _A       )
  return x
def    lowerCamelCase__ (		_A       ):
  a     ,			a		:      int       =  x.shape
  a		:      List[Any]       =  x.reshape(_A				,    in_channel // 4				,    4       )
  a		:      Union[str, Any]       =  x[:, :, [0, 2, 1, 3]].transpose(1				,    2       ).reshape(_A				,    _A       )

  return x
def    lowerCamelCase__ (		_A       ):
  a		:      Optional[int]       =  x.shape[0]
  a		:      Optional[int]       =  x.reshape(4				,    in_channel // 4       )
  a		:      List[str]       =  x[[0, 2, 1, 3], :].transpose(0				,    1       ).reshape(_A       )
  return x
def    lowerCamelCase__ (		_A       ):
  a		:      Tuple       =  x.shape[0]
  a		:      str       =  x.reshape(in_channel // 4				,    4       )
  a		:      Tuple       =  x[:, [0, 2, 1, 3]].transpose(0				,    1       ).reshape(_A       )
  return x
def    lowerCamelCase__ (		_A				,    _A				,    _A       ):
  a		:      Optional[int]       =  {
      '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',
  }
  a		:      Optional[Any]       =  model_name_to_url[model_name]
  a		:      Tuple       =  torch.hub.load_state_dict_from_url(_A				,    map_location='cpu'				,    file_name=_A       )[
      'state_dict'
  ]

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

  a		:      List[str]       =  get_upernet_config(_A       )
  a		:      Any       =  UperNetForSemanticSegmentation(_A       )
  model.eval()

  # replace "bn" => "batch_norm"
  for key in state_dict.copy().keys():
    a		:      Any       =  state_dict.pop(_A       )
    if "bn" in key:
      a		:      Dict       =  key.replace('bn'				,    'batch_norm'       )
    a		:      Tuple       =  val

  # rename keys
  a		:      Union[str, Any]       =  create_rename_keys(_A       )
  for src, dest in rename_keys:
    rename_key(_A				,    _A				,    _A       )
  read_in_q_k_v(_A				,    config.backbone_config       )

  # fix downsample parameters
  for key, value in state_dict.items():
    if "downsample" in key:
      if "reduction" in key:
        a		:      int       =  reverse_correct_unfold_reduction_order(_A       )
      if "norm" in key:
        a		:      Dict       =  reverse_correct_unfold_norm_order(_A       )

  model.load_state_dict(_A       )

  # verify on image
  a		:      int       =  'https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg'
  a		:      List[str]       =  Image.open(requests.get(_A				,    stream=_A       ).raw       ).convert('RGB'       )

  a		:      List[Any]       =  SegformerImageProcessor()
  a		:      Optional[int]       =  processor(_A				,    return_tensors='pt'       ).pixel_values

  with torch.no_grad():
    a		:      Dict       =  model(_A       )
    a		:      Tuple       =  outputs.logits

  print(logits.shape       )
  print('First values of logits:'				,    logits[0, 0, :3, :3]       )
  # assert values
  if model_name == "upernet-swin-tiny":
    a		:      Union[str, Any]       =  torch.tensor(
        [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]]       )
  elif model_name == "upernet-swin-small":
    a		:      Dict       =  torch.tensor(
        [[-7.1921, -7.1921, -6.9532], [-7.1921, -7.1921, -6.9532], [-7.0908, -7.0908, -6.8534]]       )
  elif model_name == "upernet-swin-base":
    a		:      int       =  torch.tensor(
        [[-6.5851, -6.5851, -6.4330], [-6.5851, -6.5851, -6.4330], [-6.4763, -6.4763, -6.3254]]       )
  elif model_name == "upernet-swin-large":
    a		:      Optional[Any]       =  torch.tensor(
        [[-7.5297, -7.5297, -7.3802], [-7.5297, -7.5297, -7.3802], [-7.4044, -7.4044, -7.2586]]       )
  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:      Dict		 =		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:      Tuple		 =		parser.parse_args()
     convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




'''simple docstring'''






import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging


logging.set_verbosity_info()
lowerCAmelCase:      Optional[int]		 =		logging.get_logger(__name__)

lowerCAmelCase:      Union[str, Any]		 =		{
    'post_extract_proj': 'feature_projection.projection',
    'encoder.pos_conv.0': 'encoder.pos_conv_embed.conv',
    'self_attn.k_proj': 'encoder.layers.*.attention.k_proj',
    'self_attn.v_proj': 'encoder.layers.*.attention.v_proj',
    'self_attn.q_proj': 'encoder.layers.*.attention.q_proj',
    'self_attn.out_proj': 'encoder.layers.*.attention.out_proj',
    'self_attn_layer_norm': 'encoder.layers.*.layer_norm',
    'fc1': 'encoder.layers.*.feed_forward.intermediate_dense',
    'fc2': 'encoder.layers.*.feed_forward.output_dense',
    'final_layer_norm': 'encoder.layers.*.final_layer_norm',
    'encoder.layer_norm': 'encoder.layer_norm',
    'encoder.layer_norm_for_extract': 'layer_norm_for_extract',
    'w2v_model.layer_norm': 'feature_projection.layer_norm',
    'quantizer.weight_proj': 'quantizer.weight_proj',
    'quantizer.vars': 'quantizer.codevectors',
    'project_q': 'project_q',
    'final_proj': 'project_hid',
    'w2v_encoder.proj': 'lm_head',
    'label_embs_concat': 'label_embeddings_concat',
    'mask_emb': 'masked_spec_embed',
    'spk_proj': 'speaker_proj',
}
lowerCAmelCase:      int		 =		[
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
    'label_embeddings_concat',
    'speaker_proj',
    'layer_norm_for_extract',
]
def    lowerCamelCase__ (		_A				,    _A				,    _A				,    _A				,    _A       ):
  for attribute in key.split('.'       ):
    a		:      Tuple       =  getattr(_A				,    _A       )

  if weight_type is not None:
    a		:      Union[str, Any]       =  getattr(_A				,    _A       ).shape
  else:
    a		:      Optional[Any]       =  hf_pointer.shape

  if hf_shape != value.shape:
    raise ValueError(
        f"""Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be"""
        f""" {value.shape} for {full_name}"""       )

  if weight_type == "weight":
    a		:      str       =  value
  elif weight_type == "weight_g":
    a		:      Optional[int]       =  value
  elif weight_type == "weight_v":
    a		:      List[Any]       =  value
  elif weight_type == "bias":
    a		:      Any       =  value
  else:
    a		:      Optional[Any]       =  value

  logger.info(f"""{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}."""       )
def    lowerCamelCase__ (		_A				,    _A       ):
  a		:      Optional[int]       =  []
  a		:      Tuple       =  fairseq_model.state_dict()

  a		:      List[Any]       =  hf_model.unispeech_sat.feature_extractor

  for name, value in fairseq_dict.items():
    a		:      Optional[Any]       =  False
    if "conv_layers" in name:
      load_conv_layer(
          _A				,    _A				,    _A				,    _A				,    hf_model.config.feat_extract_norm == 'group'				,    )
      a		:      int       =  True
    else:
      for key, mapped_key in MAPPING.items():
        a		:      Any       =  'unispeech_sat.' + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
        if key in name or key.split('w2v_model.'       )[-1] == name.split('.'       )[0]:
          if "layer_norm_for_extract" in name and (".".join(name.split('.'       )[:-1]       ) != key):
            # special case since naming is very similar
            continue
          a		:      str       =  True
          if "*" in mapped_key:
            a		:      Union[str, Any]       =  name.split(_A       )[0].split('.'       )[-2]
            a		:      List[str]       =  mapped_key.replace('*'				,    _A       )
          if "weight_g" in name:
            a		:      Optional[int]       =  'weight_g'
          elif "weight_v" in name:
            a		:      Any       =  'weight_v'
          elif "bias" in name:
            a		:      str       =  'bias'
          elif "weight" in name:
            # TODO: don't match quantizer.weight_proj
            a		:      List[Any]       =  'weight'
          else:
            a		:      Any       =  None
          set_recursively(_A				,    _A				,    _A				,    _A				,    _A       )
        continue
    if not is_used:
      unused_weights.append(_A       )

  logger.warning(f"""Unused weights: {unused_weights}"""       )
def    lowerCamelCase__ (		_A				,    _A				,    _A				,    _A				,    _A       ):
  a		:      List[str]       =  full_name.split('conv_layers.'       )[-1]
  a		:      Any       =  name.split('.'       )
  a		:      Optional[Any]       =  int(items[0]       )
  a		:      Union[str, Any]       =  int(items[1]       )

  if type_id == 0:
    if "bias" in name:
      if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
        raise ValueError(
            f"""{full_name} has size {value.shape}, but"""
            f""" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found."""       )
      a		:      Optional[int]       =  value
      logger.info(f"""Feat extract conv layer {layer_id} was initialized from {full_name}."""       )
    elif "weight" in name:
      if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
        raise ValueError(
            f"""{full_name} has size {value.shape}, but"""
            f""" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found."""       )
      a		:      Optional[int]       =  value
      logger.info(f"""Feat extract conv layer {layer_id} was initialized from {full_name}."""       )
  elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
    if "bias" in name:
      if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
        raise ValueError(
            f"""{full_name} has size {value.shape}, but"""
            f""" {feature_extractor[layer_id].layer_norm.bias.data.shape} was found."""       )
      a		:      Union[str, Any]       =  value
      logger.info(f"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""       )
    elif "weight" in name:
      if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
        raise ValueError(
            f"""{full_name} has size {value.shape}, but"""
            f""" {feature_extractor[layer_id].layer_norm.weight.data.shape} was found."""       )
      a		:      Union[str, Any]       =  value
      logger.info(f"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""       )
  else:
    unused_weights.append(_A       )
@torch.no_grad()
def    lowerCamelCase__ (		_A				,    _A				,    _A=None				,    _A=None				,    _A=True       ):
  if config_path is not None:
    a		:      str       =  UniSpeechSatConfig.from_pretrained(_A       )
  else:
    a		:      Dict       =  UniSpeechSatConfig()

  a		:      Any       =  ''

  if is_finetuned:
    a		:      Tuple       =  UniSpeechSatForCTC(_A       )
  else:
    a		:      List[str]       =  UniSpeechSatForPreTraining(_A       )

  a     ,			a     ,			a		:      str       =  fairseq.checkpoint_utils.load_model_ensemble_and_task(
      [checkpoint_path]				,    arg_overrides={'data': '/'.join(dict_path.split('/'       )[:-1]       )}       )
  a		:      str       =  model[0].eval()

  recursively_load_weights(_A				,    _A       )

  hf_wavavec.save_pretrained(_A       )


if __name__ == "__main__":
     lowerCAmelCase:      Tuple		 =		argparse.ArgumentParser()
     parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
     parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to fairseq checkpoint')
     parser.add_argument('--dict_path', default=None, type=str, help='Path to dict of fine-tuned model')
     parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
     parser.add_argument(
         '--not_finetuned', action='store_true', help='Whether the model to convert is a fine-tuned model or not'
     )
     lowerCAmelCase:      Optional[int]		 =		parser.parse_args()
     convert_unispeech_sat_checkpoint(
         args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
     )