


"""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__    (		lowercase       ,   unittest.TestCase    ):
       """simple docstring"""


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

       def 	lowercase__							(						self						):






         """simple docstring"""





         torch.manual_seed(0						)
         lowerCAmelCase					:					Optional[int]								=							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=snake_case__		,					use_timestep_embedding=snake_case__		,					time_embedding_type="fourier"		,					mid_block_type="UNetMidBlock1D"		,					down_block_types=("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")		,					up_block_types=("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")		,					)
         lowerCAmelCase					:					List[str]								=							IPNDMScheduler()

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

       def 	lowercase__							(						self		,					snake_case__		,					snake_case__=0						):






         """simple docstring"""





         if str(snake_case__						).startswith("mps"						):
           lowerCAmelCase					:					List[str]								=							torch.manual_seed(snake_case__						)
         else:
           lowerCAmelCase					:					Optional[int]								=							torch.Generator(device=snake_case__						).manual_seed(snake_case__						)
         lowerCAmelCase					:					Tuple								=							{
             "batch_size": 1,
             "generator": generator,
             "num_inference_steps": 4,
         }
         return inputs

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Dict								=							"cpu"  # ensure determinism for the device-dependent torch.Generator
         lowerCAmelCase					:					Tuple								=							self.get_dummy_components()
         lowerCAmelCase					:					List[Any]								=							DanceDiffusionPipeline(**snake_case__						)
         lowerCAmelCase					:					List[Any]								=							pipe.to(snake_case__						)
         pipe.set_progress_bar_config(disable=snake_case__						)

         lowerCAmelCase					:					int								=							self.get_dummy_inputs(snake_case__						)
         lowerCAmelCase					:					Optional[Any]								=							pipe(**snake_case__						)
         lowerCAmelCase					:					List[Any]								=							output.audios

         lowerCAmelCase					:					Optional[int]								=							audio[0, -3:, -3:]

         assert audio.shape == (1, 2, components["unet"].sample_size)
         lowerCAmelCase					:					List[Any]								=							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 	lowercase__							(						self						):






         """simple docstring"""





         return super().test_save_load_local()

       @skip_mps
       def 	lowercase__							(						self						):






         """simple docstring"""





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

       @skip_mps
       def 	lowercase__							(						self						):






         """simple docstring"""





         return super().test_save_load_optional_components()

       @skip_mps
       def 	lowercase__							(						self						):






         """simple docstring"""





         return super().test_attention_slicing_forward_pass()







       def 	lowercase__							(						self						):






         """simple docstring"""





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


@slow
@require_torch_gpu
class 				SCREAMING_SNAKE_CASE__    (		unittest.TestCase    ):
       """simple docstring"""


       def 	lowercase__							(						self						):






         """simple docstring"""





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

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Optional[int]								=							torch_device

         lowerCAmelCase					:					Optional[Any]								=							DanceDiffusionPipeline.from_pretrained("harmonai/maestro-150k"						)
         lowerCAmelCase					:					Optional[Any]								=							pipe.to(snake_case__						)
         pipe.set_progress_bar_config(disable=snake_case__						)

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

         lowerCAmelCase					:					Optional[int]								=							audio[0, -3:, -3:]

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






         """simple docstring"""





         lowerCAmelCase					:					Optional[Any]								=							torch_device

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

         lowerCAmelCase					:					Tuple								=							torch.manual_seed(0						)
         lowerCAmelCase					:					List[str]								=							pipe(generator=snake_case__		,					num_inference_steps=100		,					audio_length_in_s=4.096						)
         lowerCAmelCase					:					Any								=							output.audios

         lowerCAmelCase					:					Optional[Any]								=							audio[0, -3:, -3:]

         assert audio.shape == (1, 2, pipe.unet.sample_size)
         lowerCAmelCase					:					List[str]								=							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









"""simple docstring"""







import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible


class 				SCREAMING_SNAKE_CASE__    (		unittest.TestCase    ):
       """simple docstring"""


       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					int								=							[
             "safety_checker/pytorch_model.bin",
             "safety_checker/model.safetensors",
             "vae/diffusion_pytorch_model.bin",
             "vae/diffusion_pytorch_model.safetensors",
             "text_encoder/pytorch_model.bin",
             "text_encoder/model.safetensors",
             "unet/diffusion_pytorch_model.bin",
             "unet/diffusion_pytorch_model.safetensors",
         ]
         self.assertTrue(is_safetensors_compatible(snake_case__						)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Dict								=							[
             "unet/diffusion_pytorch_model.bin",
             "unet/diffusion_pytorch_model.safetensors",
         ]
         self.assertTrue(is_safetensors_compatible(snake_case__						)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Optional[Any]								=							[
             "safety_checker/pytorch_model.bin",
             "safety_checker/model.safetensors",
             "vae/diffusion_pytorch_model.bin",
             "vae/diffusion_pytorch_model.safetensors",
             "text_encoder/pytorch_model.bin",
             "text_encoder/model.safetensors",
             "unet/diffusion_pytorch_model.bin",
             # Removed: 'unet/diffusion_pytorch_model.safetensors',
         ]
         self.assertFalse(is_safetensors_compatible(snake_case__						)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					List[str]								=							[
             "text_encoder/pytorch_model.bin",
             "text_encoder/model.safetensors",
         ]
         self.assertTrue(is_safetensors_compatible(snake_case__						)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							[
             "safety_checker/pytorch_model.bin",
             "safety_checker/model.safetensors",
             "vae/diffusion_pytorch_model.bin",
             "vae/diffusion_pytorch_model.safetensors",
             "text_encoder/pytorch_model.bin",
             # Removed: 'text_encoder/model.safetensors',
             "unet/diffusion_pytorch_model.bin",
             "unet/diffusion_pytorch_model.safetensors",
         ]
         self.assertFalse(is_safetensors_compatible(snake_case__						)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					List[Any]								=							[
             "safety_checker/pytorch_model.fp16.bin",
             "safety_checker/model.fp16.safetensors",
             "vae/diffusion_pytorch_model.fp16.bin",
             "vae/diffusion_pytorch_model.fp16.safetensors",
             "text_encoder/pytorch_model.fp16.bin",
             "text_encoder/model.fp16.safetensors",
             "unet/diffusion_pytorch_model.fp16.bin",
             "unet/diffusion_pytorch_model.fp16.safetensors",
         ]
         lowerCAmelCase					:					Union[str, Any]								=							"fp16"
         self.assertTrue(is_safetensors_compatible(snake_case__		,					variant=snake_case__						)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Optional[int]								=							[
             "unet/diffusion_pytorch_model.fp16.bin",
             "unet/diffusion_pytorch_model.fp16.safetensors",
         ]
         lowerCAmelCase					:					Tuple								=							"fp16"
         self.assertTrue(is_safetensors_compatible(snake_case__		,					variant=snake_case__						)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					List[Any]								=							[
             "unet/diffusion_pytorch_model.bin",
             "unet/diffusion_pytorch_model.safetensors",
         ]
         lowerCAmelCase					:					Union[str, Any]								=							"fp16"
         self.assertTrue(is_safetensors_compatible(snake_case__		,					variant=snake_case__						)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					List[Any]								=							[
             "safety_checker/pytorch_model.fp16.bin",
             "safety_checker/model.fp16.safetensors",
             "vae/diffusion_pytorch_model.fp16.bin",
             "vae/diffusion_pytorch_model.fp16.safetensors",
             "text_encoder/pytorch_model.fp16.bin",
             "text_encoder/model.fp16.safetensors",
             "unet/diffusion_pytorch_model.fp16.bin",
             # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
         ]
         lowerCAmelCase					:					Dict								=							"fp16"
         self.assertFalse(is_safetensors_compatible(snake_case__		,					variant=snake_case__						)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					str								=							[
             "text_encoder/pytorch_model.fp16.bin",
             "text_encoder/model.fp16.safetensors",
         ]
         lowerCAmelCase					:					int								=							"fp16"
         self.assertTrue(is_safetensors_compatible(snake_case__		,					variant=snake_case__						)						)

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					str								=							[
             "text_encoder/pytorch_model.bin",
             "text_encoder/model.safetensors",
         ]
         lowerCAmelCase					:					Tuple								=							"fp16"
         self.assertTrue(is_safetensors_compatible(snake_case__		,					variant=snake_case__						)						)







       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					int								=							[
             "safety_checker/pytorch_model.fp16.bin",
             "safety_checker/model.fp16.safetensors",
             "vae/diffusion_pytorch_model.fp16.bin",
             "vae/diffusion_pytorch_model.fp16.safetensors",
             "text_encoder/pytorch_model.fp16.bin",
             # 'text_encoder/model.fp16.safetensors',
             "unet/diffusion_pytorch_model.fp16.bin",
             "unet/diffusion_pytorch_model.fp16.safetensors",
         ]
         lowerCAmelCase					:					List[Any]								=							"fp16"
         self.assertFalse(is_safetensors_compatible(snake_case__		,					variant=snake_case__						)						)





