




'''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 unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class 					_a (a__,		unittest.TestCase	):





     '''simple docstring'''



     lowerCAmelCase_  :     Dict          =					DanceDiffusionPipeline
     lowerCAmelCase_  :     int          =					UNCONDITIONAL_AUDIO_GENERATION_PARAMS
     lowerCAmelCase_  :     Optional[int]          =					PipelineTesterMixin.required_optional_params - {
         """callback""",
         """latents""",
         """callback_steps""",
         """output_type""",
         """num_images_per_prompt""",
     }
     lowerCAmelCase_  :     Tuple          =					UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
     lowerCAmelCase_  :     Dict          =					False
     lowerCAmelCase_  :     Tuple          =					False







     def 			snake_case_      (				self							)      ->				List[Any]:
          torch.manual_seed(0							)
          snake_case			:      List[str]        =			UNetaDModel(
              block_out_channels=(32, 32, 64)							,extra_in_channels=16							,sample_size=512							,sample_rate=16_000							,in_channels=2							,out_channels=2							,flip_sin_to_cos=__a							,use_timestep_embedding=__a							,time_embedding_type="""fourier"""							,mid_block_type="""UNetMidBlock1D"""							,down_block_types=("""DownBlock1DNoSkip""", """DownBlock1D""", """AttnDownBlock1D""")							,up_block_types=("""AttnUpBlock1D""", """UpBlock1D""", """UpBlock1DNoSkip""")							,)
          snake_case			:      Dict        =			IPNDMScheduler()

          snake_case			:      int        =			{
              """unet""": unet,
              """scheduler""": scheduler,
          }
          return components







     def 			snake_case_      (				self							,__a							,__a=0							)      ->				Tuple:
          if str(__a							).startswith("""mps"""							):
               snake_case			:      Optional[Any]        =			torch.manual_seed(__a							)
          else:
               snake_case			:      Union[str, Any]        =			torch.Generator(device=__a							).manual_seed(__a							)
          snake_case			:      str        =			{
              """batch_size""": 1,
              """generator""": generator,
              """num_inference_steps""": 4,
          }
          return inputs







     def 			snake_case_      (				self							)      ->				Any:
          snake_case			:      Dict        =			"""cpu"""  # ensure determinism for the device-dependent torch.Generator
          snake_case			:      str        =			self.get_dummy_components()
          snake_case			:      List[str]        =			DanceDiffusionPipeline(**__a							)
          snake_case			:      Dict        =			pipe.to(__a							)
          pipe.set_progress_bar_config(disable=__a							)

          snake_case			:      Union[str, Any]        =			self.get_dummy_inputs(__a							)
          snake_case			:      str        =			pipe(**__a							)
          snake_case			:      Union[str, Any]        =			output.audios

          snake_case			:      List[str]        =			audio[0, -3:, -3:]

          assert audio.shape == (1, 2, components["unet"].sample_size)
          snake_case			:      Optional[Any]        =			np.array([-0.7_265, 1.0_000, -0.8_388, 0.1_175, 0.9_498, -1.0_000]							)
          assert np.abs(audio_slice.flatten() - expected_slice							).max() < 1E-2







     @skip_mps
     def 			snake_case_      (				self							)      ->				List[str]:
          return super().test_save_load_local()







     @skip_mps
     def 			snake_case_      (				self							)      ->				Union[str, Any]:
          return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3							)







     @skip_mps
     def 			snake_case_      (				self							)      ->				Dict:
          return super().test_save_load_optional_components()







     @skip_mps
     def 			snake_case_      (				self							)      ->				Any:
          return super().test_attention_slicing_forward_pass()







     def 			snake_case_      (				self							)      ->				str:
          super().test_inference_batch_single_identical(expected_max_diff=3E-3							)







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





     '''simple docstring'''



     def 			snake_case_      (				self							)      ->				Optional[int]:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







     def 			snake_case_      (				self							)      ->				List[Any]:
          snake_case			:      str        =			torch_device

          snake_case			:      List[str]        =			DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""							)
          snake_case			:      Any        =			pipe.to(__a							)
          pipe.set_progress_bar_config(disable=__a							)

          snake_case			:      int        =			torch.manual_seed(0							)
          snake_case			:      int        =			pipe(generator=__a							,num_inference_steps=100							,audio_length_in_s=4.096							)
          snake_case			:      Dict        =			output.audios

          snake_case			:      Union[str, Any]        =			audio[0, -3:, -3:]

          assert audio.shape == (1, 2, pipe.unet.sample_size)
          snake_case			:      Tuple        =			np.array([-0.0_192, -0.0_231, -0.0_318, -0.0_059, 0.0_002, -0.0_020]							)

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







     def 			snake_case_      (				self							)      ->				str:
          snake_case			:      str        =			torch_device

          snake_case			:      List[Any]        =			DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k"""							,torch_dtype=torch.floataa							)
          snake_case			:      str        =			pipe.to(__a							)
          pipe.set_progress_bar_config(disable=__a							)

          snake_case			:      str        =			torch.manual_seed(0							)
          snake_case			:      List[str]        =			pipe(generator=__a							,num_inference_steps=100							,audio_length_in_s=4.096							)
          snake_case			:      Optional[int]        =			output.audios

          snake_case			:      Any        =			audio[0, -3:, -3:]

          assert audio.shape == (1, 2, pipe.unet.sample_size)
          snake_case			:      Optional[int]        =			np.array([-0.0_367, -0.0_488, -0.0_771, -0.0_525, -0.0_444, -0.0_341]							)

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



