


'''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      UpperCAmelCase__   (      )	->						Optional[Any]:
     __lowerCamelCase  :      Union[str, Any]														=						'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
     __lowerCamelCase  :      Any														=						Image.open(requests.get(UpperCAmelCase_ ,       stream=UpperCAmelCase_				).raw				).convert('RGB'				)

     return image



def      UpperCAmelCase__   (      UpperCAmelCase_       :							Tuple				)	->						Tuple:
     __lowerCamelCase  :      Optional[Any]														=						[]
     # 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_       :							Optional[int] ,       UpperCAmelCase_       :							Dict ,       UpperCAmelCase_       :							List[str]				)	->						int:
     __lowerCamelCase  :      List[Any]														=						dct.pop(UpperCAmelCase_				)
     __lowerCamelCase  :      Optional[int]														=						val



def      UpperCAmelCase__   (      UpperCAmelCase_       :							Union[str, Any] ,       UpperCAmelCase_       :							List[Any]				)	->						List[Any]:
     for i in range(config.vision_config.num_hidden_layers				):
          # read in original q and v biases
          __lowerCamelCase  :      Tuple														=						state_dict.pop(F'visual_encoder.blocks.{i}.attn.q_bias'				)
          __lowerCamelCase  :      int														=						state_dict.pop(F'visual_encoder.blocks.{i}.attn.v_bias'				)

          # next, set bias in the state dict
          __lowerCamelCase  :      List[Any]														=						torch.cat((q_bias, torch.zeros_like(UpperCAmelCase_ ,       requires_grad=UpperCAmelCase_				), v_bias)				)
          __lowerCamelCase  :      Any														=						qkv_bias



def      UpperCAmelCase__   (      UpperCAmelCase_       :							int ,       UpperCAmelCase_       :							str				)	->						Optional[int]:
     __lowerCamelCase  :      int														=						3_64 if 'coco' in model_name else 2_24
     __lowerCamelCase  :      List[Any]														=						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:
          __lowerCamelCase  :      Union[str, Any]														=						OPTConfig.from_pretrained('facebook/opt-2.7b' ,       eos_token_id=UpperCAmelCase_				).to_dict()
     elif "opt-6.7b" in model_name:
          __lowerCamelCase  :      List[str]														=						OPTConfig.from_pretrained('facebook/opt-6.7b' ,       eos_token_id=UpperCAmelCase_				).to_dict()
     elif "t5-xl" in model_name:
          __lowerCamelCase  :      Union[str, Any]														=						TaConfig.from_pretrained('google/flan-t5-xl' ,       dense_act_fn='gelu' ,       bos_token_id=1				).to_dict()
     elif "t5-xxl" in model_name:
          __lowerCamelCase  :      List[str]														=						TaConfig.from_pretrained('google/flan-t5-xxl' ,       dense_act_fn='gelu' ,       bos_token_id=1				).to_dict()

     __lowerCamelCase  :      Tuple														=						BlipaConfig(vision_config=UpperCAmelCase_ ,       text_config=UpperCAmelCase_				)

     return config, image_size



@torch.no_grad()
def      UpperCAmelCase__   (      UpperCAmelCase_       :							List[Any] ,       UpperCAmelCase_       :							Optional[int]=None ,       UpperCAmelCase_       :							Union[str, Any]=False				)	->						Dict:
     __lowerCamelCase  :      Tuple														=						(
         AutoTokenizer.from_pretrained('facebook/opt-2.7b'				)
         if 'opt' in model_name
         else AutoTokenizer.from_pretrained('google/flan-t5-xl'				)
     )
     __lowerCamelCase  :      Union[str, Any]														=						tokenizer('\n' ,       add_special_tokens=UpperCAmelCase_				).input_ids[0]
     __lowerCamelCase							,			__lowerCamelCase  :      str														=						get_blipa_config(UpperCAmelCase_ ,       eos_token_id=UpperCAmelCase_				)

     __lowerCamelCase  :      Tuple														=						BlipaForConditionalGeneration(UpperCAmelCase_				).eval()

     __lowerCamelCase  :      int														=						{
         '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  :      int														=						model_name_to_original[model_name]

     # load original model
     print('Loading original model...'				)
     __lowerCamelCase  :      Any														=						'cuda' if torch.cuda.is_available() else 'cpu'
     __lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase  :      str														=						load_model_and_preprocess(
         name=UpperCAmelCase_ ,       model_type=UpperCAmelCase_ ,       is_eval=UpperCAmelCase_ ,       device=UpperCAmelCase_				)
     original_model.eval()
     print('Done!'				)

     # update state dict keys
     __lowerCamelCase  :      Dict														=						original_model.state_dict()
     __lowerCamelCase  :      List[Any]														=						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():
          __lowerCamelCase  :      Optional[Any]														=						state_dict.pop(UpperCAmelCase_				)
          if key.startswith('Qformer.bert'				):
               __lowerCamelCase  :      Dict														=						key.replace('Qformer.bert' ,       'qformer'				)
          if "attention.self" in key:
               __lowerCamelCase  :      int														=						key.replace('self' ,       'attention'				)
          if "opt_proj" in key:
               __lowerCamelCase  :      Any														=						key.replace('opt_proj' ,       'language_projection'				)
          if "t5_proj" in key:
               __lowerCamelCase  :      Optional[Any]														=						key.replace('t5_proj' ,       'language_projection'				)
          if key.startswith('opt'				):
               __lowerCamelCase  :      Dict														=						key.replace('opt' ,       'language'				)
          if key.startswith('t5'				):
               __lowerCamelCase  :      Optional[int]														=						key.replace('t5' ,       'language'				)
          __lowerCamelCase  :      str														=						val

     # read in qv biases
     read_in_q_v_bias(UpperCAmelCase_ ,       UpperCAmelCase_				)

     __lowerCamelCase							,			__lowerCamelCase  :      List[str]														=						hf_model.load_state_dict(UpperCAmelCase_ ,       strict=UpperCAmelCase_				)
     assert len(UpperCAmelCase_				) == 0
     assert unexpected_keys == ["qformer.embeddings.position_ids"]

     __lowerCamelCase  :      Optional[Any]														=						load_demo_image()
     __lowerCamelCase  :      Any														=						vis_processors['eval'](UpperCAmelCase_				).unsqueeze(0				).to(UpperCAmelCase_				)
     __lowerCamelCase  :      Dict														=						tokenizer(['\n'] ,       return_tensors='pt'				).input_ids.to(UpperCAmelCase_				)

     # create processor
     __lowerCamelCase  :      Tuple														=						BlipImageProcessor(
         size={'height': image_size, 'width': image_size} ,       image_mean=UpperCAmelCase_ ,       image_std=UpperCAmelCase_				)
     __lowerCamelCase  :      List[str]														=						BlipaProcessor(image_processor=UpperCAmelCase_ ,       tokenizer=UpperCAmelCase_				)
     __lowerCamelCase  :      str														=						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:
               __lowerCamelCase  :      Optional[Any]														=						original_model({'image': original_pixel_values, 'text_input': ['']}				).logits
               __lowerCamelCase  :      List[str]														=						hf_model(UpperCAmelCase_ ,       UpperCAmelCase_				).logits
          else:
               __lowerCamelCase  :      Tuple														=						original_model(
                   {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}				).logits
               __lowerCamelCase  :      List[str]														=						input_ids.masked_fill(input_ids == tokenizer.pad_token_id ,       -1_00				)
               __lowerCamelCase  :      Dict														=						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":
          __lowerCamelCase  :      List[Any]														=						torch.tensor(
              [[-41.5_850, -4.4_440, -8.9_922], [-47.4_322, -5.9_143, -1.7_340]] ,       device=UpperCAmelCase_				)
          assert torch.allclose(logits[0, :3, :3] ,       UpperCAmelCase_ ,       atol=1e-4				)
     elif model_name == "blip2-flan-t5-xl-coco":
          __lowerCamelCase  :      Tuple														=						torch.tensor(
              [[-57.0_109, -9.8_967, -12.6_280], [-68.6_578, -12.7_191, -10.5_065]] ,       device=UpperCAmelCase_				)
     else:
          # cast to same type
          __lowerCamelCase  :      Optional[int]														=						logits.dtype
          assert torch.allclose(original_logits.to(UpperCAmelCase_				) ,       UpperCAmelCase_ ,       atol=1e-2				)
     print('Looks ok!'				)

     print('Generating a caption...'				)
     __lowerCamelCase  :      str														=						''
     __lowerCamelCase  :      str														=						tokenizer(UpperCAmelCase_ ,       return_tensors='pt'				).input_ids.to(UpperCAmelCase_				)

     __lowerCamelCase  :      Union[str, Any]														=						original_model.generate({'image': original_pixel_values}				)
     __lowerCamelCase  :      int														=						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_				)
     __lowerCamelCase  :      Any														=						input_ids.shape[1]
     __lowerCamelCase  :      Optional[int]														=						processor.batch_decode(outputs[:, prompt_length:] ,       skip_special_tokens=UpperCAmelCase_				)
     __lowerCamelCase  :      Dict														=						[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__":
      A__						:   int				=    argparse.ArgumentParser()
      A__						:   Optional[int]				=    [
          """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""",
      )

      A__						:   Any				=    parser.parse_args()

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














import importlib
import os

import fsspec
import pytest
from fsspec import register_implementation
from fsspec.registry import _registry as _fsspec_registry

from datasets.filesystems import COMPRESSION_FILESYSTEMS, HfFileSystem, extract_path_from_uri, is_remote_filesystem

from .utils import require_lza, require_zstandard






def 	UpperCamelCase	(							snake_case__  : int	):




     '''simple docstring'''






     assert "mock" in _fsspec_registry
     assert "bz2" in _fsspec_registry






def 	UpperCamelCase	(							):




     '''simple docstring'''






     assert "mock" not in _fsspec_registry
     assert "bz2" in _fsspec_registry






def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :Tuple 							=    """mock-s3-bucket"""
     __snake_case      :Optional[Any] 							=    f'''s3://{mock_bucket}'''
     __snake_case      :Optional[int] 							=    extract_path_from_uri(snake_case__	)
     assert dataset_path.startswith("""s3://"""	) is False

     __snake_case      :Union[str, Any] 							=    """./local/path"""
     __snake_case      :List[str] 							=    extract_path_from_uri(snake_case__	)
     assert dataset_path == new_dataset_path






