


import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class lowerCAmelCase_						(			lowerCamelCase__ ,						lowerCamelCase__ ,						unittest.TestCase	):

 '''simple docstring'''

 __snake_case												=  StableDiffusionDiffEditPipeline
 __snake_case												=  TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"height", "width", "image"} | {"image_latents"}
 __snake_case												=  TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {"image"} | {"image_latents"}
 __snake_case												=  frozenset(
     []	)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 __snake_case												=  frozenset([]	)



 def 						UpperCamelCase__					(     self ):
       torch.manual_seed(0 )
       snake_case_	=				UNetaDConditionModel(
           block_out_channels=(32, 64)		,				layers_per_block=2		,				sample_size=32		,				in_channels=4		,				out_channels=4		,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')		,				up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')		,				cross_attention_dim=32		,				attention_head_dim=(2, 4)		,				use_linear_projection=_UpperCAmelCase		,				)
       snake_case_	=				DDIMScheduler(
           beta_start=0.00_085		,				beta_end=0.012		,				beta_schedule='''scaled_linear'''		,				clip_sample=_UpperCAmelCase		,				set_alpha_to_one=_UpperCAmelCase		,				)
       snake_case_	=				DDIMInverseScheduler(
           beta_start=0.00_085		,				beta_end=0.012		,				beta_schedule='''scaled_linear'''		,				clip_sample=_UpperCAmelCase		,				set_alpha_to_zero=_UpperCAmelCase		,				)
       torch.manual_seed(0 )
       snake_case_	=				AutoencoderKL(
           block_out_channels=[32, 64]		,				in_channels=3		,				out_channels=3		,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']		,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']		,				latent_channels=4		,				sample_size=1_28		,				)
       torch.manual_seed(0 )
       snake_case_	=				CLIPTextConfig(
           bos_token_id=0		,				eos_token_id=2		,				hidden_size=32		,				intermediate_size=37		,				layer_norm_eps=1E-05		,				num_attention_heads=4		,				num_hidden_layers=5		,				pad_token_id=1		,				vocab_size=10_00		,				hidden_act='''gelu'''		,				projection_dim=5_12		,				)
       snake_case_	=				CLIPTextModel(_UpperCAmelCase )
       snake_case_	=				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )

       snake_case_	=				{
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''inverse_scheduler''': inverse_scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }

       return components



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase=0 ):
       snake_case_	=				floats_tensor((1, 16, 16)		,				rng=random.Random(_UpperCAmelCase ) ).to(_UpperCAmelCase )
       snake_case_	=				floats_tensor((1, 2, 4, 16, 16)		,				rng=random.Random(_UpperCAmelCase ) ).to(_UpperCAmelCase )
       if str(_UpperCAmelCase ).startswith('''mps''' ):
             snake_case_	=				torch.manual_seed(_UpperCAmelCase )
       else:
             snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(_UpperCAmelCase )
       snake_case_	=				{
           '''prompt''': '''a dog and a newt''',
           '''mask_image''': mask,
           '''image_latents''': latents,
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''inpaint_strength''': 1.0,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
       }

       return inputs



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase=0 ):
       snake_case_	=				floats_tensor((1, 3, 32, 32)		,				rng=random.Random(_UpperCAmelCase ) ).to(_UpperCAmelCase )
       snake_case_	=				image.cpu().permute(0		,				2		,				3		,				1 )[0]
       snake_case_	=				Image.fromarray(np.uinta(_UpperCAmelCase ) ).convert('''RGB''' )
       if str(_UpperCAmelCase ).startswith('''mps''' ):
             snake_case_	=				torch.manual_seed(_UpperCAmelCase )
       else:
             snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(_UpperCAmelCase )
       snake_case_	=				{
           '''image''': image,
           '''source_prompt''': '''a cat and a frog''',
           '''target_prompt''': '''a dog and a newt''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''num_maps_per_mask''': 2,
           '''mask_encode_strength''': 1.0,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
       }

       return inputs



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase=0 ):
       snake_case_	=				floats_tensor((1, 3, 32, 32)		,				rng=random.Random(_UpperCAmelCase ) ).to(_UpperCAmelCase )
       snake_case_	=				image.cpu().permute(0		,				2		,				3		,				1 )[0]
       snake_case_	=				Image.fromarray(np.uinta(_UpperCAmelCase ) ).convert('''RGB''' )
       if str(_UpperCAmelCase ).startswith('''mps''' ):
             snake_case_	=				torch.manual_seed(_UpperCAmelCase )
       else:
             snake_case_	=				torch.Generator(device=_UpperCAmelCase ).manual_seed(_UpperCAmelCase )
       snake_case_	=				{
           '''image''': image,
           '''prompt''': '''a cat and a frog''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''inpaint_strength''': 1.0,
           '''guidance_scale''': 6.0,
           '''decode_latents''': True,
           '''output_type''': '''numpy''',
       }
       return inputs



 def 						UpperCamelCase__					(     self ):
       if not hasattr(self.pipeline_class		,				'''_optional_components''' ):
             return

       snake_case_	=				self.get_dummy_components()
       snake_case_	=				self.pipeline_class(**_UpperCAmelCase )
       pipe.to(_UpperCAmelCase )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       # set all optional components to None and update pipeline config accordingly
       for optional_component in pipe._optional_components:
             setattr(_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase )
       pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components} )

       snake_case_	=				self.get_dummy_inputs(_UpperCAmelCase )
       snake_case_	=				pipe(**_UpperCAmelCase )[0]

       with tempfile.TemporaryDirectory() as tmpdir:
             pipe.save_pretrained(_UpperCAmelCase )
             snake_case_	=				self.pipeline_class.from_pretrained(_UpperCAmelCase )
             pipe_loaded.to(_UpperCAmelCase )
             pipe_loaded.set_progress_bar_config(disable=_UpperCAmelCase )

       for optional_component in pipe._optional_components:
             self.assertTrue(
                 getattr(_UpperCAmelCase		,				_UpperCAmelCase ) is None		,				F'''`{optional_component}` did not stay set to None after loading.'''		,				)

       snake_case_	=				self.get_dummy_inputs(_UpperCAmelCase )
       snake_case_	=				pipe_loaded(**_UpperCAmelCase )[0]

       snake_case_	=				np.abs(output - output_loaded ).max()
       self.assertLess(_UpperCAmelCase		,				1E-4 )



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				'''cpu'''

       snake_case_	=				self.get_dummy_components()
       snake_case_	=				self.pipeline_class(**_UpperCAmelCase )
       pipe.to(_UpperCAmelCase )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       snake_case_	=				self.get_dummy_mask_inputs(_UpperCAmelCase )
       snake_case_	=				pipe.generate_mask(**_UpperCAmelCase )
       snake_case_	=				mask[0, -3:, -3:]

       self.assertEqual(mask.shape		,				(1, 16, 16) )
       snake_case_	=				np.array([0] * 9 )
       snake_case_	=				np.abs(mask_slice.flatten() - expected_slice ).max()
       self.assertLessEqual(_UpperCAmelCase		,				1E-3 )
       self.assertEqual(mask[0, -3, -4]		,				0 )



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				'''cpu'''

       snake_case_	=				self.get_dummy_components()
       snake_case_	=				self.pipeline_class(**_UpperCAmelCase )
       pipe.to(_UpperCAmelCase )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       snake_case_	=				self.get_dummy_inversion_inputs(_UpperCAmelCase )
       snake_case_	=				pipe.invert(**_UpperCAmelCase ).images
       snake_case_	=				image[0, -1, -3:, -3:]

       self.assertEqual(image.shape		,				(2, 32, 32, 3) )
       snake_case_	=				np.array(
           [0.5_150, 0.5_134, 0.5_043, 0.5_376, 0.4_694, 0.51_050, 0.5_015, 0.4_407, 0.4_799]		,				)
       snake_case_	=				np.abs(image_slice.flatten() - expected_slice ).max()
       self.assertLessEqual(_UpperCAmelCase		,				1E-3 )



 def 						UpperCamelCase__					(     self ):
       super().test_inference_batch_single_identical(expected_max_diff=5E-3 )





 def 						UpperCamelCase__					(     self ):
       snake_case_	=				'''cpu'''

       snake_case_	=				self.get_dummy_components()

       snake_case_	=				{'''beta_start''': 0.00_085, '''beta_end''': 0.012, '''beta_schedule''': '''scaled_linear'''}
       snake_case_	=				DPMSolverMultistepScheduler(**_UpperCAmelCase )
       snake_case_	=				DPMSolverMultistepInverseScheduler(**_UpperCAmelCase )

       snake_case_	=				self.pipeline_class(**_UpperCAmelCase )
       pipe.to(_UpperCAmelCase )
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       snake_case_	=				self.get_dummy_inversion_inputs(_UpperCAmelCase )
       snake_case_	=				pipe.invert(**_UpperCAmelCase ).images
       snake_case_	=				image[0, -1, -3:, -3:]

       self.assertEqual(image.shape		,				(2, 32, 32, 3) )
       snake_case_	=				np.array(
           [0.5_150, 0.5_134, 0.5_043, 0.5_376, 0.4_694, 0.51_050, 0.5_015, 0.4_407, 0.4_799]		,				)
       snake_case_	=				np.abs(image_slice.flatten() - expected_slice ).max()
       self.assertLessEqual(_UpperCAmelCase		,				1E-3 )







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

 '''simple docstring'''

 def 						UpperCamelCase__					(     self ):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



 @classmethod
 def 						UpperCamelCase__					(     cls ):
       snake_case_	=				load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png''' )

       snake_case_	=				raw_image.convert('''RGB''' ).resize((7_68, 7_68) )

       snake_case_	=				raw_image



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				torch.manual_seed(0 )

       snake_case_	=				StableDiffusionDiffEditPipeline.from_pretrained(
           '''stabilityai/stable-diffusion-2-1'''		,				safety_checker=_UpperCAmelCase		,				torch_dtype=torch.floataa )
       snake_case_	=				DDIMScheduler.from_config(pipe.scheduler.config )
       snake_case_	=				DDIMInverseScheduler.from_config(pipe.scheduler.config )
       pipe.enable_model_cpu_offload()
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       snake_case_	=				'''a bowl of fruit'''
       snake_case_	=				'''a bowl of pears'''

       snake_case_	=				pipe.generate_mask(
           image=self.raw_image		,				source_prompt=_UpperCAmelCase		,				target_prompt=_UpperCAmelCase		,				generator=_UpperCAmelCase		,				)

       snake_case_	=				pipe.invert(
           prompt=_UpperCAmelCase		,				image=self.raw_image		,				inpaint_strength=0.7		,				generator=_UpperCAmelCase ).latents

       snake_case_	=				pipe(
           prompt=_UpperCAmelCase		,				mask_image=_UpperCAmelCase		,				image_latents=_UpperCAmelCase		,				generator=_UpperCAmelCase		,				negative_prompt=_UpperCAmelCase		,				inpaint_strength=0.7		,				output_type='''numpy'''		,				).images[0]

       snake_case_	=				(
           np.array(
               load_image(
                   '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                   '''/diffedit/pears.png''' ).resize((7_68, 7_68) ) )
           / 2_55
       )
       assert np.abs((expected_image - image).max() ) < 5E-1





 def 						UpperCamelCase__					(     self ):
       snake_case_	=				torch.manual_seed(0 )

       snake_case_	=				StableDiffusionDiffEditPipeline.from_pretrained(
           '''stabilityai/stable-diffusion-2-1'''		,				safety_checker=_UpperCAmelCase		,				torch_dtype=torch.floataa )
       snake_case_	=				DPMSolverMultistepScheduler.from_config(pipe.scheduler.config )
       snake_case_	=				DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config )
       pipe.enable_model_cpu_offload()
       pipe.set_progress_bar_config(disable=_UpperCAmelCase )

       snake_case_	=				'''a bowl of fruit'''
       snake_case_	=				'''a bowl of pears'''

       snake_case_	=				pipe.generate_mask(
           image=self.raw_image		,				source_prompt=_UpperCAmelCase		,				target_prompt=_UpperCAmelCase		,				generator=_UpperCAmelCase		,				)

       snake_case_	=				pipe.invert(
           prompt=_UpperCAmelCase		,				image=self.raw_image		,				inpaint_strength=0.7		,				generator=_UpperCAmelCase		,				num_inference_steps=25		,				).latents

       snake_case_	=				pipe(
           prompt=_UpperCAmelCase		,				mask_image=_UpperCAmelCase		,				image_latents=_UpperCAmelCase		,				generator=_UpperCAmelCase		,				negative_prompt=_UpperCAmelCase		,				inpaint_strength=0.7		,				num_inference_steps=25		,				output_type='''numpy'''		,				).images[0]

       snake_case_	=				(
           np.array(
               load_image(
                   '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                   '''/diffedit/pears.png''' ).resize((7_68, 7_68) ) )
           / 2_55
       )
       assert np.abs((expected_image - image).max() ) < 5E-1



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
from lavis.models import load_model_and_preprocess
from PIL import Image

from transformers import (
    AutoTokenizer,
    BlipaConfig,
    BlipaForConditionalGeneration,
    BlipaProcessor,
    BlipaVisionConfig,
    BlipImageProcessor,
    OPTConfig,
    TaConfig,
)
from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD






def 	__lowerCAmelCase   ()->  Optional[Any]:




      """simple docstring"""

      snake_case_	=				'''https://storage.googleapis.com/sfr-vision-language-research/LAVIS/assets/merlion.png'''
      snake_case_	=				Image.open(requests.get(SCREAMING_SNAKE_CASE , stream=SCREAMING_SNAKE_CASE ).raw ).convert('''RGB''' )

      return image






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE )->  Any:




      """simple docstring"""

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

      # fmt: on
      return rename_keys






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  Tuple:




      """simple docstring"""

      snake_case_	=				dct.pop(SCREAMING_SNAKE_CASE )
      snake_case_	=				val






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  Dict:




      """simple docstring"""

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

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






