

from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFSequenceClassifierOutput,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import logging
from .configuration_regnet import RegNetConfig


UpperCamelCase								=					logging.get_logger(__name__)

# General docstring
UpperCamelCase								=					'RegNetConfig'

# Base docstring
UpperCamelCase								=					'facebook/regnet-y-040'
UpperCamelCase								=					[1, 1088, 7, 7]

# Image classification docstring
UpperCamelCase								=					'facebook/regnet-y-040'
UpperCamelCase								=					'tabby, tabby cat'

UpperCamelCase								=					[
    'facebook/regnet-y-040',
    # See all regnet models at https://huggingface.co/models?filter=regnet
]





class 			__lowerCamelCase	(   tf.keras.layers.Layer   ):




    """simple docstring"""

    def __init__(	self					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					int = 3       ,     SCREAMING_SNAKE_CASE__					:					int = 1       ,     SCREAMING_SNAKE_CASE__					:					int = 1       ,     SCREAMING_SNAKE_CASE__					:					Optional[str] = "relu"       ,     **SCREAMING_SNAKE_CASE__					:					Optional[Any]       ,     )				->	Tuple:
         super().__init__(**__lowerCAmelCase       )
         # The padding and conv has been verified in
         # https://colab.research.google.com/gist/sayakpaul/854bc10eeaf21c9ee2119e0b9f3841a7/scratchpad.ipynb
         lowerCAmelCase__          =			tf.keras.layers.ZeroPaddingaD(padding=kernel_size // 2       )
         lowerCAmelCase__          =			tf.keras.layers.ConvaD(
             filters=__lowerCAmelCase       ,     kernel_size=__lowerCAmelCase       ,     strides=__lowerCAmelCase       ,     padding="VALID"       ,     groups=__lowerCAmelCase       ,     use_bias=__lowerCAmelCase       ,     name="convolution"       ,     )
         lowerCAmelCase__          =			tf.keras.layers.BatchNormalization(epsilon=1e-5       ,     momentum=0.9       ,     name="normalization"       )
         lowerCAmelCase__          =			ACTaFN[activation] if activation is not None else tf.identity


    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       )				->	int:
         lowerCAmelCase__          =			self.convolution(self.padding(__lowerCAmelCase       )       )
         lowerCAmelCase__          =			self.normalization(__lowerCAmelCase       )
         lowerCAmelCase__          =			self.activation(__lowerCAmelCase       )
         return hidden_state






class 			__lowerCamelCase	(   tf.keras.layers.Layer   ):




    """simple docstring"""

    def __init__(	self					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					RegNetConfig       ,     **SCREAMING_SNAKE_CASE__					:					Dict       )				->	Dict:
         super().__init__(**__lowerCAmelCase       )
         lowerCAmelCase__          =			config.num_channels
         lowerCAmelCase__          =			TFRegNetConvLayer(
             out_channels=config.embedding_size       ,     kernel_size=3       ,     stride=2       ,     activation=config.hidden_act       ,     name="embedder"       ,     )


    def        a				(	self					:					int       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       )				->	int:
         lowerCAmelCase__          =			shape_list(__lowerCAmelCase       )[1]
         if tf.executing_eagerly() and num_channels != self.num_channels:
              raise ValueError(
                  "Make sure that the channel dimension of the pixel values match with the one set in the configuration."       )

         # When running on CPU, `tf.keras.layers.Conv2D` doesn't support `NCHW` format.
         # So change the input format from `NCHW` to `NHWC`.
         # shape = (batch_size, in_height, in_width, in_channels=num_channels)
         lowerCAmelCase__          =			tf.transpose(__lowerCAmelCase       ,     perm=(0, 2, 3, 1)       )
         lowerCAmelCase__          =			self.embedder(__lowerCAmelCase       )
         return hidden_state






