





'''simple docstring'''






import argparse
import json

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    VideoMAEConfig,
    VideoMAEForPreTraining,
    VideoMAEForVideoClassification,
    VideoMAEImageProcessor,
)


def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=VideoMAEConfig()

  set_architecture_configs(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )

  if "finetuned" not in model_name:
    _UpperCamelCase       			=False

  if "finetuned" in model_name:
    _UpperCamelCase       			='''huggingface/label-files'''
    if "kinetics" in model_name:
      _UpperCamelCase       			=400
      _UpperCamelCase       			='''kinetics400-id2label.json'''
    elif "ssv2" in model_name:
      _UpperCamelCase       			=174
      _UpperCamelCase       			='''something-something-v2-id2label.json'''
    else:
      raise ValueError('''Model name should either contain \'kinetics\' or \'ssv2\' in case it\'s fine-tuned.'''   )
    _UpperCamelCase       			=json.load(open(hf_hub_download(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      repo_type='''dataset'''   )  ,      '''r'''   )   )
    _UpperCamelCase       			={int(__SCREAMING_SNAKE_CASE   ): v for k, v in idalabel.items()}
    _UpperCamelCase       			=idalabel
    _UpperCamelCase       			={v: k for k, v in idalabel.items()}

  return config


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  if "small" in model_name:
    _UpperCamelCase       			=384
    _UpperCamelCase       			=1536
    _UpperCamelCase       			=12
    _UpperCamelCase       			=16
    _UpperCamelCase       			=12
    _UpperCamelCase       			=3
    _UpperCamelCase       			=192
    _UpperCamelCase       			=768
  elif "large" in model_name:
    _UpperCamelCase       			=1024
    _UpperCamelCase       			=4096
    _UpperCamelCase       			=24
    _UpperCamelCase       			=16
    _UpperCamelCase       			=12
    _UpperCamelCase       			=8
    _UpperCamelCase       			=512
    _UpperCamelCase       			=2048
  elif "huge" in model_name:
    _UpperCamelCase       			=1280
    _UpperCamelCase       			=5120
    _UpperCamelCase       			=32
    _UpperCamelCase       			=16
    _UpperCamelCase       			=12
    _UpperCamelCase       			=8
    _UpperCamelCase       			=640
    _UpperCamelCase       			=2560
  elif "base" not in model_name:
    raise ValueError('''Model name should include either "small", "base", "large", or "huge"'''   )


