
'''simple docstring'''
from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
from ...modeling_outputs import (
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import logging
from .configuration_regnet import RegNetConfig


__lowerCAmelCase												=	logging.get_logger(__name__)

# General docstring
__lowerCAmelCase												=	'RegNetConfig'

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

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

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


class       _lowerCAmelCase   (   nn.Module				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase = 3  ,  UpperCAmelCase = 1  ,  UpperCAmelCase = 1  ,  UpperCAmelCase = "relu"  ,  )					->       int:
														super().__init__()
														_snake_case       =  nn.Convad(
														    UpperCAmelCase  ,  UpperCAmelCase  ,  kernel_size=UpperCAmelCase  ,  stride=UpperCAmelCase  ,  padding=kernel_size // 2  ,  groups=UpperCAmelCase  ,  bias=UpperCAmelCase  ,  )
														_snake_case       =  nn.BatchNormad(UpperCAmelCase  )
														_snake_case       =  ACTaFN[activation] if activation is not None else nn.Identity()





							def 	lowercase							(self  ,  UpperCAmelCase  )					->       Optional[int]:
														_snake_case       =  self.convolution(UpperCAmelCase  )
														_snake_case       =  self.normalization(UpperCAmelCase  )
														_snake_case       =  self.activation(UpperCAmelCase  )
														return hidden_state


class       _lowerCAmelCase   (   nn.Module				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  )					->       Any:
														super().__init__()
														_snake_case       =  RegNetConvLayer(
														    config.num_channels  ,  config.embedding_size  ,  kernel_size=3  ,  stride=2  ,  activation=config.hidden_act  )
														_snake_case       =  config.num_channels





							def 	lowercase							(self  ,  UpperCAmelCase  )					->       Any:
														_snake_case       =  pixel_values.shape[1]
														if 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."""  )
														_snake_case       =  self.embedder(UpperCAmelCase  )
														return hidden_state


class       _lowerCAmelCase   (   nn.Module				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase = 2  )					->       int:
														super().__init__()
														_snake_case       =  nn.Convad(UpperCAmelCase  ,  UpperCAmelCase  ,  kernel_size=1  ,  stride=UpperCAmelCase  ,  bias=UpperCAmelCase  )
														_snake_case       =  nn.BatchNormad(UpperCAmelCase  )





							def 	lowercase							(self  ,  UpperCAmelCase  )					->       Tensor:
														_snake_case       =  self.convolution(UpperCAmelCase  )
														_snake_case       =  self.normalization(UpperCAmelCase  )
														return hidden_state


class       _lowerCAmelCase   (   nn.Module				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  ,  UpperCAmelCase  )					->       Tuple:
														super().__init__()

														_snake_case       =  nn.AdaptiveAvgPoolad((1, 1)  )
														_snake_case       =  nn.Sequential(
														    nn.Convad(UpperCAmelCase  ,  UpperCAmelCase  ,  kernel_size=1  )  ,  nn.ReLU()  ,  nn.Convad(UpperCAmelCase  ,  UpperCAmelCase  ,  kernel_size=1  )  ,  nn.Sigmoid()  ,  )





							def 	lowercase							(self  ,  UpperCAmelCase  )					->       int:
														# b c h w -> b c 1 1
														_snake_case       =  self.pooler(UpperCAmelCase  )
														_snake_case       =  self.attention(UpperCAmelCase  )
														_snake_case       =  hidden_state * attention
														return hidden_state


class       _lowerCAmelCase   (   nn.Module				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase = 1  )					->       Optional[Any]:
														super().__init__()
														_snake_case       =  in_channels != out_channels or stride != 1
														_snake_case       =  max(1  ,  out_channels // config.groups_width  )
														_snake_case       =  (
														    RegNetShortCut(UpperCAmelCase  ,  UpperCAmelCase  ,  stride=UpperCAmelCase  ) if should_apply_shortcut else nn.Identity()
														)
														_snake_case       =  nn.Sequential(
														    RegNetConvLayer(UpperCAmelCase  ,  UpperCAmelCase  ,  kernel_size=1  ,  activation=config.hidden_act  )  ,  RegNetConvLayer(UpperCAmelCase  ,  UpperCAmelCase  ,  stride=UpperCAmelCase  ,  groups=UpperCAmelCase  ,  activation=config.hidden_act  )  ,  RegNetConvLayer(UpperCAmelCase  ,  UpperCAmelCase  ,  kernel_size=1  ,  activation=UpperCAmelCase  )  ,  )
														_snake_case       =  ACTaFN[config.hidden_act]





							def 	lowercase							(self  ,  UpperCAmelCase  )					->       Tuple:
														_snake_case       =  hidden_state
														_snake_case       =  self.layer(UpperCAmelCase  )
														_snake_case       =  self.shortcut(UpperCAmelCase  )
														hidden_state += residual
														_snake_case       =  self.activation(UpperCAmelCase  )
														return hidden_state


class       _lowerCAmelCase   (   nn.Module				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase = 1  )					->       str:
														super().__init__()
														_snake_case       =  in_channels != out_channels or stride != 1
														_snake_case       =  max(1  ,  out_channels // config.groups_width  )
														_snake_case       =  (
														    RegNetShortCut(UpperCAmelCase  ,  UpperCAmelCase  ,  stride=UpperCAmelCase  ) if should_apply_shortcut else nn.Identity()
														)
														_snake_case       =  nn.Sequential(
														    RegNetConvLayer(UpperCAmelCase  ,  UpperCAmelCase  ,  kernel_size=1  ,  activation=config.hidden_act  )  ,  RegNetConvLayer(UpperCAmelCase  ,  UpperCAmelCase  ,  stride=UpperCAmelCase  ,  groups=UpperCAmelCase  ,  activation=config.hidden_act  )  ,  RegNetSELayer(UpperCAmelCase  ,  reduced_channels=int(round(in_channels / 4  )  )  )  ,  RegNetConvLayer(UpperCAmelCase  ,  UpperCAmelCase  ,  kernel_size=1  ,  activation=UpperCAmelCase  )  ,  )
														_snake_case       =  ACTaFN[config.hidden_act]





							def 	lowercase							(self  ,  UpperCAmelCase  )					->       List[str]:
														_snake_case       =  hidden_state
														_snake_case       =  self.layer(UpperCAmelCase  )
														_snake_case       =  self.shortcut(UpperCAmelCase  )
														hidden_state += residual
														_snake_case       =  self.activation(UpperCAmelCase  )
														return hidden_state


class       _lowerCAmelCase   (   nn.Module				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase = 2  ,  UpperCAmelCase = 2  ,  )					->       int:
														super().__init__()

														_snake_case       =  RegNetXLayer if config.layer_type == """x""" else RegNetYLayer

														_snake_case       =  nn.Sequential(
														    # downsampling is done in the first layer with stride of 2
														    layer(
														        UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  stride=UpperCAmelCase  ,  )  ,  *[layer(UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ) for _ in range(depth - 1  )]  ,  )





							def 	lowercase							(self  ,  UpperCAmelCase  )					->       Union[str, Any]:
														_snake_case       =  self.layers(UpperCAmelCase  )
														return hidden_state


class       _lowerCAmelCase   (   nn.Module				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  )					->       Union[str, Any]:
														super().__init__()
														_snake_case       =  nn.ModuleList([]  )
														# based on `downsample_in_first_stage`, the first layer of the first stage may or may not downsample the input
														self.stages.append(
														    RegNetStage(
														        UpperCAmelCase  ,  config.embedding_size  ,  config.hidden_sizes[0]  ,  stride=2 if config.downsample_in_first_stage else 1  ,  depth=config.depths[0]  ,  )  )
														_snake_case       =  zip(config.hidden_sizes  ,  config.hidden_sizes[1:]  )
														for (in_channels, out_channels), depth in zip(UpperCAmelCase  ,  config.depths[1:]  ):
																					self.stages.append(RegNetStage(UpperCAmelCase  ,  UpperCAmelCase  ,  UpperCAmelCase  ,  depth=UpperCAmelCase  )  )





							def 	lowercase							(self  ,  UpperCAmelCase  ,  UpperCAmelCase = False  ,  UpperCAmelCase = True  )					->       BaseModelOutputWithNoAttention:
														_snake_case       =  () if output_hidden_states else None

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

																					_snake_case       =  stage_module(UpperCAmelCase  )

														if output_hidden_states:
																					_snake_case       =  hidden_states + (hidden_state,)

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

														return BaseModelOutputWithNoAttention(last_hidden_state=UpperCAmelCase  ,  hidden_states=UpperCAmelCase  )


class       _lowerCAmelCase   (   __snake_case				):







							'''simple docstring'''

							lowerCAmelCase_ 						=						RegNetConfig
							lowerCAmelCase_ 						=						"regnet"
							lowerCAmelCase_ 						=						"pixel_values"
							lowerCAmelCase_ 						=						True

							def 	lowercase							(self  ,  UpperCAmelCase  )					->       Any:
														if isinstance(UpperCAmelCase  ,  nn.Convad  ):
																					nn.init.kaiming_normal_(module.weight  ,  mode="""fan_out"""  ,  nonlinearity="""relu"""  )
														elif isinstance(UpperCAmelCase  ,  (nn.BatchNormad, nn.GroupNorm)  ):
																					nn.init.constant_(module.weight  ,  1  )
																					nn.init.constant_(module.bias  ,  0  )





							def 	lowercase							(self  ,  UpperCAmelCase  ,  UpperCAmelCase=False  )					->       int:
														if isinstance(UpperCAmelCase  ,  UpperCAmelCase  ):
																					_snake_case       =  value
__lowerCAmelCase												=	r'\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\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 [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'

__lowerCAmelCase												=	r'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\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 [`~file_utils.ModelOutput`] instead of a plain tuple.\n'


@add_start_docstrings(
    "The bare RegNet model outputting raw features without any specific head on top."					,							__snake_case					,							)
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class       _lowerCAmelCase   (   __snake_case				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  )					->       Optional[int]:
														super().__init__(UpperCAmelCase  )
														_snake_case       =  config
														_snake_case       =  RegNetEmbeddings(UpperCAmelCase  )
														_snake_case       =  RegNetEncoder(UpperCAmelCase  )
														_snake_case       =  nn.AdaptiveAvgPoolad((1, 1)  )
														# Initialize weights and apply final processing
														self.post_init()





							@add_start_docstrings_to_model_forward(UpperCAmelCase  )
							@add_code_sample_docstrings(
							    checkpoint=_CHECKPOINT_FOR_DOC  ,  output_type=UpperCAmelCase  ,  config_class=_CONFIG_FOR_DOC  ,  modality="""vision"""  ,  expected_output=_EXPECTED_OUTPUT_SHAPE  ,  )
							def 	lowercase							(self  ,  UpperCAmelCase  ,  UpperCAmelCase = None  ,  UpperCAmelCase = None  )					->       BaseModelOutputWithPoolingAndNoAttention:
														_snake_case       =  (
														    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
														)
														_snake_case       =  return_dict if return_dict is not None else self.config.use_return_dict

														_snake_case       =  self.embedder(UpperCAmelCase  )

														_snake_case       =  self.encoder(
														    UpperCAmelCase  ,  output_hidden_states=UpperCAmelCase  ,  return_dict=UpperCAmelCase  )

														_snake_case       =  encoder_outputs[0]

														_snake_case       =  self.pooler(UpperCAmelCase  )

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

														return BaseModelOutputWithPoolingAndNoAttention(
														    last_hidden_state=UpperCAmelCase  ,  pooler_output=UpperCAmelCase  ,  hidden_states=encoder_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    "					,							__snake_case					,							)
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class       _lowerCAmelCase   (   __snake_case				):







							'''simple docstring'''

							def __init__(self  ,  UpperCAmelCase  )					->       List[str]:
														super().__init__(UpperCAmelCase  )
														_snake_case       =  config.num_labels
														_snake_case       =  RegNetModel(UpperCAmelCase  )
														# classification head
														_snake_case       =  nn.Sequential(
														    nn.Flatten()  ,  nn.Linear(config.hidden_sizes[-1]  ,  config.num_labels  ) if config.num_labels > 0 else nn.Identity()  ,  )
														# initialize weights and apply final processing
														self.post_init()





							@add_start_docstrings_to_model_forward(UpperCAmelCase  )
							@add_code_sample_docstrings(
							    checkpoint=_IMAGE_CLASS_CHECKPOINT  ,  output_type=UpperCAmelCase  ,  config_class=_CONFIG_FOR_DOC  ,  expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT  ,  )
							def 	lowercase							(self  ,  UpperCAmelCase = None  ,  UpperCAmelCase = None  ,  UpperCAmelCase = None  ,  UpperCAmelCase = None  ,  )					->       ImageClassifierOutputWithNoAttention:
														_snake_case       =  return_dict if return_dict is not None else self.config.use_return_dict

														_snake_case       =  self.regnet(UpperCAmelCase  ,  output_hidden_states=UpperCAmelCase  ,  return_dict=UpperCAmelCase  )

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

														_snake_case       =  self.classifier(UpperCAmelCase  )

														_snake_case       =  None

														if labels is not None:
																					if self.config.problem_type is None:
																												if self.num_labels == 1:
																																			_snake_case       =  """regression"""
																												elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																																			_snake_case       =  """single_label_classification"""
																												else:
																																			_snake_case       =  """multi_label_classification"""
																					if self.config.problem_type == "regression":
																												_snake_case       =  MSELoss()
																												if self.num_labels == 1:
																																			_snake_case       =  loss_fct(logits.squeeze()  ,  labels.squeeze()  )
																												else:
																																			_snake_case       =  loss_fct(UpperCAmelCase  ,  UpperCAmelCase  )
																					elif self.config.problem_type == "single_label_classification":
																												_snake_case       =  CrossEntropyLoss()
																												_snake_case       =  loss_fct(logits.view(-1  ,  self.num_labels  )  ,  labels.view(-1  )  )
																					elif self.config.problem_type == "multi_label_classification":
																												_snake_case       =  BCEWithLogitsLoss()
																												_snake_case       =  loss_fct(UpperCAmelCase  ,  UpperCAmelCase  )

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

														return ImageClassifierOutputWithNoAttention(loss=UpperCAmelCase  ,  logits=UpperCAmelCase  ,  hidden_states=outputs.hidden_states  )

'''simple docstring'''
from ...utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_torch_available,
    is_transformers_available,
)


try:
				if not (is_transformers_available() and is_torch_available()):
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
				from .multicontrolnet import MultiControlNetModel
				from .pipeline_controlnet import StableDiffusionControlNetPipeline
				from .pipeline_controlnet_imgaimg import StableDiffusionControlNetImgaImgPipeline
				from .pipeline_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline


if is_transformers_available() and is_flax_available():
				from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline