





"""simple docstring"""





from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block
@dataclass
class 		SCREAMING_SNAKE_CASE_				(     UpperCamelCase__ ):
				"""simple docstring"""
				__snake_case :							List[str]									=		42

class 		SCREAMING_SNAKE_CASE_				(     nn.Module ):
				"""simple docstring"""




				def __init__(  self	:Any    ,     __lowercase	:str=3    ,     __lowercase	:List[Any]=3    ,     __lowercase	:str=("DownEncoderBlock2D",)    ,     __lowercase	:Optional[int]=(64,)    ,     __lowercase	:str=2    ,     __lowercase	:Optional[Any]=32    ,     __lowercase	:int="silu"    ,     __lowercase	:List[Any]=True    ,     ):
											super().__init__()
											__lowerCamelCase   :   Optional[int]       =layers_per_block

											__lowerCamelCase   :   Dict       =torch.nn.Convad(
											    __A    ,     block_out_channels[0]    ,     kernel_size=3    ,     stride=1    ,     padding=1    ,     )

											__lowerCamelCase   :   str       =None
											__lowerCamelCase   :   Tuple       =nn.ModuleList([]      )

											# down
											__lowerCamelCase   :   Optional[int]       =block_out_channels[0]
											for i, down_block_type in enumerate(__A      ):
																		__lowerCamelCase   :   int       =output_channel
																		__lowerCamelCase   :   Union[str, Any]       =block_out_channels[i]
																		__lowerCamelCase   :   List[str]       =i == len(__A      ) - 1

																		__lowerCamelCase   :   int       =get_down_block(
																		    __A    ,     num_layers=self.layers_per_block    ,     in_channels=__A    ,     out_channels=__A    ,     add_downsample=not is_final_block    ,     resnet_eps=1e-6    ,     downsample_padding=0    ,     resnet_act_fn=__A    ,     resnet_groups=__A    ,     attention_head_dim=__A    ,     temb_channels=__A    ,     )
																		self.down_blocks.append(__A      )

											# mid
											__lowerCamelCase   :   Union[str, Any]       =UNetMidBlockaD(
											    in_channels=block_out_channels[-1]    ,     resnet_eps=1e-6    ,     resnet_act_fn=__A    ,     output_scale_factor=1    ,     resnet_time_scale_shift='''default'''    ,     attention_head_dim=block_out_channels[-1]    ,     resnet_groups=__A    ,     temb_channels=__A    ,     )

											# out
											__lowerCamelCase   :   Tuple       =nn.GroupNorm(num_channels=block_out_channels[-1]    ,     num_groups=__A    ,     eps=1e-6      )
											__lowerCamelCase   :   Tuple       =nn.SiLU()

											__lowerCamelCase   :   List[Any]       =2 * out_channels if double_z else out_channels
											__lowerCamelCase   :   Tuple       =nn.Convad(block_out_channels[-1]    ,     __A    ,     3    ,     padding=1      )

											__lowerCamelCase   :   Optional[Any]       =False




				def        __lowercase	(  self	:Optional[int]    ,     __lowercase	:Tuple      ):
											__lowerCamelCase   :   List[Any]       =x
											__lowerCamelCase   :   Dict       =self.conv_in(__A      )

											if self.training and self.gradient_checkpointing:

																		def create_custom_forward(__lowercase	:Tuple      ):
																									def custom_forward(*__lowercase	:List[Any]      ):
																																return module(*__A      )

																									return custom_forward

																		# down
																		if is_torch_version('''>='''    ,     '''1.11.0'''      ):
																									for down_block in self.down_blocks:
																																__lowerCamelCase   :   Optional[Any]       =torch.utils.checkpoint.checkpoint(
																																    create_custom_forward(__A      )    ,     __A    ,     use_reentrant=__A      )
																									# middle
																									__lowerCamelCase   :   Tuple       =torch.utils.checkpoint.checkpoint(
																									    create_custom_forward(self.mid_block      )    ,     __A    ,     use_reentrant=__A      )
																		else:
																									for down_block in self.down_blocks:
																																__lowerCamelCase   :   Any       =torch.utils.checkpoint.checkpoint(create_custom_forward(__A      )    ,     __A      )
																									# middle
																									__lowerCamelCase   :   Optional[int]       =torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block      )    ,     __A      )

											else:
																		# down
																		for down_block in self.down_blocks:
																									__lowerCamelCase   :   Any       =down_block(__A      )

																		# middle
																		__lowerCamelCase   :   List[Any]       =self.mid_block(__A      )

											# post-process
											__lowerCamelCase   :   Union[str, Any]       =self.conv_norm_out(__A      )
											__lowerCamelCase   :   List[str]       =self.conv_act(__A      )
											__lowerCamelCase   :   int       =self.conv_out(__A      )

											return sample

