






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 				UpperCAmelCase_  (		__lowerCamelCase		):
     __lowerCamelCase      				=  42







class 				UpperCAmelCase_  (		__lowerCamelCase							,      __lowerCamelCase		):




     @register_to_config
     def __init__(				self	,						_lowerCAmelCase = 32	,						_lowerCAmelCase = 64	,						_lowerCAmelCase = 20	,						_lowerCAmelCase = 768	,						_lowerCAmelCase=77	,						_lowerCAmelCase=4	,						_lowerCAmelCase = 0.0	,						_lowerCAmelCase = "silu"	,						_lowerCAmelCase = None	,						_lowerCAmelCase = None	,						_lowerCAmelCase = "linear"	,						_lowerCAmelCase = "prd"	,						_lowerCAmelCase = None	,						_lowerCAmelCase = None	,						_lowerCAmelCase = None	,						):
          super().__init__()
          UpperCAmelCase__      :  Dict				      =			num_attention_heads
          UpperCAmelCase__      :  Optional[Any]				      =			attention_head_dim
          UpperCAmelCase__      :  List[str]				      =			num_attention_heads * attention_head_dim
          UpperCAmelCase__      :  List[str]				      =			additional_embeddings

          UpperCAmelCase__      :  Tuple				      =			time_embed_dim or inner_dim
          UpperCAmelCase__      :  int				      =			embedding_proj_dim or embedding_dim
          UpperCAmelCase__      :  Tuple				      =			clip_embed_dim or embedding_dim

          UpperCAmelCase__      :  Tuple				      =			Timesteps(_lowerCAmelCase	,						_lowerCAmelCase	,						0		)
          UpperCAmelCase__      :  int				      =			TimestepEmbedding(_lowerCAmelCase	,						_lowerCAmelCase	,						out_dim=_lowerCAmelCase	,						act_fn=_lowerCAmelCase		)

          UpperCAmelCase__      :  List[str]				      =			nn.Linear(_lowerCAmelCase	,						_lowerCAmelCase		)

          if embedding_proj_norm_type is None:
               UpperCAmelCase__      :  Optional[int]				      =			None
          elif embedding_proj_norm_type == "layer":
               UpperCAmelCase__      :  str				      =			nn.LayerNorm(_lowerCAmelCase		)
          else:
               raise ValueError(f"unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"		)

          UpperCAmelCase__      :  List[Any]				      =			nn.Linear(_lowerCAmelCase	,						_lowerCAmelCase		)

          if encoder_hid_proj_type is None:
               UpperCAmelCase__      :  str				      =			None
          elif encoder_hid_proj_type == "linear":
               UpperCAmelCase__      :  List[Any]				      =			nn.Linear(_lowerCAmelCase	,						_lowerCAmelCase		)
          else:
               raise ValueError(f"unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"		)

          UpperCAmelCase__      :  Dict				      =			nn.Parameter(torch.zeros(1	,						num_embeddings + additional_embeddings	,						_lowerCAmelCase		)		)

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

          UpperCAmelCase__      :  Dict				      =			nn.ModuleList(
              [
                  BasicTransformerBlock(
                      _lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase	,						dropout=_lowerCAmelCase	,						activation_fn="""gelu"""	,						attention_bias=_lowerCAmelCase	,						)
                  for d in range(_lowerCAmelCase		)
              ]		)

          if norm_in_type == "layer":
               UpperCAmelCase__      :  Optional[Any]				      =			nn.LayerNorm(_lowerCAmelCase		)
          elif norm_in_type is None:
               UpperCAmelCase__      :  Union[str, Any]				      =			None
          else:
               raise ValueError(f"Unsupported norm_in_type: {norm_in_type}."		)

          UpperCAmelCase__      :  Optional[int]				      =			nn.LayerNorm(_lowerCAmelCase		)

          UpperCAmelCase__      :  Union[str, Any]				      =			nn.Linear(_lowerCAmelCase	,						_lowerCAmelCase		)

          UpperCAmelCase__      :  str				      =			torch.full(
              [num_embeddings + additional_embeddings, num_embeddings + additional_embeddings]	,						-1_0_0_0_0.0		)
          causal_attention_mask.triu_(1		)
          UpperCAmelCase__      :  List[str]				      =			causal_attention_mask[None, ...]
          self.register_buffer("""causal_attention_mask"""	,						_lowerCAmelCase	,						persistent=_lowerCAmelCase		)

          UpperCAmelCase__      :  Tuple				      =			nn.Parameter(torch.zeros(1	,						_lowerCAmelCase		)		)
          UpperCAmelCase__      :  List[Any]				      =			nn.Parameter(torch.zeros(1	,						_lowerCAmelCase		)		)




     @property
     # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  List[Any]				      =			{}

          def fn_recursive_add_processors(_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
               if hasattr(_lowerCAmelCase	,						"""set_processor"""		):
                    UpperCAmelCase__      :  int				      =			module.processor

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

               return processors

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

          return processors




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  Union[str, Any]				      =			len(self.attn_processors.keys()		)

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

          def fn_recursive_attn_processor(_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
               if hasattr(_lowerCAmelCase	,						"""set_processor"""		):
                    if not isinstance(_lowerCAmelCase	,						_lowerCAmelCase		):
                         module.set_processor(_lowerCAmelCase		)
                    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}"	,						_lowerCAmelCase	,						_lowerCAmelCase		)

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




     def   __UpperCAmelCase  (				self		):
          self.set_attn_processor(AttnProcessor()		)




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase = None	,						_lowerCAmelCase = None	,						_lowerCAmelCase = True	,						):
          UpperCAmelCase__      :  List[str]				      =			hidden_states.shape[0]

          UpperCAmelCase__      :  List[str]				      =			timestep
          if not torch.is_tensor(_lowerCAmelCase		):
               UpperCAmelCase__      :  Tuple				      =			torch.tensor([timesteps]	,						dtype=torch.long	,						device=hidden_states.device		)
          elif torch.is_tensor(_lowerCAmelCase		) and len(timesteps.shape		) == 0:
               UpperCAmelCase__      :  List[str]				      =			timesteps[None].to(hidden_states.device		)

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

          UpperCAmelCase__      :  Tuple				      =			self.time_proj(_lowerCAmelCase		)

          # 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__      :  Dict				      =			timesteps_projected.to(dtype=self.dtype		)
          UpperCAmelCase__      :  Dict				      =			self.time_embedding(_lowerCAmelCase		)

          if self.embedding_proj_norm is not None:
               UpperCAmelCase__      :  List[str]				      =			self.embedding_proj_norm(_lowerCAmelCase		)

          UpperCAmelCase__      :  List[str]				      =			self.embedding_proj(_lowerCAmelCase		)
          if self.encoder_hidden_states_proj is not None and encoder_hidden_states is not None:
               UpperCAmelCase__      :  str				      =			self.encoder_hidden_states_proj(_lowerCAmelCase		)
          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__      :  List[Any]				      =			self.proj_in(_lowerCAmelCase		)

          UpperCAmelCase__      :  Dict				      =			self.positional_embedding.to(hidden_states.dtype		)

          UpperCAmelCase__      :  Any				      =			[]
          UpperCAmelCase__      :  Any				      =			0

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

          if len(proj_embeddings.shape		) == 2:
               UpperCAmelCase__      :  List[str]				      =			proj_embeddings[:, None, :]

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

          UpperCAmelCase__      :  Tuple				      =			additional_embeds + [
              proj_embeddings,
              time_embeddings[:, None, :],
              hidden_states,
          ]

          if self.prd_embedding is not None:
               UpperCAmelCase__      :  List[str]				      =			self.prd_embedding.to(hidden_states.dtype		).expand(_lowerCAmelCase	,						-1	,						-1		)
               additional_embeds.append(_lowerCAmelCase		)

          UpperCAmelCase__      :  Dict				      =			torch.cat(
              _lowerCAmelCase	,						dim=1	,						)

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

          UpperCAmelCase__      :  List[Any]				      =			hidden_states + positional_embeddings

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

          if self.norm_in is not None:
               UpperCAmelCase__      :  int				      =			self.norm_in(_lowerCAmelCase		)

          for block in self.transformer_blocks:
               UpperCAmelCase__      :  List[Any]				      =			block(_lowerCAmelCase	,						attention_mask=_lowerCAmelCase		)

          UpperCAmelCase__      :  int				      =			self.norm_out(_lowerCAmelCase		)

          if self.prd_embedding is not None:
               UpperCAmelCase__      :  Dict				      =			hidden_states[:, -1]
          else:
               UpperCAmelCase__      :  Optional[Any]				      =			hidden_states[:, additional_embeddings_len:]

          UpperCAmelCase__      :  str				      =			self.proj_to_clip_embeddings(_lowerCAmelCase		)

          if not return_dict:
               return (predicted_image_embedding,)

          return PriorTransformerOutput(predicted_image_embedding=_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  int				      =			(prior_latents * self.clip_std) + self.clip_mean
          return prior_latents









import unittest

import numpy as np
import torch

from diffusers import KarrasVePipeline, KarrasVeScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()






class _UpperCamelCase				(				unittest.TestCase ):
			'''simple docstring'''







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

				"""simple docstring"""

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   str					  =							UNetaDModel(
				    block_out_channels=(32, 64)   ,   layers_per_block=2   ,   sample_size=32   ,   in_channels=3   ,   out_channels=3   ,   down_block_types=("DownBlock2D", "AttnDownBlock2D")   ,   up_block_types=("AttnUpBlock2D", "UpBlock2D")   ,   )
				return model






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

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							self.dummy_uncond_unet
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							KarrasVeScheduler()

				SCREAMING_SNAKE_CASE :   Any					  =							KarrasVePipeline(unet=a   ,   scheduler=a						)
				pipe.to(a						)
				pipe.set_progress_bar_config(disable=a						)

				SCREAMING_SNAKE_CASE :   Any					  =							torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							pipe(num_inference_steps=2   ,   generator=a   ,   output_type="numpy"						).images

				SCREAMING_SNAKE_CASE :   List[str]					  =							torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   List[Any]					  =							pipe(num_inference_steps=2   ,   generator=a   ,   output_type="numpy"   ,   return_dict=a						)[0]

				SCREAMING_SNAKE_CASE :   List[Any]					  =							image[0, -3:, -3:, -1]
				SCREAMING_SNAKE_CASE :   Any					  =							image_from_tuple[0, -3:, -3:, -1]

				assert image.shape == (1, 32, 32, 3)
				SCREAMING_SNAKE_CASE :   str					  =							np.array([0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]						)

				assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2
				assert np.abs(image_from_tuple_slice.flatten() - expected_slice						).max() < 1e-2






@slow
@require_torch
class _UpperCamelCase				(				unittest.TestCase ):
			'''simple docstring'''







			def 				__UpperCamelCase			(      self				:			int						)				->							Tuple:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Dict					  =							"google/ncsnpp-celebahq-256"
				SCREAMING_SNAKE_CASE :   List[Any]					  =							UNetaDModel.from_pretrained(a						)
				SCREAMING_SNAKE_CASE :   Any					  =							KarrasVeScheduler()

				SCREAMING_SNAKE_CASE :   Optional[Any]					  =							KarrasVePipeline(unet=a   ,   scheduler=a						)
				pipe.to(a						)
				pipe.set_progress_bar_config(disable=a						)

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Any					  =							pipe(num_inference_steps=20   ,   generator=a   ,   output_type="numpy"						).images

				SCREAMING_SNAKE_CASE :   List[str]					  =							image[0, -3:, -3:, -1]
				assert image.shape == (1, 256, 256, 3)
				SCREAMING_SNAKE_CASE :   str					  =							np.array([0.578, 0.5811, 0.5924, 0.5809, 0.587, 0.5886, 0.5861, 0.5802, 0.586]						)

				assert np.abs(image_slice.flatten() - expected_slice						).max() < 1e-2