





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_float32 (there's also the fix_lavis branch)
# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
# same for Vicuna-13b
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipImageProcessor,
    InstructBlipConfig,
    InstructBlipForConditionalGeneration,
    InstructBlipProcessor,
    InstructBlipQFormerConfig,
    InstructBlipVisionConfig,
    LlamaConfig,
    LlamaTokenizerFast,
    TaConfig,
    TaTokenizerFast,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD





def 		__SCREAMING_SNAKE_CASE    (							)    ->							str:







			'''simple docstring'''

			UpperCAmelCase  						=							'''https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg'''
			UpperCAmelCase  						=							Image.open(requests.get(SCREAMING_SNAKE_CASE_     ,     stream=SCREAMING_SNAKE_CASE_					).raw					).convert('''RGB'''					)

			return image





def 		__SCREAMING_SNAKE_CASE    (							UpperCamelCase__					)    ->							Optional[Any]:







			'''simple docstring'''

			UpperCAmelCase  						=							[]
			# 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.embeddings.layernorm.weight''')					)
			rename_keys.append(('''Qformer.bert.embeddings.LayerNorm.bias''', '''qformer.embeddings.layernorm.bias''')					)

			# fmt: on
			return rename_keys





def 		__SCREAMING_SNAKE_CASE    (							UpperCamelCase__     ,     UpperCamelCase__     ,     UpperCamelCase__					)    ->							List[Any]:







			'''simple docstring'''

			UpperCAmelCase  						=							dct.pop(SCREAMING_SNAKE_CASE_					)
			UpperCAmelCase  						=							val





def 		__SCREAMING_SNAKE_CASE    (							UpperCamelCase__     ,     UpperCamelCase__					)    ->							Union[str, Any]:







			'''simple docstring'''

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

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





def 		__SCREAMING_SNAKE_CASE    (							UpperCamelCase__					)    ->							int:







			'''simple docstring'''

			UpperCAmelCase  						=							364 if '''coco''' in model_name else 224
			UpperCAmelCase  						=							InstructBlipVisionConfig(image_size=SCREAMING_SNAKE_CASE_					).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 "t5-xl" in model_name:
						UpperCAmelCase  						=							TaConfig.from_pretrained('''google/flan-t5-xl'''     ,     dense_act_fn='''gelu'''     ,     bos_token_id=1					).to_dict()
			elif "t5-xxl" in model_name:
						UpperCAmelCase  						=							TaConfig.from_pretrained('''google/flan-t5-xxl'''     ,     dense_act_fn='''gelu'''     ,     bos_token_id=1					).to_dict()
			elif "vicuna-7b" in model_name:
						UpperCAmelCase  						=							LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''     ,     vocab_size=3_2001					).to_dict()
			elif "vicuna-13b" in model_name:
						UpperCAmelCase  						=							LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf'''     ,     vocab_size=3_2001					).to_dict()
			else:
						raise ValueError('''Model name not supported'''					)

			# the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
			UpperCAmelCase  						=							InstructBlipQFormerConfig(vocab_size=3_0523					).to_dict()
			UpperCAmelCase  						=							InstructBlipConfig(vision_config=SCREAMING_SNAKE_CASE_     ,     text_config=SCREAMING_SNAKE_CASE_     ,     qformer_config=SCREAMING_SNAKE_CASE_					)

			return config, image_size





@torch.no_grad()
def 		__SCREAMING_SNAKE_CASE    (							UpperCamelCase__     ,     UpperCamelCase__=None     ,     UpperCamelCase__=False					)    ->							int:







			'''simple docstring'''

			UpperCAmelCase  						=							AutoTokenizer.from_pretrained('''bert-base-uncased'''     ,     truncation_side='''left'''					)
			qformer_tokenizer.add_special_tokens({'''bos_token''': '''[DEC]'''}					)

			if "t5" in model_name:
						UpperCAmelCase  						=							TaTokenizerFast.from_pretrained('''google/flan-t5-xl'''     ,     truncation_side='''left'''					)
			elif "vicuna" in model_name:
						# the following was used in the original implementation:
						# tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
						# tokenizer.add_special_tokens({"pad_token": "[PAD]"})
						# tokenizer.add_special_tokens({"bos_token": "</s>"})
						# tokenizer.add_special_tokens({"eos_token": "</s>"})
						# tokenizer.add_special_tokens({"unk_token": "</s>"})
						UpperCAmelCase  						=							LlamaTokenizerFast.from_pretrained(
						    '''huggyllama/llama-7b'''     ,     truncation_side='''left'''     ,     bos_token='''</s>'''     ,     unk_token='''</s>'''					)
						tokenizer.add_special_tokens({'''pad_token''': '''[PAD]'''}					)

			UpperCAmelCase						,     UpperCAmelCase  						=							get_blipa_config(SCREAMING_SNAKE_CASE_					)
			UpperCAmelCase  						=							InstructBlipForConditionalGeneration(SCREAMING_SNAKE_CASE_					).eval()

			UpperCAmelCase  						=							{
			    '''instructblip-vicuna-7b''': ('''blip2_vicuna_instruct''', '''vicuna7b'''),
			    '''instructblip-vicuna-13b''': ('''blip2_vicuna_instruct''', '''vicuna13b'''),
			    '''instructblip-flan-t5-xl''': ('''blip2_t5_instruct''', '''flant5xl'''),
			    '''instructblip-flan-t5-xxl''': ('''blip2_t5_instruct''', '''flant5xxl'''),
			}

			UpperCAmelCase						,     UpperCAmelCase  						=							model_name_to_original[model_name]

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

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

			# some keys can be renamed efficiently
			for key, val in state_dict.copy().items():
						UpperCAmelCase  						=							state_dict.pop(SCREAMING_SNAKE_CASE_					)
						if key.startswith('''Qformer.bert'''					):
									UpperCAmelCase  						=							key.replace('''Qformer.bert'''     ,     '''qformer'''					)
						if "attention.self" in key:
									UpperCAmelCase  						=							key.replace('''self'''     ,     '''attention'''					)
						if "llm_proj" in key:
									UpperCAmelCase  						=							key.replace('''llm_proj'''     ,     '''language_projection'''					)
						if "t5_proj" in key:
									UpperCAmelCase  						=							key.replace('''t5_proj'''     ,     '''language_projection'''					)
						if key.startswith('''llm_model'''					):
									UpperCAmelCase  						=							key.replace('''llm_model'''     ,     '''language_model'''					)
						if key.startswith('''t5'''					):
									UpperCAmelCase  						=							key.replace('''t5'''     ,     '''language'''					)
						UpperCAmelCase  						=							val

			# read in qv biases
			read_in_q_v_bias(SCREAMING_SNAKE_CASE_     ,     SCREAMING_SNAKE_CASE_					)

			# note: weights get loaded in torch.float32 by default
			hf_model.load_state_dict(SCREAMING_SNAKE_CASE_     ,     strict=SCREAMING_SNAKE_CASE_					)

			UpperCAmelCase  						=							load_demo_image()
			UpperCAmelCase  						=							'''What is unusual about this image?'''

			# create processor
			UpperCAmelCase  						=							BlipImageProcessor(
			    size={'''height''': image_size, '''width''': image_size}     ,     image_mean=SCREAMING_SNAKE_CASE_     ,     image_std=SCREAMING_SNAKE_CASE_					)
			UpperCAmelCase  						=							InstructBlipProcessor(
			    image_processor=SCREAMING_SNAKE_CASE_     ,     tokenizer=SCREAMING_SNAKE_CASE_     ,     qformer_tokenizer=SCREAMING_SNAKE_CASE_     ,     )
			UpperCAmelCase  						=							processor(images=SCREAMING_SNAKE_CASE_     ,     text=SCREAMING_SNAKE_CASE_     ,     return_tensors='''pt'''					).to(SCREAMING_SNAKE_CASE_					)

			# make sure processor creates exact same pixel values
			UpperCAmelCase  						=							vis_processors['''eval'''](SCREAMING_SNAKE_CASE_					).unsqueeze(0					).to(SCREAMING_SNAKE_CASE_					)
			UpperCAmelCase  						=							inputs.pixel_values
			assert torch.allclose(original_pixel_values.to(pixel_values.device					)     ,     SCREAMING_SNAKE_CASE_					)

			original_model.to(SCREAMING_SNAKE_CASE_					)
			hf_model.to(SCREAMING_SNAKE_CASE_					)
			with torch.no_grad():
						if "vicuna" in model_name:
									UpperCAmelCase  						=							original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}					).logits
									UpperCAmelCase  						=							hf_model(**SCREAMING_SNAKE_CASE_					).logits
						else:
									UpperCAmelCase  						=							original_model(
									    {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}					).logits
									UpperCAmelCase  						=							tokenizer('''\n'''     ,     return_tensors='''pt'''					).input_ids.to(SCREAMING_SNAKE_CASE_					)
									UpperCAmelCase  						=							label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id     ,     -100					)
									UpperCAmelCase  						=							hf_model(**SCREAMING_SNAKE_CASE_     ,     labels=SCREAMING_SNAKE_CASE_					).logits

			print('''First values of original logits:'''     ,     original_logits[0, :3, :3]					)
			print('''First values of HF logits:'''     ,     logits[0, :3, :3]					)

			# assert values
			assert original_logits.shape == logits.shape
			UpperCAmelCase  						=							1E-4 if '''vicuna''' in model_name else 1E-5
			assert torch.allclose(original_logits.to(logits.device					)     ,     SCREAMING_SNAKE_CASE_     ,     atol=SCREAMING_SNAKE_CASE_					)
			print('''Looks ok!'''					)

			print('''Generating with original model...'''					)
			UpperCAmelCase  						=							original_model.generate({'''image''': original_pixel_values, '''prompt''': prompt}     ,     num_beams=5					)

			# important: we need to cast the weights of the HF model to the appropriate type
			print('''Generating with HF model...'''					)
			UpperCAmelCase  						=							hf_model.generate(
			    **SCREAMING_SNAKE_CASE_     ,     do_sample=SCREAMING_SNAKE_CASE_     ,     num_beams=5     ,     max_length=256     ,     min_length=1     ,     top_p=0.9     ,     repetition_penalty=1.5     ,     length_penalty=1.0     ,     temperature=1     ,     )
			if "vicuna" in model_name:
						# convert output id 0 to 2 (eos_token_id)
						# TODO add this in the generate method?
						UpperCAmelCase  						=							2
			print('''Original generation:'''     ,     SCREAMING_SNAKE_CASE_					)
			UpperCAmelCase  						=							processor.batch_decode(SCREAMING_SNAKE_CASE_     ,     skip_special_tokens=SCREAMING_SNAKE_CASE_					)
			UpperCAmelCase  						=							[text.strip() for text in output_text]
			print('''HF generation:'''     ,     SCREAMING_SNAKE_CASE_					)

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

			if push_to_hub:
						processor.push_to_hub(F"""Salesforce/{model_name}"""					)
						hf_model.push_to_hub(F"""Salesforce/{model_name}"""					)


if __name__ == "__main__":
				__A							:       List[Any]         =  argparse.ArgumentParser()
				__A							:       Dict         =  [
				    "instructblip-vicuna-7b",
				    "instructblip-vicuna-13b",
				    "instructblip-flan-t5-xl",
				    "instructblip-flan-t5-xxl",
				]
				parser.add_argument(
				    "--model_name",
				    default="instructblip-flan-t5-xl",
				    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",
				)

				__A							:       Dict         =  parser.parse_args()

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










import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 			__UpperCamelCase    (							A__	, unittest.TestCase    ):
     __A	:  Any							  =   DanceDiffusionPipeline
     __A	:  Any							  =   UNCONDITIONAL_AUDIO_GENERATION_PARAMS
     __A	:  Tuple							  =   PipelineTesterMixin.required_optional_params - {
         """callback""",
         """latents""",
         """callback_steps""",
         """output_type""",
         """num_images_per_prompt""",
     }
     __A	:  Tuple							  =   UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
     __A	:  List[str]							  =   False
     __A	:  str							  =   False





     def 		UpperCamelCase(     self					):
        torch.manual_seed(0					)
        _UpperCAmelCase           =   UNetaDModel(
            block_out_channels=(32, 32, 64)				,							extra_in_channels=16				,							sample_size=512				,							sample_rate=16000				,							in_channels=2				,							out_channels=2				,							flip_sin_to_cos=_UpperCamelCase				,							use_timestep_embedding=_UpperCamelCase				,							time_embedding_type='''fourier'''				,							mid_block_type='''UNetMidBlock1D'''				,							down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')				,							up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')				,							)
        _UpperCAmelCase           =   IPNDMScheduler()

        _UpperCAmelCase           =   {
            '''unet''': unet,
            '''scheduler''': scheduler,
        }
        return components





     def 		UpperCamelCase(     self				,							_UpperCamelCase				,							_UpperCamelCase=0					):
        if str(_UpperCamelCase					).startswith('''mps'''					):
           _UpperCAmelCase           =   torch.manual_seed(_UpperCamelCase					)
        else:
           _UpperCAmelCase           =   torch.Generator(device=_UpperCamelCase					).manual_seed(_UpperCamelCase					)
        _UpperCAmelCase           =   {
            '''batch_size''': 1,
            '''generator''': generator,
            '''num_inference_steps''': 4,
        }
        return inputs





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _UpperCAmelCase           =   self.get_dummy_components()
        _UpperCAmelCase           =   DanceDiffusionPipeline(**_UpperCamelCase					)
        _UpperCAmelCase           =   pipe.to(_UpperCamelCase					)
        pipe.set_progress_bar_config(disable=_UpperCamelCase					)

        _UpperCAmelCase           =   self.get_dummy_inputs(_UpperCamelCase					)
        _UpperCAmelCase           =   pipe(**_UpperCamelCase					)
        _UpperCAmelCase           =   output.audios

        _UpperCAmelCase           =   audio[0, -3:, -3:]

        assert audio.shape == (1, 2, components["unet"].sample_size)
        _UpperCAmelCase           =   np.array([-0.7265, 1.0000, -0.8388, 0.1175, 0.9498, -1.0000]					)
        assert np.abs(audio_slice.flatten() - expected_slice					).max() < 1e-2





     @skip_mps
     def 		UpperCamelCase(     self					):
        return super().test_save_load_local()





     @skip_mps
     def 		UpperCamelCase(     self					):
        return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3					)





     @skip_mps
     def 		UpperCamelCase(     self					):
        return super().test_save_load_optional_components()





     @skip_mps
     def 		UpperCamelCase(     self					):
        return super().test_attention_slicing_forward_pass()





     def 		UpperCamelCase(     self					):
        super().test_inference_batch_single_identical(expected_max_diff=3e-3					)




@slow
@require_torch_gpu
class 			__UpperCamelCase    (							unittest.TestCase    ):





     def 		UpperCamelCase(     self					):
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   torch_device

        _UpperCAmelCase           =   DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''					)
        _UpperCAmelCase           =   pipe.to(_UpperCamelCase					)
        pipe.set_progress_bar_config(disable=_UpperCamelCase					)

        _UpperCAmelCase           =   torch.manual_seed(0					)
        _UpperCAmelCase           =   pipe(generator=_UpperCamelCase				,							num_inference_steps=100				,							audio_length_in_s=4.096					)
        _UpperCAmelCase           =   output.audios

        _UpperCAmelCase           =   audio[0, -3:, -3:]

        assert audio.shape == (1, 2, pipe.unet.sample_size)
        _UpperCAmelCase           =   np.array([-0.0192, -0.0231, -0.0318, -0.0059, 0.0002, -0.0020]					)

        assert np.abs(audio_slice.flatten() - expected_slice					).max() < 1e-2





     def 		UpperCamelCase(     self					):
        _UpperCAmelCase           =   torch_device

        _UpperCAmelCase           =   DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''				,							torch_dtype=torch.floataa					)
        _UpperCAmelCase           =   pipe.to(_UpperCamelCase					)
        pipe.set_progress_bar_config(disable=_UpperCamelCase					)

        _UpperCAmelCase           =   torch.manual_seed(0					)
        _UpperCAmelCase           =   pipe(generator=_UpperCamelCase				,							num_inference_steps=100				,							audio_length_in_s=4.096					)
        _UpperCAmelCase           =   output.audios

        _UpperCAmelCase           =   audio[0, -3:, -3:]

        assert audio.shape == (1, 2, pipe.unet.sample_size)
        _UpperCAmelCase           =   np.array([-0.0367, -0.0488, -0.0771, -0.0525, -0.0444, -0.0341]					)

        assert np.abs(audio_slice.flatten() - expected_slice					).max() < 1e-2