def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  if "encoder." in name:
    _UpperCamelCase       			=name.replace('''encoder.'''  ,      ''''''   )
  if "cls_token" in name:
    _UpperCamelCase       			=name.replace('''cls_token'''  ,      '''videomae.embeddings.cls_token'''   )
  if "decoder_pos_embed" in name:
    _UpperCamelCase       			=name.replace('''decoder_pos_embed'''  ,      '''decoder.decoder_pos_embed'''   )
  if "pos_embed" in name and "decoder" not in name:
    _UpperCamelCase       			=name.replace('''pos_embed'''  ,      '''videomae.embeddings.position_embeddings'''   )
  if "patch_embed.proj" in name:
    _UpperCamelCase       			=name.replace('''patch_embed.proj'''  ,      '''videomae.embeddings.patch_embeddings.projection'''   )
  if "patch_embed.norm" in name:
    _UpperCamelCase       			=name.replace('''patch_embed.norm'''  ,      '''videomae.embeddings.norm'''   )
  if "decoder.blocks" in name:
    _UpperCamelCase       			=name.replace('''decoder.blocks'''  ,      '''decoder.decoder_layers'''   )
  if "blocks" in name:
    _UpperCamelCase       			=name.replace('''blocks'''  ,      '''videomae.encoder.layer'''   )
  if "attn.proj" in name:
    _UpperCamelCase       			=name.replace('''attn.proj'''  ,      '''attention.output.dense'''   )
  if "attn" in name and "bias" not in name:
    _UpperCamelCase       			=name.replace('''attn'''  ,      '''attention.self'''   )
  if "attn" in name:
    _UpperCamelCase       			=name.replace('''attn'''  ,      '''attention.attention'''   )
  if "norm1" in name:
    _UpperCamelCase       			=name.replace('''norm1'''  ,      '''layernorm_before'''   )
  if "norm2" in name:
    _UpperCamelCase       			=name.replace('''norm2'''  ,      '''layernorm_after'''   )
  if "mlp.fc1" in name:
    _UpperCamelCase       			=name.replace('''mlp.fc1'''  ,      '''intermediate.dense'''   )
  if "mlp.fc2" in name:
    _UpperCamelCase       			=name.replace('''mlp.fc2'''  ,      '''output.dense'''   )
  if "decoder_embed" in name:
    _UpperCamelCase       			=name.replace('''decoder_embed'''  ,      '''decoder.decoder_embed'''   )
  if "decoder_norm" in name:
    _UpperCamelCase       			=name.replace('''decoder_norm'''  ,      '''decoder.decoder_norm'''   )
  if "decoder_pred" in name:
    _UpperCamelCase       			=name.replace('''decoder_pred'''  ,      '''decoder.decoder_pred'''   )
  if "norm.weight" in name and "decoder" not in name and "fc" not in name:
    _UpperCamelCase       			=name.replace('''norm.weight'''  ,      '''videomae.layernorm.weight'''   )
  if "norm.bias" in name and "decoder" not in name and "fc" not in name:
    _UpperCamelCase       			=name.replace('''norm.bias'''  ,      '''videomae.layernorm.bias'''   )
  if "head" in name and "decoder" not in name:
    _UpperCamelCase       			=name.replace('''head'''  ,      '''classifier'''   )

  return name


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  for key in orig_state_dict.copy().keys():
    _UpperCamelCase       			=orig_state_dict.pop(__SCREAMING_SNAKE_CASE   )

    if key.startswith('''encoder.'''   ):
      _UpperCamelCase       			=key.replace('''encoder.'''  ,      ''''''   )

    if "qkv" in key:
      _UpperCamelCase       			=key.split('''.'''   )
      if key.startswith('''decoder.blocks'''   ):
        _UpperCamelCase       			=config.decoder_hidden_size
        _UpperCamelCase       			=int(key_split[2]   )
        _UpperCamelCase       			='''decoder.decoder_layers.'''
        if "weight" in key:
          _UpperCamelCase       			=val[:dim, :]
          _UpperCamelCase       			=val[dim : dim * 2, :]
          _UpperCamelCase       			=val[-dim:, :]
      else:
        _UpperCamelCase       			=config.hidden_size
        _UpperCamelCase       			=int(key_split[1]   )
        _UpperCamelCase       			='''videomae.encoder.layer.'''
        if "weight" in key:
          _UpperCamelCase       			=val[:dim, :]
          _UpperCamelCase       			=val[dim : dim * 2, :]
          _UpperCamelCase       			=val[-dim:, :]
    else:
      _UpperCamelCase       			=val

  return orig_state_dict