class 			__lowerCamelCase	(   tf.keras.layers.Layer   ):




    """simple docstring"""

    def __init__(	self					:					str       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					int = 2       ,     **SCREAMING_SNAKE_CASE__					:					List[str]       )				->	Tuple:
         super().__init__(**__lowerCAmelCase       )
         lowerCAmelCase__          =			tf.keras.layers.ConvaD(
             filters=__lowerCAmelCase       ,     kernel_size=1       ,     strides=__lowerCAmelCase       ,     use_bias=__lowerCAmelCase       ,     name="convolution"       )
         lowerCAmelCase__          =			tf.keras.layers.BatchNormalization(epsilon=1e-5       ,     momentum=0.9       ,     name="normalization"       )


    def        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					tf.Tensor       ,     SCREAMING_SNAKE_CASE__					:					bool = False       )				->	str:
         return self.normalization(self.convolution(__lowerCAmelCase       )       ,     training=__lowerCAmelCase       )






class 			__lowerCamelCase	(   tf.keras.layers.Layer   ):




    """simple docstring"""

    def __init__(	self					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					int       ,     **SCREAMING_SNAKE_CASE__					:					int       )				->	List[Any]:
         super().__init__(**__lowerCAmelCase       )
         lowerCAmelCase__          =			tf.keras.layers.GlobalAveragePoolingaD(keepdims=__lowerCAmelCase       ,     name="pooler"       )
         lowerCAmelCase__          =			[
             tf.keras.layers.ConvaD(filters=__lowerCAmelCase       ,     kernel_size=1       ,     activation="relu"       ,     name="attention.0"       ),
             tf.keras.layers.ConvaD(filters=__lowerCAmelCase       ,     kernel_size=1       ,     activation="sigmoid"       ,     name="attention.2"       ),
         ]


    def        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       )				->	List[str]:
         # [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
         lowerCAmelCase__          =			self.pooler(__lowerCAmelCase       )
         for layer_module in self.attention:
              lowerCAmelCase__          =			layer_module(__lowerCAmelCase       )
         lowerCAmelCase__          =			hidden_state * pooled
         return hidden_state






class 			__lowerCamelCase	(   tf.keras.layers.Layer   ):




    """simple docstring"""

    def __init__(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					RegNetConfig       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					int = 1       ,     **SCREAMING_SNAKE_CASE__					:					List[str]       )				->	Tuple:
         super().__init__(**__lowerCAmelCase       )
         lowerCAmelCase__          =			in_channels != out_channels or stride != 1
         lowerCAmelCase__          =			max(1       ,     out_channels // config.groups_width       )
         lowerCAmelCase__          =			(
             TFRegNetShortCut(__lowerCAmelCase       ,     stride=__lowerCAmelCase       ,     name="shortcut"       )
             if should_apply_shortcut
             else tf.keras.layers.Activation("linear"       ,     name="shortcut"       )
         )
         # `self.layers` instead of `self.layer` because that is a reserved argument.
         lowerCAmelCase__          =			[
             TFRegNetConvLayer(__lowerCAmelCase       ,     kernel_size=1       ,     activation=config.hidden_act       ,     name="layer.0"       ),
             TFRegNetConvLayer(
                 __lowerCAmelCase       ,     stride=__lowerCAmelCase       ,     groups=__lowerCAmelCase       ,     activation=config.hidden_act       ,     name="layer.1"       ),
             TFRegNetConvLayer(__lowerCAmelCase       ,     kernel_size=1       ,     activation=__lowerCAmelCase       ,     name="layer.2"       ),
         ]
         lowerCAmelCase__          =			ACTaFN[config.hidden_act]


    def        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Any       )				->	List[str]:
         lowerCAmelCase__          =			hidden_state
         for layer_module in self.layers:
              lowerCAmelCase__          =			layer_module(__lowerCAmelCase       )
         lowerCAmelCase__          =			self.shortcut(__lowerCAmelCase       )
         hidden_state += residual
         lowerCAmelCase__          =			self.activation(__lowerCAmelCase       )
         return hidden_state






