





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		(      __SCREAMING_SNAKE_CASE       ):







 '''simple docstring'''
 lowercase_						=     42



class   _SCREAMING_SNAKE_CASE		(      nn.Module       ):







 '''simple docstring'''

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



   '''simple docstring'''



   super().__init__()
   lowerCamelCase__: str         =layers_per_block

   lowerCamelCase__: int         =torch.nn.Convad(
       UpperCAmelCase_			,					block_out_channels[0]			,					kernel_size=3			,					stride=1			,					padding=1			,					)

   lowerCamelCase__: Union[str, Any]         =None
   lowerCamelCase__: Any         =nn.ModuleList([])

   # down
   lowerCamelCase__: Dict         =block_out_channels[0]
   for i, down_block_type in enumerate(UpperCAmelCase_):
     lowerCamelCase__: Optional[Any]         =output_channel
     lowerCamelCase__: Optional[int]         =block_out_channels[i]
     lowerCamelCase__: Tuple         =i == len(UpperCAmelCase_) - 1

     lowerCamelCase__: int         =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__: 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"			,					attention_head_dim=block_out_channels[-1]			,					resnet_groups=UpperCAmelCase_			,					temb_channels=UpperCAmelCase_			,					)

   # out
   lowerCamelCase__: str         =nn.GroupNorm(num_channels=block_out_channels[-1]			,					num_groups=UpperCAmelCase_			,					eps=1E-6)
   lowerCamelCase__: Union[str, Any]         =nn.SiLU()

   lowerCamelCase__: str         =2 * out_channels if double_z else out_channels
   lowerCamelCase__: int         =nn.Convad(block_out_channels[-1]			,					UpperCAmelCase_			,					3			,					padding=1)

   lowerCamelCase__: Optional[int]         =False







 def 				SCREAMING_SNAKE_CASE_   (self							:  int			,					UpperCAmelCase_							:  List[Any])			->int:



   '''simple docstring'''



   lowerCamelCase__: List[str]         =x
   lowerCamelCase__: Tuple         =self.conv_in(UpperCAmelCase_)

   if self.training and self.gradient_checkpointing:

     def create_custom_forward(UpperCAmelCase_							:  str):
       def custom_forward(*UpperCAmelCase_							:  Optional[Any]):
         return module(*UpperCAmelCase_)

       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(UpperCAmelCase_)			,					UpperCAmelCase_			,					use_reentrant=UpperCAmelCase_)
       # middle
       lowerCamelCase__: Union[str, Any]         =torch.utils.checkpoint.checkpoint(
           create_custom_forward(self.mid_block)			,					UpperCAmelCase_			,					use_reentrant=UpperCAmelCase_)
     else:
       for down_block in self.down_blocks:
         lowerCamelCase__: List[Any]         =torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCAmelCase_)			,					UpperCAmelCase_)
       # middle
       lowerCamelCase__: Optional[int]         =torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block)			,					UpperCAmelCase_)

   else:
     # down
     for down_block in self.down_blocks:
       lowerCamelCase__: Optional[int]         =down_block(UpperCAmelCase_)

     # middle
     lowerCamelCase__: Any         =self.mid_block(UpperCAmelCase_)

   # post-process
   lowerCamelCase__: List[str]         =self.conv_norm_out(UpperCAmelCase_)
   lowerCamelCase__: Tuple         =self.conv_act(UpperCAmelCase_)
   lowerCamelCase__: Union[str, Any]         =self.conv_out(UpperCAmelCase_)

   return sample



