


import argparse
import json

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

from transformers import BitImageProcessor, FocalNetConfig, FocalNetForImageClassification
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, PILImageResampling






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      )->						Optional[Any]:




      """simple docstring"""

      snake_case_ =   [2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 18, 2]
      snake_case_ =   True if '''large''' in model_name or '''huge''' in model_name else False
      snake_case_ =   True if '''large''' in model_name or '''huge''' in model_name else False
      snake_case_ =   True if '''large''' in model_name or '''huge''' in model_name else False

      if "large" in model_name or "xlarge" in model_name or "huge" in model_name:
            if "fl3" in model_name:
                  snake_case_ =   [3, 3, 3, 3]
                  snake_case_ =   [5, 5, 5, 5]
            elif "fl4" in model_name:
                  snake_case_ =   [4, 4, 4, 4]
                  snake_case_ =   [3, 3, 3, 3]

      if "tiny" in model_name or "small" in model_name or "base" in model_name:
            snake_case_ =   [3, 3, 3, 3]
            if "lrf" in model_name:
                  snake_case_ =   [3, 3, 3, 3]
            else:
                  snake_case_ =   [2, 2, 2, 2]

      if "tiny" in model_name:
            snake_case_ =   96
      elif "small" in model_name:
            snake_case_ =   96
      elif "base" in model_name:
            snake_case_ =   128
      elif "large" in model_name:
            snake_case_ =   192
      elif "xlarge" in model_name:
            snake_case_ =   256
      elif "huge" in model_name:
            snake_case_ =   352

      # set label information
      snake_case_ =   '''huggingface/label-files'''
      if "large" in model_name or "huge" in model_name:
            snake_case_ =   '''imagenet-22k-id2label.json'''
      else:
            snake_case_ =   '''imagenet-1k-id2label.json'''

      snake_case_ =   json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      repo_type='''dataset'''      )      ,      '''r'''      )      )
      snake_case_ =   {int(SCREAMING_SNAKE_CASE      ): v for k, v in idalabel.items()}
      snake_case_ =   {v: k for k, v in idalabel.items()}

      snake_case_ =   FocalNetConfig(
          embed_dim=SCREAMING_SNAKE_CASE      ,      depths=SCREAMING_SNAKE_CASE      ,      focal_levels=SCREAMING_SNAKE_CASE      ,      focal_windows=SCREAMING_SNAKE_CASE      ,      use_conv_embed=SCREAMING_SNAKE_CASE      ,      idalabel=SCREAMING_SNAKE_CASE      ,      labelaid=SCREAMING_SNAKE_CASE      ,      use_post_layernorm=SCREAMING_SNAKE_CASE      ,      use_layerscale=SCREAMING_SNAKE_CASE      ,      )

      return config






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      )->						Tuple:




      """simple docstring"""

      if "patch_embed.proj" in name:
            snake_case_ =   name.replace('''patch_embed.proj'''      ,      '''embeddings.patch_embeddings.projection'''      )
      if "patch_embed.norm" in name:
            snake_case_ =   name.replace('''patch_embed.norm'''      ,      '''embeddings.norm'''      )
      if "layers" in name:
            snake_case_ =   '''encoder.''' + name
      if "encoder.layers" in name:
            snake_case_ =   name.replace('''encoder.layers'''      ,      '''encoder.stages'''      )
      if "downsample.proj" in name:
            snake_case_ =   name.replace('''downsample.proj'''      ,      '''downsample.projection'''      )
      if "blocks" in name:
            snake_case_ =   name.replace('''blocks'''      ,      '''layers'''      )
      if "modulation.f.weight" in name or "modulation.f.bias" in name:
            snake_case_ =   name.replace('''modulation.f'''      ,      '''modulation.projection_in'''      )
      if "modulation.h.weight" in name or "modulation.h.bias" in name:
            snake_case_ =   name.replace('''modulation.h'''      ,      '''modulation.projection_context'''      )
      if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
            snake_case_ =   name.replace('''modulation.proj'''      ,      '''modulation.projection_out'''      )

      if name == "norm.weight":
            snake_case_ =   '''layernorm.weight'''
      if name == "norm.bias":
            snake_case_ =   '''layernorm.bias'''

      if "head" in name:
            snake_case_ =   name.replace('''head'''      ,      '''classifier'''      )
      else:
            snake_case_ =   '''focalnet.''' + name

      return name






