



"""simple docstring"""

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

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

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin
@dataclass
class 						lowerCAmelCase__	(  _lowerCamelCase   ):


	'''simple docstring'''







	_SCREAMING_SNAKE_CASE  :		torch.FloatTensor




class 						lowerCAmelCase__	(  _lowerCamelCase       ,       _lowerCamelCase   ):


	'''simple docstring'''







	@register_to_config
	def __init__(			self							:					List[Any]    ,					_SCREAMING_SNAKE_CASE							:					int = 16    ,					_SCREAMING_SNAKE_CASE							:					int = 88    ,					_SCREAMING_SNAKE_CASE							:					Optional[int] = None    ,					_SCREAMING_SNAKE_CASE							:					Optional[int] = None    ,					_SCREAMING_SNAKE_CASE							:					int = 1    ,					_SCREAMING_SNAKE_CASE							:					float = 0.0    ,					_SCREAMING_SNAKE_CASE							:					int = 32    ,					_SCREAMING_SNAKE_CASE							:					Optional[int] = None    ,					_SCREAMING_SNAKE_CASE							:					bool = False    ,					_SCREAMING_SNAKE_CASE							:					Optional[int] = None    ,					_SCREAMING_SNAKE_CASE							:					Optional[int] = None    ,					_SCREAMING_SNAKE_CASE							:					Optional[int] = None    ,					_SCREAMING_SNAKE_CASE							:					str = "geglu"    ,					_SCREAMING_SNAKE_CASE							:					Optional[int] = None    ,					_SCREAMING_SNAKE_CASE							:					bool = False    ,					_SCREAMING_SNAKE_CASE							:					bool = False    ,					_SCREAMING_SNAKE_CASE							:					bool = False    ,					_SCREAMING_SNAKE_CASE							:					str = "layer_norm"    ,					_SCREAMING_SNAKE_CASE							:					bool = True    ,					)	->  Optional[Any]:





					"""simple docstring"""


					super().__init__()
					SCREAMING_SNAKE_CASE	:       Optional[int] 							=    use_linear_projection
					SCREAMING_SNAKE_CASE	:       Optional[int] 							=    num_attention_heads
					SCREAMING_SNAKE_CASE	:       Optional[Any] 							=    attention_head_dim
					SCREAMING_SNAKE_CASE	:       Dict 							=    num_attention_heads * attention_head_dim

					# 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
					# Define whether input is continuous or discrete depending on configuration
					SCREAMING_SNAKE_CASE	:       List[str] 							=    (in_channels is not None) and (patch_size is None)
					SCREAMING_SNAKE_CASE	:       str 							=    num_vector_embeds is not None
					SCREAMING_SNAKE_CASE	:       Dict 							=    in_channels is not None and patch_size is not None

					if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
									SCREAMING_SNAKE_CASE	:       Optional[int] 							=    (
									    f"""The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"""
									    ' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'
									    ' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'
									    ' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'
									    ' would be very nice if you could open a Pull request for the `transformer/config.json` file'
									)
									deprecate('norm_type!=num_embeds_ada_norm'    ,					'1.0.0'    ,					_SCREAMING_SNAKE_CASE    ,					standard_warn=_SCREAMING_SNAKE_CASE						)
									SCREAMING_SNAKE_CASE	:       Union[str, Any] 							=    'ada_norm'

					if self.is_input_continuous and self.is_input_vectorized:
									raise ValueError(
									    f"""Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"""
									    ' sure that either `in_channels` or `num_vector_embeds` is None.'						)
					elif self.is_input_vectorized and self.is_input_patches:
									raise ValueError(
									    f"""Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"""
									    ' sure that either `num_vector_embeds` or `num_patches` is None.'						)
					elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
									raise ValueError(
									    f"""Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"""
									    f""" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."""						)

					# 2. Define input layers
					if self.is_input_continuous:
									SCREAMING_SNAKE_CASE	:       Union[str, Any] 							=    in_channels

									SCREAMING_SNAKE_CASE	:       Dict 							=    torch.nn.GroupNorm(num_groups=_SCREAMING_SNAKE_CASE    ,					num_channels=_SCREAMING_SNAKE_CASE    ,					eps=1E-6    ,					affine=_SCREAMING_SNAKE_CASE						)
									if use_linear_projection:
													SCREAMING_SNAKE_CASE	:       List[Any] 							=    nn.Linear(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE						)
									else:
													SCREAMING_SNAKE_CASE	:       Dict 							=    nn.Convad(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					kernel_size=1    ,					stride=1    ,					padding=0						)
					elif self.is_input_vectorized:
									assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
									assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

									SCREAMING_SNAKE_CASE	:       Dict 							=    sample_size
									SCREAMING_SNAKE_CASE	:       Dict 							=    sample_size
									SCREAMING_SNAKE_CASE	:       Optional[Any] 							=    num_vector_embeds
									SCREAMING_SNAKE_CASE	:       Any 							=    self.height * self.width

									SCREAMING_SNAKE_CASE	:       int 							=    ImagePositionalEmbeddings(
									    num_embed=_SCREAMING_SNAKE_CASE    ,					embed_dim=_SCREAMING_SNAKE_CASE    ,					height=self.height    ,					width=self.width						)
					elif self.is_input_patches:
									assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

									SCREAMING_SNAKE_CASE	:       Tuple 							=    sample_size
									SCREAMING_SNAKE_CASE	:       List[str] 							=    sample_size

									SCREAMING_SNAKE_CASE	:       int 							=    patch_size
									SCREAMING_SNAKE_CASE	:       Any 							=    PatchEmbed(
									    height=_SCREAMING_SNAKE_CASE    ,					width=_SCREAMING_SNAKE_CASE    ,					patch_size=_SCREAMING_SNAKE_CASE    ,					in_channels=_SCREAMING_SNAKE_CASE    ,					embed_dim=_SCREAMING_SNAKE_CASE    ,					)

					# 3. Define transformers blocks
					SCREAMING_SNAKE_CASE	:       int 							=    nn.ModuleList(
					    [
					        BasicTransformerBlock(
					            _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					dropout=_SCREAMING_SNAKE_CASE    ,					cross_attention_dim=_SCREAMING_SNAKE_CASE    ,					activation_fn=_SCREAMING_SNAKE_CASE    ,					num_embeds_ada_norm=_SCREAMING_SNAKE_CASE    ,					attention_bias=_SCREAMING_SNAKE_CASE    ,					only_cross_attention=_SCREAMING_SNAKE_CASE    ,					upcast_attention=_SCREAMING_SNAKE_CASE    ,					norm_type=_SCREAMING_SNAKE_CASE    ,					norm_elementwise_affine=_SCREAMING_SNAKE_CASE    ,					)
					        for d in range(_SCREAMING_SNAKE_CASE						)
					    ]						)

					# 4. Define output layers
					SCREAMING_SNAKE_CASE	:       str 							=    in_channels if out_channels is None else out_channels
					if self.is_input_continuous:
									# TODO: should use out_channels for continuous projections
									if use_linear_projection:
													SCREAMING_SNAKE_CASE	:       int 							=    nn.Linear(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE						)
									else:
													SCREAMING_SNAKE_CASE	:       Dict 							=    nn.Convad(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					kernel_size=1    ,					stride=1    ,					padding=0						)
					elif self.is_input_vectorized:
									SCREAMING_SNAKE_CASE	:       Optional[int] 							=    nn.LayerNorm(_SCREAMING_SNAKE_CASE						)
									SCREAMING_SNAKE_CASE	:       Optional[int] 							=    nn.Linear(_SCREAMING_SNAKE_CASE    ,					self.num_vector_embeds - 1						)
					elif self.is_input_patches:
									SCREAMING_SNAKE_CASE	:       Tuple 							=    nn.LayerNorm(_SCREAMING_SNAKE_CASE    ,					elementwise_affine=_SCREAMING_SNAKE_CASE    ,					eps=1E-6						)
									SCREAMING_SNAKE_CASE	:       int 							=    nn.Linear(_SCREAMING_SNAKE_CASE    ,					2 * inner_dim						)
									SCREAMING_SNAKE_CASE	:       Optional[int] 							=    nn.Linear(_SCREAMING_SNAKE_CASE    ,					patch_size * patch_size * self.out_channels						)







	def  _lowerCAmelCase    (			self							:					Tuple    ,					_SCREAMING_SNAKE_CASE							:					torch.Tensor    ,					_SCREAMING_SNAKE_CASE							:					Optional[torch.Tensor] = None    ,					_SCREAMING_SNAKE_CASE							:					Optional[torch.LongTensor] = None    ,					_SCREAMING_SNAKE_CASE							:					Optional[torch.LongTensor] = None    ,					_SCREAMING_SNAKE_CASE							:					Dict[str, Any] = None    ,					_SCREAMING_SNAKE_CASE							:					Optional[torch.Tensor] = None    ,					_SCREAMING_SNAKE_CASE							:					Optional[torch.Tensor] = None    ,					_SCREAMING_SNAKE_CASE							:					bool = True    ,					)	->  int:





					"""simple docstring"""


					if attention_mask is not None and attention_mask.ndim == 2:
									# assume that mask is expressed as:
									#   (1 = keep,      0 = discard)
									# convert mask into a bias that can be added to attention scores:
									#       (keep = +0,     discard = -10000.0)
									SCREAMING_SNAKE_CASE	:       Union[str, Any] 							=    (1 - attention_mask.to(hidden_states.dtype						)) * -1_0_0_0_0.0
									SCREAMING_SNAKE_CASE	:       List[str] 							=    attention_mask.unsqueeze(1						)

					# convert encoder_attention_mask to a bias the same way we do for attention_mask
					if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
									SCREAMING_SNAKE_CASE	:       List[Any] 							=    (1 - encoder_attention_mask.to(hidden_states.dtype						)) * -1_0_0_0_0.0
									SCREAMING_SNAKE_CASE	:       int 							=    encoder_attention_mask.unsqueeze(1						)

					# 1. Input
					if self.is_input_continuous:
									SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE	:       int 							=    hidden_states.shape
									SCREAMING_SNAKE_CASE	:       Any 							=    hidden_states

									SCREAMING_SNAKE_CASE	:       Union[str, Any] 							=    self.norm(_SCREAMING_SNAKE_CASE						)
									if not self.use_linear_projection:
													SCREAMING_SNAKE_CASE	:       Optional[Any] 							=    self.proj_in(_SCREAMING_SNAKE_CASE						)
													SCREAMING_SNAKE_CASE	:       Optional[int] 							=    hidden_states.shape[1]
													SCREAMING_SNAKE_CASE	:       str 							=    hidden_states.permute(0    ,					2    ,					3    ,					1						).reshape(_SCREAMING_SNAKE_CASE    ,					height * width    ,					_SCREAMING_SNAKE_CASE						)
									else:
													SCREAMING_SNAKE_CASE	:       str 							=    hidden_states.shape[1]
													SCREAMING_SNAKE_CASE	:       int 							=    hidden_states.permute(0    ,					2    ,					3    ,					1						).reshape(_SCREAMING_SNAKE_CASE    ,					height * width    ,					_SCREAMING_SNAKE_CASE						)
													SCREAMING_SNAKE_CASE	:       List[str] 							=    self.proj_in(_SCREAMING_SNAKE_CASE						)
					elif self.is_input_vectorized:
									SCREAMING_SNAKE_CASE	:       int 							=    self.latent_image_embedding(_SCREAMING_SNAKE_CASE						)
					elif self.is_input_patches:
									SCREAMING_SNAKE_CASE	:       List[str] 							=    self.pos_embed(_SCREAMING_SNAKE_CASE						)

					# 2. Blocks
					for block in self.transformer_blocks:
									SCREAMING_SNAKE_CASE	:       List[str] 							=    block(
									    _SCREAMING_SNAKE_CASE    ,					attention_mask=_SCREAMING_SNAKE_CASE    ,					encoder_hidden_states=_SCREAMING_SNAKE_CASE    ,					encoder_attention_mask=_SCREAMING_SNAKE_CASE    ,					timestep=_SCREAMING_SNAKE_CASE    ,					cross_attention_kwargs=_SCREAMING_SNAKE_CASE    ,					class_labels=_SCREAMING_SNAKE_CASE    ,					)

					# 3. Output
					if self.is_input_continuous:
									if not self.use_linear_projection:
													SCREAMING_SNAKE_CASE	:       List[str] 							=    hidden_states.reshape(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE						).permute(0    ,					3    ,					1    ,					2						).contiguous()
													SCREAMING_SNAKE_CASE	:       Dict 							=    self.proj_out(_SCREAMING_SNAKE_CASE						)
									else:
													SCREAMING_SNAKE_CASE	:       Union[str, Any] 							=    self.proj_out(_SCREAMING_SNAKE_CASE						)
													SCREAMING_SNAKE_CASE	:       int 							=    hidden_states.reshape(_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE						).permute(0    ,					3    ,					1    ,					2						).contiguous()

									SCREAMING_SNAKE_CASE	:       Any 							=    hidden_states + residual
					elif self.is_input_vectorized:
									SCREAMING_SNAKE_CASE	:       int 							=    self.norm_out(_SCREAMING_SNAKE_CASE						)
									SCREAMING_SNAKE_CASE	:       int 							=    self.out(_SCREAMING_SNAKE_CASE						)
									# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
									SCREAMING_SNAKE_CASE	:       List[Any] 							=    logits.permute(0    ,					2    ,					1						)

									# log(p(x_0))
									SCREAMING_SNAKE_CASE	:       Optional[Any] 							=    F.log_softmax(logits.double()    ,					dim=1						).float()
					elif self.is_input_patches:
									# TODO: cleanup!
									SCREAMING_SNAKE_CASE	:       Dict 							=    self.transformer_blocks[0].norma.emb(
									    _SCREAMING_SNAKE_CASE    ,					_SCREAMING_SNAKE_CASE    ,					hidden_dtype=hidden_states.dtype						)
									SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE	:       List[Any] 							=    self.proj_out_a(F.silu(_SCREAMING_SNAKE_CASE						)						).chunk(2    ,					dim=1						)
									SCREAMING_SNAKE_CASE	:       Union[str, Any] 							=    self.norm_out(_SCREAMING_SNAKE_CASE						) * (1 + scale[:, None]) + shift[:, None]
									SCREAMING_SNAKE_CASE	:       List[Any] 							=    self.proj_out_a(_SCREAMING_SNAKE_CASE						)

									# unpatchify
									SCREAMING_SNAKE_CASE	:       List[str] 							=    int(hidden_states.shape[1] ** 0.5						)
									SCREAMING_SNAKE_CASE	:       Optional[Any] 							=    hidden_states.reshape(
									    shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)						)
									SCREAMING_SNAKE_CASE	:       Union[str, Any] 							=    torch.einsum('nhwpqc->nchpwq'    ,					_SCREAMING_SNAKE_CASE						)
									SCREAMING_SNAKE_CASE	:       str 							=    hidden_states.reshape(
									    shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)						)

					if not return_dict:
									return (output,)

					return TransformeraDModelOutput(sample=_SCREAMING_SNAKE_CASE						)










