





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 	UpperCamelCase_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       42




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Any				,					__A     :str=3				,					__A     :List[Any]=3				,					__A     :str=("DownEncoderBlock2D",)				,					__A     :Optional[int]=(64,)				,					__A     :str=2				,					__A     :Optional[Any]=32				,					__A     :int="silu"				,					__A     :List[Any]=True				,					) -> Optional[Any]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   layers_per_block

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

									SCREAMING_SNAKE_CASE__								=   None
									SCREAMING_SNAKE_CASE__								=   nn.ModuleList([]		)

									# down
									SCREAMING_SNAKE_CASE__								=   block_out_channels[0]
									for i, down_block_type in enumerate(__A		):
											SCREAMING_SNAKE_CASE__								=   output_channel
											SCREAMING_SNAKE_CASE__								=   block_out_channels[i]
											SCREAMING_SNAKE_CASE__								=   i == len(__A		) - 1

											SCREAMING_SNAKE_CASE__								=   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
									SCREAMING_SNAKE_CASE__								=   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
									SCREAMING_SNAKE_CASE__								=   nn.GroupNorm(num_channels=block_out_channels[-1]				,					num_groups=__A				,					eps=1E-6		)
									SCREAMING_SNAKE_CASE__								=   nn.SiLU()

									SCREAMING_SNAKE_CASE__								=   2 * out_channels if double_z else out_channels
									SCREAMING_SNAKE_CASE__								=   nn.Convad(block_out_channels[-1]				,					__A				,					3				,					padding=1		)

									SCREAMING_SNAKE_CASE__								=   False


							def   _snake_case	(				self     :Optional[int]				,					__A     :Tuple		) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   x
									SCREAMING_SNAKE_CASE__								=   self.conv_in(__A		)

									if self.training and self.gradient_checkpointing:

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

													return custom_forward

											# down
											if is_torch_version(""">="""				,					"""1.11.0"""		):
													for down_block in self.down_blocks:
															SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(
															    create_custom_forward(__A		)				,					__A				,					use_reentrant=__A		)
													# middle
													SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(
													    create_custom_forward(self.mid_block		)				,					__A				,					use_reentrant=__A		)
											else:
													for down_block in self.down_blocks:
															SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(create_custom_forward(__A		)				,					__A		)
													# middle
													SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block		)				,					__A		)

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

											# middle
											SCREAMING_SNAKE_CASE__								=   self.mid_block(__A		)

									# post-process
									SCREAMING_SNAKE_CASE__								=   self.conv_norm_out(__A		)
									SCREAMING_SNAKE_CASE__								=   self.conv_act(__A		)
									SCREAMING_SNAKE_CASE__								=   self.conv_out(__A		)

									return sample




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Any				,					__A     :List[Any]=3				,					__A     :str=3				,					__A     :Tuple=("UpDecoderBlock2D",)				,					__A     :Tuple=(64,)				,					__A     :Any=2				,					__A     :int=32				,					__A     :Optional[int]="silu"				,					__A     :Optional[Any]="group"				,					) -> Optional[int]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   layers_per_block

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

									SCREAMING_SNAKE_CASE__								=   None
									SCREAMING_SNAKE_CASE__								=   nn.ModuleList([]		)

									SCREAMING_SNAKE_CASE__								=   in_channels if norm_type == """spatial""" else None

									# mid
									SCREAMING_SNAKE_CASE__								=   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
									SCREAMING_SNAKE_CASE__								=   list(reversed(__A		)		)
									SCREAMING_SNAKE_CASE__								=   reversed_block_out_channels[0]
									for i, up_block_type in enumerate(__A		):
											SCREAMING_SNAKE_CASE__								=   output_channel
											SCREAMING_SNAKE_CASE__								=   reversed_block_out_channels[i]

											SCREAMING_SNAKE_CASE__								=   i == len(__A		) - 1

											SCREAMING_SNAKE_CASE__								=   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		)
											SCREAMING_SNAKE_CASE__								=   output_channel

									# out
									if norm_type == "spatial":
											SCREAMING_SNAKE_CASE__								=   SpatialNorm(block_out_channels[0]				,					__A		)
									else:
											SCREAMING_SNAKE_CASE__								=   nn.GroupNorm(num_channels=block_out_channels[0]				,					num_groups=__A				,					eps=1E-6		)
									SCREAMING_SNAKE_CASE__								=   nn.SiLU()
									SCREAMING_SNAKE_CASE__								=   nn.Convad(block_out_channels[0]				,					__A				,					3				,					padding=1		)

									SCREAMING_SNAKE_CASE__								=   False


							def   _snake_case	(				self     :str				,					__A     :Tuple				,					__A     :int=None		) -> str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   z
									SCREAMING_SNAKE_CASE__								=   self.conv_in(__A		)

									SCREAMING_SNAKE_CASE__								=   next(iter(self.up_blocks.parameters()		)		).dtype
									if self.training and self.gradient_checkpointing:

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

													return custom_forward

											if is_torch_version(""">="""				,					"""1.11.0"""		):
													# middle
													SCREAMING_SNAKE_CASE__								=   torch.utils.checkpoint.checkpoint(
													    create_custom_forward(self.mid_block		)				,					__A				,					__A				,					use_reentrant=__A		)
													SCREAMING_SNAKE_CASE__								=   sample.to(__A		)

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

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

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

        # post-process
									if latent_embeds is None:
											SCREAMING_SNAKE_CASE__								=   self.conv_norm_out(__A		)
									else:
											SCREAMING_SNAKE_CASE__								=   self.conv_norm_out(__A				,					__A		)
									SCREAMING_SNAKE_CASE__								=   self.conv_act(__A		)
									SCREAMING_SNAKE_CASE__								=   self.conv_out(__A		)

									return sample




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Optional[int]				,					__A     :Optional[Any]				,					__A     :List[Any]				,					__A     :str				,					__A     :List[str]=None				,					__A     :Optional[Any]="random"				,					__A     :Optional[int]=False				,					__A     :int=True		) -> Optional[int]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   n_e
									SCREAMING_SNAKE_CASE__								=   vq_embed_dim
									SCREAMING_SNAKE_CASE__								=   beta
									SCREAMING_SNAKE_CASE__								=   legacy

									SCREAMING_SNAKE_CASE__								=   nn.Embedding(self.n_e				,					self.vq_embed_dim		)
									self.embedding.weight.data.uniform_(-1.0 / self.n_e				,					1.0 / self.n_e		)

									SCREAMING_SNAKE_CASE__								=   remap
									if self.remap is not None:
											self.register_buffer("""used"""				,					torch.tensor(np.load(self.remap		)		)		)
											SCREAMING_SNAKE_CASE__								=   self.used.shape[0]
											SCREAMING_SNAKE_CASE__								=   unknown_index  # "random" or "extra" or integer
											if self.unknown_index == "extra":
													SCREAMING_SNAKE_CASE__								=   self.re_embed
													SCREAMING_SNAKE_CASE__								=   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:
											SCREAMING_SNAKE_CASE__								=   n_e

									SCREAMING_SNAKE_CASE__								=   sane_index_shape


							def   _snake_case	(				self     :Tuple				,					__A     :List[Any]		) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   inds.shape
									assert len(__A		) > 1
									SCREAMING_SNAKE_CASE__								=   inds.reshape(ishape[0]				,					-1		)
									SCREAMING_SNAKE_CASE__								=   self.used.to(__A		)
									SCREAMING_SNAKE_CASE__								=   (inds[:, :, None] == used[None, None, ...]).long()
									SCREAMING_SNAKE_CASE__								=   match.argmax(-1		)
									SCREAMING_SNAKE_CASE__								=   match.sum(2		) < 1
									if self.unknown_index == "random":
											SCREAMING_SNAKE_CASE__								=   torch.randint(0				,					self.re_embed				,					size=new[unknown].shape		).to(device=new.device		)
									else:
											SCREAMING_SNAKE_CASE__								=   self.unknown_index
									return new.reshape(__A		)


							def   _snake_case	(				self     :Dict				,					__A     :Optional[Any]		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   inds.shape
									assert len(__A		) > 1
									SCREAMING_SNAKE_CASE__								=   inds.reshape(ishape[0]				,					-1		)
									SCREAMING_SNAKE_CASE__								=   self.used.to(__A		)
									if self.re_embed > self.used.shape[0]:  # extra token
											SCREAMING_SNAKE_CASE__								=   0  # simply set to zero
									SCREAMING_SNAKE_CASE__								=   torch.gather(used[None, :][inds.shape[0] * [0], :]				,					1				,					__A		)
									return back.reshape(__A		)


							def   _snake_case	(				self     :Optional[int]				,					__A     :List[str]		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   z.permute(0				,					2				,					3				,					1		).contiguous()
									SCREAMING_SNAKE_CASE__								=   z.view(-1				,					self.vq_embed_dim		)

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

									SCREAMING_SNAKE_CASE__								=   self.embedding(__A		).view(z.shape		)
									SCREAMING_SNAKE_CASE__								=   None
									SCREAMING_SNAKE_CASE__								=   None

									# compute loss for embedding
									if not self.legacy:
											SCREAMING_SNAKE_CASE__								=   self.beta * torch.mean((z_q.detach() - z) ** 2		) + torch.mean((z_q - z.detach()) ** 2		)
									else:
											SCREAMING_SNAKE_CASE__								=   torch.mean((z_q.detach() - z) ** 2		) + self.beta * torch.mean((z_q - z.detach()) ** 2		)

									# preserve gradients
									SCREAMING_SNAKE_CASE__								=   z + (z_q - z).detach()

									# reshape back to match original input shape
									SCREAMING_SNAKE_CASE__								=   z_q.permute(0				,					3				,					1				,					2		).contiguous()

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

									if self.sane_index_shape:
											SCREAMING_SNAKE_CASE__								=   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   _snake_case	(				self     :Union[str, Any]				,					__A     :Dict				,					__A     :Optional[Any]		) -> Dict:



									"""simple docstring"""



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

									# get quantized latent vectors
									SCREAMING_SNAKE_CASE__								=   self.embedding(__A		)

									if shape is not None:
											SCREAMING_SNAKE_CASE__								=   z_q.view(__A		)
											# reshape back to match original input shape
											SCREAMING_SNAKE_CASE__								=   z_q.permute(0				,					3				,					1				,					2		).contiguous()

									return z_q




class 	UpperCamelCase_							(  UpperCamelCase__     ):


							def __init__(				self     :List[Any]				,					__A     :List[str]				,					__A     :Dict=False		) -> List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   parameters
									SCREAMING_SNAKE_CASE__						,    SCREAMING_SNAKE_CASE__								=   torch.chunk(__A				,					2				,					dim=1		)
									SCREAMING_SNAKE_CASE__								=   torch.clamp(self.logvar				,					-3_0.0				,					2_0.0		)
									SCREAMING_SNAKE_CASE__								=   deterministic
									SCREAMING_SNAKE_CASE__								=   torch.exp(0.5 * self.logvar		)
									SCREAMING_SNAKE_CASE__								=   torch.exp(self.logvar		)
									if self.deterministic:
											SCREAMING_SNAKE_CASE__								=   SCREAMING_SNAKE_CASE__								=   torch.zeros_like(
											    self.mean				,					device=self.parameters.device				,					dtype=self.parameters.dtype		)


							def   _snake_case	(				self     :Union[str, Any]				,					__A     :Optional[torch.Generator] = None		) -> torch.FloatTensor:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   randn_tensor(
									    self.mean.shape				,					generator=__A				,					device=self.parameters.device				,					dtype=self.parameters.dtype		)
									SCREAMING_SNAKE_CASE__								=   self.mean + self.std * sample
									return x


							def   _snake_case	(				self     :str				,					__A     :Any=None		) -> Union[str, Any]:



									"""simple docstring"""



									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   _snake_case	(				self     :str				,					__A     :Optional[int]				,					__A     :List[Any]=[1, 2, 3]		) -> Any:



									"""simple docstring"""



									if self.deterministic:
											return torch.Tensor([0.0]		)
									SCREAMING_SNAKE_CASE__								=   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   _snake_case	(				self     :Any		) -> str:



									"""simple docstring"""



									return self.mean






import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import ClassLabel, Features, Value
from .base import TaskTemplate




@dataclass(frozen=UpperCamelCase__     )
class 	UpperCamelCase_							(  UpperCamelCase__     ):
							# `task` is not a ClassVar since we want it to be part of the `asdict` output for JSON serialization
							lowerCamelCase_							     =       field(default="text-classification"					,	metadata={"include_in_asdict_even_if_is_default": True}     )
							lowerCamelCase_							     =       Features({"text": Value("string"     )}     )
							lowerCamelCase_							     =       Features({"labels": ClassLabel}     )
							lowerCamelCase_							     =       "text"
							lowerCamelCase_							     =       "labels"


							def   _snake_case	(				self     :Any				,					__A     :Dict		) -> Optional[Any]:



									"""simple docstring"""



									if self.label_column not in features:
											raise ValueError(f'''Column {self.label_column} is not present in features.'''		)
									if not isinstance(features[self.label_column]				,					__A		):
											raise ValueError(f'''Column {self.label_column} is not a ClassLabel.'''		)
									SCREAMING_SNAKE_CASE__								=   copy.deepcopy(self		)
									SCREAMING_SNAKE_CASE__								=   self.label_schema.copy()
									SCREAMING_SNAKE_CASE__								=   features[self.label_column]
									SCREAMING_SNAKE_CASE__								=   label_schema
									return task_template


							@property
							def   _snake_case	(				self     :str		) -> Dict[str, str]:



									"""simple docstring"""



									return {
									    self.text_column: "text",
									    self.label_column: "labels",
									}