
"""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  lowercase     ():
						__lowerCAmelCase     =  """https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
						__lowerCAmelCase     =  Image.open(requests.get(_lowerCAmelCase ,    stream=_lowerCAmelCase						).raw						).convert("""RGB"""						)

						return image

def  lowercase     (_lowerCAmelCase						):
						__lowerCAmelCase     =  []
						# 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  lowercase     (_lowerCAmelCase ,    _lowerCAmelCase ,    _lowerCAmelCase						):
						__lowerCAmelCase     =  dct.pop(_lowerCAmelCase						)
						__lowerCAmelCase     =  val

def  lowercase     (_lowerCAmelCase ,    _lowerCAmelCase						):
						for i in range(config.vision_config.num_hidden_layers						):
												# read in original q and v biases
												__lowerCAmelCase     =  state_dict.pop(f"""visual_encoder.blocks.{i}.attn.q_bias"""						)
												__lowerCAmelCase     =  state_dict.pop(f"""visual_encoder.blocks.{i}.attn.v_bias"""						)

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

def  lowercase     (_lowerCAmelCase ,    _lowerCAmelCase						):
						__lowerCAmelCase     =  364 if """coco""" in model_name else 224
						__lowerCAmelCase     =  BlipaVisionConfig(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 "opt-2.7b" in model_name:
												__lowerCAmelCase     =  OPTConfig.from_pretrained("""facebook/opt-2.7b""" ,    eos_token_id=_lowerCAmelCase						).to_dict()
						elif "opt-6.7b" in model_name:
												__lowerCAmelCase     =  OPTConfig.from_pretrained("""facebook/opt-6.7b""" ,    eos_token_id=_lowerCAmelCase						).to_dict()
						elif "t5-xl" in model_name:
												__lowerCAmelCase     =  TaConfig.from_pretrained("""google/flan-t5-xl""" ,    dense_act_fn="""gelu""" ,    bos_token_id=1						).to_dict()
						elif "t5-xxl" in model_name:
												__lowerCAmelCase     =  TaConfig.from_pretrained("""google/flan-t5-xxl""" ,    dense_act_fn="""gelu""" ,    bos_token_id=1						).to_dict()

						__lowerCAmelCase     =  BlipaConfig(vision_config=_lowerCAmelCase ,    text_config=_lowerCAmelCase						)

						return config, image_size

@torch.no_grad()
def  lowercase     (_lowerCAmelCase ,    _lowerCAmelCase=None ,    _lowerCAmelCase=False						):
						__lowerCAmelCase     =  (
						    AutoTokenizer.from_pretrained("""facebook/opt-2.7b"""						)
						    if """opt""" in model_name
						    else AutoTokenizer.from_pretrained("""google/flan-t5-xl"""						)
						)
						__lowerCAmelCase     =  tokenizer("""\n""" ,    add_special_tokens=_lowerCAmelCase						).input_ids[0]
						__lowerCAmelCase   ,     __lowerCAmelCase     =  get_blipa_config(_lowerCAmelCase ,    eos_token_id=_lowerCAmelCase						)

						__lowerCAmelCase     =  BlipaForConditionalGeneration(_lowerCAmelCase						).eval()

						__lowerCAmelCase     =  {
						    """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"""),
						}

						__lowerCAmelCase   ,     __lowerCAmelCase     =  model_name_to_original[model_name]

						# load original model
						print("""Loading original model..."""						)
						__lowerCAmelCase     =  """cuda""" if torch.cuda.is_available() else """cpu"""
						__lowerCAmelCase   ,     __lowerCAmelCase   ,     __lowerCAmelCase     =  load_model_and_preprocess(
						    name=_lowerCAmelCase ,    model_type=_lowerCAmelCase ,    is_eval=_lowerCAmelCase ,    device=_lowerCAmelCase						)
						original_model.eval()
						print("""Done!"""						)

						# update state dict keys
						__lowerCAmelCase     =  original_model.state_dict()
						__lowerCAmelCase     =  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():
												__lowerCAmelCase     =  state_dict.pop(_lowerCAmelCase						)
												if key.startswith("""Qformer.bert"""						):
																		__lowerCAmelCase     =  key.replace("""Qformer.bert""" ,    """qformer"""						)
												if "attention.self" in key:
																		__lowerCAmelCase     =  key.replace("""self""" ,    """attention"""						)
												if "opt_proj" in key:
																		__lowerCAmelCase     =  key.replace("""opt_proj""" ,    """language_projection"""						)
												if "t5_proj" in key:
																		__lowerCAmelCase     =  key.replace("""t5_proj""" ,    """language_projection"""						)
												if key.startswith("""opt"""						):
																		__lowerCAmelCase     =  key.replace("""opt""" ,    """language"""						)
												if key.startswith("""t5"""						):
																		__lowerCAmelCase     =  key.replace("""t5""" ,    """language"""						)
												__lowerCAmelCase     =  val

						# read in qv biases
						read_in_q_v_bias(_lowerCAmelCase ,    _lowerCAmelCase						)

						__lowerCAmelCase   ,     __lowerCAmelCase     =  hf_model.load_state_dict(_lowerCAmelCase ,    strict=_lowerCAmelCase						)
						assert len(_lowerCAmelCase						) == 0
						assert unexpected_keys == ["qformer.embeddings.position_ids"]

						__lowerCAmelCase     =  load_demo_image()
						__lowerCAmelCase     =  vis_processors["""eval"""](_lowerCAmelCase						).unsqueeze(0						).to(_lowerCAmelCase						)
						__lowerCAmelCase     =  tokenizer(["""\n"""] ,    return_tensors="""pt"""						).input_ids.to(_lowerCAmelCase						)

						# create processor
						__lowerCAmelCase     =  BlipImageProcessor(
						    size={"""height""": image_size, """width""": image_size} ,    image_mean=_lowerCAmelCase ,    image_std=_lowerCAmelCase						)
						__lowerCAmelCase     =  BlipaProcessor(image_processor=_lowerCAmelCase ,    tokenizer=_lowerCAmelCase						)
						__lowerCAmelCase     =  processor(images=_lowerCAmelCase ,    return_tensors="""pt"""						).pixel_values.to(_lowerCAmelCase						)

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

						original_model.to(_lowerCAmelCase						)
						hf_model.to(_lowerCAmelCase						)
						with torch.no_grad():
												if "opt" in model_name:
																		__lowerCAmelCase     =  original_model({"""image""": original_pixel_values, """text_input""": [""""""]}						).logits
																		__lowerCAmelCase     =  hf_model(_lowerCAmelCase ,    _lowerCAmelCase						).logits
												else:
																		__lowerCAmelCase     =  original_model(
																		    {"""image""": original_pixel_values, """text_input""": ["""\n"""], """text_output""": ["""\n"""]}						).logits
																		__lowerCAmelCase     =  input_ids.masked_fill(input_ids == tokenizer.pad_token_id ,    -100						)
																		__lowerCAmelCase     =  hf_model(_lowerCAmelCase ,    _lowerCAmelCase ,    labels=_lowerCAmelCase						).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":
												__lowerCAmelCase     =  torch.tensor(
												    [[-41.5_850, -4.4_440, -8.9_922], [-47.4_322, -5.9_143, -1.7_340]] ,    device=_lowerCAmelCase						)
												assert torch.allclose(logits[0, :3, :3] ,    _lowerCAmelCase ,    atol=1E-4						)
						elif model_name == "blip2-flan-t5-xl-coco":
												__lowerCAmelCase     =  torch.tensor(
												    [[-57.0_109, -9.8_967, -12.6_280], [-68.6_578, -12.7_191, -10.5_065]] ,    device=_lowerCAmelCase						)
						else:
												# cast to same type
												__lowerCAmelCase     =  logits.dtype
												assert torch.allclose(original_logits.to(_lowerCAmelCase						) ,    _lowerCAmelCase ,    atol=1E-2						)
						print("""Looks ok!"""						)

						print("""Generating a caption..."""						)
						__lowerCAmelCase     =  """"""
						__lowerCAmelCase     =  tokenizer(_lowerCAmelCase ,    return_tensors="""pt"""						).input_ids.to(_lowerCAmelCase						)

						__lowerCAmelCase     =  original_model.generate({"""image""": original_pixel_values}						)
						__lowerCAmelCase     =  hf_model.generate(
						    _lowerCAmelCase ,    _lowerCAmelCase ,    do_sample=_lowerCAmelCase ,    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:""" ,    _lowerCAmelCase						)
						__lowerCAmelCase     =  input_ids.shape[1]
						__lowerCAmelCase     =  processor.batch_decode(outputs[:, prompt_length:] ,    skip_special_tokens=_lowerCAmelCase						)
						__lowerCAmelCase     =  [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"""nielsr/{model_name}"""						)
												hf_model.push_to_hub(f"""nielsr/{model_name}"""						)


if __name__ == "__main__":
			SCREAMING_SNAKE_CASE_					     =	argparse.ArgumentParser()
			SCREAMING_SNAKE_CASE_					     =	[
			    '''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''',
			)

			SCREAMING_SNAKE_CASE_					     =	parser.parse_args()

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


"""simple docstring"""
import operator as op


SCREAMING_SNAKE_CASE_					     =	'''scaler.pt'''
SCREAMING_SNAKE_CASE_					     =	'''pytorch_model'''
SCREAMING_SNAKE_CASE_					     =	'''random_states'''
SCREAMING_SNAKE_CASE_					     =	'''optimizer'''
SCREAMING_SNAKE_CASE_					     =	'''scheduler'''
SCREAMING_SNAKE_CASE_					     =	'''pytorch_model.bin'''
SCREAMING_SNAKE_CASE_					     =	'''pytorch_model.bin.index.json'''
SCREAMING_SNAKE_CASE_					     =	'''model.safetensors'''
SCREAMING_SNAKE_CASE_					     =	'''model.safetensors.index.json'''
SCREAMING_SNAKE_CASE_					     =	'''1.10.2'''
SCREAMING_SNAKE_CASE_					     =	'''py38'''
SCREAMING_SNAKE_CASE_					     =	'''4.17.0'''
SCREAMING_SNAKE_CASE_					     =	['''ml.p3.16xlarge''', '''ml.p3dn.24xlarge''', '''ml.p4dn.24xlarge''']
SCREAMING_SNAKE_CASE_					     =	['''FULL_SHARD''', '''SHARD_GRAD_OP''', '''NO_SHARD''', '''HYBRID_SHARD''', '''HYBRID_SHARD_ZERO2''']
SCREAMING_SNAKE_CASE_					     =	['''TRANSFORMER_BASED_WRAP''', '''SIZE_BASED_WRAP''', '''NO_WRAP''']
SCREAMING_SNAKE_CASE_					     =	['''BACKWARD_PRE''', '''BACKWARD_POST''', '''NO_PREFETCH''']
SCREAMING_SNAKE_CASE_					     =	['''FULL_STATE_DICT''', '''LOCAL_STATE_DICT''', '''SHARDED_STATE_DICT''']
SCREAMING_SNAKE_CASE_					     =	'''2.0.1'''
SCREAMING_SNAKE_CASE_					     =	['''pdsh''', '''standard''', '''openmpi''', '''mvapich''']
SCREAMING_SNAKE_CASE_					     =	['''default''', '''reduce-overhead''', '''max-autotune''']

SCREAMING_SNAKE_CASE_					     =	{'''>''': op.gt, '''>=''': op.ge, '''==''': op.eq, '''!=''': op.ne, '''<=''': op.le, '''<''': op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
SCREAMING_SNAKE_CASE_					     =	[
    '''nnodes''',
    '''nproc_per_node''',
    '''rdzv_backend''',
    '''rdzv_endpoint''',
    '''rdzv_id''',
    '''rdzv_conf''',
    '''standalone''',
    '''max_restarts''',
    '''monitor_interval''',
    '''start_method''',
    '''role''',
    '''module''',
    '''m''',
    '''no_python''',
    '''run_path''',
    '''log_dir''',
    '''r''',
    '''redirects''',
    '''t''',
    '''tee''',
    '''node_rank''',
    '''master_addr''',
    '''master_port''',
]

SCREAMING_SNAKE_CASE_					     =	['''DEEPSPEED''', '''MULTI_GPU''', '''FSDP''', '''MEGATRON_LM''']
SCREAMING_SNAKE_CASE_					     =	['''DEEPSPEED''', '''MULTI_XPU''', '''FSDP''']
