






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 					UpperCamelCase_ (     snake_case_							):



	'''simple docstring'''




	lowerCAmelCase						=					42






class 					UpperCamelCase_ (     snake_case_			,   snake_case_							):



	'''simple docstring'''




	lowerCAmelCase						=					True




	@register_to_config
	def __init__(			self			,     a = 3			,     a = 3			,     a = ("DownEncoderBlock2D",)			,     a = ("UpDecoderBlock2D",)			,     a = (64,)			,     a = 1			,     a = "silu"			,     a = 4			,     a = 32			,     a = 32			,     a = 0.18_215			,     )			->       Optional[int]:
								super().__init__()

								# pass init params to Encoder
								snake_case_           =				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
								snake_case_           =				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			,     )

								snake_case_           =				nn.Convad(2 * latent_channels			,     2 * latent_channels			,     1   )
								snake_case_           =				nn.Convad(a			,     a			,     1   )

								snake_case_           =				False
								snake_case_           =				False

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




	def 				_UpperCamelCase     (			self			,     a			,     a=False   )			->       Dict:
								if isinstance(a			,     (Encoder, Decoder)   ):
															snake_case_           =				value




	def 				_UpperCamelCase     (			self			,     a = True   )			->       Any:
								snake_case_           =				use_tiling




	def 				_UpperCamelCase     (			self   )			->       Union[str, Any]:
								self.enable_tiling(a   )




	def 				_UpperCamelCase     (			self   )			->       str:
								snake_case_           =				True




	def 				_UpperCamelCase     (			self   )			->       int:
								snake_case_           =				False




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

								def fn_recursive_add_processors(a			,     a			,     a   ):
															if hasattr(a			,     'set_processor'   ):
																						snake_case_           =				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   )			->       List[str]:
								snake_case_           =				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   )			->       Dict:
								self.set_attn_processor(AttnProcessor()   )




	@apply_forward_hook
	def 				_UpperCamelCase     (			self			,     a			,     a = 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(a			,     return_dict=a   )

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

								snake_case_           =				self.quant_conv(a   )
								snake_case_           =				DiagonalGaussianDistribution(a   )

								if not return_dict:
															return (posterior,)

								return AutoencoderKLOutput(latent_dist=a   )




	def 				_UpperCamelCase     (			self			,     a			,     a = 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(a			,     return_dict=a   )

								snake_case_           =				self.post_quant_conv(a   )
								snake_case_           =				self.decoder(a   )

								if not return_dict:
															return (dec,)

								return DecoderOutput(sample=a   )




	@apply_forward_hook
	def 				_UpperCamelCase     (			self			,     a			,     a = True   )			->       Union[DecoderOutput, torch.FloatTensor]:
								if self.use_slicing and z.shape[0] > 1:
															snake_case_           =				[self._decode(a   ).sample for z_slice in z.split(1   )]
															snake_case_           =				torch.cat(a   )
								else:
															snake_case_           =				self._decode(a   ).sample

								if not return_dict:
															return (decoded,)

								return DecoderOutput(sample=a   )




	def 				_UpperCamelCase     (			self			,     a			,     a			,     a   )			->       str:
								snake_case_           =				min(a.shape[2]			,     b.shape[2]			,     a   )
								for y in range(a   ):
															snake_case_           =				a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
								return b




	def 				_UpperCamelCase     (			self			,     a			,     a			,     a   )			->       List[Any]:
								snake_case_           =				min(a.shape[3]			,     b.shape[3]			,     a   )
								for x in range(a   ):
															snake_case_           =				a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
								return b




	def 				_UpperCamelCase     (			self			,     a			,     a = True   )			->       AutoencoderKLOutput:
								snake_case_           =				int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)   )
								snake_case_           =				int(self.tile_latent_min_size * self.tile_overlap_factor   )
								snake_case_           =				self.tile_latent_min_size - blend_extent

								# Split the image into 512x512 tiles and encode them separately.
								snake_case_           =				[]
								for i in range(0			,     x.shape[2]			,     a   ):
															snake_case_           =				[]
															for j in range(0			,     x.shape[3]			,     a   ):
																						snake_case_           =				x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
																						snake_case_           =				self.encoder(a   )
																						snake_case_           =				self.quant_conv(a   )
																						row.append(a   )
															rows.append(a   )
								snake_case_           =				[]
								for i, row in enumerate(a   ):
															snake_case_           =				[]
															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:
																													snake_case_           =				self.blend_v(rows[i - 1][j]			,     a			,     a   )
																						if j > 0:
																													snake_case_           =				self.blend_h(row[j - 1]			,     a			,     a   )
																						result_row.append(tile[:, :, :row_limit, :row_limit]   )
															result_rows.append(torch.cat(a			,     dim=3   )   )

								snake_case_           =				torch.cat(a			,     dim=2   )
								snake_case_           =				DiagonalGaussianDistribution(a   )

								if not return_dict:
															return (posterior,)

								return AutoencoderKLOutput(latent_dist=a   )




	def 				_UpperCamelCase     (			self			,     a			,     a = True   )			->       Union[DecoderOutput, torch.FloatTensor]:
								snake_case_           =				int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)   )
								snake_case_           =				int(self.tile_sample_min_size * self.tile_overlap_factor   )
								snake_case_           =				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.
								snake_case_           =				[]
								for i in range(0			,     z.shape[2]			,     a   ):
															snake_case_           =				[]
															for j in range(0			,     z.shape[3]			,     a   ):
																						snake_case_           =				z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
																						snake_case_           =				self.post_quant_conv(a   )
																						snake_case_           =				self.decoder(a   )
																						row.append(a   )
															rows.append(a   )
								snake_case_           =				[]
								for i, row in enumerate(a   ):
															snake_case_           =				[]
															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:
																													snake_case_           =				self.blend_v(rows[i - 1][j]			,     a			,     a   )
																						if j > 0:
																													snake_case_           =				self.blend_h(row[j - 1]			,     a			,     a   )
																						result_row.append(tile[:, :, :row_limit, :row_limit]   )
															result_rows.append(torch.cat(a			,     dim=3   )   )

								snake_case_           =				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			,     )			->       Union[DecoderOutput, torch.FloatTensor]:
								snake_case_           =				sample
								snake_case_           =				self.encode(a   ).latent_dist
								if sample_posterior:
															snake_case_           =				posterior.sample(generator=a   )
								else:
															snake_case_           =				posterior.mode()
								snake_case_           =				self.decode(a   ).sample

								if not return_dict:
															return (dec,)

								return DecoderOutput(sample=a   )













