'''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    a (       __lowerCAmelCase						,  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    lowercase_  (							self					):

    '''simple docstring'''







    torch.manual_seed(0					)
    __UpperCAmelCase:					Optional[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=snake_case_   ,    )
    __UpperCAmelCase:					Union[str, Any]		= DDIMScheduler(
        beta_start=0.0_0_0_8_5   ,    beta_end=0.0_1_2   ,    beta_schedule="""scaled_linear"""   ,    clip_sample=snake_case_   ,    set_alpha_to_one=snake_case_   ,    )
    torch.manual_seed(0					)
    __UpperCAmelCase:					Optional[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					)
    __UpperCAmelCase:					Optional[Any]		= ClapTextConfig(
        bos_token_id=0   ,    eos_token_id=2   ,    hidden_size=32   ,    intermediate_size=37   ,    layer_norm_eps=1e-0_5   ,    num_attention_heads=4   ,    num_hidden_layers=5   ,    pad_token_id=1   ,    vocab_size=1000   ,    projection_dim=32   ,    )
    __UpperCAmelCase:					List[Any]		= ClapTextModelWithProjection(snake_case_					)
    __UpperCAmelCase:					Union[str, Any]		= RobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-random-roberta"""   ,    model_max_length=77					)

    __UpperCAmelCase:					Dict		= 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=snake_case_   ,    )

    __UpperCAmelCase:					Any		= SpeechTaHifiGan(snake_case_					)

    __UpperCAmelCase:					List[Any]		= {
        """unet""": unet,
        """scheduler""": scheduler,
        """vae""": vae,
        """text_encoder""": text_encoder,
        """tokenizer""": tokenizer,
        """vocoder""": vocoder,
    }
    return components







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

    '''simple docstring'''







    if str(snake_case_					).startswith("""mps"""					):
      __UpperCAmelCase:					Optional[Any]		= torch.manual_seed(snake_case_					)
    else:
      __UpperCAmelCase:					Dict		= torch.Generator(device=snake_case_					).manual_seed(snake_case_					)
    __UpperCAmelCase:					Optional[Any]		= {
        """prompt""": """A hammer hitting a wooden surface""",
        """generator""": generator,
        """num_inference_steps""": 2,
        """guidance_scale""": 6.0,
    }
    return inputs







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					Optional[Any]		= """cpu"""  # ensure determinism for the device-dependent torch.Generator

    __UpperCAmelCase:					Dict		= self.get_dummy_components()
    __UpperCAmelCase:					List[Any]		= AudioLDMPipeline(**snake_case_					)
    __UpperCAmelCase:					str		= audioldm_pipe.to(snake_case_					)
    audioldm_pipe.set_progress_bar_config(disable=snake_case_					)

    __UpperCAmelCase:					Any		= self.get_dummy_inputs(snake_case_					)
    __UpperCAmelCase:					Union[str, Any]		= audioldm_pipe(**snake_case_					)
    __UpperCAmelCase:					Union[str, Any]		= output.audios[0]

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

    __UpperCAmelCase:					Dict		= audio[:10]
    __UpperCAmelCase:					Tuple		= np.array(
        [-0.0_0_5_0, 0.0_0_5_0, -0.0_0_6_0, 0.0_0_3_3, -0.0_0_2_6, 0.0_0_3_3, -0.0_0_2_7, 0.0_0_3_3, -0.0_0_2_8, 0.0_0_3_3]					)

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







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					Any		= self.get_dummy_components()
    __UpperCAmelCase:					List[str]		= AudioLDMPipeline(**snake_case_					)
    __UpperCAmelCase:					Optional[Any]		= audioldm_pipe.to(snake_case_					)
    __UpperCAmelCase:					List[Any]		= audioldm_pipe.to(snake_case_					)
    audioldm_pipe.set_progress_bar_config(disable=snake_case_					)

    __UpperCAmelCase:					int		= self.get_dummy_inputs(snake_case_					)
    __UpperCAmelCase:					Dict		= 3 * [inputs["""prompt"""]]

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

    __UpperCAmelCase:					Tuple		= self.get_dummy_inputs(snake_case_					)
    __UpperCAmelCase:					Any		= 3 * [inputs.pop("""prompt"""					)]

    __UpperCAmelCase:					Union[str, Any]		= audioldm_pipe.tokenizer(
        snake_case_   ,    padding="""max_length"""   ,    max_length=audioldm_pipe.tokenizer.model_max_length   ,    truncation=snake_case_   ,    return_tensors="""pt"""   ,    )
    __UpperCAmelCase:					str		= text_inputs["""input_ids"""].to(snake_case_					)

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

    __UpperCAmelCase:					Union[str, Any]		= prompt_embeds

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

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







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					List[Any]		= self.get_dummy_components()
    __UpperCAmelCase:					List[Any]		= AudioLDMPipeline(**snake_case_					)
    __UpperCAmelCase:					Tuple		= audioldm_pipe.to(snake_case_					)
    __UpperCAmelCase:					Tuple		= audioldm_pipe.to(snake_case_					)
    audioldm_pipe.set_progress_bar_config(disable=snake_case_					)

    __UpperCAmelCase:					Dict		= self.get_dummy_inputs(snake_case_					)
    __UpperCAmelCase:					Union[str, Any]		= 3 * ["""this is a negative prompt"""]
    __UpperCAmelCase:					Tuple		= negative_prompt
    __UpperCAmelCase:					List[Any]		= 3 * [inputs["""prompt"""]]

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

    __UpperCAmelCase:					Optional[int]		= self.get_dummy_inputs(snake_case_					)
    __UpperCAmelCase:					Union[str, Any]		= 3 * [inputs.pop("""prompt"""					)]

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

      __UpperCAmelCase:					Union[str, Any]		= audioldm_pipe.text_encoder(
          snake_case_   ,    )
      __UpperCAmelCase:					Dict		= text_embeds.text_embeds
      # additional L_2 normalization over each hidden-state
      __UpperCAmelCase:					List[str]		= F.normalize(snake_case_   ,    dim=-1					)

      embeds.append(snake_case_					)

    __UpperCAmelCase:					Dict		= embeds

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

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







  def    lowercase_  (							self					):

    '''simple docstring'''







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

    __UpperCAmelCase:					Dict		= self.get_dummy_inputs(snake_case_					)
    __UpperCAmelCase:					Union[str, Any]		= """egg cracking"""
    __UpperCAmelCase:					str		= audioldm_pipe(**snake_case_   ,    negative_prompt=snake_case_					)
    __UpperCAmelCase:					int		= output.audios[0]

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

    __UpperCAmelCase:					Union[str, Any]		= audio[:10]
    __UpperCAmelCase:					Any		= np.array(
        [-0.0_0_5_1, 0.0_0_5_0, -0.0_0_6_0, 0.0_0_3_4, -0.0_0_2_6, 0.0_0_3_3, -0.0_0_2_7, 0.0_0_3_3, -0.0_0_2_8, 0.0_0_3_2]					)

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







  def    lowercase_  (							self					):

    '''simple docstring'''







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

    __UpperCAmelCase:					Optional[Any]		= """A hammer hitting a wooden surface"""

    # test num_waveforms_per_prompt=1 (default)
    __UpperCAmelCase:					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
    __UpperCAmelCase:					Union[str, Any]		= 2
    __UpperCAmelCase:					List[Any]		= 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:					Optional[Any]		= 2
    __UpperCAmelCase:					int		= 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
    __UpperCAmelCase:					Tuple		= 2
    __UpperCAmelCase:					List[Any]		= 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    lowercase_  (							self					):

    '''simple docstring'''







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

    __UpperCAmelCase:					Tuple		= self.get_dummy_inputs(snake_case_					)
    __UpperCAmelCase:					Tuple		= audioldm_pipe(audio_length_in_s=0.0_1_6   ,    **snake_case_					)
    __UpperCAmelCase:					Optional[Any]		= output.audios[0]

    assert audio.ndim == 1
    assert len(snake_case_					) / vocoder_sampling_rate == 0.0_1_6

    __UpperCAmelCase:					Optional[Any]		= audioldm_pipe(audio_length_in_s=0.0_3_2   ,    **snake_case_					)
    __UpperCAmelCase:					str		= output.audios[0]

    assert audio.ndim == 1
    assert len(snake_case_					) / vocoder_sampling_rate == 0.0_3_2







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					Optional[int]		= self.get_dummy_components()
    __UpperCAmelCase:					Optional[Any]		= AudioLDMPipeline(**snake_case_					)
    __UpperCAmelCase:					int		= audioldm_pipe.to(snake_case_					)
    audioldm_pipe.set_progress_bar_config(disable=snake_case_					)

    __UpperCAmelCase:					List[Any]		= ["""hey"""]

    __UpperCAmelCase:					Tuple		= audioldm_pipe(snake_case_   ,    num_inference_steps=1					)
    __UpperCAmelCase:					str		= output.audios.shape
    assert audio_shape == (1, 256)

    __UpperCAmelCase:					List[Any]		= audioldm_pipe.vocoder.config
    config.model_in_dim *= 2
    __UpperCAmelCase:					List[str]		= SpeechTaHifiGan(snake_case_					).to(snake_case_					)
    __UpperCAmelCase:					Tuple		= audioldm_pipe(snake_case_   ,    num_inference_steps=1					)
    __UpperCAmelCase:					Union[str, Any]		= 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    lowercase_  (							self					):

    '''simple docstring'''







    self._test_attention_slicing_forward_pass(test_mean_pixel_difference=snake_case_					)







  def    lowercase_  (							self					):

    '''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    lowercase_  (							self					):

    '''simple docstring'''







    self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=snake_case_					)




