




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 			_a   (		SCREAMING_SNAKE_CASE_							:       List[Any]   ,   SCREAMING_SNAKE_CASE_							:       List[Any]  ):
					__lowerCAmelCase            = "https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg"
					__lowerCAmelCase            = Image.open(requests.get(SCREAMING_SNAKE_CASE_   ,   stream=SCREAMING_SNAKE_CASE_  ).raw  ).convert("RGB"  )

					__lowerCAmelCase            = 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)  ),
					    ]  )
					__lowerCAmelCase            = transform(SCREAMING_SNAKE_CASE_  ).unsqueeze(0  ).to(SCREAMING_SNAKE_CASE_  )
					return image







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Dict  ):
					if "visual_encoder" in key:
										__lowerCAmelCase            = re.sub("visual_encoder*"   ,   "vision_model.encoder"   ,   SCREAMING_SNAKE_CASE_  )
					if "blocks" in key:
										__lowerCAmelCase            = re.sub(R"blocks"   ,   "layers"   ,   SCREAMING_SNAKE_CASE_  )
					if "attn" in key:
										__lowerCAmelCase            = re.sub(R"attn"   ,   "self_attn"   ,   SCREAMING_SNAKE_CASE_  )
					if "norm1" in key:
										__lowerCAmelCase            = re.sub(R"norm1"   ,   "layer_norm1"   ,   SCREAMING_SNAKE_CASE_  )
					if "norm2" in key:
										__lowerCAmelCase            = re.sub(R"norm2"   ,   "layer_norm2"   ,   SCREAMING_SNAKE_CASE_  )
					if "encoder.norm" in key:
										__lowerCAmelCase            = re.sub(R"encoder.norm"   ,   "post_layernorm"   ,   SCREAMING_SNAKE_CASE_  )
					if "encoder.patch_embed.proj" in key:
										__lowerCAmelCase            = re.sub(R"encoder.patch_embed.proj"   ,   "embeddings.patch_embedding"   ,   SCREAMING_SNAKE_CASE_  )

					if "encoder.pos_embed" in key:
										__lowerCAmelCase            = re.sub(R"encoder.pos_embed"   ,   "embeddings.position_embedding"   ,   SCREAMING_SNAKE_CASE_  )
					if "encoder.cls_token" in key:
										__lowerCAmelCase            = re.sub(R"encoder.cls_token"   ,   "embeddings.class_embedding"   ,   SCREAMING_SNAKE_CASE_  )

					if "self_attn" in key:
										__lowerCAmelCase            = re.sub(R"self_attn.proj"   ,   "self_attn.projection"   ,   SCREAMING_SNAKE_CASE_  )

					return key