def 					_a   ():


  """simple docstring"""

  _UpperCamelCase       			=hf_hub_download(
      repo_id='''hf-internal-testing/spaghetti-video'''  ,      filename='''eating_spaghetti.npy'''  ,      repo_type='''dataset'''   )
  _UpperCamelCase       			=np.load(__SCREAMING_SNAKE_CASE   )
  return list(__SCREAMING_SNAKE_CASE   )


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=get_videomae_config(__SCREAMING_SNAKE_CASE   )

  if "finetuned" in model_name:
    _UpperCamelCase       			=VideoMAEForVideoClassification(__SCREAMING_SNAKE_CASE   )
  else:
    _UpperCamelCase       			=VideoMAEForPreTraining(__SCREAMING_SNAKE_CASE   )

  # download original checkpoint, hosted on Google Drive
  _UpperCamelCase       			='''pytorch_model.bin'''
  gdown.cached_download(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      quiet=__SCREAMING_SNAKE_CASE   )
  _UpperCamelCase       			=torch.load(__SCREAMING_SNAKE_CASE  ,      map_location='''cpu'''   )
  if "model" in files:
    _UpperCamelCase       			=files['''model''']
  else:
    _UpperCamelCase       			=files['''module''']
  _UpperCamelCase       			=convert_state_dict(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )

  model.load_state_dict(__SCREAMING_SNAKE_CASE   )
  model.eval()

  # verify model on basic input
  _UpperCamelCase       			=VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]  ,      image_std=[0.5, 0.5, 0.5]   )
  _UpperCamelCase       			=prepare_video()
  _UpperCamelCase       			=image_processor(__SCREAMING_SNAKE_CASE  ,      return_tensors='''pt'''   )

  if "finetuned" not in model_name:
    _UpperCamelCase       			=hf_hub_download(repo_id='''hf-internal-testing/bool-masked-pos'''  ,      filename='''bool_masked_pos.pt'''   )
    _UpperCamelCase       			=torch.load(__SCREAMING_SNAKE_CASE   )

  _UpperCamelCase       			=model(**__SCREAMING_SNAKE_CASE   )
  _UpperCamelCase       			=outputs.logits

  _UpperCamelCase       			=[
      '''videomae-small-finetuned-kinetics''',
      '''videomae-small-finetuned-ssv2''',
      # Kinetics-400 checkpoints (short = pretrained only for 800 epochs instead of 1600)
      '''videomae-base-short''',
      '''videomae-base-short-finetuned-kinetics''',
      '''videomae-base''',
      '''videomae-base-finetuned-kinetics''',
      '''videomae-large''',
      '''videomae-large-finetuned-kinetics''',
      '''videomae-huge-finetuned-kinetics''',
      # Something-Something-v2 checkpoints (short = pretrained only for 800 epochs instead of 2400)
      '''videomae-base-short-ssv2''',
      '''videomae-base-short-finetuned-ssv2''',
      '''videomae-base-ssv2''',
      '''videomae-base-finetuned-ssv2''',
  ]

  # NOTE: logits were tested with image_mean and image_std equal to [0.5, 0.5, 0.5] and [0.5, 0.5, 0.5]
  if model_name == "videomae-small-finetuned-kinetics":
    _UpperCamelCase       			=torch.Size([1, 400]   )
    _UpperCamelCase       			=torch.tensor([-0.9_2_9_1, -0.4_0_6_1, -0.9_3_0_7]   )
  elif model_name == "videomae-small-finetuned-ssv2":
    _UpperCamelCase       			=torch.Size([1, 174]   )
    _UpperCamelCase       			=torch.tensor([0.2_6_7_1, -0.4_6_8_9, -0.8_2_3_5]   )
  elif model_name == "videomae-base":
    _UpperCamelCase       			=torch.Size([1, 1408, 1536]   )
    _UpperCamelCase       			=torch.tensor([[0.7_7_3_9, 0.7_9_6_8, 0.7_0_8_9], [0.6_7_0_1, 0.7_4_8_7, 0.6_2_0_9], [0.4_2_8_7, 0.5_1_5_8, 0.4_7_7_3]]   )
  elif model_name == "videomae-base-short":
    _UpperCamelCase       			=torch.Size([1, 1408, 1536]   )
    _UpperCamelCase       			=torch.tensor([[0.7_9_9_4, 0.9_6_1_2, 0.8_5_0_8], [0.7_4_0_1, 0.8_9_5_8, 0.8_3_0_2], [0.5_8_6_2, 0.7_4_6_8, 0.7_3_2_5]]   )
    # we verified the loss both for normalized and unnormalized targets for this one
    _UpperCamelCase       			=torch.tensor([0.5_1_4_2]   ) if config.norm_pix_loss else torch.tensor([0.6_4_6_9]   )
  elif model_name == "videomae-large":
    _UpperCamelCase       			=torch.Size([1, 1408, 1536]   )
    _UpperCamelCase       			=torch.tensor([[0.7_1_4_9, 0.7_9_9_7, 0.6_9_6_6], [0.6_7_6_8, 0.7_8_6_9, 0.6_9_4_8], [0.5_1_3_9, 0.6_2_2_1, 0.5_6_0_5]]   )
  elif model_name == "videomae-large-finetuned-kinetics":
    _UpperCamelCase       			=torch.Size([1, 400]   )
    _UpperCamelCase       			=torch.tensor([0.0_7_7_1, 0.0_0_1_1, -0.3_6_2_5]   )
  elif model_name == "videomae-huge-finetuned-kinetics":
    _UpperCamelCase       			=torch.Size([1, 400]   )
    _UpperCamelCase       			=torch.tensor([0.2_4_3_3, 0.1_6_3_2, -0.4_8_9_4]   )
  elif model_name == "videomae-base-short-finetuned-kinetics":
    _UpperCamelCase       			=torch.Size([1, 400]   )
    _UpperCamelCase       			=torch.tensor([0.6_5_8_8, 0.0_9_9_0, -0.2_4_9_3]   )
  elif model_name == "videomae-base-finetuned-kinetics":
    _UpperCamelCase       			=torch.Size([1, 400]   )
    _UpperCamelCase       			=torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]   )
  elif model_name == "videomae-base-short-ssv2":
    _UpperCamelCase       			=torch.Size([1, 1408, 1536]   )
    _UpperCamelCase       			=torch.tensor([[0.4_7_1_2, 0.5_2_9_6, 0.5_7_8_6], [0.2_2_7_8, 0.2_7_2_9, 0.4_0_2_6], [0.0_3_5_2, 0.0_7_3_0, 0.2_5_0_6]]   )
  elif model_name == "videomae-base-short-finetuned-ssv2":
    _UpperCamelCase       			=torch.Size([1, 174]   )
    _UpperCamelCase       			=torch.tensor([-0.0_5_3_7, -0.1_5_3_9, -0.3_2_6_6]   )
  elif model_name == "videomae-base-ssv2":
    _UpperCamelCase       			=torch.Size([1, 1408, 1536]   )
    _UpperCamelCase       			=torch.tensor([[0.8_1_3_1, 0.8_7_2_7, 0.8_5_4_6], [0.7_3_6_6, 0.9_3_7_7, 0.8_8_7_0], [0.5_9_3_5, 0.8_8_7_4, 0.8_5_6_4]]   )
  elif model_name == "videomae-base-finetuned-ssv2":
    _UpperCamelCase       			=torch.Size([1, 174]   )
    _UpperCamelCase       			=torch.tensor([0.1_9_6_1, -0.8_3_3_7, -0.6_3_8_9]   )
  else:
    raise ValueError(f'''Model name not supported. Should be one of {model_names}'''   )

  # verify logits
  assert logits.shape == expected_shape
  if "finetuned" in model_name:
    assert torch.allclose(logits[0, :3]  ,      __SCREAMING_SNAKE_CASE  ,      atol=1E-4   )
  else:
    print('''Logits:'''  ,      logits[0, :3, :3]   )
    assert torch.allclose(logits[0, :3, :3]  ,      __SCREAMING_SNAKE_CASE  ,      atol=1E-4   )
  print('''Logits ok!'''   )

  # verify loss, if applicable
  if model_name == "videomae-base-short":
    _UpperCamelCase       			=outputs.loss
    assert torch.allclose(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      atol=1E-4   )
    print('''Loss ok!'''   )

  if pytorch_dump_folder_path is not None:
    print(f'''Saving model and image processor to {pytorch_dump_folder_path}'''   )
    image_processor.save_pretrained(__SCREAMING_SNAKE_CASE   )
    model.save_pretrained(__SCREAMING_SNAKE_CASE   )

  if push_to_hub:
    print('''Pushing to the hub...'''   )
    model.push_to_hub(__SCREAMING_SNAKE_CASE  ,      organization='''nielsr'''   )


