


'''simple docstring'''






import argparse
import re

import requests
import torch

# git clone https://github.com/salesforce/BLIP.git
from models.blip import blip_decoder
from models.blip_itm import blip_itm
from models.blip_vqa import blip_vqa
from PIL import Image
from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode

from transformers import (
    BertTokenizer,
    BlipConfig,
    BlipForConditionalGeneration,
    BlipForImageTextRetrieval,
    BlipForQuestionAnswering,
)







def     UpperCamelCase_      (   A__ :       List[Any]							,   A__ :       Optional[int] ):
  '''simple docstring'''





  lowerCAmelCase_      :				List[str]      =       """https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg"""
  lowerCAmelCase_      :				Dict      =       Image.open(requests.get(A__							,   stream=A__ ).raw ).convert("""RGB""" )

  lowerCAmelCase_      :				List[str]      =       transforms.Compose(
      [
          transforms.Resize((image_size, image_size)							,   interpolation=InterpolationMode.BICUBIC ),
          transforms.ToTensor(),
          transforms.Normalize((0.48145466, 0.4578275, 0.40821073)							,   (0.26862954, 0.26130258, 0.27577711) ),
      ] )
  lowerCAmelCase_      :				List[Any]      =       transform(A__ ).unsqueeze(0 ).to(A__ )
  return image







def     UpperCamelCase_      (   A__ :       List[Any] ):
  '''simple docstring'''





  if "visual_encoder" in key:
    lowerCAmelCase_      :				Dict      =       re.sub("""visual_encoder*"""							,   """vision_model.encoder"""							,   A__ )
  if "blocks" in key:
    lowerCAmelCase_      :				List[Any]      =       re.sub(R"""blocks"""							,   """layers"""							,   A__ )
  if "attn" in key:
    lowerCAmelCase_      :				Union[str, Any]      =       re.sub(R"""attn"""							,   """self_attn"""							,   A__ )
  if "norm1" in key:
    lowerCAmelCase_      :				Dict      =       re.sub(R"""norm1"""							,   """layer_norm1"""							,   A__ )
  if "norm2" in key:
    lowerCAmelCase_      :				str      =       re.sub(R"""norm2"""							,   """layer_norm2"""							,   A__ )
  if "encoder.norm" in key:
    lowerCAmelCase_      :				Optional[Any]      =       re.sub(R"""encoder.norm"""							,   """post_layernorm"""							,   A__ )
  if "encoder.patch_embed.proj" in key:
    lowerCAmelCase_      :				Optional[int]      =       re.sub(R"""encoder.patch_embed.proj"""							,   """embeddings.patch_embedding"""							,   A__ )

  if "encoder.pos_embed" in key:
    lowerCAmelCase_      :				Tuple      =       re.sub(R"""encoder.pos_embed"""							,   """embeddings.position_embedding"""							,   A__ )
  if "encoder.cls_token" in key:
    lowerCAmelCase_      :				List[str]      =       re.sub(R"""encoder.cls_token"""							,   """embeddings.class_embedding"""							,   A__ )

  if "self_attn" in key:
    lowerCAmelCase_      :				Optional[int]      =       re.sub(R"""self_attn.proj"""							,   """self_attn.projection"""							,   A__ )

  return key