@slow
class    a (       unittest.TestCase				):



  """simple docstring"""







  def    lowercase_  (							self					):

    '''simple docstring'''







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







  def    lowercase_  (							self   ,    snake_case_   ,    snake_case_="cpu"   ,    snake_case_=torch.floataa   ,    snake_case_=0					):

    '''simple docstring'''







    __UpperCAmelCase:					List[str]		= torch.Generator(device=snake_case_					).manual_seed(snake_case_					)
    __UpperCAmelCase:					Any		= np.random.RandomState(snake_case_					).standard_normal((1, 8, 128, 16)					)
    __UpperCAmelCase:					str		= torch.from_numpy(snake_case_					).to(device=snake_case_   ,    dtype=snake_case_					)
    __UpperCAmelCase:					int		= {
        """prompt""": """A hammer hitting a wooden surface""",
        """latents""": latents,
        """generator""": generator,
        """num_inference_steps""": 3,
        """guidance_scale""": 2.5,
    }
    return inputs







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					Optional[int]		= AudioLDMPipeline.from_pretrained("""cvssp/audioldm"""					)
    __UpperCAmelCase:					Union[str, Any]		= audioldm_pipe.to(snake_case_					)
    audioldm_pipe.set_progress_bar_config(disable=snake_case_					)

    __UpperCAmelCase:					int		= self.get_inputs(snake_case_					)
    __UpperCAmelCase:					Any		= 25
    __UpperCAmelCase:					str		= audioldm_pipe(**snake_case_					).audios[0]

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

    __UpperCAmelCase:					Dict		= audio[7_7230:7_7240]
    __UpperCAmelCase:					Optional[int]		= np.array(
        [-0.4_8_8_4, -0.4_6_0_7, 0.0_0_2_3, 0.5_0_0_7, 0.5_8_9_6, 0.5_1_5_1, 0.3_8_1_3, -0.0_2_0_8, -0.3_6_8_7, -0.4_3_1_5]					)
    __UpperCAmelCase:					List[Any]		= np.abs(expected_slice - audio_slice					).max()
    assert max_diff < 1e-2







  def    lowercase_  (							self					):

    '''simple docstring'''







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

    __UpperCAmelCase:					Dict		= self.get_inputs(snake_case_					)
    __UpperCAmelCase:					Dict		= audioldm_pipe(**snake_case_					).audios[0]

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

    __UpperCAmelCase:					Tuple		= audio[2_7780:2_7790]
    __UpperCAmelCase:					Any		= np.array([-0.2_1_3_1, -0.0_8_7_3, -0.0_1_2_4, -0.0_1_8_9, 0.0_5_6_9, 0.1_3_7_3, 0.1_8_8_3, 0.2_8_8_6, 0.3_2_9_7, 0.2_2_1_2]					)
    __UpperCAmelCase:					List[str]		= np.abs(expected_slice - audio_slice					).max()
    assert max_diff < 3e-2
'''simple docstring'''






import itertools
import math






def  UpperCamelCase__					(					_lowercase    :				int     )							->	bool:

  if 1 < number < 4:
    # 2 and 3 are primes
    return True
  elif number < 2 or number % 2 == 0 or number % 3 == 0:
    # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
    return False

  # All primes number are in format of 6k +/- 1
  for i in range(5   ,    int(math.sqrt(_lowercase     ) + 1     )   ,    6     ):
    if number % i == 0 or number % (i + 2) == 0:
      return False
  return True






def  UpperCamelCase__					(					)							->	Optional[int]:
  __UpperCAmelCase:					Union[str, Any]		= 2
  while True:
    if is_prime(_lowercase     ):
      yield num
    num += 1






def  UpperCamelCase__					(					_lowercase    :				int = 1_0_0_0_1     )							->	int:
  return next(itertools.islice(prime_generator()   ,    nth - 1   ,    _lowercase     )     )


if __name__ == "__main__":
    print(F"""{solution() = }""")