

import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class 				__UpperCamelCase				(					unittest.TestCase			):







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





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







	@property
	def 	_SCREAMING_SNAKE_CASE							(   self: int      ):

						'''simple docstring'''





						torch.manual_seed(0      )
						__magic_name__							=    UNetaDModel(
						    sample_size=(32, 64)  , in_channels=1  , out_channels=1  , layers_per_block=2  , block_out_channels=(1_28, 1_28)  , down_block_types=('AttnDownBlock2D', 'DownBlock2D')  , up_block_types=('UpBlock2D', 'AttnUpBlock2D')  , )
						return model







	@property
	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						torch.manual_seed(0      )
						__magic_name__							=    UNetaDConditionModel(
						    sample_size=(64, 32)  , in_channels=1  , out_channels=1  , layers_per_block=2  , block_out_channels=(1_28, 1_28)  , down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')  , up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')  , cross_attention_dim=10  , )
						return model







	@property
	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]      ):

						'''simple docstring'''





						torch.manual_seed(0      )
						__magic_name__							=    AutoencoderKL(
						    sample_size=(1_28, 64)  , in_channels=1  , out_channels=1  , latent_channels=1  , layers_per_block=2  , block_out_channels=(1_28, 1_28)  , down_block_types=('DownEncoderBlock2D', 'DownEncoderBlock2D')  , up_block_types=('UpDecoderBlock2D', 'UpDecoderBlock2D')  , )
						__magic_name__							=    UNetaDModel(
						    sample_size=(64, 32)  , in_channels=1  , out_channels=1  , layers_per_block=2  , block_out_channels=(1_28, 1_28)  , down_block_types=('AttnDownBlock2D', 'DownBlock2D')  , up_block_types=('UpBlock2D', 'AttnUpBlock2D')  , )
						return vqvae, unet







	@slow
	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						__magic_name__							=    'cpu'  # ensure determinism for the device-dependent torch.Generator
						__magic_name__							=    Mel(
						    x_res=self.dummy_unet.config.sample_size[1]  , y_res=self.dummy_unet.config.sample_size[0]  , )

						__magic_name__							=    DDPMScheduler()
						__magic_name__							=    AudioDiffusionPipeline(vqvae=__UpperCamelCase  , unet=self.dummy_unet  , mel=__UpperCamelCase  , scheduler=__UpperCamelCase      )
						__magic_name__							=    pipe.to(__UpperCamelCase      )
						pipe.set_progress_bar_config(disable=__UpperCamelCase      )

						__magic_name__							=    torch.Generator(device=__UpperCamelCase      ).manual_seed(42      )
						__magic_name__							=    pipe(generator=__UpperCamelCase  , steps=4      )
						__magic_name__							=    output.audios[0]
						__magic_name__							=    output.images[0]

						__magic_name__							=    torch.Generator(device=__UpperCamelCase      ).manual_seed(42      )
						__magic_name__							=    pipe(generator=__UpperCamelCase  , steps=4  , return_dict=__UpperCamelCase      )
						__magic_name__							=    output[0][0]

						assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
						assert (
						    image.height == self.dummy_unet.config.sample_size[0]
						    and image.width == self.dummy_unet.config.sample_size[1]
						)
						__magic_name__							=    np.frombuffer(image.tobytes()  , dtype='uint8'      )[:10]
						__magic_name__							=    np.frombuffer(image_from_tuple.tobytes()  , dtype='uint8'      )[:10]
						__magic_name__							=    np.array([69, 2_55, 2_55, 2_55, 0, 0, 77, 1_81, 12, 1_27]      )

						assert np.abs(image_slice.flatten() - expected_slice      ).max() == 0
						assert np.abs(image_from_tuple_slice.flatten() - expected_slice      ).max() == 0

						__magic_name__							=    Mel(
						    x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]  , y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]  , )

						__magic_name__							=    DDIMScheduler()
						__magic_name__							=    self.dummy_vqvae_and_unet
						__magic_name__							=    AudioDiffusionPipeline(
						    vqvae=self.dummy_vqvae_and_unet[0]  , unet=dummy_vqvae_and_unet[1]  , mel=__UpperCamelCase  , scheduler=__UpperCamelCase      )
						__magic_name__							=    pipe.to(__UpperCamelCase      )
						pipe.set_progress_bar_config(disable=__UpperCamelCase      )

						np.random.seed(0      )
						__magic_name__							=    np.random.uniform(-1  , 1  , ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)      )
						__magic_name__							=    torch.Generator(device=__UpperCamelCase      ).manual_seed(42      )
						__magic_name__							=    pipe(raw_audio=__UpperCamelCase  , generator=__UpperCamelCase  , start_step=5  , steps=10      )
						__magic_name__							=    output.images[0]

						assert (
						    image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
						    and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
						)
						__magic_name__							=    np.frombuffer(image.tobytes()  , dtype='uint8'      )[:10]
						__magic_name__							=    np.array([1_20, 1_17, 1_10, 1_09, 1_38, 1_67, 1_38, 1_48, 1_32, 1_21]      )

						assert np.abs(image_slice.flatten() - expected_slice      ).max() == 0

						__magic_name__							=    self.dummy_unet_condition
						__magic_name__							=    AudioDiffusionPipeline(
						    vqvae=self.dummy_vqvae_and_unet[0]  , unet=__UpperCamelCase  , mel=__UpperCamelCase  , scheduler=__UpperCamelCase      )
						__magic_name__							=    pipe.to(__UpperCamelCase      )
						pipe.set_progress_bar_config(disable=__UpperCamelCase      )

						np.random.seed(0      )
						__magic_name__							=    torch.rand((1, 1, 10)      )
						__magic_name__							=    pipe(generator=__UpperCamelCase  , encoding=__UpperCamelCase      )
						__magic_name__							=    output.images[0]
						__magic_name__							=    np.frombuffer(image.tobytes()  , dtype='uint8'      )[:10]
						__magic_name__							=    np.array([1_07, 1_03, 1_20, 1_27, 1_42, 1_22, 1_13, 1_22, 97, 1_11]      )

						assert np.abs(image_slice.flatten() - expected_slice      ).max() == 0




