






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   __lowerCAmelCase			(					_UpperCamelCase		,unittest.TestCase  ):
   _UpperCamelCase :						List[str]     		=				AudioLDMPipeline
   _UpperCamelCase :						List[str]     		=				TEXT_TO_AUDIO_PARAMS
   _UpperCamelCase :						Any     		=				TEXT_TO_AUDIO_BATCH_PARAMS
   _UpperCamelCase :						Optional[Any]     		=				frozenset(
       [
           """num_inference_steps""",
           """num_waveforms_per_prompt""",
           """generator""",
           """latents""",
           """output_type""",
           """return_dict""",
           """callback""",
           """callback_steps""",
       ]  )





   def 	_snake_case ( self      )				->  Any:
      """simple docstring"""





      torch.manual_seed(0      )
      a__	:						int      		=  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=snake_case					,			)
      a__	:						List[Any]      		=  DDIMScheduler(
          beta_start=0.00_085					,			beta_end=0.012					,			beta_schedule="scaled_linear"					,			clip_sample=snake_case					,			set_alpha_to_one=snake_case					,			)
      torch.manual_seed(0      )
      a__	:						Union[str, Any]      		=  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      )
      a__	:						Union[str, 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=1_000					,			projection_dim=32					,			)
      a__	:						int      		=  ClapTextModelWithProjection(snake_case      )
      a__	:						Optional[Any]      		=  RobertaTokenizer.from_pretrained("hf-internal-testing/tiny-random-roberta"					,			model_max_length=77      )

      a__	:						List[Any]      		=  SpeechTaHifiGanConfig(
          model_in_dim=8					,			sampling_rate=16_000					,			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=snake_case					,			)

      a__	:						str      		=  SpeechTaHifiGan(snake_case      )

      a__	:						Dict      		=  {
          "unet": unet,
          "scheduler": scheduler,
          "vae": vae,
          "text_encoder": text_encoder,
          "tokenizer": tokenizer,
          "vocoder": vocoder,
      }
      return components





   def 	_snake_case ( self					,			snake_case					,			snake_case=0      )				->  List[str]:
      """simple docstring"""





      if str(snake_case      ).startswith("mps"      ):
         a__	:						int      		=  torch.manual_seed(snake_case      )
      else:
         a__	:						List[str]      		=  torch.Generator(device=snake_case      ).manual_seed(snake_case      )
      a__	:						int      		=  {
          "prompt": "A hammer hitting a wooden surface",
          "generator": generator,
          "num_inference_steps": 2,
          "guidance_scale": 6.0,
      }
      return inputs





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						Any      		=  "cpu"  # ensure determinism for the device-dependent torch.Generator

      a__	:						Union[str, Any]      		=  self.get_dummy_components()
      a__	:						Tuple      		=  AudioLDMPipeline(**snake_case      )
      a__	:						Optional[int]      		=  audioldm_pipe.to(snake_case      )
      audioldm_pipe.set_progress_bar_config(disable=snake_case      )

      a__	:						List[Any]      		=  self.get_dummy_inputs(snake_case      )
      a__	:						List[Any]      		=  audioldm_pipe(**snake_case      )
      a__	:						Optional[Any]      		=  output.audios[0]

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

      a__	:						List[Any]      		=  audio[:10]
      a__	:						int      		=  np.array(
          [-0.0_050, 0.0_050, -0.0_060, 0.0_033, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_033]      )

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





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      a__	:						Union[str, Any]      		=  self.get_dummy_components()
      a__	:						List[Any]      		=  AudioLDMPipeline(**snake_case      )
      a__	:						Optional[Any]      		=  audioldm_pipe.to(snake_case      )
      a__	:						Union[str, Any]      		=  audioldm_pipe.to(snake_case      )
      audioldm_pipe.set_progress_bar_config(disable=snake_case      )

      a__	:						Tuple      		=  self.get_dummy_inputs(snake_case      )
      a__	:						List[str]      		=  3 * [inputs["prompt"]]

      # forward
      a__	:						str      		=  audioldm_pipe(**snake_case      )
      a__	:						Any      		=  output.audios[0]

      a__	:						List[str]      		=  self.get_dummy_inputs(snake_case      )
      a__	:						List[Any]      		=  3 * [inputs.pop("prompt"      )]

      a__	:						str      		=  audioldm_pipe.tokenizer(
          snake_case					,			padding="max_length"					,			max_length=audioldm_pipe.tokenizer.model_max_length					,			truncation=snake_case					,			return_tensors="pt"					,			)
      a__	:						int      		=  text_inputs["input_ids"].to(snake_case      )

      a__	:						str      		=  audioldm_pipe.text_encoder(
          snake_case					,			)
      a__	:						Optional[int]      		=  prompt_embeds.text_embeds
      # additional L_2 normalization over each hidden-state
      a__	:						int      		=  F.normalize(snake_case					,			dim=-1      )

      a__	:						List[str]      		=  prompt_embeds

      # forward
      a__	:						str      		=  audioldm_pipe(**snake_case      )
      a__	:						Any      		=  output.audios[0]

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





   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      a__	:						Tuple      		=  self.get_dummy_components()
      a__	:						int      		=  AudioLDMPipeline(**snake_case      )
      a__	:						Optional[int]      		=  audioldm_pipe.to(snake_case      )
      a__	:						Union[str, Any]      		=  audioldm_pipe.to(snake_case      )
      audioldm_pipe.set_progress_bar_config(disable=snake_case      )

      a__	:						Optional[Any]      		=  self.get_dummy_inputs(snake_case      )
      a__	:						Dict      		=  3 * ["this is a negative prompt"]
      a__	:						Union[str, Any]      		=  negative_prompt
      a__	:						Any      		=  3 * [inputs["prompt"]]

      # forward
      a__	:						Tuple      		=  audioldm_pipe(**snake_case      )
      a__	:						str      		=  output.audios[0]

      a__	:						Dict      		=  self.get_dummy_inputs(snake_case      )
      a__	:						List[str]      		=  3 * [inputs.pop("prompt"      )]

      a__	:						List[Any]      		=  []
      for p in [prompt, negative_prompt]:
         a__	:						Optional[Any]      		=  audioldm_pipe.tokenizer(
             snake_case					,			padding="max_length"					,			max_length=audioldm_pipe.tokenizer.model_max_length					,			truncation=snake_case					,			return_tensors="pt"					,			)
         a__	:						Any      		=  text_inputs["input_ids"].to(snake_case      )

         a__	:						Any      		=  audioldm_pipe.text_encoder(
             snake_case					,			)
         a__	:						str      		=  text_embeds.text_embeds
         # additional L_2 normalization over each hidden-state
         a__	:						Tuple      		=  F.normalize(snake_case					,			dim=-1      )

         embeds.append(snake_case      )

      a__	:						Any      		=  embeds

      # forward
      a__	:						Dict      		=  audioldm_pipe(**snake_case      )
      a__	:						Optional[int]      		=  output.audios[0]

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





   def 	_snake_case ( self      )				->  Any:
      """simple docstring"""





      a__	:						Optional[int]      		=  "cpu"  # ensure determinism for the device-dependent torch.Generator
      a__	:						Union[str, Any]      		=  self.get_dummy_components()
      a__	:						int      		=  PNDMScheduler(skip_prk_steps=snake_case      )
      a__	:						str      		=  AudioLDMPipeline(**snake_case      )
      a__	:						Union[str, Any]      		=  audioldm_pipe.to(snake_case      )
      audioldm_pipe.set_progress_bar_config(disable=snake_case      )

      a__	:						Any      		=  self.get_dummy_inputs(snake_case      )
      a__	:						Optional[Any]      		=  "egg cracking"
      a__	:						List[Any]      		=  audioldm_pipe(**snake_case					,			negative_prompt=snake_case      )
      a__	:						Tuple      		=  output.audios[0]

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

      a__	:						List[Any]      		=  audio[:10]
      a__	:						Dict      		=  np.array(
          [-0.0_051, 0.0_050, -0.0_060, 0.0_034, -0.0_026, 0.0_033, -0.0_027, 0.0_033, -0.0_028, 0.0_032]      )

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





   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





      a__	:						Union[str, Any]      		=  "cpu"  # ensure determinism for the device-dependent torch.Generator
      a__	:						Optional[Any]      		=  self.get_dummy_components()
      a__	:						List[str]      		=  PNDMScheduler(skip_prk_steps=snake_case      )
      a__	:						Dict      		=  AudioLDMPipeline(**snake_case      )
      a__	:						int      		=  audioldm_pipe.to(snake_case      )
      audioldm_pipe.set_progress_bar_config(disable=snake_case      )

      a__	:						Any      		=  "A hammer hitting a wooden surface"

      # test num_waveforms_per_prompt=1 (default)
      a__	:						Union[str, Any]      		=  audioldm_pipe(snake_case					,			num_inference_steps=2      ).audios

      assert audios.shape == (1, 256)

      # test num_waveforms_per_prompt=1 (default) for batch of prompts
      a__	:						Dict      		=  2
      a__	:						List[str]      		=  audioldm_pipe([prompt] * batch_size					,			num_inference_steps=2      ).audios

      assert audios.shape == (batch_size, 256)

      # test num_waveforms_per_prompt for single prompt
      a__	:						Dict      		=  2
      a__	:						Union[str, Any]      		=  audioldm_pipe(snake_case					,			num_inference_steps=2					,			num_waveforms_per_prompt=snake_case      ).audios

      assert audios.shape == (num_waveforms_per_prompt, 256)

      # test num_waveforms_per_prompt for batch of prompts
      a__	:						Dict      		=  2
      a__	:						Optional[int]      		=  audioldm_pipe(
          [prompt] * batch_size					,			num_inference_steps=2					,			num_waveforms_per_prompt=snake_case      ).audios

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





   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      a__	:						List[Any]      		=  "cpu"  # ensure determinism for the device-dependent torch.Generator
      a__	:						int      		=  self.get_dummy_components()
      a__	:						List[str]      		=  AudioLDMPipeline(**snake_case      )
      a__	:						Union[str, Any]      		=  audioldm_pipe.to(snake_case      )
      audioldm_pipe.set_progress_bar_config(disable=snake_case      )
      a__	:						Dict      		=  audioldm_pipe.vocoder.config.sampling_rate

      a__	:						Union[str, Any]      		=  self.get_dummy_inputs(snake_case      )
      a__	:						Tuple      		=  audioldm_pipe(audio_length_in_s=0.016					,			**snake_case      )
      a__	:						Dict      		=  output.audios[0]

      assert audio.ndim == 1
      assert len(snake_case      ) / vocoder_sampling_rate == 0.016

      a__	:						Optional[int]      		=  audioldm_pipe(audio_length_in_s=0.032					,			**snake_case      )
      a__	:						str      		=  output.audios[0]

      assert audio.ndim == 1
      assert len(snake_case      ) / vocoder_sampling_rate == 0.032





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      a__	:						Dict      		=  self.get_dummy_components()
      a__	:						Any      		=  AudioLDMPipeline(**snake_case      )
      a__	:						str      		=  audioldm_pipe.to(snake_case      )
      audioldm_pipe.set_progress_bar_config(disable=snake_case      )

      a__	:						Any      		=  ["hey"]

      a__	:						int      		=  audioldm_pipe(snake_case					,			num_inference_steps=1      )
      a__	:						List[str]      		=  output.audios.shape
      assert audio_shape == (1, 256)

      a__	:						Dict      		=  audioldm_pipe.vocoder.config
      config.model_in_dim *= 2
      a__	:						List[str]      		=  SpeechTaHifiGan(snake_case      ).to(snake_case      )
      a__	:						Tuple      		=  audioldm_pipe(snake_case					,			num_inference_steps=1      )
      a__	:						Dict      		=  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      )				->  Optional[int]:
      """simple docstring"""





      self._test_attention_slicing_forward_pass(test_mean_pixel_difference=snake_case      )





   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      self._test_inference_batch_single_identical(test_mean_pixel_difference=snake_case      )





   @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[int]:
      """simple docstring"""





      self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=snake_case      )







