

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_     (		):





     '''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        UpperCAmelCase_     (		_A       ):





     '''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        UpperCAmelCase_     (		_A   ,	_A   ,	_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					dct.pop(_A       )
     SCREAMING_SNAKE_CASE__					=					val
def        UpperCAmelCase_     (		_A   ,	_A       ):





     '''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        UpperCAmelCase_     (		_A       ):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					3_64 if '''coco''' in model_name else 2_24
     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_20_01       ).to_dict()
     elif "vicuna-13b" in model_name:
          SCREAMING_SNAKE_CASE__					=					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__					=					InstructBlipQFormerConfig(vocab_size=3_05_23       ).to_dict()
     SCREAMING_SNAKE_CASE__					=					InstructBlipConfig(vision_config=_A   ,	text_config=_A   ,	qformer_config=_A       )

     return config, image_size
@torch.no_grad()
def        UpperCAmelCase_     (		_A   ,	_A=None   ,	_A=False       ):





     '''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   ,	-1_00       )
               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=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__					=					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__":
   _SCREAMING_SNAKE_CASE  :						List[Any]         =    argparse.ArgumentParser()
   _SCREAMING_SNAKE_CASE  :						Dict         =    [
       '''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''',
   )

   _SCREAMING_SNAKE_CASE  :						List[str]         =    parser.parse_args()

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





import math






class  UpperCAmelCase__	:






  """simple docstring"""


  def    lowercase_	(							self		:	int			,     __lowerCamelCase		:	list[list[float]]			,     __lowerCamelCase		:	list[int]      )	->							int:
       SCREAMING_SNAKE_CASE__					=					0.0
       SCREAMING_SNAKE_CASE__					=					0.0
       for i in range(len(__lowerCamelCase      )      ):
            da += math.pow((sample[i] - weights[0][i])			,     2      )
            da += math.pow((sample[i] - weights[1][i])			,     2      )
            return 0 if da > da else 1
       return 0







  def    lowercase_	(							self		:	Optional[int]			,     __lowerCamelCase		:	list[list[int | float]]			,     __lowerCamelCase		:	list[int]			,     __lowerCamelCase		:	int			,     __lowerCamelCase		:	float      )	->							list[list[int | float]]:
       for i in range(len(__lowerCamelCase      )      ):
            weights[j][i] += alpha * (sample[i] - weights[j][i])
       return weights



def        UpperCAmelCase_     (		):





     '''simple docstring'''




     SCREAMING_SNAKE_CASE__					=					[[1, 1, 0, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 0, 1, 1]]

     # weight initialization ( n, C )
     SCREAMING_SNAKE_CASE__					=					[[0.2, 0.6, 0.5, 0.9], [0.8, 0.4, 0.7, 0.3]]

     # training
     SCREAMING_SNAKE_CASE__					=					SelfOrganizingMap()
     SCREAMING_SNAKE_CASE__					=					3
     SCREAMING_SNAKE_CASE__					=					0.5

     for _ in range(_A       ):
          for j in range(len(_A       )       ):
               # training sample
               SCREAMING_SNAKE_CASE__					=					training_samples[j]

               # Compute the winning vector
               SCREAMING_SNAKE_CASE__					=					self_organizing_map.get_winner(_A   ,	_A       )

               # Update the winning vector
               SCREAMING_SNAKE_CASE__					=					self_organizing_map.update(_A   ,	_A   ,	_A   ,	_A       )

    # classify test sample
     SCREAMING_SNAKE_CASE__					=					[0, 0, 0, 1]
     SCREAMING_SNAKE_CASE__					=					self_organizing_map.get_winner(_A   ,	_A       )

     # results
     print(F'''Clusters that the test sample belongs to : {winner}'''       )
     print(F'''Weights that have been trained : {weights}'''       )


# running the main() function
if __name__ == "__main__":
       main()


