



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				(UpperCamelCase		:      int      ,			UpperCamelCase		:      List[Any] )  ->Union[str, Any]:



			'''simple docstring'''







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

			lowerCamelCase__ :      Union[str, Any]								=		transforms.Compose(
			    [
			        transforms.Resize((image_size, image_size)      ,			interpolation=InterpolationMode.BICUBIC ),
			        transforms.ToTensor(),
			        transforms.Normalize((0.48_145_466, 0.4_578_275, 0.40_821_073)      ,			(0.26_862_954, 0.26_130_258, 0.27_577_711) ),
			    ] )
			lowerCamelCase__ :      Optional[Any]								=		transform(UpperCamelCase ).unsqueeze(0 ).to(UpperCamelCase )
			return image






def      _A				(UpperCamelCase		:      str )  ->Tuple:



			'''simple docstring'''







			if "visual_encoder" in key:
						lowerCamelCase__ :      Tuple								=		re.sub("""visual_encoder*"""      ,			"""vision_model.encoder"""      ,			UpperCamelCase )
			if "blocks" in key:
						lowerCamelCase__ :      Any								=		re.sub(r"""blocks"""      ,			"""layers"""      ,			UpperCamelCase )
			if "attn" in key:
						lowerCamelCase__ :      Optional[int]								=		re.sub(r"""attn"""      ,			"""self_attn"""      ,			UpperCamelCase )
			if "norm1" in key:
						lowerCamelCase__ :      Optional[Any]								=		re.sub(r"""norm1"""      ,			"""layer_norm1"""      ,			UpperCamelCase )
			if "norm2" in key:
						lowerCamelCase__ :      int								=		re.sub(r"""norm2"""      ,			"""layer_norm2"""      ,			UpperCamelCase )
			if "encoder.norm" in key:
						lowerCamelCase__ :      int								=		re.sub(r"""encoder.norm"""      ,			"""post_layernorm"""      ,			UpperCamelCase )
			if "encoder.patch_embed.proj" in key:
						lowerCamelCase__ :      str								=		re.sub(r"""encoder.patch_embed.proj"""      ,			"""embeddings.patch_embedding"""      ,			UpperCamelCase )

			if "encoder.pos_embed" in key:
						lowerCamelCase__ :      Optional[int]								=		re.sub(r"""encoder.pos_embed"""      ,			"""embeddings.position_embedding"""      ,			UpperCamelCase )
			if "encoder.cls_token" in key:
						lowerCamelCase__ :      str								=		re.sub(r"""encoder.cls_token"""      ,			"""embeddings.class_embedding"""      ,			UpperCamelCase )

			if "self_attn" in key:
						lowerCamelCase__ :      Dict								=		re.sub(r"""self_attn.proj"""      ,			"""self_attn.projection"""      ,			UpperCamelCase )

			return key






