
"""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_					(		a_	,							unittest.TestCase					):




		"""simple docstring"""





		__lowercase				:    Optional[int]		  =						DanceDiffusionPipeline
		__lowercase				:    List[str]		  =						UNCONDITIONAL_AUDIO_GENERATION_PARAMS
		__lowercase				:    Any		  =						PipelineTesterMixin.required_optional_params - {
		    '''callback''',
		    '''latents''',
		    '''callback_steps''',
		    '''output_type''',
		    '''num_images_per_prompt''',
		}
		__lowercase				:    Optional[Any]		  =						UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
		__lowercase				:    str		  =						False
		__lowercase				:    Union[str, Any]		  =						False




		def       snake_case_	(	self):
				torch.manual_seed(0)
				__SCREAMING_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""")				,    )
				__SCREAMING_SNAKE_CASE				=					IPNDMScheduler()

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




		def       snake_case_	(	self				,    lowerCAmelCase__				,    lowerCAmelCase__=0):

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




		def       snake_case_	(	self):
				__SCREAMING_SNAKE_CASE				=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
				__SCREAMING_SNAKE_CASE				=					self.get_dummy_components()
				__SCREAMING_SNAKE_CASE				=					DanceDiffusionPipeline(**UpperCamelCase_)
				__SCREAMING_SNAKE_CASE				=					pipe.to(UpperCamelCase_)
				pipe.set_progress_bar_config(disable=UpperCamelCase_)

				__SCREAMING_SNAKE_CASE				=					self.get_dummy_inputs(UpperCamelCase_)
				__SCREAMING_SNAKE_CASE				=					pipe(**UpperCamelCase_)
				__SCREAMING_SNAKE_CASE				=					output.audios

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

				assert audio.shape == (1, 2, components["unet"].sample_size)
				__SCREAMING_SNAKE_CASE				=					np.array([-0.72_65, 1.00_00, -0.83_88, 0.11_75, 0.94_98, -1.00_00])
				assert np.abs(audio_slice.flatten() - expected_slice).max() < 1E-2




		@skip_mps
		def       snake_case_	(	self):
				return super().test_save_load_local()




		@skip_mps
		def       snake_case_	(	self):
				return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3)




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




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




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



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




		"""simple docstring"""





		def       snake_case_	(	self):
				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()




		def       snake_case_	(	self):
				__SCREAMING_SNAKE_CASE				=					torch_device

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

				__SCREAMING_SNAKE_CASE				=					torch.manual_seed(0)
				__SCREAMING_SNAKE_CASE				=					pipe(generator=UpperCamelCase_				,    num_inference_steps=1_0_0				,    audio_length_in_s=4.0_96)
				__SCREAMING_SNAKE_CASE				=					output.audios

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

				assert audio.shape == (1, 2, pipe.unet.sample_size)
				__SCREAMING_SNAKE_CASE				=					np.array([-0.01_92, -0.02_31, -0.03_18, -0.00_59, 0.00_02, -0.00_20])

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




		def       snake_case_	(	self):
				__SCREAMING_SNAKE_CASE				=					torch_device

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

				__SCREAMING_SNAKE_CASE				=					torch.manual_seed(0)
				__SCREAMING_SNAKE_CASE				=					pipe(generator=UpperCamelCase_				,    num_inference_steps=1_0_0				,    audio_length_in_s=4.0_96)
				__SCREAMING_SNAKE_CASE				=					output.audios

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

				assert audio.shape == (1, 2, pipe.unet.sample_size)
				__SCREAMING_SNAKE_CASE				=					np.array([-0.03_67, -0.04_88, -0.07_71, -0.05_25, -0.04_44, -0.03_41])

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












"""simple docstring"""
from typing import List, Optional, Union

import numpy as np

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


snake_case : Dict     		=		logging.get_logger(__name__)



