



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

   return image





def 		a__			(		_UpperCamelCase			: 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 		a__			(		_UpperCamelCase			: List[Any] ,_UpperCamelCase			: str ,_UpperCamelCase			: str							):
   __lowerCamelCase													=    dct.pop(_UpperCamelCase							)
   __lowerCamelCase													=    val





def 		a__			(		_UpperCamelCase			: Optional[Any] ,_UpperCamelCase			: 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(_UpperCamelCase ,requires_grad=_UpperCamelCase							), v_bias)							)
      __lowerCamelCase													=    qkv_bias





def 		a__			(		_UpperCamelCase			: List[str]							):
   __lowerCamelCase													=    3_64 if '''coco''' in model_name else 2_24
   __lowerCamelCase													=    InstructBlipVisionConfig(image_size=_UpperCamelCase							).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=3_20_01							).to_dict()
   elif "vicuna-13b" in model_name:
      __lowerCamelCase													=    LlamaConfig.from_pretrained('''decapoda-research/llama-13b-hf''' ,vocab_size=3_20_01							).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=3_05_23							).to_dict()
   __lowerCamelCase													=    InstructBlipConfig(vision_config=_UpperCamelCase ,text_config=_UpperCamelCase ,qformer_config=_UpperCamelCase							)

   return config, image_size





