


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







       """simple docstring"""







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

       return image






def 			lowercase				(_A					):







       """simple docstring"""







       _lowerCAmelCase    :							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 			lowercase				(_A   ,						_A   ,						_A					):







       """simple docstring"""







       _lowerCAmelCase    :							Optional[int]   			=       dct.pop(_A					)
       _lowerCAmelCase    :							Optional[int]   			=       val






def 			lowercase				(_A   ,						_A					):







       """simple docstring"""







       for i in range(config.vision_config.num_hidden_layers					):
              # read in original q and v biases
              _lowerCAmelCase    :							List[str]   			=       state_dict.pop(f'visual_encoder.blocks.{i}.attn.q_bias'					)
              _lowerCAmelCase    :							Dict   			=       state_dict.pop(f'visual_encoder.blocks.{i}.attn.v_bias'					)

              # next, set bias in the state dict
              _lowerCAmelCase    :							Dict   			=       torch.cat((q_bias, torch.zeros_like(_A   ,						requires_grad=_A					), v_bias)					)
              _lowerCAmelCase    :							Optional[int]   			=       qkv_bias






def 			lowercase				(_A   ,						_A					):







       """simple docstring"""







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

       _lowerCAmelCase    :							List[str]   			=       BlipaConfig(vision_config=_A   ,						text_config=_A					)

       return config, image_size






@torch.no_grad()
def 			lowercase				(_A   ,						_A=None   ,						_A=False					):







       """simple docstring"""







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

       _lowerCAmelCase    :							Any   			=       BlipaForConditionalGeneration(_A					).eval()

       _lowerCAmelCase    :							Any   			=       {
           '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    :							Optional[Any]   			=       model_name_to_original[model_name]

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

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

       # read in qv biases
       read_in_q_v_bias(_A   ,						_A					)

       _lowerCAmelCase    :							Any   			=       hf_model.load_state_dict(_A   ,						strict=_A					)
       assert len(_A					) == 0
       assert unexpected_keys == ["qformer.embeddings.position_ids"]

       _lowerCAmelCase    :							Tuple   			=       load_demo_image()
       _lowerCAmelCase    :							Optional[int]   			=       vis_processors['eval'](_A					).unsqueeze(0					).to(_A					)
       _lowerCAmelCase    :							Union[str, Any]   			=       tokenizer(['\n']   ,						return_tensors='pt'					).input_ids.to(_A					)

       # create processor
       _lowerCAmelCase    :							Tuple   			=       BlipImageProcessor(
           size={'height': image_size, 'width': image_size}   ,						image_mean=_A   ,						image_std=_A					)
       _lowerCAmelCase    :							List[str]   			=       BlipaProcessor(image_processor=_A   ,						tokenizer=_A					)
       _lowerCAmelCase    :							Union[str, Any]   			=       processor(images=_A   ,						return_tensors='pt'					).pixel_values.to(_A					)

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

       original_model.to(_A					)
       hf_model.to(_A					)
       with torch.no_grad():
              if "opt" in model_name:
                     _lowerCAmelCase    :							Optional[int]   			=       original_model({'image': original_pixel_values, 'text_input': ['']}					).logits
                     _lowerCAmelCase    :							Tuple   			=       hf_model(_A   ,						_A					).logits
              else:
                     _lowerCAmelCase    :							Optional[Any]   			=       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_0_0					)
                     _lowerCAmelCase    :							Optional[int]   			=       hf_model(_A   ,						_A   ,						labels=_A					).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    :							int   			=       torch.tensor(
                  [[-41.5_850, -4.4_440, -8.9_922], [-47.4_322, -5.9_143, -1.7_340]]   ,						device=_A					)
              assert torch.allclose(logits[0, :3, :3]   ,						_A   ,						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=_A					)
       else:
              # cast to same type
              _lowerCAmelCase    :							List[str]   			=       logits.dtype
              assert torch.allclose(original_logits.to(_A					)   ,						_A   ,						atol=1E-2					)
       print('Looks ok!'					)

       print('Generating a caption...'					)
       _lowerCAmelCase    :							List[str]   			=       ''
       _lowerCAmelCase    :							List[str]   			=       tokenizer(_A   ,						return_tensors='pt'					).input_ids.to(_A					)

       _lowerCAmelCase    :							Optional[Any]   			=       original_model.generate({'image': original_pixel_values}					)
       _lowerCAmelCase    :							Dict   			=       hf_model.generate(
           _A   ,						_A   ,						do_sample=_A   ,						num_beams=5   ,						max_length=3_0   ,						min_length=1   ,						top_p=0.9   ,						repetition_penalty=1.0   ,						length_penalty=1.0   ,						temperature=1   ,						)
       print('Original generation:'   ,						_A					)
       _lowerCAmelCase    :							int   			=       input_ids.shape[1]
       _lowerCAmelCase    :							str   			=       processor.batch_decode(outputs[:, prompt_length:]   ,						skip_special_tokens=_A					)
       _lowerCAmelCase    :							int   			=       [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'nielsr/{model_name}'					)
              hf_model.push_to_hub(f'nielsr/{model_name}'					)


if __name__ == "__main__":
       lowerCAmelCase    :     Tuple					=						argparse.ArgumentParser()
       lowerCAmelCase    :     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""",
       )

       lowerCAmelCase    :     Any					=						parser.parse_args()

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






'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCAmelCase    :     Any					=						{
    """configuration_x_clip""": [
        """XCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """XCLIPConfig""",
        """XCLIPTextConfig""",
        """XCLIPVisionConfig""",
    ],
    """processing_x_clip""": ["""XCLIPProcessor"""],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCAmelCase    :     Optional[Any]					=						[
           """XCLIP_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """XCLIPModel""",
           """XCLIPPreTrainedModel""",
           """XCLIPTextModel""",
           """XCLIPVisionModel""",
       ]

if TYPE_CHECKING:
       from .configuration_x_clip import (
           XCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
           XCLIPConfig,
           XCLIPTextConfig,
           XCLIPVisionConfig,
       )
       from .processing_x_clip import XCLIPProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_x_clip import (
                  XCLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
                  XCLIPModel,
                  XCLIPPreTrainedModel,
                  XCLIPTextModel,
                  XCLIPVisionModel,
              )

else:
       import sys

       lowerCAmelCase    :     Dict					=						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