class        UpperCamelCase__ (				a_):







    """simple docstring"""
    __UpperCAmelCase           = ["""input_values""", """padding_mask"""]







    def __init__(		self :	Optional[int]				,    UpperCamelCase_ :	int = 1				,    UpperCamelCase_ :	int = 2_4_0_0_0				,    UpperCamelCase_ :	float = 0.0				,    UpperCamelCase_ :	float = None				,    UpperCamelCase_ :	float = None				,    **UpperCamelCase_ :	str				,    ):





           '''simple docstring'''

           super().__init__(feature_size=UpperCamelCase_				,    sampling_rate=UpperCamelCase_				,    padding_value=UpperCamelCase_				,    **UpperCamelCase_ )
           __magic_name__      				=    chunk_length_s
           __magic_name__      				=    overlap







    @property
    def 				a__			(		self :	List[str] ):





           '''simple docstring'''

           if self.chunk_length_s is None:
                  return None
           else:
                  return int(self.chunk_length_s * self.sampling_rate )







    @property
    def 				a__			(		self :	Optional[Any] ):





           '''simple docstring'''

           if self.chunk_length_s is None or self.overlap is None:
                  return None
           else:
                  return max(1				,    int((1.0 - self.overlap) * self.chunk_length ) )







    def __call__(		self :	List[Any]				,    UpperCamelCase_ :	Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]				,    UpperCamelCase_ :	Optional[Union[bool, str, PaddingStrategy]] = None				,    UpperCamelCase_ :	Optional[bool] = False				,    UpperCamelCase_ :	Optional[int] = None				,    UpperCamelCase_ :	Optional[Union[str, TensorType]] = None				,    UpperCamelCase_ :	Optional[int] = None				,    ):





           '''simple docstring'''

           if sampling_rate is not None:
                  if sampling_rate != self.sampling_rate:
                         raise ValueError(
                             f"""The model corresponding to this feature extractor: {self} was trained using a sampling rate of"""
                             f""" {self.sampling_rate}. Please make sure that the provided audio input was sampled with"""
                             f""" {self.sampling_rate} and not {sampling_rate}.""" )
           else:
                  logger.warning(
                      'It is strongly recommended to pass the `sampling_rate` argument to this function. '
                      'Failing to do so can result in silent errors that might be hard to debug.' )

           if padding and truncation:
                  raise ValueError('Both padding and truncation were set. Make sure you only set one.' )
           elif padding is None:
                  # by default let's pad the inputs
                  __magic_name__      				=    True

           __magic_name__      				=    bool(
               isinstance(UpperCamelCase_				,    (list, tuple) ) and (isinstance(raw_audio[0]				,    (np.ndarray, tuple, list) )) )

           if is_batched:
                  __magic_name__      				=    [np.asarray(UpperCamelCase_				,    dtype=np.floataa ).T for audio in raw_audio]
           elif not is_batched and not isinstance(UpperCamelCase_				,    np.ndarray ):
                  __magic_name__      				=    np.asarray(UpperCamelCase_				,    dtype=np.floataa )
           elif isinstance(UpperCamelCase_				,    np.ndarray ) and raw_audio.dtype is np.dtype(np.floataa ):
                  __magic_name__      				=    raw_audio.astype(np.floataa )

           # always return batch
           if not is_batched:
                  __magic_name__      				=    [np.asarray(UpperCamelCase_ ).T]

           # verify inputs are valid
           for idx, example in enumerate(UpperCamelCase_ ):
                  if example.ndim > 2:
                         raise ValueError(f"""Expected input shape (channels, length) but got shape {example.shape}""" )
                  if self.feature_size == 1 and example.ndim != 1:
                         raise ValueError(f"""Expected mono audio but example has {example.shape[-1]} channels""" )
                  if self.feature_size == 2 and example.shape[-1] != 2:
                         raise ValueError(f"""Expected stereo audio but example has {example.shape[-1]} channels""" )

           __magic_name__      				=    None
           __magic_name__      				=    BatchFeature({'input_values': raw_audio} )
           if self.chunk_stride is not None and self.chunk_length is not None and max_length is None:
                  if truncation:
                         __magic_name__      				=    min(array.shape[0] for array in raw_audio )
                         __magic_name__      				=    int(np.floor(max_length / self.chunk_stride ) )
                         __magic_name__      				=    (nb_step - 1) * self.chunk_stride + self.chunk_length
                  elif padding:
                         __magic_name__      				=    max(array.shape[0] for array in raw_audio )
                         __magic_name__      				=    int(np.ceil(max_length / self.chunk_stride ) )
                         __magic_name__      				=    (nb_step - 1) * self.chunk_stride + self.chunk_length
                         __magic_name__      				=    'max_length'
                  else:
                         __magic_name__      				=    input_values

        # normal padding on batch
           if padded_inputs is None:
                  __magic_name__      				=    self.pad(
                      UpperCamelCase_				,    max_length=UpperCamelCase_				,    truncation=UpperCamelCase_				,    padding=UpperCamelCase_				,    return_attention_mask=UpperCamelCase_				,    )
                  if padding:
                         __magic_name__      				=    padded_inputs.pop('attention_mask' )

           __magic_name__      				=    []
           for example in padded_inputs.pop('input_values' ):
                  if self.feature_size == 1:
                         __magic_name__      				=    example[..., None]
                  input_values.append(example.T )

           __magic_name__      				=    input_values
           if return_tensors is not None:
                  __magic_name__      				=    padded_inputs.convert_to_tensors(UpperCamelCase_ )

           return padded_inputs