
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    __A(					UpperCAmelCase							):
					SCREAMING_SNAKE_CASE           =     42





class    __A(					nn.Module							):

					def __init__(						self       :  Tuple ,       __UpperCamelCase       :  List[Any]=3 ,       __UpperCamelCase       :  int=3 ,       __UpperCamelCase       :  Optional[int]=("DownEncoderBlock2D",) ,       __UpperCamelCase       :  Union[str, Any]=(6_4,) ,       __UpperCamelCase       :  Any=2 ,       __UpperCamelCase       :  Dict=3_2 ,       __UpperCamelCase       :  Union[str, Any]="silu" ,       __UpperCamelCase       :  Any=True ,       ):
												super().__init__()
												lowerCamelCase_					=       layers_per_block

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

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

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

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

												# mid
												lowerCamelCase_					=       UNetMidBlockaD(
												    in_channels=block_out_channels[-1] ,       resnet_eps=1E-6 ,       resnet_act_fn=__UpperCamelCase ,       output_scale_factor=1 ,       resnet_time_scale_shift="""default""" ,       attention_head_dim=block_out_channels[-1] ,       resnet_groups=__UpperCamelCase ,       temb_channels=__UpperCamelCase ,       )

												# out
												lowerCamelCase_					=       nn.GroupNorm(num_channels=block_out_channels[-1] ,       num_groups=__UpperCamelCase ,       eps=1E-6 )
												lowerCamelCase_					=       nn.SiLU()

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

												lowerCamelCase_					=       False



					def  lowercase__   (						self       :  Any ,       __UpperCamelCase       :  int ):
												lowerCamelCase_					=       x
												lowerCamelCase_					=       self.conv_in(__UpperCamelCase )

												if self.training and self.gradient_checkpointing:

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

																										return custom_forward

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

												else:
																			# down
																			for down_block in self.down_blocks:
																										lowerCamelCase_					=       down_block(__UpperCamelCase )

																			# middle
																			lowerCamelCase_					=       self.mid_block(__UpperCamelCase )

												# post-process
												lowerCamelCase_					=       self.conv_norm_out(__UpperCamelCase )
												lowerCamelCase_					=       self.conv_act(__UpperCamelCase )
												lowerCamelCase_					=       self.conv_out(__UpperCamelCase )

												return sample





class    __A(					nn.Module							):

					def __init__(						self       :  Any ,       __UpperCamelCase       :  List[Any]=3 ,       __UpperCamelCase       :  Optional[int]=3 ,       __UpperCamelCase       :  str=("UpDecoderBlock2D",) ,       __UpperCamelCase       :  Any=(6_4,) ,       __UpperCamelCase       :  Optional[int]=2 ,       __UpperCamelCase       :  Dict=3_2 ,       __UpperCamelCase       :  Dict="silu" ,       __UpperCamelCase       :  int="group" ,       ):
												super().__init__()
												lowerCamelCase_					=       layers_per_block

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

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

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

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

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

																			lowerCamelCase_					=       i == len(__UpperCamelCase ) - 1

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

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

												lowerCamelCase_					=       False



					def  lowercase__   (						self       :  str ,       __UpperCamelCase       :  Optional[int] ,       __UpperCamelCase       :  List[str]=None ):
												lowerCamelCase_					=       z
												lowerCamelCase_					=       self.conv_in(__UpperCamelCase )

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

																			def create_custom_forward(__UpperCamelCase       :  Dict ):
																										def custom_forward(*__UpperCamelCase       :  Optional[Any] ):
																																	return module(*__UpperCamelCase )

																										return custom_forward

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

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

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

																			# up
																			for up_block in self.up_blocks:
																										lowerCamelCase_					=       up_block(__UpperCamelCase ,       __UpperCamelCase )

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

												return sample





class    __A(					nn.Module							):

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

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

												lowerCamelCase_					=       sane_index_shape

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

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

					def  lowercase__   (						self       :  str ,       __UpperCamelCase       :  Union[str, Any] ):
												# reshape z -> (batch, height, width, channel) and flatten
												lowerCamelCase_					=       z.permute(0 ,       2 ,       3 ,       1 ).contiguous()
												lowerCamelCase_					=       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_					=       torch.argmin(torch.cdist(__UpperCamelCase ,       self.embedding.weight ) ,       dim=1 )

												lowerCamelCase_					=       self.embedding(__UpperCamelCase ).view(z.shape )
												lowerCamelCase_					=       None
												lowerCamelCase_					=       None

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

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

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

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

												if self.sane_index_shape:
																			lowerCamelCase_					=       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       :  Any ,       __UpperCamelCase       :  Optional[int] ,       __UpperCamelCase       :  Any ):
												# shape specifying (batch, height, width, channel)
												if self.remap is not None:
																			lowerCamelCase_					=       indices.reshape(shape[0] ,       -1 )  # add batch axis
																			lowerCamelCase_					=       self.unmap_to_all(__UpperCamelCase )
																			lowerCamelCase_					=       indices.reshape(-1 )  # flatten again

												# get quantized latent vectors
												lowerCamelCase_					=       self.embedding(__UpperCamelCase )

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

												return z_q