@slow
class   __lowerCAmelCase			(					unittest.TestCase  ):





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





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





   def 	_snake_case ( self					,			snake_case					,			snake_case="cpu"					,			snake_case=torch.floataa					,			snake_case=0      )				->  Union[str, Any]:
      """simple docstring"""





      a__	:						Tuple      		=  torch.Generator(device=snake_case      ).manual_seed(snake_case      )
      a__	:						Union[str, Any]      		=  np.random.RandomState(snake_case      ).standard_normal((1, 8, 128, 16)      )
      a__	:						Tuple      		=  torch.from_numpy(snake_case      ).to(device=snake_case					,			dtype=snake_case      )
      a__	:						Dict      		=  {
          "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:
      """simple docstring"""





      a__	:						Union[str, Any]      		=  AudioLDMPipeline.from_pretrained("cvssp/audioldm"      )
      a__	:						Any      		=  audioldm_pipe.to(snake_case      )
      audioldm_pipe.set_progress_bar_config(disable=snake_case      )

      a__	:						Any      		=  self.get_inputs(snake_case      )
      a__	:						str      		=  25
      a__	:						Tuple      		=  audioldm_pipe(**snake_case      ).audios[0]

      assert audio.ndim == 1
      assert len(snake_case      ) == 81_920

      a__	:						Any      		=  audio[77_230:77_240]
      a__	:						Union[str, Any]      		=  np.array(
          [-0.4_884, -0.4_607, 0.0_023, 0.5_007, 0.5_896, 0.5_151, 0.3_813, -0.0_208, -0.3_687, -0.4_315]      )
      a__	:						Optional[Any]      		=  np.abs(expected_slice - audio_slice      ).max()
      assert max_diff < 1E-2





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Optional[Any]      		=  AudioLDMPipeline.from_pretrained("cvssp/audioldm"      )
      a__	:						Union[str, Any]      		=  LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config      )
      a__	:						Dict      		=  audioldm_pipe.to(snake_case      )
      audioldm_pipe.set_progress_bar_config(disable=snake_case      )

      a__	:						str      		=  self.get_inputs(snake_case      )
      a__	:						str      		=  audioldm_pipe(**snake_case      ).audios[0]

      assert audio.ndim == 1
      assert len(snake_case      ) == 81_920

      a__	:						Optional[Any]      		=  audio[27_780:27_790]
      a__	:						List[str]      		=  np.array([-0.2_131, -0.0_873, -0.0_124, -0.0_189, 0.0_569, 0.1_373, 0.1_883, 0.2_886, 0.3_297, 0.2_212]      )
      a__	:						Any      		=  np.abs(expected_slice - audio_slice      ).max()
      assert max_diff < 3E-2








