



'''simple docstring'''



from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import BaseOutput, is_torch_available, is_transformers_available





@dataclass
class _UpperCamelCase		(  __a		):





				'''simple docstring'''





				lowerCAmelCase__            =     42
				lowerCAmelCase__            =     42


if is_transformers_available() and is_torch_available():
		from .pipeline_semantic_stable_diffusion import SemanticStableDiffusionPipeline





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__							(							):




			"""simple docstring"""


			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    :    Optional[int]							):




			"""simple docstring"""


			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    :    Optional[Any]       ,				__lowerCAmelCase    :    str       ,				__lowerCAmelCase    :    Union[str, Any]							):




			"""simple docstring"""


			lowerCAmelCase_         =      dct.pop(__lowerCAmelCase							)
			lowerCAmelCase_         =      val






def   lowerCamelCase__							(							__lowerCAmelCase    :    Any       ,				__lowerCAmelCase    :    Union[str, Any]							):




			"""simple docstring"""


			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    :    Optional[int]							):




			"""simple docstring"""


			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    :    str       ,				__lowerCAmelCase    :    int=None       ,				__lowerCAmelCase    :    Dict=False							):




			"""simple docstring"""


			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__":
			_A  =	argparse.ArgumentParser()
			_A  =	[
			    "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  =	parser.parse_args()

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



