



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__    (							snake_case__   ,       unittest.TestCase):




		'''simple docstring'''



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



		def 	lowerCAmelCase__   (self				)  ->    Optional[Any]:

						"""simple docstring"""






						torch.manual_seed(0				)
						lowerCAmelCase__  :       Optional[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=UpperCAmelCase_	,)
						lowerCAmelCase__  :       List[Any]				=  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				)
						lowerCAmelCase__  :       int				=  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				)
						lowerCAmelCase__  :       str				=  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	,)
						lowerCAmelCase__  :       Optional[Any]				=  ClapTextModelWithProjection(UpperCAmelCase_				)
						lowerCAmelCase__  :       Any				=  RobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-random-roberta'''	,model_max_length=77				)

						lowerCAmelCase__  :       Optional[Any]				=  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=UpperCAmelCase_	,)

						lowerCAmelCase__  :       List[Any]				=  SpeechTaHifiGan(UpperCAmelCase_				)

						lowerCAmelCase__  :       Union[str, Any]				=  {
						    '''unet''': unet,
						    '''scheduler''': scheduler,
						    '''vae''': vae,
						    '''text_encoder''': text_encoder,
						    '''tokenizer''': tokenizer,
						    '''vocoder''': vocoder,
						}
						return components



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase=0				)  ->    Tuple:

						"""simple docstring"""






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



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       List[str]				=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator

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

						lowerCAmelCase__  :       Union[str, Any]				=  self.get_dummy_inputs(UpperCAmelCase_				)
						lowerCAmelCase__  :       Optional[Any]				=  audioldm_pipe(**UpperCAmelCase_				)
						lowerCAmelCase__  :       List[str]				=  output.audios[0]

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

						lowerCAmelCase__  :       Dict				=  audio[:10]
						lowerCAmelCase__  :       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 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  self.get_dummy_components()
						lowerCAmelCase__  :       List[Any]				=  AudioLDMPipeline(**UpperCAmelCase_				)
						lowerCAmelCase__  :       List[Any]				=  audioldm_pipe.to(UpperCAmelCase_				)
						lowerCAmelCase__  :       int				=  audioldm_pipe.to(UpperCAmelCase_				)
						audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase_				)

						lowerCAmelCase__  :       str				=  self.get_dummy_inputs(UpperCAmelCase_				)
						lowerCAmelCase__  :       List[str]				=  3 * [inputs['''prompt''']]

						# forward
						lowerCAmelCase__  :       Dict				=  audioldm_pipe(**UpperCAmelCase_				)
						lowerCAmelCase__  :       int				=  output.audios[0]

						lowerCAmelCase__  :       Optional[int]				=  self.get_dummy_inputs(UpperCAmelCase_				)
						lowerCAmelCase__  :       Any				=  3 * [inputs.pop('''prompt'''				)]

						lowerCAmelCase__  :       int				=  audioldm_pipe.tokenizer(
						    UpperCAmelCase_	,padding='''max_length'''	,max_length=audioldm_pipe.tokenizer.model_max_length	,truncation=UpperCAmelCase_	,return_tensors='''pt'''	,)
						lowerCAmelCase__  :       List[str]				=  text_inputs['''input_ids'''].to(UpperCAmelCase_				)

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

						lowerCAmelCase__  :       Tuple				=  prompt_embeds

						# forward
						lowerCAmelCase__  :       List[str]				=  audioldm_pipe(**UpperCAmelCase_				)
						lowerCAmelCase__  :       Optional[int]				=  output.audios[0]

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



		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  self.get_dummy_components()
						lowerCAmelCase__  :       List[str]				=  AudioLDMPipeline(**UpperCAmelCase_				)
						lowerCAmelCase__  :       List[Any]				=  audioldm_pipe.to(UpperCAmelCase_				)
						lowerCAmelCase__  :       Union[str, Any]				=  audioldm_pipe.to(UpperCAmelCase_				)
						audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase_				)

						lowerCAmelCase__  :       List[Any]				=  self.get_dummy_inputs(UpperCAmelCase_				)
						lowerCAmelCase__  :       Tuple				=  3 * ['''this is a negative prompt''']
						lowerCAmelCase__  :       Tuple				=  negative_prompt
						lowerCAmelCase__  :       Union[str, Any]				=  3 * [inputs['''prompt''']]

						# forward
						lowerCAmelCase__  :       List[str]				=  audioldm_pipe(**UpperCAmelCase_				)
						lowerCAmelCase__  :       Dict				=  output.audios[0]

						lowerCAmelCase__  :       Tuple				=  self.get_dummy_inputs(UpperCAmelCase_				)
						lowerCAmelCase__  :       Union[str, Any]				=  3 * [inputs.pop('''prompt'''				)]

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

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

										embeds.append(UpperCAmelCase_				)

						lowerCAmelCase__		,	lowerCAmelCase__  :       Dict				=  embeds

						# forward
						lowerCAmelCase__  :       Dict				=  audioldm_pipe(**UpperCAmelCase_				)
						lowerCAmelCase__  :       Optional[Any]				=  output.audios[0]

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



		def 	lowerCAmelCase__   (self				)  ->    int:

						"""simple docstring"""






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

						lowerCAmelCase__  :       str				=  self.get_dummy_inputs(UpperCAmelCase_				)
						lowerCAmelCase__  :       List[str]				=  '''egg cracking'''
						lowerCAmelCase__  :       List[Any]				=  audioldm_pipe(**UpperCAmelCase_	,negative_prompt=UpperCAmelCase_				)
						lowerCAmelCase__  :       List[str]				=  output.audios[0]

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

						lowerCAmelCase__  :       Any				=  audio[:10]
						lowerCAmelCase__  :       str				=  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				)  ->    Optional[int]:

						"""simple docstring"""






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

						lowerCAmelCase__  :       Union[str, Any]				=  '''A hammer hitting a wooden surface'''

						# test num_waveforms_per_prompt=1 (default)
						lowerCAmelCase__  :       int				=  audioldm_pipe(UpperCAmelCase_	,num_inference_steps=2				).audios

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

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

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

						# test num_waveforms_per_prompt for batch of prompts
						lowerCAmelCase__  :       List[Any]				=  2
						lowerCAmelCase__  :       str				=  audioldm_pipe(
						    [prompt] * batch_size	,num_inference_steps=2	,num_waveforms_per_prompt=UpperCAmelCase_				).audios

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



		def 	lowerCAmelCase__   (self				)  ->    Optional[Any]:

						"""simple docstring"""






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

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

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

						lowerCAmelCase__  :       Union[str, Any]				=  audioldm_pipe(audio_length_in_s=0.032	,**UpperCAmelCase_				)
						lowerCAmelCase__  :       Union[str, Any]				=  output.audios[0]

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



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






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

						lowerCAmelCase__  :       Dict				=  ['''hey''']

						lowerCAmelCase__  :       str				=  audioldm_pipe(UpperCAmelCase_	,num_inference_steps=1				)
						lowerCAmelCase__  :       str				=  output.audios.shape
						assert audio_shape == (1, 2_56)

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

						"""simple docstring"""






						self._test_attention_slicing_forward_pass(test_mean_pixel_difference=UpperCAmelCase_				)



		def 	lowerCAmelCase__   (self				)  ->    List[Any]:

						"""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 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=UpperCAmelCase_				)



@slow
class  lowerCamelCase__    (							unittest.TestCase):




		'''simple docstring'''



		def 	lowerCAmelCase__   (self				)  ->    int:

						"""simple docstring"""






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



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase="cpu"	,__lowerCamelCase=torch.floataa	,__lowerCamelCase=0				)  ->    List[str]:

						"""simple docstring"""






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



		def 	lowerCAmelCase__   (self				)  ->    List[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''				)
						lowerCAmelCase__  :       Optional[int]				=  audioldm_pipe.to(UpperCAmelCase_				)
						audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase_				)

						lowerCAmelCase__  :       List[str]				=  self.get_inputs(UpperCAmelCase_				)
						lowerCAmelCase__  :       int				=  25
						lowerCAmelCase__  :       Union[str, Any]				=  audioldm_pipe(**UpperCAmelCase_				).audios[0]

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

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



		def 	lowerCAmelCase__   (self				)  ->    Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  AudioLDMPipeline.from_pretrained('''cvssp/audioldm'''				)
						lowerCAmelCase__  :       Any				=  LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config				)
						lowerCAmelCase__  :       str				=  audioldm_pipe.to(UpperCAmelCase_				)
						audioldm_pipe.set_progress_bar_config(disable=UpperCAmelCase_				)

						lowerCAmelCase__  :       Tuple				=  self.get_inputs(UpperCAmelCase_				)
						lowerCAmelCase__  :       Tuple				=  audioldm_pipe(**UpperCAmelCase_				).audios[0]

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

						lowerCAmelCase__  :       List[Any]				=  audio[2_77_80:2_77_90]
						lowerCAmelCase__  :       Tuple				=  np.array([-0.2131, -0.0873, -0.0124, -0.0189, 0.0569, 0.1373, 0.1883, 0.2886, 0.3297, 0.2212]				)
						lowerCAmelCase__  :       List[str]				=  np.abs(expected_slice - audio_slice				).max()
						assert max_diff < 3e-2






import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()



class  lowerCamelCase__    (							lowerCamelCase__   ,       lowerCamelCase__   ,       unittest.TestCase):




		'''simple docstring'''



		snake_case_			=AutoencoderKL
		snake_case_			="""sample"""
		snake_case_			=1e-2



		@property
		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  4
						lowerCAmelCase__  :       int				=  3
						lowerCAmelCase__  :       List[Any]				=  (32, 32)

						lowerCAmelCase__  :       Tuple				=  floats_tensor((batch_size, num_channels) + sizes				).to(__lowerCamelCase				)

						return {"sample": image}



		@property
		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						return (3, 32, 32)



		@property
		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						return (3, 32, 32)



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  {
						    '''block_out_channels''': [32, 64],
						    '''in_channels''': 3,
						    '''out_channels''': 3,
						    '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
						    '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
						    '''latent_channels''': 4,
						}
						lowerCAmelCase__  :       str				=  self.dummy_input
						return init_dict, inputs_dict



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						pass



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						pass



		@unittest.skipIf(torch_device == '''mps'''	,'''Gradient checkpointing skipped on MPS'''				)
		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__		,	lowerCAmelCase__  :       int				=  self.prepare_init_args_and_inputs_for_common()
						lowerCAmelCase__  :       Tuple				=  self.model_class(**__lowerCamelCase				)
						model.to(__lowerCamelCase				)

						assert not model.is_gradient_checkpointing and model.training

						lowerCAmelCase__  :       Tuple				=  model(**__lowerCamelCase				).sample
						# run the backwards pass on the model. For backwards pass, for simplicity purpose,
						# we won't calculate the loss and rather backprop on out.sum()
						model.zero_grad()

						lowerCAmelCase__  :       Optional[Any]				=  torch.randn_like(__lowerCamelCase				)
						lowerCAmelCase__  :       Tuple				=  (out - labels).mean()
						loss.backward()

						# re-instantiate the model now enabling gradient checkpointing
						lowerCAmelCase__  :       str				=  self.model_class(**__lowerCamelCase				)
						# clone model
						model_a.load_state_dict(model.state_dict()				)
						model_a.to(__lowerCamelCase				)
						model_a.enable_gradient_checkpointing()

						assert model_a.is_gradient_checkpointing and model_a.training

						lowerCAmelCase__  :       Dict				=  model_a(**__lowerCamelCase				).sample
						# run the backwards pass on the model. For backwards pass, for simplicity purpose,
						# we won't calculate the loss and rather backprop on out.sum()
						model_a.zero_grad()
						lowerCAmelCase__  :       Any				=  (out_a - labels).mean()
						loss_a.backward()

						# compare the output and parameters gradients
						self.assertTrue((loss - loss_a).abs() < 1e-5				)
						lowerCAmelCase__  :       List[str]				=  dict(model.named_parameters()				)
						lowerCAmelCase__  :       Optional[Any]				=  dict(model_a.named_parameters()				)
						for name, param in named_params.items():
										self.assertTrue(torch_all_close(param.grad.data	,named_params_a[name].grad.data	,atol=5e-5				)				)



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__		,	lowerCAmelCase__  :       str				=  AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''	,output_loading_info=__lowerCamelCase				)
						self.assertIsNotNone(__lowerCamelCase				)
						self.assertEqual(len(loading_info['''missing_keys''']				)	,0				)

						model.to(__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[int]				=  model(**self.dummy_input				)

						assert image is not None, "Make sure output is not None"



		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  AutoencoderKL.from_pretrained('''fusing/autoencoder-kl-dummy'''				)
						lowerCAmelCase__  :       Tuple				=  model.to(__lowerCamelCase				)
						model.eval()

						if torch_device == "mps":
										lowerCAmelCase__  :       Union[str, Any]				=  torch.manual_seed(0				)
						else:
										lowerCAmelCase__  :       Any				=  torch.Generator(device=__lowerCamelCase				).manual_seed(0				)

						lowerCAmelCase__  :       Optional[int]				=  torch.randn(
						    1	,model.config.in_channels	,model.config.sample_size	,model.config.sample_size	,generator=torch.manual_seed(0				)	,)
						lowerCAmelCase__  :       Dict				=  image.to(__lowerCamelCase				)
						with torch.no_grad():
										lowerCAmelCase__  :       List[Any]				=  model(__lowerCamelCase	,sample_posterior=__lowerCamelCase	,generator=__lowerCamelCase				).sample

						lowerCAmelCase__  :       Optional[Any]				=  output[0, -1, -3:, -3:].flatten().cpu()

						# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
						# the expected output slices are not the same for CPU and GPU.
						if torch_device == "mps":
										lowerCAmelCase__  :       List[str]				=  torch.tensor(
										    [
										        -4.0078e-01,
										        -3.8323e-04,
										        -1.2681e-01,
										        -1.1462e-01,
										        2.0095e-01,
										        1.0893e-01,
										        -8.8247e-02,
										        -3.0361e-01,
										        -9.8644e-03,
										    ]				)
						elif torch_device == "cpu":
										lowerCAmelCase__  :       str				=  torch.tensor(
										    [-0.1352, 0.0878, 0.0419, -0.0818, -0.1069, 0.0688, -0.1458, -0.4446, -0.0026]				)
						else:
										lowerCAmelCase__  :       Any				=  torch.tensor(
										    [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]				)

						self.assertTrue(torch_all_close(__lowerCamelCase	,__lowerCamelCase	,rtol=1e-2				)				)



@slow
class  lowerCamelCase__    (							unittest.TestCase):




		'''simple docstring'''



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    List[Any]:

						"""simple docstring"""






						return f"""gaussian_noise_s={seed}_shape={'_'.join([str(__lowerCamelCase				) for s in shape]				)}.npy"""



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






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



		def 	lowerCAmelCase__   (self	,__lowerCamelCase=0	,__lowerCamelCase=(4, 3, 5_12, 5_12)	,__lowerCamelCase=False				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  torch.floataa if fpaa else torch.floataa
						lowerCAmelCase__  :       Dict				=  torch.from_numpy(load_hf_numpy(self.get_file_format(__lowerCamelCase	,__lowerCamelCase				)				)				).to(__lowerCamelCase				).to(__lowerCamelCase				)
						return image



		def 	lowerCAmelCase__   (self	,__lowerCamelCase="CompVis/stable-diffusion-v1-4"	,__lowerCamelCase=False				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  '''fp16''' if fpaa else None
						lowerCAmelCase__  :       Optional[Any]				=  torch.floataa if fpaa else torch.floataa

						lowerCAmelCase__  :       List[Any]				=  AutoencoderKL.from_pretrained(
						    __lowerCamelCase	,subfolder='''vae'''	,torch_dtype=__lowerCamelCase	,revision=__lowerCamelCase	,)
						model.to(__lowerCamelCase				).eval()

						return model



		def 	lowerCAmelCase__   (self	,__lowerCamelCase=0				)  ->    str:

						"""simple docstring"""






						if torch_device == "mps":
										return torch.manual_seed(__lowerCamelCase				)
						return torch.Generator(device=__lowerCamelCase				).manual_seed(__lowerCamelCase				)



		@parameterized.expand(
		    [
		        # fmt: off
		        [33, [-0.1603, 0.9878, -0.0495, -0.0790, -0.2709, 0.8375, -0.2060, -0.0824], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
		        [47, [-0.2376, 0.1168, 0.1332, -0.4840, -0.2508, -0.0791, -0.0493, -0.4089], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
		        # fmt: on
		    ]				)
		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.get_sd_vae_model()
						lowerCAmelCase__  :       Any				=  self.get_sd_image(__lowerCamelCase				)
						lowerCAmelCase__  :       List[str]				=  self.get_generator(__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       Dict				=  model(__lowerCamelCase	,generator=__lowerCamelCase	,sample_posterior=__lowerCamelCase				).sample

						assert sample.shape == image.shape

						lowerCAmelCase__  :       Optional[Any]				=  sample[-1, -2:, -2:, :2].flatten().float().cpu()
						lowerCAmelCase__  :       Any				=  torch.tensor(expected_slice_mps if torch_device == '''mps''' else expected_slice				)

						assert torch_all_close(__lowerCamelCase	,__lowerCamelCase	,atol=3e-3				)



		@parameterized.expand(
		    [
		        # fmt: off
		        [33, [-0.0513, 0.0289, 1.3799, 0.2166, -0.2573, -0.0871, 0.5103, -0.0999]],
		        [47, [-0.4128, -0.1320, -0.3704, 0.1965, -0.4116, -0.2332, -0.3340, 0.2247]],
		        # fmt: on
		    ]				)
		@require_torch_gpu
		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  self.get_sd_vae_model(fpaa=__lowerCamelCase				)
						lowerCAmelCase__  :       List[str]				=  self.get_sd_image(__lowerCamelCase	,fpaa=__lowerCamelCase				)
						lowerCAmelCase__  :       List[Any]				=  self.get_generator(__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       Optional[Any]				=  model(__lowerCamelCase	,generator=__lowerCamelCase	,sample_posterior=__lowerCamelCase				).sample

						assert sample.shape == image.shape

						lowerCAmelCase__  :       Optional[Any]				=  sample[-1, -2:, :2, -2:].flatten().float().cpu()
						lowerCAmelCase__  :       Any				=  torch.tensor(__lowerCamelCase				)

						assert torch_all_close(__lowerCamelCase	,__lowerCamelCase	,atol=1e-2				)



		@parameterized.expand(
		    [
		        # fmt: off
		        [33, [-0.1609, 0.9866, -0.0487, -0.0777, -0.2716, 0.8368, -0.2055, -0.0814], [-0.2395, 0.0098, 0.0102, -0.0709, -0.2840, -0.0274, -0.0718, -0.1824]],
		        [47, [-0.2377, 0.1147, 0.1333, -0.4841, -0.2506, -0.0805, -0.0491, -0.4085], [0.0350, 0.0847, 0.0467, 0.0344, -0.0842, -0.0547, -0.0633, -0.1131]],
		        # fmt: on
		    ]				)
		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  self.get_sd_vae_model()
						lowerCAmelCase__  :       Dict				=  self.get_sd_image(__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       Tuple				=  model(__lowerCamelCase				).sample

						assert sample.shape == image.shape

						lowerCAmelCase__  :       Dict				=  sample[-1, -2:, -2:, :2].flatten().float().cpu()
						lowerCAmelCase__  :       List[Any]				=  torch.tensor(expected_slice_mps if torch_device == '''mps''' else expected_slice				)

						assert torch_all_close(__lowerCamelCase	,__lowerCamelCase	,atol=3e-3				)



		@parameterized.expand(
		    [
		        # fmt: off
		        [13, [-0.2051, -0.1803, -0.2311, -0.2114, -0.3292, -0.3574, -0.2953, -0.3323]],
		        [37, [-0.2632, -0.2625, -0.2199, -0.2741, -0.4539, -0.4990, -0.3720, -0.4925]],
		        # fmt: on
		    ]				)
		@require_torch_gpu
		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       Tuple				=  self.get_sd_vae_model()
						lowerCAmelCase__  :       Union[str, Any]				=  self.get_sd_image(__lowerCamelCase	,shape=(3, 4, 64, 64)				)

						with torch.no_grad():
										lowerCAmelCase__  :       str				=  model.decode(__lowerCamelCase				).sample

						assert list(sample.shape				) == [3, 3, 5_12, 5_12]

						lowerCAmelCase__  :       Dict				=  sample[-1, -2:, :2, -2:].flatten().cpu()
						lowerCAmelCase__  :       Any				=  torch.tensor(__lowerCamelCase				)

						assert torch_all_close(__lowerCamelCase	,__lowerCamelCase	,atol=1e-3				)



		@parameterized.expand(
		    [
		        # fmt: off
		        [27, [-0.0369, 0.0207, -0.0776, -0.0682, -0.1747, -0.1930, -0.1465, -0.2039]],
		        [16, [-0.1628, -0.2134, -0.2747, -0.2642, -0.3774, -0.4404, -0.3687, -0.4277]],
		        # fmt: on
		    ]				)
		@require_torch_gpu
		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  self.get_sd_vae_model(fpaa=__lowerCamelCase				)
						lowerCAmelCase__  :       Dict				=  self.get_sd_image(__lowerCamelCase	,shape=(3, 4, 64, 64)	,fpaa=__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       Optional[Any]				=  model.decode(__lowerCamelCase				).sample

						assert list(sample.shape				) == [3, 3, 5_12, 5_12]

						lowerCAmelCase__  :       Dict				=  sample[-1, -2:, :2, -2:].flatten().float().cpu()
						lowerCAmelCase__  :       Union[str, Any]				=  torch.tensor(__lowerCamelCase				)

						assert torch_all_close(__lowerCamelCase	,__lowerCamelCase	,atol=5e-3				)



		@parameterized.expand([(13,), (16,), (27,)]				)
		@require_torch_gpu
		@unittest.skipIf(not is_xformers_available()	,reason='''xformers is not required when using PyTorch 2.0.'''				)
		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  self.get_sd_vae_model(fpaa=__lowerCamelCase				)
						lowerCAmelCase__  :       str				=  self.get_sd_image(__lowerCamelCase	,shape=(3, 4, 64, 64)	,fpaa=__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       Any				=  model.decode(__lowerCamelCase				).sample

						model.enable_xformers_memory_efficient_attention()
						with torch.no_grad():
										lowerCAmelCase__  :       Optional[Any]				=  model.decode(__lowerCamelCase				).sample

						assert list(sample.shape				) == [3, 3, 5_12, 5_12]

						assert torch_all_close(__lowerCamelCase	,__lowerCamelCase	,atol=1e-1				)



		@parameterized.expand([(13,), (16,), (37,)]				)
		@require_torch_gpu
		@unittest.skipIf(not is_xformers_available()	,reason='''xformers is not required when using PyTorch 2.0.'''				)
		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.get_sd_vae_model()
						lowerCAmelCase__  :       Dict				=  self.get_sd_image(__lowerCamelCase	,shape=(3, 4, 64, 64)				)

						with torch.no_grad():
										lowerCAmelCase__  :       Tuple				=  model.decode(__lowerCamelCase				).sample

						model.enable_xformers_memory_efficient_attention()
						with torch.no_grad():
										lowerCAmelCase__  :       Optional[Any]				=  model.decode(__lowerCamelCase				).sample

						assert list(sample.shape				) == [3, 3, 5_12, 5_12]

						assert torch_all_close(__lowerCamelCase	,__lowerCamelCase	,atol=1e-2				)



		@parameterized.expand(
		    [
		        # fmt: off
		        [33, [-0.3001, 0.0918, -2.6984, -3.9720, -3.2099, -5.0353, 1.7338, -0.2065, 3.4267]],
		        [47, [-1.5030, -4.3871, -6.0355, -9.1157, -1.6661, -2.7853, 2.1607, -5.0823, 2.5633]],
		        # fmt: on
		    ]				)
		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Tuple				=  self.get_sd_vae_model()
						lowerCAmelCase__  :       Tuple				=  self.get_sd_image(__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[Any]				=  self.get_generator(__lowerCamelCase				)

						with torch.no_grad():
										lowerCAmelCase__  :       List[Any]				=  model.encode(__lowerCamelCase				).latent_dist
										lowerCAmelCase__  :       int				=  dist.sample(generator=__lowerCamelCase				)

						assert list(sample.shape				) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

						lowerCAmelCase__  :       Optional[int]				=  sample[0, -1, -3:, -3:].flatten().cpu()
						lowerCAmelCase__  :       int				=  torch.tensor(__lowerCamelCase				)

						lowerCAmelCase__  :       Any				=  3e-3 if torch_device != '''mps''' else 1e-2
						assert torch_all_close(__lowerCamelCase	,__lowerCamelCase	,atol=__lowerCamelCase				)

