






from dataclasses import dataclass
from typing import Dict, Optional, Union

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .attention_processor import AttentionProcessor, AttnProcessor
from .embeddings import TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
@dataclass
class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE__	):


			'''simple docstring'''





			__a							:							List[Any]					     =					42
class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__	):


			'''simple docstring'''





			@register_to_config
			def __init__(   self							:						Any ,      lowercase							:						Union[str, Any] = 3_2 ,      lowercase							:						Dict = 6_4 ,      lowercase							:						Optional[int] = 2_0 ,      lowercase							:						Dict = 7_6_8 ,      lowercase							:						List[str]=7_7 ,      lowercase							:						Dict=4 ,      lowercase							:						Optional[Any] = 0.0 ,      lowercase							:						List[Any] = "silu" ,      lowercase							:						Any = None ,      lowercase							:						str = None ,      lowercase							:						Optional[Any] = "linear" ,      lowercase							:						Union[str, Any] = "prd" ,      lowercase							:						Union[str, Any] = None ,      lowercase							:						List[Any] = None ,      lowercase							:						str = None ,      )							->    Dict:

									'''simple docstring'''







									super().__init__()
									UpperCamelCase__         = num_attention_heads
									UpperCamelCase__         = attention_head_dim
									UpperCamelCase__         = num_attention_heads * attention_head_dim
									UpperCamelCase__         = additional_embeddings

									UpperCamelCase__         = time_embed_dim or inner_dim
									UpperCamelCase__         = embedding_proj_dim or embedding_dim
									UpperCamelCase__         = clip_embed_dim or embedding_dim

									UpperCamelCase__         = Timesteps(_lowercase ,      _lowercase ,      0			)
									UpperCamelCase__         = TimestepEmbedding(_lowercase ,      _lowercase ,      out_dim=_lowercase ,      act_fn=_lowercase			)

									UpperCamelCase__         = nn.Linear(_lowercase ,      _lowercase			)

									if embedding_proj_norm_type is None:
															UpperCamelCase__         = None
									elif embedding_proj_norm_type == "layer":
															UpperCamelCase__         = nn.LayerNorm(_lowercase			)
									else:
															raise ValueError(f"unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"			)

									UpperCamelCase__         = nn.Linear(_lowercase ,      _lowercase			)

									if encoder_hid_proj_type is None:
															UpperCamelCase__         = None
									elif encoder_hid_proj_type == "linear":
															UpperCamelCase__         = nn.Linear(_lowercase ,      _lowercase			)
									else:
															raise ValueError(f"unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"			)

									UpperCamelCase__         = nn.Parameter(torch.zeros(1 ,      num_embeddings + additional_embeddings ,      _lowercase			)			)

									if added_emb_type == "prd":
															UpperCamelCase__         = nn.Parameter(torch.zeros(1 ,      1 ,      _lowercase			)			)
									elif added_emb_type is None:
															UpperCamelCase__         = None
									else:
															raise ValueError(
															    f"`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `\'prd\'` or `None`."			)

									UpperCamelCase__         = nn.ModuleList(
									    [
									        BasicTransformerBlock(
									            _lowercase ,      _lowercase ,      _lowercase ,      dropout=_lowercase ,      activation_fn="""gelu""" ,      attention_bias=_lowercase ,      )
									        for d in range(_lowercase			)
									    ]			)

									if norm_in_type == "layer":
															UpperCamelCase__         = nn.LayerNorm(_lowercase			)
									elif norm_in_type is None:
															UpperCamelCase__         = None
									else:
															raise ValueError(f"Unsupported norm_in_type: {norm_in_type}."			)

									UpperCamelCase__         = nn.LayerNorm(_lowercase			)

									UpperCamelCase__         = nn.Linear(_lowercase ,      _lowercase			)

									UpperCamelCase__         = torch.full(
									    [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings] ,      -1_0_0_0_0.0			)
									causal_attention_mask.triu_(1			)
									UpperCamelCase__         = causal_attention_mask[None, ...]
									self.register_buffer("""causal_attention_mask""" ,      _lowercase ,      persistent=_lowercase			)

									UpperCamelCase__         = nn.Parameter(torch.zeros(1 ,      _lowercase			)			)
									UpperCamelCase__         = nn.Parameter(torch.zeros(1 ,      _lowercase			)			)






			@property
			# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
			def 				A      (   self							:						Dict			)							->    Dict[str, AttentionProcessor]:

									'''simple docstring'''







									UpperCamelCase__         = {}

									def fn_recursive_add_processors(lowercase							:						Tuple ,      lowercase							:						Dict ,      lowercase							:						str			):
															if hasattr(_lowercase ,      """set_processor"""			):
																					UpperCamelCase__         = module.processor

															for sub_name, child in module.named_children():
																					fn_recursive_add_processors(f"{name}.{sub_name}" ,      _lowercase ,      _lowercase			)

															return processors

									for name, module in self.named_children():
															fn_recursive_add_processors(_lowercase ,      _lowercase ,      _lowercase			)

									return processors






			def 				A      (   self							:						Tuple ,      lowercase							:						Any			)							->    Union[str, Any]:

									'''simple docstring'''







									UpperCamelCase__         = len(self.attn_processors.keys()			)

									if isinstance(_lowercase ,      _lowercase			) and len(_lowercase			) != count:
															raise ValueError(
															    f"A dict of processors was passed, but the number of processors {len(_lowercase			)} does not match the"
															    f" number of attention layers: {count}. Please make sure to pass {count} processor classes."			)

									def fn_recursive_attn_processor(lowercase							:						int ,      lowercase							:						Optional[int] ,      lowercase							:						Union[str, Any]			):
															if hasattr(_lowercase ,      """set_processor"""			):
																					if not isinstance(_lowercase ,      _lowercase			):
																											module.set_processor(_lowercase			)
																					else:
																											module.set_processor(processor.pop(f"{name}.processor"			)			)

															for sub_name, child in module.named_children():
																					fn_recursive_attn_processor(f"{name}.{sub_name}" ,      _lowercase ,      _lowercase			)

									for name, module in self.named_children():
															fn_recursive_attn_processor(_lowercase ,      _lowercase ,      _lowercase			)






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

									'''simple docstring'''







									self.set_attn_processor(AttnProcessor()			)






			def 				A      (   self							:						Optional[Any] ,      lowercase							:						Optional[Any] ,      lowercase							:						List[Any] ,      lowercase							:						List[Any] ,      lowercase							:						Dict = None ,      lowercase							:						int = None ,      lowercase							:						Optional[int] = True ,      )							->    Any:

									'''simple docstring'''







									UpperCamelCase__         = hidden_states.shape[0]

									UpperCamelCase__         = timestep
									if not torch.is_tensor(_lowercase			):
															UpperCamelCase__         = torch.tensor([timesteps] ,      dtype=torch.long ,      device=hidden_states.device			)
									elif torch.is_tensor(_lowercase			) and len(timesteps.shape			) == 0:
															UpperCamelCase__         = timesteps[None].to(hidden_states.device			)

									# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
									UpperCamelCase__         = timesteps * torch.ones(_lowercase ,      dtype=timesteps.dtype ,      device=timesteps.device			)

									UpperCamelCase__         = self.time_proj(_lowercase			)

									# timesteps does not contain any weights and will always return f32 tensors
									# but time_embedding might be fp16, so we need to cast here.
									UpperCamelCase__         = timesteps_projected.to(dtype=self.dtype			)
									UpperCamelCase__         = self.time_embedding(_lowercase			)

									if self.embedding_proj_norm is not None:
															UpperCamelCase__         = self.embedding_proj_norm(_lowercase			)

									UpperCamelCase__         = self.embedding_proj(_lowercase			)
									if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
															UpperCamelCase__         = self.encoder_hidden_states_proj(_lowercase			)
									elif self.encoder_hidden_states_proj is not None and encoder_hidden_states is None:
															raise ValueError("""`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set"""			)

									UpperCamelCase__         = self.proj_in(_lowercase			)

									UpperCamelCase__         = self.positional_embedding.to(hidden_states.dtype			)

									UpperCamelCase__         = []
									UpperCamelCase__         = 0

									if encoder_hidden_states is not None:
															additional_embeds.append(_lowercase			)
															additional_embeddings_len += encoder_hidden_states.shape[1]

									if len(proj_embeddings.shape			) == 2:
															UpperCamelCase__         = proj_embeddings[:, None, :]

									if len(hidden_states.shape			) == 2:
															UpperCamelCase__         = hidden_states[:, None, :]

									UpperCamelCase__         = additional_embeds + [
									    proj_embeddings,
									    time_embeddings[:, None, :],
									    hidden_states,
									]

									if self.prd_embedding is not None:
															UpperCamelCase__         = self.prd_embedding.to(hidden_states.dtype			).expand(_lowercase ,      -1 ,      -1			)
															additional_embeds.append(_lowercase			)

									UpperCamelCase__         = torch.cat(
									    _lowercase ,      dim=1 ,      )

									# Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
									UpperCamelCase__         = additional_embeddings_len + proj_embeddings.shape[1] + 1
									if positional_embeddings.shape[1] < hidden_states.shape[1]:
															UpperCamelCase__         = F.pad(
															    _lowercase ,      (
															        0,
															        0,
															        additional_embeddings_len,
															        self.prd_embedding.shape[1] if self.prd_embedding is not None else 0,
															    ) ,      value=0.0 ,      )

									UpperCamelCase__         = hidden_states + positional_embeddings

									if attention_mask is not None:
															UpperCamelCase__         = (1 - attention_mask.to(hidden_states.dtype			)) * -1_0_0_0_0.0
															UpperCamelCase__         = F.pad(_lowercase ,      (0, self.additional_embeddings) ,      value=0.0			)
															UpperCamelCase__         = (attention_mask[:, None, :] + self.causal_attention_mask).to(hidden_states.dtype			)
															UpperCamelCase__         = attention_mask.repeat_interleave(self.config.num_attention_heads ,      dim=0			)

									if self.norm_in is not None:
															UpperCamelCase__         = self.norm_in(_lowercase			)

									for block in self.transformer_blocks:
															UpperCamelCase__         = block(_lowercase ,      attention_mask=_lowercase			)

									UpperCamelCase__         = self.norm_out(_lowercase			)

									if self.prd_embedding is not None:
															UpperCamelCase__         = hidden_states[:, -1]
									else:
															UpperCamelCase__         = hidden_states[:, additional_embeddings_len:]

									UpperCamelCase__         = self.proj_to_clip_embeddings(_lowercase			)

									if not return_dict:
															return (predicted_image_embedding,)

									return PriorTransformerOutput(predicted_image_embedding=_lowercase			)







			def 				A      (   self							:						Tuple ,      lowercase							:						Union[str, Any]			)							->    Union[str, Any]:

									'''simple docstring'''







									UpperCamelCase__         = (prior_latents * self.clip_std) + self.clip_mean
									return prior_latents