class 		SCREAMING_SNAKE_CASE_				(     nn.Module ):
				"""simple docstring"""




				def __init__(  self	:Any    ,     __lowercase	:List[Any]=3    ,     __lowercase	:str=3    ,     __lowercase	:Tuple=("UpDecoderBlock2D",)    ,     __lowercase	:Tuple=(64,)    ,     __lowercase	:Any=2    ,     __lowercase	:int=32    ,     __lowercase	:Optional[int]="silu"    ,     __lowercase	:Optional[Any]="group"    ,     ):
											super().__init__()
											__lowerCamelCase   :   Union[str, Any]       =layers_per_block

											__lowerCamelCase   :   Tuple       =nn.Convad(
											    __A    ,     block_out_channels[-1]    ,     kernel_size=3    ,     stride=1    ,     padding=1    ,     )

											__lowerCamelCase   :   Any       =None
											__lowerCamelCase   :   List[str]       =nn.ModuleList([]      )

											__lowerCamelCase   :   Optional[int]       =in_channels if norm_type == '''spatial''' else None

											# mid
											__lowerCamelCase   :   List[Any]       =UNetMidBlockaD(
											    in_channels=block_out_channels[-1]    ,     resnet_eps=1e-6    ,     resnet_act_fn=__A    ,     output_scale_factor=1    ,     resnet_time_scale_shift='''default''' if norm_type == '''group''' else norm_type    ,     attention_head_dim=block_out_channels[-1]    ,     resnet_groups=__A    ,     temb_channels=__A    ,     )

											# up
											__lowerCamelCase   :   Optional[Any]       =list(reversed(__A      )      )
											__lowerCamelCase   :   List[str]       =reversed_block_out_channels[0]
											for i, up_block_type in enumerate(__A      ):
																		__lowerCamelCase   :   Any       =output_channel
																		__lowerCamelCase   :   List[str]       =reversed_block_out_channels[i]

																		__lowerCamelCase   :   Optional[int]       =i == len(__A      ) - 1

																		__lowerCamelCase   :   Tuple       =get_up_block(
																		    __A    ,     num_layers=self.layers_per_block + 1    ,     in_channels=__A    ,     out_channels=__A    ,     prev_output_channel=__A    ,     add_upsample=not is_final_block    ,     resnet_eps=1e-6    ,     resnet_act_fn=__A    ,     resnet_groups=__A    ,     attention_head_dim=__A    ,     temb_channels=__A    ,     resnet_time_scale_shift=__A    ,     )
																		self.up_blocks.append(__A      )
																		__lowerCamelCase   :   str       =output_channel

											# out
											if norm_type == "spatial":
																		__lowerCamelCase   :   Optional[Any]       =SpatialNorm(block_out_channels[0]    ,     __A      )
											else:
																		__lowerCamelCase   :   Optional[int]       =nn.GroupNorm(num_channels=block_out_channels[0]    ,     num_groups=__A    ,     eps=1e-6      )
											__lowerCamelCase   :   int       =nn.SiLU()
											__lowerCamelCase   :   List[str]       =nn.Convad(block_out_channels[0]    ,     __A    ,     3    ,     padding=1      )

											__lowerCamelCase   :   List[Any]       =False




				def        __lowercase	(  self	:str    ,     __lowercase	:Tuple    ,     __lowercase	:int=None      ):
											__lowerCamelCase   :   Optional[Any]       =z
											__lowerCamelCase   :   Dict       =self.conv_in(__A      )

											__lowerCamelCase   :   Union[str, Any]       =next(iter(self.up_blocks.parameters()      )      ).dtype
											if self.training and self.gradient_checkpointing:

																		def create_custom_forward(__lowercase	:Tuple      ):
																									def custom_forward(*__lowercase	:List[Any]      ):
																																return module(*__A      )

																									return custom_forward

																		if is_torch_version('''>='''    ,     '''1.11.0'''      ):
																									# middle
																									__lowerCamelCase   :   Optional[int]       =torch.utils.checkpoint.checkpoint(
																									    create_custom_forward(self.mid_block      )    ,     __A    ,     __A    ,     use_reentrant=__A      )
																									__lowerCamelCase   :   List[str]       =sample.to(__A      )

																									# up
																									for up_block in self.up_blocks:
																																__lowerCamelCase   :   int       =torch.utils.checkpoint.checkpoint(
																																    create_custom_forward(__A      )    ,     __A    ,     __A    ,     use_reentrant=__A      )
																		else:
																									# middle
																									__lowerCamelCase   :   Any       =torch.utils.checkpoint.checkpoint(
																									    create_custom_forward(self.mid_block      )    ,     __A    ,     __A      )
																									__lowerCamelCase   :   Any       =sample.to(__A      )

																									# up
																									for up_block in self.up_blocks:
																																__lowerCamelCase   :   Any       =torch.utils.checkpoint.checkpoint(create_custom_forward(__A      )    ,     __A    ,     __A      )
											else:
																		# middle
																		__lowerCamelCase   :   List[Any]       =self.mid_block(__A    ,     __A      )
																		__lowerCamelCase   :   str       =sample.to(__A      )

																		# up
																		for up_block in self.up_blocks:
																									__lowerCamelCase   :   Dict       =up_block(__A    ,     __A      )

        # post-process
											if latent_embeds is None:
																		__lowerCamelCase   :   Dict       =self.conv_norm_out(__A      )
											else:
																		__lowerCamelCase   :   str       =self.conv_norm_out(__A    ,     __A      )
											__lowerCamelCase   :   Dict       =self.conv_act(__A      )
											__lowerCamelCase   :   List[Any]       =self.conv_out(__A      )

											return sample

