




'''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__				(      __lowercase				):
     snake_case			:      Dict        =			[2, 2, 6, 2] if """tiny""" in model_name else [2, 2, 18, 2]
     snake_case			:      Any        =			True if """large""" in model_name or """huge""" in model_name else False
     snake_case			:      Optional[Any]        =			True if """large""" in model_name or """huge""" in model_name else False
     snake_case			:      Dict        =			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			:      Optional[int]        =			[3, 3, 3, 3]
               snake_case			:      Tuple        =			[5, 5, 5, 5]
          elif "fl4" in model_name:
               snake_case			:      List[Any]        =			[4, 4, 4, 4]
               snake_case			:      Dict        =			[3, 3, 3, 3]

     if "tiny" in model_name or "small" in model_name or "base" in model_name:
          snake_case			:      Optional[Any]        =			[3, 3, 3, 3]
          if "lrf" in model_name:
               snake_case			:      List[Any]        =			[3, 3, 3, 3]
          else:
               snake_case			:      List[str]        =			[2, 2, 2, 2]

     if "tiny" in model_name:
          snake_case			:      Optional[int]        =			96
     elif "small" in model_name:
          snake_case			:      int        =			96
     elif "base" in model_name:
          snake_case			:      List[str]        =			128
     elif "large" in model_name:
          snake_case			:      Optional[Any]        =			192
     elif "xlarge" in model_name:
          snake_case			:      int        =			256
     elif "huge" in model_name:
          snake_case			:      Dict        =			352

     # set label information
     snake_case			:      Tuple        =			"""huggingface/label-files"""
     if "large" in model_name or "huge" in model_name:
          snake_case			:      Dict        =			"""imagenet-22k-id2label.json"""
     else:
          snake_case			:      List[Any]        =			"""imagenet-1k-id2label.json"""

     snake_case			:      Union[str, Any]        =			json.load(open(hf_hub_download(__lowercase     ,   __lowercase     ,   repo_type="""dataset"""				)     ,   """r"""				)				)
     snake_case			:      str        =			{int(__lowercase				): v for k, v in idalabel.items()}
     snake_case			:      Any        =			{v: k for k, v in idalabel.items()}

     snake_case			:      Dict        =			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__				(      __lowercase				):
     if "patch_embed.proj" in name:
          snake_case			:      Optional[int]        =			name.replace("""patch_embed.proj"""     ,   """embeddings.patch_embeddings.projection"""				)
     if "patch_embed.norm" in name:
          snake_case			:      Optional[Any]        =			name.replace("""patch_embed.norm"""     ,   """embeddings.norm"""				)
     if "layers" in name:
          snake_case			:      Optional[int]        =			"""encoder.""" + name
     if "encoder.layers" in name:
          snake_case			:      Optional[Any]        =			name.replace("""encoder.layers"""     ,   """encoder.stages"""				)
     if "downsample.proj" in name:
          snake_case			:      str        =			name.replace("""downsample.proj"""     ,   """downsample.projection"""				)
     if "blocks" in name:
          snake_case			:      Dict        =			name.replace("""blocks"""     ,   """layers"""				)
     if "modulation.f.weight" in name or "modulation.f.bias" in name:
          snake_case			:      Dict        =			name.replace("""modulation.f"""     ,   """modulation.projection_in"""				)
     if "modulation.h.weight" in name or "modulation.h.bias" in name:
          snake_case			:      List[Any]        =			name.replace("""modulation.h"""     ,   """modulation.projection_context"""				)
     if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
          snake_case			:      List[str]        =			name.replace("""modulation.proj"""     ,   """modulation.projection_out"""				)

     if name == "norm.weight":
          snake_case			:      List[str]        =			"""layernorm.weight"""
     if name == "norm.bias":
          snake_case			:      Tuple        =			"""layernorm.bias"""

     if "head" in name:
          snake_case			:      Dict        =			name.replace("""head"""     ,   """classifier"""				)
     else:
          snake_case			:      List[Any]        =			"""focalnet.""" + name

     return name



def        lowerCamelCase__				(      __lowercase     ,   __lowercase     ,   __lowercase=False				):
     # fmt: off
     snake_case			:      List[Any]        =			{
         """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			:      List[Any]        =			model_name_to_url[model_name]
     print("""Checkpoint URL: """     ,   __lowercase				)
     snake_case			:      Optional[int]        =			torch.hub.load_state_dict_from_url(__lowercase     ,   map_location="""cpu"""				)["""model"""]

     # rename keys
     for key in state_dict.copy().keys():
          snake_case			:      Union[str, Any]        =			state_dict.pop(__lowercase				)
          snake_case			:      List[Any]        =			val

     snake_case			:      Dict        =			get_focalnet_config(__lowercase				)
     snake_case			:      Union[str, Any]        =			FocalNetForImageClassification(__lowercase				)
     model.eval()

     # load state dict
     model.load_state_dict(__lowercase				)

     # verify conversion
     snake_case			:      Union[str, Any]        =			"""http://images.cocodataset.org/val2017/000000039769.jpg"""

     snake_case			:      Any        =			BitImageProcessor(
         do_resize=__lowercase     ,   size={"""shortest_edge""": 256}     ,   resample=PILImageResampling.BILINEAR     ,   do_center_crop=__lowercase     ,   crop_size=224     ,   do_normalize=__lowercase     ,   image_mean=__lowercase     ,   image_std=__lowercase     ,   )
     snake_case			:      Any        =			Image.open(requests.get(__lowercase     ,   stream=__lowercase				).raw				)
     snake_case			:      Optional[Any]        =			processor(images=__lowercase     ,   return_tensors="""pt"""				)

     snake_case			:      Dict        =			transforms.Compose(
         [
             transforms.Resize(256				),
             transforms.CenterCrop(224				),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406]     ,   std=[0.229, 0.224, 0.225]				),
         ]				)

     snake_case			:      Tuple        =			image_transforms(__lowercase				).unsqueeze(0				)

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

     snake_case			:      Dict        =			model(**__lowercase				)

     snake_case			:      Any        =			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			:      int        =			torch.tensor([0.2166, -0.4368, 0.2191]				)
     elif model_name == "focalnet-tiny-lrf":
          snake_case			:      int        =			torch.tensor([1.1669, 0.0125, -0.1695]				)
     elif model_name == "focalnet-small":
          snake_case			:      Tuple        =			torch.tensor([0.4917, -0.0430, 0.1341]				)
     elif model_name == "focalnet-small-lrf":
          snake_case			:      Union[str, Any]        =			torch.tensor([-0.2588, -0.5342, -0.2331]				)
     elif model_name == "focalnet-base":
          snake_case			:      List[Any]        =			torch.tensor([-0.1655, -0.4090, -0.1730]				)
     elif model_name == "focalnet-base-lrf":
          snake_case			:      Any        =			torch.tensor([0.5306, -0.0483, -0.3928]				)
     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      :   Tuple      =       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      :   Optional[Any]      =       parser.parse_args()
      convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









'''simple docstring'''







import os


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_doctest_list.py
lowercase      :   List[str]      =       """."""


if __name__ == "__main__":
      lowercase      :   List[str]      =       os.path.join(REPO_PATH, """utils/documentation_tests.txt""")
      lowercase      :   Union[str, Any]      =       []
      lowercase      :   List[str]      =       []
      with open(doctest_file_path) as fp:
            for line in fp:
                  lowercase      :   List[str]      =       line.strip()
                  lowercase      :   Optional[Any]      =       os.path.join(REPO_PATH, line)
                  if not (os.path.isfile(path) or os.path.isdir(path)):
                        non_existent_paths.append(line)
                  all_paths.append(path)
      if len(non_existent_paths) > 0:
            lowercase      :   int      =       """\n""".join(non_existent_paths)
            raise ValueError(F"""`utils/documentation_tests.txt` contains non-existent paths:\n{non_existent_paths}""")
      if all_paths != sorted(all_paths):
            raise ValueError("""Files in `utils/documentation_tests.txt` are not in alphabetical order.""")



