



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_float32 (there's also the fix_lavis branch)
# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
# same for Vicuna-13b
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipImageProcessor,
    InstructBlipConfig,
    InstructBlipForConditionalGeneration,
    InstructBlipProcessor,
    InstructBlipQFormerConfig,
    InstructBlipVisionConfig,
    LlamaConfig,
    LlamaTokenizerFast,
    TaConfig,
    TaTokenizerFast,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
def 	__lowerCAmelCase						(		):


						'''simple docstring'''






						snake_case_   : Optional[int]						=			"""https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"""
						snake_case_   : Optional[Any]						=			Image.open(requests.get(__UpperCamelCase					, stream=__UpperCamelCase    ).raw    ).convert("""RGB"""    )

						return image
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Optional[Any]    ):


						'''simple docstring'''






						snake_case_   : Tuple						=			[]
						# 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.embeddings.layernorm.weight""")    )
						rename_keys.append(("""Qformer.bert.embeddings.LayerNorm.bias""", """qformer.embeddings.layernorm.bias""")    )

						# fmt: on
						return rename_keys
def 	__lowerCAmelCase						(		__UpperCamelCase      :			List[Any]					, __UpperCamelCase      :			Optional[int]					, __UpperCamelCase      :			List[Any]    ):


						'''simple docstring'''






						snake_case_   : Union[str, Any]						=			dct.pop(__UpperCamelCase    )
						snake_case_   : Optional[int]						=			val
