





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     __magic_name__     (							__A							):





     """simple docstring"""



     __UpperCamelCase    			=   42




class     __magic_name__     (							nn.Module							):





     """simple docstring"""



     def __init__(							self				:str	,					snake_case				:List[Any]=3	,					snake_case				:List[str]=3	,					snake_case				:int=("DownEncoderBlock2D",)	,					snake_case				:str=(64,)	,					snake_case				:List[Any]=2	,					snake_case				:List[str]=32	,					snake_case				:Union[str, Any]="silu"	,					snake_case				:Any=True	,					):





        '''simple docstring'''




        super().__init__()
        A_    : List[str]								=	layers_per_block

        A_    : Tuple								=	torch.nn.Convad(
            __UpperCAmelCase	,					block_out_channels[0]	,					kernel_size=3	,					stride=1	,					padding=1	,					)

        A_    : Any								=	None
        A_    : Tuple								=	nn.ModuleList([]     )

        # down
        A_    : Optional[Any]								=	block_out_channels[0]
        for i, down_block_type in enumerate(__UpperCAmelCase     ):
           A_    : List[str]								=	output_channel
           A_    : str								=	block_out_channels[i]
           A_    : Dict								=	i == len(__UpperCAmelCase     ) - 1

           A_    : str								=	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
        A_    : Tuple								=	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
        A_    : List[str]								=	nn.GroupNorm(num_channels=block_out_channels[-1]	,					num_groups=__UpperCAmelCase	,					eps=1e-6     )
        A_    : int								=	nn.SiLU()

        A_    : List[Any]								=	2 * out_channels if double_z else out_channels
        A_    : int								=	nn.Convad(block_out_channels[-1]	,					__UpperCAmelCase	,					3	,					padding=1     )

        A_    : Any								=	False





     def   SCREAMING_SNAKE_CASE						(							self				:str	,					snake_case				:List[str]     ):





        '''simple docstring'''




        A_    : Dict								=	x
        A_    : Optional[int]								=	self.conv_in(__UpperCAmelCase     )

        if self.training and self.gradient_checkpointing:

           def create_custom_forward(snake_case				:int     ):
              def custom_forward(*snake_case				:int     ):
                 return module(*__UpperCAmelCase     )

              return custom_forward

           # down
           if is_torch_version(">="	,					"1.11.0"     ):
              for down_block in self.down_blocks:
                 A_    : Optional[int]								=	torch.utils.checkpoint.checkpoint(
                     create_custom_forward(__UpperCAmelCase     )	,					__UpperCAmelCase	,					use_reentrant=__UpperCAmelCase     )
              # middle
              A_    : Optional[Any]								=	torch.utils.checkpoint.checkpoint(
                  create_custom_forward(self.mid_block     )	,					__UpperCAmelCase	,					use_reentrant=__UpperCAmelCase     )
           else:
              for down_block in self.down_blocks:
                 A_    : Union[str, Any]								=	torch.utils.checkpoint.checkpoint(create_custom_forward(__UpperCAmelCase     )	,					__UpperCAmelCase     )
              # middle
              A_    : Optional[Any]								=	torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block     )	,					__UpperCAmelCase     )

        else:
           # down
           for down_block in self.down_blocks:
              A_    : Union[str, Any]								=	down_block(__UpperCAmelCase     )

           # middle
           A_    : int								=	self.mid_block(__UpperCAmelCase     )

        # post-process
        A_    : str								=	self.conv_norm_out(__UpperCAmelCase     )
        A_    : Optional[int]								=	self.conv_act(__UpperCAmelCase     )
        A_    : List[str]								=	self.conv_out(__UpperCAmelCase     )

        return sample




