






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






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

				return image






def     a     (		lowerCamelCase_							):






				'''simple docstring'''
				lowercase__  					=     []
				# 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     a     (		lowerCamelCase_						,  lowerCamelCase_						,  lowerCamelCase_							):






				'''simple docstring'''
				lowercase__  					=     dct.pop(lowerCamelCase_							)
				lowercase__  					=     val






def     a     (		lowerCamelCase_						,  lowerCamelCase_							):






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

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






def     a     (		lowerCamelCase_							):






				'''simple docstring'''
				lowercase__  					=     364 if '''coco''' in model_name else 224
				lowercase__  					=     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:
								lowercase__  					=     TaConfig.from_pretrained('''google/flan-t5-xl'''						,  dense_act_fn='''gelu'''						,  bos_token_id=1							).to_dict()
				elif "t5-xxl" in model_name:
								lowercase__  					=     TaConfig.from_pretrained('''google/flan-t5-xxl'''						,  dense_act_fn='''gelu'''						,  bos_token_id=1							).to_dict()
				elif "vicuna-7b" in model_name:
								lowercase__  					=     LlamaConfig.from_pretrained('''decapoda-research/llama-7b-hf'''						,  vocab_size=3_2001							).to_dict()
				elif "vicuna-13b" in model_name:
								lowercase__  					=     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
				lowercase__  					=     InstructBlipQFormerConfig(vocab_size=3_0523							).to_dict()
				lowercase__  					=     InstructBlipConfig(vision_config=lowerCamelCase_						,  text_config=lowerCamelCase_						,  qformer_config=lowerCamelCase_							)

				return config, image_size






