





'''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       UpperCAmelCase							(     _lowerCamelCase      :  int			,   _lowerCamelCase      :  Dict=False     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				Tuple   =					[]
						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"
												SCREAMING_SNAKE_CASE__     :				Dict   =					[(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       UpperCAmelCase							(     _lowerCamelCase      :  Union[str, Any]			,   _lowerCamelCase      :  int			,   _lowerCamelCase      :  List[str]=False     ):





						'''simple docstring'''


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







def       UpperCAmelCase							(     _lowerCamelCase      :  List[str]     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				Dict   =					["head.weight", "head.bias"]
						for k in ignore_keys:
												state_dict.pop(_lowerCamelCase			,   _lowerCamelCase     )







def       UpperCAmelCase							(     _lowerCamelCase      :  str     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				List[Any]   =					[
						    "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(_lowerCamelCase			,   _lowerCamelCase     )







def       UpperCAmelCase							(     _lowerCamelCase      :  Optional[int]			,   _lowerCamelCase      :  List[str]			,   _lowerCamelCase      :  int     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				List[Any]   =					dct.pop(_lowerCamelCase     )
						SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					val







def       UpperCAmelCase							(     _lowerCamelCase      :  List[str]			,   _lowerCamelCase      :  Optional[int]     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				List[str]   =					ViTMSNConfig()
						SCREAMING_SNAKE_CASE__     :				Any   =					1_000

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

						if "s16" in checkpoint_url:
												SCREAMING_SNAKE_CASE__     :				Optional[int]   =					384
												SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					1_536
												SCREAMING_SNAKE_CASE__     :				List[Any]   =					6
						elif "l16" in checkpoint_url:
												SCREAMING_SNAKE_CASE__     :				str   =					1_024
												SCREAMING_SNAKE_CASE__     :				str   =					4_096
												SCREAMING_SNAKE_CASE__     :				Dict   =					24
												SCREAMING_SNAKE_CASE__     :				Optional[int]   =					16
												SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					0.1
						elif "b4" in checkpoint_url:
												SCREAMING_SNAKE_CASE__     :				Optional[int]   =					4
						elif "l7" in checkpoint_url:
												SCREAMING_SNAKE_CASE__     :				int   =					7
												SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					1_024
												SCREAMING_SNAKE_CASE__     :				List[Any]   =					4_096
												SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					24
												SCREAMING_SNAKE_CASE__     :				Tuple   =					16
												SCREAMING_SNAKE_CASE__     :				int   =					0.1

						SCREAMING_SNAKE_CASE__     :				str   =					ViTMSNModel(_lowerCamelCase     )

						SCREAMING_SNAKE_CASE__     :				List[str]   =					torch.hub.load_state_dict_from_url(_lowerCamelCase			,   map_location="cpu"     )["target_encoder"]

						SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					ViTImageProcessor(size=config.image_size     )

						remove_projection_head(_lowerCamelCase     )
						SCREAMING_SNAKE_CASE__     :				str   =					create_rename_keys(_lowerCamelCase			,   base_model=_lowerCamelCase     )

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

						model.load_state_dict(_lowerCamelCase     )
						model.eval()

						SCREAMING_SNAKE_CASE__     :				Tuple   =					"http://images.cocodataset.org/val2017/000000039769.jpg"

						SCREAMING_SNAKE_CASE__     :				Tuple   =					Image.open(requests.get(_lowerCamelCase			,   stream=_lowerCamelCase     ).raw     )
						SCREAMING_SNAKE_CASE__     :				Dict   =					ViTImageProcessor(
						    size=config.image_size			,   image_mean=_lowerCamelCase			,   image_std=_lowerCamelCase     )
						SCREAMING_SNAKE_CASE__     :				str   =					image_processor(images=_lowerCamelCase			,   return_tensors="pt"     )

						# forward pass
						torch.manual_seed(2     )
						SCREAMING_SNAKE_CASE__     :				Tuple   =					model(**_lowerCamelCase     )
						SCREAMING_SNAKE_CASE__     :				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:
												SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					torch.tensor([[-1.0_9_1_5, -1.4_8_7_6, -1.1_8_0_9]]     )
						elif "b16" in checkpoint_url:
												SCREAMING_SNAKE_CASE__     :				Dict   =					torch.tensor([[14.2_889, -18.9_045, 11.7_281]]     )
						elif "l16" in checkpoint_url:
												SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					torch.tensor([[41.5_028, -22.8_681, 45.6_475]]     )
						elif "b4" in checkpoint_url:
												SCREAMING_SNAKE_CASE__     :				int   =					torch.tensor([[-4.3_8_6_8, 5.2_9_3_2, -0.4_1_3_7]]     )
						else:
												SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					torch.tensor([[-0.1_7_9_2, -0.6_4_6_5, 2.4_2_6_3]]     )

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

						print(f"""Saving model to {pytorch_dump_folder_path}"""     )
						model.save_pretrained(_lowerCamelCase     )

						print(f"""Saving image processor to {pytorch_dump_folder_path}"""     )
						image_processor.save_pretrained(_lowerCamelCase     )


if __name__ == "__main__":
				__lowercase       :List[str]       =  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."
				)

				__lowercase       :Dict       =  parser.parse_args()
				convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
				import jax
				import jax.numpy as jnp
				from flax.jax_utils import replicate
				from flax.training.common_utils import shard



@slow
@require_flax
class    _a     (				unittest.TestCase							):






							"""simple docstring"""

							def      A_  (      self				:     Dict    )	->List[Any]:
													# clean up the VRAM after each test
													super().tearDown()
													gc.collect()

							def      A_  (      self				:     Dict    )	->Tuple:
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Dict   =					FlaxControlNetModel.from_pretrained(
													    "lllyasviel/sd-controlnet-canny"							,	from_pt=a							,	dtype=jnp.bfloataa    )
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Dict   =					FlaxStableDiffusionControlNetPipeline.from_pretrained(
													    "runwayml/stable-diffusion-v1-5"							,	controlnet=a							,	from_pt=a							,	dtype=jnp.bfloataa    )
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					controlnet_params

													SCREAMING_SNAKE_CASE__     :				Dict   =					"bird"
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					jax.device_count()
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					pipe.prepare_text_inputs([prompts] * num_samples    )

													SCREAMING_SNAKE_CASE__     :				Dict   =					load_image(
													    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"    )
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					pipe.prepare_image_inputs([canny_image] * num_samples    )

													SCREAMING_SNAKE_CASE__     :				List[Any]   =					jax.random.PRNGKey(0    )
													SCREAMING_SNAKE_CASE__     :				int   =					jax.random.split(a							,	jax.device_count()    )

													SCREAMING_SNAKE_CASE__     :				List[Any]   =					replicate(a    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					shard(a    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					shard(a    )

													SCREAMING_SNAKE_CASE__     :				Dict   =					pipe(
													    prompt_ids=a							,	image=a							,	params=a							,	prng_seed=a							,	num_inference_steps=50							,	jit=a							,	).images
													assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]    )
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					images[0, 2_53:2_56, 2_53:2_56, -1]

													SCREAMING_SNAKE_CASE__     :				Tuple   =					jnp.asarray(jax.device_get(image_slice.flatten()    )    )
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					jnp.array(
													    [0.16_7969, 0.11_6699, 0.08_1543, 0.15_4297, 0.13_2812, 0.10_8887, 0.16_9922, 0.16_9922, 0.20_5078]    )
													print(f"""output_slice: {output_slice}"""    )
													assert jnp.abs(output_slice - expected_slice    ).max() < 1E-2






							def      A_  (      self				:     List[Any]    )	->Optional[Any]:
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				int   =					FlaxControlNetModel.from_pretrained(
													    "lllyasviel/sd-controlnet-openpose"							,	from_pt=a							,	dtype=jnp.bfloataa    )
													SCREAMING_SNAKE_CASE__,							SCREAMING_SNAKE_CASE__     :				Optional[int]   =					FlaxStableDiffusionControlNetPipeline.from_pretrained(
													    "runwayml/stable-diffusion-v1-5"							,	controlnet=a							,	from_pt=a							,	dtype=jnp.bfloataa    )
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					controlnet_params

													SCREAMING_SNAKE_CASE__     :				Any   =					"Chef in the kitchen"
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jax.device_count()
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					pipe.prepare_text_inputs([prompts] * num_samples    )

													SCREAMING_SNAKE_CASE__     :				Dict   =					load_image(
													    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png"    )
													SCREAMING_SNAKE_CASE__     :				str   =					pipe.prepare_image_inputs([pose_image] * num_samples    )

													SCREAMING_SNAKE_CASE__     :				Any   =					jax.random.PRNGKey(0    )
													SCREAMING_SNAKE_CASE__     :				List[str]   =					jax.random.split(a							,	jax.device_count()    )

													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					replicate(a    )
													SCREAMING_SNAKE_CASE__     :				Tuple   =					shard(a    )
													SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					shard(a    )

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					pipe(
													    prompt_ids=a							,	image=a							,	params=a							,	prng_seed=a							,	num_inference_steps=50							,	jit=a							,	).images
													assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]    )
													SCREAMING_SNAKE_CASE__     :				str   =					images[0, 2_53:2_56, 2_53:2_56, -1]

													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jnp.asarray(jax.device_get(image_slice.flatten()    )    )
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					jnp.array(
													    [[0.27_1484, 0.26_1719, 0.27_5391, 0.27_7344, 0.27_9297, 0.29_1016, 0.29_4922, 0.30_2734, 0.30_2734]]    )
													print(f"""output_slice: {output_slice}"""    )
													assert jnp.abs(output_slice - expected_slice    ).max() < 1E-2