





'''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       UpperCamelCase__  (		__lowerCAmelCase     ):
	lowerCAmelCase__					:			torch.FloatTensor






class       UpperCamelCase__  (		nn.Module     ):







	def __init__(				self			:   List[str] , lowerCamelCase			:   int=3 , lowerCamelCase			:   Dict=3 , lowerCamelCase			:   int=("DownEncoderBlock2D",) , lowerCamelCase			:   Tuple=(6_4,) , lowerCamelCase			:   List[Any]=2 , lowerCamelCase			:   List[Any]=3_2 , lowerCamelCase			:   Tuple="silu" , lowerCamelCase			:   str=True , ):




				'''simple docstring'''
				super().__init__()
				a__     = layers_per_block

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

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

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

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

				# mid
				a__     = UNetMidBlockaD(
				    in_channels=block_out_channels[-1] , resnet_eps=1e-6 , resnet_act_fn=lowerCamelCase , output_scale_factor=1 , resnet_time_scale_shift="default" , attention_head_dim=block_out_channels[-1] , resnet_groups=lowerCamelCase , temb_channels=lowerCamelCase , )

				# out
				a__     = nn.GroupNorm(num_channels=block_out_channels[-1] , num_groups=lowerCamelCase , eps=1e-6		)
				a__     = nn.SiLU()

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

				a__     = False







	def 						__a   (				self			:   Optional[Any] , lowerCamelCase			:   Optional[int]		):




				'''simple docstring'''
				a__     = x
				a__     = self.conv_in(lowerCamelCase		)

				if self.training and self.gradient_checkpointing:

							def create_custom_forward(lowerCamelCase			:   Optional[int]		):
										def custom_forward(*lowerCamelCase			:   Union[str, Any]		):
													return module(*lowerCamelCase		)

										return custom_forward

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

				else:
							# down
							for down_block in self.down_blocks:
										a__     = down_block(lowerCamelCase		)

							# middle
							a__     = self.mid_block(lowerCamelCase		)

				# post-process
				a__     = self.conv_norm_out(lowerCamelCase		)
				a__     = self.conv_act(lowerCamelCase		)
				a__     = self.conv_out(lowerCamelCase		)

				return sample






class       UpperCamelCase__  (		nn.Module     ):







	def __init__(				self			:   str , lowerCamelCase			:   Dict=3 , lowerCamelCase			:   List[Any]=3 , lowerCamelCase			:   Optional[int]=("UpDecoderBlock2D",) , lowerCamelCase			:   Tuple=(6_4,) , lowerCamelCase			:   str=2 , lowerCamelCase			:   Union[str, Any]=3_2 , lowerCamelCase			:   int="silu" , lowerCamelCase			:   Union[str, Any]="group" , ):




				'''simple docstring'''
				super().__init__()
				a__     = layers_per_block

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

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

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

				# mid
				a__     = UNetMidBlockaD(
				    in_channels=block_out_channels[-1] , resnet_eps=1e-6 , resnet_act_fn=lowerCamelCase , 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=lowerCamelCase , temb_channels=lowerCamelCase , )

				# up
				a__     = list(reversed(lowerCamelCase		)		)
				a__     = reversed_block_out_channels[0]
				for i, up_block_type in enumerate(lowerCamelCase		):
							a__     = output_channel
							a__     = reversed_block_out_channels[i]

							a__     = i == len(lowerCamelCase		) - 1

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

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

				a__     = False







	def 						__a   (				self			:   Tuple , lowerCamelCase			:   int , lowerCamelCase			:   int=None		):




				'''simple docstring'''
				a__     = z
				a__     = self.conv_in(lowerCamelCase		)

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

							def create_custom_forward(lowerCamelCase			:   List[Any]		):
										def custom_forward(*lowerCamelCase			:   Optional[int]		):
													return module(*lowerCamelCase		)

										return custom_forward

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

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

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

							# up
							for up_block in self.up_blocks:
										a__     = up_block(lowerCamelCase , lowerCamelCase		)

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

				return sample






class       UpperCamelCase__  (		nn.Module     ):







	def __init__(				self			:   List[str] , lowerCamelCase			:   int , lowerCamelCase			:   List[Any] , lowerCamelCase			:   List[str] , lowerCamelCase			:   Tuple=None , lowerCamelCase			:   Tuple="random" , lowerCamelCase			:   int=False , lowerCamelCase			:   List[str]=True		):




				'''simple docstring'''
				super().__init__()
				a__     = n_e
				a__     = vq_embed_dim
				a__     = beta
				a__     = legacy

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

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

				a__     = sane_index_shape







	def 						__a   (				self			:   int , lowerCamelCase			:   int		):




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







	def 						__a   (				self			:   List[Any] , lowerCamelCase			:   Tuple		):




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







	def 						__a   (				self			:   Optional[Any] , lowerCamelCase			:   List[str]		):




				'''simple docstring'''
				# reshape z -> (batch, height, width, channel) and flatten
				a__     = z.permute(0 , 2 , 3 , 1		).contiguous()
				a__     = z.view(-1 , self.vq_embed_dim		)

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

				a__     = self.embedding(lowerCamelCase		).view(z.shape		)
				a__     = None
				a__     = None

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

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

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

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

				if self.sane_index_shape:
							a__     = 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 						__a   (				self			:   Tuple , lowerCamelCase			:   Dict , lowerCamelCase			:   str		):




				'''simple docstring'''
				# shape specifying (batch, height, width, channel)
				if self.remap is not None:
							a__     = indices.reshape(shape[0] , -1		)  # add batch axis
							a__     = self.unmap_to_all(lowerCamelCase		)
							a__     = indices.reshape(-1		)  # flatten again

				# get quantized latent vectors
				a__     = self.embedding(lowerCamelCase		)

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

				return z_q






class       UpperCamelCase__  (		__lowerCAmelCase     ):







	def __init__(				self			:   Dict , lowerCamelCase			:   List[str] , lowerCamelCase			:   Dict=False		):




				'''simple docstring'''
				a__     = parameters
				a__   ,				a__     = torch.chunk(lowerCamelCase , 2 , dim=1		)
				a__     = torch.clamp(self.logvar , -30.0 , 20.0		)
				a__     = deterministic
				a__     = torch.exp(0.5 * self.logvar		)
				a__     = torch.exp(self.logvar		)
				if self.deterministic:
							a__     = a__     = torch.zeros_like(
							    self.mean , device=self.parameters.device , dtype=self.parameters.dtype		)







	def 						__a   (				self			:   Any , lowerCamelCase			:   Optional[torch.Generator] = None		):




				'''simple docstring'''
				# make sure sample is on the same device as the parameters and has same dtype
				a__     = randn_tensor(
				    self.mean.shape , generator=lowerCamelCase , device=self.parameters.device , dtype=self.parameters.dtype		)
				a__     = self.mean + self.std * sample
				return x







	def 						__a   (				self			:   List[str] , lowerCamelCase			:   int=None		):




				'''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 						__a   (				self			:   Dict , lowerCamelCase			:   List[str] , lowerCamelCase			:   List[Any]=[1, 2, 3]		):




				'''simple docstring'''
				if self.deterministic:
							return torch.Tensor([0.0]		)
				a__     = np.log(2.0 * np.pi		)
				return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean , 2		) / self.var , dim=lowerCamelCase		)







	def 						__a   (				self			:   Dict		):




				'''simple docstring'''
				return self.mean










