






"""simple docstring"""
import argparse

import requests
import torch

# pip3 install salesforce-lavis
# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD






def 			__UpperCamelCase	(		):
	A_		:						str	    =   """https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
	A_		:						List[str]	    =   Image.open(requests.get(snake_case__   ,   stream=snake_case__					).raw					).convert("""RGB"""					)

	return image






def 			__UpperCamelCase	(		snake_case__					):
	A_		:						str	    =   []
	# fmt: off

	# vision encoder
	rename_keys.append(("""visual_encoder.cls_token""", """vision_model.embeddings.class_embedding""")					)
	rename_keys.append(("""visual_encoder.pos_embed""", """vision_model.embeddings.position_embedding""")					)
	rename_keys.append(("""visual_encoder.patch_embed.proj.weight""", """vision_model.embeddings.patch_embedding.weight""")					)
	rename_keys.append(("""visual_encoder.patch_embed.proj.bias""", """vision_model.embeddings.patch_embedding.bias""")					)
	rename_keys.append(("""ln_vision.weight""", """vision_model.post_layernorm.weight""")					)
	rename_keys.append(("""ln_vision.bias""", """vision_model.post_layernorm.bias""")					)

	for i in range(config.vision_config.num_hidden_layers					):
		rename_keys.append((F"""visual_encoder.blocks.{i}.norm1.weight""", F"""vision_model.encoder.layers.{i}.layer_norm1.weight""")					)
		rename_keys.append((F"""visual_encoder.blocks.{i}.norm1.bias""", F"""vision_model.encoder.layers.{i}.layer_norm1.bias""")					)
		rename_keys.append((F"""visual_encoder.blocks.{i}.norm2.weight""", F"""vision_model.encoder.layers.{i}.layer_norm2.weight""")					)
		rename_keys.append((F"""visual_encoder.blocks.{i}.norm2.bias""", F"""vision_model.encoder.layers.{i}.layer_norm2.bias""")					)
		rename_keys.append((F"""visual_encoder.blocks.{i}.attn.qkv.weight""", F"""vision_model.encoder.layers.{i}.self_attn.qkv.weight""")					)
		rename_keys.append((F"""visual_encoder.blocks.{i}.attn.proj.weight""", F"""vision_model.encoder.layers.{i}.self_attn.projection.weight""",)					)
		rename_keys.append((F"""visual_encoder.blocks.{i}.attn.proj.bias""", F"""vision_model.encoder.layers.{i}.self_attn.projection.bias""")					)
		rename_keys.append((F"""visual_encoder.blocks.{i}.mlp.fc1.weight""", F"""vision_model.encoder.layers.{i}.mlp.fc1.weight""")					)
		rename_keys.append((F"""visual_encoder.blocks.{i}.mlp.fc1.bias""", F"""vision_model.encoder.layers.{i}.mlp.fc1.bias""")					)
		rename_keys.append((F"""visual_encoder.blocks.{i}.mlp.fc2.weight""", F"""vision_model.encoder.layers.{i}.mlp.fc2.weight""")					)
		rename_keys.append((F"""visual_encoder.blocks.{i}.mlp.fc2.bias""", F"""vision_model.encoder.layers.{i}.mlp.fc2.bias""")					)

	# QFormer
	rename_keys.append(("""Qformer.bert.embeddings.LayerNorm.weight""", """qformer.layernorm.weight""")					)
	rename_keys.append(("""Qformer.bert.embeddings.LayerNorm.bias""", """qformer.layernorm.bias""")					)

	# fmt: on
	return rename_keys






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__					):
	A_		:						str	    =   dct.pop(snake_case__					)
	A_		:						str	    =   val






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__					):
	for i in range(config.vision_config.num_hidden_layers					):
		# read in original q and v biases
		A_		:						Any	    =   state_dict.pop(F"""visual_encoder.blocks.{i}.attn.q_bias"""					)
		A_		:						Any	    =   state_dict.pop(F"""visual_encoder.blocks.{i}.attn.v_bias"""					)

		# next, set bias in the state dict
		A_		:						List[str]	    =   torch.cat((q_bias, torch.zeros_like(snake_case__   ,   requires_grad=snake_case__					), v_bias)					)
		A_		:						Optional[Any]	    =   qkv_bias






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__					):
	A_		:						str	    =   364 if """coco""" in model_name else 224
	A_		:						Dict	    =   BlipaVisionConfig(image_size=snake_case__					).to_dict()

	# make sure the models have proper bos_token_id and eos_token_id set (important for generation)
	# seems like flan-T5 models don't have bos_token_id properly set?
	if "opt-2.7b" in model_name:
		A_		:						List[Any]	    =   OPTConfig.from_pretrained("""facebook/opt-2.7b"""   ,   eos_token_id=snake_case__					).to_dict()
	elif "opt-6.7b" in model_name:
		A_		:						Optional[Any]	    =   OPTConfig.from_pretrained("""facebook/opt-6.7b"""   ,   eos_token_id=snake_case__					).to_dict()
	elif "t5-xl" in model_name:
		A_		:						Optional[int]	    =   TaConfig.from_pretrained("""google/flan-t5-xl"""   ,   dense_act_fn="""gelu"""   ,   bos_token_id=1					).to_dict()
	elif "t5-xxl" in model_name:
		A_		:						Optional[Any]	    =   TaConfig.from_pretrained("""google/flan-t5-xxl"""   ,   dense_act_fn="""gelu"""   ,   bos_token_id=1					).to_dict()

	A_		:						Optional[Any]	    =   BlipaConfig(vision_config=snake_case__   ,   text_config=snake_case__					)

	return config, image_size






@torch.no_grad()
def 			__UpperCamelCase	(		snake_case__   ,   snake_case__=None   ,   snake_case__=False					):
	A_		:						List[Any]	    =   (
	    AutoTokenizer.from_pretrained("""facebook/opt-2.7b"""					)
	    if """opt""" in model_name
	    else AutoTokenizer.from_pretrained("""google/flan-t5-xl"""					)
	)
	A_		:						Optional[Any]	    =   tokenizer("""\n"""   ,   add_special_tokens=snake_case__					).input_ids[0]
	A_    ,     A_		:						str	    =   get_blipa_config(snake_case__   ,   eos_token_id=snake_case__					)

	A_		:						Dict	    =   BlipaForConditionalGeneration(snake_case__					).eval()

	A_		:						str	    =   {
	    """blip2-opt-2.7b""": ("""blip2_opt""", """pretrain_opt2.7b"""),
	    """blip2-opt-6.7b""": ("""blip2_opt""", """pretrain_opt6.7b"""),
	    """blip2-opt-2.7b-coco""": ("""blip2_opt""", """caption_coco_opt2.7b"""),
	    """blip2-opt-6.7b-coco""": ("""blip2_opt""", """caption_coco_opt6.7b"""),
	    """blip2-flan-t5-xl""": ("""blip2_t5""", """pretrain_flant5xl"""),
	    """blip2-flan-t5-xl-coco""": ("""blip2_t5""", """caption_coco_flant5xl"""),
	    """blip2-flan-t5-xxl""": ("""blip2_t5""", """pretrain_flant5xxl"""),
	}

	A_    ,     A_		:						Tuple	    =   model_name_to_original[model_name]

	# load original model
	print("""Loading original model..."""					)
	A_		:						Any	    =   """cuda""" if torch.cuda.is_available() else """cpu"""
	A_    ,     A_    ,     A_		:						List[Any]	    =   load_model_and_preprocess(
	    name=snake_case__   ,   model_type=snake_case__   ,   is_eval=snake_case__   ,   device=snake_case__					)
	original_model.eval()
	print("""Done!"""					)

	# update state dict keys
	A_		:						Dict	    =   original_model.state_dict()
	A_		:						Any	    =   create_rename_keys(snake_case__					)
	for src, dest in rename_keys:
		rename_key(snake_case__   ,   snake_case__   ,   snake_case__					)

	# some keys can be renamed efficiently
	for key, val in state_dict.copy().items():
		A_		:						Optional[Any]	    =   state_dict.pop(snake_case__					)
		if key.startswith("""Qformer.bert"""					):
			A_		:						List[Any]	    =   key.replace("""Qformer.bert"""   ,   """qformer"""					)
		if "attention.self" in key:
			A_		:						Tuple	    =   key.replace("""self"""   ,   """attention"""					)
		if "opt_proj" in key:
			A_		:						Dict	    =   key.replace("""opt_proj"""   ,   """language_projection"""					)
		if "t5_proj" in key:
			A_		:						Union[str, Any]	    =   key.replace("""t5_proj"""   ,   """language_projection"""					)
		if key.startswith("""opt"""					):
			A_		:						Tuple	    =   key.replace("""opt"""   ,   """language"""					)
		if key.startswith("""t5"""					):
			A_		:						List[Any]	    =   key.replace("""t5"""   ,   """language"""					)
		A_		:						Optional[Any]	    =   val

	# read in qv biases
	read_in_q_v_bias(snake_case__   ,   snake_case__					)

	A_    ,     A_		:						Dict	    =   hf_model.load_state_dict(snake_case__   ,   strict=snake_case__					)
	assert len(snake_case__					) == 0
	assert unexpected_keys == ["qformer.embeddings.position_ids"]

	A_		:						Tuple	    =   load_demo_image()
	A_		:						Tuple	    =   vis_processors["""eval"""](snake_case__					).unsqueeze(0					).to(snake_case__					)
	A_		:						Any	    =   tokenizer(["""\n"""]   ,   return_tensors="""pt"""					).input_ids.to(snake_case__					)

	# create processor
	A_		:						List[Any]	    =   BlipImageProcessor(
	    size={"""height""": image_size, """width""": image_size}   ,   image_mean=snake_case__   ,   image_std=snake_case__					)
	A_		:						Any	    =   BlipaProcessor(image_processor=snake_case__   ,   tokenizer=snake_case__					)
	A_		:						Union[str, Any]	    =   processor(images=snake_case__   ,   return_tensors="""pt"""					).pixel_values.to(snake_case__					)

	# make sure processor creates exact same pixel values
	assert torch.allclose(snake_case__   ,   snake_case__					)

	original_model.to(snake_case__					)
	hf_model.to(snake_case__					)
	with torch.no_grad():
		if "opt" in model_name:
			A_		:						Optional[int]	    =   original_model({"""image""": original_pixel_values, """text_input""": [""""""]}					).logits
			A_		:						List[Any]	    =   hf_model(snake_case__   ,   snake_case__					).logits
		else:
			A_		:						Optional[int]	    =   original_model(
			    {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]}					).logits
			A_		:						str	    =   input_ids.masked_fill(input_ids == tokenizer.pad_token_id   ,   -100					)
			A_		:						Any	    =   hf_model(snake_case__   ,   snake_case__   ,   labels=snake_case__					).logits

	assert original_logits.shape == logits.shape
	print("""First values of original logits:"""   ,   original_logits[0, :3, :3]					)
	print("""First values of HF logits:"""   ,   logits[0, :3, :3]					)

	# assert values
	if model_name == "blip2-flan-t5-xl":
		A_		:						Tuple	    =   torch.tensor(
		    [[-41.5_850, -4.4_440, -8.9_922], [-47.4_322, -5.9_143, -1.7_340]]   ,   device=snake_case__					)
		assert torch.allclose(logits[0, :3, :3]   ,   snake_case__   ,   atol=1E-4					)
	elif model_name == "blip2-flan-t5-xl-coco":
		A_		:						List[Any]	    =   torch.tensor(
		    [[-57.0_109, -9.8_967, -12.6_280], [-68.6_578, -12.7_191, -10.5_065]]   ,   device=snake_case__					)
	else:
		# cast to same type
		A_		:						Dict	    =   logits.dtype
		assert torch.allclose(original_logits.to(snake_case__					)   ,   snake_case__   ,   atol=1E-2					)
	print("""Looks ok!"""					)

	print("""Generating a caption..."""					)
	A_		:						str	    =   """"""
	A_		:						str	    =   tokenizer(snake_case__   ,   return_tensors="""pt"""					).input_ids.to(snake_case__					)

	A_		:						Any	    =   original_model.generate({"""image""": original_pixel_values}					)
	A_		:						Any	    =   hf_model.generate(
	    snake_case__   ,   snake_case__   ,   do_sample=snake_case__   ,   num_beams=5   ,   max_length=30   ,   min_length=1   ,   top_p=0.9   ,   repetition_penalty=1.0   ,   length_penalty=1.0   ,   temperature=1   ,   )
	print("""Original generation:"""   ,   snake_case__					)
	A_		:						Union[str, Any]	    =   input_ids.shape[1]
	A_		:						Dict	    =   processor.batch_decode(outputs[:, prompt_length:]   ,   skip_special_tokens=snake_case__					)
	A_		:						Optional[Any]	    =   [text.strip() for text in output_text]
	print("""HF generation:"""   ,   snake_case__					)

	if pytorch_dump_folder_path is not None:
		processor.save_pretrained(snake_case__					)
		hf_model.save_pretrained(snake_case__					)

	if push_to_hub:
		processor.push_to_hub(F"""nielsr/{model_name}"""					)
		hf_model.push_to_hub(F"""nielsr/{model_name}"""					)


if __name__ == "__main__":
							_lowerCAmelCase         =							argparse.ArgumentParser()
							_lowerCAmelCase         =							[
							    "blip2-opt-2.7b",
							    "blip2-opt-6.7b",
							    "blip2-opt-2.7b-coco",
							    "blip2-opt-6.7b-coco",
							    "blip2-flan-t5-xl",
							    "blip2-flan-t5-xl-coco",
							    "blip2-flan-t5-xxl",
							]
							parser.add_argument(
							    "--model_name",
							    default="blip2-opt-2.7b",
							    choices=choices,
							    type=str,
							    help="Path to hf config.json of model to convert",
							)
							parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
							parser.add_argument(
							    "--push_to_hub",
							    action="store_true",
							    help="Whether to push the model and processor to the hub after converting",
							)

							_lowerCAmelCase         =							parser.parse_args()

							convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









"""simple docstring"""
import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()



class      SCREAMING_SNAKE_CASE  (				_SCREAMING_SNAKE_CASE       ,       unittest.TestCase						):





	"""simple docstring"""




	_A       :				int    =				PriorTransformer
	_A       :				Dict    =				"""hidden_states"""





	@property
	def    lowerCamelCase(self							):
		A_		:						Tuple	    =   4
		A_		:						Dict	    =   8
		A_		:						str	    =   7

		A_		:						Dict	    =   floats_tensor((batch_size, embedding_dim)							).to(lowerCAmelCase_							)

		A_		:						str	    =   floats_tensor((batch_size, embedding_dim)							).to(lowerCAmelCase_							)
		A_		:						List[str]	    =   floats_tensor((batch_size, num_embeddings, embedding_dim)							).to(lowerCAmelCase_							)

		return {
		    "hidden_states": hidden_states,
		    "timestep": 2,
		    "proj_embedding": proj_embedding,
		    "encoder_hidden_states": encoder_hidden_states,
		}





	def    lowerCamelCase(self     ,     lowerCAmelCase_=0							):
		torch.manual_seed(lowerCAmelCase_							)
		A_		:						Optional[Any]	    =   4
		A_		:						Tuple	    =   8
		A_		:						Tuple	    =   7

		A_		:						int	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)

		A_		:						List[str]	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)
		A_		:						Tuple	    =   torch.randn((batch_size, num_embeddings, embedding_dim)							).to(lowerCAmelCase_							)

		return {
		    "hidden_states": hidden_states,
		    "timestep": 2,
		    "proj_embedding": proj_embedding,
		    "encoder_hidden_states": encoder_hidden_states,
		}





	@property
	def    lowerCamelCase(self							):
		return (4, 8)





	@property
	def    lowerCamelCase(self							):
		return (4, 8)





	def    lowerCamelCase(self							):
		A_		:						int	    =   {
		    """num_attention_heads""": 2,
		    """attention_head_dim""": 4,
		    """num_layers""": 2,
		    """embedding_dim""": 8,
		    """num_embeddings""": 7,
		    """additional_embeddings""": 4,
		}
		A_		:						Dict	    =   self.dummy_input
		return init_dict, inputs_dict





	def    lowerCamelCase(self							):
		A_    ,     A_		:						Optional[int]	    =   PriorTransformer.from_pretrained(
		    """hf-internal-testing/prior-dummy"""     ,     output_loading_info=lowerCAmelCase_							)
		self.assertIsNotNone(lowerCAmelCase_							)
		self.assertEqual(len(loading_info["""missing_keys"""]							)     ,     0							)

		model.to(lowerCAmelCase_							)
		A_		:						List[Any]	    =   model(**self.dummy_input							)[0]

		assert hidden_states is not None, "Make sure output is not None"





	def    lowerCamelCase(self							):
		A_    ,     A_		:						Union[str, Any]	    =   self.prepare_init_args_and_inputs_for_common()

		A_		:						Union[str, Any]	    =   self.model_class(**lowerCAmelCase_							)
		A_		:						Tuple	    =   inspect.signature(model.forward							)
		# signature.parameters is an OrderedDict => so arg_names order is deterministic
		A_		:						Tuple	    =   [*signature.parameters.keys()]

		A_		:						List[str]	    =   ["""hidden_states""", """timestep"""]
		self.assertListEqual(arg_names[:2]     ,     lowerCAmelCase_							)





	def    lowerCamelCase(self							):
		A_		:						int	    =   PriorTransformer.from_pretrained("""hf-internal-testing/prior-dummy"""							)
		A_		:						Optional[Any]	    =   model.to(lowerCAmelCase_							)

		if hasattr(lowerCAmelCase_     ,     """set_default_attn_processor"""							):
			model.set_default_attn_processor()

		A_		:						int	    =   self.get_dummy_seed_input()

		with torch.no_grad():
			A_		:						int	    =   model(**lowerCAmelCase_							)[0]

		A_		:						Any	    =   output[0, :5].flatten().cpu()
		print(lowerCAmelCase_							)

		# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
		# the expected output slices are not the same for CPU and GPU.
		A_		:						Optional[int]	    =   torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239]							)
		self.assertTrue(torch_all_close(lowerCAmelCase_     ,     lowerCAmelCase_     ,     rtol=1e-2							)							)






@slow
class      SCREAMING_SNAKE_CASE  (				unittest.TestCase						):





	"""simple docstring"""




	def    lowerCamelCase(self     ,     lowerCAmelCase_=1     ,     lowerCAmelCase_=768     ,     lowerCAmelCase_=77     ,     lowerCAmelCase_=0							):
		torch.manual_seed(lowerCAmelCase_							)
		A_		:						Tuple	    =   batch_size
		A_		:						Any	    =   embedding_dim
		A_		:						Tuple	    =   num_embeddings

		A_		:						Any	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)

		A_		:						Union[str, Any]	    =   torch.randn((batch_size, embedding_dim)							).to(lowerCAmelCase_							)
		A_		:						Tuple	    =   torch.randn((batch_size, num_embeddings, embedding_dim)							).to(lowerCAmelCase_							)

		return {
		    "hidden_states": hidden_states,
		    "timestep": 2,
		    "proj_embedding": proj_embedding,
		    "encoder_hidden_states": encoder_hidden_states,
		}





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





	@parameterized.expand(
	    [
	        # fmt: off
	        [13, [-0.5861,  0.1283, -0.0931,  0.0882,  0.4476,  0.1329, -0.0498,  0.0640]],
	        [37, [-0.4913,  0.0110, -0.0483,  0.0541,  0.4954, -0.0170,  0.0354,  0.1651]],
	        # fmt: on
	    ]							)
	def    lowerCamelCase(self     ,     lowerCAmelCase_     ,     lowerCAmelCase_							):
		A_		:						List[Any]	    =   PriorTransformer.from_pretrained("""kandinsky-community/kandinsky-2-1-prior"""     ,     subfolder="""prior"""							)
		model.to(lowerCAmelCase_							)
		A_		:						int	    =   self.get_dummy_seed_input(seed=lowerCAmelCase_							)

		with torch.no_grad():
			A_		:						Tuple	    =   model(**lowerCAmelCase_							)[0]

		assert list(sample.shape							) == [1, 768]

		A_		:						Tuple	    =   sample[0, :8].flatten().cpu()
		print(lowerCAmelCase_							)
		A_		:						Optional[int]	    =   torch.tensor(lowerCAmelCase_							)

		assert torch_all_close(lowerCAmelCase_     ,     lowerCAmelCase_     ,     atol=1e-3							)