class   _SCREAMING_SNAKE_CASE		(      nn.Module       ):







 '''simple docstring'''

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



   '''simple docstring'''



   super().__init__()
   lowerCamelCase__: Optional[int]         =layers_per_block

   lowerCamelCase__: str         =nn.Convad(
       UpperCAmelCase_			,					block_out_channels[-1]			,					kernel_size=3			,					stride=1			,					padding=1			,					)

   lowerCamelCase__: List[Any]         =None
   lowerCamelCase__: Optional[Any]         =nn.ModuleList([])

   lowerCamelCase__: Any         =in_channels if norm_type == "spatial" else None

   # mid
   lowerCamelCase__: Optional[Any]         =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__: Optional[Any]         =list(reversed(UpperCAmelCase_))
   lowerCamelCase__: Union[str, Any]         =reversed_block_out_channels[0]
   for i, up_block_type in enumerate(UpperCAmelCase_):
     lowerCamelCase__: str         =output_channel
     lowerCamelCase__: Any         =reversed_block_out_channels[i]

     lowerCamelCase__: Dict         =i == len(UpperCAmelCase_) - 1

     lowerCamelCase__: Dict         =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__: Union[str, Any]         =output_channel

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

   lowerCamelCase__: Optional[int]         =False







 def 				SCREAMING_SNAKE_CASE_   (self							:  Union[str, Any]			,					UpperCAmelCase_							:  List[Any]			,					UpperCAmelCase_							:  str=None)			->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__: Tuple         =z
   lowerCamelCase__: List[Any]         =self.conv_in(UpperCAmelCase_)

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

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

       return custom_forward

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

       # up
       for up_block in self.up_blocks:
         lowerCamelCase__: Union[str, Any]         =torch.utils.checkpoint.checkpoint(
             create_custom_forward(UpperCAmelCase_)			,					UpperCAmelCase_			,					UpperCAmelCase_			,					use_reentrant=UpperCAmelCase_)
     else:
       # middle
       lowerCamelCase__: List[Any]         =torch.utils.checkpoint.checkpoint(
           create_custom_forward(self.mid_block)			,					UpperCAmelCase_			,					UpperCAmelCase_)
       lowerCamelCase__: Optional[Any]         =sample.to(UpperCAmelCase_)

       # up
       for up_block in self.up_blocks:
         lowerCamelCase__: Optional[Any]         =torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCAmelCase_)			,					UpperCAmelCase_			,					UpperCAmelCase_)
   else:
     # middle
     lowerCamelCase__: Optional[int]         =self.mid_block(UpperCAmelCase_			,					UpperCAmelCase_)
     lowerCamelCase__: Union[str, Any]         =sample.to(UpperCAmelCase_)

     # up
     for up_block in self.up_blocks:
       lowerCamelCase__: List[Any]         =up_block(UpperCAmelCase_			,					UpperCAmelCase_)

        # post-process
   if latent_embeds is None:
     lowerCamelCase__: Union[str, Any]         =self.conv_norm_out(UpperCAmelCase_)
   else:
     lowerCamelCase__: List[Any]         =self.conv_norm_out(UpperCAmelCase_			,					UpperCAmelCase_)
   lowerCamelCase__: Optional[Any]         =self.conv_act(UpperCAmelCase_)
   lowerCamelCase__: str         =self.conv_out(UpperCAmelCase_)

   return sample



class   _SCREAMING_SNAKE_CASE		(      nn.Module       ):







 '''simple docstring'''

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



   '''simple docstring'''



   super().__init__()
   lowerCamelCase__: str         =n_e
   lowerCamelCase__: str         =vq_embed_dim
   lowerCamelCase__: str         =beta
   lowerCamelCase__: List[str]         =legacy

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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Optional[Any]			,					UpperCAmelCase_							:  List[Any])			->List[Any]:



   '''simple docstring'''



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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Dict			,					UpperCAmelCase_							:  Union[str, Any])			->Optional[Any]:



   '''simple docstring'''



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

 def 				SCREAMING_SNAKE_CASE_   (self							:  str			,					UpperCAmelCase_							:  Optional[int])			->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__: Any         =z.permute(0			,					2			,					3			,					1).contiguous()
   lowerCamelCase__: Dict         =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__: Tuple         =torch.argmin(torch.cdist(UpperCAmelCase_			,					self.embedding.weight)			,					dim=1)

   lowerCamelCase__: List[str]         =self.embedding(UpperCAmelCase_).view(z.shape)
   lowerCamelCase__: Optional[int]         =None
   lowerCamelCase__: Dict         =None

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

   # preserve gradients
   lowerCamelCase__: Optional[int]         =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__: Union[str, Any]         =min_encoding_indices.reshape(z.shape[0]			,					-1)  # add batch axis
     lowerCamelCase__: Any         =self.remap_to_used(UpperCAmelCase_)
     lowerCamelCase__: Dict         =min_encoding_indices.reshape(-1			,					1)  # flatten

   if self.sane_index_shape:
     lowerCamelCase__: Any         =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							:  Tuple			,					UpperCAmelCase_							:  Optional[Any]			,					UpperCAmelCase_							:  int)			->int:



   '''simple docstring'''



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

   # get quantized latent vectors
   lowerCamelCase__: Dict         =self.embedding(UpperCAmelCase_)

   if shape is not None:
     lowerCamelCase__: Optional[Any]         =z_q.view(UpperCAmelCase_)
     # reshape back to match original input shape
     lowerCamelCase__: Union[str, Any]         =z_q.permute(0			,					3			,					1			,					2).contiguous()

   return z_q



