



'''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_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_     (			):
		__SCREAMING_SNAKE_CASE :	Any					 =	"""https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"""
		__SCREAMING_SNAKE_CASE :	List[Any]					 =	Image.open(requests.get(_lowerCamelCase ,   stream=_lowerCamelCase      ).raw      ).convert("""RGB"""      )

		return image


def        lowerCAmelCase_     (			_lowerCamelCase:     Tuple      ):
		__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	[]
		# 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_     (			_lowerCamelCase:     Dict ,   _lowerCamelCase:     Optional[int] ,   _lowerCamelCase:     List[Any]      ):
		__SCREAMING_SNAKE_CASE :	Tuple					 =	dct.pop(_lowerCamelCase      )
		__SCREAMING_SNAKE_CASE :	List[Any]					 =	val


def        lowerCAmelCase_     (			_lowerCamelCase:     List[Any] ,   _lowerCamelCase:     Optional[int]      ):
		for i in range(config.vision_config.num_hidden_layers      ):
				# read in original q and v biases
				__SCREAMING_SNAKE_CASE :	List[Any]					 =	state_dict.pop(F"visual_encoder.blocks.{i}.attn.q_bias"      )
				__SCREAMING_SNAKE_CASE :	Tuple					 =	state_dict.pop(F"visual_encoder.blocks.{i}.attn.v_bias"      )

				# next, set bias in the state dict
				__SCREAMING_SNAKE_CASE :	Dict					 =	torch.cat((q_bias, torch.zeros_like(_lowerCamelCase ,   requires_grad=_lowerCamelCase      ), v_bias)      )
				__SCREAMING_SNAKE_CASE :	Optional[int]					 =	qkv_bias


def        lowerCAmelCase_     (			_lowerCamelCase:     Tuple      ):
		__SCREAMING_SNAKE_CASE :	int					 =	3_64 if """coco""" in model_name else 2_24
		__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	InstructBlipVisionConfig(image_size=_lowerCamelCase      ).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:
				__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	TaConfig.from_pretrained("""google/flan-t5-xl""" ,   dense_act_fn="""gelu""" ,   bos_token_id=1      ).to_dict()
		elif "t5-xxl" in model_name:
				__SCREAMING_SNAKE_CASE :	Optional[int]					 =	TaConfig.from_pretrained("""google/flan-t5-xxl""" ,   dense_act_fn="""gelu""" ,   bos_token_id=1      ).to_dict()
		elif "vicuna-7b" in model_name:
				__SCREAMING_SNAKE_CASE :	Optional[int]					 =	LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf""" ,   vocab_size=3_20_01      ).to_dict()
		elif "vicuna-13b" in model_name:
				__SCREAMING_SNAKE_CASE :	Tuple					 =	LlamaConfig.from_pretrained("""decapoda-research/llama-13b-hf""" ,   vocab_size=3_20_01      ).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
		__SCREAMING_SNAKE_CASE :	Any					 =	InstructBlipQFormerConfig(vocab_size=3_05_23      ).to_dict()
		__SCREAMING_SNAKE_CASE :	Dict					 =	InstructBlipConfig(vision_config=_lowerCamelCase ,   text_config=_lowerCamelCase ,   qformer_config=_lowerCamelCase      )

		return config, image_size


@torch.no_grad()
def        lowerCAmelCase_     (			_lowerCamelCase:     Union[str, Any] ,   _lowerCamelCase:     Dict=None ,   _lowerCamelCase:     Union[str, Any]=False      ):
		__SCREAMING_SNAKE_CASE :	List[Any]					 =	AutoTokenizer.from_pretrained("""bert-base-uncased""" ,   truncation_side="""left"""      )
		qformer_tokenizer.add_special_tokens({"""bos_token""": """[DEC]"""}      )

		if "t5" in model_name:
				__SCREAMING_SNAKE_CASE :	List[Any]					 =	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>"})
				__SCREAMING_SNAKE_CASE :	Optional[int]					 =	LlamaTokenizerFast.from_pretrained(
				    """huggyllama/llama-7b""" ,   truncation_side="""left""" ,   bos_token="""</s>""" ,   unk_token="""</s>"""      )
				tokenizer.add_special_tokens({"""pad_token""": """[PAD]"""}      )

		__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Optional[int]					 =	get_blipa_config(_lowerCamelCase      )
		__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	InstructBlipForConditionalGeneration(_lowerCamelCase      ).eval()

		__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	{
		    """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"""),
		}

		__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Optional[int]					 =	model_name_to_original[model_name]

		# load original model
		print("""Loading original model..."""      )
		__SCREAMING_SNAKE_CASE :	Any					 =	"""cuda:1""" if torch.cuda.is_available() else """cpu"""
		__SCREAMING_SNAKE_CASE :	Tuple					 =	"""cuda:2""" if torch.cuda.is_available() else """cpu"""
		__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE ,		__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	load_model_and_preprocess(
		    name=_lowerCamelCase ,   model_type=_lowerCamelCase ,   is_eval=_lowerCamelCase ,   device=_lowerCamelCase      )
		original_model.eval()
		print("""Done!"""      )

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

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

		# read in qv biases
		read_in_q_v_bias(_lowerCamelCase ,   _lowerCamelCase      )

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

		__SCREAMING_SNAKE_CASE :	List[str]					 =	load_demo_image()
		__SCREAMING_SNAKE_CASE :	List[Any]					 =	"""What is unusual about this image?"""

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

		# make sure processor creates exact same pixel values
		__SCREAMING_SNAKE_CASE :	Any					 =	vis_processors["""eval"""](_lowerCamelCase      ).unsqueeze(0      ).to(_lowerCamelCase      )
		__SCREAMING_SNAKE_CASE :	Optional[int]					 =	inputs.pixel_values
		assert torch.allclose(original_pixel_values.to(pixel_values.device      ) ,   _lowerCamelCase      )

		original_model.to(_lowerCamelCase      )
		hf_model.to(_lowerCamelCase      )
		with torch.no_grad():
				if "vicuna" in model_name:
						__SCREAMING_SNAKE_CASE :	Optional[int]					 =	original_model({"""image""": original_pixel_values, """text_input""": [prompt]}      ).logits
						__SCREAMING_SNAKE_CASE :	List[str]					 =	hf_model(**_lowerCamelCase      ).logits
				else:
						__SCREAMING_SNAKE_CASE :	str					 =	original_model(
						    {"""image""": original_pixel_values, """text_input""": [prompt], """text_output""": ["""\n"""]}      ).logits
						__SCREAMING_SNAKE_CASE :	Any					 =	tokenizer("""\n""" ,   return_tensors="""pt"""      ).input_ids.to(_lowerCamelCase      )
						__SCREAMING_SNAKE_CASE :	str					 =	label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id ,   -1_00      )
						__SCREAMING_SNAKE_CASE :	str					 =	hf_model(**_lowerCamelCase ,   labels=_lowerCamelCase      ).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
		__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	1E-4 if """vicuna""" in model_name else 1E-5
		assert torch.allclose(original_logits.to(logits.device      ) ,   _lowerCamelCase ,   atol=_lowerCamelCase      )
		print("""Looks ok!"""      )

		print("""Generating with original model..."""      )
		__SCREAMING_SNAKE_CASE :	Any					 =	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..."""      )
		__SCREAMING_SNAKE_CASE :	str					 =	hf_model.generate(
		    **_lowerCamelCase ,   do_sample=_lowerCamelCase ,   num_beams=5 ,   max_length=2_56 ,   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?
				__SCREAMING_SNAKE_CASE :	int					 =	2
		print("""Original generation:""" ,   _lowerCamelCase      )
		__SCREAMING_SNAKE_CASE :	Tuple					 =	processor.batch_decode(_lowerCamelCase ,   skip_special_tokens=_lowerCamelCase      )
		__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	[text.strip() for text in output_text]
		print("""HF generation:""" ,   _lowerCamelCase      )

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

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


if __name__ == "__main__":
						UpperCamelCase__ :			List[str]												=				argparse.ArgumentParser()
						UpperCamelCase__ :			Union[str, 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''',
						)

						UpperCamelCase__ :			str												=				parser.parse_args()

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