from __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

lowercase    				=   TypeVar("T")
lowercase    				=   TypeVar("U")




class 					UpperCamelCase_ (     Generic[T, U]							):



	'''simple docstring'''




	def __init__(			self			,     a			,     a   )			->       List[Any]:
								snake_case_           =				key
								snake_case_           =				val
								snake_case_           =				None
								snake_case_           =				None






	def __repr__(			self   )			->       str:
								return (
								    F'''Node: key: {self.key}, val: {self.val}, '''
								    F'''has next: {bool(self.next   )}, has prev: {bool(self.prev   )}'''
								)






class 					UpperCamelCase_ (     Generic[T, U]							):



	'''simple docstring'''




	def __init__(			self   )			->       None:
								snake_case_           =				DoubleLinkedListNode(a			,     a   )
								snake_case_           =				DoubleLinkedListNode(a			,     a   )
								snake_case_     ,						snake_case_           =				self.rear, self.head




	def __repr__(			self   )			->       str:
								snake_case_           =				['DoubleLinkedList']
								snake_case_           =				self.head
								while node.next is not None:
															rep.append(str(a   )   )
															snake_case_           =				node.next
								rep.append(str(self.rear   )   )
								return ",\n    ".join(a   )




	def 				_UpperCamelCase     (			self			,     a   )			->       None:
								snake_case_           =				self.rear.prev

								# All nodes other than self.head are guaranteed to have non-None previous
								assert previous is not None

								snake_case_           =				node
								snake_case_           =				previous
								snake_case_           =				node
								snake_case_           =				self.rear






	def 				_UpperCamelCase     (			self			,     a   )			->       DoubleLinkedListNode[T, U] | None:

								if node.prev is None or node.next is None:
															return None

								snake_case_           =				node.next
								snake_case_           =				node.prev
								snake_case_           =				None
								snake_case_           =				None
								return node






class 					UpperCamelCase_ (     Generic[T, U]							):



	'''simple docstring'''




	lowerCAmelCase						=					{}




	def __init__(			self			,     a   )			->       str:
								snake_case_           =				DoubleLinkedList()
								snake_case_           =				capacity
								snake_case_           =				0
								snake_case_           =				0
								snake_case_           =				0
								snake_case_           =				{}




	def __repr__(			self   )			->       str:
								return (
								    F'''CacheInfo(hits={self.hits}, misses={self.miss}, '''
								    F'''capacity={self.capacity}, current size={self.num_keys})'''
								)




	def __contains__(			self			,     a   )			->       bool:
								return key in self.cache




	def 				_UpperCamelCase     (			self			,     a   )			->       U | None:
								# Note: pythonic interface would throw KeyError rather than return None

								if key in self.cache:
															self.hits += 1
															snake_case_           =				self.cache[key]
															snake_case_           =				self.list.remove(self.cache[key]   )
															assert node == value_node

															# node is guaranteed not None because it is in self.cache
															assert node is not None
															self.list.add(a   )
															return node.val
								self.miss += 1
								return None




	def 				_UpperCamelCase     (			self			,     a			,     a   )			->       None:

								if key not in self.cache:
															if self.num_keys >= self.capacity:
																						# delete first node (oldest) when over capacity
																						snake_case_           =				self.list.head.next

																						# guaranteed to have a non-None first node when num_keys > 0
																						# explain to type checker via assertions
																						assert first_node is not None
																						assert first_node.key is not None
																						assert (
																						    self.list.remove(a   ) is not None
																						)  # node guaranteed to be in list assert node.key is not None

																						del self.cache[first_node.key]
																						self.num_keys -= 1
															snake_case_           =				DoubleLinkedListNode(a			,     a   )
															self.list.add(self.cache[key]   )
															self.num_keys += 1

								else:
															# bump node to the end of the list, update value
															snake_case_           =				self.list.remove(self.cache[key]   )
															assert node is not None  # node guaranteed to be in list
															snake_case_           =				value
															self.list.add(a   )






	@classmethod
	def 				_UpperCamelCase     (			cls			,     a = 1_28   )			->       Callable[[Callable[[T], U]], Callable[..., U]]:

								def cache_decorator_inner(a   ) -> Callable[..., U]:
															def cache_decorator_wrapper(*a   ) -> U:
																						if func not in cls.decorator_function_to_instance_map:
																													snake_case_           =				LRUCache(a   )

																						snake_case_           =				cls.decorator_function_to_instance_map[func].get(args[0]   )
																						if result is None:
																													snake_case_           =				func(*a   )
																													cls.decorator_function_to_instance_map[func].put(args[0]			,     a   )
																						return result

															def cache_info() -> LRUCache[T, U]:
																						return cls.decorator_function_to_instance_map[func]

															setattr(a			,     'cache_info'			,     a   )  # noqa: B010

															return cache_decorator_wrapper

								return cache_decorator_inner


if __name__ == "__main__":
	import doctest

	doctest.testmod()