@torch.no_grad()
def      _A				(UpperCamelCase		:      Any      ,			UpperCamelCase		:      str=None )  ->Union[str, Any]:



			'''simple docstring'''







			if config_path is not None:
						lowerCamelCase__ :      List[Any]								=		BlipConfig.from_pretrained(UpperCamelCase )
			else:
						lowerCamelCase__ :      Optional[Any]								=		BlipConfig(projection_dim=512      ,			text_config={}      ,			vision_config={} )

			lowerCamelCase__ :      Dict								=		BlipForConditionalGeneration(UpperCamelCase ).eval()

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

			lowerCamelCase__ :      int								=		blip_decoder(pretrained=UpperCamelCase      ,			image_size=384      ,			vit="""base""" )
			lowerCamelCase__ :      Dict								=		pt_model.eval()

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

			hf_model.load_state_dict(UpperCamelCase )

			lowerCamelCase__ :      Dict								=		384
			lowerCamelCase__ :      Union[str, Any]								=		load_demo_image(image_size=UpperCamelCase      ,			device="""cpu""" )
			lowerCamelCase__ :      Union[str, Any]								=		BertTokenizer.from_pretrained("""bert-base-uncased""" )
			lowerCamelCase__ :      Optional[Any]								=		tokenizer(["""a picture of"""] ).input_ids

			lowerCamelCase__ :      Optional[int]								=		hf_model.generate(UpperCamelCase      ,			UpperCamelCase )

			assert out[0].tolist() == [30522, 1037, 3861, 1997, 1037, 2450, 3564, 2006, 1996, 3509, 2007, 2014, 3899, 102]

			lowerCamelCase__ :      Tuple								=		hf_model.generate(UpperCamelCase )

			assert out[0].tolist() == [30522, 1037, 2450, 3564, 2006, 1996, 3509, 2007, 2014, 3899, 102]

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

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

			lowerCamelCase__ :      int								=		blip_vqa(pretrained=UpperCamelCase      ,			image_size=UpperCamelCase      ,			vit="""base""" )
			vqa_model.eval()

			lowerCamelCase__ :      str								=		vqa_model.state_dict()
			for key in modified_state_dict.copy():
						lowerCamelCase__ :      Tuple								=		modified_state_dict.pop(UpperCamelCase )
						lowerCamelCase__ :      Optional[int]								=		rename_key(UpperCamelCase )
						lowerCamelCase__ :      Optional[int]								=		value

			lowerCamelCase__ :      List[Any]								=		BlipForQuestionAnswering(UpperCamelCase )

			hf_vqa_model.load_state_dict(UpperCamelCase )

			lowerCamelCase__ :      Any								=		["""How many dogs are in this image?"""]
			lowerCamelCase__ :      int								=		tokenizer(UpperCamelCase      ,			return_tensors="""pt""" ).input_ids

			lowerCamelCase__ :      Any								=		hf_vqa_model.generate(UpperCamelCase      ,			UpperCamelCase )
			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__ :      Union[str, Any]								=		blip_itm(pretrained=UpperCamelCase      ,			image_size=UpperCamelCase      ,			vit="""base""" )
			itm_model.eval()

			lowerCamelCase__ :      Union[str, Any]								=		itm_model.state_dict()
			for key in modified_state_dict.copy():
						lowerCamelCase__ :      List[str]								=		modified_state_dict.pop(UpperCamelCase )
						lowerCamelCase__ :      Any								=		rename_key(UpperCamelCase )
						lowerCamelCase__ :      Tuple								=		value

			lowerCamelCase__ :      Tuple								=		BlipForImageTextRetrieval(UpperCamelCase )

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

			hf_itm_model.load_state_dict(UpperCamelCase )
			hf_itm_model.eval()

			lowerCamelCase__ :      Tuple								=		hf_itm_model(UpperCamelCase      ,			UpperCamelCase      ,			use_itm_head=UpperCamelCase )
			lowerCamelCase__ :      str								=		hf_itm_model(UpperCamelCase      ,			UpperCamelCase      ,			use_itm_head=UpperCamelCase )

			assert out[0].item() == 0.2_110_687_494_277_954
			assert torch.nn.functional.softmax(out_itm[0]      ,			dim=1 )[:, 1].item() == 0.45_698_845_386_505_127

			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)








import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


_lowercase     							=						logging.get_logger(__name__)






def      _A				(UpperCamelCase		:      Dict      ,			UpperCamelCase		:      Optional[int] )  ->Tuple:



			'''simple docstring'''







			lowerCamelCase__ :      List[str]								=		nn.functional.normalize(UpperCamelCase )
			lowerCamelCase__ :      Tuple								=		nn.functional.normalize(UpperCamelCase )
			return torch.mm(UpperCamelCase      ,			normalized_text_embeds.t() )