'''simple docstring'''


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    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):


			'''simple docstring'''





			def 				A      (   self							:						Any			)							->    Any:

									'''simple docstring'''







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






			@property
			def 				A      (   self							:						Optional[Any]			)							->    Union[str, Any]:

									'''simple docstring'''







									torch.manual_seed(0			)
									UpperCamelCase__         = UNetaDModel(
									    sample_size=(3_2, 6_4) ,      in_channels=1 ,      out_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""AttnDownBlock2D""", """DownBlock2D""") ,      up_block_types=("""UpBlock2D""", """AttnUpBlock2D""") ,      )
									return model






			@property
			def 				A      (   self							:						List[Any]			)							->    Union[str, Any]:

									'''simple docstring'''







									torch.manual_seed(0			)
									UpperCamelCase__         = UNetaDConditionModel(
									    sample_size=(6_4, 3_2) ,      in_channels=1 ,      out_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""") ,      up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""") ,      cross_attention_dim=1_0 ,      )
									return model






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

									'''simple docstring'''







									torch.manual_seed(0			)
									UpperCamelCase__         = AutoencoderKL(
									    sample_size=(1_2_8, 6_4) ,      in_channels=1 ,      out_channels=1 ,      latent_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""") ,      up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""") ,      )
									UpperCamelCase__         = UNetaDModel(
									    sample_size=(6_4, 3_2) ,      in_channels=1 ,      out_channels=1 ,      layers_per_block=2 ,      block_out_channels=(1_2_8, 1_2_8) ,      down_block_types=("""AttnDownBlock2D""", """DownBlock2D""") ,      up_block_types=("""UpBlock2D""", """AttnUpBlock2D""") ,      )
									return vqvae, unet







			@slow
			def 				A      (   self							:						Tuple			)							->    List[str]:

									'''simple docstring'''







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

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

									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(generator=lowercase ,      steps=4			)
									UpperCamelCase__         = output.audios[0]
									UpperCamelCase__         = output.images[0]

									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(generator=lowercase ,      steps=4 ,      return_dict=lowercase			)
									UpperCamelCase__         = 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]
									)
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.frombuffer(image_from_tuple.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([6_9, 2_5_5, 2_5_5, 2_5_5, 0, 0, 7_7, 1_8_1, 1_2, 1_2_7]			)

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

									UpperCamelCase__         = 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] ,      )

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

									np.random.seed(0			)
									UpperCamelCase__         = np.random.uniform(-1 ,      1 ,      ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)			)
									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(raw_audio=lowercase ,      generator=lowercase ,      start_step=5 ,      steps=1_0			)
									UpperCamelCase__         = 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]
									)
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([1_2_0, 1_1_7, 1_1_0, 1_0_9, 1_3_8, 1_6_7, 1_3_8, 1_4_8, 1_3_2, 1_2_1]			)

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

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

									np.random.seed(0			)
									UpperCamelCase__         = torch.rand((1, 1, 1_0)			)
									UpperCamelCase__         = pipe(generator=lowercase ,      encoding=lowercase			)
									UpperCamelCase__         = output.images[0]
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([1_0_7, 1_0_3, 1_2_0, 1_2_7, 1_4_2, 1_2_2, 1_1_3, 1_2_2, 9_7, 1_1_1]			)

									assert np.abs(image_slice.flatten() - expected_slice			).max() == 0
@slow
@require_torch_gpu
class    _SCREAMING_SNAKE_CASE							(		unittest.TestCase	):


			'''simple docstring'''





			def 				A      (   self							:						int			)							->    str:

									'''simple docstring'''







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







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

									'''simple docstring'''







									UpperCamelCase__         = torch_device

									UpperCamelCase__         = DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""			)
									UpperCamelCase__         = pipe.to(lowercase			)
									pipe.set_progress_bar_config(disable=lowercase			)

									UpperCamelCase__         = torch.Generator(device=lowercase			).manual_seed(4_2			)
									UpperCamelCase__         = pipe(generator=lowercase			)
									UpperCamelCase__         = output.audios[0]
									UpperCamelCase__         = 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]
									UpperCamelCase__         = np.frombuffer(image.tobytes() ,      dtype="""uint8"""			)[:1_0]
									UpperCamelCase__         = np.array([1_5_1, 1_6_7, 1_5_4, 1_4_4, 1_2_2, 1_3_4, 1_2_1, 1_0_5, 7_0, 2_6]			)

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