if __name__ == "__main__":
      __lowerCamelCase						: int		=   argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          '--checkpoint_url',
          default='https://drive.google.com/u/1/uc?id=1tEhLyskjb755TJ65ptsrafUG2llSwQE1&amp;export=download&amp;confirm=t&amp;uuid=aa3276eb-fb7e-482a-adec-dc7171df14c4',
          type=str,
          help=(
              'URL of the original PyTorch checkpoint (on Google Drive) you\'d like to convert. Should be a direct'
              ' download link.'
          ),
      )
      parser.add_argument(
          '--pytorch_dump_folder_path',
          default='/Users/nielsrogge/Documents/VideoMAE/Test',
          type=str,
          help='Path to the output PyTorch model directory.',
      )
      parser.add_argument('--model_name', default='videomae-base', type=str, help='Name of the model.')
      parser.add_argument(
          '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
      )

      __lowerCamelCase						: int		=   parser.parse_args()
      convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)













'''simple docstring'''






import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np







class 		UpperCAmelCase   :


  """simple docstring"""







  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Union[str, Any]:
    torch.manual_seed(0  )
    _UpperCamelCase       			=TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=UNetaDConditionModel(
        sample_size=32							,	layers_per_block=1							,	block_out_channels=[32, 64]							,	down_block_types=[
            '''ResnetDownsampleBlock2D''',
            '''SimpleCrossAttnDownBlock2D''',
        ]							,	mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,	up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,	in_channels=3							,	out_channels=6							,	cross_attention_dim=32							,	encoder_hid_dim=32							,	attention_head_dim=8							,	addition_embed_type='''text'''							,	addition_embed_type_num_heads=2							,	cross_attention_norm='''group_norm'''							,	resnet_time_scale_shift='''scale_shift'''							,	act_fn='''gelu'''							,	)
    unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    torch.manual_seed(0  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	thresholding=UpperCamelCase__							,	dynamic_thresholding_ratio=0.95							,	sample_max_value=1.0							,	prediction_type='''epsilon'''							,	variance_type='''learned_range'''							,	)

    torch.manual_seed(0  )
    _UpperCamelCase       			=IFWatermarker()

    return {
        "text_encoder": text_encoder,
        "tokenizer": tokenizer,
        "unet": unet,
        "scheduler": scheduler,
        "watermarker": watermarker,
        "safety_checker": None,
        "feature_extractor": None,
    }


  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Optional[int]:
    torch.manual_seed(0  )
    _UpperCamelCase       			=TaEncoderModel.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-t5'''  )

    torch.manual_seed(0  )
    _UpperCamelCase       			=UNetaDConditionModel(
        sample_size=32							,	layers_per_block=[1, 2]							,	block_out_channels=[32, 64]							,	down_block_types=[
            '''ResnetDownsampleBlock2D''',
            '''SimpleCrossAttnDownBlock2D''',
        ]							,	mid_block_type='''UNetMidBlock2DSimpleCrossAttn'''							,	up_block_types=['''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D''']							,	in_channels=6							,	out_channels=6							,	cross_attention_dim=32							,	encoder_hid_dim=32							,	attention_head_dim=8							,	addition_embed_type='''text'''							,	addition_embed_type_num_heads=2							,	cross_attention_norm='''group_norm'''							,	resnet_time_scale_shift='''scale_shift'''							,	act_fn='''gelu'''							,	class_embed_type='''timestep'''							,	mid_block_scale_factor=1.414							,	time_embedding_act_fn='''gelu'''							,	time_embedding_dim=32							,	)
    unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    torch.manual_seed(0  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	thresholding=UpperCamelCase__							,	dynamic_thresholding_ratio=0.95							,	sample_max_value=1.0							,	prediction_type='''epsilon'''							,	variance_type='''learned_range'''							,	)

    torch.manual_seed(0  )
    _UpperCamelCase       			=DDPMScheduler(
        num_train_timesteps=1000							,	beta_schedule='''squaredcos_cap_v2'''							,	beta_start=0.0001							,	beta_end=0.02							,	)

    torch.manual_seed(0  )
    _UpperCamelCase       			=IFWatermarker()

    return {
        "text_encoder": text_encoder,
        "tokenizer": tokenizer,
        "unet": unet,
        "scheduler": scheduler,
        "image_noising_scheduler": image_noising_scheduler,
        "watermarker": watermarker,
        "safety_checker": None,
        "feature_extractor": None,
    }


  def   UpperCamelCase__      (				self    :					Any  )     ->						Any:
    _UpperCamelCase       			=self.get_dummy_components()
    _UpperCamelCase       			=self.pipeline_class(**UpperCamelCase__  )
    pipe.to(UpperCamelCase__  )
    pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )

    _UpperCamelCase       			=inputs['''prompt''']
    _UpperCamelCase       			=inputs['''generator''']
    _UpperCamelCase       			=inputs['''num_inference_steps''']
    _UpperCamelCase       			=inputs['''output_type''']

    if "image" in inputs:
      _UpperCamelCase       			=inputs['''image''']
    else:
      _UpperCamelCase       			=None

    if "mask_image" in inputs:
      _UpperCamelCase       			=inputs['''mask_image''']
    else:
      _UpperCamelCase       			=None

    if "original_image" in inputs:
      _UpperCamelCase       			=inputs['''original_image''']
    else:
      _UpperCamelCase       			=None

    _UpperCamelCase			,     _UpperCamelCase       			=pipe.encode_prompt(UpperCamelCase__  )

    # inputs with prompt converted to embeddings
    _UpperCamelCase       			={
        '''prompt_embeds''': prompt_embeds,
        '''negative_prompt_embeds''': negative_prompt_embeds,
        '''generator''': generator,
        '''num_inference_steps''': num_inference_steps,
        '''output_type''': output_type,
    }

    if image is not None:
      _UpperCamelCase       			=image

    if mask_image is not None:
      _UpperCamelCase       			=mask_image

    if original_image is not None:
      _UpperCamelCase       			=original_image

    # set all optional components to None
    for optional_component in pipe._optional_components:
      setattr(UpperCamelCase__							,	UpperCamelCase__							,	UpperCamelCase__  )

    _UpperCamelCase       			=pipe(**UpperCamelCase__  )[0]

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

    pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    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.'''							,	)

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )

    _UpperCamelCase       			=inputs['''generator''']
    _UpperCamelCase       			=inputs['''num_inference_steps''']
    _UpperCamelCase       			=inputs['''output_type''']

    # inputs with prompt converted to embeddings
    _UpperCamelCase       			={
        '''prompt_embeds''': prompt_embeds,
        '''negative_prompt_embeds''': negative_prompt_embeds,
        '''generator''': generator,
        '''num_inference_steps''': num_inference_steps,
        '''output_type''': output_type,
    }

    if image is not None:
      _UpperCamelCase       			=image

    if mask_image is not None:
      _UpperCamelCase       			=mask_image

    if original_image is not None:
      _UpperCamelCase       			=original_image

    _UpperCamelCase       			=pipe_loaded(**UpperCamelCase__  )[0]

    _UpperCamelCase       			=np.abs(to_np(UpperCamelCase__  ) - to_np(UpperCamelCase__  )  ).max()
    self.assertLess(UpperCamelCase__							,	1E-4  )




  def   UpperCamelCase__      (				self    :					int  )     ->						str:
    _UpperCamelCase       			=self.get_dummy_components()
    _UpperCamelCase       			=self.pipeline_class(**UpperCamelCase__  )
    pipe.to(UpperCamelCase__  )
    pipe.set_progress_bar_config(disable=UpperCamelCase__  )

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )
    _UpperCamelCase       			=pipe(**UpperCamelCase__  )[0]

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

    pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()  )  # For reproducibility tests

    _UpperCamelCase       			=self.get_dummy_inputs(UpperCamelCase__  )
    _UpperCamelCase       			=pipe_loaded(**UpperCamelCase__  )[0]

    _UpperCamelCase       			=np.abs(to_np(UpperCamelCase__  ) - to_np(UpperCamelCase__  )  ).max()
    self.assertLess(UpperCamelCase__							,	1E-4  )






