





"""simple docstring"""




import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)






def       _SCREAMING_SNAKE_CASE		(	_lowercase						:    List[Any]    ,       _lowercase						:    Tuple=False					) ->List[str]:






	'''simple docstring'''






	a      :		Dict	=				[]
	for i in range(config.num_hidden_layers					):
		# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
		rename_keys.append((F"""module.blocks.{i}.norm1.weight""", F"""vit.encoder.layer.{i}.layernorm_before.weight""")					)
		rename_keys.append((F"""module.blocks.{i}.norm1.bias""", F"""vit.encoder.layer.{i}.layernorm_before.bias""")					)
		rename_keys.append(
		    (F"""module.blocks.{i}.attn.proj.weight""", F"""vit.encoder.layer.{i}.attention.output.dense.weight""")					)
		rename_keys.append((F"""module.blocks.{i}.attn.proj.bias""", F"""vit.encoder.layer.{i}.attention.output.dense.bias""")					)
		rename_keys.append((F"""module.blocks.{i}.norm2.weight""", F"""vit.encoder.layer.{i}.layernorm_after.weight""")					)
		rename_keys.append((F"""module.blocks.{i}.norm2.bias""", F"""vit.encoder.layer.{i}.layernorm_after.bias""")					)
		rename_keys.append((F"""module.blocks.{i}.mlp.fc1.weight""", F"""vit.encoder.layer.{i}.intermediate.dense.weight""")					)
		rename_keys.append((F"""module.blocks.{i}.mlp.fc1.bias""", F"""vit.encoder.layer.{i}.intermediate.dense.bias""")					)
		rename_keys.append((F"""module.blocks.{i}.mlp.fc2.weight""", F"""vit.encoder.layer.{i}.output.dense.weight""")					)
		rename_keys.append((F"""module.blocks.{i}.mlp.fc2.bias""", F"""vit.encoder.layer.{i}.output.dense.bias""")					)

	# projection layer + position embeddings
	rename_keys.extend(
	    [
	        ("module.cls_token", "vit.embeddings.cls_token"),
	        ("module.patch_embed.proj.weight", "vit.embeddings.patch_embeddings.projection.weight"),
	        ("module.patch_embed.proj.bias", "vit.embeddings.patch_embeddings.projection.bias"),
	        ("module.pos_embed", "vit.embeddings.position_embeddings"),
	    ]					)

	if base_model:
		# layernorm + pooler
		rename_keys.extend(
		    [
		        ("module.norm.weight", "layernorm.weight"),
		        ("module.norm.bias", "layernorm.bias"),
		    ]					)

		# if just the base model, we should remove "vit" from all keys that start with "vit"
		a      :		int	=				[(pair[0], pair[1][4:]) if pair[1].startswith("vit"					) else pair for pair in rename_keys]
	else:
		# layernorm + classification head
		rename_keys.extend(
		    [
		        ("norm.weight", "vit.layernorm.weight"),
		        ("norm.bias", "vit.layernorm.bias"),
		        ("head.weight", "classifier.weight"),
		        ("head.bias", "classifier.bias"),
		    ]					)

	return rename_keys






def       _SCREAMING_SNAKE_CASE		(	_lowercase						:    Optional[int]    ,       _lowercase						:    int    ,       _lowercase						:    List[Any]=False					) ->int:






	'''simple docstring'''






	for i in range(config.num_hidden_layers					):
		if base_model:
			a      :		List[Any]	=				""
		else:
			a      :		int	=				"vit."
		# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
		a      :		Optional[int]	=				state_dict.pop(F"""module.blocks.{i}.attn.qkv.weight"""					)
		a      :		Tuple	=				state_dict.pop(F"""module.blocks.{i}.attn.qkv.bias"""					)
		# next, add query, keys and values (in that order) to the state dict
		a      :		int	=				in_proj_weight[
		    : config.hidden_size, :
		]
		a      :		Optional[Any]	=				in_proj_bias[: config.hidden_size]
		a      :		str	=				in_proj_weight[
		    config.hidden_size : config.hidden_size * 2, :
		]
		a      :		Optional[int]	=				in_proj_bias[
		    config.hidden_size : config.hidden_size * 2
		]
		a      :		Dict	=				in_proj_weight[
		    -config.hidden_size :, :
		]
		a      :		Optional[int]	=				in_proj_bias[-config.hidden_size :]






def       _SCREAMING_SNAKE_CASE		(	_lowercase						:    Any					) ->Any:






	'''simple docstring'''






	a      :		Union[str, Any]	=				["head.weight", "head.bias"]
	for k in ignore_keys:
		state_dict.pop(_lowercase    ,       _lowercase					)






