




'''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   lowercase__(			__UpperCamelCase:					Optional[int]    ):

      """simple docstring"""






      SCREAMING_SNAKE_CASE   :      Any    					=						VideoMAEConfig()

      set_architecture_configs(__UpperCamelCase      ,__UpperCamelCase    )

      if "finetuned" not in model_name:
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						False

      if "finetuned" in model_name:
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						'huggingface/label-files'
            if "kinetics" in model_name:
                  SCREAMING_SNAKE_CASE   :      Tuple    					=						4_00
                  SCREAMING_SNAKE_CASE   :      str    					=						'kinetics400-id2label.json'
            elif "ssv2" in model_name:
                  SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						1_74
                  SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						'something-something-v2-id2label.json'
            else:
                  raise ValueError('Model name should either contain \'kinetics\' or \'ssv2\' in case it\'s fine-tuned.'    )
            SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						json.load(open(hf_hub_download(__UpperCamelCase      ,__UpperCamelCase      ,repo_type='dataset'    )      ,'r'    )    )
            SCREAMING_SNAKE_CASE   :      Any    					=						{int(__UpperCamelCase    ): v for k, v in idalabel.items()}
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						idalabel
            SCREAMING_SNAKE_CASE   :      str    					=						{v: k for k, v in idalabel.items()}

      return config




def   lowercase__(			__UpperCamelCase:					List[str]      ,__UpperCamelCase:					Any    ):

      """simple docstring"""






      if "small" in model_name:
            SCREAMING_SNAKE_CASE   :      Tuple    					=						3_84
            SCREAMING_SNAKE_CASE   :      int    					=						15_36
            SCREAMING_SNAKE_CASE   :      int    					=						12
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						16
            SCREAMING_SNAKE_CASE   :      int    					=						12
            SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						3
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						1_92
            SCREAMING_SNAKE_CASE   :      int    					=						7_68
      elif "large" in model_name:
            SCREAMING_SNAKE_CASE   :      Any    					=						10_24
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						40_96
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						24
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						16
            SCREAMING_SNAKE_CASE   :      List[str]    					=						12
            SCREAMING_SNAKE_CASE   :      Any    					=						8
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						5_12
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						20_48
      elif "huge" in model_name:
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						12_80
            SCREAMING_SNAKE_CASE   :      List[str]    					=						51_20
            SCREAMING_SNAKE_CASE   :      Tuple    					=						32
            SCREAMING_SNAKE_CASE   :      List[Any]    					=						16
            SCREAMING_SNAKE_CASE   :      Any    					=						12
            SCREAMING_SNAKE_CASE   :      str    					=						8
            SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						6_40
            SCREAMING_SNAKE_CASE   :      Tuple    					=						25_60
      elif "base" not in model_name:
            raise ValueError('Model name should include either "small", "base", "large", or "huge"'    )




def   lowercase__(			__UpperCamelCase:					str    ):

      """simple docstring"""






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

      return name




def   lowercase__(			__UpperCamelCase:					str      ,__UpperCamelCase:					str    ):

      """simple docstring"""






      for key in orig_state_dict.copy().keys():
            SCREAMING_SNAKE_CASE   :      int    					=						orig_state_dict.pop(__UpperCamelCase    )

            if key.startswith('encoder.'    ):
                  SCREAMING_SNAKE_CASE   :      List[str]    					=						key.replace('encoder.'      ,''    )

            if "qkv" in key:
                  SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						key.split('.'    )
                  if key.startswith('decoder.blocks'    ):
                        SCREAMING_SNAKE_CASE   :      Any    					=						config.decoder_hidden_size
                        SCREAMING_SNAKE_CASE   :      str    					=						int(key_split[2]    )
                        SCREAMING_SNAKE_CASE   :      Dict    					=						'decoder.decoder_layers.'
                        if "weight" in key:
                              SCREAMING_SNAKE_CASE   :      int    					=						val[:dim, :]
                              SCREAMING_SNAKE_CASE   :      Any    					=						val[dim : dim * 2, :]
                              SCREAMING_SNAKE_CASE   :      str    					=						val[-dim:, :]
                  else:
                        SCREAMING_SNAKE_CASE   :      Optional[int]    					=						config.hidden_size
                        SCREAMING_SNAKE_CASE   :      Optional[int]    					=						int(key_split[1]    )
                        SCREAMING_SNAKE_CASE   :      Tuple    					=						'videomae.encoder.layer.'
                        if "weight" in key:
                              SCREAMING_SNAKE_CASE   :      Any    					=						val[:dim, :]
                              SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						val[dim : dim * 2, :]
                              SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						val[-dim:, :]
            else:
                  SCREAMING_SNAKE_CASE   :      List[Any]    					=						val

      return orig_state_dict




def   lowercase__(			):

      """simple docstring"""






      SCREAMING_SNAKE_CASE   :      Any    					=						hf_hub_download(
          repo_id='hf-internal-testing/spaghetti-video'      ,filename='eating_spaghetti.npy'      ,repo_type='dataset'    )
      SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						np.load(__UpperCamelCase    )
      return list(__UpperCamelCase    )




def   lowercase__(			__UpperCamelCase:					Optional[int]      ,__UpperCamelCase:					List[str]      ,__UpperCamelCase:					Union[str, Any]      ,__UpperCamelCase:					Dict    ):

      """simple docstring"""






      SCREAMING_SNAKE_CASE   :      Tuple    					=						get_videomae_config(__UpperCamelCase    )

      if "finetuned" in model_name:
            SCREAMING_SNAKE_CASE   :      Tuple    					=						VideoMAEForVideoClassification(__UpperCamelCase    )
      else:
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						VideoMAEForPreTraining(__UpperCamelCase    )

      # download original checkpoint, hosted on Google Drive
      SCREAMING_SNAKE_CASE   :      List[str]    					=						'pytorch_model.bin'
      gdown.cached_download(__UpperCamelCase      ,__UpperCamelCase      ,quiet=__UpperCamelCase    )
      SCREAMING_SNAKE_CASE   :      Dict    					=						torch.load(__UpperCamelCase      ,map_location='cpu'    )
      if "model" in files:
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						files['model']
      else:
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						files['module']
      SCREAMING_SNAKE_CASE   :      Optional[int]    					=						convert_state_dict(__UpperCamelCase      ,__UpperCamelCase    )

      model.load_state_dict(__UpperCamelCase    )
      model.eval()

      # verify model on basic input
      SCREAMING_SNAKE_CASE   :      Tuple    					=						VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5]      ,image_std=[0.5, 0.5, 0.5]    )
      SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						prepare_video()
      SCREAMING_SNAKE_CASE   :      Any    					=						image_processor(__UpperCamelCase      ,return_tensors='pt'    )

      if "finetuned" not in model_name:
            SCREAMING_SNAKE_CASE   :      int    					=						hf_hub_download(repo_id='hf-internal-testing/bool-masked-pos'      ,filename='bool_masked_pos.pt'    )
            SCREAMING_SNAKE_CASE   :      List[Any]    					=						torch.load(__UpperCamelCase    )

      SCREAMING_SNAKE_CASE   :      int    					=						model(**__UpperCamelCase    )
      SCREAMING_SNAKE_CASE   :      Optional[int]    					=						outputs.logits

      SCREAMING_SNAKE_CASE   :      int    					=						[
          '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":
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						torch.Size([1, 4_00]    )
            SCREAMING_SNAKE_CASE   :      int    					=						torch.tensor([-0.9_2_9_1, -0.4_0_6_1, -0.9_3_0_7]    )
      elif model_name == "videomae-small-finetuned-ssv2":
            SCREAMING_SNAKE_CASE   :      Any    					=						torch.Size([1, 1_74]    )
            SCREAMING_SNAKE_CASE   :      str    					=						torch.tensor([0.2_6_7_1, -0.4_6_8_9, -0.8_2_3_5]    )
      elif model_name == "videomae-base":
            SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.Size([1, 14_08, 15_36]    )
            SCREAMING_SNAKE_CASE   :      Any    					=						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":
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						torch.Size([1, 14_08, 15_36]    )
            SCREAMING_SNAKE_CASE   :      List[Any]    					=						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
            SCREAMING_SNAKE_CASE   :      str    					=						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":
            SCREAMING_SNAKE_CASE   :      List[str]    					=						torch.Size([1, 14_08, 15_36]    )
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						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":
            SCREAMING_SNAKE_CASE   :      List[str]    					=						torch.Size([1, 4_00]    )
            SCREAMING_SNAKE_CASE   :      List[Any]    					=						torch.tensor([0.0_7_7_1, 0.0_0_1_1, -0.3_6_2_5]    )
      elif model_name == "videomae-huge-finetuned-kinetics":
            SCREAMING_SNAKE_CASE   :      int    					=						torch.Size([1, 4_00]    )
            SCREAMING_SNAKE_CASE   :      List[str]    					=						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":
            SCREAMING_SNAKE_CASE   :      Any    					=						torch.Size([1, 4_00]    )
            SCREAMING_SNAKE_CASE   :      List[str]    					=						torch.tensor([0.6_5_8_8, 0.0_9_9_0, -0.2_4_9_3]    )
      elif model_name == "videomae-base-finetuned-kinetics":
            SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.Size([1, 4_00]    )
            SCREAMING_SNAKE_CASE   :      List[Any]    					=						torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]    )
      elif model_name == "videomae-base-short-ssv2":
            SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.Size([1, 14_08, 15_36]    )
            SCREAMING_SNAKE_CASE   :      Dict    					=						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":
            SCREAMING_SNAKE_CASE   :      Any    					=						torch.Size([1, 1_74]    )
            SCREAMING_SNAKE_CASE   :      Any    					=						torch.tensor([-0.0_5_3_7, -0.1_5_3_9, -0.3_2_6_6]    )
      elif model_name == "videomae-base-ssv2":
            SCREAMING_SNAKE_CASE   :      Optional[int]    					=						torch.Size([1, 14_08, 15_36]    )
            SCREAMING_SNAKE_CASE   :      str    					=						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":
            SCREAMING_SNAKE_CASE   :      List[str]    					=						torch.Size([1, 1_74]    )
            SCREAMING_SNAKE_CASE   :      int    					=						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]      ,__UpperCamelCase      ,atol=1e-4    )
      else:
            print('Logits:'      ,logits[0, :3, :3]    )
            assert torch.allclose(logits[0, :3, :3]      ,__UpperCamelCase      ,atol=1e-4    )
      print('Logits ok!'    )

      # verify loss, if applicable
      if model_name == "videomae-base-short":
            SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						outputs.loss
            assert torch.allclose(__UpperCamelCase      ,__UpperCamelCase      ,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(__UpperCamelCase    )
            model.save_pretrained(__UpperCamelCase    )

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


if __name__ == "__main__":
      UpperCamelCase_    	=				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."
      )

      UpperCamelCase_    	=				parser.parse_args()
      convert_videomae_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)








'''simple docstring'''



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 							_a				(			SCREAMING_SNAKE_CASE ,      SCREAMING_SNAKE_CASE ,      unittest.TestCase ):





       '''simple docstring'''






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







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             torch.manual_seed(0  )
             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						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=A,							)
             SCREAMING_SNAKE_CASE   :      int    					=						DDIMScheduler(
                 beta_start=0.0_00_85,							beta_end=0.0_12,							beta_schedule='scaled_linear',							clip_sample=A,							set_alpha_to_one=A,							)
             SCREAMING_SNAKE_CASE   :      str    					=						DDIMInverseScheduler(
                 beta_start=0.0_00_85,							beta_end=0.0_12,							beta_schedule='scaled_linear',							clip_sample=A,							set_alpha_to_zero=A,							)
             torch.manual_seed(0  )
             SCREAMING_SNAKE_CASE   :      Dict    					=						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  )
             SCREAMING_SNAKE_CASE   :      Tuple    					=						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=1_000,							hidden_act='gelu',							projection_dim=512,							)
             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						CLIPTextModel(A  )
             SCREAMING_SNAKE_CASE   :      str    					=						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )

             SCREAMING_SNAKE_CASE   :      int    					=						{
                 '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,							A,							A=0  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						floats_tensor((1, 16, 16),							rng=random.Random(A  )  ).to(A  )
             SCREAMING_SNAKE_CASE   :      List[str]    					=						floats_tensor((1, 2, 4, 16, 16),							rng=random.Random(A  )  ).to(A  )
             if str(A  ).startswith('mps'  ):
                   SCREAMING_SNAKE_CASE   :      List[str]    					=						torch.manual_seed(A  )
             else:
                   SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.Generator(device=A  ).manual_seed(A  )
             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						{
                 '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,							A,							A=0  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      str    					=						floats_tensor((1, 3, 32, 32),							rng=random.Random(A  )  ).to(A  )
             SCREAMING_SNAKE_CASE   :      Any    					=						image.cpu().permute(0,							2,							3,							1  )[0]
             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						Image.fromarray(np.uinta(A  )  ).convert('RGB'  )
             if str(A  ).startswith('mps'  ):
                   SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.manual_seed(A  )
             else:
                   SCREAMING_SNAKE_CASE   :      int    					=						torch.Generator(device=A  ).manual_seed(A  )
             SCREAMING_SNAKE_CASE   :      Dict    					=						{
                 '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,							A,							A=0  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Any    					=						floats_tensor((1, 3, 32, 32),							rng=random.Random(A  )  ).to(A  )
             SCREAMING_SNAKE_CASE   :      List[Any]    					=						image.cpu().permute(0,							2,							3,							1  )[0]
             SCREAMING_SNAKE_CASE   :      int    					=						Image.fromarray(np.uinta(A  )  ).convert('RGB'  )
             if str(A  ).startswith('mps'  ):
                   SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						torch.manual_seed(A  )
             else:
                   SCREAMING_SNAKE_CASE   :      Any    					=						torch.Generator(device=A  ).manual_seed(A  )
             SCREAMING_SNAKE_CASE   :      Any    					=						{
                 '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  ):







             '''simple docstring'''







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

             SCREAMING_SNAKE_CASE   :      Tuple    					=						self.get_dummy_components()
             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						self.pipeline_class(**A  )
             pipe.to(A  )
             pipe.set_progress_bar_config(disable=A  )

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

             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						self.get_dummy_inputs(A  )
             SCREAMING_SNAKE_CASE   :      Dict    					=						pipe(**A  )[0]

             with tempfile.TemporaryDirectory() as tmpdir:
                   pipe.save_pretrained(A  )
                   SCREAMING_SNAKE_CASE   :      List[Any]    					=						self.pipeline_class.from_pretrained(A  )
                   pipe_loaded.to(A  )
                   pipe_loaded.set_progress_bar_config(disable=A  )

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

             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						self.get_dummy_inputs(A  )
             SCREAMING_SNAKE_CASE   :      Tuple    					=						pipe_loaded(**A  )[0]

             SCREAMING_SNAKE_CASE   :      List[str]    					=						np.abs(output - output_loaded  ).max()
             self.assertLess(A,							1E-4  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      str    					=						'cpu'

             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						self.get_dummy_components()
             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						self.pipeline_class(**A  )
             pipe.to(A  )
             pipe.set_progress_bar_config(disable=A  )

             SCREAMING_SNAKE_CASE   :      str    					=						self.get_dummy_mask_inputs(A  )
             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						pipe.generate_mask(**A  )
             SCREAMING_SNAKE_CASE   :      Dict    					=						mask[0, -3:, -3:]

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







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Tuple    					=						'cpu'

             SCREAMING_SNAKE_CASE   :      Dict    					=						self.get_dummy_components()
             SCREAMING_SNAKE_CASE   :      Dict    					=						self.pipeline_class(**A  )
             pipe.to(A  )
             pipe.set_progress_bar_config(disable=A  )

             SCREAMING_SNAKE_CASE   :      Tuple    					=						self.get_dummy_inversion_inputs(A  )
             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						pipe.invert(**A  ).images
             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						image[0, -1, -3:, -3:]

             self.assertEqual(image.shape,							(2, 32, 32, 3)  )
             SCREAMING_SNAKE_CASE   :      Tuple    					=						np.array(
                 [0.51_50, 0.51_34, 0.50_43, 0.53_76, 0.46_94, 0.5_10_50, 0.50_15, 0.44_07, 0.47_99],							)
             SCREAMING_SNAKE_CASE   :      Dict    					=						np.abs(image_slice.flatten() - expected_slice  ).max()
             self.assertLessEqual(A,							1E-3  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             super().test_inference_batch_single_identical(expected_max_diff=5E-3  )







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Tuple    					=						'cpu'

             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						self.get_dummy_components()

             SCREAMING_SNAKE_CASE   :      Dict    					=						{'beta_start': 0.0_00_85, 'beta_end': 0.0_12, 'beta_schedule': 'scaled_linear'}
             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						DPMSolverMultistepScheduler(**A  )
             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						DPMSolverMultistepInverseScheduler(**A  )

             SCREAMING_SNAKE_CASE   :      Tuple    					=						self.pipeline_class(**A  )
             pipe.to(A  )
             pipe.set_progress_bar_config(disable=A  )

             SCREAMING_SNAKE_CASE   :      Tuple    					=						self.get_dummy_inversion_inputs(A  )
             SCREAMING_SNAKE_CASE   :      List[str]    					=						pipe.invert(**A  ).images
             SCREAMING_SNAKE_CASE   :      Optional[Any]    					=						image[0, -1, -3:, -3:]

             self.assertEqual(image.shape,							(2, 32, 32, 3)  )
             SCREAMING_SNAKE_CASE   :      Tuple    					=						np.array(
                 [0.51_50, 0.51_34, 0.50_43, 0.53_76, 0.46_94, 0.5_10_50, 0.50_15, 0.44_07, 0.47_99],							)
             SCREAMING_SNAKE_CASE   :      Any    					=						np.abs(image_slice.flatten() - expected_slice  ).max()
             self.assertLessEqual(A,							1E-3  )






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





       '''simple docstring'''






       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()







       @classmethod
       def        UpperCamelCase_							(      cls  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						load_image(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'  )

             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						raw_image.convert('RGB'  ).resize((768, 768)  )

             SCREAMING_SNAKE_CASE   :      List[str]    					=						raw_image







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Tuple    					=						torch.manual_seed(0  )

             SCREAMING_SNAKE_CASE   :      Dict    					=						StableDiffusionDiffEditPipeline.from_pretrained(
                 'stabilityai/stable-diffusion-2-1',							safety_checker=A,							torch_dtype=torch.floataa  )
             SCREAMING_SNAKE_CASE   :      List[Any]    					=						DDIMScheduler.from_config(pipe.scheduler.config  )
             SCREAMING_SNAKE_CASE   :      int    					=						DDIMInverseScheduler.from_config(pipe.scheduler.config  )
             pipe.enable_model_cpu_offload()
             pipe.set_progress_bar_config(disable=A  )

             SCREAMING_SNAKE_CASE   :      List[Any]    					=						'a bowl of fruit'
             SCREAMING_SNAKE_CASE   :      List[str]    					=						'a bowl of pears'

             SCREAMING_SNAKE_CASE   :      Dict    					=						pipe.generate_mask(
                 image=self.raw_image,							source_prompt=A,							target_prompt=A,							generator=A,							)

             SCREAMING_SNAKE_CASE   :      Optional[int]    					=						pipe.invert(
                 prompt=A,							image=self.raw_image,							inpaint_strength=0.7,							generator=A  ).latents

             SCREAMING_SNAKE_CASE   :      List[str]    					=						pipe(
                 prompt=A,							mask_image=A,							image_latents=A,							generator=A,							negative_prompt=A,							inpaint_strength=0.7,							output_type='numpy',							).images[0]

             SCREAMING_SNAKE_CASE   :      List[Any]    					=						(
                 np.array(
                     load_image(
                         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                         '/diffedit/pears.png'  ).resize((768, 768)  )  )
                 / 255
             )
             assert np.abs((expected_image - image).max()  ) < 5E-1







       def        UpperCamelCase_							(      self  ):







             '''simple docstring'''







             SCREAMING_SNAKE_CASE   :      Dict    					=						torch.manual_seed(0  )

             SCREAMING_SNAKE_CASE   :      int    					=						StableDiffusionDiffEditPipeline.from_pretrained(
                 'stabilityai/stable-diffusion-2-1',							safety_checker=A,							torch_dtype=torch.floataa  )
             SCREAMING_SNAKE_CASE   :      List[str]    					=						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config  )
             SCREAMING_SNAKE_CASE   :      List[str]    					=						DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config  )
             pipe.enable_model_cpu_offload()
             pipe.set_progress_bar_config(disable=A  )

             SCREAMING_SNAKE_CASE   :      str    					=						'a bowl of fruit'
             SCREAMING_SNAKE_CASE   :      Tuple    					=						'a bowl of pears'

             SCREAMING_SNAKE_CASE   :      List[Any]    					=						pipe.generate_mask(
                 image=self.raw_image,							source_prompt=A,							target_prompt=A,							generator=A,							)

             SCREAMING_SNAKE_CASE   :      Union[str, Any]    					=						pipe.invert(
                 prompt=A,							image=self.raw_image,							inpaint_strength=0.7,							generator=A,							num_inference_steps=25,							).latents

             SCREAMING_SNAKE_CASE   :      str    					=						pipe(
                 prompt=A,							mask_image=A,							image_latents=A,							generator=A,							negative_prompt=A,							inpaint_strength=0.7,							num_inference_steps=25,							output_type='numpy',							).images[0]

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