class 		SCREAMING_SNAKE_CASE_				(     nn.Module ):
				"""simple docstring"""




				def __init__(  self	:Optional[int]    ,     __lowercase	:Optional[Any]    ,     __lowercase	:List[Any]    ,     __lowercase	:str    ,     __lowercase	:List[str]=None    ,     __lowercase	:Optional[Any]="random"    ,     __lowercase	:Optional[int]=False    ,     __lowercase	:int=True      ):
											super().__init__()
											__lowerCamelCase   :   int       =n_e
											__lowerCamelCase   :   List[Any]       =vq_embed_dim
											__lowerCamelCase   :   Tuple       =beta
											__lowerCamelCase   :   Dict       =legacy

											__lowerCamelCase   :   Optional[int]       =nn.Embedding(self.n_e    ,     self.vq_embed_dim      )
											self.embedding.weight.data.uniform_(-1.0 / self.n_e    ,     1.0 / self.n_e      )

											__lowerCamelCase   :   Dict       =remap
											if self.remap is not None:
																		self.register_buffer('''used'''    ,     torch.tensor(np.load(self.remap      )      )      )
																		__lowerCamelCase   :   str       =self.used.shape[0]
																		__lowerCamelCase   :   str       =unknown_index  # "random" or "extra" or integer
																		if self.unknown_index == "extra":
																									__lowerCamelCase   :   List[Any]       =self.re_embed
																									__lowerCamelCase   :   Optional[int]       =self.re_embed + 1
																		print(
																		    f'Remapping {self.n_e} indices to {self.re_embed} indices. '
																		    f'Using {self.unknown_index} for unknown indices.'      )
											else:
																		__lowerCamelCase   :   int       =n_e

											__lowerCamelCase   :   Any       =sane_index_shape




				def        __lowercase	(  self	:Tuple    ,     __lowercase	:List[Any]      ):
											__lowerCamelCase   :   str       =inds.shape
											assert len(__A      ) > 1
											__lowerCamelCase   :   Optional[int]       =inds.reshape(ishape[0]    ,     -1      )
											__lowerCamelCase   :   Any       =self.used.to(__A      )
											__lowerCamelCase   :   Union[str, Any]       =(inds[:, :, None] == used[None, None, ...]).long()
											__lowerCamelCase   :   Optional[int]       =match.argmax(-1      )
											__lowerCamelCase   :   Union[str, Any]       =match.sum(2      ) < 1
											if self.unknown_index == "random":
																		__lowerCamelCase   :   Any       =torch.randint(0    ,     self.re_embed    ,     size=new[unknown].shape      ).to(device=new.device      )
											else:
																		__lowerCamelCase   :   Optional[Any]       =self.unknown_index
											return new.reshape(__A      )




				def        __lowercase	(  self	:Dict    ,     __lowercase	:Optional[Any]      ):
											__lowerCamelCase   :   Tuple       =inds.shape
											assert len(__A      ) > 1
											__lowerCamelCase   :   str       =inds.reshape(ishape[0]    ,     -1      )
											__lowerCamelCase   :   List[Any]       =self.used.to(__A      )
											if self.re_embed > self.used.shape[0]:  # extra token
																		__lowerCamelCase   :   Optional[Any]       =0  # simply set to zero
											__lowerCamelCase   :   Dict       =torch.gather(used[None, :][inds.shape[0] * [0], :]    ,     1    ,     __A      )
											return back.reshape(__A      )




				def        __lowercase	(  self	:Optional[int]    ,     __lowercase	:List[str]      ):
											__lowerCamelCase   :   List[Any]       =z.permute(0    ,     2    ,     3    ,     1      ).contiguous()
											__lowerCamelCase   :   List[Any]       =z.view(-1    ,     self.vq_embed_dim      )

											# distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
											__lowerCamelCase   :   str       =torch.argmin(torch.cdist(__A    ,     self.embedding.weight      )    ,     dim=1      )

											__lowerCamelCase   :   Dict       =self.embedding(__A      ).view(z.shape      )
											__lowerCamelCase   :   Optional[Any]       =None
											__lowerCamelCase   :   Union[str, Any]       =None

											# compute loss for embedding
											if not self.legacy:
																		__lowerCamelCase   :   Optional[Any]       =self.beta * torch.mean((z_q.detach() - z) ** 2      ) + torch.mean((z_q - z.detach()) ** 2      )
											else:
																		__lowerCamelCase   :   Tuple       =torch.mean((z_q.detach() - z) ** 2      ) + self.beta * torch.mean((z_q - z.detach()) ** 2      )

											# preserve gradients
											__lowerCamelCase   :   str       =z + (z_q - z).detach()

											# reshape back to match original input shape
											__lowerCamelCase   :   Optional[Any]       =z_q.permute(0    ,     3    ,     1    ,     2      ).contiguous()

											if self.remap is not None:
																		__lowerCamelCase   :   Dict       =min_encoding_indices.reshape(z.shape[0]    ,     -1      )  # add batch axis
																		__lowerCamelCase   :   int       =self.remap_to_used(__A      )
																		__lowerCamelCase   :   Optional[Any]       =min_encoding_indices.reshape(-1    ,     1      )  # flatten

											if self.sane_index_shape:
																		__lowerCamelCase   :   Dict       =min_encoding_indices.reshape(z_q.shape[0]    ,     z_q.shape[2]    ,     z_q.shape[3]      )

											return z_q, loss, (perplexity, min_encodings, min_encoding_indices)




				def        __lowercase	(  self	:Union[str, Any]    ,     __lowercase	:Dict    ,     __lowercase	:Optional[Any]      ):



											if self.remap is not None:
																		__lowerCamelCase   :   Any       =indices.reshape(shape[0]    ,     -1      )  # add batch axis
																		__lowerCamelCase   :   List[str]       =self.unmap_to_all(__A      )
																		__lowerCamelCase   :   Union[str, Any]       =indices.reshape(-1      )  # flatten again

											# get quantized latent vectors
											__lowerCamelCase   :   Dict       =self.embedding(__A      )

											if shape is not None:
																		__lowerCamelCase   :   Dict       =z_q.view(__A      )
																		# reshape back to match original input shape
																		__lowerCamelCase   :   List[Any]       =z_q.permute(0    ,     3    ,     1    ,     2      ).contiguous()

											return z_q

