
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     (				lowerCAmelCase__							:				List[str]  ,   lowerCAmelCase__							:				List[str]  ):
 __a :  str						=  '''https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg'''
 __a :  Dict						=  Image.open(requests.get(lowerCAmelCase__  ,   stream=lowerCAmelCase__  ).raw  ).convert('''RGB'''  )

 __a :  Union[str, Any]						=  transforms.Compose(
     [
         transforms.Resize((image_size, image_size)  ,   interpolation=InterpolationMode.BICUBIC  ),
         transforms.ToTensor(),
         transforms.Normalize((0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73)  ,   (0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11)  ),
     ]  )
 __a :  List[str]						=  transform(lowerCAmelCase__  ).unsqueeze(0  ).to(lowerCAmelCase__  )
 return image




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

 if "encoder.pos_embed" in key:
  __a :  Any						=  re.sub(R'''encoder.pos_embed'''  ,   '''embeddings.position_embedding'''  ,   lowerCAmelCase__  )
 if "encoder.cls_token" in key:
  __a :  str						=  re.sub(R'''encoder.cls_token'''  ,   '''embeddings.class_embedding'''  ,   lowerCAmelCase__  )

 if "self_attn" in key:
  __a :  List[str]						=  re.sub(R'''self_attn.proj'''  ,   '''self_attn.projection'''  ,   lowerCAmelCase__  )

 return key