@torch.no_grad()
def     a     (		lowerCamelCase_						,  lowerCamelCase_=None						,  lowerCamelCase_=False							):






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

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

				lowercase__					,		lowercase__  					=     get_blipa_config(lowerCamelCase_							)
				lowercase__  					=     InstructBlipForConditionalGeneration(lowerCamelCase_							).eval()

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

				lowercase__					,		lowercase__  					=     model_name_to_original[model_name]

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

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

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

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

				# make sure processor creates exact same pixel values
				lowercase__  					=     vis_processors['''eval'''](lowerCamelCase_							).unsqueeze(0							).to(lowerCamelCase_							)
				lowercase__  					=     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:
												lowercase__  					=     original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}							).logits
												lowercase__  					=     hf_model(**lowerCamelCase_							).logits
								else:
												lowercase__  					=     original_model(
												    {'''image''': original_pixel_values, '''text_input''': [prompt], '''text_output''': ['''\n''']}							).logits
												lowercase__  					=     tokenizer('''\n'''						,  return_tensors='''pt'''							).input_ids.to(lowerCamelCase_							)
												lowercase__  					=     label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id						,  -100							)
												lowercase__  					=     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
				lowercase__  					=     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...'''							)
				lowercase__  					=     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...'''							)
				lowercase__  					=     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?
								lowercase__  					=     2
				print('''Original generation:'''						,  lowerCamelCase_							)
				lowercase__  					=     processor.batch_decode(lowerCamelCase_						,  skip_special_tokens=lowerCamelCase_							)
				lowercase__  					=     [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__      :     Union[str, 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__      :     Optional[int]       =					parser.parse_args()

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












import tempfile

import torch

from diffusers import IPNDMScheduler

from .test_schedulers import SchedulerCommonTest
class  _UpperCAmelCase   (			A__      ):

				"""simple docstring"""




				lowercase__											=			(IPNDMScheduler,)
				lowercase__											=			(("""num_inference_steps""", 50),)




				def 							lowercase__  (	self		:       Union[str, Any],    **lowerCamelCase		:       Tuple    ):

								'''simple docstring'''
								lowercase__  					=     {'''num_train_timesteps''': 1_000}
								config.update(**lowerCamelCase    )
								return config




				def 							lowercase__  (	self		:       Any,    lowerCamelCase		:       Any=0,    **lowerCamelCase		:       List[str]    ):

								'''simple docstring'''
								lowercase__  					=     dict(self.forward_default_kwargs    )
								lowercase__  					=     kwargs.pop('''num_inference_steps''',    lowerCamelCase    )
								lowercase__  					=     self.dummy_sample
								lowercase__  					=     0.1 * sample
								lowercase__  					=     [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

								for scheduler_class in self.scheduler_classes:
												lowercase__  					=     self.get_scheduler_config(**lowerCamelCase    )
												lowercase__  					=     scheduler_class(**lowerCamelCase    )
												scheduler.set_timesteps(lowerCamelCase    )
												# copy over dummy past residuals
												lowercase__  					=     dummy_past_residuals[:]

												if time_step is None:
																lowercase__  					=     scheduler.timesteps[len(scheduler.timesteps    ) // 2]

												with tempfile.TemporaryDirectory() as tmpdirname:
																scheduler.save_config(lowerCamelCase    )
																lowercase__  					=     scheduler_class.from_pretrained(lowerCamelCase    )
																new_scheduler.set_timesteps(lowerCamelCase    )
																# copy over dummy past residuals
																lowercase__  					=     dummy_past_residuals[:]

												lowercase__  					=     scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample
												lowercase__  					=     new_scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample

												assert torch.sum(torch.abs(output - new_output    )    ) < 1E-5, "Scheduler outputs are not identical"

												lowercase__  					=     scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample
												lowercase__  					=     new_scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample

												assert torch.sum(torch.abs(output - new_output    )    ) < 1E-5, "Scheduler outputs are not identical"




				def 							lowercase__  (	self		:       str    ):

								'''simple docstring'''
								pass




				def 							lowercase__  (	self		:       Dict,    lowerCamelCase		:       Optional[int]=0,    **lowerCamelCase		:       Dict    ):

								'''simple docstring'''
								lowercase__  					=     dict(self.forward_default_kwargs    )
								lowercase__  					=     kwargs.pop('''num_inference_steps''',    lowerCamelCase    )
								lowercase__  					=     self.dummy_sample
								lowercase__  					=     0.1 * sample
								lowercase__  					=     [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

								for scheduler_class in self.scheduler_classes:
												lowercase__  					=     self.get_scheduler_config()
												lowercase__  					=     scheduler_class(**lowerCamelCase    )
												scheduler.set_timesteps(lowerCamelCase    )

												# copy over dummy past residuals (must be after setting timesteps)
												lowercase__  					=     dummy_past_residuals[:]

												if time_step is None:
																lowercase__  					=     scheduler.timesteps[len(scheduler.timesteps    ) // 2]

												with tempfile.TemporaryDirectory() as tmpdirname:
																scheduler.save_config(lowerCamelCase    )
																lowercase__  					=     scheduler_class.from_pretrained(lowerCamelCase    )
																# copy over dummy past residuals
																new_scheduler.set_timesteps(lowerCamelCase    )

																# copy over dummy past residual (must be after setting timesteps)
																lowercase__  					=     dummy_past_residuals[:]

												lowercase__  					=     scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample
												lowercase__  					=     new_scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample

												assert torch.sum(torch.abs(output - new_output    )    ) < 1E-5, "Scheduler outputs are not identical"

												lowercase__  					=     scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample
												lowercase__  					=     new_scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample

												assert torch.sum(torch.abs(output - new_output    )    ) < 1E-5, "Scheduler outputs are not identical"




				def 							lowercase__  (	self		:       int,    **lowerCamelCase		:       List[str]    ):

								'''simple docstring'''
								lowercase__  					=     self.scheduler_classes[0]
								lowercase__  					=     self.get_scheduler_config(**lowerCamelCase    )
								lowercase__  					=     scheduler_class(**lowerCamelCase    )

								lowercase__  					=     10
								lowercase__  					=     self.dummy_model()
								lowercase__  					=     self.dummy_sample_deter
								scheduler.set_timesteps(lowerCamelCase    )

								for i, t in enumerate(scheduler.timesteps    ):
												lowercase__  					=     model(lowerCamelCase,    lowerCamelCase    )
												lowercase__  					=     scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase    ).prev_sample

								for i, t in enumerate(scheduler.timesteps    ):
												lowercase__  					=     model(lowerCamelCase,    lowerCamelCase    )
												lowercase__  					=     scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase    ).prev_sample

								return sample




				def 							lowercase__  (	self		:       List[Any]    ):

								'''simple docstring'''
								lowercase__  					=     dict(self.forward_default_kwargs    )

								lowercase__  					=     kwargs.pop('''num_inference_steps''',    lowerCamelCase    )

								for scheduler_class in self.scheduler_classes:
												lowercase__  					=     self.get_scheduler_config()
												lowercase__  					=     scheduler_class(**lowerCamelCase    )

												lowercase__  					=     self.dummy_sample
												lowercase__  					=     0.1 * sample

												if num_inference_steps is not None and hasattr(lowerCamelCase,    '''set_timesteps'''    ):
																scheduler.set_timesteps(lowerCamelCase    )
												elif num_inference_steps is not None and not hasattr(lowerCamelCase,    '''set_timesteps'''    ):
																lowercase__  					=     num_inference_steps

												# copy over dummy past residuals (must be done after set_timesteps)
												lowercase__  					=     [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
												lowercase__  					=     dummy_past_residuals[:]

												lowercase__  					=     scheduler.timesteps[5]
												lowercase__  					=     scheduler.timesteps[6]

												lowercase__  					=     scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample
												lowercase__  					=     scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample

												self.assertEqual(output_a.shape,    sample.shape    )
												self.assertEqual(output_a.shape,    output_a.shape    )

												lowercase__  					=     scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample
												lowercase__  					=     scheduler.step(lowerCamelCase,    lowerCamelCase,    lowerCamelCase,    **lowerCamelCase    ).prev_sample

												self.assertEqual(output_a.shape,    sample.shape    )
												self.assertEqual(output_a.shape,    output_a.shape    )




				def 							lowercase__  (	self		:       Dict    ):

								'''simple docstring'''
								for timesteps in [100, 1_000]:
												self.check_over_configs(num_train_timesteps=lowerCamelCase,    time_step=lowerCamelCase    )




				def 							lowercase__  (	self		:       str    ):

								'''simple docstring'''
								for t, num_inference_steps in zip([1, 5, 10],    [10, 50, 100]    ):
												self.check_over_forward(num_inference_steps=lowerCamelCase,    time_step=lowerCamelCase    )




				def 							lowercase__  (	self		:       List[str]    ):

								'''simple docstring'''
								lowercase__  					=     self.full_loop()
								lowercase__  					=     torch.mean(torch.abs(lowerCamelCase    )    )

								assert abs(result_mean.item() - 2_540_529    ) < 10