class 				__A					( A_      ):
			UpperCamelCase					:Optional[Any]     =			CLIPConfig

			UpperCamelCase					:Optional[Any]     =			['''CLIPEncoderLayer''']





			def __init__(self							,    __magic_name__				):
						super().__init__(__magic_name__				)

						lowerCamelCase__ :      Union[str, Any]								=		CLIPVisionModel(config.vision_config				)
						lowerCamelCase__ :      Any								=		nn.Linear(config.vision_config.hidden_size							,    config.projection_dim							,    bias=__magic_name__				)

						lowerCamelCase__ :      Any								=		nn.Parameter(torch.ones(17							,    config.projection_dim				)							,    requires_grad=__magic_name__				)
						lowerCamelCase__ :      Union[str, Any]								=		nn.Parameter(torch.ones(3							,    config.projection_dim				)							,    requires_grad=__magic_name__				)

						lowerCamelCase__ :      Any								=		nn.Parameter(torch.ones(17				)							,    requires_grad=__magic_name__				)
						lowerCamelCase__ :      Optional[Any]								=		nn.Parameter(torch.ones(3				)							,    requires_grad=__magic_name__				)





			@torch.no_grad()
			def        _snake_case    (self							,    __magic_name__							,    __magic_name__				):
						lowerCamelCase__ :      Any								=		self.vision_model(__magic_name__				)[1]  # pooled_output
						lowerCamelCase__ :      Tuple								=		self.visual_projection(__magic_name__				)

						# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
						lowerCamelCase__ :      Any								=		cosine_distance(__magic_name__							,    self.special_care_embeds				).cpu().float().numpy()
						lowerCamelCase__ :      str								=		cosine_distance(__magic_name__							,    self.concept_embeds				).cpu().float().numpy()

						lowerCamelCase__ :      int								=		[]
						lowerCamelCase__ :      str								=		image_embeds.shape[0]
						for i in range(__magic_name__				):
									lowerCamelCase__ :      str								=		{"""special_scores""": {}, """special_care""": [], """concept_scores""": {}, """bad_concepts""": []}

									# increase this value to create a stronger `nfsw` filter
									# at the cost of increasing the possibility of filtering benign images
									lowerCamelCase__ :      Optional[Any]								=		0.0

									for concept_idx in range(len(special_cos_dist[0]				)				):
												lowerCamelCase__ :      List[str]								=		special_cos_dist[i][concept_idx]
												lowerCamelCase__ :      Tuple								=		self.special_care_embeds_weights[concept_idx].item()
												lowerCamelCase__ :      str								=		round(concept_cos - concept_threshold + adjustment							,    3				)
												if result_img["special_scores"][concept_idx] > 0:
															result_img["special_care"].append({concept_idx, result_img["""special_scores"""][concept_idx]}				)
															lowerCamelCase__ :      List[str]								=		0.01

									for concept_idx in range(len(cos_dist[0]				)				):
												lowerCamelCase__ :      str								=		cos_dist[i][concept_idx]
												lowerCamelCase__ :      Dict								=		self.concept_embeds_weights[concept_idx].item()
												lowerCamelCase__ :      int								=		round(concept_cos - concept_threshold + adjustment							,    3				)
												if result_img["concept_scores"][concept_idx] > 0:
															result_img["bad_concepts"].append(__magic_name__				)

									result.append(__magic_name__				)

						lowerCamelCase__ :      Tuple								=		[len(res["""bad_concepts"""]				) > 0 for res in result]

						return images, has_nsfw_concepts





			@torch.no_grad()
			def        _snake_case    (self							,    __magic_name__							,    __magic_name__				):
						lowerCamelCase__ :      Union[str, Any]								=		self.vision_model(__magic_name__				)[1]  # pooled_output
						lowerCamelCase__ :      str								=		self.visual_projection(__magic_name__				)

						lowerCamelCase__ :      Optional[Any]								=		cosine_distance(__magic_name__							,    self.special_care_embeds				)
						lowerCamelCase__ :      Dict								=		cosine_distance(__magic_name__							,    self.concept_embeds				)

						# increase this value to create a stronger `nsfw` filter
						# at the cost of increasing the possibility of filtering benign images
						lowerCamelCase__ :      Any								=		0.0

						lowerCamelCase__ :      Dict								=		special_cos_dist - self.special_care_embeds_weights + adjustment
						# special_scores = special_scores.round(decimals=3)
						lowerCamelCase__ :      Union[str, Any]								=		torch.any(special_scores > 0							,    dim=1				)
						lowerCamelCase__ :      Optional[Any]								=		special_care * 0.01
						lowerCamelCase__ :      Dict								=		special_adjustment.unsqueeze(1				).expand(-1							,    cos_dist.shape[1]				)

						lowerCamelCase__ :      Tuple								=		(cos_dist - self.concept_embeds_weights) + special_adjustment
						# concept_scores = concept_scores.round(decimals=3)
						lowerCamelCase__ :      Any								=		torch.any(concept_scores > 0							,    dim=1				)

						return images, has_nsfw_concepts



