


'''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       a				(      __a       )  ->     List[str]:



					'''simple docstring'''






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

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

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

					# set label information
					UpperCamelCase__       :str							 =			'''huggingface/label-files'''
					if "large" in model_name or "huge" in model_name:
										UpperCamelCase__       :Tuple							 =			'''imagenet-22k-id2label.json'''
					else:
										UpperCamelCase__       :Union[str, Any]							 =			'''imagenet-1k-id2label.json'''

					UpperCamelCase__       :List[str]							 =			json.load(open(hf_hub_download(__lowerCAmelCase			, __lowerCAmelCase			, repo_type='''dataset'''       )			, '''r'''       )       )
					UpperCamelCase__       :List[Any]							 =			{int(__lowerCAmelCase       ): v for k, v in idalabel.items()}
					UpperCamelCase__       :List[str]							 =			{v: k for k, v in idalabel.items()}

					UpperCamelCase__       :Any							 =			FocalNetConfig(
					    embed_dim=__lowerCAmelCase			, depths=__lowerCAmelCase			, focal_levels=__lowerCAmelCase			, focal_windows=__lowerCAmelCase			, use_conv_embed=__lowerCAmelCase			, idalabel=__lowerCAmelCase			, labelaid=__lowerCAmelCase			, use_post_layernorm=__lowerCAmelCase			, use_layerscale=__lowerCAmelCase			, )

					return config


def       a				(      __a       )  ->     Union[str, Any]:



					'''simple docstring'''






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

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

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

					return name


def       a				(      __a			, __a			, __a=False       )  ->     Any:



					'''simple docstring'''






					UpperCamelCase__       :Optional[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

					UpperCamelCase__       :Tuple							 =			model_name_to_url[model_name]
					print('''Checkpoint URL: '''			, __lowerCAmelCase       )
					UpperCamelCase__       :Any							 =			torch.hub.load_state_dict_from_url(__lowerCAmelCase			, map_location='''cpu'''       )['''model''']

					# rename keys
					for key in state_dict.copy().keys():
										UpperCamelCase__       :Optional[Any]							 =			state_dict.pop(__lowerCAmelCase       )
										UpperCamelCase__       :List[str]							 =			val

					UpperCamelCase__       :List[str]							 =			get_focalnet_config(__lowerCAmelCase       )
					UpperCamelCase__       :Dict							 =			FocalNetForImageClassification(__lowerCAmelCase       )
					model.eval()

					# load state dict
					model.load_state_dict(__lowerCAmelCase       )

					# verify conversion
					UpperCamelCase__       :List[str]							 =			'''http://images.cocodataset.org/val2017/000000039769.jpg'''

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

					UpperCamelCase__       :int							 =			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]       ),
					    ]       )

					UpperCamelCase__       :Optional[Any]							 =			image_transforms(__lowerCAmelCase       ).unsqueeze(0       )

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

					UpperCamelCase__       :List[Any]							 =			model(**__lowerCAmelCase       )

					UpperCamelCase__       :Dict							 =			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__       :Any							 =			torch.tensor([0.2_1_6_6, -0.4_3_6_8, 0.2_1_9_1]       )
					elif model_name == "focalnet-tiny-lrf":
										UpperCamelCase__       :List[Any]							 =			torch.tensor([1.1_6_6_9, 0.0_1_2_5, -0.1_6_9_5]       )
					elif model_name == "focalnet-small":
										UpperCamelCase__       :Dict							 =			torch.tensor([0.4_9_1_7, -0.0_4_3_0, 0.1_3_4_1]       )
					elif model_name == "focalnet-small-lrf":
										UpperCamelCase__       :Tuple							 =			torch.tensor([-0.2_5_8_8, -0.5_3_4_2, -0.2_3_3_1]       )
					elif model_name == "focalnet-base":
										UpperCamelCase__       :Tuple							 =			torch.tensor([-0.1_6_5_5, -0.4_0_9_0, -0.1_7_3_0]       )
					elif model_name == "focalnet-base-lrf":
										UpperCamelCase__       :List[str]							 =			torch.tensor([0.5_3_0_6, -0.0_4_8_3, -0.3_9_2_8]       )
					assert torch.allclose(outputs.logits[0, :3]			, __lowerCAmelCase			, 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(__lowerCAmelCase       )
										processor.save_pretrained(__lowerCAmelCase       )

					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__":
		__snake_case											=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.''',
		)

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



'''simple docstring'''
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class        lowercase     (							A__     ,		A__     ,		A__     ,		unittest.TestCase ):

		"""simple docstring"""

		_a   			=    AltDiffusionPipeline
		_a   			=    TEXT_TO_IMAGE_PARAMS
		_a   			=    TEXT_TO_IMAGE_BATCH_PARAMS
		_a   			=    TEXT_TO_IMAGE_IMAGE_PARAMS
		_a   			=    TEXT_TO_IMAGE_IMAGE_PARAMS




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							torch.manual_seed(0							)
							UpperCamelCase__       :List[str]							 =			UNetaDConditionModel(
							    block_out_channels=(32, 64)   ,       layers_per_block=2   ,       sample_size=32   ,       in_channels=4   ,       out_channels=4   ,       down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,       up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,       cross_attention_dim=32   ,       )
							UpperCamelCase__       :int							 =			DDIMScheduler(
							    beta_start=0.00085   ,       beta_end=0.012   ,       beta_schedule='''scaled_linear'''   ,       clip_sample=UpperCamelCase_   ,       set_alpha_to_one=UpperCamelCase_   ,       )
							torch.manual_seed(0							)
							UpperCamelCase__       :List[str]							 =			AutoencoderKL(
							    block_out_channels=[32, 64]   ,       in_channels=3   ,       out_channels=3   ,       down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,       up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,       latent_channels=4   ,       )

							# TODO: address the non-deterministic text encoder (fails for save-load tests)
							# torch.manual_seed(0)
							# text_encoder_config = RobertaSeriesConfig(
							#     hidden_size=32,
							#     project_dim=32,
							#     intermediate_size=37,
							#     layer_norm_eps=1e-05,
							#     num_attention_heads=4,
							#     num_hidden_layers=5,
							#     vocab_size=5002,
							# )
							# text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

							torch.manual_seed(0							)
							UpperCamelCase__       :Union[str, Any]							 =			CLIPTextConfig(
							    bos_token_id=0   ,       eos_token_id=2   ,       hidden_size=32   ,       projection_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       pad_token_id=1   ,       vocab_size=5002   ,       )
							UpperCamelCase__       :Optional[Any]							 =			CLIPTextModel(UpperCamelCase_							)

							UpperCamelCase__       :int							 =			XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''							)
							UpperCamelCase__       :Dict							 =			77

							UpperCamelCase__       :Dict							 =			{
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''vae''': vae,
							    '''text_encoder''': text_encoder,
							    '''tokenizer''': tokenizer,
							    '''safety_checker''': None,
							    '''feature_extractor''': None,
							}
							return components




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_=0							):

							'''simple docstring'''


							if str(UpperCamelCase_							).startswith('''mps'''							):
												UpperCamelCase__       :Dict							 =			torch.manual_seed(UpperCamelCase_							)
							else:
												UpperCamelCase__       :List[str]							 =			torch.Generator(device=UpperCamelCase_							).manual_seed(UpperCamelCase_							)
							UpperCamelCase__       :Optional[int]							 =			{
							    '''prompt''': '''A painting of a squirrel eating a burger''',
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}
							return inputs




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							super().test_attention_slicing_forward_pass(expected_max_diff=3e-3							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							super().test_inference_batch_single_identical(expected_max_diff=3e-3							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :str							 =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator

							UpperCamelCase__       :Tuple							 =			self.get_dummy_components()
							torch.manual_seed(0							)
							UpperCamelCase__       :Dict							 =			RobertaSeriesConfig(
							    hidden_size=32   ,       project_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       vocab_size=5002   ,       )
							# TODO: remove after fixing the non-deterministic text encoder
							UpperCamelCase__       :Union[str, Any]							 =			RobertaSeriesModelWithTransformation(UpperCamelCase_							)
							UpperCamelCase__       :Union[str, Any]							 =			text_encoder

							UpperCamelCase__       :Any							 =			AltDiffusionPipeline(**UpperCamelCase_							)
							UpperCamelCase__       :List[Any]							 =			alt_pipe.to(UpperCamelCase_							)
							alt_pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__       :int							 =			self.get_dummy_inputs(UpperCamelCase_							)
							UpperCamelCase__       :Dict							 =			'''A photo of an astronaut'''
							UpperCamelCase__       :Union[str, Any]							 =			alt_pipe(**UpperCamelCase_							)
							UpperCamelCase__       :Any							 =			output.images
							UpperCamelCase__       :List[Any]							 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							UpperCamelCase__       :List[Any]							 =			np.array(
							    [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :str							 =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator

							UpperCamelCase__       :int							 =			self.get_dummy_components()
							UpperCamelCase__       :Union[str, Any]							 =			PNDMScheduler(skip_prk_steps=UpperCamelCase_							)
							torch.manual_seed(0							)
							UpperCamelCase__       :Dict							 =			RobertaSeriesConfig(
							    hidden_size=32   ,       project_dim=32   ,       intermediate_size=37   ,       layer_norm_eps=1e-05   ,       num_attention_heads=4   ,       num_hidden_layers=5   ,       vocab_size=5002   ,       )
							# TODO: remove after fixing the non-deterministic text encoder
							UpperCamelCase__       :str							 =			RobertaSeriesModelWithTransformation(UpperCamelCase_							)
							UpperCamelCase__       :Union[str, Any]							 =			text_encoder
							UpperCamelCase__       :Union[str, Any]							 =			AltDiffusionPipeline(**UpperCamelCase_							)
							UpperCamelCase__       :Optional[int]							 =			alt_pipe.to(UpperCamelCase_							)
							alt_pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__       :List[str]							 =			self.get_dummy_inputs(UpperCamelCase_							)
							UpperCamelCase__       :Optional[Any]							 =			alt_pipe(**UpperCamelCase_							)
							UpperCamelCase__       :str							 =			output.images
							UpperCamelCase__       :int							 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 64, 64, 3)
							UpperCamelCase__       :List[Any]							 =			np.array(
							    [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2



@slow
@require_torch_gpu
class        lowercase     (							unittest.TestCase ):

		"""simple docstring"""

		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Optional[Any]							 =			AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,       safety_checker=UpperCamelCase_							)
							UpperCamelCase__       :str							 =			alt_pipe.to(UpperCamelCase_							)
							alt_pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__       :int							 =			'''A painting of a squirrel eating a burger'''
							UpperCamelCase__       :int							 =			torch.manual_seed(0							)
							UpperCamelCase__       :Optional[Any]							 =			alt_pipe([prompt]   ,       generator=UpperCamelCase_   ,       guidance_scale=6.0   ,       num_inference_steps=20   ,       output_type='''np'''							)

							UpperCamelCase__       :Union[str, Any]							 =			output.images

							UpperCamelCase__       :Optional[Any]							 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							UpperCamelCase__       :Union[str, Any]							 =			np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Any							 =			DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''   ,       subfolder='''scheduler'''							)

							UpperCamelCase__       :List[str]							 =			AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''   ,       scheduler=UpperCamelCase_   ,       safety_checker=UpperCamelCase_							)
							UpperCamelCase__       :Dict							 =			alt_pipe.to(UpperCamelCase_							)
							alt_pipe.set_progress_bar_config(disable=UpperCamelCase_							)

							UpperCamelCase__       :List[Any]							 =			'''A painting of a squirrel eating a burger'''
							UpperCamelCase__       :str							 =			torch.manual_seed(0							)

							UpperCamelCase__       :Optional[int]							 =			alt_pipe([prompt]   ,       generator=UpperCamelCase_   ,       num_inference_steps=2   ,       output_type='''numpy'''							)
							UpperCamelCase__       :Any							 =			output.images

							UpperCamelCase__       :int							 =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 512, 512, 3)
							UpperCamelCase__       :int							 =			np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]							)

							assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-2