'''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_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  UpperCamelCase__					(					)							->	Any:
		__UpperCAmelCase:					Optional[int]		= """https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"""
		__UpperCAmelCase:					Optional[int]		= Image.open(requests.get(_lowercase   ,    stream=_lowercase     ).raw     ).convert("""RGB"""     )

		return image






def  UpperCamelCase__					(					_lowercase    :				Dict     )							->	Optional[Any]:
		__UpperCAmelCase:					Optional[Any]		= []
		# 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  UpperCamelCase__					(					_lowercase    :				Any   ,    _lowercase    :				Any   ,    _lowercase    :				Dict     )							->	Union[str, Any]:
		__UpperCAmelCase:					List[str]		= dct.pop(_lowercase     )
		__UpperCAmelCase:					Dict		= val






def  UpperCamelCase__					(					_lowercase    :				List[str]   ,    _lowercase    :				Dict     )							->	Optional[int]:
		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:					Any		= state_dict.pop(F'''visual_encoder.blocks.{i}.attn.v_bias'''     )

				# next, set bias in the state dict
				__UpperCAmelCase:					Optional[int]		= torch.cat((q_bias, torch.zeros_like(_lowercase   ,    requires_grad=_lowercase     ), v_bias)     )
				__UpperCAmelCase:					Tuple		= qkv_bias






def  UpperCamelCase__					(					_lowercase    :				str     )							->	List[Any]:
		__UpperCAmelCase:					List[Any]		= 3_6_4 if """coco""" in model_name else 2_2_4
		__UpperCAmelCase:					List[str]		= InstructBlipVisionConfig(image_size=_lowercase     ).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:					Union[str, Any]		= TaConfig.from_pretrained("""google/flan-t5-xl"""   ,    dense_act_fn="""gelu"""   ,    bos_token_id=1     ).to_dict()
		elif "t5-xxl" in model_name:
				__UpperCAmelCase:					str		= TaConfig.from_pretrained("""google/flan-t5-xxl"""   ,    dense_act_fn="""gelu"""   ,    bos_token_id=1     ).to_dict()
		elif "vicuna-7b" in model_name:
				__UpperCAmelCase:					Optional[int]		= LlamaConfig.from_pretrained("""decapoda-research/llama-7b-hf"""   ,    vocab_size=3_2_0_0_1     ).to_dict()
		elif "vicuna-13b" in model_name:
				__UpperCAmelCase:					Optional[Any]		= LlamaConfig.from_pretrained("""decapoda-research/llama-13b-hf"""   ,    vocab_size=3_2_0_0_1     ).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:					List[str]		= InstructBlipQFormerConfig(vocab_size=3_0_5_2_3     ).to_dict()
		__UpperCAmelCase:					Dict		= InstructBlipConfig(vision_config=_lowercase   ,    text_config=_lowercase   ,    qformer_config=_lowercase     )

		return config, image_size






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

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

		__UpperCAmelCase:					Optional[Any]		= get_blipa_config(_lowercase     )
		__UpperCAmelCase:					Optional[int]		= InstructBlipForConditionalGeneration(_lowercase     ).eval()

		__UpperCAmelCase:					int		= {
		    """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:					Any		= model_name_to_original[model_name]

		# load original model
		print("""Loading original model..."""     )
		__UpperCAmelCase:					List[Any]		= """cuda:1""" if torch.cuda.is_available() else """cpu"""
		__UpperCAmelCase:					Tuple		= """cuda:2""" if torch.cuda.is_available() else """cpu"""
		__UpperCAmelCase:					str		= load_model_and_preprocess(
		    name=_lowercase   ,    model_type=_lowercase   ,    is_eval=_lowercase   ,    device=_lowercase     )
		original_model.eval()
		print("""Done!"""     )

		# update state dict keys
		__UpperCAmelCase:					Optional[int]		= original_model.state_dict()
		__UpperCAmelCase:					Optional[Any]		= create_rename_keys(_lowercase     )
		for src, dest in rename_keys:
				rename_key(_lowercase   ,    _lowercase   ,    _lowercase     )

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

		# read in qv biases
		read_in_q_v_bias(_lowercase   ,    _lowercase     )

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

		__UpperCAmelCase:					List[str]		= load_demo_image()
		__UpperCAmelCase:					Optional[Any]		= """What is unusual about this image?"""

		# create processor
		__UpperCAmelCase:					Union[str, Any]		= BlipImageProcessor(
		    size={"""height""": image_size, """width""": image_size}   ,    image_mean=_lowercase   ,    image_std=_lowercase     )
		__UpperCAmelCase:					Any		= InstructBlipProcessor(
		    image_processor=_lowercase   ,    tokenizer=_lowercase   ,    qformer_tokenizer=_lowercase   ,    )
		__UpperCAmelCase:					List[str]		= processor(images=_lowercase   ,    text=_lowercase   ,    return_tensors="""pt"""     ).to(_lowercase     )

		# make sure processor creates exact same pixel values
		__UpperCAmelCase:					str		= vis_processors["""eval"""](_lowercase     ).unsqueeze(0     ).to(_lowercase     )
		__UpperCAmelCase:					List[Any]		= inputs.pixel_values
		assert torch.allclose(original_pixel_values.to(pixel_values.device     )   ,    _lowercase     )

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

		print("""Generating with original model..."""     )
		__UpperCAmelCase:					str		= 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:					Tuple		= hf_model.generate(
		    **_lowercase   ,    do_sample=_lowercase   ,    num_beams=5   ,    max_length=2_5_6   ,    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:					str		= 2
		print("""Original generation:"""   ,    _lowercase     )
		__UpperCAmelCase:					List[str]		= processor.batch_decode(_lowercase   ,    skip_special_tokens=_lowercase     )
		__UpperCAmelCase:					Dict		= [text.strip() for text in output_text]
		print("""HF generation:"""   ,    _lowercase     )

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

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


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

				SCREAMING_SNAKE_CASE_	   =			parser.parse_args()

				convert_blipa_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE_	   =			{
    'configuration_electra': ['ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP', 'ElectraConfig', 'ElectraOnnxConfig'],
    'tokenization_electra': ['ElectraTokenizer'],
}

try:
    if not is_tokenizers_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			['ElectraTokenizerFast']

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			[
        'ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST',
        'ElectraForCausalLM',
        'ElectraForMaskedLM',
        'ElectraForMultipleChoice',
        'ElectraForPreTraining',
        'ElectraForQuestionAnswering',
        'ElectraForSequenceClassification',
        'ElectraForTokenClassification',
        'ElectraModel',
        'ElectraPreTrainedModel',
        'load_tf_weights_in_electra',
    ]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			[
        'TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST',
        'TFElectraForMaskedLM',
        'TFElectraForMultipleChoice',
        'TFElectraForPreTraining',
        'TFElectraForQuestionAnswering',
        'TFElectraForSequenceClassification',
        'TFElectraForTokenClassification',
        'TFElectraModel',
        'TFElectraPreTrainedModel',
    ]

try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			[
        'FlaxElectraForCausalLM',
        'FlaxElectraForMaskedLM',
        'FlaxElectraForMultipleChoice',
        'FlaxElectraForPreTraining',
        'FlaxElectraForQuestionAnswering',
        'FlaxElectraForSequenceClassification',
        'FlaxElectraForTokenClassification',
        'FlaxElectraModel',
        'FlaxElectraPreTrainedModel',
    ]


if TYPE_CHECKING:
    from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig, ElectraOnnxConfig
    from .tokenization_electra import ElectraTokenizer

    try:
        if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_electra_fast import ElectraTokenizerFast

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_electra import (
            ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
            ElectraForCausalLM,
            ElectraForMaskedLM,
            ElectraForMultipleChoice,
            ElectraForPreTraining,
            ElectraForQuestionAnswering,
            ElectraForSequenceClassification,
            ElectraForTokenClassification,
            ElectraModel,
            ElectraPreTrainedModel,
            load_tf_weights_in_electra,
        )

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_electra import (
            TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
            TFElectraForMaskedLM,
            TFElectraForMultipleChoice,
            TFElectraForPreTraining,
            TFElectraForQuestionAnswering,
            TFElectraForSequenceClassification,
            TFElectraForTokenClassification,
            TFElectraModel,
            TFElectraPreTrainedModel,
        )

    try:
        if not is_flax_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_flax_electra import (
            FlaxElectraForCausalLM,
            FlaxElectraForMaskedLM,
            FlaxElectraForMultipleChoice,
            FlaxElectraForPreTraining,
            FlaxElectraForQuestionAnswering,
            FlaxElectraForSequenceClassification,
            FlaxElectraForTokenClassification,
            FlaxElectraModel,
            FlaxElectraPreTrainedModel,
        )

else:
    import sys

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