


"""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       (   )      ->		Any:







		"""simple docstring"""





		A					       = 'https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'
		A					       = Image.open(requests.get(UpperCamelCase__ ,			stream=UpperCamelCase__	).raw	).convert('RGB'	)

		return image





def  __snake_case       (   UpperCamelCase__	)      ->		Any:







		"""simple docstring"""





		A					       = []
		# 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       (   UpperCamelCase__ ,			UpperCamelCase__ ,			UpperCamelCase__	)      ->		Tuple:







		"""simple docstring"""





		A					       = dct.pop(UpperCamelCase__	)
		A					       = val





def  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__	)      ->		Dict:







		"""simple docstring"""





		for i in range(config.vision_config.num_hidden_layers	):
				# read in original q and v biases
				A					       = state_dict.pop(f'visual_encoder.blocks.{i}.attn.q_bias'	)
				A					       = state_dict.pop(f'visual_encoder.blocks.{i}.attn.v_bias'	)

				# next, set bias in the state dict
				A					       = torch.cat((q_bias, torch.zeros_like(UpperCamelCase__ ,			requires_grad=UpperCamelCase__	), v_bias)	)
				A					       = qkv_bias





def  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__	)      ->		str:







		"""simple docstring"""





		A					       = 364 if 'coco' in model_name else 224
		A					       = BlipaVisionConfig(image_size=UpperCamelCase__	).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:
				A					       = OPTConfig.from_pretrained('facebook/opt-2.7b' ,			eos_token_id=UpperCamelCase__	).to_dict()
		elif "opt-6.7b" in model_name:
				A					       = OPTConfig.from_pretrained('facebook/opt-6.7b' ,			eos_token_id=UpperCamelCase__	).to_dict()
		elif "t5-xl" in model_name:
				A					       = TaConfig.from_pretrained('google/flan-t5-xl' ,			dense_act_fn='gelu' ,			bos_token_id=1	).to_dict()
		elif "t5-xxl" in model_name:
				A					       = TaConfig.from_pretrained('google/flan-t5-xxl' ,			dense_act_fn='gelu' ,			bos_token_id=1	).to_dict()

		A					       = BlipaConfig(vision_config=UpperCamelCase__ ,			text_config=UpperCamelCase__	)

		return config, image_size





@torch.no_grad()
def  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__=None ,			UpperCamelCase__=False	)      ->		Tuple:







		"""simple docstring"""





		A					       = (
		    AutoTokenizer.from_pretrained('facebook/opt-2.7b'	)
		    if 'opt' in model_name
		    else AutoTokenizer.from_pretrained('google/flan-t5-xl'	)
		)
		A					       = tokenizer('\n' ,			add_special_tokens=UpperCamelCase__	).input_ids[0]
		A       , A					       = get_blipa_config(UpperCamelCase__ ,			eos_token_id=UpperCamelCase__	)

		A					       = BlipaForConditionalGeneration(UpperCamelCase__	).eval()

		A					       = {
		    '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'),
		}

		A       , A					       = model_name_to_original[model_name]

		# load original model
		print('Loading original model...'	)
		A					       = 'cuda' if torch.cuda.is_available() else 'cpu'
		A       , A       , A					       = load_model_and_preprocess(
		    name=UpperCamelCase__ ,			model_type=UpperCamelCase__ ,			is_eval=UpperCamelCase__ ,			device=UpperCamelCase__	)
		original_model.eval()
		print('Done!'	)

		# update state dict keys
		A					       = original_model.state_dict()
		A					       = create_rename_keys(UpperCamelCase__	)
		for src, dest in rename_keys:
				rename_key(UpperCamelCase__ ,			UpperCamelCase__ ,			UpperCamelCase__	)

		# some keys can be renamed efficiently
		for key, val in state_dict.copy().items():
				A					       = state_dict.pop(UpperCamelCase__	)
				if key.startswith('Qformer.bert'	):
						A					       = key.replace('Qformer.bert' ,			'qformer'	)
				if "attention.self" in key:
						A					       = key.replace('self' ,			'attention'	)
				if "opt_proj" in key:
						A					       = key.replace('opt_proj' ,			'language_projection'	)
				if "t5_proj" in key:
						A					       = key.replace('t5_proj' ,			'language_projection'	)
				if key.startswith('opt'	):
						A					       = key.replace('opt' ,			'language'	)
				if key.startswith('t5'	):
						A					       = key.replace('t5' ,			'language'	)
				A					       = val

		# read in qv biases
		read_in_q_v_bias(UpperCamelCase__ ,			UpperCamelCase__	)

		A       , A					       = hf_model.load_state_dict(UpperCamelCase__ ,			strict=UpperCamelCase__	)
		assert len(UpperCamelCase__	) == 0
		assert unexpected_keys == ["qformer.embeddings.position_ids"]

		A					       = load_demo_image()
		A					       = vis_processors['eval'](UpperCamelCase__	).unsqueeze(0	).to(UpperCamelCase__	)
		A					       = tokenizer(['\n'] ,			return_tensors='pt'	).input_ids.to(UpperCamelCase__	)

		# create processor
		A					       = BlipImageProcessor(
		    size={'height': image_size, 'width': image_size} ,			image_mean=UpperCamelCase__ ,			image_std=UpperCamelCase__	)
		A					       = BlipaProcessor(image_processor=UpperCamelCase__ ,			tokenizer=UpperCamelCase__	)
		A					       = processor(images=UpperCamelCase__ ,			return_tensors='pt'	).pixel_values.to(UpperCamelCase__	)

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

		original_model.to(UpperCamelCase__	)
		hf_model.to(UpperCamelCase__	)
		with torch.no_grad():
				if "opt" in model_name:
						A					       = original_model({'image': original_pixel_values, 'text_input': ['']}	).logits
						A					       = hf_model(UpperCamelCase__ ,			UpperCamelCase__	).logits
				else:
						A					       = original_model(
						    {'image': original_pixel_values, 'text_input': ['\n'], 'text_output': ['\n']}	).logits
						A					       = input_ids.masked_fill(input_ids == tokenizer.pad_token_id ,			-100	)
						A					       = hf_model(UpperCamelCase__ ,			UpperCamelCase__ ,			labels=UpperCamelCase__	).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":
				A					       = torch.tensor(
				    [[-41.5850, -4.4_4_4_0, -8.9_9_2_2], [-47.4322, -5.9_1_4_3, -1.7_3_4_0]] ,			device=UpperCamelCase__	)
				assert torch.allclose(logits[0, :3, :3] ,			UpperCamelCase__ ,			atol=1E-4	)
		elif model_name == "blip2-flan-t5-xl-coco":
				A					       = torch.tensor(
				    [[-57.0109, -9.8_9_6_7, -12.6280], [-68.6578, -12.7191, -10.5065]] ,			device=UpperCamelCase__	)
		else:
				# cast to same type
				A					       = logits.dtype
				assert torch.allclose(original_logits.to(UpperCamelCase__	) ,			UpperCamelCase__ ,			atol=1E-2	)
		print('Looks ok!'	)

		print('Generating a caption...'	)
		A					       = ''
		A					       = tokenizer(UpperCamelCase__ ,			return_tensors='pt'	).input_ids.to(UpperCamelCase__	)

		A					       = original_model.generate({'image': original_pixel_values}	)
		A					       = hf_model.generate(
		    UpperCamelCase__ ,			UpperCamelCase__ ,			do_sample=UpperCamelCase__ ,			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:' ,			UpperCamelCase__	)
		A					       = input_ids.shape[1]
		A					       = processor.batch_decode(outputs[:, prompt_length:] ,			skip_special_tokens=UpperCamelCase__	)
		A					       = [text.strip() for text in output_text]
		print('HF generation:' ,			UpperCamelCase__	)

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

		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						:				Optional[int]             =      argparse.ArgumentParser()
	UpperCamelCase						:				str             =      [
	    "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						:				str             =      parser.parse_args()

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









"""simple docstring"""


import unittest

from huggingface_hub import hf_hub_download

from transformers import MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING, VideoMAEFeatureExtractor
from transformers.pipelines import VideoClassificationPipeline, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_decord,
    require_tf,
    require_torch,
    require_torch_or_tf,
    require_vision,
)

from .test_pipelines_common import ANY



@is_pipeline_test
@require_torch_or_tf
@require_vision
@require_decord
class 			lowerCamelCase__						(   unittest.TestCase			):
						lowerCAmelCase       =     MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING







						def __a						(							self   :    Dict			,			_lowercase   :    int			,			_lowercase   :    Any			,			_lowercase   :    int		):
								A					       = hf_hub_download(
								    repo_id='nateraw/video-demo'			,			filename='archery.mp4'			,			repo_type='dataset'		)
								A					       = VideoClassificationPipeline(model=_lowercase			,			image_processor=_lowercase			,			top_k=2		)
								A					       = [
								    example_video_filepath,
								    'https://huggingface.co/datasets/nateraw/video-demo/resolve/main/archery.mp4',
								]
								return video_classifier, examples







						def __a						(							self   :    Tuple			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    List[Any]		):
								for example in examples:
										A					       = video_classifier(_lowercase		)

										self.assertEqual(
										    _lowercase			,			[
										        {'score': ANY(_lowercase		), 'label': ANY(_lowercase		)},
										        {'score': ANY(_lowercase		), 'label': ANY(_lowercase		)},
										    ]			,			)







						@require_torch
						def __a						(							self   :    str		):
								A					       = 'hf-internal-testing/tiny-random-VideoMAEForVideoClassification'
								A					       = VideoMAEFeatureExtractor(
								    size={'shortest_edge': 10}			,			crop_size={'height': 10, 'width': 10}		)
								A					       = pipeline(
								    'video-classification'			,			model=_lowercase			,			feature_extractor=_lowercase			,			frame_sampling_rate=4		)

								A					       = hf_hub_download(repo_id='nateraw/video-demo'			,			filename='archery.mp4'			,			repo_type='dataset'		)
								A					       = video_classifier(_lowercase			,			top_k=2		)
								self.assertEqual(
								    nested_simplify(_lowercase			,			decimals=4		)			,			[{'score': 0.5_1_9_9, 'label': 'LABEL_0'}, {'score': 0.4_8_0_1, 'label': 'LABEL_1'}]			,			)

								A					       = video_classifier(
								    [
								        video_file_path,
								        video_file_path,
								    ]			,			top_k=2			,			)
								self.assertEqual(
								    nested_simplify(_lowercase			,			decimals=4		)			,			[
								        [{'score': 0.5_1_9_9, 'label': 'LABEL_0'}, {'score': 0.4_8_0_1, 'label': 'LABEL_1'}],
								        [{'score': 0.5_1_9_9, 'label': 'LABEL_0'}, {'score': 0.4_8_0_1, 'label': 'LABEL_1'}],
								    ]			,			)







						@require_tf
						def __a						(							self   :    Dict		):
								pass