def  __lowerCAmelCase					(SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE=False      )->						Union[str, Any]:




      """simple docstring"""

      snake_case_ =   {
          '''focalnet-tiny''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_srf.pth''',
          '''focalnet-tiny-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_tiny_lrf.pth''',
          '''focalnet-small''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_srf.pth''',
          '''focalnet-small-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_small_lrf.pth''',
          '''focalnet-base''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_srf.pth''',
          '''focalnet-base-lrf''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_lrf.pth''',
          '''focalnet-large-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384.pth''',
          '''focalnet-large-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_large_lrf_384_fl4.pth''',
          '''focalnet-xlarge-lrf-fl3''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384.pth''',
          '''focalnet-xlarge-lrf-fl4''': '''https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_xlarge_lrf_384_fl4.pth''',
      }
      # fmt: on

      snake_case_ =   model_name_to_url[model_name]
      print('''Checkpoint URL: '''      ,      SCREAMING_SNAKE_CASE      )
      snake_case_ =   torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE      ,      map_location='''cpu'''      )['''model''']

      # rename keys
      for key in state_dict.copy().keys():
            snake_case_ =   state_dict.pop(SCREAMING_SNAKE_CASE      )
            snake_case_ =   val

      snake_case_ =   get_focalnet_config(SCREAMING_SNAKE_CASE      )
      snake_case_ =   FocalNetForImageClassification(SCREAMING_SNAKE_CASE      )
      model.eval()

      # load state dict
      model.load_state_dict(SCREAMING_SNAKE_CASE      )

      # verify conversion
      snake_case_ =   '''http://images.cocodataset.org/val2017/000000039769.jpg'''

      snake_case_ =   BitImageProcessor(
          do_resize=SCREAMING_SNAKE_CASE      ,      size={'''shortest_edge''': 256}      ,      resample=PILImageResampling.BILINEAR      ,      do_center_crop=SCREAMING_SNAKE_CASE      ,      crop_size=224      ,      do_normalize=SCREAMING_SNAKE_CASE      ,      image_mean=SCREAMING_SNAKE_CASE      ,      image_std=SCREAMING_SNAKE_CASE      ,      )
      snake_case_ =   Image.open(requests.get(SCREAMING_SNAKE_CASE      ,      stream=SCREAMING_SNAKE_CASE      ).raw      )
      snake_case_ =   processor(images=SCREAMING_SNAKE_CASE      ,      return_tensors='''pt'''      )

      snake_case_ =   transforms.Compose(
          [
              transforms.Resize(256      ),
              transforms.CenterCrop(224      ),
              transforms.ToTensor(),
              transforms.Normalize(mean=[0.4_8_5, 0.4_5_6, 0.4_0_6]      ,      std=[0.2_2_9, 0.2_2_4, 0.2_2_5]      ),
          ]      )

      snake_case_ =   image_transforms(SCREAMING_SNAKE_CASE      ).unsqueeze(0      )

      # verify pixel_values
      assert torch.allclose(inputs.pixel_values      ,      SCREAMING_SNAKE_CASE      ,      atol=1E-4      )

      snake_case_ =   model(**SCREAMING_SNAKE_CASE      )

      snake_case_ =   outputs.logits.argmax(-1      ).item()
      print('''Predicted class:'''      ,      model.config.idalabel[predicted_class_idx]      )

      print('''First values of logits:'''      ,      outputs.logits[0, :3]      )

      if model_name == "focalnet-tiny":
            snake_case_ =   torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]      )
      elif model_name == "focalnet-tiny-lrf":
            snake_case_ =   torch.tensor([1.1_6_6_9, 0.0_1_2_5, -0.1_6_9_5]      )
      elif model_name == "focalnet-small":
            snake_case_ =   torch.tensor([0.4_9_1_7, -0.0_4_3_0, 0.1_3_4_1]      )
      elif model_name == "focalnet-small-lrf":
            snake_case_ =   torch.tensor([-0.2_5_8_8, -0.5_3_4_2, -0.2_3_3_1]      )
      elif model_name == "focalnet-base":
            snake_case_ =   torch.tensor([-0.1_6_5_5, -0.4_0_9_0, -0.1_7_3_0]      )
      elif model_name == "focalnet-base-lrf":
            snake_case_ =   torch.tensor([0.5_3_0_6, -0.0_4_8_3, -0.3_9_2_8]      )
      assert torch.allclose(outputs.logits[0, :3]      ,      SCREAMING_SNAKE_CASE      ,      atol=1E-4      )
      print('''Looks ok!'''      )

      if pytorch_dump_folder_path is not None:
            print(f'''Saving model and processor of {model_name} to {pytorch_dump_folder_path}'''      )
            model.save_pretrained(SCREAMING_SNAKE_CASE      )
            processor.save_pretrained(SCREAMING_SNAKE_CASE      )

      if push_to_hub:
            print(f'''Pushing model and processor of {model_name} to the hub...'''      )
            model.push_to_hub(f'''{model_name}'''      )
            processor.push_to_hub(f'''{model_name}'''      )


if __name__ == "__main__":
   UpperCAmelCase    =					argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       """--model_name""",
       default="""focalnet-tiny""",
       type=str,
       help="""Name of the FocalNet 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 to push the model and processor to the hub.""",
   )

   UpperCAmelCase    =					parser.parse_args()
   convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)



from ...utils import is_note_seq_available, is_transformers_available, is_torch_available
from ...utils import OptionalDependencyNotAvailable


try:
   if not (is_transformers_available() and is_torch_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
   from .notes_encoder import SpectrogramNotesEncoder
   from .continous_encoder import SpectrogramContEncoder
   from .pipeline_spectrogram_diffusion import (
       SpectrogramContEncoder,
       SpectrogramDiffusionPipeline,
       TaFilmDecoder,
   )

try:
   if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
   from .midi_utils import MidiProcessor