



"""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"""
import os
from collections.abc import Iterator


def 			_SCREAMING_SNAKE_CASE						(__lowerCAmelCase = "."							) ->						Iterator[str]:







  '''simple docstring'''


  for dir_path, dir_names, filenames in os.walk(__lowerCAmelCase							):
    lowercase_								= [d for d in dir_names if d != """scripts""" and d[0] not in """._"""]
    for filename in filenames:
      if filename == "__init__.py":
        continue
      if os.path.splitext(__lowerCAmelCase							)[1] in (".py", ".ipynb"):
        yield os.path.join(__lowerCAmelCase ,  __lowerCAmelCase							).lstrip("""./"""							)


def 			_SCREAMING_SNAKE_CASE						(__lowerCAmelCase							) ->						Any:







  '''simple docstring'''


  return F'''{i * "  "}*''' if i else "\n##"


def 			_SCREAMING_SNAKE_CASE						(__lowerCAmelCase ,  __lowerCAmelCase							) ->						str:







  '''simple docstring'''


  lowercase_								= old_path.split(os.sep							)
  for i, new_part in enumerate(new_path.split(os.sep							)							):
    if (i + 1 > len(__lowerCAmelCase							) or old_parts[i] != new_part) and new_part:
      print(F'''{md_prefix(__lowerCAmelCase							)} {new_part.replace("_" ,  " "							).title()}'''							)
  return new_path


def 			_SCREAMING_SNAKE_CASE						(__lowerCAmelCase = "."							) ->						None:







  '''simple docstring'''


  lowercase_								= """"""
  for filepath in sorted(good_file_paths(__lowerCAmelCase							)							):
    lowercase_					,    lowercase_								= os.path.split(__lowerCAmelCase							)
    if filepath != old_path:
      lowercase_								= print_path(__lowerCAmelCase ,  __lowerCAmelCase							)
    lowercase_								= (filepath.count(os.sep							) + 1) if filepath else 0
    lowercase_								= F'''{filepath}/{filename}'''.replace(""" """ ,  """%20"""							)
    lowercase_								= os.path.splitext(filename.replace("""_""" ,  """ """							).title()							)[0]
    print(F'''{md_prefix(__lowerCAmelCase							)} [{filename}]({url})'''							)


if __name__ == "__main__":
       print_directory_md(".")