def 	UpperCamelCase	(							snake_case__  : str	):




     '''simple docstring'''






     __snake_case      :List[str] 							=    is_remote_filesystem(snake_case__	)
     assert is_remote is True

     __snake_case      :Any 							=    fsspec.filesystem("""file"""	)

     __snake_case      :Optional[Any] 							=    is_remote_filesystem(snake_case__	)
     assert is_remote is False






@pytest.mark.parametrize("""compression_fs_class"""						,snake_case__	)
def 	UpperCamelCase	(							snake_case__  : Optional[Any]						,snake_case__  : List[str]						,snake_case__  : int						,snake_case__  : List[Any]						,snake_case__  : Tuple						,snake_case__  : str						,snake_case__  : Optional[Any]	):




     '''simple docstring'''






     __snake_case      :Optional[int] 							=    {"""gzip""": gz_file, """xz""": xz_file, """zstd""": zstd_file, """bz2""": bza_file, """lz4""": lza_file}
     __snake_case      :Optional[Any] 							=    input_paths[compression_fs_class.protocol]
     if input_path is None:
          __snake_case      :Optional[int] 							=    f'''for \'{compression_fs_class.protocol}\' compression protocol, '''
          if compression_fs_class.protocol == "lz4":
               reason += require_lza.kwargs["reason"]
          elif compression_fs_class.protocol == "zstd":
               reason += require_zstandard.kwargs["reason"]
          pytest.skip(snake_case__	)
     __snake_case      :Optional[Any] 							=    fsspec.filesystem(compression_fs_class.protocol						,fo=snake_case__	)
     assert isinstance(snake_case__						,snake_case__	)
     __snake_case      :Dict 							=    os.path.basename(snake_case__	)
     __snake_case      :Dict 							=    expected_filename[: expected_filename.rindex("""."""	)]
     assert fs.glob("""*"""	) == [expected_filename]
     with fs.open(snake_case__						,"""r"""						,encoding="""utf-8"""	) as f, open(snake_case__						,encoding="""utf-8"""	) as expected_file:
          assert f.read() == expected_file.read()






