





'''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__   (   __lowercase						:  List[str]     ) ->      Tuple:



       """simple docstring"""






       __UpperCamelCase       =  384
       if "tiny" in model_name:
              __UpperCamelCase       =  [3, 3, 9, 3]
              __UpperCamelCase       =  [96, 192, 384, 768]
       if "small" in model_name:
              __UpperCamelCase       =  [3, 3, 27, 3]
              __UpperCamelCase       =  [96, 192, 384, 768]
       if "base" in model_name:
              __UpperCamelCase       =  [3, 3, 27, 3]
              __UpperCamelCase       =  [128, 256, 512, 1024]
              __UpperCamelCase       =  512
       if "large" in model_name:
              __UpperCamelCase       =  [3, 3, 27, 3]
              __UpperCamelCase       =  [192, 384, 768, 1536]
              __UpperCamelCase       =  768
       if "xlarge" in model_name:
              __UpperCamelCase       =  [3, 3, 27, 3]
              __UpperCamelCase       =  [256, 512, 1024, 2048]
              __UpperCamelCase       =  1024

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

       __UpperCamelCase       =  ConvNextConfig(
           depths=__lowercase				,  hidden_sizes=__lowercase				,  out_features=['stage1', 'stage2', 'stage3', 'stage4']     )
       __UpperCamelCase       =  UperNetConfig(
           backbone_config=__lowercase				,  auxiliary_in_channels=__lowercase				,  num_labels=__lowercase				,  idalabel=__lowercase				,  labelaid=__lowercase				,  )

       return config
def 				lowercase__   (   __lowercase						:  int     ) ->      str:



       """simple docstring"""






       __UpperCamelCase       =  []

       # 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__   (   __lowercase						:  Dict				,  __lowercase						:  Optional[int]				,  __lowercase						:  Optional[Any]     ) ->      Union[str, Any]:



       """simple docstring"""






       __UpperCamelCase       =  dct.pop(__lowercase     )
       __UpperCamelCase       =  val
def 				lowercase__   (   __lowercase						:  Union[str, Any]				,  __lowercase						:  List[Any]				,  __lowercase						:  Dict     ) ->      Any:



       """simple docstring"""






       __UpperCamelCase       =  {
           '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',
       }
       __UpperCamelCase       =  model_name_to_url[model_name]
       __UpperCamelCase       =  torch.hub.load_state_dict_from_url(__lowercase				,  map_location='cpu'     )['state_dict']

       __UpperCamelCase       =  get_upernet_config(__lowercase     )
       __UpperCamelCase       =  UperNetForSemanticSegmentation(__lowercase     )
       model.eval()

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

       # rename keys
       __UpperCamelCase       =  create_rename_keys(__lowercase     )
       for src, dest in rename_keys:
              rename_key(__lowercase				,  __lowercase				,  __lowercase     )

       model.load_state_dict(__lowercase     )

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

       __UpperCamelCase       =  SegformerImageProcessor()
       __UpperCamelCase       =  processor(__lowercase				,  return_tensors='pt'     ).pixel_values

       with torch.no_grad():
              __UpperCamelCase       =  model(__lowercase     )

       if model_name == "upernet-convnext-tiny":
              __UpperCamelCase       =  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":
              __UpperCamelCase       =  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":
              __UpperCamelCase       =  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":
              __UpperCamelCase       =  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":
              __UpperCamelCase       =  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]				,  __lowercase				,  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(__lowercase     )
              print(F'''Saving processor to {pytorch_dump_folder_path}'''     )
              processor.save_pretrained(__lowercase     )

       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__":
    a__							:							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.'''
    )

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







'''simple docstring'''
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


a__							:							List[Any] ={
    '''configuration_timesformer''': ['''TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''TimesformerConfig'''],
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    a__							:							Optional[int] =[
        '''TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
        '''TimesformerModel''',
        '''TimesformerForVideoClassification''',
        '''TimesformerPreTrainedModel''',
    ]

if TYPE_CHECKING:
    from .configuration_timesformer import TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, TimesformerConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_timesformer import (
            TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
            TimesformerForVideoClassification,
            TimesformerModel,
            TimesformerPreTrainedModel,
        )

else:
    import sys

    a__							:							Optional[int] =_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)
