
"""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 	__snake_case  (    SCREAMING_SNAKE_CASE:					Optional[Any]     ,  SCREAMING_SNAKE_CASE:					List[Any]=False						):






							"""simple docstring"""



							_lowerCAmelCase  	=			[]
							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"
														_lowerCAmelCase  	=			[(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 	__snake_case  (    SCREAMING_SNAKE_CASE:					List[Any]     ,  SCREAMING_SNAKE_CASE:					Tuple     ,  SCREAMING_SNAKE_CASE:					int=False						):






							"""simple docstring"""



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






def 	__snake_case  (    SCREAMING_SNAKE_CASE:					Any						):






							"""simple docstring"""



							_lowerCAmelCase  	=			['head.weight', 'head.bias']
							for k in ignore_keys:
														state_dict.pop(SCREAMING_SNAKE_CASE     ,  SCREAMING_SNAKE_CASE						)






def 	__snake_case  (    SCREAMING_SNAKE_CASE:					Tuple						):






							"""simple docstring"""



							_lowerCAmelCase  	=			[
							    '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(SCREAMING_SNAKE_CASE     ,  SCREAMING_SNAKE_CASE						)






def 	__snake_case  (    SCREAMING_SNAKE_CASE:					List[Any]     ,  SCREAMING_SNAKE_CASE:					int     ,  SCREAMING_SNAKE_CASE:					Optional[int]						):






							"""simple docstring"""



							_lowerCAmelCase  	=			dct.pop(SCREAMING_SNAKE_CASE						)
							_lowerCAmelCase  	=			val






def 	__snake_case  (    SCREAMING_SNAKE_CASE:					str     ,  SCREAMING_SNAKE_CASE:					List[str]						):






							"""simple docstring"""



							_lowerCAmelCase  	=			ViTMSNConfig()
							_lowerCAmelCase  	=			1000

							_lowerCAmelCase  	=			'datasets/huggingface/label-files'
							_lowerCAmelCase  	=			'imagenet-1k-id2label.json'
							_lowerCAmelCase  	=			json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE     ,  SCREAMING_SNAKE_CASE						)     ,  'r'						)						)
							_lowerCAmelCase  	=			{int(SCREAMING_SNAKE_CASE						): v for k, v in idalabel.items()}
							_lowerCAmelCase  	=			idalabel
							_lowerCAmelCase  	=			{v: k for k, v in idalabel.items()}

							if "s16" in checkpoint_url:
														_lowerCAmelCase  	=			384
														_lowerCAmelCase  	=			1536
														_lowerCAmelCase  	=			6
							elif "l16" in checkpoint_url:
														_lowerCAmelCase  	=			1024
														_lowerCAmelCase  	=			4096
														_lowerCAmelCase  	=			24
														_lowerCAmelCase  	=			16
														_lowerCAmelCase  	=			0.1
							elif "b4" in checkpoint_url:
														_lowerCAmelCase  	=			4
							elif "l7" in checkpoint_url:
														_lowerCAmelCase  	=			7
														_lowerCAmelCase  	=			1024
														_lowerCAmelCase  	=			4096
														_lowerCAmelCase  	=			24
														_lowerCAmelCase  	=			16
														_lowerCAmelCase  	=			0.1

							_lowerCAmelCase  	=			ViTMSNModel(SCREAMING_SNAKE_CASE						)

							_lowerCAmelCase  	=			torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE     ,  map_location='cpu'						)['target_encoder']

							_lowerCAmelCase  	=			ViTImageProcessor(size=config.image_size						)

							remove_projection_head(SCREAMING_SNAKE_CASE						)
							_lowerCAmelCase  	=			create_rename_keys(SCREAMING_SNAKE_CASE     ,  base_model=SCREAMING_SNAKE_CASE						)

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

							model.load_state_dict(SCREAMING_SNAKE_CASE						)
							model.eval()

							_lowerCAmelCase  	=			'http://images.cocodataset.org/val2017/000000039769.jpg'

							_lowerCAmelCase  	=			Image.open(requests.get(SCREAMING_SNAKE_CASE     ,  stream=SCREAMING_SNAKE_CASE						).raw						)
							_lowerCAmelCase  	=			ViTImageProcessor(
							    size=config.image_size     ,  image_mean=SCREAMING_SNAKE_CASE     ,  image_std=SCREAMING_SNAKE_CASE						)
							_lowerCAmelCase  	=			image_processor(images=SCREAMING_SNAKE_CASE     ,  return_tensors='pt'						)

							# forward pass
							torch.manual_seed(2						)
							_lowerCAmelCase  	=			model(**SCREAMING_SNAKE_CASE						)
							_lowerCAmelCase  	=			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:
														_lowerCAmelCase  	=			torch.tensor([[-1.09_15, -1.48_76, -1.18_09]]						)
							elif "b16" in checkpoint_url:
														_lowerCAmelCase  	=			torch.tensor([[14.28_89, -18.90_45, 11.72_81]]						)
							elif "l16" in checkpoint_url:
														_lowerCAmelCase  	=			torch.tensor([[41.50_28, -22.86_81, 45.64_75]]						)
							elif "b4" in checkpoint_url:
														_lowerCAmelCase  	=			torch.tensor([[-4.38_68, 5.29_32, -0.41_37]]						)
							else:
														_lowerCAmelCase  	=			torch.tensor([[-0.17_92, -0.64_65, 2.42_63]]						)

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

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

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


if __name__ == "__main__":
	_snake_case          =						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.'''
	)

	_snake_case          =						parser.parse_args()
	convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)








"""simple docstring"""



import argparse

import intel_extension_for_pytorch as ipex
import torch

from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline


_snake_case          =						argparse.ArgumentParser('''Stable Diffusion script with intel optimization''', add_help=False)
parser.add_argument('''--dpm''', action='''store_true''', help='''Enable DPMSolver or not''')
parser.add_argument('''--steps''', default=None, type=int, help='''Num inference steps''')
_snake_case          =						parser.parse_args()


_snake_case          =						'''cpu'''
_snake_case          =						'''a lovely <dicoo> in red dress and hat, in the snowly and brightly night, with many brighly buildings'''

_snake_case          =						'''path-to-your-trained-model'''
_snake_case          =						StableDiffusionPipeline.from_pretrained(model_id)
if args.dpm:
	_snake_case          =						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
_snake_case          =						pipe.to(device)

# to channels last
_snake_case          =						pipe.unet.to(memory_format=torch.channels_last)
_snake_case          =						pipe.vae.to(memory_format=torch.channels_last)
_snake_case          =						pipe.text_encoder.to(memory_format=torch.channels_last)
if pipe.requires_safety_checker:
	_snake_case          =						pipe.safety_checker.to(memory_format=torch.channels_last)

# optimize with ipex
_snake_case          =						torch.randn(2, 4, 6_4, 6_4)
_snake_case          =						torch.rand(1) * 9_9_9
_snake_case          =						torch.randn(2, 7_7, 7_6_8)
_snake_case          =						(sample, timestep, encoder_hidden_status)
try:
	_snake_case          =						ipex.optimize(pipe.unet.eval(), dtype=torch.bfloataa, inplace=True, sample_input=input_example)
except Exception:
	_snake_case          =						ipex.optimize(pipe.unet.eval(), dtype=torch.bfloataa, inplace=True)
_snake_case          =						ipex.optimize(pipe.vae.eval(), dtype=torch.bfloataa, inplace=True)
_snake_case          =						ipex.optimize(pipe.text_encoder.eval(), dtype=torch.bfloataa, inplace=True)
if pipe.requires_safety_checker:
	_snake_case          =						ipex.optimize(pipe.safety_checker.eval(), dtype=torch.bfloataa, inplace=True)

# compute
_snake_case          =						6_6_6
_snake_case          =						torch.Generator(device).manual_seed(seed)
_snake_case          =						{'''generator''': generator}
if args.steps is not None:
	_snake_case          =						args.steps

with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloataa):
	_snake_case          =						pipe(prompt, **generate_kwargs).images[0]

# save image
image.save('''generated.png''')