'''simple docstring'''


def    _lowerCamelCase  (__lowerCamelCase  :     list[list[float]]     )   ->    list[list[float]]:
			a__     = []
			for data in source_data:
						for i, el in enumerate(__lowerCamelCase     ):
									if len(__lowerCamelCase     ) < i + 1:
												data_lists.append([]     )
									data_lists[i].append(float(__lowerCamelCase     )     )
			return data_lists







def    _lowerCamelCase  (__lowerCamelCase  :     list[list[float]]					,  __lowerCamelCase  :     list[int]     )   ->    list[list[float]]:
			a__     = []
			for dlist, weight in zip(__lowerCamelCase					,  __lowerCamelCase     ):
						a__     = min(__lowerCamelCase     )
						a__     = max(__lowerCamelCase     )

						a__     = []
						# for weight 0 score is 1 - actual score
						if weight == 0:
									for item in dlist:
												try:
															score.append(1 - ((item - mind) / (maxd - mind))     )
												except ZeroDivisionError:
															score.append(1     )

						elif weight == 1:
									for item in dlist:
												try:
															score.append((item - mind) / (maxd - mind)     )
												except ZeroDivisionError:
															score.append(0     )

        # weight not 0 or 1
						else:
									a__     = f'''Invalid weight of {weight:f} provided'''
									raise ValueError(__lowerCamelCase     )

						score_lists.append(__lowerCamelCase     )

			return score_lists







def    _lowerCamelCase  (__lowerCamelCase  :     list[list[float]]     )   ->    list[float]:
			a__     = [0 for i in range(len(score_lists[0]     )     )]

			for slist in score_lists:
						for j, ele in enumerate(__lowerCamelCase     ):
									a__     = final_scores[j] + ele

			return final_scores







def    _lowerCamelCase  (__lowerCamelCase  :     list[list[float]]					,  __lowerCamelCase  :     list[int]     )   ->    list[list[float]]:
			a__     = get_data(__lowerCamelCase     )
			a__     = calculate_each_score(__lowerCamelCase					,  __lowerCamelCase     )
			a__     = generate_final_scores(__lowerCamelCase     )

			# append scores to source data
			for i, ele in enumerate(__lowerCamelCase     ):
						source_data[i].append(__lowerCamelCase     )

			return source_data



