





'''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 lowerCamelCase						(snake_case_						,      unittest.TestCase					):
  _lowercase  :	Union[str, Any]       					=				DanceDiffusionPipeline
  _lowercase  :	Optional[Any]       					=				UNCONDITIONAL_AUDIO_GENERATION_PARAMS
  _lowercase  :	List[Any]       					=				PipelineTesterMixin.required_optional_params - {
      """callback""",
      """latents""",
      """callback_steps""",
      """output_type""",
      """num_images_per_prompt""",
  }
  _lowercase  :	List[str]       					=				UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
  _lowercase  :	Any       					=				False
  _lowercase  :	List[Any]       					=				False







  def 				UpperCAmelCase_				(		self		)		->						Optional[int]:






        """simple docstring"""





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

        _snake_case		:	Union[str, Any]					     =     {
            '''unet''': unet,
            '''scheduler''': scheduler,
        }
        return components







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__=0		)		->						List[str]:






        """simple docstring"""





        if str(lowercase__		).startswith('''mps'''		):
              _snake_case		:	Any					     =     torch.manual_seed(lowercase__		)
        else:
              _snake_case		:	List[Any]					     =     torch.Generator(device=lowercase__		).manual_seed(lowercase__		)
        _snake_case		:	Dict					     =     {
            '''batch_size''': 1,
            '''generator''': generator,
            '''num_inference_steps''': 4,
        }
        return inputs







  def 				UpperCAmelCase_				(		self		)		->						List[Any]:






        """simple docstring"""





        _snake_case		:	Union[str, Any]					     =     '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _snake_case		:	int					     =     self.get_dummy_components()
        _snake_case		:	Dict					     =     DanceDiffusionPipeline(**lowercase__		)
        _snake_case		:	Union[str, Any]					     =     pipe.to(lowercase__		)
        pipe.set_progress_bar_config(disable=lowercase__		)

        _snake_case		:	Tuple					     =     self.get_dummy_inputs(lowercase__		)
        _snake_case		:	Any					     =     pipe(**lowercase__		)
        _snake_case		:	Tuple					     =     output.audios

        _snake_case		:	int					     =     audio[0, -3:, -3:]

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






        """simple docstring"""





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







  @skip_mps
  def 				UpperCAmelCase_				(		self		)		->						Any:






        """simple docstring"""





        return super().test_save_load_optional_components()







  @skip_mps
  def 				UpperCAmelCase_				(		self		)		->						Tuple:






        """simple docstring"""





        return super().test_attention_slicing_forward_pass()







  def 				UpperCAmelCase_				(		self		)		->						List[str]:






        """simple docstring"""





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





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







  def 				UpperCAmelCase_				(		self		)		->						Union[str, Any]:






        """simple docstring"""





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







  def 				UpperCAmelCase_				(		self		)		->						Any:






        """simple docstring"""





        _snake_case		:	List[str]					     =     torch_device

        _snake_case		:	Union[str, Any]					     =     DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''		)
        _snake_case		:	Tuple					     =     pipe.to(lowercase__		)
        pipe.set_progress_bar_config(disable=lowercase__		)

        _snake_case		:	List[Any]					     =     torch.manual_seed(0		)
        _snake_case		:	Dict					     =     pipe(generator=lowercase__     ,     num_inference_steps=100     ,     audio_length_in_s=4.096		)
        _snake_case		:	Optional[Any]					     =     output.audios

        _snake_case		:	Union[str, Any]					     =     audio[0, -3:, -3:]

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






        """simple docstring"""





        _snake_case		:	List[str]					     =     torch_device

        _snake_case		:	List[str]					     =     DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''     ,     torch_dtype=torch.floataa		)
        _snake_case		:	Tuple					     =     pipe.to(lowercase__		)
        pipe.set_progress_bar_config(disable=lowercase__		)

        _snake_case		:	int					     =     torch.manual_seed(0		)
        _snake_case		:	int					     =     pipe(generator=lowercase__     ,     num_inference_steps=100     ,     audio_length_in_s=4.096		)
        _snake_case		:	List[Any]					     =     output.audios

        _snake_case		:	List[str]					     =     audio[0, -3:, -3:]

        assert audio.shape == (1, 2, pipe.unet.sample_size)
        _snake_case		:	Dict					     =     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 functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase					:		List[str]		=							logging.get_logger(__name__)

UpperCAmelCase					:		Dict		=							{
    'asapp/sew-d-tiny-100k': 'https://huggingface.co/asapp/sew-d-tiny-100k/resolve/main/config.json',
    # See all SEW-D models at https://huggingface.co/models?filter=sew-d
}





class lowerCamelCase						(a__					):
  _lowercase  :	List[str]       					=				"""sew-d"""







  def __init__(		self     ,     lowercase__=32     ,     lowercase__=768     ,     lowercase__=12     ,     lowercase__=12     ,     lowercase__=3_072     ,     lowercase__=2     ,     lowercase__=512     ,     lowercase__=256     ,     lowercase__=True     ,     lowercase__=True     ,     lowercase__=("p2c", "c2p")     ,     lowercase__="layer_norm"     ,     lowercase__="gelu_python"     ,     lowercase__=0.1     ,     lowercase__=0.1     ,     lowercase__=0.1     ,     lowercase__=0.0     ,     lowercase__=0.1     ,     lowercase__=0.02     ,     lowercase__=1E-7     ,     lowercase__=1E-5     ,     lowercase__="group"     ,     lowercase__="gelu"     ,     lowercase__=(64, 128, 128, 128, 128, 256, 256, 256, 256, 512, 512, 512, 512)     ,     lowercase__=(5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1)     ,     lowercase__=(10, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1)     ,     lowercase__=False     ,     lowercase__=128     ,     lowercase__=16     ,     lowercase__=True     ,     lowercase__=0.05     ,     lowercase__=10     ,     lowercase__=2     ,     lowercase__=0.0     ,     lowercase__=10     ,     lowercase__=0     ,     lowercase__="mean"     ,     lowercase__=False     ,     lowercase__=False     ,     lowercase__=256     ,     lowercase__=0     ,     lowercase__=1     ,     lowercase__=2     ,     **lowercase__     ,     )		->						Dict:






        """simple docstring"""





        super().__init__(**lowercase__     ,     pad_token_id=lowercase__     ,     bos_token_id=lowercase__     ,     eos_token_id=lowercase__		)
        _snake_case		:	List[str]					     =     hidden_size
        _snake_case		:	Optional[Any]					     =     feat_extract_norm
        _snake_case		:	Tuple					     =     feat_extract_activation
        _snake_case		:	Tuple					     =     list(lowercase__		)
        _snake_case		:	Any					     =     list(lowercase__		)
        _snake_case		:	Any					     =     list(lowercase__		)
        _snake_case		:	Any					     =     conv_bias
        _snake_case		:	List[Any]					     =     num_conv_pos_embeddings
        _snake_case		:	Any					     =     num_conv_pos_embedding_groups
        _snake_case		:	Union[str, Any]					     =     len(self.conv_dim		)
        _snake_case		:	Optional[Any]					     =     num_hidden_layers
        _snake_case		:	Optional[int]					     =     intermediate_size
        _snake_case		:	Any					     =     squeeze_factor
        _snake_case		:	Optional[Any]					     =     max_position_embeddings
        _snake_case		:	Tuple					     =     position_buckets
        _snake_case		:	Tuple					     =     share_att_key
        _snake_case		:	Any					     =     relative_attention
        _snake_case		:	Optional[int]					     =     norm_rel_ebd
        _snake_case		:	Optional[Any]					     =     list(lowercase__		)
        _snake_case		:	List[Any]					     =     hidden_act
        _snake_case		:	List[Any]					     =     num_attention_heads
        _snake_case		:	Dict					     =     hidden_dropout
        _snake_case		:	Tuple					     =     attention_dropout
        _snake_case		:	Union[str, Any]					     =     activation_dropout
        _snake_case		:	List[Any]					     =     feat_proj_dropout
        _snake_case		:	Optional[int]					     =     final_dropout
        _snake_case		:	Optional[Any]					     =     layer_norm_eps
        _snake_case		:	Dict					     =     feature_layer_norm_eps
        _snake_case		:	List[Any]					     =     initializer_range
        _snake_case		:	Dict					     =     vocab_size

        if (
            (len(self.conv_stride		) != self.num_feat_extract_layers)
            or (len(self.conv_kernel		) != self.num_feat_extract_layers)
            or (len(self.conv_dim		) != self.num_feat_extract_layers)
        ):
              raise ValueError(
                  '''Configuration for convolutional layers is incorrect.'''
                  '''It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`,'''
                  F'''but is `len(config.conv_dim) = {len(self.conv_dim		)}`, `len(config.conv_stride)'''
                  F'''= {len(self.conv_stride		)}`, `len(config.conv_kernel) = {len(self.conv_kernel		)}`.'''		)

        # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
        _snake_case		:	Union[str, Any]					     =     apply_spec_augment
        _snake_case		:	Any					     =     mask_time_prob
        _snake_case		:	List[str]					     =     mask_time_length
        _snake_case		:	Dict					     =     mask_time_min_masks
        _snake_case		:	Union[str, Any]					     =     mask_feature_prob
        _snake_case		:	Tuple					     =     mask_feature_length
        _snake_case		:	Union[str, Any]					     =     mask_feature_min_masks

        # ctc loss
        _snake_case		:	Optional[Any]					     =     ctc_loss_reduction
        _snake_case		:	Optional[Any]					     =     ctc_zero_infinity

        # sequence classification
        _snake_case		:	List[Any]					     =     use_weighted_layer_sum
        _snake_case		:	Any					     =     classifier_proj_size







  @property
  def 				UpperCAmelCase_				(		self		)		->						Any:






        """simple docstring"""





        return functools.reduce(operator.mul     ,     self.conv_stride     ,     1		)