@torch.no_grad()
def 			_a   (		SCREAMING_SNAKE_CASE_							:       Dict   ,   SCREAMING_SNAKE_CASE_							:       List[str]=None  ):
					if config_path is not None:
										__lowerCAmelCase            = BlipConfig.from_pretrained(SCREAMING_SNAKE_CASE_  )
					else:
										__lowerCAmelCase            = BlipConfig(projection_dim=5_12   ,   text_config={}   ,   vision_config={}  )

					__lowerCAmelCase            = BlipForConditionalGeneration(SCREAMING_SNAKE_CASE_  ).eval()

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

					__lowerCAmelCase            = blip_decoder(pretrained=SCREAMING_SNAKE_CASE_   ,   image_size=3_84   ,   vit="base"  )
					__lowerCAmelCase            = pt_model.eval()

					__lowerCAmelCase            = pt_model.state_dict()
					for key in modified_state_dict.copy():
										__lowerCAmelCase            = modified_state_dict.pop(SCREAMING_SNAKE_CASE_  )
										__lowerCAmelCase            = rename_key(SCREAMING_SNAKE_CASE_  )
										__lowerCAmelCase            = value

					hf_model.load_state_dict(SCREAMING_SNAKE_CASE_  )

					__lowerCAmelCase            = 3_84
					__lowerCAmelCase            = load_demo_image(image_size=SCREAMING_SNAKE_CASE_   ,   device="cpu"  )
					__lowerCAmelCase            = BertTokenizer.from_pretrained("bert-base-uncased"  )
					__lowerCAmelCase            = tokenizer(["a picture of"]  ).input_ids

					__lowerCAmelCase            = hf_model.generate(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  )

					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            = hf_model.generate(SCREAMING_SNAKE_CASE_  )

					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(SCREAMING_SNAKE_CASE_  )

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

					__lowerCAmelCase            = blip_vqa(pretrained=SCREAMING_SNAKE_CASE_   ,   image_size=SCREAMING_SNAKE_CASE_   ,   vit="base"  )
					vqa_model.eval()

					__lowerCAmelCase            = vqa_model.state_dict()
					for key in modified_state_dict.copy():
										__lowerCAmelCase            = modified_state_dict.pop(SCREAMING_SNAKE_CASE_  )
										__lowerCAmelCase            = rename_key(SCREAMING_SNAKE_CASE_  )
										__lowerCAmelCase            = value

					__lowerCAmelCase            = BlipForQuestionAnswering(SCREAMING_SNAKE_CASE_  )

					hf_vqa_model.load_state_dict(SCREAMING_SNAKE_CASE_  )

					__lowerCAmelCase            = ["How many dogs are in this image?"]
					__lowerCAmelCase            = tokenizer(SCREAMING_SNAKE_CASE_   ,   return_tensors="pt"  ).input_ids

					__lowerCAmelCase            = hf_vqa_model.generate(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_  )
					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            = "https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_retrieval_coco.pth"

					__lowerCAmelCase            = blip_itm(pretrained=SCREAMING_SNAKE_CASE_   ,   image_size=SCREAMING_SNAKE_CASE_   ,   vit="base"  )
					itm_model.eval()

					__lowerCAmelCase            = itm_model.state_dict()
					for key in modified_state_dict.copy():
										__lowerCAmelCase            = modified_state_dict.pop(SCREAMING_SNAKE_CASE_  )
										__lowerCAmelCase            = rename_key(SCREAMING_SNAKE_CASE_  )
										__lowerCAmelCase            = value

					__lowerCAmelCase            = BlipForImageTextRetrieval(SCREAMING_SNAKE_CASE_  )

					__lowerCAmelCase            = ["A picture of a woman with a dog sitting in a beach"]
					__lowerCAmelCase            = tokenizer(
					    SCREAMING_SNAKE_CASE_   ,   return_tensors="pt"   ,   padding="max_length"   ,   truncation=SCREAMING_SNAKE_CASE_   ,   max_length=35   ,   ).input_ids

					hf_itm_model.load_state_dict(SCREAMING_SNAKE_CASE_  )
					hf_itm_model.eval()

					__lowerCAmelCase            = hf_itm_model(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_   ,   use_itm_head=SCREAMING_SNAKE_CASE_  )
					__lowerCAmelCase            = hf_itm_model(SCREAMING_SNAKE_CASE_   ,   SCREAMING_SNAKE_CASE_   ,   use_itm_head=SCREAMING_SNAKE_CASE_  )

					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__":
					UpperCamelCase__					    =    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""")
					UpperCamelCase__					    =    parser.parse_args()

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










from math import ceil, sqrt







def 			_a   (		SCREAMING_SNAKE_CASE_							:       int = 1_00_00_00  ):
					__lowerCAmelCase            = 0

					for outer_width in range(3   ,   (limit // 4) + 2  ):
										if outer_width**2 > limit:
															__lowerCAmelCase            = max(ceil(sqrt(outer_width**2 - limit  )  )   ,   1  )
										else:
															__lowerCAmelCase            = 1
										if (outer_width - hole_width_lower_bound) % 2:
															hole_width_lower_bound += 1

										answer += (outer_width - hole_width_lower_bound - 2) // 2 + 1

					return answer


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




