



'''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 					snake_case_	(   ) ->  Tuple:
							UpperCAmelCase					:			List[Any]									=						"""https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png"""
							UpperCAmelCase					:			List[Any]									=						Image.open(requests.get(lowerCAmelCase__  , stream=lowerCAmelCase__    ).raw    ).convert('''RGB'''    )

							return image


def 					snake_case_	(   _lowerCAmelCase    :		Any    ) ->  Tuple:
							UpperCAmelCase					:			List[str]									=						[]
							# 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 					snake_case_	(   _lowerCAmelCase    :		Optional[int]  , _lowerCAmelCase    :		Dict  , _lowerCAmelCase    :		str    ) ->  Any:
							UpperCAmelCase					:			int									=						dct.pop(lowerCAmelCase__    )
							UpperCAmelCase					:			Optional[int]									=						val


def 					snake_case_	(   _lowerCAmelCase    :		List[str]  , _lowerCAmelCase    :		str    ) ->  List[Any]:


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

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


def 					snake_case_	(   _lowerCAmelCase    :		Optional[Any]  , _lowerCAmelCase    :		Union[str, Any]    ) ->  Optional[Any]:
							UpperCAmelCase					:			List[str]									=						364 if """coco""" in model_name else 224
							UpperCAmelCase					:			List[str]									=						BlipaVisionConfig(image_size=lowerCAmelCase__    ).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:
														UpperCAmelCase					:			List[Any]									=						OPTConfig.from_pretrained('''facebook/opt-2.7b'''  , eos_token_id=lowerCAmelCase__    ).to_dict()
							elif "opt-6.7b" in model_name:
														UpperCAmelCase					:			Union[str, Any]									=						OPTConfig.from_pretrained('''facebook/opt-6.7b'''  , eos_token_id=lowerCAmelCase__    ).to_dict()
							elif "t5-xl" in model_name:
														UpperCAmelCase					:			Dict									=						TaConfig.from_pretrained('''google/flan-t5-xl'''  , dense_act_fn='''gelu'''  , bos_token_id=1    ).to_dict()
							elif "t5-xxl" in model_name:
														UpperCAmelCase					:			List[Any]									=						TaConfig.from_pretrained('''google/flan-t5-xxl'''  , dense_act_fn='''gelu'''  , bos_token_id=1    ).to_dict()

							UpperCAmelCase					:			Dict									=						BlipaConfig(vision_config=lowerCAmelCase__  , text_config=lowerCAmelCase__    )

							return config, image_size


@torch.no_grad()
def 					snake_case_	(   _lowerCAmelCase    :		str  , _lowerCAmelCase    :		Optional[Any]=None  , _lowerCAmelCase    :		Any=False    ) ->  int:
							UpperCAmelCase					:			Any									=						(
							    AutoTokenizer.from_pretrained('''facebook/opt-2.7b'''    )
							    if """opt""" in model_name
							    else AutoTokenizer.from_pretrained('''google/flan-t5-xl'''    )
							)
							UpperCAmelCase					:			Dict									=						tokenizer('''\n'''  , add_special_tokens=lowerCAmelCase__    ).input_ids[0]
							UpperCAmelCase					:			Optional[Any]									=						get_blipa_config(lowerCAmelCase__  , eos_token_id=lowerCAmelCase__    )

							UpperCAmelCase					:			Tuple									=						BlipaForConditionalGeneration(lowerCAmelCase__    ).eval()

							UpperCAmelCase					:			Optional[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"""),
							}

							UpperCAmelCase					:			List[Any]									=						model_name_to_original[model_name]

							# load original model
							print('''Loading original model...'''    )
							UpperCAmelCase					:			Any									=						"""cuda""" if torch.cuda.is_available() else """cpu"""
							UpperCAmelCase					:			int									=						load_model_and_preprocess(
							    name=lowerCAmelCase__  , model_type=lowerCAmelCase__  , is_eval=lowerCAmelCase__  , device=lowerCAmelCase__    )
							original_model.eval()
							print('''Done!'''    )

							# update state dict keys
							UpperCAmelCase					:			List[Any]									=						original_model.state_dict()
							UpperCAmelCase					:			Union[str, Any]									=						create_rename_keys(lowerCAmelCase__    )
							for src, dest in rename_keys:
														rename_key(lowerCAmelCase__  , lowerCAmelCase__  , lowerCAmelCase__    )

							# some keys can be renamed efficiently
							for key, val in state_dict.copy().items():
														UpperCAmelCase					:			List[str]									=						state_dict.pop(lowerCAmelCase__    )
														if key.startswith('''Qformer.bert'''    ):
																					UpperCAmelCase					:			Optional[Any]									=						key.replace('''Qformer.bert'''  , '''qformer'''    )
														if "attention.self" in key:
																					UpperCAmelCase					:			int									=						key.replace('''self'''  , '''attention'''    )
														if "opt_proj" in key:
																					UpperCAmelCase					:			Optional[Any]									=						key.replace('''opt_proj'''  , '''language_projection'''    )
														if "t5_proj" in key:
																					UpperCAmelCase					:			Optional[int]									=						key.replace('''t5_proj'''  , '''language_projection'''    )
														if key.startswith('''opt'''    ):
																					UpperCAmelCase					:			List[Any]									=						key.replace('''opt'''  , '''language'''    )
														if key.startswith('''t5'''    ):
																					UpperCAmelCase					:			str									=						key.replace('''t5'''  , '''language'''    )
														UpperCAmelCase					:			int									=						val

							# read in qv biases
							read_in_q_v_bias(lowerCAmelCase__  , lowerCAmelCase__    )

							UpperCAmelCase					:			int									=						hf_model.load_state_dict(lowerCAmelCase__  , strict=lowerCAmelCase__    )
							assert len(lowerCAmelCase__    ) == 0
							assert unexpected_keys == ["qformer.embeddings.position_ids"]

							UpperCAmelCase					:			str									=						load_demo_image()
							UpperCAmelCase					:			str									=						vis_processors["""eval"""](lowerCAmelCase__    ).unsqueeze(0    ).to(lowerCAmelCase__    )
							UpperCAmelCase					:			Tuple									=						tokenizer(['''\n''']  , return_tensors='''pt'''    ).input_ids.to(lowerCAmelCase__    )

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

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

							original_model.to(lowerCAmelCase__    )
							hf_model.to(lowerCAmelCase__    )
							with torch.no_grad():
														if "opt" in model_name:
																					UpperCAmelCase					:			Tuple									=						original_model({'''image''': original_pixel_values, '''text_input''': ['''''']}    ).logits
																					UpperCAmelCase					:			Tuple									=						hf_model(lowerCAmelCase__  , lowerCAmelCase__    ).logits
														else:
																					UpperCAmelCase					:			List[str]									=						original_model(
																					    {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']}    ).logits
																					UpperCAmelCase					:			List[Any]									=						input_ids.masked_fill(input_ids == tokenizer.pad_token_id  , -100    )
																					UpperCAmelCase					:			List[Any]									=						hf_model(lowerCAmelCase__  , lowerCAmelCase__  , labels=lowerCAmelCase__    ).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":
														UpperCAmelCase					:			Tuple									=						torch.tensor(
														    [[-4_1.5_8_5_0, -4.4_4_4_0, -8.9_9_2_2], [-4_7.4_3_2_2, -5.9_1_4_3, -1.7_3_4_0]]  , device=lowerCAmelCase__    )
														assert torch.allclose(logits[0, :3, :3]  , lowerCAmelCase__  , atol=1e-4    )
							elif model_name == "blip2-flan-t5-xl-coco":
														UpperCAmelCase					:			int									=						torch.tensor(
														    [[-5_7.0_1_0_9, -9.8_9_6_7, -1_2.6_2_8_0], [-6_8.6_5_7_8, -1_2.7_1_9_1, -1_0.5_0_6_5]]  , device=lowerCAmelCase__    )
							else:
														# cast to same type
														UpperCAmelCase					:			Union[str, Any]									=						logits.dtype
														assert torch.allclose(original_logits.to(lowerCAmelCase__    )  , lowerCAmelCase__  , atol=1e-2    )
							print('''Looks ok!'''    )

							print('''Generating a caption...'''    )
							UpperCAmelCase					:			Tuple									=						""""""
							UpperCAmelCase					:			int									=						tokenizer(lowerCAmelCase__  , return_tensors='''pt'''    ).input_ids.to(lowerCAmelCase__    )

							UpperCAmelCase					:			Tuple									=						original_model.generate({'''image''': original_pixel_values}    )
							UpperCAmelCase					:			Tuple									=						hf_model.generate(
							    lowerCAmelCase__  , lowerCAmelCase__  , do_sample=lowerCAmelCase__  , num_beams=5  , max_length=30  , min_length=1  , top_p=0.9  , repetition_penalty=1.0  , length_penalty=1.0  , temperature=1  , )
							print('''Original generation:'''  , lowerCAmelCase__    )
							UpperCAmelCase					:			str									=						input_ids.shape[1]
							UpperCAmelCase					:			int									=						processor.batch_decode(outputs[:, prompt_length:]  , skip_special_tokens=lowerCAmelCase__    )
							UpperCAmelCase					:			Dict									=						[text.strip() for text in output_text]
							print('''HF generation:'''  , lowerCAmelCase__    )

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

							if push_to_hub:
														processor.push_to_hub(f"""nielsr/{model_name}"""    )
														hf_model.push_to_hub(f"""nielsr/{model_name}"""    )


if __name__ == "__main__":
						UpperCamelCase__:  Union[str, Any]				      =     argparse.ArgumentParser()
						UpperCamelCase__:  int				      =     [
						    '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",
						)

						UpperCamelCase__:  int				      =     parser.parse_args()

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










import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel

from ...utils import logging


a      :  Union[str, Any]			     =			logging.get_logger(__name__)







class 				_a				(   _lowerCAmelCase					):
			A      		=	CLIPConfig

			A      		=	['''CLIPEncoderLayer''']

			def __init__(self,   SCREAMING_SNAKE_CASE_					)       ->   Optional[Any]:
						super().__init__(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Tuple							      =     CLIPVisionModelWithProjection(config.vision_config					)

						UpperCAmelCase_:		Tuple							      =     nn.Linear(config.vision_config.projection_dim,   1					)
						UpperCAmelCase_:		Tuple							      =     nn.Linear(config.vision_config.projection_dim,   1					)

			@torch.no_grad()
			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=0.5,   SCREAMING_SNAKE_CASE_=0.5					)       ->   Tuple:
						UpperCAmelCase_:		Optional[Any]							      =     self.vision_model(SCREAMING_SNAKE_CASE_					)[0]

						UpperCAmelCase_:		List[str]							      =     self.p_head(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Any							      =     nsfw_detected.flatten()
						UpperCAmelCase_:		List[Any]							      =     nsfw_detected > p_threshold
						UpperCAmelCase_:		Any							      =     nsfw_detected.tolist()

						if any(SCREAMING_SNAKE_CASE_					):
									logger.warning(
									    """Potential NSFW content was detected in one or more images. A black image will be returned instead."""
									    """ Try again with a different prompt and/or seed."""					)

						for idx, nsfw_detected_ in enumerate(SCREAMING_SNAKE_CASE_					):
									if nsfw_detected_:
												UpperCAmelCase_:		Tuple							      =     np.zeros(images[idx].shape					)

						UpperCAmelCase_:		str							      =     self.w_head(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		List[Any]							      =     watermark_detected.flatten()
						UpperCAmelCase_:		Tuple							      =     watermark_detected > w_threshold
						UpperCAmelCase_:		int							      =     watermark_detected.tolist()

						if any(SCREAMING_SNAKE_CASE_					):
									logger.warning(
									    """Potential watermarked content was detected in one or more images. A black image will be returned instead."""
									    """ Try again with a different prompt and/or seed."""					)

						for idx, watermark_detected_ in enumerate(SCREAMING_SNAKE_CASE_					):
									if watermark_detected_:
												UpperCAmelCase_:		Any							      =     np.zeros(images[idx].shape					)

						return images, nsfw_detected, watermark_detected






