
'''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 	__lowercase				(  _lowercase    ):
			lowerCamelCase			:		"DiagonalGaussianDistribution"

class 	__lowercase				(  _lowercase      ,     _lowercase    ):
			lowerCamelCase			:		Optional[int]	      =			True

			@register_to_config
			def __init__(self			,   A = 3			,   A = 3			,   A = ("DownEncoderBlock2D",)			,   A = ("UpDecoderBlock2D",)			,   A = (6_4,)			,   A = 1			,   A = "silu"			,   A = 4			,   A = 3_2			,   A = 3_2			,   A = 0.1_82_15			,   ):
								super().__init__()

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

								# pass init params to Decoder
								lowerCamelCase_					:      List[str]     =		Decoder(
								    in_channels=A			,   out_channels=A			,   up_block_types=A			,   block_out_channels=A			,   layers_per_block=A			,   norm_num_groups=A			,   act_fn=A			,   )

								lowerCamelCase_					:      List[Any]     =		nn.Convad(2 * latent_channels			,   2 * latent_channels			,   1      )
								lowerCamelCase_					:      str     =		nn.Convad(A			,   A			,   1      )

								lowerCamelCase_					:      Any     =		False
								lowerCamelCase_					:      List[str]     =		False

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

			def        UpperCAmelCase__ (self			,   A			,   A=False      ):
								if isinstance(A			,   (Encoder, Decoder)      ):
													lowerCamelCase_					:      Optional[int]     =		value

			def        UpperCAmelCase__ (self			,   A = True      ):
								lowerCamelCase_					:      Optional[int]     =		use_tiling

			def        UpperCAmelCase__ (self      ):
								self.enable_tiling(A      )

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Optional[Any]     =		True

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Union[str, Any]     =		False

			@property
			# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Optional[int]     =		{}

								def fn_recursive_add_processors(A			,   A			,   A      ):
													if hasattr(A			,   '''set_processor'''      ):
																		lowerCamelCase_					:      Any     =		module.processor

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

													return processors

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

								return processors

			def        UpperCAmelCase__ (self			,   A      ):
								lowerCamelCase_					:      List[str]     =		len(self.attn_processors.keys()      )

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

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

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

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

			@apply_forward_hook
			def        UpperCAmelCase__ (self			,   A			,   A = True      ):
								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(A			,   return_dict=A      )

								if self.use_slicing and x.shape[0] > 1:
													lowerCamelCase_					:      int     =		[self.encoder(A      ) for x_slice in x.split(1      )]
													lowerCamelCase_					:      str     =		torch.cat(A      )
								else:
													lowerCamelCase_					:      str     =		self.encoder(A      )

								lowerCamelCase_					:      List[Any]     =		self.quant_conv(A      )
								lowerCamelCase_					:      Tuple     =		DiagonalGaussianDistribution(A      )

								if not return_dict:
													return (posterior,)

								return AutoencoderKLOutput(latent_dist=A      )

			def        UpperCAmelCase__ (self			,   A			,   A = True      ):
								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(A			,   return_dict=A      )

								lowerCamelCase_					:      Dict     =		self.post_quant_conv(A      )
								lowerCamelCase_					:      int     =		self.decoder(A      )

								if not return_dict:
													return (dec,)

								return DecoderOutput(sample=A      )

			@apply_forward_hook
			def        UpperCAmelCase__ (self			,   A			,   A = True      ):
								if self.use_slicing and z.shape[0] > 1:
													lowerCamelCase_					:      int     =		[self._decode(A      ).sample for z_slice in z.split(1      )]
													lowerCamelCase_					:      Any     =		torch.cat(A      )
								else:
													lowerCamelCase_					:      List[Any]     =		self._decode(A      ).sample

								if not return_dict:
													return (decoded,)

								return DecoderOutput(sample=A      )

			def        UpperCAmelCase__ (self			,   A			,   A			,   A      ):
								lowerCamelCase_					:      List[Any]     =		min(a.shape[2]			,   b.shape[2]			,   A      )
								for y in range(A      ):
													lowerCamelCase_					:      Dict     =		a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
								return b

			def        UpperCAmelCase__ (self			,   A			,   A			,   A      ):
								lowerCamelCase_					:      List[Any]     =		min(a.shape[3]			,   b.shape[3]			,   A      )
								for x in range(A      ):
													lowerCamelCase_					:      Optional[int]     =		a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
								return b

			def        UpperCAmelCase__ (self			,   A			,   A = True      ):
								lowerCamelCase_					:      Optional[int]     =		int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)      )
								lowerCamelCase_					:      str     =		int(self.tile_latent_min_size * self.tile_overlap_factor      )
								lowerCamelCase_					:      Tuple     =		self.tile_latent_min_size - blend_extent

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

								lowerCamelCase_					:      int     =		torch.cat(A			,   dim=2      )
								lowerCamelCase_					:      str     =		DiagonalGaussianDistribution(A      )

								if not return_dict:
													return (posterior,)

								return AutoencoderKLOutput(latent_dist=A      )

			def        UpperCAmelCase__ (self			,   A			,   A = True      ):
								lowerCamelCase_					:      Optional[Any]     =		int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)      )
								lowerCamelCase_					:      Optional[Any]     =		int(self.tile_sample_min_size * self.tile_overlap_factor      )
								lowerCamelCase_					:      str     =		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.
								lowerCamelCase_					:      Dict     =		[]
								for i in range(0			,   z.shape[2]			,   A      ):
													lowerCamelCase_					:      str     =		[]
													for j in range(0			,   z.shape[3]			,   A      ):
																		lowerCamelCase_					:      Union[str, Any]     =		z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
																		lowerCamelCase_					:      List[Any]     =		self.post_quant_conv(A      )
																		lowerCamelCase_					:      List[Any]     =		self.decoder(A      )
																		row.append(A      )
													rows.append(A      )
								lowerCamelCase_					:      Optional[int]     =		[]
								for i, row in enumerate(A      ):
													lowerCamelCase_					:      Tuple     =		[]
													for j, tile in enumerate(A      ):
																		# blend the above tile and the left tile
																		# to the current tile and add the current tile to the result row
																		if i > 0:
																							lowerCamelCase_					:      Optional[Any]     =		self.blend_v(rows[i - 1][j]			,   A			,   A      )
																		if j > 0:
																							lowerCamelCase_					:      List[str]     =		self.blend_h(row[j - 1]			,   A			,   A      )
																		result_row.append(tile[:, :, :row_limit, :row_limit]      )
													result_rows.append(torch.cat(A			,   dim=3      )      )

								lowerCamelCase_					:      List[str]     =		torch.cat(A			,   dim=2      )
								if not return_dict:
													return (dec,)

								return DecoderOutput(sample=A      )






			def        UpperCAmelCase__ (self			,   A			,   A = False			,   A = True			,   A = None			,   ):
								lowerCamelCase_					:      int     =		sample
								lowerCamelCase_					:      str     =		self.encode(A      ).latent_dist
								if sample_posterior:
													lowerCamelCase_					:      Union[str, Any]     =		posterior.sample(generator=A      )
								else:
													lowerCamelCase_					:      Tuple     =		posterior.mode()
								lowerCamelCase_					:      Dict     =		self.decode(A      ).sample

								if not return_dict:
													return (dec,)

								return DecoderOutput(sample=A      )




'''simple docstring'''


def lowercase_	(   _lowercase = 1_000				)     ->							int:
					'''simple docstring'''

					lowerCamelCase_					:      Any     =		-1
					lowerCamelCase_					:      Optional[Any]     =		0
					for a in range(1			,       n // 3				):
										# Solving the two equations a**2+b**2=c**2 and a+b+c=N eliminating c
										lowerCamelCase_					:      Any     =		(n * n - 2 * a * n) // (2 * n - 2 * a)
										lowerCamelCase_					:      int     =		n - a - b
										if c * c == (a * a + b * b):
															lowerCamelCase_					:      int     =		a * b * c
															if candidate >= product:
																				lowerCamelCase_					:      Any     =		candidate
					return product


if __name__ == "__main__":
					print(f'{solution() = }')


