



import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    UNetaDConditionModel,
    VideoToVideoSDPipeline,
)
from diffusers.utils import floats_tensor, is_xformers_available, skip_mps
from diffusers.utils.testing_utils import enable_full_determinism, slow, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()

@skip_mps
class        UpperCAmelCase_  (       UpperCamelCase_   ,   unittest.TestCase	):

						'''simple docstring'''




						UpperCamelCase__      :			int      =  VideoToVideoSDPipeline
						UpperCamelCase__      :			Union[str, Any]      =  TEXT_GUIDED_IMAGE_VARIATION_PARAMS.union({'''video'''}	) - {'''image''', '''width''', '''height'''}
						UpperCamelCase__      :			List[Any]      =  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'''video'''}	) - {'''image'''}
						UpperCamelCase__      :			str      =  PipelineTesterMixin.required_optional_params - {'''latents'''}
						UpperCamelCase__      :			Optional[Any]      =  False

						# No `output_type`.
						UpperCamelCase__      :			str      =  frozenset(
						    [
						        '''num_inference_steps''',
						        '''generator''',
						        '''latents''',
						        '''return_dict''',
						        '''callback''',
						        '''callback_steps''',
						    ]	)






						def 	_A     (   self		):







													'''simple docstring'''

													torch.manual_seed(0		)
													__SCREAMING_SNAKE_CASE									= UNetaDConditionModel(
													    block_out_channels=(32, 64, 64, 64)				,							layers_per_block=2				,							sample_size=32				,							in_channels=4				,							out_channels=4				,							down_block_types=('CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'CrossAttnDownBlock3D', 'DownBlock3D')				,							up_block_types=('UpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D', 'CrossAttnUpBlock3D')				,							cross_attention_dim=32				,							attention_head_dim=4				,							)
													__SCREAMING_SNAKE_CASE									= DDIMScheduler(
													    beta_start=0.0_0_0_8_5				,							beta_end=0.0_1_2				,							beta_schedule='scaled_linear'				,							clip_sample=_A				,							set_alpha_to_one=_A				,							)
													torch.manual_seed(0		)
													__SCREAMING_SNAKE_CASE									= 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				,							sample_size=128				,							)
													torch.manual_seed(0		)
													__SCREAMING_SNAKE_CASE									= CLIPTextConfig(
													    bos_token_id=0				,							eos_token_id=2				,							hidden_size=32				,							intermediate_size=37				,							layer_norm_eps=1e-05				,							num_attention_heads=4				,							num_hidden_layers=5				,							pad_token_id=1				,							vocab_size=1_000				,							hidden_act='gelu'				,							projection_dim=512				,							)
													__SCREAMING_SNAKE_CASE									= CLIPTextModel(_A		)
													__SCREAMING_SNAKE_CASE									= CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'		)

													__SCREAMING_SNAKE_CASE									= {
													    'unet': unet,
													    'scheduler': scheduler,
													    'vae': vae,
													    'text_encoder': text_encoder,
													    'tokenizer': tokenizer,
													}
													return components






						def 	_A     (   self				,							_A				,							_A=0		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= floats_tensor((1, 3, 3, 32, 32)				,							rng=random.Random(_A		)		).to(_A		)

													if str(_A		).startswith('mps'		):
																				__SCREAMING_SNAKE_CASE									= torch.manual_seed(_A		)
													else:
																				__SCREAMING_SNAKE_CASE									= torch.Generator(device=_A		).manual_seed(_A		)
													__SCREAMING_SNAKE_CASE									= {
													    'prompt': 'A painting of a squirrel eating a burger',
													    'video': video,
													    'generator': generator,
													    'num_inference_steps': 2,
													    'guidance_scale': 6.0,
													    'output_type': 'pt',
													}
													return inputs






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= 'cpu'  # ensure determinism for the device-dependent torch.Generator
													__SCREAMING_SNAKE_CASE									= self.get_dummy_components()
													__SCREAMING_SNAKE_CASE									= VideoToVideoSDPipeline(**_A		)
													__SCREAMING_SNAKE_CASE									= sd_pipe.to(_A		)
													sd_pipe.set_progress_bar_config(disable=_A		)

													__SCREAMING_SNAKE_CASE									= self.get_dummy_inputs(_A		)
													__SCREAMING_SNAKE_CASE									= 'np'
													__SCREAMING_SNAKE_CASE									= sd_pipe(**_A		).frames
													__SCREAMING_SNAKE_CASE									= frames[0][-3:, -3:, -1]

													assert frames[0].shape == (32, 32, 3)
													__SCREAMING_SNAKE_CASE									= np.array([106, 117, 113, 174, 137, 112, 148, 151, 131]		)

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






						@unittest.skipIf(
						    torch_device != 'cuda' or not is_xformers_available()				,							reason='XFormers attention is only available with CUDA and `xformers` installed'				,							)
						def 	_A     (   self		):







													'''simple docstring'''

													self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=_A				,							expected_max_diff=5e-3		)






						@unittest.skip(reason='Batching needs to be properly figured out first for this pipeline.'		)
						def 	_A     (   self		):







													'''simple docstring'''

													pass






						@unittest.skip(reason='Batching needs to be properly figured out first for this pipeline.'		)
						def 	_A     (   self		):







													'''simple docstring'''

													pass






						@unittest.skip(reason='`num_images_per_prompt` argument is not supported for this pipeline.'		)
						def 	_A     (   self		):







													'''simple docstring'''

													pass






						def 	_A     (   self		):







													'''simple docstring'''

													return super().test_progress_bar()







@slow
@skip_mps
class        UpperCAmelCase_  (       unittest.TestCase	):

						'''simple docstring'''




						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= VideoToVideoSDPipeline.from_pretrained('cerspense/zeroscope_v2_XL'				,							torch_dtype=torch.floataa		)
													pipe.enable_model_cpu_offload()

													# 10 frames
													__SCREAMING_SNAKE_CASE									= torch.Generator(device='cpu'		).manual_seed(0		)
													__SCREAMING_SNAKE_CASE									= torch.randn((1, 10, 3, 1_024, 576)				,							generator=_A		)
													__SCREAMING_SNAKE_CASE									= video.to('cuda'		)

													__SCREAMING_SNAKE_CASE									= 'Spiderman is surfing'

													__SCREAMING_SNAKE_CASE									= pipe(_A				,							video=_A				,							generator=_A				,							num_inference_steps=3				,							output_type='pt'		).frames

													__SCREAMING_SNAKE_CASE									= np.array([-1.0_4_5_8_9_8_4, -1.1_2_7_9_2_9_7, -0.9_6_6_3_0_8_6, -0.9_1_5_0_3_9_0_6, -0.7_5_0_9_7_6_5_6]		)
													assert np.abs(video_frames.cpu().numpy()[0, 0, 0, 0, -5:] - expected_array		).sum() < 1e-2











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       (			a__       )    ->			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(a__							,       a__							,       repo_type='dataset'       )							,       'r'       )       )
							__SCREAMING_SNAKE_CASE									= {int(a__       ): v for k, v in idalabel.items()}
							__SCREAMING_SNAKE_CASE									= {v: k for k, v in idalabel.items()}

							__SCREAMING_SNAKE_CASE									= FocalNetConfig(
							    embed_dim=a__							,       depths=a__							,       focal_levels=a__							,       focal_windows=a__							,       use_conv_embed=a__							,       idalabel=a__							,       labelaid=a__							,       use_post_layernorm=a__							,       use_layerscale=a__							,       )

							return config






