



"""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    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  unittest.TestCase   ):
    lowercase__      =   DanceDiffusionPipeline
    lowercase__      =   UNCONDITIONAL_AUDIO_GENERATION_PARAMS
    lowercase__      =   PipelineTesterMixin.required_optional_params - {
        "callback",
        "latents",
        "callback_steps",
        "output_type",
        "num_images_per_prompt",
    }
    lowercase__      =   UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
    lowercase__      =   False
    lowercase__      =   False

    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      torch.manual_seed(0)
      lowercase_								= 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=lowerCAmelCase_ ,  use_timestep_embedding=lowerCAmelCase_ ,  time_embedding_type="""fourier""" ,  mid_block_type="""UNetMidBlock1D""" ,  down_block_types=("""DownBlock1DNoSkip""", """DownBlock1D""", """AttnDownBlock1D""") ,  up_block_types=("""AttnUpBlock1D""", """UpBlock1D""", """UpBlock1DNoSkip""") ,  )
      lowercase_								= IPNDMScheduler()

      lowercase_								= {
          """unet""": unet,
          """scheduler""": scheduler,
      }
      return components

    def 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				Dict ,  lowerCAmelCase_			:				List[Any]=0):







      """simple docstring"""







      if str(lowerCAmelCase_).startswith("""mps"""):
        lowercase_								= torch.manual_seed(lowerCAmelCase_)
      else:
        lowercase_								= torch.Generator(device=lowerCAmelCase_).manual_seed(lowerCAmelCase_)
      lowercase_								= {
          """batch_size""": 1,
          """generator""": generator,
          """num_inference_steps""": 4,
      }
      return inputs

    def 					_UpperCAmelCase	(				self			:				Optional[int]):







      """simple docstring"""







      lowercase_								= """cpu"""  # ensure determinism for the device-dependent torch.Generator
      lowercase_								= self.get_dummy_components()
      lowercase_								= DanceDiffusionPipeline(**lowerCAmelCase_)
      lowercase_								= pipe.to(lowerCAmelCase_)
      pipe.set_progress_bar_config(disable=lowerCAmelCase_)

      lowercase_								= self.get_dummy_inputs(lowerCAmelCase_)
      lowercase_								= pipe(**lowerCAmelCase_)
      lowercase_								= output.audios

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

      assert audio.shape == (1, 2, components["unet"].sample_size)
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      return super().test_save_load_local()

    @skip_mps
    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3)

    @skip_mps
    def 					_UpperCAmelCase	(				self			:				List[str]):







      """simple docstring"""







      return super().test_save_load_optional_components()

    @skip_mps
    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







      return super().test_attention_slicing_forward_pass()




    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







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







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

    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







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

    def 					_UpperCAmelCase	(				self			:				Union[str, Any]):







      """simple docstring"""







      lowercase_								= torch_device

      lowercase_								= DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k""")
      lowercase_								= pipe.to(lowerCAmelCase_)
      pipe.set_progress_bar_config(disable=lowerCAmelCase_)

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

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

      assert audio.shape == (1, 2, pipe.unet.sample_size)
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      lowercase_								= torch_device

      lowercase_								= DanceDiffusionPipeline.from_pretrained("""harmonai/maestro-150k""" ,  torch_dtype=torch.floataa)
      lowercase_								= pipe.to(lowerCAmelCase_)
      pipe.set_progress_bar_config(disable=lowerCAmelCase_)

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

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

      assert audio.shape == (1, 2, pipe.unet.sample_size)
      lowercase_								= 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






"""simple docstring"""
from typing import List

import datasets
from datasets.tasks import AudioClassification

from ..folder_based_builder import folder_based_builder


UpperCAmelCase						:		Union[str, Any]			= datasets.utils.logging.get_logger(__name__)


class    SCREAMING_SNAKE_CASE__ (							folder_based_builder.FolderBasedBuilderConfig   ):
    lowercase__      =   None
    lowercase__      =   None







class    SCREAMING_SNAKE_CASE__ (							folder_based_builder.FolderBasedBuilder   ):
    lowercase__      =   datasets.Audio()
    lowercase__      =   "audio"
    lowercase__      =   AudioFolderConfig
    lowercase__      =   42  # definition at the bottom of the script
    lowercase__      =   AudioClassification(audio_column="audio"      ,  label_column="label"   )







UpperCAmelCase						:		Tuple			= [
    ".aiff",
    ".au",
    ".avr",
    ".caf",
    ".flac",
    ".htk",
    ".svx",
    ".mat4",
    ".mat5",
    ".mpc2k",
    ".ogg",
    ".paf",
    ".pvf",
    ".raw",
    ".rf64",
    ".sd2",
    ".sds",
    ".ircam",
    ".voc",
    ".w64",
    ".wav",
    ".nist",
    ".wavex",
    ".wve",
    ".xi",
    ".mp3",
    ".opus",
]
UpperCAmelCase						:		Any			= AUDIO_EXTENSIONS

