'''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      UpperCamelCase_						(					)					->     List[Any]:






 '''simple docstring'''

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

 return image







def      UpperCamelCase_						(					snake_case_       :       List[str]      )					->     str:






 '''simple docstring'''

 __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.embeddings.layernorm.weight""")      )
 rename_keys.append(("""Qformer.bert.embeddings.LayerNorm.bias""", """qformer.embeddings.layernorm.bias""")      )

 # fmt: on
 return rename_keys







def      UpperCamelCase_						(					snake_case_       :       Dict						,		snake_case_       :       List[Any]						,		snake_case_       :       Tuple      )					->     Union[str, Any]:






 '''simple docstring'''

 __lowerCAmelCase											=						dct.pop(snake_case_      )
 __lowerCAmelCase											=						val







def      UpperCamelCase_						(					snake_case_       :       int						,		snake_case_       :       Tuple      )					->     Tuple:






 '''simple docstring'''

 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(snake_case_						,		requires_grad=snake_case_      ), v_bias)      )
  __lowerCAmelCase											=						qkv_bias







def      UpperCamelCase_						(					snake_case_       :       str      )					->     Optional[Any]:






 '''simple docstring'''

 __lowerCAmelCase											=						3_64 if """coco""" in model_name else 2_24
 __lowerCAmelCase											=						InstructBlipVisionConfig(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 "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()
 elif "vicuna-7b" in model_name:
  __lowerCAmelCase											=						LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf"""						,		vocab_size=3_20_01      ).to_dict()
 elif "vicuna-13b" in model_name:
  __lowerCAmelCase											=						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
 __lowerCAmelCase											=						InstructBlipQFormerConfig(vocab_size=3_05_23      ).to_dict()
 __lowerCAmelCase											=						InstructBlipConfig(vision_config=snake_case_						,		text_config=snake_case_						,		qformer_config=snake_case_      )

 return config, image_size







@torch.no_grad()
def      UpperCamelCase_						(					snake_case_       :       Optional[int]						,		snake_case_       :       Dict=None						,		snake_case_       :       str=False      )					->     Optional[Any]:






 '''simple docstring'''

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

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

 __lowerCAmelCase							,       __lowerCAmelCase											=						get_blipa_config(snake_case_      )
 __lowerCAmelCase											=						InstructBlipForConditionalGeneration(snake_case_      ).eval()

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

 __lowerCAmelCase							,       __lowerCAmelCase											=						model_name_to_original[model_name]

 # load original model
 print("""Loading original model..."""      )
 __lowerCAmelCase											=						"""cuda:1""" if torch.cuda.is_available() else """cpu"""
 __lowerCAmelCase											=						"""cuda:2""" if torch.cuda.is_available() else """cpu"""
 __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase											=						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
 __lowerCAmelCase											=						original_model.state_dict()
 __lowerCAmelCase											=						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():
  __lowerCAmelCase											=						state_dict.pop(snake_case_      )
  if key.startswith("""Qformer.bert"""      ):
   __lowerCAmelCase											=						key.replace("""Qformer.bert"""						,		"""qformer"""      )
  if "attention.self" in key:
   __lowerCAmelCase											=						key.replace("""self"""						,		"""attention"""      )
  if "llm_proj" in key:
   __lowerCAmelCase											=						key.replace("""llm_proj"""						,		"""language_projection"""      )
  if "t5_proj" in key:
   __lowerCAmelCase											=						key.replace("""t5_proj"""						,		"""language_projection"""      )
  if key.startswith("""llm_model"""      ):
   __lowerCAmelCase											=						key.replace("""llm_model"""						,		"""language_model"""      )
  if key.startswith("""t5"""      ):
   __lowerCAmelCase											=						key.replace("""t5"""						,		"""language"""      )
  __lowerCAmelCase											=						val

 # read in qv biases
 read_in_q_v_bias(snake_case_						,		snake_case_      )

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

 __lowerCAmelCase											=						load_demo_image()
 __lowerCAmelCase											=						"""What is unusual about this image?"""

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

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

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

 print("""Generating with original model..."""      )
 __lowerCAmelCase											=						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..."""      )
 __lowerCAmelCase											=						hf_model.generate(
     **snake_case_						,		do_sample=snake_case_						,		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?
  __lowerCAmelCase											=						2
 print("""Original generation:"""						,		snake_case_      )
 __lowerCAmelCase											=						processor.batch_decode(snake_case_						,		skip_special_tokens=snake_case_      )
 __lowerCAmelCase											=						[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"""Salesforce/{model_name}"""      )
  hf_model.push_to_hub(f"""Salesforce/{model_name}"""      )


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

 _A			:      Dict					      =				parser.parse_args()

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







'''simple docstring'''

from math import sqrt

import numpy as np
from sympy import symbols

# Coefficient
# Speed of light (m/s)
_A			:      List[Any]					      =				299792458

# Symbols
_A      ,     _A      ,     _A      ,     _A			:      Union[str, Any]					      =				symbols('''ct x y z''')







def      UpperCamelCase_						(					snake_case_       :       float      )					->     float:






 '''simple docstring'''

 if velocity > c:
  raise ValueError("""Speed must not exceed light speed 299,792,458 [m/s]!"""      )
 elif velocity < 1:
  # Usually the speed should be much higher than 1 (c order of magnitude)
  raise ValueError("""Speed must be greater than or equal to 1!"""      )

 return velocity / c







def      UpperCamelCase_						(					snake_case_       :       float      )					->     float:






 '''simple docstring'''

 return 1 / sqrt(1 - beta(snake_case_      ) ** 2      )







def      UpperCamelCase_						(					snake_case_       :       float      )					->     np.ndarray:






 '''simple docstring'''

 return np.array(
     [
         [gamma(snake_case_      ), -gamma(snake_case_      ) * beta(snake_case_      ), 0, 0],
         [-gamma(snake_case_      ) * beta(snake_case_      ), gamma(snake_case_      ), 0, 0],
         [0, 0, 1, 0],
         [0, 0, 0, 1],
     ]      )







def      UpperCamelCase_						(					snake_case_       :       float						,		snake_case_       :       np.ndarray | None = None      )					->     np.ndarray:






 '''simple docstring'''

 if event is None:
  __lowerCAmelCase											=						np.array([ct, x, y, z]      )  # Symbolic four vector
 else:
  event[0] *= c  # x0 is ct (speed of light * time)

 return transformation_matrix(snake_case_      ) @ event


if __name__ == "__main__":
 import doctest

 doctest.testmod()

 # Example of symbolic vector:
 _A			:      str					      =				transform(29979245)
 print('''Example of four vector: ''')
 print(f'ct\' = {four_vector[0]}')
 print(f'x\' = {four_vector[1]}')
 print(f'y\' = {four_vector[2]}')
 print(f'z\' = {four_vector[3]}')

 # Substitute symbols with numerical values
 _A			:      int					      =				{ct: c, x: 1, y: 1, z: 1}
 _A			:      Any					      =				[four_vector[i].subs(sub_dict) for i in range(4)]

 print(f'\n{numerical_vector}')