class    __A(					UpperCAmelCase							):

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

					def  lowercase__   (						self       :  List[str] ,       __UpperCamelCase       :  Optional[torch.Generator] = None ):
												# make sure sample is on the same device as the parameters and has same dtype
												lowerCamelCase_					=       randn_tensor(
												    self.mean.shape ,       generator=__UpperCamelCase ,       device=self.parameters.device ,       dtype=self.parameters.dtype )
												lowerCamelCase_					=       self.mean + self.std * sample
												return x

					def  lowercase__   (						self       :  str ,       __UpperCamelCase       :  Dict=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       :  int ,       __UpperCamelCase       :  int ,       __UpperCamelCase       :  str=[1, 2, 3] ):
												if self.deterministic:
																			return torch.Tensor([0.0] )
												lowerCamelCase_					=       np.log(2.0 * np.pi )
												return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean ,       2 ) / self.var ,       dim=__UpperCamelCase )



					def  lowercase__   (						self       :  Optional[int] ):
												return self.mean







import importlib.metadata
import operator
import re
import sys
from typing import Optional

from packaging import version


lowercase   =	{
    '''<''': operator.lt,
    '''<=''': operator.le,
    '''==''': operator.eq,
    '''!=''': operator.ne,
    '''>=''': operator.ge,
    '''>''': operator.gt,
}



def   __lowerCAmelCase    (				UpperCAmelCase__					:   Tuple      ,					UpperCAmelCase__					:   int      ,					UpperCAmelCase__					:   Optional[int]      ,					UpperCAmelCase__					:   Optional[int]      ,					UpperCAmelCase__					:   Optional[int]      ,					UpperCAmelCase__					:   int							)	->     Union[str, Any]:
							if got_ver is None or want_ver is None:
														raise ValueError(
														    F'''Unable to compare versions for {requirement}: need={want_ver} found={got_ver}. This is unusual. Consider'''
														    F''' reinstalling {pkg}.'''							)
							if not ops[op](version.parse(UpperCAmelCase__							)      ,					version.parse(UpperCAmelCase__							)							):
														raise ImportError(
														    F'''{requirement} is required for a normal functioning of this module, but found {pkg}=={got_ver}.{hint}'''							)



def   __lowerCAmelCase    (				UpperCAmelCase__					:   str      ,					UpperCAmelCase__					:   Optional[str] = None							)	->     None:
							lowerCamelCase_					=       F'''\n{hint}''' if hint is not None else """"""

							# non-versioned check
							if re.match(R"""^[\w_\-\d]+$"""      ,					UpperCAmelCase__							):
														lowerCamelCase_	, lowerCamelCase_	, lowerCamelCase_					=       requirement, None, None
							else:
														lowerCamelCase_					=       re.findall(R"""^([^!=<>\s]+)([\s!=<>]{1,2}.+)"""      ,					UpperCAmelCase__							)
														if not match:
																					raise ValueError(
																					    """requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23, but"""
																					    F''' got {requirement}'''							)
														lowerCamelCase_	, lowerCamelCase_					=       match[0]
														lowerCamelCase_					=       want_full.split(""","""							)  # there could be multiple requirements
														lowerCamelCase_					=       {}
														for w in want_range:
																					lowerCamelCase_					=       re.findall(R"""^([\s!=<>]{1,2})(.+)"""      ,					UpperCAmelCase__							)
																					if not match:
																												raise ValueError(
																												    """requirement needs to be in the pip package format, .e.g., package_a==1.23, or package_b>=1.23,"""
																												    F''' but got {requirement}'''							)
																					lowerCamelCase_	, lowerCamelCase_					=       match[0]
																					lowerCamelCase_					=       want_ver
																					if op not in ops:
																												raise ValueError(F'''{requirement}: need one of {list(ops.keys()							)}, but got {op}'''							)

    # special case
							if pkg == "python":
														lowerCamelCase_					=       """.""".join([str(UpperCAmelCase__							) for x in sys.version_info[:3]]							)
														for op, want_ver in wanted.items():
																					_compare_versions(UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__							)
														return

							# check if any version is installed
							try:
														lowerCamelCase_					=       importlib.metadata.version(UpperCAmelCase__							)
							except importlib.metadata.PackageNotFoundError:
														raise importlib.metadata.PackageNotFoundError(
														    F'''The \'{requirement}\' distribution was not found and is required by this application. {hint}'''							)

							# check that the right version is installed if version number or a range was provided
							if want_ver is not None:
														for op, want_ver in wanted.items():
																					_compare_versions(UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__      ,					UpperCAmelCase__							)



def   __lowerCAmelCase    (				UpperCAmelCase__					:   List[str]							)	->     Dict:
							lowerCamelCase_					=       """Try: pip install transformers -U or pip install -e '.[dev]' if you're working with git main"""
							return require_version(UpperCAmelCase__      ,					UpperCAmelCase__							)





