



'''simple docstring'''






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

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .attention_processor import AttentionProcessor, AttnProcessor
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder

@dataclass
class        SCREAMING_SNAKE_CASE(							A__    ):







	"""simple docstring"""






	lowerCamelCase__    				= 42

class        SCREAMING_SNAKE_CASE(							A__							,	A__    ):







	"""simple docstring"""






	lowerCamelCase__    				= True







	@register_to_config
	def __init__(		self    :  Optional[int]	,    __snake_case    :  int = 3	,    __snake_case    :  int = 3	,    __snake_case    :  Tuple[str] = ("DownEncoderBlock2D",)	,    __snake_case    :  Tuple[str] = ("UpDecoderBlock2D",)	,    __snake_case    :  Tuple[int] = (64,)	,    __snake_case    :  int = 1	,    __snake_case    :  str = "silu"	,    __snake_case    :  int = 4	,    __snake_case    :  int = 32	,    __snake_case    :  int = 32	,    __snake_case    :  float = 0.1_82_15	,    )    ->	Dict:
								super().__init__()

								# pass init params to Encoder
								UpperCAmelCase				: Union[str, Any]											=	Encoder(
								    in_channels=__snake_case	,    out_channels=__snake_case	,    down_block_types=__snake_case	,    block_out_channels=__snake_case	,    layers_per_block=__snake_case	,    act_fn=__snake_case	,    norm_num_groups=__snake_case	,    double_z=__snake_case	,    )

								# pass init params to Decoder
								UpperCAmelCase				: Dict											=	Decoder(
								    in_channels=__snake_case	,    out_channels=__snake_case	,    up_block_types=__snake_case	,    block_out_channels=__snake_case	,    layers_per_block=__snake_case	,    norm_num_groups=__snake_case	,    act_fn=__snake_case	,    )

								UpperCAmelCase				: Optional[int]											=	nn.Convad(2 * latent_channels	,    2 * latent_channels	,    1	)
								UpperCAmelCase				: Tuple											=	nn.Convad(__snake_case	,    __snake_case	,    1	)

								UpperCAmelCase				: Union[str, Any]											=	False
								UpperCAmelCase				: Optional[Any]											=	False

								# only relevant if vae tiling is enabled
								UpperCAmelCase				: List[str]											=	self.config.sample_size
								UpperCAmelCase				: List[str]											=	(
								    self.config.sample_size[0]
								    if isinstance(self.config.sample_size	,    (list, tuple)	)
								    else self.config.sample_size
								)
								UpperCAmelCase				: int											=	int(sample_size / (2 ** (len(self.config.block_out_channels	) - 1))	)
								UpperCAmelCase				: Dict											=	0.25







	def   A			(		self    :  int	,    __snake_case    :  Dict	,    __snake_case    :  List[str]=False	)    ->	Dict:
								if isinstance(__snake_case	,    (Encoder, Decoder)	):
															UpperCAmelCase				: List[str]											=	value







	def   A			(		self    :  Optional[Any]	,    __snake_case    :  bool = True	)    ->	Dict:
								UpperCAmelCase				: Dict											=	use_tiling







	def   A			(		self    :  Tuple	)    ->	Optional[Any]:
								self.enable_tiling(__snake_case	)







	def   A			(		self    :  Dict	)    ->	Tuple:
								UpperCAmelCase				: Optional[int]											=	True







	def   A			(		self    :  List[Any]	)    ->	List[str]:
								UpperCAmelCase				: Dict											=	False







	@property
	# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
	def   A			(		self    :  List[Any]	)    ->	Dict[str, AttentionProcessor]:
								UpperCAmelCase				: int											=	{}

								def fn_recursive_add_processors(__snake_case    :  str	,    __snake_case    :  torch.nn.Module	,    __snake_case    :  Dict[str, AttentionProcessor]	):
															if hasattr(__snake_case	,    '''set_processor'''	):
																						UpperCAmelCase				: str											=	module.processor

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

															return processors

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

								return processors







	def   A			(		self    :  Optional[Any]	,    __snake_case    :  Union[AttentionProcessor, Dict[str, AttentionProcessor]]	)    ->	Dict:
								UpperCAmelCase				: Union[str, Any]											=	len(self.attn_processors.keys()	)

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

								def fn_recursive_attn_processor(__snake_case    :  str	,    __snake_case    :  torch.nn.Module	,    __snake_case    :  List[str]	):
															if hasattr(__snake_case	,    '''set_processor'''	):
																						if not isinstance(__snake_case	,    __snake_case	):
																													module.set_processor(__snake_case	)
																						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}"""	,    __snake_case	,    __snake_case	)

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







	def   A			(		self    :  Any	)    ->	str:
								self.set_attn_processor(AttnProcessor()	)







	@apply_forward_hook
	def   A			(		self    :  List[str]	,    __snake_case    :  torch.FloatTensor	,    __snake_case    :  bool = True	)    ->	AutoencoderKLOutput:
								if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
															return self.tiled_encode(__snake_case	,    return_dict=__snake_case	)

								if self.use_slicing and x.shape[0] > 1:
															UpperCAmelCase				: List[Any]											=	[self.encoder(__snake_case	) for x_slice in x.split(1	)]
															UpperCAmelCase				: Dict											=	torch.cat(__snake_case	)
								else:
															UpperCAmelCase				: Union[str, Any]											=	self.encoder(__snake_case	)

								UpperCAmelCase				: Optional[int]											=	self.quant_conv(__snake_case	)
								UpperCAmelCase				: Any											=	DiagonalGaussianDistribution(__snake_case	)

								if not return_dict:
															return (posterior,)

								return AutoencoderKLOutput(latent_dist=__snake_case	)







	def   A			(		self    :  Optional[Any]	,    __snake_case    :  torch.FloatTensor	,    __snake_case    :  bool = True	)    ->	Union[DecoderOutput, torch.FloatTensor]:
								if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
															return self.tiled_decode(__snake_case	,    return_dict=__snake_case	)

								UpperCAmelCase				: int											=	self.post_quant_conv(__snake_case	)
								UpperCAmelCase				: Tuple											=	self.decoder(__snake_case	)

								if not return_dict:
															return (dec,)

								return DecoderOutput(sample=__snake_case	)







	@apply_forward_hook
	def   A			(		self    :  Tuple	,    __snake_case    :  torch.FloatTensor	,    __snake_case    :  bool = True	)    ->	Union[DecoderOutput, torch.FloatTensor]:
								if self.use_slicing and z.shape[0] > 1:
															UpperCAmelCase				: Optional[int]											=	[self._decode(__snake_case	).sample for z_slice in z.split(1	)]
															UpperCAmelCase				: List[str]											=	torch.cat(__snake_case	)
								else:
															UpperCAmelCase				: Union[str, Any]											=	self._decode(__snake_case	).sample

								if not return_dict:
															return (decoded,)

								return DecoderOutput(sample=__snake_case	)







	def   A			(		self    :  Tuple	,    __snake_case    :  List[str]	,    __snake_case    :  Tuple	,    __snake_case    :  Optional[int]	)    ->	Optional[int]:
								UpperCAmelCase				: List[Any]											=	min(a.shape[2]	,    b.shape[2]	,    __snake_case	)
								for y in range(__snake_case	):
															UpperCAmelCase				: List[Any]											=	a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
								return b







	def   A			(		self    :  int	,    __snake_case    :  Dict	,    __snake_case    :  str	,    __snake_case    :  Any	)    ->	Any:
								UpperCAmelCase				: Tuple											=	min(a.shape[3]	,    b.shape[3]	,    __snake_case	)
								for x in range(__snake_case	):
															UpperCAmelCase				: Optional[int]											=	a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
								return b







	def   A			(		self    :  Tuple	,    __snake_case    :  torch.FloatTensor	,    __snake_case    :  bool = True	)    ->	AutoencoderKLOutput:
								UpperCAmelCase				: Dict											=	int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)	)
								UpperCAmelCase				: Any											=	int(self.tile_latent_min_size * self.tile_overlap_factor	)
								UpperCAmelCase				: List[Any]											=	self.tile_latent_min_size - blend_extent

								# Split the image into 512x512 tiles and encode them separately.
								UpperCAmelCase				: str											=	[]
								for i in range(0	,    x.shape[2]	,    __snake_case	):
															UpperCAmelCase				: Union[str, Any]											=	[]
															for j in range(0	,    x.shape[3]	,    __snake_case	):
																						UpperCAmelCase				: Dict											=	x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
																						UpperCAmelCase				: List[Any]											=	self.encoder(__snake_case	)
																						UpperCAmelCase				: Optional[Any]											=	self.quant_conv(__snake_case	)
																						row.append(__snake_case	)
															rows.append(__snake_case	)
								UpperCAmelCase				: Any											=	[]
								for i, row in enumerate(__snake_case	):
															UpperCAmelCase				: int											=	[]
															for j, tile in enumerate(__snake_case	):
																						# blend the above tile and the left tile
																						# to the current tile and add the current tile to the result row
																						if i > 0:
																													UpperCAmelCase				: str											=	self.blend_v(rows[i - 1][j]	,    __snake_case	,    __snake_case	)
																						if j > 0:
																													UpperCAmelCase				: Optional[Any]											=	self.blend_h(row[j - 1]	,    __snake_case	,    __snake_case	)
																						result_row.append(tile[:, :, :row_limit, :row_limit]	)
															result_rows.append(torch.cat(__snake_case	,    dim=3	)	)

								UpperCAmelCase				: Optional[Any]											=	torch.cat(__snake_case	,    dim=2	)
								UpperCAmelCase				: int											=	DiagonalGaussianDistribution(__snake_case	)

								if not return_dict:
															return (posterior,)

								return AutoencoderKLOutput(latent_dist=__snake_case	)







	def   A			(		self    :  Dict	,    __snake_case    :  torch.FloatTensor	,    __snake_case    :  bool = True	)    ->	Union[DecoderOutput, torch.FloatTensor]:
								UpperCAmelCase				: Optional[int]											=	int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)	)
								UpperCAmelCase				: List[Any]											=	int(self.tile_sample_min_size * self.tile_overlap_factor	)
								UpperCAmelCase				: Union[str, Any]											=	self.tile_sample_min_size - blend_extent

								# Split z into overlapping 64x64 tiles and decode them separately.
								# The tiles have an overlap to avoid seams between tiles.
								UpperCAmelCase				: Any											=	[]
								for i in range(0	,    z.shape[2]	,    __snake_case	):
															UpperCAmelCase				: str											=	[]
															for j in range(0	,    z.shape[3]	,    __snake_case	):
																						UpperCAmelCase				: int											=	z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
																						UpperCAmelCase				: Optional[int]											=	self.post_quant_conv(__snake_case	)
																						UpperCAmelCase				: Optional[int]											=	self.decoder(__snake_case	)
																						row.append(__snake_case	)
															rows.append(__snake_case	)
								UpperCAmelCase				: List[Any]											=	[]
								for i, row in enumerate(__snake_case	):
															UpperCAmelCase				: Optional[Any]											=	[]
															for j, tile in enumerate(__snake_case	):
																						# blend the above tile and the left tile
																						# to the current tile and add the current tile to the result row
																						if i > 0:
																													UpperCAmelCase				: str											=	self.blend_v(rows[i - 1][j]	,    __snake_case	,    __snake_case	)
																						if j > 0:
																													UpperCAmelCase				: str											=	self.blend_h(row[j - 1]	,    __snake_case	,    __snake_case	)
																						result_row.append(tile[:, :, :row_limit, :row_limit]	)
															result_rows.append(torch.cat(__snake_case	,    dim=3	)	)

								UpperCAmelCase				: Optional[int]											=	torch.cat(__snake_case	,    dim=2	)
								if not return_dict:
															return (dec,)

								return DecoderOutput(sample=__snake_case	)







	def   A			(		self    :  Dict	,    __snake_case    :  torch.FloatTensor	,    __snake_case    :  bool = False	,    __snake_case    :  bool = True	,    __snake_case    :  Optional[torch.Generator] = None	,    )    ->	Union[DecoderOutput, torch.FloatTensor]:
								UpperCAmelCase				: List[str]											=	sample
								UpperCAmelCase				: Optional[Any]											=	self.encode(__snake_case	).latent_dist
								if sample_posterior:
															UpperCAmelCase				: List[str]											=	posterior.sample(generator=__snake_case	)
								else:
															UpperCAmelCase				: Optional[Any]											=	posterior.mode()
								UpperCAmelCase				: List[Any]											=	self.decode(__snake_case	).sample

								if not return_dict:
															return (dec,)

								return DecoderOutput(sample=__snake_case	)










'''simple docstring'''






import math

import tensorflow as tf
from packaging import version


def 				snake_case_  (      _lowerCAmelCase						:		List[str]						)   ->     Dict:
							UpperCAmelCase				: Tuple											=	tf.convert_to_tensor(_lowerCAmelCase						)
							UpperCAmelCase				: Dict											=	0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0						)     ,   x.dtype						)						))

							return x * cdf


def 				snake_case_  (      _lowerCAmelCase						:		int						)   ->     List[Any]:
							UpperCAmelCase				: str											=	tf.convert_to_tensor(_lowerCAmelCase						)
							UpperCAmelCase				: Union[str, Any]											=	tf.cast(math.pi     ,   x.dtype						)
							UpperCAmelCase				: Tuple											=	tf.cast(0.0_4_4_7_1_5     ,   x.dtype						)
							UpperCAmelCase				: List[Any]											=	0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi						) * (x + coeff * tf.pow(_lowerCAmelCase     ,   3						))						))

							return x * cdf


def 				snake_case_  (      _lowerCAmelCase						:		Optional[Any]						)   ->     Any:
							UpperCAmelCase				: List[Any]											=	tf.convert_to_tensor(_lowerCAmelCase						)

							return x * tf.tanh(tf.math.softplus(_lowerCAmelCase						)						)


def 				snake_case_  (      _lowerCAmelCase						:		str						)   ->     Tuple:
							UpperCAmelCase				: Union[str, Any]											=	tf.convert_to_tensor(_lowerCAmelCase						)
							UpperCAmelCase				: Union[str, Any]											=	tf.cast(0.0_4_4_7_1_5     ,   x.dtype						)
							UpperCAmelCase				: Tuple											=	tf.cast(0.7_9_7_8_8_4_5_6_0_8     ,   x.dtype						)

							return 0.5 * x * (1.0 + tf.tanh(x * coeffa * (1.0 + coeffa * x * x)						))


def 				snake_case_  (      _lowerCAmelCase						:		Optional[Any]						)   ->     Optional[int]:
							UpperCAmelCase				: str											=	tf.convert_to_tensor(_lowerCAmelCase						)
							UpperCAmelCase				: Optional[int]											=	tf.cast(1.7_0_2     ,   x.dtype						)
							return x * tf.math.sigmoid(coeff * x						)


def 				snake_case_  (      _lowerCAmelCase						:		Any						)   ->     Any:
							return tf.clip_by_value(_gelu(_lowerCAmelCase						)     ,   -10     ,   10						)


def 				snake_case_  (      _lowerCAmelCase						:		Dict     ,   _lowerCAmelCase						:		Optional[Any]=-1						)   ->     Union[str, Any]:
							UpperCAmelCase							,	UpperCAmelCase				: Optional[Any]											=	tf.split(_lowerCAmelCase     ,   2     ,   axis=_lowerCAmelCase						)
							return a * tf.math.sigmoid(_lowerCAmelCase						)


if version.parse(tf.version.VERSION) >= version.parse("2.4"):


						def 				snake_case_  (      _lowerCAmelCase						:		Tuple						)   ->     Any:
													return tf.keras.activations.gelu(_lowerCAmelCase     ,   approximate=_lowerCAmelCase						)

						UpperCamelCase__:     List[Any]										=    tf.keras.activations.gelu




						UpperCamelCase__:     List[Any]										=    approximate_gelu_wrap
else:
						UpperCamelCase__:     List[str]										=    _gelu
						UpperCamelCase__:     Tuple										=    _gelu_new


UpperCamelCase__:     List[Any]										=    {
    "gelu": gelu,
    "gelu_10": gelu_aa,
    "gelu_fast": gelu_fast,
    "gelu_new": gelu_new,
    "glu": glu,
    "mish": mish,
    "quick_gelu": quick_gelu,
    "relu": tf.keras.activations.relu,
    "sigmoid": tf.keras.activations.sigmoid,
    "silu": tf.keras.activations.swish,
    "swish": tf.keras.activations.swish,
    "tanh": tf.keras.activations.tanh,
}




def 				snake_case_  (      _lowerCAmelCase						:		Dict						)   ->     List[Any]:
							if activation_string in ACTaFN:
														return ACTaFN[activation_string]
							else:
														raise KeyError(f"""function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()						)}"""						)





