
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    a      ( )     ->str:




			'''simple docstring'''


			SCREAMING_SNAKE_CASE     						= '''https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'''
			SCREAMING_SNAKE_CASE     						= Image.open(requests.get(a		,						stream=a	).raw	).convert('''RGB'''	)

			return image


def    a      ( a	)     ->Any:




			'''simple docstring'''


			SCREAMING_SNAKE_CASE     						= []
			# 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    a      ( a		,						a		,						a	)     ->Optional[int]:




			'''simple docstring'''


			SCREAMING_SNAKE_CASE     						= dct.pop(a	)
			SCREAMING_SNAKE_CASE     						= val


def    a      ( a		,						a	)     ->Optional[Any]:




			'''simple docstring'''


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

						# next, set bias in the state dict
						SCREAMING_SNAKE_CASE     						= torch.cat((q_bias, torch.zeros_like(a		,						requires_grad=a	), v_bias)	)
						SCREAMING_SNAKE_CASE     						= qkv_bias


def    a      ( a	)     ->Optional[int]:




			'''simple docstring'''


			SCREAMING_SNAKE_CASE     						= 364 if '''coco''' in model_name else 224
			SCREAMING_SNAKE_CASE     						= InstructBlipVisionConfig(image_size=a	).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     						= 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     						= 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     						= LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''		,						vocab_size=3_2001	).to_dict()
			elif "vicuna-13b" in model_name:
						SCREAMING_SNAKE_CASE     						= LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''		,						vocab_size=3_2001	).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     						= InstructBlipQFormerConfig(vocab_size=3_0523	).to_dict()
			SCREAMING_SNAKE_CASE     						= InstructBlipConfig(vision_config=a		,						text_config=a		,						qformer_config=a	)

			return config, image_size