class 		SCREAMING_SNAKE_CASE_				(     UpperCamelCase__ ):
				"""simple docstring"""




				def __init__(  self	:List[Any]    ,     __lowercase	:List[str]    ,     __lowercase	:Dict=False      ):
											__lowerCamelCase   :   List[str]       =parameters
											__lowerCamelCase     ,    __lowerCamelCase   :   Union[str, Any]       =torch.chunk(__A    ,     2    ,     dim=1      )
											__lowerCamelCase   :   int       =torch.clamp(self.logvar    ,     -30.0    ,     20.0      )
											__lowerCamelCase   :   Union[str, Any]       =deterministic
											__lowerCamelCase   :   Any       =torch.exp(0.5 * self.logvar      )
											__lowerCamelCase   :   Union[str, Any]       =torch.exp(self.logvar      )
											if self.deterministic:
																		__lowerCamelCase   :   Optional[Any]       =torch.zeros_like(
																		    self.mean    ,     device=self.parameters.device    ,     dtype=self.parameters.dtype      )




				def        __lowercase	(  self	:Union[str, Any]    ,     __lowercase	:Optional[torch.Generator] = None      ):
											__lowerCamelCase   :   List[Any]       =randn_tensor(
											    self.mean.shape    ,     generator=__A    ,     device=self.parameters.device    ,     dtype=self.parameters.dtype      )
											__lowerCamelCase   :   str       =self.mean + self.std * sample
											return x




				def        __lowercase	(  self	:str    ,     __lowercase	:Any=None      ):



											if self.deterministic:
																		return torch.Tensor([0.0]      )
											else:
																		if other is None:
																									return 0.5 * torch.sum(torch.pow(self.mean    ,     2      ) + self.var - 1.0 - self.logvar    ,     dim=[1, 2, 3]      )
																		else:
																									return 0.5 * torch.sum(
																									    torch.pow(self.mean - other.mean    ,     2      ) / other.var
																									    + self.var / other.var
																									    - 1.0
																									    - self.logvar
																									    + other.logvar    ,     dim=[1, 2, 3]    ,     )




				def        __lowercase	(  self	:str    ,     __lowercase	:Optional[int]    ,     __lowercase	:List[Any]=[1, 2, 3]      ):



											if self.deterministic:
																		return torch.Tensor([0.0]      )
											__lowerCamelCase   :   Union[str, Any]       =np.log(2.0 * np.pi      )
											return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean    ,     2      ) / self.var    ,     dim=__A      )




				def        __lowercase	(  self	:Any      ):
											return self.mean