@torch.no_grad()
def    __UpperCamelCase     (				lowerCAmelCase__							:				int  ,   lowerCAmelCase__							:				Any=None  ):
 if config_path is not None:
  __a :  Dict						=  BlipConfig.from_pretrained(lowerCAmelCase__  )
 else:
  __a :  List[str]						=  BlipConfig(projection_dim=5_1_2  ,   text_config={}  ,   vision_config={}  )

 __a :  Dict						=  BlipForConditionalGeneration(lowerCAmelCase__  ).eval()

 __a :  Optional[Any]						=  '''https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth'''

 __a :  List[str]						=  blip_decoder(pretrained=lowerCAmelCase__  ,   image_size=3_8_4  ,   vit='''base'''  )
 __a :  str						=  pt_model.eval()

 __a :  Union[str, Any]						=  pt_model.state_dict()
 for key in modified_state_dict.copy():
  __a :  Any						=  modified_state_dict.pop(lowerCAmelCase__  )
  __a :  Dict						=  rename_key(lowerCAmelCase__  )
  __a :  List[str]						=  value

 hf_model.load_state_dict(lowerCAmelCase__  )

 __a :  List[str]						=  3_8_4
 __a :  Any						=  load_demo_image(image_size=lowerCAmelCase__  ,   device='''cpu'''  )
 __a :  int						=  BertTokenizer.from_pretrained('''bert-base-uncased'''  )
 __a :  Optional[Any]						=  tokenizer(['''a picture of''']  ).input_ids

 __a :  Union[str, Any]						=  hf_model.generate(lowerCAmelCase__  ,   lowerCAmelCase__  )

 assert out[0].tolist() == [3_0_5_2_2, 1_0_3_7, 3_8_6_1, 1_9_9_7, 1_0_3_7, 2_4_5_0, 3_5_6_4, 2_0_0_6, 1_9_9_6, 3_5_0_9, 2_0_0_7, 2_0_1_4, 3_8_9_9, 1_0_2]

 __a :  int						=  hf_model.generate(lowerCAmelCase__  )

 assert out[0].tolist() == [3_0_5_2_2, 1_0_3_7, 2_4_5_0, 3_5_6_4, 2_0_0_6, 1_9_9_6, 3_5_0_9, 2_0_0_7, 2_0_1_4, 3_8_9_9, 1_0_2]

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

 # model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_vqa.pth'
 __a :  Optional[Any]						=  (
     '''https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_vqa_capfilt_large.pth'''
 )

 __a :  str						=  blip_vqa(pretrained=lowerCAmelCase__  ,   image_size=lowerCAmelCase__  ,   vit='''base'''  )
 vqa_model.eval()

 __a :  str						=  vqa_model.state_dict()
 for key in modified_state_dict.copy():
  __a :  Tuple						=  modified_state_dict.pop(lowerCAmelCase__  )
  __a :  str						=  rename_key(lowerCAmelCase__  )
  __a :  str						=  value

 __a :  List[str]						=  BlipForQuestionAnswering(lowerCAmelCase__  )

 hf_vqa_model.load_state_dict(lowerCAmelCase__  )

 __a :  int						=  ['''How many dogs are in this image?''']
 __a :  List[Any]						=  tokenizer(lowerCAmelCase__  ,   return_tensors='''pt'''  ).input_ids

 __a :  str						=  hf_vqa_model.generate(lowerCAmelCase__  ,   lowerCAmelCase__  )
 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'''  )

 __a :  Union[str, Any]						=  '''https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_coco.pth'''

 __a :  int						=  blip_itm(pretrained=lowerCAmelCase__  ,   image_size=lowerCAmelCase__  ,   vit='''base'''  )
 itm_model.eval()

 __a :  int						=  itm_model.state_dict()
 for key in modified_state_dict.copy():
  __a :  List[Any]						=  modified_state_dict.pop(lowerCAmelCase__  )
  __a :  Optional[Any]						=  rename_key(lowerCAmelCase__  )
  __a :  Dict						=  value

 __a :  Dict						=  BlipForImageTextRetrieval(lowerCAmelCase__  )

 __a :  List[str]						=  ['''A picture of a woman with a dog sitting in a beach''']
 __a :  Tuple						=  tokenizer(
     lowerCAmelCase__  ,   return_tensors='''pt'''  ,   padding='''max_length'''  ,   truncation=lowerCAmelCase__  ,   max_length=3_5  ,   ).input_ids

 hf_itm_model.load_state_dict(lowerCAmelCase__  )
 hf_itm_model.eval()

 __a :  Union[str, Any]						=  hf_itm_model(lowerCAmelCase__  ,   lowerCAmelCase__  ,   use_itm_head=lowerCAmelCase__  )
 __a :  List[Any]						=  hf_itm_model(lowerCAmelCase__  ,   lowerCAmelCase__  ,   use_itm_head=lowerCAmelCase__  )

 assert out[0].item() == 0.21_10_68_74_94_27_79_54
 assert torch.nn.functional.softmax(out_itm[0]  ,   dim=1  )[:, 1].item() == 0.4_56_98_84_53_86_50_51_27

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


if __name__ == "__main__":
      lowercase__       =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')
      lowercase__       =parser.parse_args()

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


def    __UpperCamelCase     (				lowerCAmelCase__							:				int = 5_0_0_0_0_0_0_0  ):
 __a :  int						=  set()
 __a :  str						=  int((limit - 2_4) ** (1 / 2)  )

 __a :  int						=  set(range(3  ,   prime_square_limit + 1  ,   2  )  )
 primes.add(2  )
 for p in range(3  ,   prime_square_limit + 1  ,   2  ):
  if p not in primes:
   continue
  primes.difference_update(set(range(p * p  ,   prime_square_limit + 1  ,   lowerCAmelCase__  )  )  )

 for primea in primes:
  __a :  Union[str, Any]						=  primea * primea
  for primea in primes:
   __a :  Union[str, Any]						=  primea * primea * primea
   if square + cube >= limit - 1_6:
    break
   for primea in primes:
    __a :  int						=  primea * primea * primea * primea
    __a :  Union[str, Any]						=  square + cube + tetr
    if total >= limit:
     break
    ret.add(lowerCAmelCase__  )

 return len(lowerCAmelCase__  )


if __name__ == "__main__":
      print(F"""{solution() = }""")