class 			__lowerCamelCase	(   tf.keras.layers.Layer   ):




    """simple docstring"""

    def __init__(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					RegNetConfig       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					int = 1       ,     **SCREAMING_SNAKE_CASE__					:					List[Any]       )				->	Optional[Any]:
         super().__init__(**__lowerCAmelCase       )
         lowerCAmelCase__          =			in_channels != out_channels or stride != 1
         lowerCAmelCase__          =			max(1       ,     out_channels // config.groups_width       )
         lowerCAmelCase__          =			(
             TFRegNetShortCut(__lowerCAmelCase       ,     stride=__lowerCAmelCase       ,     name="shortcut"       )
             if should_apply_shortcut
             else tf.keras.layers.Activation("linear"       ,     name="shortcut"       )
         )
         lowerCAmelCase__          =			[
             TFRegNetConvLayer(__lowerCAmelCase       ,     kernel_size=1       ,     activation=config.hidden_act       ,     name="layer.0"       ),
             TFRegNetConvLayer(
                 __lowerCAmelCase       ,     stride=__lowerCAmelCase       ,     groups=__lowerCAmelCase       ,     activation=config.hidden_act       ,     name="layer.1"       ),
             TFRegNetSELayer(__lowerCAmelCase       ,     reduced_channels=int(round(in_channels / 4       )       )       ,     name="layer.2"       ),
             TFRegNetConvLayer(__lowerCAmelCase       ,     kernel_size=1       ,     activation=__lowerCAmelCase       ,     name="layer.3"       ),
         ]
         lowerCAmelCase__          =			ACTaFN[config.hidden_act]


    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Dict       )				->	Tuple:
         lowerCAmelCase__          =			hidden_state
         for layer_module in self.layers:
              lowerCAmelCase__          =			layer_module(__lowerCAmelCase       )
         lowerCAmelCase__          =			self.shortcut(__lowerCAmelCase       )
         hidden_state += residual
         lowerCAmelCase__          =			self.activation(__lowerCAmelCase       )
         return hidden_state






class 			__lowerCamelCase	(   tf.keras.layers.Layer   ):




    """simple docstring"""

    def __init__(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					RegNetConfig       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					int = 2       ,     SCREAMING_SNAKE_CASE__					:					int = 2       ,     **SCREAMING_SNAKE_CASE__					:					Dict       )				->	int:
         super().__init__(**__lowerCAmelCase       )

         lowerCAmelCase__          =			TFRegNetXLayer if config.layer_type == "x" else TFRegNetYLayer
         lowerCAmelCase__          =			[
             # downsampling is done in the first layer with stride of 2
             layer(__lowerCAmelCase       ,     __lowerCAmelCase       ,     __lowerCAmelCase       ,     stride=__lowerCAmelCase       ,     name="layers.0"       ),
             *[layer(__lowerCAmelCase       ,     __lowerCAmelCase       ,     __lowerCAmelCase       ,     name=f'layers.{i+1}'       ) for i in range(depth - 1       )],
         ]


    def        a				(	self					:					int       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       )				->	str:
         for layer_module in self.layers:
              lowerCAmelCase__          =			layer_module(__lowerCAmelCase       )
         return hidden_state






class 			__lowerCamelCase	(   tf.keras.layers.Layer   ):




    """simple docstring"""

    def __init__(	self					:					str       ,     SCREAMING_SNAKE_CASE__					:					RegNetConfig       ,     **SCREAMING_SNAKE_CASE__					:					Optional[Any]       )				->	Tuple:
         super().__init__(**__lowerCAmelCase       )
         lowerCAmelCase__          =			[]
         # based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
         self.stages.append(
             TFRegNetStage(
                 __lowerCAmelCase       ,     config.embedding_size       ,     config.hidden_sizes[0]       ,     stride=2 if config.downsample_in_first_stage else 1       ,     depth=config.depths[0]       ,     name="stages.0"       ,     )       )
         lowerCAmelCase__          =			zip(config.hidden_sizes       ,     config.hidden_sizes[1:]       )
         for i, ((in_channels, out_channels), depth) in enumerate(zip(__lowerCAmelCase       ,     config.depths[1:]       )       ):
              self.stages.append(TFRegNetStage(__lowerCAmelCase       ,     __lowerCAmelCase       ,     __lowerCAmelCase       ,     depth=__lowerCAmelCase       ,     name=f'stages.{i+1}'       )       )


    def        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					tf.Tensor       ,     SCREAMING_SNAKE_CASE__					:					bool = False       ,     SCREAMING_SNAKE_CASE__					:					bool = True       )				->	List[str]:
         lowerCAmelCase__          =			() if output_hidden_states else None

         for stage_module in self.stages:
              if output_hidden_states:
                   lowerCAmelCase__          =			hidden_states + (hidden_state,)

              lowerCAmelCase__          =			stage_module(__lowerCAmelCase       )

         if output_hidden_states:
              lowerCAmelCase__          =			hidden_states + (hidden_state,)

         if not return_dict:
              return tuple(v for v in [hidden_state, hidden_states] if v is not None       )

         return TFBaseModelOutputWithNoAttention(last_hidden_state=__lowerCAmelCase       ,     hidden_states=__lowerCAmelCase       )