class     __magic_name__     (							nn.Module							):





     """simple docstring"""



     def __init__(							self				:str	,					snake_case				:Union[str, Any]=3	,					snake_case				:int=3	,					snake_case				:Dict=("UpDecoderBlock2D",)	,					snake_case				:List[Any]=(64,)	,					snake_case				:Any=2	,					snake_case				:str=32	,					snake_case				:int="silu"	,					snake_case				:Optional[Any]="group"	,					):





        '''simple docstring'''




        super().__init__()
        A_    : Any								=	layers_per_block

        A_    : int								=	nn.Convad(
            __UpperCAmelCase	,					block_out_channels[-1]	,					kernel_size=3	,					stride=1	,					padding=1	,					)

        A_    : str								=	None
        A_    : int								=	nn.ModuleList([]     )

        A_    : Optional[int]								=	in_channels if norm_type == "spatial" else None

        # mid
        A_    : int								=	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
        A_    : Optional[int]								=	list(reversed(__UpperCAmelCase     )     )
        A_    : Tuple								=	reversed_block_out_channels[0]
        for i, up_block_type in enumerate(__UpperCAmelCase     ):
           A_    : str								=	output_channel
           A_    : Any								=	reversed_block_out_channels[i]

           A_    : Dict								=	i == len(__UpperCAmelCase     ) - 1

           A_    : Any								=	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     )
           A_    : Tuple								=	output_channel

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

        A_    : List[Any]								=	False





     def   SCREAMING_SNAKE_CASE						(							self				:List[Any]	,					snake_case				:int	,					snake_case				:int=None     ):





        '''simple docstring'''




        A_    : List[Any]								=	z
        A_    : List[Any]								=	self.conv_in(__UpperCAmelCase     )

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

           def create_custom_forward(snake_case				:Union[str, Any]     ):
              def custom_forward(*snake_case				:List[str]     ):
                 return module(*__UpperCAmelCase     )

              return custom_forward

           if is_torch_version(">="	,					"1.11.0"     ):
              # middle
              A_    : Union[str, Any]								=	torch.utils.checkpoint.checkpoint(
                  create_custom_forward(self.mid_block     )	,					__UpperCAmelCase	,					__UpperCAmelCase	,					use_reentrant=__UpperCAmelCase     )
              A_    : List[str]								=	sample.to(__UpperCAmelCase     )

              # up
              for up_block in self.up_blocks:
                 A_    : str								=	torch.utils.checkpoint.checkpoint(
                     create_custom_forward(__UpperCAmelCase     )	,					__UpperCAmelCase	,					__UpperCAmelCase	,					use_reentrant=__UpperCAmelCase     )
           else:
              # middle
              A_    : Dict								=	torch.utils.checkpoint.checkpoint(
                  create_custom_forward(self.mid_block     )	,					__UpperCAmelCase	,					__UpperCAmelCase     )
              A_    : int								=	sample.to(__UpperCAmelCase     )

              # up
              for up_block in self.up_blocks:
                 A_    : str								=	torch.utils.checkpoint.checkpoint(create_custom_forward(__UpperCAmelCase     )	,					__UpperCAmelCase	,					__UpperCAmelCase     )
        else:
           # middle
           A_    : Tuple								=	self.mid_block(__UpperCAmelCase	,					__UpperCAmelCase     )
           A_    : Optional[Any]								=	sample.to(__UpperCAmelCase     )

           # up
           for up_block in self.up_blocks:
              A_    : int								=	up_block(__UpperCAmelCase	,					__UpperCAmelCase     )

        # post-process
        if latent_embeds is None:
           A_    : Union[str, Any]								=	self.conv_norm_out(__UpperCAmelCase     )
        else:
           A_    : str								=	self.conv_norm_out(__UpperCAmelCase	,					__UpperCAmelCase     )
        A_    : int								=	self.conv_act(__UpperCAmelCase     )
        A_    : str								=	self.conv_out(__UpperCAmelCase     )

        return sample




class     __magic_name__     (							nn.Module							):





     """simple docstring"""



     def __init__(							self				:int	,					snake_case				:Any	,					snake_case				:Union[str, Any]	,					snake_case				:List[Any]	,					snake_case				:Optional[int]=None	,					snake_case				:Dict="random"	,					snake_case				:Optional[int]=False	,					snake_case				:int=True     ):





        '''simple docstring'''




        super().__init__()
        A_    : Union[str, Any]								=	n_e
        A_    : str								=	vq_embed_dim
        A_    : Union[str, Any]								=	beta
        A_    : Dict								=	legacy

        A_    : 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     )

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

        A_    : Tuple								=	sane_index_shape



     def   SCREAMING_SNAKE_CASE						(							self				:Union[str, Any]	,					snake_case				:Union[str, Any]     ):





        '''simple docstring'''




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



     def   SCREAMING_SNAKE_CASE						(							self				:List[Any]	,					snake_case				:Tuple     ):





        '''simple docstring'''




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



     def   SCREAMING_SNAKE_CASE						(							self				:Optional[Any]	,					snake_case				:Union[str, Any]     ):





        '''simple docstring'''




        A_    : Tuple								=	z.permute(0	,					2	,					3	,					1     ).contiguous()
        A_    : 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
        A_    : str								=	torch.argmin(torch.cdist(__UpperCAmelCase	,					self.embedding.weight     )	,					dim=1     )

        A_    : Dict								=	self.embedding(__UpperCAmelCase     ).view(z.shape     )
        A_    : int								=	None
        A_    : Tuple								=	None

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

        # preserve gradients
        A_    : Any								=	z + (z_q - z).detach()

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

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

        if self.sane_index_shape:
           A_    : str								=	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   SCREAMING_SNAKE_CASE						(							self				:List[str]	,					snake_case				:Tuple	,					snake_case				:Dict     ):





        '''simple docstring'''




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

        # get quantized latent vectors
        A_    : Dict								=	self.embedding(__UpperCAmelCase     )

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

        return z_q




