


'''simple docstring'''



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__   ( A     :    int  ):




      '''simple docstring'''






      UpperCAmelCase         =			[2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 18, 2]
      UpperCAmelCase         =			True if '''large''' in model_name or '''huge''' in model_name else False
      UpperCAmelCase         =			True if '''large''' in model_name or '''huge''' in model_name else False
      UpperCAmelCase         =			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:
                  UpperCAmelCase         =			[3, 3, 3, 3]
                  UpperCAmelCase         =			[5, 5, 5, 5]
            elif "fl4" in model_name:
                  UpperCAmelCase         =			[4, 4, 4, 4]
                  UpperCAmelCase         =			[3, 3, 3, 3]

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

      if "tiny" in model_name:
            UpperCAmelCase         =			96
      elif "small" in model_name:
            UpperCAmelCase         =			96
      elif "base" in model_name:
            UpperCAmelCase         =			1_28
      elif "large" in model_name:
            UpperCAmelCase         =			1_92
      elif "xlarge" in model_name:
            UpperCAmelCase         =			2_56
      elif "huge" in model_name:
            UpperCAmelCase         =			3_52

      # set label information
      UpperCAmelCase         =			'''huggingface/label-files'''
      if "large" in model_name or "huge" in model_name:
            UpperCAmelCase         =			'''imagenet-22k-id2label.json'''
      else:
            UpperCAmelCase         =			'''imagenet-1k-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         =			FocalNetConfig(
          embed_dim=_lowercase   ,     depths=_lowercase   ,     focal_levels=_lowercase   ,     focal_windows=_lowercase   ,     use_conv_embed=_lowercase   ,     idalabel=_lowercase   ,     labelaid=_lowercase   ,     use_post_layernorm=_lowercase   ,     use_layerscale=_lowercase   ,     )

      return config


def    lowerCamelCase__   ( A     :    Tuple  ):




      '''simple docstring'''






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

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

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

      return name


def    lowerCamelCase__   ( A     :    Optional[Any]   ,     A     :    int   ,     A     :    List[Any]=False  ):




      '''simple docstring'''






      UpperCAmelCase         =			{
          '''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

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

      # rename keys
      for key in state_dict.copy().keys():
            UpperCAmelCase         =			state_dict.pop(_lowercase  )
            UpperCAmelCase         =			val

      UpperCAmelCase         =			get_focalnet_config(_lowercase  )
      UpperCAmelCase         =			FocalNetForImageClassification(_lowercase  )
      model.eval()

      # load state dict
      model.load_state_dict(_lowercase  )

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

      UpperCAmelCase         =			BitImageProcessor(
          do_resize=_lowercase   ,     size={'''shortest_edge''': 2_56}   ,     resample=PILImageResampling.BILINEAR   ,     do_center_crop=_lowercase   ,     crop_size=2_24   ,     do_normalize=_lowercase   ,     image_mean=_lowercase   ,     image_std=_lowercase   ,     )
      UpperCAmelCase         =			Image.open(requests.get(_lowercase   ,     stream=_lowercase  ).raw  )
      UpperCAmelCase         =			processor(images=_lowercase   ,     return_tensors='''pt'''  )

      UpperCAmelCase         =			transforms.Compose(
          [
              transforms.Resize(2_56  ),
              transforms.CenterCrop(2_24  ),
              transforms.ToTensor(),
              transforms.Normalize(mean=[0.485, 0.456, 0.406]   ,     std=[0.229, 0.224, 0.225]  ),
          ]  )

      UpperCAmelCase         =			image_transforms(_lowercase  ).unsqueeze(0  )

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

      UpperCAmelCase         =			model(**_lowercase  )

      UpperCAmelCase         =			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":
            UpperCAmelCase         =			torch.tensor([0.2_166, -0.4_368, 0.2_191]  )
      elif model_name == "focalnet-tiny-lrf":
            UpperCAmelCase         =			torch.tensor([1.1_669, 0.0_125, -0.1_695]  )
      elif model_name == "focalnet-small":
            UpperCAmelCase         =			torch.tensor([0.4_917, -0.0_430, 0.1_341]  )
      elif model_name == "focalnet-small-lrf":
            UpperCAmelCase         =			torch.tensor([-0.2_588, -0.5_342, -0.2_331]  )
      elif model_name == "focalnet-base":
            UpperCAmelCase         =			torch.tensor([-0.1_655, -0.4_090, -0.1_730]  )
      elif model_name == "focalnet-base-lrf":
            UpperCAmelCase         =			torch.tensor([0.5_306, -0.0_483, -0.3_928]  )
      assert torch.allclose(outputs.logits[0, :3]   ,     _lowercase   ,     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(_lowercase  )
            processor.save_pretrained(_lowercase  )

      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__":
 _lowercase    :    List[Any]			     =  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.""",
 )

 _lowercase    :    Dict			     =  parser.parse_args()
 convert_focalnet_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


_lowercase    :    Union[str, Any]			     =  logging.get_logger(__name__)

_lowercase    :    List[str]			     =  {
    """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__(     lowerCAmelCase ):
				__magic_name__       :   List[str]							= "encodec"




				def __init__(	self    :				List[str]   ,					lowerCAmelCase    :				int=[1.5, 3.0, 6.0, 12.0, 24.0]   ,					lowerCAmelCase    :				Tuple=24000   ,					lowerCAmelCase    :				List[Any]=1   ,					lowerCAmelCase    :				Union[str, Any]=False   ,					lowerCAmelCase    :				str=None   ,					lowerCAmelCase    :				Dict=None   ,					lowerCAmelCase    :				str=128   ,					lowerCAmelCase    :				Any=32   ,					lowerCAmelCase    :				Any=1   ,					lowerCAmelCase    :				List[Any]=[8, 5, 4, 2]   ,					lowerCAmelCase    :				Union[str, Any]="weight_norm"   ,					lowerCAmelCase    :				str=7   ,					lowerCAmelCase    :				Optional[int]=7   ,					lowerCAmelCase    :				Any=3   ,					lowerCAmelCase    :				Tuple=2   ,					lowerCAmelCase    :				Optional[Any]=True   ,					lowerCAmelCase    :				List[str]="reflect"   ,					lowerCAmelCase    :				Optional[int]=2   ,					lowerCAmelCase    :				List[Any]=2   ,					lowerCAmelCase    :				Union[str, Any]=1.0   ,					lowerCAmelCase    :				Optional[Any]=1024   ,					lowerCAmelCase    :				Optional[Any]=None   ,					lowerCAmelCase    :				str=True   ,					**lowerCAmelCase    :				str   ,					)->  List[str]:

										"""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__(**lowerCAmelCase   )




				@property
				def   a__(	self    :				str   )->  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   a__(	self    :				List[str]   )->  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   a__(	self    :				List[Any]   )->  int:

										"""simple docstring"""




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





				@property
				def   a__(	self    :				int   )->  int:

										"""simple docstring"""




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