@keras_serializable
class 			__lowerCamelCase	(   tf.keras.layers.Layer   ):




    """simple docstring"""

    snake_case__ =				RegNetConfig


    def __init__(	self					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       ,     **SCREAMING_SNAKE_CASE__					:					Optional[int]       )				->	Any:
         super().__init__(**__lowerCAmelCase       )
         lowerCAmelCase__          =			config
         lowerCAmelCase__          =			TFRegNetEmbeddings(__lowerCAmelCase       ,     name="embedder"       )
         lowerCAmelCase__          =			TFRegNetEncoder(__lowerCAmelCase       ,     name="encoder"       )
         lowerCAmelCase__          =			tf.keras.layers.GlobalAveragePoolingaD(keepdims=__lowerCAmelCase       ,     name="pooler"       )


    @unpack_inputs
    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					tf.Tensor       ,     SCREAMING_SNAKE_CASE__					:					Optional[bool] = None       ,     SCREAMING_SNAKE_CASE__					:					Optional[bool] = None       ,     SCREAMING_SNAKE_CASE__					:					bool = False       ,     )				->	Tuple:
         lowerCAmelCase__          =			(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         lowerCAmelCase__          =			return_dict if return_dict is not None else self.config.use_return_dict

         lowerCAmelCase__          =			self.embedder(__lowerCAmelCase       ,     training=__lowerCAmelCase       )

         lowerCAmelCase__          =			self.encoder(
             __lowerCAmelCase       ,     output_hidden_states=__lowerCAmelCase       ,     return_dict=__lowerCAmelCase       ,     training=__lowerCAmelCase       )

         lowerCAmelCase__          =			encoder_outputs[0]
         lowerCAmelCase__          =			self.pooler(__lowerCAmelCase       )

         # Change to NCHW output format have uniformity in the modules
         lowerCAmelCase__          =			tf.transpose(__lowerCAmelCase       ,     perm=(0, 3, 1, 2)       )
         lowerCAmelCase__          =			tf.transpose(__lowerCAmelCase       ,     perm=(0, 3, 1, 2)       )

         # Change the other hidden state outputs to NCHW as well
         if output_hidden_states:
              lowerCAmelCase__          =			tuple([tf.transpose(__lowerCAmelCase       ,     perm=(0, 3, 1, 2)       ) for h in encoder_outputs[1]]       )

         if not return_dict:
              return (last_hidden_state, pooled_output) + encoder_outputs[1:]

         return TFBaseModelOutputWithPoolingAndNoAttention(
             last_hidden_state=__lowerCAmelCase       ,     pooler_output=__lowerCAmelCase       ,     hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states       ,     )






class 			__lowerCamelCase	(   lowerCAmelCase_   ):




    """simple docstring"""

    snake_case__ =				RegNetConfig
    snake_case__ =				'regnet'
    snake_case__ =				'pixel_values'


    @property
    def        a				(	self					:					Any       )				->	int:
         return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)       ,     dtype=tf.floataa       )}






UpperCamelCase								=					R'\n    Parameters:\n    This model is a Tensorflow\n    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a\n    regular Tensorflow Module and refer to the Tensorflow documentation for all matter related to general usage and\n    behavior.\n        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.\n'

UpperCamelCase								=					R'\n    Args:\n        pixel_values (`tf.Tensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConveNextImageProcessor.__call__`] for details.\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n'





@add_start_docstrings(
    "The bare RegNet model outputting raw features without any specific head on top."     ,    lowerCAmelCase_     ,    )