class     __magic_name__     (							__A							):





     """simple docstring"""



     def __init__(							self				:Optional[Any]	,					snake_case				:str	,					snake_case				:Optional[int]=False     ):





        '''simple docstring'''




        A_    : Union[str, Any]								=	parameters
        A_					,  A_    : Any								=	torch.chunk(__UpperCAmelCase	,					2	,					dim=1     )
        A_    : Union[str, Any]								=	torch.clamp(self.logvar	,					-30.0	,					20.0     )
        A_    : int								=	deterministic
        A_    : Dict								=	torch.exp(0.5 * self.logvar     )
        A_    : List[str]								=	torch.exp(self.logvar     )
        if self.deterministic:
           A_    : int								=	torch.zeros_like(
               self.mean	,					device=self.parameters.device	,					dtype=self.parameters.dtype     )



     def   SCREAMING_SNAKE_CASE						(							self				:Optional[int]	,					snake_case				:str = None     ):





        '''simple docstring'''




        A_    : Optional[int]								=	randn_tensor(
            self.mean.shape	,					generator=__UpperCAmelCase	,					device=self.parameters.device	,					dtype=self.parameters.dtype     )
        A_    : Optional[Any]								=	self.mean + self.std * sample
        return x



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





        '''simple docstring'''




        if self.deterministic:
           return torch.Tensor([0.0]     )
        A_    : Dict								=	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   SCREAMING_SNAKE_CASE						(							self				:Union[str, Any]     ):





        '''simple docstring'''




        return self.mean

"""simple docstring"""




import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
				from deepspeed import DeepSpeedEngine

if is_tpu_available(check_device=False):
				import torch_xla.core.xla_model as xm






def      UpperCAmelCase    (		UpperCamelCase__ ):






							"""simple docstring"""





							if is_torch_version('<'				,			'2.0.0' ) or not hasattr(UpperCamelCase__				,			'_dynamo' ):
														return False
							return isinstance(UpperCamelCase__				,			torch._dynamo.eval_frame.OptimizedModule )






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ = True ):






							"""simple docstring"""





							A__		   =							(torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

							A__		   =							is_compiled_module(UpperCamelCase__ )
							if is_compiled:
														A__		   =							model
														A__		   =							model._orig_mod

							if is_deepspeed_available():
														options += (DeepSpeedEngine,)

							while isinstance(UpperCamelCase__				,			UpperCamelCase__ ):
														A__		   =							model.module

							if not keep_fpaa_wrapper:
														A__		   =							getattr(UpperCamelCase__				,			'forward' )
														A__		   =							model.__dict__.pop('_original_forward'				,			UpperCamelCase__ )
														if original_forward is not None:
																					while hasattr(UpperCamelCase__				,			'__wrapped__' ):
																												A__		   =							forward.__wrapped__
																												if forward == original_forward:
																																			break
																					A__		   =							forward
														if getattr(UpperCamelCase__				,			'_converted_to_transformer_engine'				,			UpperCamelCase__ ):
																					convert_model(UpperCamelCase__				,			to_transformer_engine=UpperCamelCase__ )

							if is_compiled:
														A__		   =							model
														A__		   =							compiled_model

							return model






def      UpperCAmelCase    (		):






							"""simple docstring"""





							PartialState().wait_for_everyone()






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





							if PartialState().distributed_type == DistributedType.TPU:
														xm.save(UpperCamelCase__				,			UpperCamelCase__ )
							elif PartialState().local_process_index == 0:
														torch.save(UpperCamelCase__				,			UpperCamelCase__ )






@contextmanager
def      UpperCAmelCase    (		**UpperCamelCase__ ):






							"""simple docstring"""





							for key, value in kwargs.items():
														A__		   =							str(UpperCamelCase__ )

							yield

							for key in kwargs:
														if key.upper() in os.environ:
																					del os.environ[key.upper()]






def      UpperCAmelCase    (		UpperCamelCase__ ):






							"""simple docstring"""





							if not hasattr(UpperCamelCase__				,			'__qualname__' ) and not hasattr(UpperCamelCase__				,			'__name__' ):
														A__		   =							getattr(UpperCamelCase__				,			'__class__'				,			UpperCamelCase__ )
							if hasattr(UpperCamelCase__				,			'__qualname__' ):
														return obj.__qualname__
							if hasattr(UpperCamelCase__				,			'__name__' ):
														return obj.__name__
							return str(UpperCamelCase__ )






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





							for key, value in source.items():
														if isinstance(UpperCamelCase__				,			UpperCamelCase__ ):
																					A__		   =							destination.setdefault(UpperCamelCase__				,			{} )
																					merge_dicts(UpperCamelCase__				,			UpperCamelCase__ )
														else:
																					A__		   =							value

							return destination






def      UpperCAmelCase    (		UpperCamelCase__ = None ):






							"""simple docstring"""





							if port is None:
														A__		   =							29_500
							with socket.socket(socket.AF_INET				,			socket.SOCK_STREAM ) as s:
														return s.connect_ex(('localhost', port) ) == 0





