



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 					__lowerCAmelCase		(  )  ->						str:
					__lowerCAmelCase	   ="""https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"""
					__lowerCAmelCase	   =Image.open(requests.get(__lowerCamelCase				,   stream=__lowerCamelCase	).raw	).convert("""RGB"""	)

					return image






def 					__lowerCAmelCase		(  __lowerCamelCase	:    int	)  ->						Any:
					__lowerCAmelCase	   =[]
					# 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 					__lowerCAmelCase		(  __lowerCamelCase	:    Dict				,   __lowerCamelCase	:    Optional[Any]				,   __lowerCamelCase	:    Any	)  ->						Optional[Any]:
					__lowerCAmelCase	   =dct.pop(__lowerCamelCase	)
					__lowerCAmelCase	   =val






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

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






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Union[str, Any]	)  ->						int:
					__lowerCAmelCase	   =364 if """coco""" in model_name else 224
					__lowerCAmelCase	   =InstructBlipVisionConfig(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 "t5-xl" in model_name:
										__lowerCAmelCase	   =TaConfig.from_pretrained("""google/flan-t5-xl"""				,   dense_act_fn="""gelu"""				,   bos_token_id=1	).to_dict()
					elif "t5-xxl" in model_name:
										__lowerCAmelCase	   =TaConfig.from_pretrained("""google/flan-t5-xxl"""				,   dense_act_fn="""gelu"""				,   bos_token_id=1	).to_dict()
					elif "vicuna-7b" in model_name:
										__lowerCAmelCase	   =LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf"""				,   vocab_size=32001	).to_dict()
					elif "vicuna-13b" in model_name:
										__lowerCAmelCase	   =LlamaConfig.from_pretrained("""decapoda-research/llama-13b-hf"""				,   vocab_size=32001	).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
					__lowerCAmelCase	   =InstructBlipQFormerConfig(vocab_size=30523	).to_dict()
					__lowerCAmelCase	   =InstructBlipConfig(vision_config=__lowerCamelCase				,   text_config=__lowerCamelCase				,   qformer_config=__lowerCamelCase	)

					return config, image_size






@torch.no_grad()
def 					__lowerCAmelCase		(  __lowerCamelCase	:    Dict				,   __lowerCamelCase	:    Union[str, Any]=None				,   __lowerCamelCase	:    List[str]=False	)  ->						int:
					__lowerCAmelCase	   =AutoTokenizer.from_pretrained("""bert-base-uncased"""				,   truncation_side="""left"""	)
					qformer_tokenizer.add_special_tokens({"""bos_token""": """[DEC]"""}	)

					if "t5" in model_name:
										__lowerCAmelCase	   =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>"})
										__lowerCAmelCase	   =LlamaTokenizerFast.from_pretrained(
										    """huggyllama/llama-7b"""				,   truncation_side="""left"""				,   bos_token="""</s>"""				,   unk_token="""</s>"""	)
										tokenizer.add_special_tokens({"""pad_token""": """[PAD]"""}	)

					__lowerCAmelCase   ,	__lowerCAmelCase	   =get_blipa_config(__lowerCamelCase	)
					__lowerCAmelCase	   =InstructBlipForConditionalGeneration(__lowerCamelCase	).eval()

					__lowerCAmelCase	   ={
					    """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"""),
					}

					__lowerCAmelCase   ,	__lowerCAmelCase	   =model_name_to_original[model_name]

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

					# update state dict keys
					__lowerCAmelCase	   =original_model.state_dict()
					__lowerCAmelCase	   =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():
										__lowerCAmelCase	   =state_dict.pop(__lowerCamelCase	)
										if key.startswith("""Qformer.bert"""	):
															__lowerCAmelCase	   =key.replace("""Qformer.bert"""				,   """qformer"""	)
										if "attention.self" in key:
															__lowerCAmelCase	   =key.replace("""self"""				,   """attention"""	)
										if "llm_proj" in key:
															__lowerCAmelCase	   =key.replace("""llm_proj"""				,   """language_projection"""	)
										if "t5_proj" in key:
															__lowerCAmelCase	   =key.replace("""t5_proj"""				,   """language_projection"""	)
										if key.startswith("""llm_model"""	):
															__lowerCAmelCase	   =key.replace("""llm_model"""				,   """language_model"""	)
										if key.startswith("""t5"""	):
															__lowerCAmelCase	   =key.replace("""t5"""				,   """language"""	)
										__lowerCAmelCase	   =val

					# read in qv biases
					read_in_q_v_bias(__lowerCamelCase				,   __lowerCamelCase	)

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

					__lowerCAmelCase	   =load_demo_image()
					__lowerCAmelCase	   ="""What is unusual about this image?"""

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

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

					original_model.to(__lowerCamelCase	)
					hf_model.to(__lowerCamelCase	)
					with torch.no_grad():
										if "vicuna" in model_name:
															__lowerCAmelCase	   =original_model({"""image""": original_pixel_values, """text_input""": [prompt]}	).logits
															__lowerCAmelCase	   =hf_model(**__lowerCamelCase	).logits
										else:
															__lowerCAmelCase	   =original_model(
															    {"""image""": original_pixel_values, """text_input""": [prompt], """text_output""": ["""\n"""]}	).logits
															__lowerCAmelCase	   =tokenizer("""\n"""				,   return_tensors="""pt"""	).input_ids.to(__lowerCamelCase	)
															__lowerCAmelCase	   =label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id				,   -100	)
															__lowerCAmelCase	   =hf_model(**__lowerCamelCase				,   labels=__lowerCamelCase	).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
					__lowerCAmelCase	   =1E-4 if """vicuna""" in model_name else 1E-5
					assert torch.allclose(original_logits.to(logits.device	)				,   __lowerCamelCase				,   atol=__lowerCamelCase	)
					print("""Looks ok!"""	)

					print("""Generating with original model..."""	)
					__lowerCAmelCase	   =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..."""	)
					__lowerCAmelCase	   =hf_model.generate(
					    **__lowerCamelCase				,   do_sample=__lowerCamelCase				,   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?
										__lowerCAmelCase	   =2
					print("""Original generation:"""				,   __lowerCamelCase	)
					__lowerCAmelCase	   =processor.batch_decode(__lowerCamelCase				,   skip_special_tokens=__lowerCamelCase	)
					__lowerCAmelCase	   =[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"""Salesforce/{model_name}"""	)
										hf_model.push_to_hub(f"""Salesforce/{model_name}"""	)


if __name__ == "__main__":
				lowercase_	       =			argparse.ArgumentParser()
				lowercase_	       =			[
				    '''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''',
				)

				lowercase_	       =			parser.parse_args()

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









from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


lowercase_	       =			{
    '''configuration_encodec''': [
        '''ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''EncodecConfig''',
    ],
    '''feature_extraction_encodec''': ['''EncodecFeatureExtractor'''],
}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				lowercase_	       =			[
				    '''ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST''',
				    '''EncodecModel''',
				    '''EncodecPreTrainedModel''',
				]

if TYPE_CHECKING:
				from .configuration_encodec import (
				    ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP,
				    EncodecConfig,
				)
				from .feature_extraction_encodec import EncodecFeatureExtractor

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_encodec import (
								    ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST,
								    EncodecModel,
								    EncodecPreTrainedModel,
								)

else:
				import sys

				lowercase_	       =			_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




