"""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	(    __lowercase		,   unittest.TestCase						):
   lowercase							=						AudioLDMPipeline
   lowercase							=						TEXT_TO_AUDIO_PARAMS
   lowercase							=						TEXT_TO_AUDIO_BATCH_PARAMS
   lowercase							=						frozenset(
       [
           'num_inference_steps',
           'num_waveforms_per_prompt',
           'generator',
           'latents',
           'output_type',
           'return_dict',
           'callback',
           'callback_steps',
       ]						)
   def     _UpperCAmelCase     (    self    )			->   int:


        '''simple docstring'''


        torch.manual_seed(0    )
        lowercase_			:       Tuple									=			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=__UpperCamelCase  ,)
        lowercase_			:       str									=			DDIMScheduler(
            beta_start=0.0_0085  ,beta_end=0.012  ,beta_schedule='scaled_linear'  ,clip_sample=__UpperCamelCase  ,set_alpha_to_one=__UpperCamelCase  ,)
        torch.manual_seed(0    )
        lowercase_			:       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    )
        lowercase_			:       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  ,)
        lowercase_			:       Dict									=			ClapTextModelWithProjection(__UpperCamelCase    )
        lowercase_			:       Any									=			RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'  ,model_max_length=77    )

        lowercase_			:       Union[str, Any]									=			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=__UpperCamelCase  ,)

        lowercase_			:       List[Any]									=			SpeechTaHifiGan(__UpperCamelCase    )

        lowercase_			:       Tuple									=			{
            "unet": unet,
            "scheduler": scheduler,
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "vocoder": vocoder,
        }
        return components
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase=0    )			->   int:


        '''simple docstring'''


        if str(__UpperCamelCase    ).startswith('mps'    ):
             lowercase_			:       List[str]									=			torch.manual_seed(__UpperCamelCase    )
        else:
             lowercase_			:       Dict									=			torch.Generator(device=__UpperCamelCase    ).manual_seed(__UpperCamelCase    )
        lowercase_			:       str									=			{
            "prompt": "A hammer hitting a wooden surface",
            "generator": generator,
            "num_inference_steps": 2,
            "guidance_scale": 6.0,
        }
        return inputs
   def     _UpperCAmelCase     (    self    )			->   Any:


        '''simple docstring'''


        lowercase_			:       Tuple									=			"cpu"  # ensure determinism for the device-dependent torch.Generator

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

        lowercase_			:       Optional[int]									=			self.get_dummy_inputs(__UpperCamelCase    )
        lowercase_			:       Union[str, Any]									=			audioldm_pipe(**__UpperCamelCase    )
        lowercase_			:       Tuple									=			output.audios[0]

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

        lowercase_			:       Optional[Any]									=			audio[:10]
        lowercase_			:       Tuple									=			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     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


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

        lowercase_			:       Optional[Any]									=			self.get_dummy_inputs(__UpperCamelCase    )
        lowercase_			:       Any									=			3 * [inputs["prompt"]]

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

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

        lowercase_			:       Optional[Any]									=			audioldm_pipe.tokenizer(
            __UpperCamelCase  ,padding='max_length'  ,max_length=audioldm_pipe.tokenizer.model_max_length  ,truncation=__UpperCamelCase  ,return_tensors='pt'  ,)
        lowercase_			:       Optional[Any]									=			text_inputs["input_ids"].to(__UpperCamelCase    )

        lowercase_			:       Any									=			audioldm_pipe.text_encoder(
            __UpperCamelCase  ,)
        lowercase_			:       str									=			prompt_embeds.text_embeds
        # additional L_2 normalization over each hidden-state
        lowercase_			:       List[Any]									=			F.normalize(__UpperCamelCase  ,dim=-1    )

        lowercase_			:       Optional[int]									=			prompt_embeds

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

        assert np.abs(audio_a - audio_a    ).max() < 1e-2
   def     _UpperCAmelCase     (    self    )			->   Tuple:


        '''simple docstring'''


        lowercase_			:       str									=			self.get_dummy_components()
        lowercase_			:       str									=			AudioLDMPipeline(**__UpperCamelCase    )
        lowercase_			:       Dict									=			audioldm_pipe.to(__UpperCamelCase    )
        lowercase_			:       Optional[int]									=			audioldm_pipe.to(__UpperCamelCase    )
        audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase    )

        lowercase_			:       Any									=			self.get_dummy_inputs(__UpperCamelCase    )
        lowercase_			:       List[Any]									=			3 * ["this is a negative prompt"]
        lowercase_			:       str									=			negative_prompt
        lowercase_			:       List[Any]									=			3 * [inputs["prompt"]]

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

        lowercase_			:       List[Any]									=			self.get_dummy_inputs(__UpperCamelCase    )
        lowercase_			:       Optional[int]									=			3 * [inputs.pop('prompt'    )]

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

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

             embeds.append(__UpperCamelCase    )

        lowercase_			:       Tuple									=			embeds

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

        assert np.abs(audio_a - audio_a    ).max() < 1e-2
   def     _UpperCAmelCase     (    self    )			->   int:


        '''simple docstring'''


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

        lowercase_			:       str									=			self.get_dummy_inputs(__UpperCamelCase    )
        lowercase_			:       Optional[int]									=			"egg cracking"
        lowercase_			:       Optional[int]									=			audioldm_pipe(**__UpperCamelCase  ,negative_prompt=__UpperCamelCase    )
        lowercase_			:       Any									=			output.audios[0]

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

        lowercase_			:       int									=			audio[:10]
        lowercase_			:       Optional[int]									=			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     _UpperCAmelCase     (    self    )			->   List[str]:


        '''simple docstring'''


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

        lowercase_			:       int									=			"A hammer hitting a wooden surface"

        # test num_waveforms_per_prompt=1 (default)
        lowercase_			:       Optional[int]									=			audioldm_pipe(__UpperCamelCase  ,num_inference_steps=2    ).audios

        assert audios.shape == (1, 256)

        # test num_waveforms_per_prompt=1 (default) for batch of prompts
        lowercase_			:       List[Any]									=			2
        lowercase_			:       Union[str, 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
        lowercase_			:       Dict									=			2
        lowercase_			:       str									=			audioldm_pipe(__UpperCamelCase  ,num_inference_steps=2  ,num_waveforms_per_prompt=__UpperCamelCase    ).audios

        assert audios.shape == (num_waveforms_per_prompt, 256)

        # test num_waveforms_per_prompt for batch of prompts
        lowercase_			:       Any									=			2
        lowercase_			:       Dict									=			audioldm_pipe(
            [prompt] * batch_size  ,num_inference_steps=2  ,num_waveforms_per_prompt=__UpperCamelCase    ).audios

        assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)
   def     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


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

        lowercase_			:       Union[str, Any]									=			self.get_dummy_inputs(__UpperCamelCase    )
        lowercase_			:       List[str]									=			audioldm_pipe(audio_length_in_s=0.016  ,**__UpperCamelCase    )
        lowercase_			:       Optional[int]									=			output.audios[0]

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

        lowercase_			:       Optional[Any]									=			audioldm_pipe(audio_length_in_s=0.032  ,**__UpperCamelCase    )
        lowercase_			:       Tuple									=			output.audios[0]

        assert audio.ndim == 1
        assert len(__UpperCamelCase    ) / vocoder_sampling_rate == 0.032
   def     _UpperCAmelCase     (    self    )			->   Any:


        '''simple docstring'''


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

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

        lowercase_			:       Any									=			audioldm_pipe(__UpperCamelCase  ,num_inference_steps=1    )
        lowercase_			:       Any									=			output.audios.shape
        assert audio_shape == (1, 256)

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


        '''simple docstring'''


        self._test_attention_slicing_forward_pass(test_mean_pixel_difference=__UpperCamelCase    )
   def     _UpperCAmelCase     (    self    )			->   List[str]:


        '''simple docstring'''


        self._test_inference_batch_single_identical(test_mean_pixel_difference=__UpperCamelCase    )







   @unittest.skipIf(
       torch_device != 'cuda' or not is_xformers_available()  ,reason='XFormers attention is only available with CUDA and `xformers` installed'  ,)
   def     _UpperCAmelCase     (    self    )			->   Union[str, Any]:


        '''simple docstring'''


        self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=__UpperCamelCase    )







@slow
class  UpperCamelCase	(    unittest.TestCase						):
   def     _UpperCAmelCase     (    self    )			->   List[Any]:


        '''simple docstring'''


        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase="cpu"  ,__UpperCamelCase=torch.floataa  ,__UpperCamelCase=0    )			->   List[str]:


        '''simple docstring'''


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


        '''simple docstring'''


        lowercase_			:       Union[str, Any]									=			AudioLDMPipeline.from_pretrained('cvssp/audioldm'    )
        lowercase_			:       Union[str, Any]									=			audioldm_pipe.to(__UpperCamelCase    )
        audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase    )

        lowercase_			:       Dict									=			self.get_inputs(__UpperCamelCase    )
        lowercase_			:       List[str]									=			25
        lowercase_			:       Dict									=			audioldm_pipe(**__UpperCamelCase    ).audios[0]

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

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







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


        '''simple docstring'''


        lowercase_			:       List[str]									=			AudioLDMPipeline.from_pretrained('cvssp/audioldm'    )
        lowercase_			:       Optional[int]									=			LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config    )
        lowercase_			:       Union[str, Any]									=			audioldm_pipe.to(__UpperCamelCase    )
        audioldm_pipe.set_progress_bar_config(disable=__UpperCamelCase    )

        lowercase_			:       int									=			self.get_inputs(__UpperCamelCase    )
        lowercase_			:       List[str]									=			audioldm_pipe(**__UpperCamelCase    ).audios[0]

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

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


"""simple docstring"""







import os
import random
import sys

from . import cryptomath_module as cryptoMath  # noqa: N812
from . import rabin_miller as rabinMiller  # noqa: N812



def     lowercase__( ):
     print('Making key files...'   )
     make_key_files('rsa'  ,    10_24   )
     print('Key files generation successful.'   )



def     lowercase__( __SCREAMING_SNAKE_CASE							:  int   ):
     print('Generating prime p...'   )
     lowercase_			:       List[str]									=			rabinMiller.generate_large_prime(__SCREAMING_SNAKE_CASE   )
     print('Generating prime q...'   )
     lowercase_			:       int									=			rabinMiller.generate_large_prime(__SCREAMING_SNAKE_CASE   )
     lowercase_			:       Union[str, Any]									=			p * q

     print('Generating e that is relatively prime to (p - 1) * (q - 1)...'   )
     while True:
          lowercase_			:       str									=			random.randrange(2 ** (key_size - 1)  ,    2 ** (key_size)   )
          if cryptoMath.gcd(__SCREAMING_SNAKE_CASE  ,    (p - 1) * (q - 1)   ) == 1:
               break

     print('Calculating d that is mod inverse of e...'   )
     lowercase_			:       List[str]									=			cryptoMath.find_mod_inverse(__SCREAMING_SNAKE_CASE  ,    (p - 1) * (q - 1)   )

     lowercase_			:       Any									=			(n, e)
     lowercase_			:       int									=			(n, d)
     return (public_key, private_key)



def     lowercase__( __SCREAMING_SNAKE_CASE							:  str  ,    __SCREAMING_SNAKE_CASE							:  int   ):
     if os.path.exists(F'''{name}_pubkey.txt'''   ) or os.path.exists(F'''{name}_privkey.txt'''   ):
          print('\nWARNING:'   )
          print(
              F'''"{name}_pubkey.txt" or "{name}_privkey.txt" already exists. \n'''
              'Use a different name or delete these files and re-run this program.'   )
          sys.exit()

     lowercase_							,						lowercase_			:       int									=			generate_key(__SCREAMING_SNAKE_CASE   )
     print(F'''\nWriting public key to file {name}_pubkey.txt...'''   )
     with open(F'''{name}_pubkey.txt'''  ,    'w'   ) as out_file:
          out_file.write(F'''{key_size},{public_key[0]},{public_key[1]}'''   )

     print(F'''Writing private key to file {name}_privkey.txt...'''   )
     with open(F'''{name}_privkey.txt'''  ,    'w'   ) as out_file:
          out_file.write(F'''{key_size},{private_key[0]},{private_key[1]}'''   )


if __name__ == "__main__":
       main()

