






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__ (   __SCREAMING_SNAKE_CASE	, unittest.TestCase     ):







    """simple docstring"""

    __magic_name__        =		DanceDiffusionPipeline
    __magic_name__        =		UNCONDITIONAL_AUDIO_GENERATION_PARAMS
    __magic_name__        =		PipelineTesterMixin.required_optional_params - {
        'callback',
        'latents',
        'callback_steps',
        'output_type',
        'num_images_per_prompt',
    }
    __magic_name__        =		UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
    __magic_name__        =		False
    __magic_name__        =		False
    def 		a_      (	self   ):
         torch.manual_seed(0   )
         snake_case					=    UNetaDModel(
             block_out_channels=(3_2, 3_2, 6_4)      , extra_in_channels=1_6      , sample_size=5_1_2      , sample_rate=1_6_0_0_0      , in_channels=2      , out_channels=2      , flip_sin_to_cos=UpperCamelCase__      , use_timestep_embedding=UpperCamelCase__      , time_embedding_type='''fourier'''      , mid_block_type='''UNetMidBlock1D'''      , down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')      , up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')      , )
         snake_case					=    IPNDMScheduler()

         snake_case					=    {
             '''unet''': unet,
             '''scheduler''': scheduler,
         }
         return components
    def 		a_      (	self      , __snake_case      , __snake_case=0   ):







         if str(UpperCamelCase__   ).startswith('''mps'''   ):
              snake_case					=    torch.manual_seed(UpperCamelCase__   )
         else:
              snake_case					=    torch.Generator(device=UpperCamelCase__   ).manual_seed(UpperCamelCase__   )
         snake_case					=    {
             '''batch_size''': 1,
             '''generator''': generator,
             '''num_inference_steps''': 4,
         }
         return inputs
    def 		a_      (	self   ):
         snake_case					=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
         snake_case					=    self.get_dummy_components()
         snake_case					=    DanceDiffusionPipeline(**UpperCamelCase__   )
         snake_case					=    pipe.to(UpperCamelCase__   )
         pipe.set_progress_bar_config(disable=UpperCamelCase__   )

         snake_case					=    self.get_dummy_inputs(UpperCamelCase__   )
         snake_case					=    pipe(**UpperCamelCase__   )
         snake_case					=    output.audios

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

         assert audio.shape == (1, 2, components["unet"].sample_size)
         snake_case					=    np.array([-0.7265, 1.0000, -0.8388, 0.1175, 0.9498, -1.0000]   )
         assert np.abs(audio_slice.flatten() - expected_slice   ).max() < 1E-2
    @skip_mps
    def 		a_      (	self   ):
         return super().test_save_load_local()
    @skip_mps
    def 		a_      (	self   ):
         return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3   )
    @skip_mps
    def 		a_      (	self   ):
         return super().test_save_load_optional_components()
    @skip_mps
    def 		a_      (	self   ):
         return super().test_attention_slicing_forward_pass()




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






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







    """simple docstring"""

    def 		a_      (	self   ):
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()
    def 		a_      (	self   ):
         snake_case					=    torch_device

         snake_case					=    DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''   )
         snake_case					=    pipe.to(UpperCamelCase__   )
         pipe.set_progress_bar_config(disable=UpperCamelCase__   )

         snake_case					=    torch.manual_seed(0   )
         snake_case					=    pipe(generator=UpperCamelCase__      , num_inference_steps=1_0_0      , audio_length_in_s=4.096   )
         snake_case					=    output.audios

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

         assert audio.shape == (1, 2, pipe.unet.sample_size)
         snake_case					=    np.array([-0.0192, -0.0231, -0.0318, -0.0059, 0.0002, -0.0020]   )

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




    def 		a_      (	self   ):
         snake_case					=    torch_device

         snake_case					=    DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''      , torch_dtype=torch.floataa   )
         snake_case					=    pipe.to(UpperCamelCase__   )
         pipe.set_progress_bar_config(disable=UpperCamelCase__   )

         snake_case					=    torch.manual_seed(0   )
         snake_case					=    pipe(generator=UpperCamelCase__      , num_inference_steps=1_0_0      , audio_length_in_s=4.096   )
         snake_case					=    output.audios

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

         assert audio.shape == (1, 2, pipe.unet.sample_size)
         snake_case					=    np.array([-0.0367, -0.0488, -0.0771, -0.0525, -0.0444, -0.0341]   )

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








from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


_SCREAMING_SNAKE_CASE      :			Dict      				=   {"configuration_van": ["VAN_PRETRAINED_CONFIG_ARCHIVE_MAP", "VanConfig"]}


try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _SCREAMING_SNAKE_CASE      :			Dict      				=   [
     "VAN_PRETRAINED_MODEL_ARCHIVE_LIST",
     "VanForImageClassification",
     "VanModel",
     "VanPreTrainedModel",
 ]

if TYPE_CHECKING:
 from .configuration_van import VAN_PRETRAINED_CONFIG_ARCHIVE_MAP, VanConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_van import (
      VAN_PRETRAINED_MODEL_ARCHIVE_LIST,
      VanForImageClassification,
      VanModel,
      VanPreTrainedModel,
  )

else:
 import sys

 _SCREAMING_SNAKE_CASE      :			str      				=   _LazyModule(__name__, globals()["__file__"], _import_structure)