@torch.no_grad()
def    a      ( a		,						a=None		,						a=False	)     ->Tuple:




			'''simple docstring'''


			SCREAMING_SNAKE_CASE     						= 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     						= 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     						= 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     						= get_blipa_config(a	)
			SCREAMING_SNAKE_CASE     						= InstructBlipForConditionalGeneration(a	).eval()

			SCREAMING_SNAKE_CASE     						= {
			    '''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     						= model_name_to_original[model_name]

			# load original model
			print('''Loading original model...'''	)
			SCREAMING_SNAKE_CASE     						= '''cuda:1''' if torch.cuda.is_available() else '''cpu'''
			SCREAMING_SNAKE_CASE     						= '''cuda:2''' if torch.cuda.is_available() else '''cpu'''
			SCREAMING_SNAKE_CASE						,							SCREAMING_SNAKE_CASE						,							SCREAMING_SNAKE_CASE     						= load_model_and_preprocess(
			    name=a		,						model_type=a		,						is_eval=a		,						device=a	)
			original_model.eval()
			print('''Done!'''	)

			# update state dict keys
			SCREAMING_SNAKE_CASE     						= original_model.state_dict()
			SCREAMING_SNAKE_CASE     						= create_rename_keys(a	)
			for src, dest in rename_keys:
						rename_key(a		,						a		,						a	)

			# some keys can be renamed efficiently
			for key, val in state_dict.copy().items():
						SCREAMING_SNAKE_CASE     						= state_dict.pop(a	)
						if key.startswith('''Qformer.bert'''	):
									SCREAMING_SNAKE_CASE     						= key.replace('''Qformer.bert'''		,						'''qformer'''	)
						if "attention.self" in key:
									SCREAMING_SNAKE_CASE     						= key.replace('''self'''		,						'''attention'''	)
						if "llm_proj" in key:
									SCREAMING_SNAKE_CASE     						= key.replace('''llm_proj'''		,						'''language_projection'''	)
						if "t5_proj" in key:
									SCREAMING_SNAKE_CASE     						= key.replace('''t5_proj'''		,						'''language_projection'''	)
						if key.startswith('''llm_model'''	):
									SCREAMING_SNAKE_CASE     						= key.replace('''llm_model'''		,						'''language_model'''	)
						if key.startswith('''t5'''	):
									SCREAMING_SNAKE_CASE     						= key.replace('''t5'''		,						'''language'''	)
						SCREAMING_SNAKE_CASE     						= val

			# read in qv biases
			read_in_q_v_bias(a		,						a	)

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

			SCREAMING_SNAKE_CASE     						= load_demo_image()
			SCREAMING_SNAKE_CASE     						= '''What is unusual about this image?'''

			# create processor
			SCREAMING_SNAKE_CASE     						= BlipImageProcessor(
			    size={'''height''': image_size, '''width''': image_size}		,						image_mean=a		,						image_std=a	)
			SCREAMING_SNAKE_CASE     						= InstructBlipProcessor(
			    image_processor=a		,						tokenizer=a		,						qformer_tokenizer=a		,						)
			SCREAMING_SNAKE_CASE     						= processor(images=a		,						text=a		,						return_tensors='''pt'''	).to(a	)

			# make sure processor creates exact same pixel values
			SCREAMING_SNAKE_CASE     						= vis_processors['''eval'''](a	).unsqueeze(0	).to(a	)
			SCREAMING_SNAKE_CASE     						= inputs.pixel_values
			assert torch.allclose(original_pixel_values.to(pixel_values.device	)		,						a	)

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

			print('''Generating with original model...'''	)
			SCREAMING_SNAKE_CASE     						= 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     						= hf_model.generate(
			    **a		,						do_sample=a		,						num_beams=5		,						max_length=256		,						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     						= 2
			print('''Original generation:'''		,						a	)
			SCREAMING_SNAKE_CASE     						= processor.batch_decode(a		,						skip_special_tokens=a	)
			SCREAMING_SNAKE_CASE     						= [text.strip() for text in output_text]
			print('''HF generation:'''		,						a	)

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

			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					   =       argparse.ArgumentParser()
							__lowerCAmelCase					   =       [
							    '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					   =       parser.parse_args()

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

import os
from argparse import ArgumentParser
from typing import List

import torch.utils.data

from datasets import Dataset, IterableDataset
from datasets.distributed import split_dataset_by_node


__lowerCAmelCase					   =       4
__lowerCAmelCase					   =       3





class     lowerCamelCase      (       __lowerCamelCase					):
	pass



def    a      ( a	)     ->Union[str, Any]:




			'''simple docstring'''


			for shard in shards:
						for i in range(a	):
									yield {"i": i, "shard": shard}


def    a      ( )     ->List[str]:




			'''simple docstring'''


			SCREAMING_SNAKE_CASE     						= int(os.environ['''RANK''']	)
			SCREAMING_SNAKE_CASE     						= int(os.environ['''WORLD_SIZE''']	)

			SCREAMING_SNAKE_CASE     						= ArgumentParser()
			parser.add_argument('''--streaming'''		,						type=a	)
			parser.add_argument('''--local_rank'''		,						type=a	)
			parser.add_argument('''--num_workers'''		,						type=a		,						default=0	)
			SCREAMING_SNAKE_CASE     						= parser.parse_args()
			SCREAMING_SNAKE_CASE     						= args.streaming
			SCREAMING_SNAKE_CASE     						= args.num_workers

			SCREAMING_SNAKE_CASE     						= {'''shards''': [F"""shard_{shard_idx}""" for shard_idx in range(a	)]}
			SCREAMING_SNAKE_CASE     						= IterableDataset.from_generator(a		,						gen_kwargs=a	)
			if not streaming:
						SCREAMING_SNAKE_CASE     						= Dataset.from_list(list(a	)	)

			SCREAMING_SNAKE_CASE     						= split_dataset_by_node(a		,						rank=a		,						world_size=a	)
			SCREAMING_SNAKE_CASE     						= torch.utils.data.DataLoader(a		,						num_workers=a	)

			SCREAMING_SNAKE_CASE     						= NUM_SHARDS * NUM_ITEMS_PER_SHARD
			SCREAMING_SNAKE_CASE     						= full_size // world_size
			expected_local_size += int(rank < (full_size % world_size)	)

			SCREAMING_SNAKE_CASE     						= sum(1 for _ in dataloader	)
			if local_size != expected_local_size:
						raise FailedTestError(F"""local_size {local_size} != expected_local_size {expected_local_size}"""	)


if __name__ == "__main__":
							main()