def       _SCREAMING_SNAKE_CASE		(	_lowercase						:    Dict					) ->Tuple:






	'''simple docstring'''






	a      :		Optional[int]	=				[
	    "module.fc.fc1.weight",
	    "module.fc.fc1.bias",
	    "module.fc.bn1.weight",
	    "module.fc.bn1.bias",
	    "module.fc.bn1.running_mean",
	    "module.fc.bn1.running_var",
	    "module.fc.bn1.num_batches_tracked",
	    "module.fc.fc2.weight",
	    "module.fc.fc2.bias",
	    "module.fc.bn2.weight",
	    "module.fc.bn2.bias",
	    "module.fc.bn2.running_mean",
	    "module.fc.bn2.running_var",
	    "module.fc.bn2.num_batches_tracked",
	    "module.fc.fc3.weight",
	    "module.fc.fc3.bias",
	]
	for k in ignore_keys:
		state_dict.pop(_lowercase    ,       _lowercase					)






def       _SCREAMING_SNAKE_CASE		(	_lowercase						:    List[Any]    ,       _lowercase						:    str    ,       _lowercase						:    Any					) ->List[Any]:






	'''simple docstring'''






	a      :		Union[str, Any]	=				dct.pop(_lowercase					)
	a      :		Optional[Any]	=				val






def       _SCREAMING_SNAKE_CASE		(	_lowercase						:    str    ,       _lowercase						:    Dict					) ->List[str]:






	'''simple docstring'''






	a      :		Tuple	=				ViTMSNConfig()
	a      :		List[Any]	=				1000

	a      :		str	=				"datasets/huggingface/label-files"
	a      :		Optional[Any]	=				"imagenet-1k-id2label.json"
	a      :		str	=				json.load(open(hf_hub_download(_lowercase    ,       _lowercase					)    ,       "r"					)					)
	a      :		Optional[Any]	=				{int(_lowercase					): v for k, v in idalabel.items()}
	a      :		Dict	=				idalabel
	a      :		Union[str, Any]	=				{v: k for k, v in idalabel.items()}

	if "s16" in checkpoint_url:
		a      :		Any	=				384
		a      :		Any	=				1536
		a      :		Optional[int]	=				6
	elif "l16" in checkpoint_url:
		a      :		str	=				1024
		a      :		Any	=				4096
		a      :		str	=				24
		a      :		List[str]	=				16
		a      :		Any	=				0.1
	elif "b4" in checkpoint_url:
		a      :		Optional[Any]	=				4
	elif "l7" in checkpoint_url:
		a      :		Dict	=				7
		a      :		int	=				1024
		a      :		List[str]	=				4096
		a      :		Optional[int]	=				24
		a      :		int	=				16
		a      :		List[Any]	=				0.1

	a      :		List[str]	=				ViTMSNModel(_lowercase					)

	a      :		Union[str, Any]	=				torch.hub.load_state_dict_from_url(_lowercase    ,       map_location="cpu"					)["target_encoder"]

	a      :		List[str]	=				ViTImageProcessor(size=config.image_size					)

	remove_projection_head(_lowercase					)
	a      :		int	=				create_rename_keys(_lowercase    ,       base_model=_lowercase					)

	for src, dest in rename_keys:
		rename_key(_lowercase    ,       _lowercase    ,       _lowercase					)
	read_in_q_k_v(_lowercase    ,       _lowercase    ,       base_model=_lowercase					)

	model.load_state_dict(_lowercase					)
	model.eval()

	a      :		str	=				"http://images.cocodataset.org/val2017/000000039769.jpg"

	a      :		Union[str, Any]	=				Image.open(requests.get(_lowercase    ,       stream=_lowercase					).raw					)
	a      :		List[Any]	=				ViTImageProcessor(
	    size=config.image_size    ,       image_mean=_lowercase    ,       image_std=_lowercase					)
	a      :		List[Any]	=				image_processor(images=_lowercase    ,       return_tensors="pt"					)

	# forward pass
	torch.manual_seed(2					)
	a      :		List[Any]	=				model(**_lowercase					)
	a      :		Optional[Any]	=				outputs.last_hidden_state

	# The following Colab Notebook was used to generate these outputs:
	# https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
	if "s16" in checkpoint_url:
		a      :		Dict	=				torch.tensor([[-1.0915, -1.4876, -1.1809]]					)
	elif "b16" in checkpoint_url:
		a      :		Union[str, Any]	=				torch.tensor([[14.2889, -18.9045, 11.7281]]					)
	elif "l16" in checkpoint_url:
		a      :		Any	=				torch.tensor([[41.5028, -22.8681, 45.6475]]					)
	elif "b4" in checkpoint_url:
		a      :		Optional[Any]	=				torch.tensor([[-4.3868, 5.2932, -0.4137]]					)
	else:
		a      :		Tuple	=				torch.tensor([[-0.1792, -0.6465, 2.4263]]					)

	# verify logits
	assert torch.allclose(last_hidden_state[:, 0, :3]    ,       _lowercase    ,       atol=1E-4					)

	print(F"""Saving model to {pytorch_dump_folder_path}"""					)
	model.save_pretrained(_lowercase					)

	print(F"""Saving image processor to {pytorch_dump_folder_path}"""					)
	image_processor.save_pretrained(_lowercase					)


