



"""simple docstring"""

import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class  _UpperCAmelCase (    lowerCAmelCase__					,  unittest.TestCase):
 _lowerCAmelCase     :    Dict	      =	AudioLDMPipeline
 _lowerCAmelCase     :    Optional[Any]	      =	TEXT_TO_AUDIO_PARAMS
 _lowerCAmelCase     :    str	      =	TEXT_TO_AUDIO_BATCH_PARAMS
 _lowerCAmelCase     :    List[Any]	      =	frozenset(
     [
         """num_inference_steps""",
         """num_waveforms_per_prompt""",
         """generator""",
         """latents""",
         """output_type""",
         """return_dict""",
         """callback""",
         """callback_steps""",
     ])

 def 		_snake_case      (  self	: Dict       ):
     torch.manual_seed(0       )
     snake_case_				:  Union[str, Any]       =			UNetaDConditionModel(
         block_out_channels=(32, 64)						,  layers_per_block=2						,  sample_size=32						,  in_channels=4						,  out_channels=4						,  down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')						,  up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')						,  cross_attention_dim=(32, 64)						,  class_embed_type='''simple_projection'''						,  projection_class_embeddings_input_dim=32						,  class_embeddings_concat=lowercase_						,  )
     snake_case_				:  Dict       =			DDIMScheduler(
         beta_start=0.0_00_85						,  beta_end=0.0_12						,  beta_schedule='''scaled_linear'''						,  clip_sample=lowercase_						,  set_alpha_to_one=lowercase_						,  )
     torch.manual_seed(0       )
     snake_case_				:  Tuple       =			AutoencoderKL(
         block_out_channels=[32, 64]						,  in_channels=1						,  out_channels=1						,  down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']						,  up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']						,  latent_channels=4						,  )
     torch.manual_seed(0       )
     snake_case_				:  Optional[Any]       =			ClapTextConfig(
         bos_token_id=0						,  eos_token_id=2						,  hidden_size=32						,  intermediate_size=37						,  layer_norm_eps=1E-05						,  num_attention_heads=4						,  num_hidden_layers=5						,  pad_token_id=1						,  vocab_size=1000						,  projection_dim=32						,  )
     snake_case_				:  List[Any]       =			ClapTextModelWithProjection(lowercase_       )
     snake_case_				:  List[Any]       =			RobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-roberta'''						,  model_max_length=77       )

     snake_case_				:  str       =			SpeechTaHifiGanConfig(
         model_in_dim=8						,  sampling_rate=16000						,  upsample_initial_channel=16						,  upsample_rates=[2, 2]						,  upsample_kernel_sizes=[4, 4]						,  resblock_kernel_sizes=[3, 7]						,  resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]						,  normalize_before=lowercase_						,  )

     snake_case_				:  str       =			SpeechTaHifiGan(lowercase_       )

     snake_case_				:  Optional[int]       =			{
         '''unet''': unet,
         '''scheduler''': scheduler,
         '''vae''': vae,
         '''text_encoder''': text_encoder,
         '''tokenizer''': tokenizer,
         '''vocoder''': vocoder,
     }
     return components

 def 		_snake_case      (  self	: Union[str, Any]						,  lowercase_	: Optional[int]						,  lowercase_	: Dict=0       ):
     if str(lowercase_       ).startswith('''mps'''       ):
         snake_case_				:  str       =			torch.manual_seed(lowercase_       )
     else:
         snake_case_				:  Union[str, Any]       =			torch.Generator(device=lowercase_       ).manual_seed(lowercase_       )
     snake_case_				:  Dict       =			{
         '''prompt''': '''A hammer hitting a wooden surface''',
         '''generator''': generator,
         '''num_inference_steps''': 2,
         '''guidance_scale''': 6.0,
     }
     return inputs

 def 		_snake_case      (  self	: int       ):
     snake_case_				:  str       =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator

     snake_case_				:  int       =			self.get_dummy_components()
     snake_case_				:  str       =			AudioLDMPipeline(**lowercase_       )
     snake_case_				:  List[str]       =			audioldm_pipe.to(lowercase_       )
     audioldm_pipe.set_progress_bar_config(disable=lowercase_       )

     snake_case_				:  Optional[Any]       =			self.get_dummy_inputs(lowercase_       )
     snake_case_				:  Tuple       =			audioldm_pipe(**lowercase_       )
     snake_case_				:  Optional[int]       =			output.audios[0]

     assert audio.ndim == 1
     assert len(lowercase_       ) == 256

     snake_case_				:  Any       =			audio[:10]
     snake_case_				:  Optional[Any]       =			np.array(
         [-0.00_50, 0.00_50, -0.00_60, 0.00_33, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_33]       )

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

 def 		_snake_case      (  self	: Tuple       ):
     snake_case_				:  int       =			self.get_dummy_components()
     snake_case_				:  List[str]       =			AudioLDMPipeline(**lowercase_       )
     snake_case_				:  Any       =			audioldm_pipe.to(lowercase_       )
     snake_case_				:  Optional[int]       =			audioldm_pipe.to(lowercase_       )
     audioldm_pipe.set_progress_bar_config(disable=lowercase_       )

     snake_case_				:  Tuple       =			self.get_dummy_inputs(lowercase_       )
     snake_case_				:  str       =			3 * [inputs['''prompt''']]

     # forward
     snake_case_				:  List[Any]       =			audioldm_pipe(**lowercase_       )
     snake_case_				:  Tuple       =			output.audios[0]

     snake_case_				:  Any       =			self.get_dummy_inputs(lowercase_       )
     snake_case_				:  Any       =			3 * [inputs.pop('''prompt'''       )]

     snake_case_				:  List[str]       =			audioldm_pipe.tokenizer(
         lowercase_						,  padding='''max_length'''						,  max_length=audioldm_pipe.tokenizer.model_max_length						,  truncation=lowercase_						,  return_tensors='''pt'''						,  )
     snake_case_				:  List[str]       =			text_inputs['''input_ids'''].to(lowercase_       )

     snake_case_				:  Tuple       =			audioldm_pipe.text_encoder(
         lowercase_						,  )
     snake_case_				:  int       =			prompt_embeds.text_embeds
     # additional L_2 normalization over each hidden-state
     snake_case_				:  int       =			F.normalize(lowercase_						,  dim=-1       )

     snake_case_				:  List[Any]       =			prompt_embeds

     # forward
     snake_case_				:  str       =			audioldm_pipe(**lowercase_       )
     snake_case_				:  Union[str, Any]       =			output.audios[0]

     assert np.abs(audio_a - audio_a       ).max() < 1E-2

 def 		_snake_case      (  self	: Union[str, Any]       ):
     snake_case_				:  Tuple       =			self.get_dummy_components()
     snake_case_				:  str       =			AudioLDMPipeline(**lowercase_       )
     snake_case_				:  str       =			audioldm_pipe.to(lowercase_       )
     snake_case_				:  Tuple       =			audioldm_pipe.to(lowercase_       )
     audioldm_pipe.set_progress_bar_config(disable=lowercase_       )

     snake_case_				:  int       =			self.get_dummy_inputs(lowercase_       )
     snake_case_				:  Optional[Any]       =			3 * ['''this is a negative prompt''']
     snake_case_				:  int       =			negative_prompt
     snake_case_				:  int       =			3 * [inputs['''prompt''']]

     # forward
     snake_case_				:  List[str]       =			audioldm_pipe(**lowercase_       )
     snake_case_				:  str       =			output.audios[0]

     snake_case_				:  str       =			self.get_dummy_inputs(lowercase_       )
     snake_case_				:  List[str]       =			3 * [inputs.pop('''prompt'''       )]

     snake_case_				:  str       =			[]
     for p in [prompt, negative_prompt]:
         snake_case_				:  str       =			audioldm_pipe.tokenizer(
             lowercase_						,  padding='''max_length'''						,  max_length=audioldm_pipe.tokenizer.model_max_length						,  truncation=lowercase_						,  return_tensors='''pt'''						,  )
         snake_case_				:  int       =			text_inputs['''input_ids'''].to(lowercase_       )

         snake_case_				:  str       =			audioldm_pipe.text_encoder(
             lowercase_						,  )
         snake_case_				:  Dict       =			text_embeds.text_embeds
         # additional L_2 normalization over each hidden-state
         snake_case_				:  str       =			F.normalize(lowercase_						,  dim=-1       )

         embeds.append(lowercase_       )

     snake_case_,			snake_case_				:  Any       =			embeds

     # forward
     snake_case_				:  Optional[Any]       =			audioldm_pipe(**lowercase_       )
     snake_case_				:  str       =			output.audios[0]

     assert np.abs(audio_a - audio_a       ).max() < 1E-2

 def 		_snake_case      (  self	: Dict       ):
     snake_case_				:  int       =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
     snake_case_				:  List[Any]       =			self.get_dummy_components()
     snake_case_				:  Optional[int]       =			PNDMScheduler(skip_prk_steps=lowercase_       )
     snake_case_				:  List[Any]       =			AudioLDMPipeline(**lowercase_       )
     snake_case_				:  Dict       =			audioldm_pipe.to(lowercase_       )
     audioldm_pipe.set_progress_bar_config(disable=lowercase_       )

     snake_case_				:  List[str]       =			self.get_dummy_inputs(lowercase_       )
     snake_case_				:  Any       =			'''egg cracking'''
     snake_case_				:  int       =			audioldm_pipe(**lowercase_						,  negative_prompt=lowercase_       )
     snake_case_				:  Dict       =			output.audios[0]

     assert audio.ndim == 1
     assert len(lowercase_       ) == 256

     snake_case_				:  Any       =			audio[:10]
     snake_case_				:  Dict       =			np.array(
         [-0.00_51, 0.00_50, -0.00_60, 0.00_34, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_32]       )

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

 def 		_snake_case      (  self	: Optional[Any]       ):
     snake_case_				:  Optional[Any]       =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
     snake_case_				:  Optional[Any]       =			self.get_dummy_components()
     snake_case_				:  Tuple       =			PNDMScheduler(skip_prk_steps=lowercase_       )
     snake_case_				:  str       =			AudioLDMPipeline(**lowercase_       )
     snake_case_				:  Dict       =			audioldm_pipe.to(lowercase_       )
     audioldm_pipe.set_progress_bar_config(disable=lowercase_       )

     snake_case_				:  str       =			'''A hammer hitting a wooden surface'''

     # test num_waveforms_per_prompt=1 (default)
     snake_case_				:  List[str]       =			audioldm_pipe(lowercase_						,  num_inference_steps=2       ).audios

     assert audios.shape == (1, 256)

     # test num_waveforms_per_prompt=1 (default) for batch of prompts
     snake_case_				:  Any       =			2
     snake_case_				:  Optional[int]       =			audioldm_pipe([prompt] * batch_size						,  num_inference_steps=2       ).audios

     assert audios.shape == (batch_size, 256)

     # test num_waveforms_per_prompt for single prompt
     snake_case_				:  Optional[int]       =			2
     snake_case_				:  Tuple       =			audioldm_pipe(lowercase_						,  num_inference_steps=2						,  num_waveforms_per_prompt=lowercase_       ).audios

     assert audios.shape == (num_waveforms_per_prompt, 256)

     # test num_waveforms_per_prompt for batch of prompts
     snake_case_				:  List[str]       =			2
     snake_case_				:  Tuple       =			audioldm_pipe(
         [prompt] * batch_size						,  num_inference_steps=2						,  num_waveforms_per_prompt=lowercase_       ).audios

     assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)

 def 		_snake_case      (  self	: List[str]       ):
     snake_case_				:  str       =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
     snake_case_				:  List[str]       =			self.get_dummy_components()
     snake_case_				:  Optional[int]       =			AudioLDMPipeline(**lowercase_       )
     snake_case_				:  List[Any]       =			audioldm_pipe.to(lowercase_       )
     audioldm_pipe.set_progress_bar_config(disable=lowercase_       )
     snake_case_				:  str       =			audioldm_pipe.vocoder.config.sampling_rate

     snake_case_				:  str       =			self.get_dummy_inputs(lowercase_       )
     snake_case_				:  Dict       =			audioldm_pipe(audio_length_in_s=0.0_16						,  **lowercase_       )
     snake_case_				:  List[str]       =			output.audios[0]

     assert audio.ndim == 1
     assert len(lowercase_       ) / vocoder_sampling_rate == 0.0_16

     snake_case_				:  List[str]       =			audioldm_pipe(audio_length_in_s=0.0_32						,  **lowercase_       )
     snake_case_				:  Tuple       =			output.audios[0]

     assert audio.ndim == 1
     assert len(lowercase_       ) / vocoder_sampling_rate == 0.0_32

 def 		_snake_case      (  self	: int       ):
     snake_case_				:  List[Any]       =			self.get_dummy_components()
     snake_case_				:  Dict       =			AudioLDMPipeline(**lowercase_       )
     snake_case_				:  Union[str, Any]       =			audioldm_pipe.to(lowercase_       )
     audioldm_pipe.set_progress_bar_config(disable=lowercase_       )

     snake_case_				:  Optional[int]       =			['''hey''']

     snake_case_				:  List[Any]       =			audioldm_pipe(lowercase_						,  num_inference_steps=1       )
     snake_case_				:  List[str]       =			output.audios.shape
     assert audio_shape == (1, 256)

     snake_case_				:  str       =			audioldm_pipe.vocoder.config
     config.model_in_dim *= 2
     snake_case_				:  Dict       =			SpeechTaHifiGan(lowercase_       ).to(lowercase_       )
     snake_case_				:  Any       =			audioldm_pipe(lowercase_						,  num_inference_steps=1       )
     snake_case_				:  List[str]       =			output.audios.shape
     # waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
     assert audio_shape == (1, 256)

 def 		_snake_case      (  self	: Any       ):
     self._test_attention_slicing_forward_pass(test_mean_pixel_difference=lowercase_       )

 def 		_snake_case      (  self	: int       ):
     self._test_inference_batch_single_identical(test_mean_pixel_difference=lowercase_       )






 @unittest.skipIf(
     torch_device != '''cuda''' or not is_xformers_available()						,  reason='''XFormers attention is only available with CUDA and `xformers` installed'''						,  )
 def 		_snake_case      (  self	: Optional[Any]       ):
     self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowercase_       )





@slow
class  _UpperCAmelCase (    unittest.TestCase):

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

 def 		_snake_case      (  self	: Optional[Any]						,  lowercase_	: Dict						,  lowercase_	: Union[str, Any]="cpu"						,  lowercase_	: Optional[int]=torch.floataa						,  lowercase_	: Optional[int]=0       ):
     snake_case_				:  Union[str, Any]       =			torch.Generator(device=lowercase_       ).manual_seed(lowercase_       )
     snake_case_				:  str       =			np.random.RandomState(lowercase_       ).standard_normal((1, 8, 128, 16)       )
     snake_case_				:  Optional[Any]       =			torch.from_numpy(lowercase_       ).to(device=lowercase_						,  dtype=lowercase_       )
     snake_case_				:  Any       =			{
         '''prompt''': '''A hammer hitting a wooden surface''',
         '''latents''': latents,
         '''generator''': generator,
         '''num_inference_steps''': 3,
         '''guidance_scale''': 2.5,
     }
     return inputs

 def 		_snake_case      (  self	: int       ):
     snake_case_				:  List[Any]       =			AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''       )
     snake_case_				:  Tuple       =			audioldm_pipe.to(lowercase_       )
     audioldm_pipe.set_progress_bar_config(disable=lowercase_       )

     snake_case_				:  Tuple       =			self.get_inputs(lowercase_       )
     snake_case_				:  str       =			25
     snake_case_				:  Optional[int]       =			audioldm_pipe(**lowercase_       ).audios[0]

     assert audio.ndim == 1
     assert len(lowercase_       ) == 81920

     snake_case_				:  Tuple       =			audio[77230:77240]
     snake_case_				:  str       =			np.array(
         [-0.48_84, -0.46_07, 0.00_23, 0.50_07, 0.58_96, 0.51_51, 0.38_13, -0.02_08, -0.36_87, -0.43_15]       )
     snake_case_				:  Tuple       =			np.abs(expected_slice - audio_slice       ).max()
     assert max_diff < 1E-2






 def 		_snake_case      (  self	: Union[str, Any]       ):
     snake_case_				:  str       =			AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''       )
     snake_case_				:  Union[str, Any]       =			LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config       )
     snake_case_				:  int       =			audioldm_pipe.to(lowercase_       )
     audioldm_pipe.set_progress_bar_config(disable=lowercase_       )

     snake_case_				:  Any       =			self.get_inputs(lowercase_       )
     snake_case_				:  List[Any]       =			audioldm_pipe(**lowercase_       ).audios[0]

     assert audio.ndim == 1
     assert len(lowercase_       ) == 81920

     snake_case_				:  Optional[int]       =			audio[27780:27790]
     snake_case_				:  Union[str, Any]       =			np.array([-0.21_31, -0.08_73, -0.01_24, -0.01_89, 0.05_69, 0.13_73, 0.18_83, 0.28_86, 0.32_97, 0.22_12]       )
     snake_case_				:  Any       =			np.abs(expected_slice - audio_slice       ).max()
     assert max_diff < 3E-2










"""simple docstring"""

from argparse import ArgumentParser

from datasets.commands.convert import ConvertCommand
from datasets.commands.dummy_data import DummyDataCommand
from datasets.commands.env import EnvironmentCommand
from datasets.commands.run_beam import RunBeamCommand
from datasets.commands.test import TestCommand
from datasets.utils.logging import set_verbosity_info




def __lowercase    (							_a					):
    return {key.lstrip('''-'''					): value for key, value in zip(unknown_args[::2]	,    unknown_args[1::2]					)}




def __lowercase    (							):
    snake_case_				:  List[str]       =			ArgumentParser(
        '''HuggingFace Datasets CLI tool'''	,    usage='''datasets-cli <command> [<args>]'''	,    allow_abbrev=_a					)
    snake_case_				:  List[Any]       =			parser.add_subparsers(help='''datasets-cli command helpers'''					)
    set_verbosity_info()

    # Register commands
    ConvertCommand.register_subcommand(_a					)
    EnvironmentCommand.register_subcommand(_a					)
    TestCommand.register_subcommand(_a					)
    RunBeamCommand.register_subcommand(_a					)
    DummyDataCommand.register_subcommand(_a					)

    # Parse args
    snake_case_,			snake_case_				:  Optional[Any]       =			parser.parse_known_args()
    if not hasattr(_a	,    '''func'''					):
        parser.print_help()
        exit(1					)
    snake_case_				:  Optional[int]       =			parse_unknown_args(_a					)

    # Run
    snake_case_				:  Optional[int]       =			args.func(_a	,    **_a					)
    service.run()


if __name__ == "__main__":
     main()





