


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       lowerCamelCase__   (       lowerCAmelCase):
     SCREAMING_SNAKE_CASE__      			=			42







class       lowerCamelCase__   (       nn.Module):





     def __init__(self		,     UpperCAmelCase=3		,     UpperCAmelCase=3		,     UpperCAmelCase=("DownEncoderBlock2D",)		,     UpperCAmelCase=(6_4,)		,     UpperCAmelCase=2		,     UpperCAmelCase=3_2		,     UpperCAmelCase="silu"		,     UpperCAmelCase=True		,     )     ->   Any:
      super().__init__()
      _lowercase							      =layers_per_block

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

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

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

       _lowercase							      =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
      _lowercase							      =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
      _lowercase							      =nn.GroupNorm(num_channels=block_out_channels[-1]		,     num_groups=UpperCAmelCase		,     eps=1e-6       )
      _lowercase							      =nn.SiLU()

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

      _lowercase							      =False





     def      __A					(self		,     UpperCAmelCase       )     ->   Any:
      _lowercase							      =x
      _lowercase							      =self.conv_in(UpperCAmelCase       )

      if self.training and self.gradient_checkpointing:

       def create_custom_forward(UpperCAmelCase       ):
        def custom_forward(*UpperCAmelCase       ):
         return module(*UpperCAmelCase       )

        return custom_forward

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

      else:
       # down
       for down_block in self.down_blocks:
        _lowercase							      =down_block(UpperCAmelCase       )

       # middle
       _lowercase							      =self.mid_block(UpperCAmelCase       )

      # post-process
      _lowercase							      =self.conv_norm_out(UpperCAmelCase       )
      _lowercase							      =self.conv_act(UpperCAmelCase       )
      _lowercase							      =self.conv_out(UpperCAmelCase       )

      return sample







class       lowerCamelCase__   (       nn.Module):





     def __init__(self		,     UpperCAmelCase=3		,     UpperCAmelCase=3		,     UpperCAmelCase=("UpDecoderBlock2D",)		,     UpperCAmelCase=(6_4,)		,     UpperCAmelCase=2		,     UpperCAmelCase=3_2		,     UpperCAmelCase="silu"		,     UpperCAmelCase="group"		,     )     ->   int:
      super().__init__()
      _lowercase							      =layers_per_block

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

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

      _lowercase							      =in_channels if norm_type == '''spatial''' else None

      # mid
      _lowercase							      =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
      _lowercase							      =list(reversed(UpperCAmelCase       )       )
      _lowercase							      =reversed_block_out_channels[0]
      for i, up_block_type in enumerate(UpperCAmelCase       ):
       _lowercase							      =output_channel
       _lowercase							      =reversed_block_out_channels[i]

       _lowercase							      =i == len(UpperCAmelCase       ) - 1

       _lowercase							      =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       )
       _lowercase							      =output_channel

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

      _lowercase							      =False





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase=None       )     ->   Any:
      _lowercase							      =z
      _lowercase							      =self.conv_in(UpperCAmelCase       )

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

       def create_custom_forward(UpperCAmelCase       ):
        def custom_forward(*UpperCAmelCase       ):
         return module(*UpperCAmelCase       )

        return custom_forward

       if is_torch_version('''>='''		,     '''1.11.0'''       ):
        # middle
        _lowercase							      =torch.utils.checkpoint.checkpoint(
            create_custom_forward(self.mid_block       )		,     UpperCAmelCase		,     UpperCAmelCase		,     use_reentrant=UpperCAmelCase       )
        _lowercase							      =sample.to(UpperCAmelCase       )

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

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

       # up
       for up_block in self.up_blocks:
        _lowercase							      =up_block(UpperCAmelCase		,     UpperCAmelCase       )

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

      return sample







class       lowerCamelCase__   (       nn.Module):





     def __init__(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase=None		,     UpperCAmelCase="random"		,     UpperCAmelCase=False		,     UpperCAmelCase=True       )     ->   Tuple:
      super().__init__()
      _lowercase							      =n_e
      _lowercase							      =vq_embed_dim
      _lowercase							      =beta
      _lowercase							      =legacy

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

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

      _lowercase							      =sane_index_shape





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





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





     def      __A					(self		,     UpperCAmelCase       )     ->   str:
      # reshape z -> (batch, height, width, channel) and flatten
      _lowercase							      =z.permute(0		,     2		,     3		,     1       ).contiguous()
      _lowercase							      =z.view(-1		,     self.vq_embed_dim       )

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

      _lowercase							      =self.embedding(UpperCAmelCase       ).view(z.shape       )
      _lowercase							      =None
      _lowercase							      =None

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

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

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

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

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

      # get quantized latent vectors
      _lowercase							      =self.embedding(UpperCAmelCase       )

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

      return z_q







class       lowerCamelCase__   (       lowerCAmelCase):





     def __init__(self		,     UpperCAmelCase		,     UpperCAmelCase=False       )     ->   Tuple:
      _lowercase							      =parameters
      _lowercase      ,    _lowercase							      =torch.chunk(UpperCAmelCase		,     2		,     dim=1       )
      _lowercase							      =torch.clamp(self.logvar		,     -30.0		,     20.0       )
      _lowercase							      =deterministic
      _lowercase							      =torch.exp(0.5 * self.logvar       )
      _lowercase							      =torch.exp(self.logvar       )
      if self.deterministic:
       _lowercase							      =_lowercase							      =torch.zeros_like(
           self.mean		,     device=self.parameters.device		,     dtype=self.parameters.dtype       )





     def      __A					(self		,     UpperCAmelCase = None       )     ->   torch.FloatTensor:
      # make sure sample is on the same device as the parameters and has same dtype
      _lowercase							      =randn_tensor(
          self.mean.shape		,     generator=UpperCAmelCase		,     device=self.parameters.device		,     dtype=self.parameters.dtype       )
      _lowercase							      =self.mean + self.std * sample
      return x





     def      __A					(self		,     UpperCAmelCase=None       )     ->   Dict:
      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		,     UpperCAmelCase		,     UpperCAmelCase=[1, 2, 3]       )     ->   Dict:
      if self.deterministic:
       return torch.Tensor([0.0]       )
      _lowercase							      =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      __A					(self       )     ->   Optional[Any]:
      return self.mean





# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
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)