if __name__ == "__main__":
		a						: Optional[Any]     	=			argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--checkpoint_url''',
		    default='''https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar''',
		    type=str,
		    help='''URL of the checkpoint you\'d like to convert.''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
		)

		a						: Optional[Any]     	=			parser.parse_args()
		convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)










import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference

class     _SCREAMING_SNAKE_CASE       (			_a		,   unittest.TestCase      ):
	snake_case__							:     Tuple     	=  ShapEImgaImgPipeline
	snake_case__							:     Optional[Any]     	=  ["""image"""]
	snake_case__							:     Union[str, Any]     	=  ["""image"""]
	snake_case__							:     Optional[Any]     	=  [
	    """num_images_per_prompt""",
	    """num_inference_steps""",
	    """generator""",
	    """latents""",
	    """guidance_scale""",
	    """frame_size""",
	    """output_type""",
	    """return_dict""",
	]
	snake_case__							:     List[str]     	=  False





	@property
	def 						_A (		self				:					Any	):
		return 32





	@property
	def 						_A (		self				:					Any	):
		return 32





	@property
	def 						_A (		self				:					Optional[Any]	):
		return self.time_input_dim * 4





	@property
	def 						_A (		self				:					Union[str, Any]	):
		return 8





	@property
	def 						_A (		self				:					int	):
		torch.manual_seed(0	)
		UpperCamelCase     :Union[str, Any]          =	CLIPVisionConfig(
		    hidden_size=self.text_embedder_hidden_size ,  image_size=64 ,  projection_dim=self.text_embedder_hidden_size ,  intermediate_size=37 ,  num_attention_heads=4 ,  num_channels=3 ,  num_hidden_layers=5 ,  patch_size=1 ,  )

		UpperCamelCase     :Optional[int]          =	CLIPVisionModel(__lowerCamelCase	)
		return model





	@property
	def 						_A (		self				:					str	):
		UpperCamelCase     :Optional[int]          =	CLIPImageProcessor(
		    crop_size=224 ,  do_center_crop=__lowerCamelCase ,  do_normalize=__lowerCamelCase ,  do_resize=__lowerCamelCase ,  image_mean=[0.48145466, 0.4578275, 0.40821073] ,  image_std=[0.26862954, 0.26130258, 0.27577711] ,  resample=3 ,  size=224 ,  )

		return image_processor





	@property
	def 						_A (		self				:					Tuple	):
		torch.manual_seed(0	)

		UpperCamelCase     :Dict          =	{
		    """num_attention_heads""": 2,
		    """attention_head_dim""": 16,
		    """embedding_dim""": self.time_input_dim,
		    """num_embeddings""": 32,
		    """embedding_proj_dim""": self.text_embedder_hidden_size,
		    """time_embed_dim""": self.time_embed_dim,
		    """num_layers""": 1,
		    """clip_embed_dim""": self.time_input_dim * 2,
		    """additional_embeddings""": 0,
		    """time_embed_act_fn""": """gelu""",
		    """norm_in_type""": """layer""",
		    """embedding_proj_norm_type""": """layer""",
		    """encoder_hid_proj_type""": None,
		    """added_emb_type""": None,
		}

		UpperCamelCase     :int          =	PriorTransformer(**__lowerCamelCase	)
		return model





	@property
	def 						_A (		self				:					Optional[int]	):
		torch.manual_seed(0	)

		UpperCamelCase     :str          =	{
		    """param_shapes""": (
		        (self.renderer_dim, 93),
		        (self.renderer_dim, 8),
		        (self.renderer_dim, 8),
		        (self.renderer_dim, 8),
		    ),
		    """d_latent""": self.time_input_dim,
		    """d_hidden""": self.renderer_dim,
		    """n_output""": 12,
		    """background""": (
		        0.1,
		        0.1,
		        0.1,
		    ),
		}
		UpperCamelCase     :List[str]          =	ShapERenderer(**__lowerCamelCase	)
		return model





	def 						_A (		self				:					str	):
		UpperCamelCase     :int          =	self.dummy_prior
		UpperCamelCase     :Any          =	self.dummy_image_encoder
		UpperCamelCase     :Dict          =	self.dummy_image_processor
		UpperCamelCase     :List[Any]          =	self.dummy_renderer

		UpperCamelCase     :int          =	HeunDiscreteScheduler(
		    beta_schedule="""exp""" ,  num_train_timesteps=1_024 ,  prediction_type="""sample""" ,  use_karras_sigmas=__lowerCamelCase ,  clip_sample=__lowerCamelCase ,  clip_sample_range=1.0 ,  )
		UpperCamelCase     :Optional[Any]          =	{
		    """prior""": prior,
		    """image_encoder""": image_encoder,
		    """image_processor""": image_processor,
		    """renderer""": renderer,
		    """scheduler""": scheduler,
		}

		return components





	def 						_A (		self				:					int ,  __lowerCamelCase				:					int ,  __lowerCamelCase				:					Any=0	):
		UpperCamelCase     :Any          =	floats_tensor((1, 3, 64, 64) ,  rng=random.Random(__lowerCamelCase	)	).to(__lowerCamelCase	)

		if str(__lowerCamelCase	).startswith("""mps"""	):
			UpperCamelCase     :List[Any]          =	torch.manual_seed(__lowerCamelCase	)
		else:
			UpperCamelCase     :Optional[int]          =	torch.Generator(device=__lowerCamelCase	).manual_seed(__lowerCamelCase	)
		UpperCamelCase     :Optional[Any]          =	{
		    """image""": input_image,
		    """generator""": generator,
		    """num_inference_steps""": 1,
		    """frame_size""": 32,
		    """output_type""": """np""",
		}
		return inputs





	def 						_A (		self				:					List[str]	):
		UpperCamelCase     :Dict          =	"""cpu"""

		UpperCamelCase     :List[Any]          =	self.get_dummy_components()

		UpperCamelCase     :Optional[int]          =	self.pipeline_class(**__lowerCamelCase	)
		UpperCamelCase     :int          =	pipe.to(__lowerCamelCase	)

		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :Optional[Any]          =	pipe(**self.get_dummy_inputs(__lowerCamelCase	)	)
		UpperCamelCase     :Dict          =	output.images[0]
		UpperCamelCase     :List[Any]          =	image[0, -3:, -3:, -1]

		assert image.shape == (20, 32, 32, 3)

		UpperCamelCase     :Dict          =	np.array(
		    [
		        0.00039216,
		        0.00039216,
		        0.00039216,
		        0.00039216,
		        0.00039216,
		        0.00039216,
		        0.00039216,
		        0.00039216,
		        0.00039216,
		    ]	)

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





	def 						_A (		self				:					List[Any]	):
		# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
		self._test_inference_batch_consistent(batch_sizes=[1, 2]	)





	def 						_A (		self				:					List[Any]	):
		UpperCamelCase     :str          =	torch_device == """cpu"""
		UpperCamelCase     :int          =	True
		self._test_inference_batch_single_identical(
		    batch_size=2 ,  test_max_difference=__lowerCamelCase ,  relax_max_difference=__lowerCamelCase ,  )





	def 						_A (		self				:					List[Any]	):
		UpperCamelCase     :List[Any]          =	self.get_dummy_components()
		UpperCamelCase     :Optional[int]          =	self.pipeline_class(**__lowerCamelCase	)
		UpperCamelCase     :List[Any]          =	pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :Any          =	1
		UpperCamelCase     :int          =	2

		UpperCamelCase     :Union[str, Any]          =	self.get_dummy_inputs(__lowerCamelCase	)

		for key in inputs.keys():
			if key in self.batch_params:
				UpperCamelCase     :str          =	batch_size * [inputs[key]]

		UpperCamelCase     :Optional[int]          =	pipe(**__lowerCamelCase ,  num_images_per_prompt=__lowerCamelCase	)[0]

		assert images.shape[0] == batch_size * num_images_per_prompt

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





	def 						_A (		self				:					Any	):
		# clean up the VRAM after each test
		super().tearDown()
		gc.collect()
		torch.cuda.empty_cache()





	def 						_A (		self				:					Any	):
		UpperCamelCase     :Optional[Any]          =	load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/shap_e/corgi.png"""	)
		UpperCamelCase     :Any          =	load_numpy(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
		    """/shap_e/test_shap_e_img2img_out.npy"""	)
		UpperCamelCase     :Union[str, Any]          =	ShapEImgaImgPipeline.from_pretrained("""openai/shap-e-img2img"""	)
		UpperCamelCase     :List[str]          =	pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :Optional[Any]          =	torch.Generator(device=__lowerCamelCase	).manual_seed(0	)

		UpperCamelCase     :Optional[int]          =	pipe(
		    __lowerCamelCase ,  generator=__lowerCamelCase ,  guidance_scale=3.0 ,  num_inference_steps=64 ,  frame_size=64 ,  output_type="""np""" ,  ).images[0]

		assert images.shape == (20, 64, 64, 3)

		assert_mean_pixel_difference(__lowerCamelCase ,  __lowerCamelCase	)