'''simple docstring'''




from unittest.mock import Mock, patch

from file_transfer.send_file import send_file


@patch("""socket.socket"""      )
@patch("""builtins.open"""      )
def        lowerCAmelCase_     (			_lowerCamelCase:     Dict ,   _lowerCamelCase:     List[Any]      ):
		# ===== initialization =====
		__SCREAMING_SNAKE_CASE :	Optional[int]					 =	Mock()
		__SCREAMING_SNAKE_CASE :	Optional[int]					 =	conn, Mock()
		__SCREAMING_SNAKE_CASE :	str					 =	iter([1, None]      )
		__SCREAMING_SNAKE_CASE :	Any					 =	lambda _lowerCamelCase      : next(_lowerCamelCase      )

		# ===== invoke =====
		send_file(filename="""mytext.txt""" ,   testing=_lowerCamelCase      )

		# ===== ensurance =====
		sock.assert_called_once()
		sock.return_value.bind.assert_called_once()
		sock.return_value.listen.assert_called_once()
		sock.return_value.accept.assert_called_once()
		conn.recv.assert_called_once()

		file.return_value.__enter__.assert_called_once()
		file.return_value.__enter__.return_value.read.assert_called()

		conn.send.assert_called_once()
		conn.close.assert_called_once()
		sock.return_value.shutdown.assert_called_once()
		sock.return_value.close.assert_called_once()