from dataclasses import dataclass
from typing import 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 .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer




@dataclass
class 	UpperCamelCase_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       42




class 	UpperCamelCase_							(  UpperCamelCase__					,	UpperCamelCase__     ):


							@register_to_config
							def __init__(				self     :Union[str, Any]				,					__A     :int = 3				,					__A     :int = 3				,					__A     :Tuple[str] = ("DownEncoderBlock2D",)				,					__A     :Tuple[str] = ("UpDecoderBlock2D",)				,					__A     :Tuple[int] = (64,)				,					__A     :int = 1				,					__A     :str = "silu"				,					__A     :int = 3				,					__A     :int = 32				,					__A     :int = 256				,					__A     :int = 32				,					__A     :Optional[int] = None				,					__A     :float = 0.1_8_2_1_5				,					__A     :str = "group"				,					) -> Any:



									"""simple docstring"""



									super().__init__()

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

									SCREAMING_SNAKE_CASE__								=   vq_embed_dim if vq_embed_dim is not None else latent_channels

									SCREAMING_SNAKE_CASE__								=   nn.Convad(__A				,					__A				,					1		)
									SCREAMING_SNAKE_CASE__								=   VectorQuantizer(__A				,					__A				,					beta=0.2_5				,					remap=__A				,					sane_index_shape=__A		)
									SCREAMING_SNAKE_CASE__								=   nn.Convad(__A				,					__A				,					1		)

									# pass init params to Decoder
									SCREAMING_SNAKE_CASE__								=   Decoder(
									    in_channels=__A				,					out_channels=__A				,					up_block_types=__A				,					block_out_channels=__A				,					layers_per_block=__A				,					act_fn=__A				,					norm_num_groups=__A				,					norm_type=__A				,					)


							@apply_forward_hook
							def   _snake_case	(				self     :Union[str, Any]				,					__A     :torch.FloatTensor				,					__A     :bool = True		) -> VQEncoderOutput:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.encoder(__A		)
									SCREAMING_SNAKE_CASE__								=   self.quant_conv(__A		)

									if not return_dict:
											return (h,)

									return VQEncoderOutput(latents=__A		)


							@apply_forward_hook
							def   _snake_case	(				self     :Tuple				,					__A     :torch.FloatTensor				,					__A     :bool = False				,					__A     :bool = True		) -> Union[DecoderOutput, torch.FloatTensor]:



									"""simple docstring"""



									if not force_not_quantize:
											SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   self.quantize(__A		)
									else:
											SCREAMING_SNAKE_CASE__								=   h
									SCREAMING_SNAKE_CASE__								=   self.post_quant_conv(__A		)
									SCREAMING_SNAKE_CASE__								=   self.decoder(__A				,					quant if self.config.norm_type == """spatial""" else None		)

									if not return_dict:
											return (dec,)

									return DecoderOutput(sample=__A		)


							def   _snake_case	(				self     :int				,					__A     :torch.FloatTensor				,					__A     :bool = True		) -> Union[DecoderOutput, torch.FloatTensor]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   sample
									SCREAMING_SNAKE_CASE__								=   self.encode(__A		).latents
									SCREAMING_SNAKE_CASE__								=   self.decode(__A		).sample

									if not return_dict:
											return (dec,)

									return DecoderOutput(sample=__A		)