@torch.no_grad()
def 		a__			(		_UpperCamelCase			: int ,_UpperCamelCase			: Tuple=None ,_UpperCamelCase			: Union[str, Any]=False							):
   __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(_UpperCamelCase							)
   __lowerCamelCase													=    InstructBlipForConditionalGeneration(_UpperCamelCase							).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=_UpperCamelCase ,model_type=_UpperCamelCase ,is_eval=_UpperCamelCase ,device=_UpperCamelCase							)
   original_model.eval()
   print('''Done!'''							)

   # update state dict keys
   __lowerCamelCase													=    original_model.state_dict()
   __lowerCamelCase													=    create_rename_keys(_UpperCamelCase							)
   for src, dest in rename_keys:
      rename_key(_UpperCamelCase ,_UpperCamelCase ,_UpperCamelCase							)

   # some keys can be renamed efficiently
   for key, val in state_dict.copy().items():
      __lowerCamelCase													=    state_dict.pop(_UpperCamelCase							)
      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(_UpperCamelCase ,_UpperCamelCase							)

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

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

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

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

   original_model.to(_UpperCamelCase							)
   hf_model.to(_UpperCamelCase							)
   with torch.no_grad():
      if "vicuna" in model_name:
         __lowerCamelCase													=    original_model({'''image''': original_pixel_values, '''text_input''': [prompt]}							).logits
         __lowerCamelCase													=    hf_model(**_UpperCamelCase							).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(_UpperCamelCase							)
         __lowerCamelCase													=    label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id ,-1_00							)
         __lowerCamelCase													=    hf_model(**_UpperCamelCase ,labels=_UpperCamelCase							).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							) ,_UpperCamelCase ,atol=_UpperCamelCase							)
   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(
       **_UpperCamelCase ,do_sample=_UpperCamelCase ,num_beams=5 ,max_length=2_56 ,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:''' ,_UpperCamelCase							)
   __lowerCamelCase													=    processor.batch_decode(_UpperCamelCase ,skip_special_tokens=_UpperCamelCase							)
   __lowerCamelCase													=    [text.strip() for text in output_text]
   print('''HF generation:''' ,_UpperCamelCase							)

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

   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)











import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class  __lowerCAmelCase     (	lowerCAmelCase__  ,	lowerCAmelCase__  ,	unittest.TestCase   ):
 lowerCAmelCase__       					=  StableDiffusionXLImgaImgPipeline
 lowerCAmelCase__       					=  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
 lowerCAmelCase__       					=  PipelineTesterMixin.required_optional_params - {"""latents"""}
 lowerCAmelCase__       					=  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 lowerCAmelCase__       					=  IMAGE_TO_IMAGE_IMAGE_PARAMS
 lowerCAmelCase__       					=  IMAGE_TO_IMAGE_IMAGE_PARAMS




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    torch.manual_seed(0				)
    __lowerCamelCase													=    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''')	,				attention_head_dim=(2, 4)	,				use_linear_projection=__UpperCAmelCase	,				addition_embed_type='''text_time'''	,				addition_time_embed_dim=8	,				transformer_layers_per_block=(1, 2)	,				projection_class_embeddings_input_dim=80	,				cross_attention_dim=64	,				)
    __lowerCamelCase													=    EulerDiscreteScheduler(
        beta_start=0.00_085	,				beta_end=0.012	,				steps_offset=1	,				beta_schedule='''scaled_linear'''	,				timestep_spacing='''leading'''	,				)
    torch.manual_seed(0				)
    __lowerCamelCase													=    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=128	,				)
    torch.manual_seed(0				)
    __lowerCamelCase													=    CLIPTextConfig(
        bos_token_id=0	,				eos_token_id=2	,				hidden_size=32	,				intermediate_size=37	,				layer_norm_eps=1E-0_5	,				num_attention_heads=4	,				num_hidden_layers=5	,				pad_token_id=1	,				vocab_size=1000	,				hidden_act='''gelu'''	,				projection_dim=32	,				)
    __lowerCamelCase													=    CLIPTextModel(__UpperCAmelCase				)
    __lowerCamelCase													=    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	,				local_files_only=__UpperCAmelCase				)

    __lowerCamelCase													=    CLIPTextModelWithProjection(__UpperCAmelCase				)
    __lowerCamelCase													=    CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	,				local_files_only=__UpperCAmelCase				)

    __lowerCamelCase													=    {
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''vae''': vae,
        '''text_encoder''': text_encoder,
        '''tokenizer''': tokenizer,
        '''text_encoder_2''': text_encoder_a,
        '''tokenizer_2''': tokenizer_a,
        # "safety_checker": None,
        # "feature_extractor": None,
    }
    return components




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase=0				):



    '''simple docstring'''
    __lowerCamelCase													=    floats_tensor((1, 3, 32, 32)	,				rng=random.Random(__UpperCAmelCase				)				).to(__UpperCAmelCase				)
    __lowerCamelCase													=    image / 2 + 0.5
    if str(__UpperCAmelCase				).startswith('''mps'''				):
       __lowerCamelCase													=    torch.manual_seed(__UpperCAmelCase				)
    else:
       __lowerCamelCase													=    torch.Generator(device=__UpperCAmelCase				).manual_seed(__UpperCAmelCase				)
    __lowerCamelCase													=    {
        '''prompt''': '''A painting of a squirrel eating a burger''',
        '''image''': image,
        '''generator''': generator,
        '''num_inference_steps''': 2,
        '''guidance_scale''': 5.0,
        '''output_type''': '''numpy''',
        '''strength''': 0.75,
    }
    return inputs




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
    __lowerCamelCase													=    self.get_dummy_components()
    __lowerCamelCase													=    StableDiffusionXLImgaImgPipeline(**__UpperCAmelCase				)
    __lowerCamelCase													=    sd_pipe.to(__UpperCAmelCase				)
    sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

    __lowerCamelCase													=    self.get_dummy_inputs(__UpperCAmelCase				)
    __lowerCamelCase													=    sd_pipe(**__UpperCAmelCase				).images
    __lowerCamelCase													=    image[0, -3:, -3:, -1]

    assert image.shape == (1, 32, 32, 3)

    __lowerCamelCase													=    np.array([0.4_656, 0.4_840, 0.4_439, 0.6_698, 0.5_574, 0.4_524, 0.5_799, 0.5_943, 0.5_165]				)

    assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    super().test_attention_slicing_forward_pass(expected_max_diff=3E-3				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    super().test_inference_batch_single_identical(expected_max_diff=3E-3				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    pass




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.get_dummy_components()
    __lowerCamelCase													=    StableDiffusionXLImgaImgPipeline(**__UpperCAmelCase				)
    __lowerCamelCase													=    sd_pipe.to(__UpperCAmelCase				)
    __lowerCamelCase													=    sd_pipe.to(__UpperCAmelCase				)
    sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

    # forward without prompt embeds
    __lowerCamelCase													=    self.get_dummy_inputs(__UpperCAmelCase				)
    __lowerCamelCase													=    3 * ['''this is a negative prompt''']
    __lowerCamelCase													=    negative_prompt
    __lowerCamelCase													=    3 * [inputs['''prompt''']]

    __lowerCamelCase													=    sd_pipe(**__UpperCAmelCase				)
    __lowerCamelCase													=    output.images[0, -3:, -3:, -1]

    # forward with prompt embeds
    __lowerCamelCase													=    self.get_dummy_inputs(__UpperCAmelCase				)
    __lowerCamelCase													=    3 * ['''this is a negative prompt''']
    __lowerCamelCase													=    3 * [inputs.pop('''prompt'''				)]

    (
        (
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,
    )													=    sd_pipe.encode_prompt(__UpperCAmelCase	,				negative_prompt=__UpperCAmelCase				)

    __lowerCamelCase													=    sd_pipe(
        **__UpperCAmelCase	,				prompt_embeds=__UpperCAmelCase	,				negative_prompt_embeds=__UpperCAmelCase	,				pooled_prompt_embeds=__UpperCAmelCase	,				negative_pooled_prompt_embeds=__UpperCAmelCase	,				)
    __lowerCamelCase													=    output.images[0, -3:, -3:, -1]

    # make sure that it's equal
    assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()				).max() < 1E-4



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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase="cpu"	,				__UpperCAmelCase=torch.floataa	,				__UpperCAmelCase=0				):



    '''simple docstring'''
    __lowerCamelCase													=    torch.Generator(device=__UpperCAmelCase				).manual_seed(__UpperCAmelCase				)
    __lowerCamelCase													=    np.random.RandomState(__UpperCAmelCase				).standard_normal((1, 4, 64, 64)				)
    __lowerCamelCase													=    torch.from_numpy(__UpperCAmelCase				).to(device=__UpperCAmelCase	,				dtype=__UpperCAmelCase				)
    __lowerCamelCase													=    {
        '''prompt''': '''a photograph of an astronaut riding a horse''',
        '''latents''': latents,
        '''generator''': generator,
        '''num_inference_steps''': 3,
        '''guidance_scale''': 7.5,
        '''output_type''': '''numpy''',
    }
    return inputs




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''				)
    pipe.to(__UpperCAmelCase				)
    pipe.set_progress_bar_config(disable=__UpperCAmelCase				)

    __lowerCamelCase													=    self.get_inputs(__UpperCAmelCase				)
    __lowerCamelCase													=    pipe(**__UpperCAmelCase				).images
    __lowerCamelCase													=    image[0, -3:, -3:, -1].flatten()

    assert image.shape == (1, 512, 512, 3)
    __lowerCamelCase													=    np.array([0.49_493, 0.47_896, 0.40_798, 0.54_214, 0.53_212, 0.48_202, 0.47_656, 0.46_329, 0.48_506]				)
    assert np.abs(image_slice - expected_slice				).max() < 7E-3