"""simple docstring"""

import torch
import torch.nn.functional as F

from diffusers import VQDiffusionScheduler

from .test_schedulers import SchedulerCommonTest
class 						lowerCAmelCase__	(  _lowerCamelCase   ):


	'''simple docstring'''







	_SCREAMING_SNAKE_CASE  :		Dict     	=   (VQDiffusionScheduler,)





	def  _lowerCAmelCase    (			self							:					Tuple    ,					**_SCREAMING_SNAKE_CASE							:					Optional[int]						)	->  int:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE	:       List[str] 							=    {
					    'num_vec_classes': 4_097,
					    'num_train_timesteps': 100,
					}

					config.update(**_SCREAMING_SNAKE_CASE						)
					return config





	def  _lowerCAmelCase    (			self							:					Dict    ,					_SCREAMING_SNAKE_CASE							:					Tuple						)	->  List[str]:





					"""simple docstring"""


					SCREAMING_SNAKE_CASE	:       Tuple 							=    4
					SCREAMING_SNAKE_CASE	:       List[str] 							=    8
					SCREAMING_SNAKE_CASE	:       List[Any] 							=    8

					SCREAMING_SNAKE_CASE	:       Tuple 							=    torch.randint(0    ,					_SCREAMING_SNAKE_CASE    ,					(batch_size, height * width)						)

					return sample





	@property
	def  _lowerCAmelCase    (			self							:					str						)	->  int:





					"""simple docstring"""


					assert False





	def  _lowerCAmelCase    (			self							:					List[str]    ,					_SCREAMING_SNAKE_CASE							:					Optional[Any]						)	->  Optional[int]:





					"""simple docstring"""


					def model(_SCREAMING_SNAKE_CASE							:					Optional[Any]    ,					_SCREAMING_SNAKE_CASE							:					Optional[Any]    ,					*_SCREAMING_SNAKE_CASE							:					str						):
									SCREAMING_SNAKE_CASE							,	SCREAMING_SNAKE_CASE	:       Any 							=    sample.shape
									SCREAMING_SNAKE_CASE	:       Optional[int] 							=    torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels)						)
									SCREAMING_SNAKE_CASE	:       List[str] 							=    F.log_softmax(logits.double()    ,					dim=1						).float()
									return return_value

					return model





	def  _lowerCAmelCase    (			self							:					List[str]						)	->  Dict:





					"""simple docstring"""


					for timesteps in [2, 5, 100, 1_000]:
									self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE						)





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





					"""simple docstring"""


					for num_vec_classes in [5, 100, 1_000, 4_000]:
									self.check_over_configs(num_vec_classes=_SCREAMING_SNAKE_CASE						)





	def  _lowerCAmelCase    (			self							:					Dict						)	->  List[Any]:





					"""simple docstring"""


					for t in [0, 50, 99]:
									self.check_over_forward(time_step=_SCREAMING_SNAKE_CASE						)







	def  _lowerCAmelCase    (			self							:					Any						)	->  Dict:





					"""simple docstring"""


					pass





