


"""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     lowercase			(SCREAMING_SNAKE_CASE_		:				int  )   ->						Optional[Any]:
						SCREAMING_SNAKE_CASE													=					[2, 2, 6, 2] if 'tiny' in model_name else [2, 2, 18, 2]
						SCREAMING_SNAKE_CASE													=					True if 'large' in model_name or 'huge' in model_name else False
						SCREAMING_SNAKE_CASE													=					True if 'large' in model_name or 'huge' in model_name else False
						SCREAMING_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:
																		SCREAMING_SNAKE_CASE													=					[3, 3, 3, 3]
																		SCREAMING_SNAKE_CASE													=					[5, 5, 5, 5]
												elif "fl4" in model_name:
																		SCREAMING_SNAKE_CASE													=					[4, 4, 4, 4]
																		SCREAMING_SNAKE_CASE													=					[3, 3, 3, 3]

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

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

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

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

						SCREAMING_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     lowercase			(SCREAMING_SNAKE_CASE_		:				int  )   ->						Optional[int]:
						if "patch_embed.proj" in name:
												SCREAMING_SNAKE_CASE													=					name.replace('patch_embed.proj' ,			'embeddings.patch_embeddings.projection'  )
						if "patch_embed.norm" in name:
												SCREAMING_SNAKE_CASE													=					name.replace('patch_embed.norm' ,			'embeddings.norm'  )
						if "layers" in name:
												SCREAMING_SNAKE_CASE													=					'encoder.' + name
						if "encoder.layers" in name:
												SCREAMING_SNAKE_CASE													=					name.replace('encoder.layers' ,			'encoder.stages'  )
						if "downsample.proj" in name:
												SCREAMING_SNAKE_CASE													=					name.replace('downsample.proj' ,			'downsample.projection'  )
						if "blocks" in name:
												SCREAMING_SNAKE_CASE													=					name.replace('blocks' ,			'layers'  )
						if "modulation.f.weight" in name or "modulation.f.bias" in name:
												SCREAMING_SNAKE_CASE													=					name.replace('modulation.f' ,			'modulation.projection_in'  )
						if "modulation.h.weight" in name or "modulation.h.bias" in name:
												SCREAMING_SNAKE_CASE													=					name.replace('modulation.h' ,			'modulation.projection_context'  )
						if "modulation.proj.weight" in name or "modulation.proj.bias" in name:
												SCREAMING_SNAKE_CASE													=					name.replace('modulation.proj' ,			'modulation.projection_out'  )

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

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

						return name


def     lowercase			(SCREAMING_SNAKE_CASE_		:				List[str] ,			SCREAMING_SNAKE_CASE_		:				Optional[int] ,			SCREAMING_SNAKE_CASE_		:				List[str]=False  )   ->						List[Any]:
						# fmt: off
						SCREAMING_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

						SCREAMING_SNAKE_CASE													=					model_name_to_url[model_name]
						print('Checkpoint URL: ' ,			SCREAMING_SNAKE_CASE_  )
						SCREAMING_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():
												SCREAMING_SNAKE_CASE													=					state_dict.pop(SCREAMING_SNAKE_CASE_  )
												SCREAMING_SNAKE_CASE													=					val

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

						# load state dict
						model.load_state_dict(SCREAMING_SNAKE_CASE_  )

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

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

						SCREAMING_SNAKE_CASE													=					transforms.Compose(
						    [
						        transforms.Resize(2_56  ),
						        transforms.CenterCrop(2_24  ),
						        transforms.ToTensor(),
						        transforms.Normalize(mean=[0.4_85, 0.4_56, 0.4_06] ,			std=[0.2_29, 0.2_24, 0.2_25]  ),
						    ]  )

						SCREAMING_SNAKE_CASE													=					image_transforms(SCREAMING_SNAKE_CASE_  ).unsqueeze(0  )

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

						SCREAMING_SNAKE_CASE													=					model(**SCREAMING_SNAKE_CASE_  )

						SCREAMING_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":
												SCREAMING_SNAKE_CASE													=					torch.tensor([0.21_66, -0.43_68, 0.21_91]  )
						elif model_name == "focalnet-tiny-lrf":
												SCREAMING_SNAKE_CASE													=					torch.tensor([1.16_69, 0.01_25, -0.16_95]  )
						elif model_name == "focalnet-small":
												SCREAMING_SNAKE_CASE													=					torch.tensor([0.49_17, -0.04_30, 0.13_41]  )
						elif model_name == "focalnet-small-lrf":
												SCREAMING_SNAKE_CASE													=					torch.tensor([-0.25_88, -0.53_42, -0.23_31]  )
						elif model_name == "focalnet-base":
												SCREAMING_SNAKE_CASE													=					torch.tensor([-0.16_55, -0.40_90, -0.17_30]  )
						elif model_name == "focalnet-base-lrf":
												SCREAMING_SNAKE_CASE													=					torch.tensor([0.53_06, -0.04_83, -0.39_28]  )
						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)






"""simple docstring"""




import unittest

import numpy as np

from diffusers import OnnxStableDiffusionInpaintPipelineLegacy
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    load_numpy,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)


if is_onnx_available():
	import onnxruntime as ort
@nightly
@require_onnxruntime
@require_torch_gpu
class 					lowerCAmelCase	(					unittest.TestCase    ):




		'''simple docstring'''







		@property
		def 							__A						(	self				)     ->							List[Any]:
								return (
								    "CUDAExecutionProvider",
								    {
								        "gpu_mem_limit": "15000000000",  # 15GB
								        "arena_extend_strategy": "kSameAsRequested",
								    },
								)







		@property
		def 							__A						(	self				)     ->							List[str]:
								SCREAMING_SNAKE_CASE													=					ort.SessionOptions()
								SCREAMING_SNAKE_CASE													=					False
								return options







		def 							__A						(	self				)     ->							Dict:
								SCREAMING_SNAKE_CASE													=					load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/overture-creations-5sI6fQgYIuo.png'				)
								SCREAMING_SNAKE_CASE													=					load_image(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/overture-creations-5sI6fQgYIuo_mask.png'				)
								SCREAMING_SNAKE_CASE													=					load_numpy(
								    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
								    '/in_paint/red_cat_sitting_on_a_park_bench_onnx.npy'				)

								# using the PNDM scheduler by default
								SCREAMING_SNAKE_CASE													=					OnnxStableDiffusionInpaintPipelineLegacy.from_pretrained(
								    'CompVis/stable-diffusion-v1-4'						,       revision='onnx'						,       safety_checker=lowerCAmelCase__						,       feature_extractor=lowerCAmelCase__						,       provider=self.gpu_provider						,       sess_options=self.gpu_options						,       )
								pipe.set_progress_bar_config(disable=lowerCAmelCase__				)

								SCREAMING_SNAKE_CASE													=					'A red cat sitting on a park bench'

								SCREAMING_SNAKE_CASE													=					np.random.RandomState(0				)
								SCREAMING_SNAKE_CASE													=					pipe(
								    prompt=lowerCAmelCase__						,       image=lowerCAmelCase__						,       mask_image=lowerCAmelCase__						,       strength=0.75						,       guidance_scale=7.5						,       num_inference_steps=15						,       generator=lowerCAmelCase__						,       output_type='np'						,       )

								SCREAMING_SNAKE_CASE													=					output.images[0]

								assert image.shape == (512, 512, 3)
								assert np.abs(expected_image - image				).max() < 1e-2


