



from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCamelCase_	   =			{
    'configuration_efficientnet': [
        'EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'EfficientNetConfig',
        'EfficientNetOnnxConfig',
    ]
}

try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCamelCase_	   =			['EfficientNetImageProcessor']

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCamelCase_	   =			[
           'EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST',
           'EfficientNetForImageClassification',
           'EfficientNetModel',
           'EfficientNetPreTrainedModel',
       ]

if TYPE_CHECKING:
       from .configuration_efficientnet import (
           EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
           EfficientNetConfig,
           EfficientNetOnnxConfig,
       )

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .image_processing_efficientnet import EfficientNetImageProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_efficientnet import (
                  EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST,
                  EfficientNetForImageClassification,
                  EfficientNetModel,
                  EfficientNetPreTrainedModel,
              )

else:
       import sys

       UpperCamelCase_	   =			_LazyModule(__name__, globals()['__file__'], _import_structure)





'''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_							(    A				):

 '''simple docstring'''
 lowercase_					:    torch.FloatTensor






class 							UpperCAmelCase_							(    nn.Module				):

 '''simple docstring'''







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



   '''simple docstring'''




   super().__init__()
   UpperCAmelCase__      :					Tuple				       =							layers_per_block

   UpperCAmelCase__      :					Optional[int]				       =							torch.nn.Convad(
       snake_case__			,				block_out_channels[0]			,				kernel_size=3			,				stride=1			,				padding=1			,				)

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

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

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

   # mid
   UpperCAmelCase__      :					Optional[Any]				       =							UNetMidBlockaD(
       in_channels=block_out_channels[-1]			,				resnet_eps=1e-6			,				resnet_act_fn=snake_case__			,				output_scale_factor=1			,				resnet_time_scale_shift="default"			,				attention_head_dim=block_out_channels[-1]			,				resnet_groups=snake_case__			,				temb_channels=snake_case__			,				)

   # out
   UpperCAmelCase__      :					Tuple				       =							nn.GroupNorm(num_channels=block_out_channels[-1]			,				num_groups=snake_case__			,				eps=1e-6 )
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.SiLU()

   UpperCAmelCase__      :					Dict				       =							2 * out_channels if double_z else out_channels
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.Convad(block_out_channels[-1]			,				snake_case__			,				3			,				padding=1 )

   UpperCAmelCase__      :					Union[str, Any]				       =							False







 def 							UpperCamelCase  (					self		:  int			,				snake_case__		:  Tuple ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[str]				       =							x
   UpperCAmelCase__      :					Dict				       =							self.conv_in(snake_case__ )

   if self.training and self.gradient_checkpointing:

     def create_custom_forward(snake_case__		:  Dict ):
       def custom_forward(*snake_case__		:  List[str] ):
         return module(*snake_case__ )

       return custom_forward

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

   else:
     # down
     for down_block in self.down_blocks:
       UpperCAmelCase__      :					Dict				       =							down_block(snake_case__ )

     # middle
     UpperCAmelCase__      :					Optional[Any]				       =							self.mid_block(snake_case__ )

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

   return sample






class 							UpperCAmelCase_							(    nn.Module				):

 '''simple docstring'''







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



   '''simple docstring'''




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

   UpperCAmelCase__      :					Any				       =							nn.Convad(
       snake_case__			,				block_out_channels[-1]			,				kernel_size=3			,				stride=1			,				padding=1			,				)

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

   UpperCAmelCase__      :					str				       =							in_channels if norm_type == "spatial" else None

   # mid
   UpperCAmelCase__      :					Optional[Any]				       =							UNetMidBlockaD(
       in_channels=block_out_channels[-1]			,				resnet_eps=1e-6			,				resnet_act_fn=snake_case__			,				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=snake_case__			,				temb_channels=snake_case__			,				)

   # up
   UpperCAmelCase__      :					Tuple				       =							list(reversed(snake_case__ ) )
   UpperCAmelCase__      :					Optional[Any]				       =							reversed_block_out_channels[0]
   for i, up_block_type in enumerate(snake_case__ ):
     UpperCAmelCase__      :					Dict				       =							output_channel
     UpperCAmelCase__      :					List[Any]				       =							reversed_block_out_channels[i]

     UpperCAmelCase__      :					List[Any]				       =							i == len(snake_case__ ) - 1

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

   # out
   if norm_type == "spatial":
     UpperCAmelCase__      :					Optional[Any]				       =							SpatialNorm(block_out_channels[0]			,				snake_case__ )
   else:
     UpperCAmelCase__      :					Optional[int]				       =							nn.GroupNorm(num_channels=block_out_channels[0]			,				num_groups=snake_case__			,				eps=1e-6 )
   UpperCAmelCase__      :					Dict				       =							nn.SiLU()
   UpperCAmelCase__      :					Union[str, Any]				       =							nn.Convad(block_out_channels[0]			,				snake_case__			,				3			,				padding=1 )

   UpperCAmelCase__      :					Union[str, Any]				       =							False







 def 							UpperCamelCase  (					self		:  List[str]			,				snake_case__		:  List[str]			,				snake_case__		:  Tuple=None ):



   '''simple docstring'''




   UpperCAmelCase__      :					str				       =							z
   UpperCAmelCase__      :					List[str]				       =							self.conv_in(snake_case__ )

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

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

       return custom_forward

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

       # up
       for up_block in self.up_blocks:
         UpperCAmelCase__      :					Tuple				       =							torch.utils.checkpoint.checkpoint(
             create_custom_forward(snake_case__ )			,				snake_case__			,				snake_case__			,				use_reentrant=snake_case__ )
     else:
       # middle
       UpperCAmelCase__      :					List[Any]				       =							torch.utils.checkpoint.checkpoint(
           create_custom_forward(self.mid_block )			,				snake_case__			,				snake_case__ )
       UpperCAmelCase__      :					int				       =							sample.to(snake_case__ )

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

     # up
     for up_block in self.up_blocks:
       UpperCAmelCase__      :					int				       =							up_block(snake_case__			,				snake_case__ )

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

   return sample






class 							UpperCAmelCase_							(    nn.Module				):

 '''simple docstring'''







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



   '''simple docstring'''




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

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

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

   UpperCAmelCase__      :					Dict				       =							sane_index_shape







 def 							UpperCamelCase  (					self		:  Tuple			,				snake_case__		:  int ):



   '''simple docstring'''




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







 def 							UpperCamelCase  (					self		:  int			,				snake_case__		:  str ):



   '''simple docstring'''




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







 def 							UpperCamelCase  (					self		:  Any			,				snake_case__		:  Any ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[str]				       =							z.permute(0			,				2			,				3			,				1 ).contiguous()
   UpperCAmelCase__      :					str				       =							z.view(-1			,				self.vq_embed_dim )

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

   UpperCAmelCase__      :					Any				       =							self.embedding(snake_case__ ).view(z.shape )
   UpperCAmelCase__      :					Union[str, Any]				       =							None
   UpperCAmelCase__      :					Any				       =							None

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

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

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

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

   if self.sane_index_shape:
     UpperCAmelCase__      :					Union[str, 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 							UpperCamelCase  (					self		:  int			,				snake_case__		:  str			,				snake_case__		:  int ):



   '''simple docstring'''




   if self.remap is not None:
     UpperCAmelCase__      :					List[Any]				       =							indices.reshape(shape[0]			,				-1 )  # add batch axis
     UpperCAmelCase__      :					Dict				       =							self.unmap_to_all(snake_case__ )
     UpperCAmelCase__      :					Dict				       =							indices.reshape(-1 )  # flatten again

   # get quantized latent vectors
   UpperCAmelCase__      :					Optional[Any]				       =							self.embedding(snake_case__ )

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

   return z_q






class 							UpperCAmelCase_							(    A				):

 '''simple docstring'''







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



   '''simple docstring'''




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







 def 							UpperCamelCase  (					self		:  List[Any]			,				snake_case__		:  Optional[torch.Generator] = None ):



   '''simple docstring'''




   UpperCAmelCase__      :					int				       =							randn_tensor(
       self.mean.shape			,				generator=snake_case__			,				device=self.parameters.device			,				dtype=self.parameters.dtype )
   UpperCAmelCase__      :					List[str]				       =							self.mean + self.std * sample
   return x







 def 							UpperCamelCase  (					self		:  Optional[Any]			,				snake_case__		:  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 							UpperCamelCase  (					self		:  List[Any]			,				snake_case__		:  Tuple			,				snake_case__		:  Tuple=[1, 2, 3] ):



   '''simple docstring'''




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







 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   return self.mean



