
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        A_   (			tf.keras.layers.Layer   ):




	'''simple docstring'''


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





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case       ):
		lowercase       =    self.convolution(self.padding(a__       )       )
		lowercase       =    self.normalization(a__       )
		lowercase       =    self.activation(a__       )
		return hidden_state







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




	'''simple docstring'''


	def __init__(       self		,    snake_case		,    **snake_case       ):
		super().__init__(**a__       )
		lowercase       =    config.num_channels
		lowercase       =    TFRegNetConvLayer(
		    out_channels=config.embedding_size		,    kernel_size=3		,    stride=2		,    activation=config.hidden_act		,    name='embedder'		,    )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case       ):
		lowercase       =    shape_list(a__       )[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)
		lowercase       =    tf.transpose(a__		,    perm=(0, 2, 3, 1)       )
		lowercase       =    self.embedder(a__       )
		return hidden_state







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




	'''simple docstring'''


	def __init__(       self		,    snake_case		,    snake_case = 2		,    **snake_case       ):
		super().__init__(**a__       )
		lowercase       =    tf.keras.layers.ConvaD(
		    filters=a__		,    kernel_size=1		,    strides=a__		,    use_bias=a__		,    name='convolution'       )
		lowercase       =    tf.keras.layers.BatchNormalization(epsilon=1E-5		,    momentum=0.9		,    name='normalization'       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case = False       ):
		return self.normalization(self.convolution(a__       )		,    training=a__       )







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




	'''simple docstring'''


	def __init__(       self		,    snake_case		,    snake_case		,    **snake_case       ):
		super().__init__(**a__       )
		lowercase       =    tf.keras.layers.GlobalAveragePoolingaD(keepdims=a__		,    name='pooler'       )
		lowercase       =    [
		    tf.keras.layers.ConvaD(filters=a__		,    kernel_size=1		,    activation='relu'		,    name='attention.0'       ),
		    tf.keras.layers.ConvaD(filters=a__		,    kernel_size=1		,    activation='sigmoid'		,    name='attention.2'       ),
		]





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case       ):
		# [batch_size, h, w, num_channels] -> [batch_size, 1, 1, num_channels]
		lowercase       =    self.pooler(a__       )
		for layer_module in self.attention:
			lowercase       =    layer_module(a__       )
		lowercase       =    hidden_state * pooled
		return hidden_state







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




	'''simple docstring'''


	def __init__(       self		,    snake_case		,    snake_case		,    snake_case		,    snake_case = 1		,    **snake_case       ):
		super().__init__(**a__       )
		lowercase       =    in_channels != out_channels or stride != 1
		lowercase       =    max(1		,    out_channels // config.groups_width       )
		lowercase       =    (
		    TFRegNetShortCut(a__		,    stride=a__		,    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.
		lowercase       =    [
		    TFRegNetConvLayer(a__		,    kernel_size=1		,    activation=config.hidden_act		,    name='layer.0'       ),
		    TFRegNetConvLayer(
		        a__		,    stride=a__		,    groups=a__		,    activation=config.hidden_act		,    name='layer.1'       ),
		    TFRegNetConvLayer(a__		,    kernel_size=1		,    activation=a__		,    name='layer.2'       ),
		]
		lowercase       =    ACTaFN[config.hidden_act]





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case       ):
		lowercase       =    hidden_state
		for layer_module in self.layers:
			lowercase       =    layer_module(a__       )
		lowercase       =    self.shortcut(a__       )
		hidden_state += residual
		lowercase       =    self.activation(a__       )
		return hidden_state







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




	'''simple docstring'''


	def __init__(       self		,    snake_case		,    snake_case		,    snake_case		,    snake_case = 1		,    **snake_case       ):
		super().__init__(**a__       )
		lowercase       =    in_channels != out_channels or stride != 1
		lowercase       =    max(1		,    out_channels // config.groups_width       )
		lowercase       =    (
		    TFRegNetShortCut(a__		,    stride=a__		,    name='shortcut'       )
		    if should_apply_shortcut
		    else tf.keras.layers.Activation('linear'		,    name='shortcut'       )
		)
		lowercase       =    [
		    TFRegNetConvLayer(a__		,    kernel_size=1		,    activation=config.hidden_act		,    name='layer.0'       ),
		    TFRegNetConvLayer(
		        a__		,    stride=a__		,    groups=a__		,    activation=config.hidden_act		,    name='layer.1'       ),
		    TFRegNetSELayer(a__		,    reduced_channels=int(round(in_channels / 4       )       )		,    name='layer.2'       ),
		    TFRegNetConvLayer(a__		,    kernel_size=1		,    activation=a__		,    name='layer.3'       ),
		]
		lowercase       =    ACTaFN[config.hidden_act]





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case       ):
		lowercase       =    hidden_state
		for layer_module in self.layers:
			lowercase       =    layer_module(a__       )
		lowercase       =    self.shortcut(a__       )
		hidden_state += residual
		lowercase       =    self.activation(a__       )
		return hidden_state







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




	'''simple docstring'''


	def __init__(       self		,    snake_case		,    snake_case		,    snake_case		,    snake_case = 2		,    snake_case = 2		,    **snake_case       ):
		super().__init__(**a__       )

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





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case       ):
		for layer_module in self.layers:
			lowercase       =    layer_module(a__       )
		return hidden_state







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




	'''simple docstring'''


	def __init__(       self		,    snake_case		,    **snake_case       ):
		super().__init__(**a__       )
		lowercase       =    []
		# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
		self.stages.append(
		    TFRegNetStage(
		        a__		,    config.embedding_size		,    config.hidden_sizes[0]		,    stride=2 if config.downsample_in_first_stage else 1		,    depth=config.depths[0]		,    name='stages.0'		,    )       )
		lowercase       =    zip(config.hidden_sizes		,    config.hidden_sizes[1:]       )
		for i, ((in_channels, out_channels), depth) in enumerate(zip(a__		,    config.depths[1:]       )       ):
			self.stages.append(TFRegNetStage(a__		,    a__		,    a__		,    depth=a__		,    name=F'''stages.{i+1}'''       )       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case = False		,    snake_case = True       ):
		lowercase       =    () if output_hidden_states else None

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

			lowercase       =    stage_module(a__       )

		if output_hidden_states:
			lowercase       =    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=a__		,    hidden_states=a__       )







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




	'''simple docstring'''


	_UpperCamelCase     :     Tuple							 =       RegNetConfig





	def __init__(       self		,    snake_case		,    **snake_case       ):
		super().__init__(**a__       )
		lowercase       =    config
		lowercase       =    TFRegNetEmbeddings(a__		,    name='embedder'       )
		lowercase       =    TFRegNetEncoder(a__		,    name='encoder'       )
		lowercase       =    tf.keras.layers.GlobalAveragePoolingaD(keepdims=a__		,    name='pooler'       )





	@unpack_inputs
	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case = None		,    snake_case = None		,    snake_case = False		,    ):
		lowercase       =    (
		    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
		)
		lowercase       =    return_dict if return_dict is not None else self.config.use_return_dict

		lowercase       =    self.embedder(a__		,    training=a__       )

		lowercase       =    self.encoder(
		    a__		,    output_hidden_states=a__		,    return_dict=a__		,    training=a__       )

		lowercase       =    encoder_outputs[0]
		lowercase       =    self.pooler(a__       )

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

		# Change the other hidden state outputs to NCHW as well
		if output_hidden_states:
			lowercase       =    tuple([tf.transpose(a__		,    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=a__		,    pooler_output=a__		,    hidden_states=hidden_states if output_hidden_states else encoder_outputs.hidden_states		,    )







class        A_   (			SCREAMING_SNAKE_CASE_   ):




	'''simple docstring'''


	_UpperCamelCase     :     Optional[Any]							 =       RegNetConfig
	_UpperCamelCase     :     List[str]							 =       "regnet"
	_UpperCamelCase     :     Dict							 =       "pixel_values"





	@property
	def 	SCREAMING_SNAKE_CASE__						(       self       ):
		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."""      ,	SCREAMING_SNAKE_CASE_      ,	)
class        A_   (			SCREAMING_SNAKE_CASE_   ):




	'''simple docstring'''


	def __init__(       self		,    snake_case		,    *snake_case		,    **snake_case       ):
		super().__init__(a__		,    *a__		,    **a__       )
		lowercase       =    TFRegNetMainLayer(a__		,    name='regnet'       )





	@unpack_inputs
	@add_start_docstrings_to_model_forward(a__       )
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC		,    output_type=a__		,    config_class=_CONFIG_FOR_DOC		,    modality='vision'		,    expected_output=_EXPECTED_OUTPUT_SHAPE		,    )
	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case = None		,    snake_case = None		,    snake_case=False		,    ):
		lowercase       =    (
		    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
		)
		lowercase       =    return_dict if return_dict is not None else self.config.use_return_dict

		lowercase       =    self.regnet(
		    pixel_values=a__		,    output_hidden_states=a__		,    return_dict=a__		,    training=a__		,    )
		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    """      ,	SCREAMING_SNAKE_CASE_      ,	)
class        A_   (			SCREAMING_SNAKE_CASE_      ,	SCREAMING_SNAKE_CASE_   ):




	'''simple docstring'''


	def __init__(       self		,    snake_case		,    *snake_case		,    **snake_case       ):
		super().__init__(a__		,    *a__		,    **a__       )
		lowercase       =    config.num_labels
		lowercase       =    TFRegNetMainLayer(a__		,    name='regnet'       )
		# classification head
		lowercase       =    [
		    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(a__       )
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT		,    output_type=a__		,    config_class=_CONFIG_FOR_DOC		,    expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT		,    )
	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case = None		,    snake_case=False		,    ):
		lowercase       =    (
		    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
		)
		lowercase       =    return_dict if return_dict is not None else self.config.use_return_dict

		lowercase       =    self.regnet(
		    a__		,    output_hidden_states=a__		,    return_dict=a__		,    training=a__       )

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

		lowercase       =    self.classifier[0](a__       )
		lowercase       =    self.classifier[1](a__       )

		lowercase       =    None if labels is None else self.hf_compute_loss(labels=a__		,    logits=a__       )

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

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


"""simple docstring"""




# NOTE: This file is deprecated and will be removed in a future version.
# It only exists so that temporarely `from diffusers.pipelines import DiffusionPipeline` works
from ...utils import deprecate
from ..controlnet.multicontrolnet import MultiControlNetModel  # noqa: F401
from ..controlnet.pipeline_controlnet import StableDiffusionControlNetPipeline  # noqa: F401


deprecate(
    'stable diffusion controlnet',
    '0.22.0',
    'Importing `StableDiffusionControlNetPipeline` or `MultiControlNetModel` from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet is deprecated. Please import `from diffusers import StableDiffusionControlNetPipeline` instead.',
    standard_warn=False,
    stacklevel=3,
)




