


'''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        __lowercase						(      )				->					str:
     """simple docstring"""




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

     return image




def        __lowercase						(      __SCREAMING_SNAKE_CASE		)				->					Optional[Any]:
     """simple docstring"""




     __a       =   []
     # 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        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)				->					int:
     """simple docstring"""




     __a       =   dct.pop(__SCREAMING_SNAKE_CASE		)
     __a       =   val




def        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)				->					Tuple:
     """simple docstring"""




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

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




def        __lowercase						(      __SCREAMING_SNAKE_CASE		)				->					List[Any]:
     """simple docstring"""




     __a       =   364 if """coco""" in model_name else 224
     __a       =   InstructBlipVisionConfig(image_size=__SCREAMING_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:
          __a       =   TaConfig.from_pretrained("""google/flan-t5-xl"""				,  dense_act_fn="""gelu"""				,  bos_token_id=1		).to_dict()
     elif "t5-xxl" in model_name:
          __a       =   TaConfig.from_pretrained("""google/flan-t5-xxl"""				,  dense_act_fn="""gelu"""				,  bos_token_id=1		).to_dict()
     elif "vicuna-7b" in model_name:
          __a       =   LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf"""				,  vocab_size=3_2001		).to_dict()
     elif "vicuna-13b" in model_name:
          __a       =   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
     __a       =   InstructBlipQFormerConfig(vocab_size=3_0523		).to_dict()
     __a       =   InstructBlipConfig(vision_config=__SCREAMING_SNAKE_CASE				,  text_config=__SCREAMING_SNAKE_CASE				,  qformer_config=__SCREAMING_SNAKE_CASE		)

     return config, image_size




@torch.no_grad()
def        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE=None				,  __SCREAMING_SNAKE_CASE=False		)				->					Tuple:
     """simple docstring"""




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

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

     __a ,						__a       =   get_blipa_config(__SCREAMING_SNAKE_CASE		)
     __a       =   InstructBlipForConditionalGeneration(__SCREAMING_SNAKE_CASE		).eval()

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

     __a ,						__a       =   model_name_to_original[model_name]

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

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

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

     # read in qv biases
     read_in_q_v_bias(__SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)

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

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

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

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

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

     print("""Generating with original model..."""		)
     __a       =   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..."""		)
     __a       =   hf_model.generate(
         **__SCREAMING_SNAKE_CASE				,  do_sample=__SCREAMING_SNAKE_CASE				,  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?
          __a       =   2
     print("""Original generation:"""				,  __SCREAMING_SNAKE_CASE		)
     __a       =   processor.batch_decode(__SCREAMING_SNAKE_CASE				,  skip_special_tokens=__SCREAMING_SNAKE_CASE		)
     __a       =   [text.strip() for text in output_text]
     print("""HF generation:"""				,  __SCREAMING_SNAKE_CASE		)

     if pytorch_dump_folder_path is not None:
          processor.save_pretrained(__SCREAMING_SNAKE_CASE		)
          hf_model.save_pretrained(__SCREAMING_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__":
 SCREAMING_SNAKE_CASE_      =			argparse.ArgumentParser()
 SCREAMING_SNAKE_CASE_      =			[
     '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_      =			parser.parse_args()

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










'''simple docstring'''







import math
import random
from typing import Any

from .hill_climbing import SearchProblem




def        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE = True				,  __SCREAMING_SNAKE_CASE = math.inf				,  __SCREAMING_SNAKE_CASE = -math.inf				,  __SCREAMING_SNAKE_CASE = math.inf				,  __SCREAMING_SNAKE_CASE = -math.inf				,  __SCREAMING_SNAKE_CASE = False				,  __SCREAMING_SNAKE_CASE = 100				,  __SCREAMING_SNAKE_CASE = 0.01				,  __SCREAMING_SNAKE_CASE = 1				,  )				->					Any:
     """simple docstring"""




     __a       =   False
     __a       =   search_prob
     __a       =   start_temperate
     __a       =   []
     __a       =   0
     __a       =   None

     while not search_end:
          __a       =   current_state.score()
          if best_state is None or current_score > best_state.score():
               __a       =   current_state
          scores.append(__SCREAMING_SNAKE_CASE		)
          iterations += 1
          __a       =   None
          __a       =   current_state.get_neighbors()
          while (
              next_state is None and neighbors
          ):  # till we do not find a neighbor that we can move to
               __a       =   random.randint(0				,  len(__SCREAMING_SNAKE_CASE		) - 1		)  # picking a random neighbor
               __a       =   neighbors.pop(__SCREAMING_SNAKE_CASE		)
               __a       =   picked_neighbor.score() - current_score

               if (
                   picked_neighbor.x > max_x
                   or picked_neighbor.x < min_x
                   or picked_neighbor.y > max_y
                   or picked_neighbor.y < min_y
               ):
                    continue  # neighbor outside our bounds

               if not find_max:
                    __a       =   change * -1  # in case we are finding minimum
               if change > 0:  # improves the solution
                    __a       =   picked_neighbor
               else:
                    __a       =   (math.e) ** (
                        change / current_temp
                    )  # probability generation function
                    if random.random() < probability:  # random number within probability
                         __a       =   picked_neighbor
          __a       =   current_temp - (current_temp * rate_of_decrease)

          if current_temp < threshold_temp or next_state is None:
               # temperature below threshold, or could not find a suitable neighbor
               __a       =   True
          else:
               __a       =   next_state

     if visualization:
          from matplotlib import pyplot as plt

          plt.plot(range(__SCREAMING_SNAKE_CASE		)				,  __SCREAMING_SNAKE_CASE		)
          plt.xlabel("""Iterations"""		)
          plt.ylabel("""Function values"""		)
          plt.show()
     return best_state


if __name__ == "__main__":




 def        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)				->					int:
      """simple docstring"""




      return (x**2) + (y**2)

 # starting the problem with initial coordinates (12, 47)
 SCREAMING_SNAKE_CASE_      =			SearchProblem(x=12, y=47, step_size=1, function_to_optimize=test_fa)
 SCREAMING_SNAKE_CASE_      =			simulated_annealing(
     prob, find_max=False, max_x=1_00, min_x=5, max_y=50, min_y=-5, visualization=True
 )
 print(
     'The minimum score for f(x, y) = x^2 + y^2 with the domain 100 > x > 5 '
     f"""and 50 > y > - 5 found via hill climbing: {local_min.score()}"""
 )

 # starting the problem with initial coordinates (12, 47)
 SCREAMING_SNAKE_CASE_      =			SearchProblem(x=12, y=47, step_size=1, function_to_optimize=test_fa)
 SCREAMING_SNAKE_CASE_      =			simulated_annealing(
     prob, find_max=True, max_x=1_00, min_x=5, max_y=50, min_y=-5, visualization=True
 )
 print(
     'The maximum score for f(x, y) = x^2 + y^2 with the domain 100 > x > 5 '
     f"""and 50 > y > - 5 found via hill climbing: {local_min.score()}"""
 )




 def        __lowercase						(      __SCREAMING_SNAKE_CASE				,  __SCREAMING_SNAKE_CASE		)				->					int:
      """simple docstring"""




      return (3 * x**2) - (6 * y)

 SCREAMING_SNAKE_CASE_      =			SearchProblem(x=3, y=4, step_size=1, function_to_optimize=test_fa)
 SCREAMING_SNAKE_CASE_      =			simulated_annealing(prob, find_max=False, visualization=True)
 print(
     'The minimum score for f(x, y) = 3*x^2 - 6*y found via hill climbing: '
     f"""{local_min.score()}"""
 )

 SCREAMING_SNAKE_CASE_      =			SearchProblem(x=3, y=4, step_size=1, function_to_optimize=test_fa)
 SCREAMING_SNAKE_CASE_      =			simulated_annealing(prob, find_max=True, visualization=True)



 print(
     'The maximum score for f(x, y) = 3*x^2 - 6*y found via hill climbing: '
     f"""{local_min.score()}"""
 )