def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )->  List[str]:




      """simple docstring"""

      snake_case_	=				364 if '''coco''' in model_name else 224
      snake_case_	=				BlipaVisionConfig(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 "opt-2.7b" in model_name:
            snake_case_	=				OPTConfig.from_pretrained('''facebook/opt-2.7b''' , eos_token_id=SCREAMING_SNAKE_CASE ).to_dict()
      elif "opt-6.7b" in model_name:
            snake_case_	=				OPTConfig.from_pretrained('''facebook/opt-6.7b''' , eos_token_id=SCREAMING_SNAKE_CASE ).to_dict()
      elif "t5-xl" in model_name:
            snake_case_	=				TaConfig.from_pretrained('''google/flan-t5-xl''' , dense_act_fn='''gelu''' , bos_token_id=1 ).to_dict()
      elif "t5-xxl" in model_name:
            snake_case_	=				TaConfig.from_pretrained('''google/flan-t5-xxl''' , dense_act_fn='''gelu''' , bos_token_id=1 ).to_dict()

      snake_case_	=				BlipaConfig(vision_config=SCREAMING_SNAKE_CASE , text_config=SCREAMING_SNAKE_CASE )

      return config, image_size






@torch.no_grad()
def 	__lowerCAmelCase   (SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE=None , SCREAMING_SNAKE_CASE=False )->  List[Any]:




      """simple docstring"""

      snake_case_	=				(
          AutoTokenizer.from_pretrained('''facebook/opt-2.7b''' )
          if '''opt''' in model_name
          else AutoTokenizer.from_pretrained('''google/flan-t5-xl''' )
      )
      snake_case_	=				tokenizer('''\n''' , add_special_tokens=SCREAMING_SNAKE_CASE ).input_ids[0]
      snake_case_	,   snake_case_	=				get_blipa_config(SCREAMING_SNAKE_CASE , eos_token_id=SCREAMING_SNAKE_CASE )

      snake_case_	=				BlipaForConditionalGeneration(SCREAMING_SNAKE_CASE ).eval()

      snake_case_	=				{
          '''blip2-opt-2.7b''': ('''blip2_opt''', '''pretrain_opt2.7b'''),
          '''blip2-opt-6.7b''': ('''blip2_opt''', '''pretrain_opt6.7b'''),
          '''blip2-opt-2.7b-coco''': ('''blip2_opt''', '''caption_coco_opt2.7b'''),
          '''blip2-opt-6.7b-coco''': ('''blip2_opt''', '''caption_coco_opt6.7b'''),
          '''blip2-flan-t5-xl''': ('''blip2_t5''', '''pretrain_flant5xl'''),
          '''blip2-flan-t5-xl-coco''': ('''blip2_t5''', '''caption_coco_flant5xl'''),
          '''blip2-flan-t5-xxl''': ('''blip2_t5''', '''pretrain_flant5xxl'''),
      }

      snake_case_	,   snake_case_	=				model_name_to_original[model_name]

      # load original model
      print('''Loading original model...''' )
      snake_case_	=				'''cuda''' if torch.cuda.is_available() else '''cpu'''
      snake_case_	,   snake_case_	,   snake_case_	=				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
      snake_case_	=				original_model.state_dict()
      snake_case_	=				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():
            snake_case_	=				state_dict.pop(SCREAMING_SNAKE_CASE )
            if key.startswith('''Qformer.bert''' ):
                  snake_case_	=				key.replace('''Qformer.bert''' , '''qformer''' )
            if "attention.self" in key:
                  snake_case_	=				key.replace('''self''' , '''attention''' )
            if "opt_proj" in key:
                  snake_case_	=				key.replace('''opt_proj''' , '''language_projection''' )
            if "t5_proj" in key:
                  snake_case_	=				key.replace('''t5_proj''' , '''language_projection''' )
            if key.startswith('''opt''' ):
                  snake_case_	=				key.replace('''opt''' , '''language''' )
            if key.startswith('''t5''' ):
                  snake_case_	=				key.replace('''t5''' , '''language''' )
            snake_case_	=				val

      # read in qv biases
      read_in_q_v_bias(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )

      snake_case_	,   snake_case_	=				hf_model.load_state_dict(SCREAMING_SNAKE_CASE , strict=SCREAMING_SNAKE_CASE )
      assert len(SCREAMING_SNAKE_CASE ) == 0
      assert unexpected_keys == ["qformer.embeddings.position_ids"]

      snake_case_	=				load_demo_image()
      snake_case_	=				vis_processors['''eval'''](SCREAMING_SNAKE_CASE ).unsqueeze(0 ).to(SCREAMING_SNAKE_CASE )
      snake_case_	=				tokenizer(['''\n'''] , return_tensors='''pt''' ).input_ids.to(SCREAMING_SNAKE_CASE )

      # create processor
      snake_case_	=				BlipImageProcessor(
          size={'''height''': image_size, '''width''': image_size} , image_mean=SCREAMING_SNAKE_CASE , image_std=SCREAMING_SNAKE_CASE )
      snake_case_	=				BlipaProcessor(image_processor=SCREAMING_SNAKE_CASE , tokenizer=SCREAMING_SNAKE_CASE )
      snake_case_	=				processor(images=SCREAMING_SNAKE_CASE , return_tensors='''pt''' ).pixel_values.to(SCREAMING_SNAKE_CASE )

      # make sure processor creates exact same pixel values
      assert torch.allclose(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE )

      original_model.to(SCREAMING_SNAKE_CASE )
      hf_model.to(SCREAMING_SNAKE_CASE )
      with torch.no_grad():
            if "opt" in model_name:
                  snake_case_	=				original_model({'''image''': original_pixel_values, '''text_input''': ['''''']} ).logits
                  snake_case_	=				hf_model(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE ).logits
            else:
                  snake_case_	=				original_model(
                      {'''image''': original_pixel_values, '''text_input''': ['''\n'''], '''text_output''': ['''\n''']} ).logits
                  snake_case_	=				input_ids.masked_fill(input_ids == tokenizer.pad_token_id , -100 )
                  snake_case_	=				hf_model(SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , labels=SCREAMING_SNAKE_CASE ).logits

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

      # assert values
      if model_name == "blip2-flan-t5-xl":
            snake_case_	=				torch.tensor(
                [[-4_1.5_8_5_0, -4.4_4_4_0, -8.9_9_2_2], [-4_7.4_3_2_2, -5.9_1_4_3, -1.7_3_4_0]] , device=SCREAMING_SNAKE_CASE )
            assert torch.allclose(logits[0, :3, :3] , SCREAMING_SNAKE_CASE , atol=1E-4 )
      elif model_name == "blip2-flan-t5-xl-coco":
            snake_case_	=				torch.tensor(
                [[-5_7.0_1_0_9, -9.8_9_6_7, -1_2.6_2_8_0], [-6_8.6_5_7_8, -1_2.7_1_9_1, -1_0.5_0_6_5]] , device=SCREAMING_SNAKE_CASE )
      else:
            # cast to same type
            snake_case_	=				logits.dtype
            assert torch.allclose(original_logits.to(SCREAMING_SNAKE_CASE ) , SCREAMING_SNAKE_CASE , atol=1E-2 )
      print('''Looks ok!''' )

      print('''Generating a caption...''' )
      snake_case_	=				''''''
      snake_case_	=				tokenizer(SCREAMING_SNAKE_CASE , return_tensors='''pt''' ).input_ids.to(SCREAMING_SNAKE_CASE )

      snake_case_	=				original_model.generate({'''image''': original_pixel_values} )
      snake_case_	=				hf_model.generate(
          SCREAMING_SNAKE_CASE , SCREAMING_SNAKE_CASE , do_sample=SCREAMING_SNAKE_CASE , num_beams=5 , max_length=30 , min_length=1 , top_p=0.9 , repetition_penalty=1.0 , length_penalty=1.0 , temperature=1 , )
      print('''Original generation:''' , SCREAMING_SNAKE_CASE )
      snake_case_	=				input_ids.shape[1]
      snake_case_	=				processor.batch_decode(outputs[:, prompt_length:] , skip_special_tokens=SCREAMING_SNAKE_CASE )
      snake_case_	=				[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'''nielsr/{model_name}''' )
            hf_model.push_to_hub(f'''nielsr/{model_name}''' )


if __name__ == "__main__":
   UpperCAmelCase      =   argparse.ArgumentParser()
   UpperCAmelCase      =   [
       """blip2-opt-2.7b""",
       """blip2-opt-6.7b""",
       """blip2-opt-2.7b-coco""",
       """blip2-opt-6.7b-coco""",
       """blip2-flan-t5-xl""",
       """blip2-flan-t5-xl-coco""",
       """blip2-flan-t5-xxl""",
   ]
   parser.add_argument(
       """--model_name""",
       default="""blip2-opt-2.7b""",
       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""",
   )

   UpperCAmelCase      =   parser.parse_args()

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