
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 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Optional[Any]				)			->						Any:
       SCREAMING_SNAKE_CASE_							:				Optional[int]		=[2, 2, 6, 2] if '''tiny''' in model_name else [2, 2, 1_8, 2]
       SCREAMING_SNAKE_CASE_							:				Optional[Any]		=True if '''large''' in model_name or '''huge''' in model_name else False
       SCREAMING_SNAKE_CASE_							:				str		=True if '''large''' in model_name or '''huge''' in model_name else False
       SCREAMING_SNAKE_CASE_							:				Tuple		=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:
                     SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=[3, 3, 3, 3]
                     SCREAMING_SNAKE_CASE_							:				Optional[Any]		=[5, 5, 5, 5]
              elif "fl4" in model_name:
                     SCREAMING_SNAKE_CASE_							:				Dict		=[4, 4, 4, 4]
                     SCREAMING_SNAKE_CASE_							:				str		=[3, 3, 3, 3]

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

       if "tiny" in model_name:
              SCREAMING_SNAKE_CASE_							:				int		=9_6
       elif "small" in model_name:
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=9_6
       elif "base" in model_name:
              SCREAMING_SNAKE_CASE_							:				int		=1_2_8
       elif "large" in model_name:
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=1_9_2
       elif "xlarge" in model_name:
              SCREAMING_SNAKE_CASE_							:				Dict		=2_5_6
       elif "huge" in model_name:
              SCREAMING_SNAKE_CASE_							:				str		=3_5_2

       # set label information
       SCREAMING_SNAKE_CASE_							:				Any		='''huggingface/label-files'''
       if "large" in model_name or "huge" in model_name:
              SCREAMING_SNAKE_CASE_							:				List[str]		='''imagenet-22k-id2label.json'''
       else:
              SCREAMING_SNAKE_CASE_							:				Optional[int]		='''imagenet-1k-id2label.json'''

       SCREAMING_SNAKE_CASE_							:				Optional[Any]		=json.load(open(hf_hub_download(UpperCAmelCase_			,							UpperCAmelCase_			,							repo_type='''dataset'''				)			,							'''r'''				)				)
       SCREAMING_SNAKE_CASE_							:				int		={int(UpperCAmelCase_				): v for k, v in idalabel.items()}
       SCREAMING_SNAKE_CASE_							:				Tuple		={v: k for k, v in idalabel.items()}

       SCREAMING_SNAKE_CASE_							:				str		=FocalNetConfig(
           embed_dim=UpperCAmelCase_			,							depths=UpperCAmelCase_			,							focal_levels=UpperCAmelCase_			,							focal_windows=UpperCAmelCase_			,							use_conv_embed=UpperCAmelCase_			,							idalabel=UpperCAmelCase_			,							labelaid=UpperCAmelCase_			,							use_post_layernorm=UpperCAmelCase_			,							use_layerscale=UpperCAmelCase_			,							)

       return config
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Dict				)			->						Dict:
       if "patch_embed.proj" in name:
              SCREAMING_SNAKE_CASE_							:				List[Any]		=name.replace('''patch_embed.proj'''			,							'''embeddings.patch_embeddings.projection'''				)
       if "patch_embed.norm" in name:
              SCREAMING_SNAKE_CASE_							:				List[Any]		=name.replace('''patch_embed.norm'''			,							'''embeddings.norm'''				)
       if "layers" in name:
              SCREAMING_SNAKE_CASE_							:				int		='''encoder.''' + name
       if "encoder.layers" in name:
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=name.replace('''encoder.layers'''			,							'''encoder.stages'''				)
       if "downsample.proj" in name:
              SCREAMING_SNAKE_CASE_							:				Tuple		=name.replace('''downsample.proj'''			,							'''downsample.projection'''				)
       if "blocks" in name:
              SCREAMING_SNAKE_CASE_							:				str		=name.replace('''blocks'''			,							'''layers'''				)
       if "modulation.f.weight" in name or "modulation.f.bias" in name:
              SCREAMING_SNAKE_CASE_							:				int		=name.replace('''modulation.f'''			,							'''modulation.projection_in'''				)
       if "modulation.h.weight" in name or "modulation.h.bias" in name:
              SCREAMING_SNAKE_CASE_							:				str		=name.replace('''modulation.h'''			,							'''modulation.projection_context'''				)
       if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=name.replace('''modulation.proj'''			,							'''modulation.projection_out'''				)

       if name == "norm.weight":
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		='''layernorm.weight'''
       if name == "norm.bias":
              SCREAMING_SNAKE_CASE_							:				Tuple		='''layernorm.bias'''

       if "head" in name:
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=name.replace('''head'''			,							'''classifier'''				)
       else:
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		='''focalnet.''' + name

       return name
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     str			,							UpperCAmelCase_	:     List[Any]			,							UpperCAmelCase_	:     Any=False				)			->						List[str]:
       # fmt: off
       SCREAMING_SNAKE_CASE_							:				Dict		={
           '''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

       SCREAMING_SNAKE_CASE_							:				int		=model_name_to_url[model_name]
       print('''Checkpoint URL: '''			,							UpperCAmelCase_				)
       SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=torch.hub.load_state_dict_from_url(UpperCAmelCase_			,							map_location='''cpu'''				)['''model''']

       # rename keys
       for key in state_dict.copy().keys():
              SCREAMING_SNAKE_CASE_							:				Dict		=state_dict.pop(UpperCAmelCase_				)
              SCREAMING_SNAKE_CASE_							:				List[Any]		=val

       SCREAMING_SNAKE_CASE_							:				List[str]		=get_focalnet_config(UpperCAmelCase_				)
       SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=FocalNetForImageClassification(UpperCAmelCase_				)
       model.eval()

       # load state dict
       model.load_state_dict(UpperCAmelCase_				)

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

       SCREAMING_SNAKE_CASE_							:				int		=BitImageProcessor(
           do_resize=UpperCAmelCase_			,							size={'''shortest_edge''': 2_5_6}			,							resample=PILImageResampling.BILINEAR			,							do_center_crop=UpperCAmelCase_			,							crop_size=2_2_4			,							do_normalize=UpperCAmelCase_			,							image_mean=UpperCAmelCase_			,							image_std=UpperCAmelCase_			,							)
       SCREAMING_SNAKE_CASE_							:				Tuple		=Image.open(requests.get(UpperCAmelCase_			,							stream=UpperCAmelCase_				).raw				)
       SCREAMING_SNAKE_CASE_							:				Tuple		=processor(images=UpperCAmelCase_			,							return_tensors='''pt'''				)

       SCREAMING_SNAKE_CASE_							:				Dict		=transforms.Compose(
           [
               transforms.Resize(2_5_6				),
               transforms.CenterCrop(2_2_4				),
               transforms.ToTensor(),
               transforms.Normalize(mean=[0.485, 0.456, 0.406]			,							std=[0.229, 0.224, 0.225]				),
           ]				)

       SCREAMING_SNAKE_CASE_							:				Any		=image_transforms(UpperCAmelCase_				).unsqueeze(0				)

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

       SCREAMING_SNAKE_CASE_							:				List[Any]		=model(**UpperCAmelCase_				)

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

       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											=    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											=    parser.parse_args()
 convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)


import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()


class 			lowercase_		(      A  ,      unittest.TestCase       ):
       __lowerCamelCase				       =						PriorTransformer
       __lowerCamelCase				       =						"hidden_states"


       @property
       def 	_snake_case				(    self				)						->    List[str]:
              SCREAMING_SNAKE_CASE_							:				List[str]		=4
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=8
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=7

              SCREAMING_SNAKE_CASE_							:				Dict		=floats_tensor((batch_size, embedding_dim)				).to(__A				)

              SCREAMING_SNAKE_CASE_							:				Dict		=floats_tensor((batch_size, embedding_dim)				).to(__A				)
              SCREAMING_SNAKE_CASE_							:				Dict		=floats_tensor((batch_size, num_embeddings, embedding_dim)				).to(__A				)

              return {
                  "hidden_states": hidden_states,
                  "timestep": 2,
                  "proj_embedding": proj_embedding,
                  "encoder_hidden_states": encoder_hidden_states,
              }


       def 	_snake_case				(    self						,       __A=0				)						->    int:
              torch.manual_seed(__A				)
              SCREAMING_SNAKE_CASE_							:				str		=4
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=8
              SCREAMING_SNAKE_CASE_							:				List[Any]		=7

              SCREAMING_SNAKE_CASE_							:				Tuple		=torch.randn((batch_size, embedding_dim)				).to(__A				)

              SCREAMING_SNAKE_CASE_							:				int		=torch.randn((batch_size, embedding_dim)				).to(__A				)
              SCREAMING_SNAKE_CASE_							:				List[Any]		=torch.randn((batch_size, num_embeddings, embedding_dim)				).to(__A				)

              return {
                  "hidden_states": hidden_states,
                  "timestep": 2,
                  "proj_embedding": proj_embedding,
                  "encoder_hidden_states": encoder_hidden_states,
              }


       @property
       def 	_snake_case				(    self				)						->    Union[str, Any]:
              return (4, 8)


       @property
       def 	_snake_case				(    self				)						->    int:
              return (4, 8)


       def 	_snake_case				(    self				)						->    List[Any]:
              SCREAMING_SNAKE_CASE_							:				List[Any]		={
                  '''num_attention_heads''': 2,
                  '''attention_head_dim''': 4,
                  '''num_layers''': 2,
                  '''embedding_dim''': 8,
                  '''num_embeddings''': 7,
                  '''additional_embeddings''': 4,
              }
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=self.dummy_input
              return init_dict, inputs_dict


       def 	_snake_case				(    self				)						->    Optional[int]:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				Any		=PriorTransformer.from_pretrained(
                  '''hf-internal-testing/prior-dummy'''						,       output_loading_info=__A				)
              self.assertIsNotNone(__A				)
              self.assertEqual(len(loading_info['''missing_keys''']				)						,       0				)

              model.to(__A				)
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=model(**self.dummy_input				)[0]

              assert hidden_states is not None, "Make sure output is not None"


       def 	_snake_case				(    self				)						->    str:
              SCREAMING_SNAKE_CASE_	,     SCREAMING_SNAKE_CASE_							:				List[Any]		=self.prepare_init_args_and_inputs_for_common()

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=self.model_class(**__A				)
              SCREAMING_SNAKE_CASE_							:				List[Any]		=inspect.signature(model.forward				)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=[*signature.parameters.keys()]

              SCREAMING_SNAKE_CASE_							:				Optional[int]		=['''hidden_states''', '''timestep''']
              self.assertListEqual(arg_names[:2]						,       __A				)





       def 	_snake_case				(    self				)						->    Dict:
              SCREAMING_SNAKE_CASE_							:				Dict		=PriorTransformer.from_pretrained('''hf-internal-testing/prior-dummy'''				)
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=model.to(__A				)

              if hasattr(__A						,       '''set_default_attn_processor'''				):
                     model.set_default_attn_processor()

              SCREAMING_SNAKE_CASE_							:				List[Any]		=self.get_dummy_seed_input()

              with torch.no_grad():
                     SCREAMING_SNAKE_CASE_							:				str		=model(**__A				)[0]

              SCREAMING_SNAKE_CASE_							:				Any		=output[0, :5].flatten().cpu()
              print(__A				)

              # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
              # the expected output slices are not the same for CPU and GPU.
              SCREAMING_SNAKE_CASE_							:				int		=torch.tensor([-1.3_436, -0.2_870, 0.7_538, 0.4_368, -0.0_239]				)
              self.assertTrue(torch_all_close(__A						,       __A						,       rtol=1e-2				)				)






@slow
class 			lowercase_		(      unittest.TestCase       ):


       def 	_snake_case				(    self						,       __A=1						,       __A=768						,       __A=77						,       __A=0				)						->    str:
              torch.manual_seed(__A				)
              SCREAMING_SNAKE_CASE_							:				Dict		=batch_size
              SCREAMING_SNAKE_CASE_							:				List[str]		=embedding_dim
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=num_embeddings

              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=torch.randn((batch_size, embedding_dim)				).to(__A				)

              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=torch.randn((batch_size, embedding_dim)				).to(__A				)
              SCREAMING_SNAKE_CASE_							:				List[str]		=torch.randn((batch_size, num_embeddings, embedding_dim)				).to(__A				)

              return {
                  "hidden_states": hidden_states,
                  "timestep": 2,
                  "proj_embedding": proj_embedding,
                  "encoder_hidden_states": encoder_hidden_states,
              }


       def 	_snake_case				(    self				)						->    Dict:
              # clean up the VRAM after each test
              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()





       @parameterized.expand(
           [
               # fmt: off
               [13, [-0.5_861,  0.1_283, -0.0_931,  0.0_882,  0.4_476,  0.1_329, -0.0_498,  0.0_640]],
               [37, [-0.4_913,  0.0_110, -0.0_483,  0.0_541,  0.4_954, -0.0_170,  0.0_354,  0.1_651]],
               # fmt: on
           ]				)
       def 	_snake_case				(    self						,       __A						,       __A				)						->    Optional[int]:
              SCREAMING_SNAKE_CASE_							:				Dict		=PriorTransformer.from_pretrained('''kandinsky-community/kandinsky-2-1-prior'''						,       subfolder='''prior'''				)
              model.to(__A				)
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=self.get_dummy_seed_input(seed=__A				)

              with torch.no_grad():
                     SCREAMING_SNAKE_CASE_							:				Dict		=model(**__A				)[0]

              assert list(sample.shape				) == [1, 768]

              SCREAMING_SNAKE_CASE_							:				Dict		=sample[0, :8].flatten().cpu()
              print(__A				)
              SCREAMING_SNAKE_CASE_							:				Optional[Any]		=torch.tensor(__A				)

              assert torch_all_close(__A						,       __A						,       atol=1e-3				)
