


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 	snake_case_       (_UpperCAmelCase     ,						unittest.TestCase				):
					UpperCAmelCase__ :   Any   				=      AudioLDMPipeline
					UpperCAmelCase__ :   Union[str, Any]   				=      TEXT_TO_AUDIO_PARAMS
					UpperCAmelCase__ :   List[str]   				=      TEXT_TO_AUDIO_BATCH_PARAMS
					UpperCAmelCase__ :   Tuple   				=      frozenset(
					    [
					        '''num_inference_steps''',
					        '''num_waveforms_per_prompt''',
					        '''generator''',
					        '''latents''',
					        '''output_type''',
					        '''return_dict''',
					        '''callback''',
					        '''callback_steps''',
					    ]				)
					def 			lowerCamelCase__(				self       :int				) ->   Any:
												torch.manual_seed(0				)
												a__					=					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=SCREAMING_SNAKE_CASE_      ,)
												a__					=					DDIMScheduler(
												    beta_start=0.0_00_85      ,beta_end=0.0_12      ,beta_schedule='scaled_linear'      ,clip_sample=SCREAMING_SNAKE_CASE_      ,set_alpha_to_one=SCREAMING_SNAKE_CASE_      ,)
												torch.manual_seed(0				)
												a__					=					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__					=					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=10_00      ,projection_dim=32      ,)
												a__					=					ClapTextModelWithProjection(SCREAMING_SNAKE_CASE_				)
												a__					=					RobertaTokenizer.from_pretrained('hf-internal-testing/tiny-random-roberta'      ,model_max_length=77				)

												a__					=					SpeechTaHifiGanConfig(
												    model_in_dim=8      ,sampling_rate=1_60_00      ,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=SCREAMING_SNAKE_CASE_      ,)

												a__					=					SpeechTaHifiGan(SCREAMING_SNAKE_CASE_				)

												a__					=					{
												    """unet""": unet,
												    """scheduler""": scheduler,
												    """vae""": vae,
												    """text_encoder""": text_encoder,
												    """tokenizer""": tokenizer,
												    """vocoder""": vocoder,
												}
												return components
					def 			lowerCamelCase__(				self       :Tuple      ,__snake_case       :Optional[Any]      ,__snake_case       :int=0				) ->   Optional[int]:
												if str(SCREAMING_SNAKE_CASE_				).startswith('mps'				):
																			a__					=					torch.manual_seed(SCREAMING_SNAKE_CASE_				)
												else:
																			a__					=					torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(SCREAMING_SNAKE_CASE_				)
												a__					=					{
												    """prompt""": """A hammer hitting a wooden surface""",
												    """generator""": generator,
												    """num_inference_steps""": 2,
												    """guidance_scale""": 6.0,
												}
												return inputs
					def 			lowerCamelCase__(				self       :Tuple				) ->   Optional[Any]:
												a__					=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator

												a__					=					self.get_dummy_components()
												a__					=					AudioLDMPipeline(**SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

												a__					=					self.get_dummy_inputs(SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe(**SCREAMING_SNAKE_CASE_				)
												a__					=					output.audios[0]

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE_				) == 2_56

												a__					=					audio[:10]
												a__					=					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 			lowerCamelCase__(				self       :Tuple				) ->   Optional[Any]:
												a__					=					self.get_dummy_components()
												a__					=					AudioLDMPipeline(**SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

												a__					=					self.get_dummy_inputs(SCREAMING_SNAKE_CASE_				)
												a__					=					3 * [inputs["""prompt"""]]

												# forward
												a__					=					audioldm_pipe(**SCREAMING_SNAKE_CASE_				)
												a__					=					output.audios[0]

												a__					=					self.get_dummy_inputs(SCREAMING_SNAKE_CASE_				)
												a__					=					3 * [inputs.pop('prompt'				)]

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

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

												a__					=					prompt_embeds

												# forward
												a__					=					audioldm_pipe(**SCREAMING_SNAKE_CASE_				)
												a__					=					output.audios[0]

												assert np.abs(audio_a - audio_a				).max() < 1E-2
					def 			lowerCamelCase__(				self       :Dict				) ->   str:
												a__					=					self.get_dummy_components()
												a__					=					AudioLDMPipeline(**SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

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

												# forward
												a__					=					audioldm_pipe(**SCREAMING_SNAKE_CASE_				)
												a__					=					output.audios[0]

												a__					=					self.get_dummy_inputs(SCREAMING_SNAKE_CASE_				)
												a__					=					3 * [inputs.pop('prompt'				)]

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

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

																			embeds.append(SCREAMING_SNAKE_CASE_				)

												a__					=					embeds

												# forward
												a__					=					audioldm_pipe(**SCREAMING_SNAKE_CASE_				)
												a__					=					output.audios[0]

												assert np.abs(audio_a - audio_a				).max() < 1E-2
					def 			lowerCamelCase__(				self       :Any				) ->   List[Any]:
												a__					=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
												a__					=					self.get_dummy_components()
												a__					=					PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE_				)
												a__					=					AudioLDMPipeline(**SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

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

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE_				) == 2_56

												a__					=					audio[:10]
												a__					=					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 			lowerCamelCase__(				self       :Optional[int]				) ->   List[Any]:
												a__					=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
												a__					=					self.get_dummy_components()
												a__					=					PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE_				)
												a__					=					AudioLDMPipeline(**SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

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

												# test num_waveforms_per_prompt=1 (default)
												a__					=					audioldm_pipe(SCREAMING_SNAKE_CASE_      ,num_inference_steps=2				).audios

												assert audios.shape == (1, 2_56)

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

												assert audios.shape == (batch_size, 2_56)

												# test num_waveforms_per_prompt for single prompt
												a__					=					2
												a__					=					audioldm_pipe(SCREAMING_SNAKE_CASE_      ,num_inference_steps=2      ,num_waveforms_per_prompt=SCREAMING_SNAKE_CASE_				).audios

												assert audios.shape == (num_waveforms_per_prompt, 2_56)

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

												assert audios.shape == (batch_size * num_waveforms_per_prompt, 2_56)
					def 			lowerCamelCase__(				self       :Optional[int]				) ->   Tuple:
												a__					=					"""cpu"""  # ensure determinism for the device-dependent torch.Generator
												a__					=					self.get_dummy_components()
												a__					=					AudioLDMPipeline(**SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe.vocoder.config.sampling_rate

												a__					=					self.get_dummy_inputs(SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe(audio_length_in_s=0.0_16      ,**SCREAMING_SNAKE_CASE_				)
												a__					=					output.audios[0]

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

												a__					=					audioldm_pipe(audio_length_in_s=0.0_32      ,**SCREAMING_SNAKE_CASE_				)
												a__					=					output.audios[0]

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE_				) / vocoder_sampling_rate == 0.0_32
					def 			lowerCamelCase__(				self       :List[Any]				) ->   str:
												a__					=					self.get_dummy_components()
												a__					=					AudioLDMPipeline(**SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

												a__					=					["""hey"""]

												a__					=					audioldm_pipe(SCREAMING_SNAKE_CASE_      ,num_inference_steps=1				)
												a__					=					output.audios.shape
												assert audio_shape == (1, 2_56)

												a__					=					audioldm_pipe.vocoder.config
												config.model_in_dim *= 2
												a__					=					SpeechTaHifiGan(SCREAMING_SNAKE_CASE_				).to(SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe(SCREAMING_SNAKE_CASE_      ,num_inference_steps=1				)
												a__					=					output.audios.shape
												# waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
												assert audio_shape == (1, 2_56)
					def 			lowerCamelCase__(				self       :str				) ->   Dict:
												self._test_attention_slicing_forward_pass(test_mean_pixel_difference=SCREAMING_SNAKE_CASE_				)
					def 			lowerCamelCase__(				self       :int				) ->   Optional[Any]:
												self._test_inference_batch_single_identical(test_mean_pixel_difference=SCREAMING_SNAKE_CASE_				)



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







@slow
class 	snake_case_       (unittest.TestCase				):
					def 			lowerCamelCase__(				self       :Optional[int]				) ->   int:
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()
					def 			lowerCamelCase__(				self       :Optional[Any]      ,__snake_case       :Tuple      ,__snake_case       :Union[str, Any]="cpu"      ,__snake_case       :Union[str, Any]=torch.floataa      ,__snake_case       :Optional[int]=0				) ->   List[Any]:
												a__					=					torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(SCREAMING_SNAKE_CASE_				)
												a__					=					np.random.RandomState(SCREAMING_SNAKE_CASE_				).standard_normal((1, 8, 1_28, 16)				)
												a__					=					torch.from_numpy(SCREAMING_SNAKE_CASE_				).to(device=SCREAMING_SNAKE_CASE_      ,dtype=SCREAMING_SNAKE_CASE_				)
												a__					=					{
												    """prompt""": """A hammer hitting a wooden surface""",
												    """latents""": latents,
												    """generator""": generator,
												    """num_inference_steps""": 3,
												    """guidance_scale""": 2.5,
												}
												return inputs
					def 			lowerCamelCase__(				self       :str				) ->   str:
												a__					=					AudioLDMPipeline.from_pretrained('cvssp/audioldm'				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

												a__					=					self.get_inputs(SCREAMING_SNAKE_CASE_				)
												a__					=					25
												a__					=					audioldm_pipe(**SCREAMING_SNAKE_CASE_				).audios[0]

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE_				) == 8_19_20

												a__					=					audio[7_72_30:7_72_40]
												a__					=					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]				)
												a__					=					np.abs(expected_slice - audio_slice				).max()
												assert max_diff < 1E-2



					def 			lowerCamelCase__(				self       :Union[str, Any]				) ->   Optional[Any]:
												a__					=					AudioLDMPipeline.from_pretrained('cvssp/audioldm'				)
												a__					=					LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config				)
												a__					=					audioldm_pipe.to(SCREAMING_SNAKE_CASE_				)
												audioldm_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

												a__					=					self.get_inputs(SCREAMING_SNAKE_CASE_				)
												a__					=					audioldm_pipe(**SCREAMING_SNAKE_CASE_				).audios[0]

												assert audio.ndim == 1
												assert len(SCREAMING_SNAKE_CASE_				) == 8_19_20

												a__					=					audio[2_77_80:2_77_90]
												a__					=					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]				)
												a__					=					np.abs(expected_slice - audio_slice				).max()
												assert max_diff < 3E-2





from collections import defaultdict
from math import ceil, sqrt






def 							__lowercase       (			__lowerCAmelCase							:    int = 1_0_0_0_0_0_0 ,   __lowerCAmelCase							:    int = 1_0       ):
							a__					=					defaultdict(__lowerCAmelCase       )

							for outer_width in range(3 ,   (t_limit // 4) + 2       ):
														if outer_width * outer_width > t_limit:
																					a__					=					max(
																					    ceil(sqrt(outer_width * outer_width - t_limit       )       ) ,   1       )
														else:
																					a__					=					1

														hole_width_lower_bound += (outer_width - hole_width_lower_bound) % 2

														for hole_width in range(__lowerCAmelCase ,   outer_width - 1 ,   2       ):
																					count[outer_width * outer_width - hole_width * hole_width] += 1

							return sum(1 for n in count.values() if 1 <= n <= 1_0       )


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