def      __lowercase       (			a__       )    ->			Any:
							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       (			a__							,       a__							,       a__=False       )    ->			Dict:
							# 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: '							,       a__       )
							__SCREAMING_SNAKE_CASE									= torch.hub.load_state_dict_from_url(a__							,       map_location='cpu'       )['model']

							# rename keys
							for key in state_dict.copy().keys():
														__SCREAMING_SNAKE_CASE									= state_dict.pop(a__       )
														__SCREAMING_SNAKE_CASE									= val

							__SCREAMING_SNAKE_CASE									= get_focalnet_config(a__       )
							__SCREAMING_SNAKE_CASE									= FocalNetForImageClassification(a__       )
							model.eval()

							# load state dict
							model.load_state_dict(a__       )

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

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

							__SCREAMING_SNAKE_CASE									= 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]       ),
							    ]       )

							__SCREAMING_SNAKE_CASE									= image_transforms(a__       ).unsqueeze(0       )

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

							__SCREAMING_SNAKE_CASE									= model(**a__       )

							__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.2166, -0.4368, 0.2191]       )
							elif model_name == "focalnet-tiny-lrf":
														__SCREAMING_SNAKE_CASE									= torch.tensor([1.1669, 0.0125, -0.1695]       )
							elif model_name == "focalnet-small":
														__SCREAMING_SNAKE_CASE									= torch.tensor([0.4917, -0.0430, 0.1341]       )
							elif model_name == "focalnet-small-lrf":
														__SCREAMING_SNAKE_CASE									= torch.tensor([-0.2588, -0.5342, -0.2331]       )
							elif model_name == "focalnet-base":
														__SCREAMING_SNAKE_CASE									= torch.tensor([-0.1655, -0.4090, -0.1730]       )
							elif model_name == "focalnet-base-lrf":
														__SCREAMING_SNAKE_CASE									= torch.tensor([0.5306, -0.0483, -0.3928]       )
							assert torch.allclose(outputs.logits[0, :3]							,       a__							,       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(a__       )
														processor.save_pretrained(a__       )

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

			lowerCAmelCase__		:					List[Any]							=parser.parse_args()
			convert_focalnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