class 			__lowerCamelCase	(   lowerCAmelCase_   ):




    """simple docstring"""

    def __init__(	self					:					int       ,     SCREAMING_SNAKE_CASE__					:					RegNetConfig       ,     *SCREAMING_SNAKE_CASE__					:					Dict       ,     **SCREAMING_SNAKE_CASE__					:					List[str]       )				->	List[str]:
         super().__init__(__lowerCAmelCase       ,     *__lowerCAmelCase       ,     **__lowerCAmelCase       )
         lowerCAmelCase__          =			TFRegNetMainLayer(__lowerCAmelCase       ,     name="regnet"       )


    @unpack_inputs
    @add_start_docstrings_to_model_forward(__lowerCAmelCase       )
    @add_code_sample_docstrings(
        checkpoint=_CHECKPOINT_FOR_DOC       ,     output_type=__lowerCAmelCase       ,     config_class=_CONFIG_FOR_DOC       ,     modality="vision"       ,     expected_output=_EXPECTED_OUTPUT_SHAPE       ,     )
    def        a				(	self					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					tf.Tensor       ,     SCREAMING_SNAKE_CASE__					:					Optional[bool] = None       ,     SCREAMING_SNAKE_CASE__					:					Optional[bool] = None       ,     SCREAMING_SNAKE_CASE__					:					Any=False       ,     )				->	Dict:
         lowerCAmelCase__          =			(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         lowerCAmelCase__          =			return_dict if return_dict is not None else self.config.use_return_dict

         lowerCAmelCase__          =			self.regnet(
             pixel_values=__lowerCAmelCase       ,     output_hidden_states=__lowerCAmelCase       ,     return_dict=__lowerCAmelCase       ,     training=__lowerCAmelCase       ,     )
         if not return_dict:
              return (outputs[0],) + outputs[1:]

         return TFBaseModelOutputWithPoolingAndNoAttention(
             last_hidden_state=outputs.last_hidden_state       ,     pooler_output=outputs.pooler_output       ,     hidden_states=outputs.hidden_states       ,     )






@add_start_docstrings(
    "\n    RegNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for\n    ImageNet.\n    "     ,    lowerCAmelCase_     ,    )
class 			__lowerCamelCase	(   lowerCAmelCase_     ,    lowerCAmelCase_   ):




    """simple docstring"""

    def __init__(	self					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					RegNetConfig       ,     *SCREAMING_SNAKE_CASE__					:					List[Any]       ,     **SCREAMING_SNAKE_CASE__					:					List[Any]       )				->	Dict:
         super().__init__(__lowerCAmelCase       ,     *__lowerCAmelCase       ,     **__lowerCAmelCase       )
         lowerCAmelCase__          =			config.num_labels
         lowerCAmelCase__          =			TFRegNetMainLayer(__lowerCAmelCase       ,     name="regnet"       )
         # classification head
         lowerCAmelCase__          =			[
             tf.keras.layers.Flatten(),
             tf.keras.layers.Dense(config.num_labels       ,     name="classifier.1"       ) if config.num_labels > 0 else tf.identity,
         ]


    @unpack_inputs
    @add_start_docstrings_to_model_forward(__lowerCAmelCase       )
    @add_code_sample_docstrings(
        checkpoint=_IMAGE_CLASS_CHECKPOINT       ,     output_type=__lowerCAmelCase       ,     config_class=_CONFIG_FOR_DOC       ,     expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       ,     )
    def        a				(	self					:					str       ,     SCREAMING_SNAKE_CASE__					:					tf.Tensor = None       ,     SCREAMING_SNAKE_CASE__					:					tf.Tensor = None       ,     SCREAMING_SNAKE_CASE__					:					bool = None       ,     SCREAMING_SNAKE_CASE__					:					bool = None       ,     SCREAMING_SNAKE_CASE__					:					str=False       ,     )				->	str:
         lowerCAmelCase__          =			(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         lowerCAmelCase__          =			return_dict if return_dict is not None else self.config.use_return_dict

         lowerCAmelCase__          =			self.regnet(
             __lowerCAmelCase       ,     output_hidden_states=__lowerCAmelCase       ,     return_dict=__lowerCAmelCase       ,     training=__lowerCAmelCase       )

         lowerCAmelCase__          =			outputs.pooler_output if return_dict else outputs[1]

         lowerCAmelCase__          =			self.classifier[0](__lowerCAmelCase       )
         lowerCAmelCase__          =			self.classifier[1](__lowerCAmelCase       )

         lowerCAmelCase__          =			None if labels is None else self.hf_compute_loss(labels=__lowerCAmelCase       ,     logits=__lowerCAmelCase       )

         if not return_dict:
              lowerCAmelCase__          =			(logits,) + outputs[2:]
              return ((loss,) + output) if loss is not None else output

         return TFSequenceClassifierOutput(loss=__lowerCAmelCase       ,     logits=__lowerCAmelCase       ,     hidden_states=outputs.hidden_states       )






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




    """simple docstring"""

    snake_case__ =				42






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Any=3       ,     SCREAMING_SNAKE_CASE__					:					int=3       ,     SCREAMING_SNAKE_CASE__					:					List[str]=("DownEncoderBlock2D",)       ,     SCREAMING_SNAKE_CASE__					:					Tuple=(64,)       ,     SCREAMING_SNAKE_CASE__					:					str=2       ,     SCREAMING_SNAKE_CASE__					:					List[Any]=32       ,     SCREAMING_SNAKE_CASE__					:					List[str]="silu"       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]=True       ,     )				->	str:
         super().__init__()
         lowerCAmelCase__          =			layers_per_block

         lowerCAmelCase__          =			torch.nn.Convad(
             SCREAMING_SNAKE_CASE__       ,     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(SCREAMING_SNAKE_CASE__       ):
              lowerCAmelCase__          =			output_channel
              lowerCAmelCase__          =			block_out_channels[i]
              lowerCAmelCase__          =			i == len(SCREAMING_SNAKE_CASE__       ) - 1

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

         # mid
         lowerCAmelCase__          =			UNetMidBlockaD(
             in_channels=block_out_channels[-1]       ,     resnet_eps=1e-6       ,     resnet_act_fn=SCREAMING_SNAKE_CASE__       ,     output_scale_factor=1       ,     resnet_time_scale_shift="default"       ,     attention_head_dim=block_out_channels[-1]       ,     resnet_groups=SCREAMING_SNAKE_CASE__       ,     temb_channels=SCREAMING_SNAKE_CASE__       ,     )

         # out
         lowerCAmelCase__          =			nn.GroupNorm(num_channels=block_out_channels[-1]       ,     num_groups=SCREAMING_SNAKE_CASE__       ,     eps=1e-6       )
         lowerCAmelCase__          =			nn.SiLU()

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

         lowerCAmelCase__          =			False


    def        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Dict       )				->	int:
         lowerCAmelCase__          =			x
         lowerCAmelCase__          =			self.conv_in(SCREAMING_SNAKE_CASE__       )

         if self.training and self.gradient_checkpointing:

              def create_custom_forward(SCREAMING_SNAKE_CASE__					:					Optional[int]       ):
                   def custom_forward(*SCREAMING_SNAKE_CASE__					:					Dict       ):
                        return module(*SCREAMING_SNAKE_CASE__       )

                   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(SCREAMING_SNAKE_CASE__       )       ,     SCREAMING_SNAKE_CASE__       ,     use_reentrant=SCREAMING_SNAKE_CASE__       )
                   # middle
                   lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(
                       create_custom_forward(self.mid_block       )       ,     SCREAMING_SNAKE_CASE__       ,     use_reentrant=SCREAMING_SNAKE_CASE__       )
              else:
                   for down_block in self.down_blocks:
                        lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(create_custom_forward(SCREAMING_SNAKE_CASE__       )       ,     SCREAMING_SNAKE_CASE__       )
                   # middle
                   lowerCAmelCase__          =			torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block       )       ,     SCREAMING_SNAKE_CASE__       )

         else:
              # down
              for down_block in self.down_blocks:
                   lowerCAmelCase__          =			down_block(SCREAMING_SNAKE_CASE__       )

              # middle
              lowerCAmelCase__          =			self.mid_block(SCREAMING_SNAKE_CASE__       )

         # post-process
         lowerCAmelCase__          =			self.conv_norm_out(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.conv_act(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.conv_out(SCREAMING_SNAKE_CASE__       )

         return sample






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					str=3       ,     SCREAMING_SNAKE_CASE__					:					Dict=3       ,     SCREAMING_SNAKE_CASE__					:					int=("UpDecoderBlock2D",)       ,     SCREAMING_SNAKE_CASE__					:					Tuple=(64,)       ,     SCREAMING_SNAKE_CASE__					:					Tuple=2       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]=32       ,     SCREAMING_SNAKE_CASE__					:					List[str]="silu"       ,     SCREAMING_SNAKE_CASE__					:					List[str]="group"       ,     )				->	Union[str, Any]:
         super().__init__()
         lowerCAmelCase__          =			layers_per_block

         lowerCAmelCase__          =			nn.Convad(
             SCREAMING_SNAKE_CASE__       ,     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=SCREAMING_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=SCREAMING_SNAKE_CASE__       ,     temb_channels=SCREAMING_SNAKE_CASE__       ,     )

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

              lowerCAmelCase__          =			i == len(SCREAMING_SNAKE_CASE__       ) - 1

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

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

         lowerCAmelCase__          =			False


    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					Dict=None       )				->	Dict:
         lowerCAmelCase__          =			z
         lowerCAmelCase__          =			self.conv_in(SCREAMING_SNAKE_CASE__       )

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

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

                   return custom_forward

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

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

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

              # up
              for up_block in self.up_blocks:
                   lowerCAmelCase__          =			up_block(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )

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

         return sample






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					Any       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]=None       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]="random"       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]=False       ,     SCREAMING_SNAKE_CASE__					:					Any=True       )				->	List[Any]:
         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        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       )				->	Optional[int]:
         lowerCAmelCase__          =			inds.shape
         assert len(SCREAMING_SNAKE_CASE__       ) > 1
         lowerCAmelCase__          =			inds.reshape(ishape[0]       ,     -1       )
         lowerCAmelCase__          =			self.used.to(SCREAMING_SNAKE_CASE__       )
         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(SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Dict       )				->	Optional[int]:
         lowerCAmelCase__          =			inds.shape
         assert len(SCREAMING_SNAKE_CASE__       ) > 1
         lowerCAmelCase__          =			inds.reshape(ishape[0]       ,     -1       )
         lowerCAmelCase__          =			self.used.to(SCREAMING_SNAKE_CASE__       )
         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       ,     SCREAMING_SNAKE_CASE__       )
         return back.reshape(SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Tuple       )				->	Tuple:
         # 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(SCREAMING_SNAKE_CASE__       ,     self.embedding.weight       )       ,     dim=1       )

         lowerCAmelCase__          =			self.embedding(SCREAMING_SNAKE_CASE__       ).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(SCREAMING_SNAKE_CASE__       )
              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        a				(	self					:					str       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					str       )				->	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(SCREAMING_SNAKE_CASE__       )
              lowerCAmelCase__          =			indices.reshape(-1       )  # flatten again

         # get quantized latent vectors
         lowerCAmelCase__          =			self.embedding(SCREAMING_SNAKE_CASE__       )

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

         return z_q






class 			__lowerCamelCase	(   UpperCamelCase__   ):




    """simple docstring"""

    def __init__(	self					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Any=False       )				->	Optional[int]:
         lowerCAmelCase__          =			parameters
         lowerCAmelCase__     ,		lowerCAmelCase__          =			torch.chunk(SCREAMING_SNAKE_CASE__       ,     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        a				(	self					:					str       ,     SCREAMING_SNAKE_CASE__					:					Optional[torch.Generator] = None       )				->	torch.FloatTensor:
         # make sure sample is on the same device as the parameters and has same dtype
         lowerCAmelCase__          =			randn_tensor(
             self.mean.shape       ,     generator=SCREAMING_SNAKE_CASE__       ,     device=self.parameters.device       ,     dtype=self.parameters.dtype       )
         lowerCAmelCase__          =			self.mean + self.std * sample
         return x


    def        a				(	self					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					List[Any]=None       )				->	Union[str, Any]:
         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					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					List[str]=[1, 2, 3]       )				->	Tuple:
         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=SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Union[str, Any]       )				->	Dict:
         return self.mean