@torch.no_grad()
def     UpperCamelCase_      (   A__ :       int							,   A__ :       Tuple=None ):
  '''simple docstring'''





  if config_path is not None:
    lowerCAmelCase_      :				List[Any]      =       BlipConfig.from_pretrained(A__ )
  else:
    lowerCAmelCase_      :				Tuple      =       BlipConfig(projection_dim=5_12							,   text_config={}							,   vision_config={} )

  lowerCAmelCase_      :				Optional[Any]      =       BlipForConditionalGeneration(A__ ).eval()

  lowerCAmelCase_      :				Any      =       """https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth"""

  lowerCAmelCase_      :				str      =       blip_decoder(pretrained=A__							,   image_size=3_84							,   vit="""base""" )
  lowerCAmelCase_      :				Optional[Any]      =       pt_model.eval()

  lowerCAmelCase_      :				int      =       pt_model.state_dict()
  for key in modified_state_dict.copy():
    lowerCAmelCase_      :				Dict      =       modified_state_dict.pop(A__ )
    lowerCAmelCase_      :				Optional[Any]      =       rename_key(A__ )
    lowerCAmelCase_      :				Dict      =       value

  hf_model.load_state_dict(A__ )

  lowerCAmelCase_      :				List[Any]      =       3_84
  lowerCAmelCase_      :				Optional[Any]      =       load_demo_image(image_size=A__							,   device="""cpu""" )
  lowerCAmelCase_      :				List[Any]      =       BertTokenizer.from_pretrained("""bert-base-uncased""" )
  lowerCAmelCase_      :				Optional[int]      =       tokenizer(["""a picture of"""] ).input_ids

  lowerCAmelCase_      :				Any      =       hf_model.generate(A__							,   A__ )

  assert out[0].tolist() == [3_05_22, 10_37, 38_61, 19_97, 10_37, 24_50, 35_64, 20_06, 19_96, 35_09, 20_07, 20_14, 38_99, 1_02]

  lowerCAmelCase_      :				Tuple      =       hf_model.generate(A__ )

  assert out[0].tolist() == [3_05_22, 10_37, 24_50, 35_64, 20_06, 19_96, 35_09, 20_07, 20_14, 38_99, 1_02]

  if pytorch_dump_folder_path is not None:
    hf_model.save_pretrained(A__ )

  # model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_vqa.pth'
  lowerCAmelCase_      :				List[str]      =       (
      """https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_vqa_capfilt_large.pth"""
  )

  lowerCAmelCase_      :				str      =       blip_vqa(pretrained=A__							,   image_size=A__							,   vit="""base""" )
  vqa_model.eval()

  lowerCAmelCase_      :				Optional[Any]      =       vqa_model.state_dict()
  for key in modified_state_dict.copy():
    lowerCAmelCase_      :				Optional[int]      =       modified_state_dict.pop(A__ )
    lowerCAmelCase_      :				List[Any]      =       rename_key(A__ )
    lowerCAmelCase_      :				Dict      =       value

  lowerCAmelCase_      :				Any      =       BlipForQuestionAnswering(A__ )

  hf_vqa_model.load_state_dict(A__ )

  lowerCAmelCase_      :				List[str]      =       ["""How many dogs are in this image?"""]
  lowerCAmelCase_      :				int      =       tokenizer(A__							,   return_tensors="""pt""" ).input_ids

  lowerCAmelCase_      :				Optional[int]      =       hf_vqa_model.generate(A__							,   A__ )
  print(tokenizer.decode(answer[0] ) )

  assert tokenizer.decode(answer[0] ) == "[UNK] 1 [SEP]"
  if pytorch_dump_folder_path is not None:
    hf_vqa_model.save_pretrained(pytorch_dump_folder_path + """_vqa""" )

  lowerCAmelCase_      :				List[str]      =       """https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_coco.pth"""

  lowerCAmelCase_      :				Any      =       blip_itm(pretrained=A__							,   image_size=A__							,   vit="""base""" )
  itm_model.eval()

  lowerCAmelCase_      :				Tuple      =       itm_model.state_dict()
  for key in modified_state_dict.copy():
    lowerCAmelCase_      :				int      =       modified_state_dict.pop(A__ )
    lowerCAmelCase_      :				Optional[int]      =       rename_key(A__ )
    lowerCAmelCase_      :				Tuple      =       value

  lowerCAmelCase_      :				Any      =       BlipForImageTextRetrieval(A__ )

  lowerCAmelCase_      :				Union[str, Any]      =       ["""A picture of a woman with a dog sitting in a beach"""]
  lowerCAmelCase_      :				int      =       tokenizer(
      A__							,   return_tensors="""pt"""							,   padding="""max_length"""							,   truncation=A__							,   max_length=35							,   ).input_ids

  hf_itm_model.load_state_dict(A__ )
  hf_itm_model.eval()

  lowerCAmelCase_      :				Optional[int]      =       hf_itm_model(A__							,   A__							,   use_itm_head=A__ )
  lowerCAmelCase_      :				Optional[Any]      =       hf_itm_model(A__							,   A__							,   use_itm_head=A__ )

  assert out[0].item() == 0.2110687494277954
  assert torch.nn.functional.softmax(out_itm[0]							,   dim=1 )[:, 1].item() == 0.45698845386505127

  if pytorch_dump_folder_path is not None:
    hf_itm_model.save_pretrained(pytorch_dump_folder_path + """_itm""" )


if __name__ == "__main__":
  __A :  List[str]    							=				argparse.ArgumentParser()
  parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
  parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
  __A :  Union[str, Any]    							=				parser.parse_args()

  convert_blip_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)







'''simple docstring'''






import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible






class        __snake_case	(       unittest.TestCase):







    """simple docstring"""



    def __lowercase       (       self   :  Tuple		)       ->  Dict:
      lowerCAmelCase_      :				str      =       [
          """safety_checker/pytorch_model.bin""",
          """safety_checker/model.safetensors""",
          """vae/diffusion_pytorch_model.bin""",
          """vae/diffusion_pytorch_model.safetensors""",
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      self.assertTrue(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  List[Any]		)       ->  int:
      lowerCAmelCase_      :				Tuple      =       [
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      self.assertTrue(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[int]		)       ->  Optional[Any]:
      lowerCAmelCase_      :				int      =       [
          """safety_checker/pytorch_model.bin""",
          """safety_checker/model.safetensors""",
          """vae/diffusion_pytorch_model.bin""",
          """vae/diffusion_pytorch_model.safetensors""",
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
          """unet/diffusion_pytorch_model.bin""",
          # Removed: 'unet/diffusion_pytorch_model.safetensors',
      ]
      self.assertFalse(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  int		)       ->  List[Any]:
      lowerCAmelCase_      :				Dict      =       [
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
      ]
      self.assertTrue(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  str		)       ->  List[str]:
      lowerCAmelCase_      :				Union[str, Any]      =       [
          """safety_checker/pytorch_model.bin""",
          """safety_checker/model.safetensors""",
          """vae/diffusion_pytorch_model.bin""",
          """vae/diffusion_pytorch_model.safetensors""",
          """text_encoder/pytorch_model.bin""",
          # Removed: 'text_encoder/model.safetensors',
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      self.assertFalse(is_safetensors_compatible(lowerCamelCase		)		)

    def __lowercase       (       self   :  List[Any]		)       ->  Tuple:
      lowerCAmelCase_      :				Union[str, Any]      =       [
          """safety_checker/pytorch_model.fp16.bin""",
          """safety_checker/model.fp16.safetensors""",
          """vae/diffusion_pytorch_model.fp16.bin""",
          """vae/diffusion_pytorch_model.fp16.safetensors""",
          """text_encoder/pytorch_model.fp16.bin""",
          """text_encoder/model.fp16.safetensors""",
          """unet/diffusion_pytorch_model.fp16.bin""",
          """unet/diffusion_pytorch_model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				Union[str, Any]      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[Any]		)       ->  List[str]:
      lowerCAmelCase_      :				str      =       [
          """unet/diffusion_pytorch_model.fp16.bin""",
          """unet/diffusion_pytorch_model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				Optional[int]      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Tuple		)       ->  Optional[Any]:
      # pass variant but use the non-variant filenames
      lowerCAmelCase_      :				Dict      =       [
          """unet/diffusion_pytorch_model.bin""",
          """unet/diffusion_pytorch_model.safetensors""",
      ]
      lowerCAmelCase_      :				str      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Optional[int]		)       ->  List[str]:
      lowerCAmelCase_      :				str      =       [
          """safety_checker/pytorch_model.fp16.bin""",
          """safety_checker/model.fp16.safetensors""",
          """vae/diffusion_pytorch_model.fp16.bin""",
          """vae/diffusion_pytorch_model.fp16.safetensors""",
          """text_encoder/pytorch_model.fp16.bin""",
          """text_encoder/model.fp16.safetensors""",
          """unet/diffusion_pytorch_model.fp16.bin""",
          # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
      ]
      lowerCAmelCase_      :				List[str]      =       """fp16"""
      self.assertFalse(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  Union[str, Any]		)       ->  Optional[int]:
      lowerCAmelCase_      :				str      =       [
          """text_encoder/pytorch_model.fp16.bin""",
          """text_encoder/model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				Optional[Any]      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)

    def __lowercase       (       self   :  List[Any]		)       ->  List[Any]:
      # pass variant but use the non-variant filenames
      lowerCAmelCase_      :				Dict      =       [
          """text_encoder/pytorch_model.bin""",
          """text_encoder/model.safetensors""",
      ]
      lowerCAmelCase_      :				Any      =       """fp16"""
      self.assertTrue(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)





    def __lowercase       (       self   :  Dict		)       ->  Any:
      lowerCAmelCase_      :				Optional[int]      =       [
          """safety_checker/pytorch_model.fp16.bin""",
          """safety_checker/model.fp16.safetensors""",
          """vae/diffusion_pytorch_model.fp16.bin""",
          """vae/diffusion_pytorch_model.fp16.safetensors""",
          """text_encoder/pytorch_model.fp16.bin""",
          # 'text_encoder/model.fp16.safetensors',
          """unet/diffusion_pytorch_model.fp16.bin""",
          """unet/diffusion_pytorch_model.fp16.safetensors""",
      ]
      lowerCAmelCase_      :				int      =       """fp16"""
      self.assertFalse(is_safetensors_compatible(lowerCamelCase   ,						variant=lowerCamelCase		)		)



