


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



     """simple docstring"""






     snake_case					=    '''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
     snake_case					=    Image.open(requests.get(UpperCamelCase_							,stream=UpperCamelCase_    ).raw    ).convert('''RGB'''    )

     return image
def       UpperCAmelCase__						(UpperCamelCase_    ):



     """simple docstring"""






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

     # fmt: on
     return rename_keys
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    dct.pop(UpperCamelCase_    )
     snake_case					=    val
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






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

          # next, set bias in the state dict
          snake_case					=    torch.cat((q_bias, torch.zeros_like(UpperCamelCase_							,requires_grad=UpperCamelCase_    ), v_bias)    )
          snake_case					=    qkv_bias
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_    ):



     """simple docstring"""






     snake_case					=    3_64 if '''coco''' in model_name else 2_24
     snake_case					=    BlipaVisionConfig(image_size=UpperCamelCase_    ).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:
          snake_case					=    OPTConfig.from_pretrained('''facebook/opt-2.7b'''							,eos_token_id=UpperCamelCase_    ).to_dict()
     elif "opt-6.7b" in model_name:
          snake_case					=    OPTConfig.from_pretrained('''facebook/opt-6.7b'''							,eos_token_id=UpperCamelCase_    ).to_dict()
     elif "t5-xl" in model_name:
          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:
          snake_case					=    TaConfig.from_pretrained('''google/flan-t5-xxl'''							,dense_act_fn='''gelu'''							,bos_token_id=1    ).to_dict()

     snake_case					=    BlipaConfig(vision_config=UpperCamelCase_							,text_config=UpperCamelCase_    )

     return config, image_size
@torch.no_grad()
def       UpperCAmelCase__						(UpperCamelCase_							,UpperCamelCase_=None							,UpperCamelCase_=False    ):



     """simple docstring"""






     snake_case					=    (
         AutoTokenizer.from_pretrained('''facebook/opt-2.7b'''    )
         if '''opt''' in model_name
         else AutoTokenizer.from_pretrained('''google/flan-t5-xl'''    )
     )
     snake_case					=    tokenizer('''\n'''							,add_special_tokens=UpperCamelCase_    ).input_ids[0]
     snake_case			,  snake_case					=    get_blipa_config(UpperCamelCase_							,eos_token_id=UpperCamelCase_    )

     snake_case					=    BlipaForConditionalGeneration(UpperCamelCase_    ).eval()

     snake_case					=    {
         '''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'''),
     }

     snake_case			,  snake_case					=    model_name_to_original[model_name]

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

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

     # some keys can be renamed efficiently
     for key, val in state_dict.copy().items():
          snake_case					=    state_dict.pop(UpperCamelCase_    )
          if key.startswith('''Qformer.bert'''    ):
               snake_case					=    key.replace('''Qformer.bert'''							,'''qformer'''    )
          if "attention.self" in key:
               snake_case					=    key.replace('''self'''							,'''attention'''    )
          if "opt_proj" in key:
               snake_case					=    key.replace('''opt_proj'''							,'''language_projection'''    )
          if "t5_proj" in key:
               snake_case					=    key.replace('''t5_proj'''							,'''language_projection'''    )
          if key.startswith('''opt'''    ):
               snake_case					=    key.replace('''opt'''							,'''language'''    )
          if key.startswith('''t5'''    ):
               snake_case					=    key.replace('''t5'''							,'''language'''    )
          snake_case					=    val

     # read in qv biases
     read_in_q_v_bias(UpperCamelCase_							,UpperCamelCase_    )

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

     snake_case					=    load_demo_image()
     snake_case					=    vis_processors['''eval'''](UpperCamelCase_    ).unsqueeze(0    ).to(UpperCamelCase_    )
     snake_case					=    tokenizer(['''\n''']							,return_tensors='''pt'''    ).input_ids.to(UpperCamelCase_    )

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

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

     original_model.to(UpperCamelCase_    )
     hf_model.to(UpperCamelCase_    )
     with torch.no_grad():
          if "opt" in model_name:
               snake_case					=    original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}    ).logits
               snake_case					=    hf_model(UpperCamelCase_							,UpperCamelCase_    ).logits
          else:
               snake_case					=    original_model(
                   {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}    ).logits
               snake_case					=    input_ids.masked_fill(input_ids == tokenizer.pad_token_id							,-1_00    )
               snake_case					=    hf_model(UpperCamelCase_							,UpperCamelCase_							,labels=UpperCamelCase_    ).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":
          snake_case					=    torch.tensor(
              [[-41.5850, -4.4440, -8.9922], [-47.4322, -5.9143, -1.7340]]							,device=UpperCamelCase_    )
          assert torch.allclose(logits[0, :3, :3]							,UpperCamelCase_							,atol=1e-4    )
     elif model_name == "blip2-flan-t5-xl-coco":
          snake_case					=    torch.tensor(
              [[-57.0109, -9.8967, -12.6280], [-68.6578, -12.7191, -10.5065]]							,device=UpperCamelCase_    )
     else:
          # cast to same type
          snake_case					=    logits.dtype
          assert torch.allclose(original_logits.to(UpperCamelCase_    )							,UpperCamelCase_							,atol=1e-2    )
     print('''Looks ok!'''    )

     print('''Generating a caption...'''    )
     snake_case					=    ''''''
     snake_case					=    tokenizer(UpperCamelCase_							,return_tensors='''pt'''    ).input_ids.to(UpperCamelCase_    )

     snake_case					=    original_model.generate({'''image''': original_pixel_values}    )
     snake_case					=    hf_model.generate(
         UpperCamelCase_							,UpperCamelCase_							,do_sample=UpperCamelCase_							,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:'''							,UpperCamelCase_    )
     snake_case					=    input_ids.shape[1]
     snake_case					=    processor.batch_decode(outputs[:, prompt_length:]							,skip_special_tokens=UpperCamelCase_    )
     snake_case					=    [text.strip() for text in output_text]
     print('''HF generation:'''							,UpperCamelCase_    )

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

     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      :			Dict      				=   argparse.ArgumentParser()
 _SCREAMING_SNAKE_CASE      :			str      				=   [
     "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      :			Optional[int]      				=   parser.parse_args()

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








from __future__ import annotations

import time

_SCREAMING_SNAKE_CASE      :			List[Any]      				=   list[tuple[int, int]]

_SCREAMING_SNAKE_CASE      :			Any      				=   [
    [0, 0, 0, 0, 0, 0, 0],
    [0, 1, 0, 0, 0, 0, 0],  # 0 are free path whereas 1's are obstacles
    [0, 0, 0, 0, 0, 0, 0],
    [0, 0, 1, 0, 0, 0, 0],
    [1, 0, 1, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 1, 0, 0],
]

_SCREAMING_SNAKE_CASE      :			Any      				=   [[-1, 0], [0, -1], [1, 0], [0, 1]]  # up, left, down, right

class   A__ :







    """simple docstring"""

    def __init__(	self      , __snake_case      , __snake_case      , __snake_case      , __snake_case      , __snake_case   ):
         snake_case					=    pos_x
         snake_case					=    pos_y
         snake_case					=    (pos_y, pos_x)
         snake_case					=    goal_x
         snake_case					=    goal_y
         snake_case					=    parent






class   A__ :







    """simple docstring"""

    def __init__(	self      , __snake_case      , __snake_case   ):
         snake_case					=    Node(start[1]      , start[0]      , goal[1]      , goal[0]      , __snake_case   )
         snake_case					=    Node(goal[1]      , goal[0]      , goal[1]      , goal[0]      , __snake_case   )

         snake_case					=    [self.start]
         snake_case					=    False
    def 		a_      (	self   ):
         while self.node_queue:
              snake_case					=    self.node_queue.pop(0   )

              if current_node.pos == self.target.pos:
                   snake_case					=    True
                   return self.retrace_path(__snake_case   )

              snake_case					=    self.get_successors(__snake_case   )

              for node in successors:
                   self.node_queue.append(__snake_case   )

         if not self.reached:
              return [self.start.pos]
         return None
    def 		a_      (	self      , __snake_case   ):
         snake_case					=    []
         for action in delta:
              snake_case					=    parent.pos_x + action[1]
              snake_case					=    parent.pos_y + action[0]
              if not (0 <= pos_x <= len(grid[0]   ) - 1 and 0 <= pos_y <= len(__snake_case   ) - 1):
                   continue

              if grid[pos_y][pos_x] != 0:
                   continue

              successors.append(
                  Node(__snake_case      , __snake_case      , self.target.pos_y      , self.target.pos_x      , __snake_case   )   )
         return successors




    def 		a_      (	self      , __snake_case   ):
         snake_case					=    node
         snake_case					=    []
         while current_node is not None:
              path.append((current_node.pos_y, current_node.pos_x)   )
              snake_case					=    current_node.parent
         path.reverse()
         return path






class   A__ :







    """simple docstring"""

    def __init__(	self      , __snake_case      , __snake_case   ):
         snake_case					=    BreadthFirstSearch(__snake_case      , __snake_case   )
         snake_case					=    BreadthFirstSearch(__snake_case      , __snake_case   )
         snake_case					=    False
    def 		a_      (	self   ):
         while self.fwd_bfs.node_queue or self.bwd_bfs.node_queue:
              snake_case					=    self.fwd_bfs.node_queue.pop(0   )
              snake_case					=    self.bwd_bfs.node_queue.pop(0   )

              if current_bwd_node.pos == current_fwd_node.pos:
                   snake_case					=    True
                   return self.retrace_bidirectional_path(
                       __snake_case      , __snake_case   )

              snake_case					=    current_bwd_node
              snake_case					=    current_fwd_node

              snake_case					=    {
                  self.fwd_bfs: self.fwd_bfs.get_successors(__snake_case   ),
                  self.bwd_bfs: self.bwd_bfs.get_successors(__snake_case   ),
              }

              for bfs in [self.fwd_bfs, self.bwd_bfs]:
                   for node in successors[bfs]:
                        bfs.node_queue.append(__snake_case   )

         if not self.reached:
              return [self.fwd_bfs.start.pos]
         return None




    def 		a_      (	self      , __snake_case      , __snake_case   ):
         snake_case					=    self.fwd_bfs.retrace_path(__snake_case   )
         snake_case					=    self.bwd_bfs.retrace_path(__snake_case   )
         bwd_path.pop()
         bwd_path.reverse()
         snake_case					=    fwd_path + bwd_path
         return path


if __name__ == "__main__":
 # all coordinates are given in format [y,x]
 import doctest

 doctest.testmod()
 _SCREAMING_SNAKE_CASE      :			Optional[Any]      				=   (0, 0)
 _SCREAMING_SNAKE_CASE      :			List[Any]      				=   (len(grid) - 1, len(grid[0]) - 1)
 for elem in grid:
  print(elem)

 _SCREAMING_SNAKE_CASE      :			List[Any]      				=   time.time()
 _SCREAMING_SNAKE_CASE      :			List[Any]      				=   BreadthFirstSearch(init, goal)
 _SCREAMING_SNAKE_CASE      :			List[str]      				=   bfs.search()
 _SCREAMING_SNAKE_CASE      :			int      				=   time.time() - start_bfs_time

 print("Unidirectional BFS computation time : ", bfs_time)

 _SCREAMING_SNAKE_CASE      :			Any      				=   time.time()
 _SCREAMING_SNAKE_CASE      :			Union[str, Any]      				=   BidirectionalBreadthFirstSearch(init, goal)
 _SCREAMING_SNAKE_CASE      :			Union[str, Any]      				=   bd_bfs.search()
 _SCREAMING_SNAKE_CASE      :			Tuple      				=   time.time() - start_bd_bfs_time

 print("Bidirectional BFS computation time : ", bd_bfs_time)