def 	__lowerCAmelCase						(		__UpperCamelCase      :			str					, __UpperCamelCase      :			str    ):


						'''simple docstring'''






						for i in range(config.vision_config.num_hidden_layers    ):
												# read in original q and v biases
												snake_case_   : Union[str, Any]						=			state_dict.pop(F'visual_encoder.blocks.{i}.attn.q_bias'    )
												snake_case_   : Optional[Any]						=			state_dict.pop(F'visual_encoder.blocks.{i}.attn.v_bias'    )

												# next, set bias in the state dict
												snake_case_   : Union[str, Any]						=			torch.cat((q_bias, torch.zeros_like(__UpperCamelCase					, requires_grad=__UpperCamelCase    ), v_bias)    )
												snake_case_   : Optional[Any]						=			qkv_bias
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Union[str, Any]    ):


						'''simple docstring'''






						snake_case_   : Dict						=			3_6_4 if """coco""" in model_name else 2_2_4
						snake_case_   : str						=			InstructBlipVisionConfig(image_size=__UpperCamelCase    ).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 "t5-xl" in model_name:
												snake_case_   : 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:
												snake_case_   : int						=			TaConfig.from_pretrained("""google/flan-t5-xxl"""					, dense_act_fn="""gelu"""					, bos_token_id=1    ).to_dict()
						elif "vicuna-7b" in model_name:
												snake_case_   : int						=			LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf"""					, vocab_size=3_2_0_0_1    ).to_dict()
						elif "vicuna-13b" in model_name:
												snake_case_   : Union[str, Any]						=			LlamaConfig.from_pretrained("""decapoda-research/llama-13b-hf"""					, vocab_size=3_2_0_0_1    ).to_dict()
						else:
												raise ValueError("""Model name not supported"""    )

						# the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
						snake_case_   : str						=			InstructBlipQFormerConfig(vocab_size=3_0_5_2_3    ).to_dict()
						snake_case_   : List[str]						=			InstructBlipConfig(vision_config=__UpperCamelCase					, text_config=__UpperCamelCase					, qformer_config=__UpperCamelCase    )

						return config, image_size
@torch.no_grad()
def 	__lowerCAmelCase						(		__UpperCamelCase      :			List[str]					, __UpperCamelCase      :			List[Any]=None					, __UpperCamelCase      :			Optional[Any]=False    ):


						'''simple docstring'''






						snake_case_   : Optional[int]						=			AutoTokenizer.from_pretrained("""bert-base-uncased"""					, truncation_side="""left"""    )
						qformer_tokenizer.add_special_tokens({"""bos_token""": """[DEC]"""}    )

						if "t5" in model_name:
												snake_case_   : Tuple						=			TaTokenizerFast.from_pretrained("""google/flan-t5-xl"""					, truncation_side="""left"""    )
						elif "vicuna" in model_name:
												# the following was used in the original implementation:
												# tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
												# tokenizer.add_special_tokens({"pad_token": "[PAD]"})
												# tokenizer.add_special_tokens({"bos_token": "</s>"})
												# tokenizer.add_special_tokens({"eos_token": "</s>"})
												# tokenizer.add_special_tokens({"unk_token": "</s>"})
												snake_case_   : Any						=			LlamaTokenizerFast.from_pretrained(
												    """huggyllama/llama-7b"""					, truncation_side="""left"""					, bos_token="""</s>"""					, unk_token="""</s>"""    )
												tokenizer.add_special_tokens({"""pad_token""": """[PAD]"""}    )

						snake_case_   : Union[str, Any]						=			get_blipa_config(__UpperCamelCase    )
						snake_case_   : Optional[Any]						=			InstructBlipForConditionalGeneration(__UpperCamelCase    ).eval()

						snake_case_   : Tuple						=			{
						    """instructblip-vicuna-7b""": ("""blip2_vicuna_instruct""", """vicuna7b"""),
						    """instructblip-vicuna-13b""": ("""blip2_vicuna_instruct""", """vicuna13b"""),
						    """instructblip-flan-t5-xl""": ("""blip2_t5_instruct""", """flant5xl"""),
						    """instructblip-flan-t5-xxl""": ("""blip2_t5_instruct""", """flant5xxl"""),
						}

						snake_case_   : int						=			model_name_to_original[model_name]

						# load original model
						print("""Loading original model..."""    )
						snake_case_   : List[Any]						=			"""cuda:1""" if torch.cuda.is_available() else """cpu"""
						snake_case_   : Dict						=			"""cuda:2""" if torch.cuda.is_available() else """cpu"""
						snake_case_   : Union[str, Any]						=			load_model_and_preprocess(
						    name=__UpperCamelCase					, model_type=__UpperCamelCase					, is_eval=__UpperCamelCase					, device=__UpperCamelCase    )
						original_model.eval()
						print("""Done!"""    )

						# update state dict keys
						snake_case_   : Any						=			original_model.state_dict()
						snake_case_   : List[str]						=			create_rename_keys(__UpperCamelCase    )
						for src, dest in rename_keys:
												rename_key(__UpperCamelCase					, __UpperCamelCase					, __UpperCamelCase    )

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

						# read in qv biases
						read_in_q_v_bias(__UpperCamelCase					, __UpperCamelCase    )

						# note: weights get loaded in torch.float32 by default
						hf_model.load_state_dict(__UpperCamelCase					, strict=__UpperCamelCase    )

						snake_case_   : Dict						=			load_demo_image()
						snake_case_   : Optional[Any]						=			"""What is unusual about this image?"""

						# create processor
						snake_case_   : int						=			BlipImageProcessor(
						    size={"""height""": image_size, """width""": image_size}					, image_mean=__UpperCamelCase					, image_std=__UpperCamelCase    )
						snake_case_   : List[str]						=			InstructBlipProcessor(
						    image_processor=__UpperCamelCase					, tokenizer=__UpperCamelCase					, qformer_tokenizer=__UpperCamelCase					, )
						snake_case_   : int						=			processor(images=__UpperCamelCase					, text=__UpperCamelCase					, return_tensors="""pt"""    ).to(__UpperCamelCase    )

						# make sure processor creates exact same pixel values
						snake_case_   : Tuple						=			vis_processors["""eval"""](__UpperCamelCase    ).unsqueeze(0    ).to(__UpperCamelCase    )
						snake_case_   : Dict						=			inputs.pixel_values
						assert torch.allclose(original_pixel_values.to(pixel_values.device    )					, __UpperCamelCase    )

						original_model.to(__UpperCamelCase    )
						hf_model.to(__UpperCamelCase    )
						with torch.no_grad():
												if "vicuna" in model_name:
																		snake_case_   : Any						=			original_model({"""image""": original_pixel_values, """text_input""": [prompt]}    ).logits
																		snake_case_   : Dict						=			hf_model(**__UpperCamelCase    ).logits
												else:
																		snake_case_   : Union[str, Any]						=			original_model(
																		    {"""image""": original_pixel_values, """text_input""": [prompt], """text_output""": ["""\n"""]}    ).logits
																		snake_case_   : List[Any]						=			tokenizer("""\n"""					, return_tensors="""pt"""    ).input_ids.to(__UpperCamelCase    )
																		snake_case_   : int						=			label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id					, -1_0_0    )
																		snake_case_   : int						=			hf_model(**__UpperCamelCase					, labels=__UpperCamelCase    ).logits

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

						# assert values
						assert original_logits.shape == logits.shape
						snake_case_   : str						=			1E-4 if """vicuna""" in model_name else 1E-5
						assert torch.allclose(original_logits.to(logits.device    )					, __UpperCamelCase					, atol=__UpperCamelCase    )
						print("""Looks ok!"""    )

						print("""Generating with original model..."""    )
						snake_case_   : int						=			original_model.generate({"""image""": original_pixel_values, """prompt""": prompt}					, num_beams=5    )

						# important: we need to cast the weights of the HF model to the appropriate type
						print("""Generating with HF model..."""    )
						snake_case_   : str						=			hf_model.generate(
						    **__UpperCamelCase					, do_sample=__UpperCamelCase					, num_beams=5					, max_length=2_5_6					, min_length=1					, top_p=0.9					, repetition_penalty=1.5					, length_penalty=1.0					, temperature=1					, )
						if "vicuna" in model_name:
												# convert output id 0 to 2 (eos_token_id)
												# TODO add this in the generate method?
												snake_case_   : Any						=			2
						print("""Original generation:"""					, __UpperCamelCase    )
						snake_case_   : List[Any]						=			processor.batch_decode(__UpperCamelCase					, skip_special_tokens=__UpperCamelCase    )
						snake_case_   : str						=			[text.strip() for text in output_text]
						print("""HF generation:"""					, __UpperCamelCase    )

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

						if push_to_hub:
												processor.push_to_hub(F'Salesforce/{model_name}'    )
												hf_model.push_to_hub(F'Salesforce/{model_name}'    )


if __name__ == "__main__":
					__lowerCAmelCase       :					List[str]      					=					argparse.ArgumentParser()
					__lowerCAmelCase       :					Optional[Any]      					=					[
					    '''instructblip-vicuna-7b''',
					    '''instructblip-vicuna-13b''',
					    '''instructblip-flan-t5-xl''',
					    '''instructblip-flan-t5-xxl''',
					]
					parser.add_argument(
					    '''--model_name''',
					    default='''instructblip-flan-t5-xl''',
					    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       :					str      					=					parser.parse_args()

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






"""simple docstring"""





import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


__lowerCAmelCase       :					Optional[int]      					=					logging.get_logger(__name__)
def 	__lowerCAmelCase						(		__UpperCamelCase      :			Tuple					, __UpperCamelCase      :			Optional[Any]    ):


						'''simple docstring'''






						snake_case_   : Union[str, Any]						=			nn.functional.normalize(__UpperCamelCase    )
						snake_case_   : Tuple						=			nn.functional.normalize(__UpperCamelCase    )
						return torch.mm(__UpperCamelCase					, normalized_text_embeds.t()    )






class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__		):





		"""simple docstring"""






		_lowerCamelCase			  =       CLIPConfig

		_lowerCamelCase			  =       ['''CLIPEncoderLayer''']





		def __init__(					self							,    _lowercase   ) ->						Union[str, Any]:







								'''simple docstring'''

								super().__init__(_lowercase   )

								snake_case_   : Tuple						=			CLIPVisionModel(config.vision_config   )
								snake_case_   : int						=			nn.Linear(config.vision_config.hidden_size							,    config.projection_dim							,    bias=_lowercase   )

								snake_case_   : Optional[Any]						=			nn.Parameter(torch.ones(1_7							,    config.projection_dim   )							,    requires_grad=_lowercase   )
								snake_case_   : Dict						=			nn.Parameter(torch.ones(3							,    config.projection_dim   )							,    requires_grad=_lowercase   )

								snake_case_   : Any						=			nn.Parameter(torch.ones(1_7   )							,    requires_grad=_lowercase   )
								snake_case_   : List[str]						=			nn.Parameter(torch.ones(3   )							,    requires_grad=_lowercase   )





		@torch.no_grad()
		def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase   ) ->						Any:







								'''simple docstring'''

								snake_case_   : int						=			self.vision_model(_lowercase   )[1]  # pooled_output
								snake_case_   : str						=			self.visual_projection(_lowercase   )

								# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
								snake_case_   : Dict						=			cosine_distance(_lowercase							,    self.special_care_embeds   ).cpu().float().numpy()
								snake_case_   : List[str]						=			cosine_distance(_lowercase							,    self.concept_embeds   ).cpu().float().numpy()

								snake_case_   : Any						=			[]
								snake_case_   : Any						=			image_embeds.shape[0]
								for i in range(_lowercase   ):
														snake_case_   : List[Any]						=			{"""special_scores""": {}, """special_care""": [], """concept_scores""": {}, """bad_concepts""": []}

														# increase this value to create a stronger `nfsw` filter
														# at the cost of increasing the possibility of filtering benign images
														snake_case_   : int						=			0.0

														for concept_idx in range(len(special_cos_dist[0]   )   ):
																				snake_case_   : List[str]						=			special_cos_dist[i][concept_idx]
																				snake_case_   : Union[str, Any]						=			self.special_care_embeds_weights[concept_idx].item()
																				snake_case_   : Tuple						=			round(concept_cos - concept_threshold + adjustment							,    3   )
																				if result_img["special_scores"][concept_idx] > 0:
																										result_img["special_care"].append({concept_idx, result_img["""special_scores"""][concept_idx]}   )
																										snake_case_   : Dict						=			0.01

														for concept_idx in range(len(cos_dist[0]   )   ):
																				snake_case_   : int						=			cos_dist[i][concept_idx]
																				snake_case_   : List[Any]						=			self.concept_embeds_weights[concept_idx].item()
																				snake_case_   : List[str]						=			round(concept_cos - concept_threshold + adjustment							,    3   )
																				if result_img["concept_scores"][concept_idx] > 0:
																										result_img["bad_concepts"].append(_lowercase   )

														result.append(_lowercase   )

								snake_case_   : Union[str, Any]						=			[len(res["""bad_concepts"""]   ) > 0 for res in result]

								return images, has_nsfw_concepts





		@torch.no_grad()
		def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase   ) ->						Tuple:







								'''simple docstring'''

								snake_case_   : Optional[Any]						=			self.vision_model(_lowercase   )[1]  # pooled_output
								snake_case_   : List[str]						=			self.visual_projection(_lowercase   )

								snake_case_   : str						=			cosine_distance(_lowercase							,    self.special_care_embeds   )
								snake_case_   : Optional[int]						=			cosine_distance(_lowercase							,    self.concept_embeds   )

								# increase this value to create a stronger `nsfw` filter
								# at the cost of increasing the possibility of filtering benign images
								snake_case_   : Tuple						=			0.0

								snake_case_   : List[Any]						=			special_cos_dist - self.special_care_embeds_weights + adjustment
								# special_scores = special_scores.round(decimals=3)
								snake_case_   : str						=			torch.any(special_scores > 0							,    dim=1   )
								snake_case_   : List[str]						=			special_care * 0.01
								snake_case_   : Optional[int]						=			special_adjustment.unsqueeze(1   ).expand(-1							,    cos_dist.shape[1]   )

								snake_case_   : Optional[Any]						=			(cos_dist - self.concept_embeds_weights) + special_adjustment
								# concept_scores = concept_scores.round(decimals=3)
								snake_case_   : str						=			torch.any(concept_scores > 0							,    dim=1   )

								return images, has_nsfw_concepts