class   _SCREAMING_SNAKE_CASE		(      __SCREAMING_SNAKE_CASE       ):







 '''simple docstring'''

 def __init__(self							:  Optional[Any]			,					UpperCAmelCase_							:  int			,					UpperCAmelCase_							:  Dict=False)			->Dict:



   '''simple docstring'''



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

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



   '''simple docstring'''



   lowerCamelCase__: str         =randn_tensor(
       self.mean.shape			,					generator=UpperCAmelCase_			,					device=self.parameters.device			,					dtype=self.parameters.dtype)
   lowerCamelCase__: str         =self.mean + self.std * sample
   return x

 def 				SCREAMING_SNAKE_CASE_   (self							:  Optional[int]			,					UpperCAmelCase_							:  List[Any]=None)			->int:



   '''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							:  Tuple			,					UpperCAmelCase_							:  Tuple			,					UpperCAmelCase_							:  Tuple=[1, 2, 3])			->Dict:



   '''simple docstring'''



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



   '''simple docstring'''



   return self.mean


'''simple docstring'''






from . import (
    albert,
    align,
    altclip,
    audio_spectrogram_transformer,
    auto,
    autoformer,
    bark,
    bart,
    barthez,
    bartpho,
    beit,
    bert,
    bert_generation,
    bert_japanese,
    bertweet,
    big_bird,
    bigbird_pegasus,
    biogpt,
    bit,
    blenderbot,
    blenderbot_small,
    blip,
    blip_a,
    bloom,
    bridgetower,
    byta,
    camembert,
    canine,
    chinese_clip,
    clap,
    clip,
    clipseg,
    codegen,
    conditional_detr,
    convbert,
    convnext,
    convnextva,
    cpm,
    cpmant,
    ctrl,
    cvt,
    dataavec,
    deberta,
    deberta_va,
    decision_transformer,
    deformable_detr,
    deit,
    deprecated,
    deta,
    detr,
    dialogpt,
    dinat,
    distilbert,
    dit,
    donut,
    dpr,
    dpt,
    efficientformer,
    efficientnet,
    electra,
    encodec,
    encoder_decoder,
    ernie,
    ernie_m,
    esm,
    falcon,
    flaubert,
    flava,
    fnet,
    focalnet,
    fsmt,
    funnel,
    git,
    glpn,
    gpta,
    gpt_bigcode,
    gpt_neo,
    gpt_neox,
    gpt_neox_japanese,
    gpt_swa,
    gptj,
    gptsan_japanese,
    graphormer,
    groupvit,
    herbert,
    hubert,
    ibert,
    imagegpt,
    informer,
    instructblip,
    jukebox,
    layoutlm,
    layoutlmva,
    layoutlmva,
    layoutxlm,
    led,
    levit,
    lilt,
    llama,
    longformer,
    longta,
    luke,
    lxmert,
    mam_aaa,
    marian,
    markuplm,
    maskaformer,
    maskformer,
    mbart,
    mbartaa,
    mega,
    megatron_bert,
    megatron_gpta,
    mgp_str,
    mluke,
    mobilebert,
    mobilenet_va,
    mobilenet_va,
    mobilevit,
    mobilevitva,
    mpnet,
    mra,
    mta,
    musicgen,
    mvp,
    nat,
    nezha,
    nllb,
    nllb_moe,
    nystromformer,
    oneformer,
    open_llama,
    openai,
    opt,
    owlvit,
    pegasus,
    pegasus_x,
    perceiver,
    phobert,
    pixastruct,
    plbart,
    poolformer,
    prophetnet,
    qdqbert,
    rag,
    realm,
    reformer,
    regnet,
    rembert,
    resnet,
    roberta,
    roberta_prelayernorm,
    roc_bert,
    roformer,
    rwkv,
    sam,
    segformer,
    sew,
    sew_d,
    speech_encoder_decoder,
    speech_to_text,
    speech_to_text_a,
    speechta,
    splinter,
    squeezebert,
    swiftformer,
    swin,
    swinasr,
    swinva,
    switch_transformers,
    ta,
    table_transformer,
    tapas,
    time_series_transformer,
    timesformer,
    timm_backbone,
    transfo_xl,
    trocr,
    tvlt,
    umta,
    unispeech,
    unispeech_sat,
    upernet,
    videomae,
    vilt,
    vision_encoder_decoder,
    vision_text_dual_encoder,
    visual_bert,
    vit,
    vit_hybrid,
    vit_mae,
    vit_msn,
    vivit,
    wavaveca,
    wavaveca_conformer,
    wavaveca_phoneme,
    wavaveca_with_lm,
    wavlm,
    whisper,
    x_clip,
    xglm,
    xlm,
    xlm_prophetnet,
    xlm_roberta,
    xlm_roberta_xl,
    xlnet,
    xmod,
    yolos,
    yoso,
)



