'''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 lowerCamelCase							(	_A     ,		unittest.TestCase					):




  '''simple docstring'''



  lowerCAmelCase__										=	AudioLDMPipeline
  lowerCAmelCase__										=	TEXT_TO_AUDIO_PARAMS
  lowerCAmelCase__										=	TEXT_TO_AUDIO_BATCH_PARAMS
  lowerCAmelCase__										=	frozenset(
      [
          '''num_inference_steps''',
          '''num_waveforms_per_prompt''',
          '''generator''',
          '''latents''',
          '''output_type''',
          '''return_dict''',
          '''callback''',
          '''callback_steps''',
      ]					)






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Optional[Any]:
        torch.manual_seed(0  )
        UpperCAmelCase_            =  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=__lowerCamelCase  ,							)
        UpperCAmelCase_            =  DDIMScheduler(
            beta_start=0.0_0085  ,							beta_end=0.012  ,							beta_schedule='''scaled_linear'''  ,							clip_sample=__lowerCamelCase  ,							set_alpha_to_one=__lowerCamelCase  ,							)
        torch.manual_seed(0  )
        UpperCAmelCase_            =  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  )
        UpperCAmelCase_            =  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  ,							)
        UpperCAmelCase_            =  ClapTextModelWithProjection(__lowerCamelCase  )
        UpperCAmelCase_            =  RobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-roberta'''  ,							model_max_length=77  )

        UpperCAmelCase_            =  SpeechTaHifiGanConfig(
            model_in_dim=8  ,							sampling_rate=1_6000  ,							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=__lowerCamelCase  ,							)

        UpperCAmelCase_            =  SpeechTaHifiGan(__lowerCamelCase  )

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






  def  lowerCAmelCase__ (      self			:       Optional[Any]  ,							UpperCAmelCase__			:       Optional[Any]  ,							UpperCAmelCase__			:       Optional[Any]=0  )				->int:
        if str(__lowerCamelCase  ).startswith('''mps'''  ):
              UpperCAmelCase_            =  torch.manual_seed(__lowerCamelCase  )
        else:
              UpperCAmelCase_            =  torch.Generator(device=__lowerCamelCase  ).manual_seed(__lowerCamelCase  )
        UpperCAmelCase_            =  {
            "prompt": "A hammer hitting a wooden surface",
            "generator": generator,
            "num_inference_steps": 2,
            "guidance_scale": 6.0,
        }
        return inputs






  def  lowerCAmelCase__ (      self			:       Dict  )				->Union[str, Any]:
        UpperCAmelCase_            =  "cpu"  # ensure determinism for the device-dependent torch.Generator

        UpperCAmelCase_            =  self.get_dummy_components()
        UpperCAmelCase_            =  AudioLDMPipeline(**__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

        UpperCAmelCase_            =  self.get_dummy_inputs(__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe(**__lowerCamelCase  )
        UpperCAmelCase_            =  output.audios[0]

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

        UpperCAmelCase_            =  audio[:10]
        UpperCAmelCase_            =  np.array(
            [-0.0050, 0.0050, -0.0060, 0.0033, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0033]  )

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






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->List[str]:
        UpperCAmelCase_            =  self.get_dummy_components()
        UpperCAmelCase_            =  AudioLDMPipeline(**__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

        UpperCAmelCase_            =  self.get_dummy_inputs(__lowerCamelCase  )
        UpperCAmelCase_            =  3 * [inputs["prompt"]]

        # forward
        UpperCAmelCase_            =  audioldm_pipe(**__lowerCamelCase  )
        UpperCAmelCase_            =  output.audios[0]

        UpperCAmelCase_            =  self.get_dummy_inputs(__lowerCamelCase  )
        UpperCAmelCase_            =  3 * [inputs.pop('''prompt'''  )]

        UpperCAmelCase_            =  audioldm_pipe.tokenizer(
            __lowerCamelCase  ,							padding='''max_length'''  ,							max_length=audioldm_pipe.tokenizer.model_max_length  ,							truncation=__lowerCamelCase  ,							return_tensors='''pt'''  ,							)
        UpperCAmelCase_            =  text_inputs["input_ids"].to(__lowerCamelCase  )

        UpperCAmelCase_            =  audioldm_pipe.text_encoder(
            __lowerCamelCase  ,							)
        UpperCAmelCase_            =  prompt_embeds.text_embeds
        # additional L_2 normalization over each hidden-state
        UpperCAmelCase_            =  F.normalize(__lowerCamelCase  ,							dim=-1  )

        UpperCAmelCase_            =  prompt_embeds

        # forward
        UpperCAmelCase_            =  audioldm_pipe(**__lowerCamelCase  )
        UpperCAmelCase_            =  output.audios[0]

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






  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->List[str]:
        UpperCAmelCase_            =  self.get_dummy_components()
        UpperCAmelCase_            =  AudioLDMPipeline(**__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

        UpperCAmelCase_            =  self.get_dummy_inputs(__lowerCamelCase  )
        UpperCAmelCase_            =  3 * ["this is a negative prompt"]
        UpperCAmelCase_            =  negative_prompt
        UpperCAmelCase_            =  3 * [inputs["prompt"]]

        # forward
        UpperCAmelCase_            =  audioldm_pipe(**__lowerCamelCase  )
        UpperCAmelCase_            =  output.audios[0]

        UpperCAmelCase_            =  self.get_dummy_inputs(__lowerCamelCase  )
        UpperCAmelCase_            =  3 * [inputs.pop('''prompt'''  )]

        UpperCAmelCase_            =  []
        for p in [prompt, negative_prompt]:
              UpperCAmelCase_            =  audioldm_pipe.tokenizer(
                  __lowerCamelCase  ,							padding='''max_length'''  ,							max_length=audioldm_pipe.tokenizer.model_max_length  ,							truncation=__lowerCamelCase  ,							return_tensors='''pt'''  ,							)
              UpperCAmelCase_            =  text_inputs["input_ids"].to(__lowerCamelCase  )

              UpperCAmelCase_            =  audioldm_pipe.text_encoder(
                  __lowerCamelCase  ,							)
              UpperCAmelCase_            =  text_embeds.text_embeds
              # additional L_2 normalization over each hidden-state
              UpperCAmelCase_            =  F.normalize(__lowerCamelCase  ,							dim=-1  )

              embeds.append(__lowerCamelCase  )

        UpperCAmelCase_            =  embeds

        # forward
        UpperCAmelCase_            =  audioldm_pipe(**__lowerCamelCase  )
        UpperCAmelCase_            =  output.audios[0]

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






  def  lowerCAmelCase__ (      self			:       Optional[Any]  )				->str:
        UpperCAmelCase_            =  "cpu"  # ensure determinism for the device-dependent torch.Generator
        UpperCAmelCase_            =  self.get_dummy_components()
        UpperCAmelCase_            =  PNDMScheduler(skip_prk_steps=__lowerCamelCase  )
        UpperCAmelCase_            =  AudioLDMPipeline(**__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

        UpperCAmelCase_            =  self.get_dummy_inputs(__lowerCamelCase  )
        UpperCAmelCase_            =  "egg cracking"
        UpperCAmelCase_            =  audioldm_pipe(**__lowerCamelCase  ,							negative_prompt=__lowerCamelCase  )
        UpperCAmelCase_            =  output.audios[0]

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

        UpperCAmelCase_            =  audio[:10]
        UpperCAmelCase_            =  np.array(
            [-0.0051, 0.0050, -0.0060, 0.0034, -0.0026, 0.0033, -0.0027, 0.0033, -0.0028, 0.0032]  )

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






  def  lowerCAmelCase__ (      self			:       int  )				->Dict:
        UpperCAmelCase_            =  "cpu"  # ensure determinism for the device-dependent torch.Generator
        UpperCAmelCase_            =  self.get_dummy_components()
        UpperCAmelCase_            =  PNDMScheduler(skip_prk_steps=__lowerCamelCase  )
        UpperCAmelCase_            =  AudioLDMPipeline(**__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

        UpperCAmelCase_            =  "A hammer hitting a wooden surface"

        # test num_waveforms_per_prompt=1 (default)
        UpperCAmelCase_            =  audioldm_pipe(__lowerCamelCase  ,							num_inference_steps=2  ).audios

        assert audios.shape == (1, 256)

        # test num_waveforms_per_prompt=1 (default) for batch of prompts
        UpperCAmelCase_            =  2
        UpperCAmelCase_            =  audioldm_pipe([prompt] * batch_size  ,							num_inference_steps=2  ).audios

        assert audios.shape == (batch_size, 256)

        # test num_waveforms_per_prompt for single prompt
        UpperCAmelCase_            =  2
        UpperCAmelCase_            =  audioldm_pipe(__lowerCamelCase  ,							num_inference_steps=2  ,							num_waveforms_per_prompt=__lowerCamelCase  ).audios

        assert audios.shape == (num_waveforms_per_prompt, 256)

        # test num_waveforms_per_prompt for batch of prompts
        UpperCAmelCase_            =  2
        UpperCAmelCase_            =  audioldm_pipe(
            [prompt] * batch_size  ,							num_inference_steps=2  ,							num_waveforms_per_prompt=__lowerCamelCase  ).audios

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






  def  lowerCAmelCase__ (      self			:       str  )				->Optional[int]:
        UpperCAmelCase_            =  "cpu"  # ensure determinism for the device-dependent torch.Generator
        UpperCAmelCase_            =  self.get_dummy_components()
        UpperCAmelCase_            =  AudioLDMPipeline(**__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe.vocoder.config.sampling_rate

        UpperCAmelCase_            =  self.get_dummy_inputs(__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe(audio_length_in_s=0.016  ,							**__lowerCamelCase  )
        UpperCAmelCase_            =  output.audios[0]

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

        UpperCAmelCase_            =  audioldm_pipe(audio_length_in_s=0.032  ,							**__lowerCamelCase  )
        UpperCAmelCase_            =  output.audios[0]

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






  def  lowerCAmelCase__ (      self			:       Any  )				->Union[str, Any]:
        UpperCAmelCase_            =  self.get_dummy_components()
        UpperCAmelCase_            =  AudioLDMPipeline(**__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

        UpperCAmelCase_            =  ["hey"]

        UpperCAmelCase_            =  audioldm_pipe(__lowerCamelCase  ,							num_inference_steps=1  )
        UpperCAmelCase_            =  output.audios.shape
        assert audio_shape == (1, 256)

        UpperCAmelCase_            =  audioldm_pipe.vocoder.config
        config.model_in_dim *= 2
        UpperCAmelCase_            =  SpeechTaHifiGan(__lowerCamelCase  ).to(__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe(__lowerCamelCase  ,							num_inference_steps=1  )
        UpperCAmelCase_            =  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  lowerCAmelCase__ (      self			:       List[str]  )				->int:
        self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__lowerCamelCase  )






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Optional[Any]:
        self._test_inference_batch_single_identical(test_mean_pixel_difference=__lowerCamelCase  )






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







@slow
class lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



  def  lowerCAmelCase__ (      self			:       int  )				->int:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()






  def  lowerCAmelCase__ (      self			:       List[str]  ,							UpperCAmelCase__			:       str  ,							UpperCAmelCase__			:       str="cpu"  ,							UpperCAmelCase__			:       int=torch.floataa  ,							UpperCAmelCase__			:       Optional[Any]=0  )				->Any:
        UpperCAmelCase_            =  torch.Generator(device=__lowerCamelCase  ).manual_seed(__lowerCamelCase  )
        UpperCAmelCase_            =  np.random.RandomState(__lowerCamelCase  ).standard_normal((1, 8, 128, 16)  )
        UpperCAmelCase_            =  torch.from_numpy(__lowerCamelCase  ).to(device=__lowerCamelCase  ,							dtype=__lowerCamelCase  )
        UpperCAmelCase_            =  {
            "prompt": "A hammer hitting a wooden surface",
            "latents": latents,
            "generator": generator,
            "num_inference_steps": 3,
            "guidance_scale": 2.5,
        }
        return inputs






  def  lowerCAmelCase__ (      self			:       Optional[Any]  )				->List[str]:
        UpperCAmelCase_            =  AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

        UpperCAmelCase_            =  self.get_inputs(__lowerCamelCase  )
        UpperCAmelCase_            =  25
        UpperCAmelCase_            =  audioldm_pipe(**__lowerCamelCase  ).audios[0]

        assert audio.ndim == 1
        assert len(__lowerCamelCase  ) == 8_1920

        UpperCAmelCase_            =  audio[7_7230:7_7240]
        UpperCAmelCase_            =  np.array(
            [-0.4884, -0.4607, 0.0023, 0.5007, 0.5896, 0.5151, 0.3813, -0.0208, -0.3687, -0.4315]  )
        UpperCAmelCase_            =  np.abs(expected_slice - audio_slice  ).max()
        assert max_diff < 1e-2






  def  lowerCAmelCase__ (      self			:       str  )				->Tuple:
        UpperCAmelCase_            =  AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''  )
        UpperCAmelCase_            =  LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config  )
        UpperCAmelCase_            =  audioldm_pipe.to(__lowerCamelCase  )
        audioldm_pipe.set_progress_bar_config(disable=__lowerCamelCase  )

        UpperCAmelCase_            =  self.get_inputs(__lowerCamelCase  )
        UpperCAmelCase_            =  audioldm_pipe(**__lowerCamelCase  ).audios[0]

        assert audio.ndim == 1
        assert len(__lowerCamelCase  ) == 8_1920

        UpperCAmelCase_            =  audio[2_7780:2_7790]
        UpperCAmelCase_            =  np.array([-0.2131, -0.0873, -0.0124, -0.0189, 0.0569, 0.1373, 0.1883, 0.2886, 0.3297, 0.2212]  )
        UpperCAmelCase_            =  np.abs(expected_slice - audio_slice  ).max()
        assert max_diff < 3e-2




'''simple docstring'''
from heapq import heappop, heappush

import numpy as np



def  __lowerCamelCase       (      _UpperCamelCase  :  np.ndarray							,    _UpperCamelCase  :  tuple[int, int]							,    _UpperCamelCase  :  tuple[int, int]							,    _UpperCamelCase  :  bool							,    ):


      '''simple docstring'''


      UpperCAmelCase_      ,		UpperCAmelCase_            =  grid.shape
      UpperCAmelCase_            =  [-1, 1, 0, 0]
      UpperCAmelCase_            =  [0, 0, -1, 1]
      if allow_diagonal:
            dx += [-1, -1, 1, 1]
            dy += [-1, 1, -1, 1]

      UpperCAmelCase_      ,		UpperCAmelCase_            =  [(0, source)], set()
      UpperCAmelCase_            =  np.full((rows, cols)							,    np.inf					)
      UpperCAmelCase_            =  0
      UpperCAmelCase_            =  np.empty((rows, cols)							,    dtype=_UpperCamelCase					)
      UpperCAmelCase_            =  None

      while queue:
            ((UpperCAmelCase_)      ,		(UpperCAmelCase_))            =  heappop(_UpperCamelCase					)
            if (x, y) in visited:
                  continue
            visited.add((x, y)					)

            if (x, y) == destination:
                  UpperCAmelCase_            =  []
                  while (x, y) != source:
                        path.append((x, y)					)
                        UpperCAmelCase_      ,		UpperCAmelCase_            =  predecessors[x, y]
                  path.append(_UpperCamelCase					)  # add the source manually
                  path.reverse()
                  return matrix[destination], path

            for i in range(len(_UpperCamelCase					)					):
                  UpperCAmelCase_      ,		UpperCAmelCase_            =  x + dx[i], y + dy[i]
                  if 0 <= nx < rows and 0 <= ny < cols:
                        UpperCAmelCase_            =  grid[nx][ny]
                        if next_node == 1 and matrix[nx, ny] > dist + 1:
                              heappush(_UpperCamelCase							,    (dist + 1, (nx, ny))					)
                              UpperCAmelCase_            =  dist + 1
                              UpperCAmelCase_            =  (x, y)

      return np.inf, []


if __name__ == "__main__":
      import doctest

      doctest.testmod()