@slow
@require_torch_gpu
class 				__UpperCamelCase				(					unittest.TestCase			):







	def 	_SCREAMING_SNAKE_CASE							(   self: str      ):

						'''simple docstring'''





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







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						__magic_name__							=    torch_device

						__magic_name__							=    DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256'      )
						__magic_name__							=    pipe.to(__UpperCamelCase      )
						pipe.set_progress_bar_config(disable=__UpperCamelCase      )

						__magic_name__							=    torch.Generator(device=__UpperCamelCase      ).manual_seed(42      )
						__magic_name__							=    pipe(generator=__UpperCamelCase      )
						__magic_name__							=    output.audios[0]
						__magic_name__							=    output.images[0]

						assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
						assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
						__magic_name__							=    np.frombuffer(image.tobytes()  , dtype='uint8'      )[:10]
						__magic_name__							=    np.array([1_51, 1_67, 1_54, 1_44, 1_22, 1_34, 1_21, 1_05, 70, 26]      )

						assert np.abs(image_slice.flatten() - expected_slice      ).max() == 0







import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin

class 				__UpperCamelCase				(					SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE			):







	@register_to_config
	def __init__(   self: Optional[Any]  , *,
	    __UpperCamelCase: int = 4  , __UpperCamelCase: int = 7_68  , __UpperCamelCase: int  , __UpperCamelCase: Any  , ):

						'''simple docstring'''





						super().__init__()

						__magic_name__							=    nn.Parameter(torch.zeros(__UpperCamelCase      )      )

						# parameters for additional clip time embeddings
						__magic_name__							=    nn.Linear(__UpperCamelCase  , __UpperCamelCase      )
						__magic_name__							=    nn.Linear(__UpperCamelCase  , __UpperCamelCase      )

						# parameters for encoder hidden states
						__magic_name__							=    clip_extra_context_tokens
						__magic_name__							=    nn.Linear(
						    __UpperCamelCase  , self.clip_extra_context_tokens * cross_attention_dim      )
						__magic_name__							=    nn.Linear(__UpperCamelCase  , __UpperCamelCase      )
						__magic_name__							=    nn.LayerNorm(__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[str]  , *, __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: str  , __UpperCamelCase: Optional[Any]  , __UpperCamelCase: Any      ):

						'''simple docstring'''





						if do_classifier_free_guidance:
											# Add the classifier free guidance embeddings to the image embeddings
											__magic_name__							=    image_embeddings.shape[0]
											__magic_name__							=    self.learned_classifier_free_guidance_embeddings.unsqueeze(0      )
											__magic_name__							=    classifier_free_guidance_embeddings.expand(
											    __UpperCamelCase  , -1      )
											__magic_name__							=    torch.cat([classifier_free_guidance_embeddings, image_embeddings]  , dim=0      )

						# The image embeddings batch size and the text embeddings batch size are equal
						assert image_embeddings.shape[0] == prompt_embeds.shape[0]

						__magic_name__							=    prompt_embeds.shape[0]

						# "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
						# adding CLIP embeddings to the existing timestep embedding, ...
						__magic_name__							=    self.embedding_proj(__UpperCamelCase      )
						__magic_name__							=    self.clip_image_embeddings_project_to_time_embeddings(__UpperCamelCase      )
						__magic_name__							=    time_projected_image_embeddings + time_projected_prompt_embeds

						# ... and by projecting CLIP embeddings into four
						# extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
						__magic_name__							=    self.clip_extra_context_tokens_proj(__UpperCamelCase      )
						__magic_name__							=    clip_extra_context_tokens.reshape(__UpperCamelCase  , -1  , self.clip_extra_context_tokens      )
						__magic_name__							=    clip_extra_context_tokens.permute(0  , 2  , 1      )

						__magic_name__							=    self.encoder_hidden_states_proj(__UpperCamelCase      )
						__magic_name__							=    self.text_encoder_hidden_states_norm(__UpperCamelCase      )
						__magic_name__							=    torch.cat([clip_extra_context_tokens, text_encoder_hidden_states]  , dim=1      )

						return text_encoder_hidden_states, additive_clip_time_embeddings