@pytest.mark.parametrize("""protocol"""						,["""zip""", """gzip"""]	)
def 	UpperCamelCase	(							snake_case__  : List[Any]						,snake_case__  : List[Any]						,snake_case__  : Dict	):




     '''simple docstring'''






     __snake_case      :Optional[int] 							=    {"""zip""": zip_jsonl_path, """gzip""": jsonl_gz_path}
     __snake_case      :List[str] 							=    compressed_file_paths[protocol]
     __snake_case      :Optional[int] 							=    """dataset.jsonl"""
     __snake_case      :Optional[int] 							=    f'''{protocol}://{member_file_path}::{compressed_file_path}'''
     __snake_case							,			*__snake_case      :Optional[Any] 							=    fsspec.get_fs_token_paths(snake_case__	)
     assert fs.isfile(snake_case__	)
     assert not fs.isfile("""non_existing_""" + member_file_path	)






@pytest.mark.integration
def 	UpperCamelCase	(							snake_case__  : Optional[int]						,snake_case__  : Dict						,snake_case__  : List[str]						,snake_case__  : Dict	):




     '''simple docstring'''






     __snake_case      :List[Any] 							=    hf_api.dataset_info(snake_case__						,token=snake_case__	)
     __snake_case      :List[str] 							=    HfFileSystem(repo_info=snake_case__						,token=snake_case__	)
     assert sorted(hffs.glob("""*"""	)	) == [".gitattributes", "data"]
     assert hffs.isdir("""data"""	)
     assert hffs.isfile(""".gitattributes"""	) and hffs.isfile("""data/text_data.txt"""	)
     with open(snake_case__	) as f:
          assert hffs.open("""data/text_data.txt"""						,"""r"""	).read() == f.read()






def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :str 							=    """bz2"""

     # Import module
     import datasets.filesystems

     # Overwrite protocol and reload
     register_implementation(snake_case__						,snake_case__						,clobber=snake_case__	)
     with pytest.warns(snake_case__	) as warning_info:
          importlib.reload(datasets.filesystems	)

     assert len(snake_case__	) == 1
     assert (
         str(warning_info[0].message	)
         == f'''A filesystem protocol was already set for {protocol} and will be overwritten.'''
     )