import tempfile
import unittest

from make_student import create_student_by_copying_alternating_layers

from transformers import AutoConfig
from transformers.file_utils import cached_property
from transformers.testing_utils import require_torch


SCREAMING_SNAKE_CASE__       :			List[Any]			      =						"""sshleifer/bart-tiny-random"""
SCREAMING_SNAKE_CASE__       :			Tuple			      =						"""patrickvonplaten/t5-tiny-random"""







@require_torch
class   __lowerCAmelCase			(					unittest.TestCase  ):





   @cached_property
   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





      return AutoConfig.from_pretrained(snake_case      )





   def 	_snake_case ( self      )				->  Any:
      """simple docstring"""





      a__		,     *a__	:						int      		=  create_student_by_copying_alternating_layers(snake_case					,			tempfile.mkdtemp()					,			e=1					,			d=1      )
      self.assertEqual(student.config.num_hidden_layers					,			1      )





   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      a__		,     *a__	:						int      		=  create_student_by_copying_alternating_layers(snake_case					,			tempfile.mkdtemp()					,			e=1					,			d=snake_case      )





   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





      a__		,     *a__	:						Any      		=  create_student_by_copying_alternating_layers(snake_case					,			tempfile.mkdtemp()					,			e=1					,			d=snake_case      )
      self.assertEqual(student.config.encoder_layers					,			1      )
      self.assertEqual(student.config.decoder_layers					,			self.teacher_config.encoder_layers      )





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__		,     *a__	:						List[str]      		=  create_student_by_copying_alternating_layers(snake_case					,			tempfile.mkdtemp()					,			e=1					,			d=1      )
      self.assertEqual(student.config.encoder_layers					,			1      )
      self.assertEqual(student.config.decoder_layers					,			1      )





   def 	_snake_case ( self      )				->  int:
      """simple docstring"""





      with self.assertRaises(snake_case      ):
         create_student_by_copying_alternating_layers(snake_case					,			tempfile.mkdtemp()					,			e=snake_case					,			d=snake_case      )

