

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_0_8_8, 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 	lowercase    (				nn.Module							):


  """simple docstring"""




  def __init__(				self	,       __snake_case	,       __snake_case	,       __snake_case = 3	,       __snake_case = 1	,       __snake_case = 1	,       __snake_case = "relu"	,       ):
   super().__init__()
   _UpperCamelCase   :		List[str]			=						nn.Convad(
       __snake_case	,       __snake_case	,       kernel_size=__snake_case	,       stride=__snake_case	,       padding=kernel_size // 2	,       groups=__snake_case	,       bias=__snake_case	,       )
   _UpperCamelCase   :		Any			=						nn.BatchNormad(__snake_case)
   _UpperCamelCase   :		Tuple			=						ACTaFN[activation] if activation is not None else nn.Identity()






  def 		A__					(				self	,       __snake_case):
   _UpperCamelCase   :		List[Any]			=						self.convolution(__snake_case)
   _UpperCamelCase   :		Union[str, Any]			=						self.normalization(__snake_case)
   _UpperCamelCase   :		str			=						self.activation(__snake_case)
   return hidden_state







class 	lowercase    (				nn.Module							):


  """simple docstring"""




  def __init__(				self	,       __snake_case):
   super().__init__()
   _UpperCamelCase   :		str			=						RegNetConvLayer(
       config.num_channels	,       config.embedding_size	,       kernel_size=3	,       stride=2	,       activation=config.hidden_act)
   _UpperCamelCase   :		str			=						config.num_channels






  def 		A__					(				self	,       __snake_case):
   _UpperCamelCase   :		List[str]			=						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.')
   _UpperCamelCase   :		List[str]			=						self.embedder(__snake_case)
   return hidden_state







class 	lowercase    (				nn.Module							):


  """simple docstring"""




  def __init__(				self	,       __snake_case	,       __snake_case	,       __snake_case = 2):
   super().__init__()
   _UpperCamelCase   :		int			=						nn.Convad(__snake_case	,       __snake_case	,       kernel_size=1	,       stride=__snake_case	,       bias=__snake_case)
   _UpperCamelCase   :		Optional[Any]			=						nn.BatchNormad(__snake_case)






  def 		A__					(				self	,       __snake_case):
   _UpperCamelCase   :		str			=						self.convolution(__snake_case)
   _UpperCamelCase   :		int			=						self.normalization(__snake_case)
   return hidden_state







class 	lowercase    (				nn.Module							):


  """simple docstring"""




  def __init__(				self	,       __snake_case	,       __snake_case):
   super().__init__()

   _UpperCamelCase   :		Optional[int]			=						nn.AdaptiveAvgPoolad((1, 1))
   _UpperCamelCase   :		List[str]			=						nn.Sequential(
       nn.Convad(__snake_case	,       __snake_case	,       kernel_size=1)	,       nn.ReLU()	,       nn.Convad(__snake_case	,       __snake_case	,       kernel_size=1)	,       nn.Sigmoid()	,       )






  def 		A__					(				self	,       __snake_case):
   _UpperCamelCase   :		str			=						self.pooler(__snake_case)
   _UpperCamelCase   :		Any			=						self.attention(__snake_case)
   _UpperCamelCase   :		List[str]			=						hidden_state * attention
   return hidden_state







class 	lowercase    (				nn.Module							):


  """simple docstring"""




  def __init__(				self	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case = 1):
   super().__init__()
   _UpperCamelCase   :		Dict			=						in_channels != out_channels or stride != 1
   _UpperCamelCase   :		Optional[int]			=						max(1	,       out_channels // config.groups_width)
   _UpperCamelCase   :		Optional[int]			=						(
       RegNetShortCut(__snake_case	,       __snake_case	,       stride=__snake_case) if should_apply_shortcut else nn.Identity()
   )
   _UpperCamelCase   :		Union[str, Any]			=						nn.Sequential(
       RegNetConvLayer(__snake_case	,       __snake_case	,       kernel_size=1	,       activation=config.hidden_act)	,       RegNetConvLayer(__snake_case	,       __snake_case	,       stride=__snake_case	,       groups=__snake_case	,       activation=config.hidden_act)	,       RegNetConvLayer(__snake_case	,       __snake_case	,       kernel_size=1	,       activation=__snake_case)	,       )
   _UpperCamelCase   :		List[Any]			=						ACTaFN[config.hidden_act]






  def 		A__					(				self	,       __snake_case):
   _UpperCamelCase   :		List[Any]			=						hidden_state
   _UpperCamelCase   :		List[Any]			=						self.layer(__snake_case)
   _UpperCamelCase   :		Optional[int]			=						self.shortcut(__snake_case)
   hidden_state += residual
   _UpperCamelCase   :		Any			=						self.activation(__snake_case)
   return hidden_state







class 	lowercase    (				nn.Module							):


  """simple docstring"""




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






  def 		A__					(				self	,       __snake_case):
   _UpperCamelCase   :		Tuple			=						hidden_state
   _UpperCamelCase   :		int			=						self.layer(__snake_case)
   _UpperCamelCase   :		int			=						self.shortcut(__snake_case)
   hidden_state += residual
   _UpperCamelCase   :		Optional[int]			=						self.activation(__snake_case)
   return hidden_state







class 	lowercase    (				nn.Module							):


  """simple docstring"""




  def __init__(				self	,       __snake_case	,       __snake_case	,       __snake_case	,       __snake_case = 2	,       __snake_case = 2	,       ):
   super().__init__()

   _UpperCamelCase   :		Dict			=						RegNetXLayer if config.layer_type == 'x' else RegNetYLayer

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






  def 		A__					(				self	,       __snake_case):
   _UpperCamelCase   :		List[Any]			=						self.layers(__snake_case)
   return hidden_state







class 	lowercase    (				nn.Module							):


  """simple docstring"""




  def __init__(				self	,       __snake_case):
   super().__init__()
   _UpperCamelCase   :		Optional[int]			=						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(
           __snake_case	,       config.embedding_size	,       config.hidden_sizes[0]	,       stride=2 if config.downsample_in_first_stage else 1	,       depth=config.depths[0]	,       ))
   _UpperCamelCase   :		str			=						zip(config.hidden_sizes	,       config.hidden_sizes[1:])
   for (in_channels, out_channels), depth in zip(__snake_case	,       config.depths[1:]):
    self.stages.append(RegNetStage(__snake_case	,       __snake_case	,       __snake_case	,       depth=__snake_case))






  def 		A__					(				self	,       __snake_case	,       __snake_case = False	,       __snake_case = True):
   _UpperCamelCase   :		List[Any]			=						() if output_hidden_states else None

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

    _UpperCamelCase   :		List[str]			=						stage_module(__snake_case)

   if output_hidden_states:
    _UpperCamelCase   :		int			=						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=__snake_case	,       hidden_states=__snake_case)







class 	lowercase    (				_lowercase							):


  """simple docstring"""




  a__					       =  RegNetConfig
  a__					       =  """regnet"""
  a__					       =  """pixel_values"""
  a__					       =  True


  def 		A__					(				self	,       __snake_case):

   if isinstance(__snake_case	,       nn.Convad):
    nn.init.kaiming_normal_(module.weight	,       mode='fan_out'	,       nonlinearity='relu')
   elif isinstance(__snake_case	,       (nn.BatchNormad, nn.GroupNorm)):
    nn.init.constant_(module.weight	,       1)
    nn.init.constant_(module.bias	,       0)






  def 		A__					(				self	,       __snake_case	,       __snake_case=False):

   if isinstance(__snake_case	,       __snake_case):
    _UpperCamelCase   :		Optional[Any]			=						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."    ,							_lowercase    ,							)
# Copied from transformers.models.resnet.modeling_resnet.ResNetModel with RESNET->REGNET,ResNet->RegNet
class 	lowercase    (				_lowercase							):


  """simple docstring"""




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






  @add_start_docstrings_to_model_forward(__snake_case)
  @add_code_sample_docstrings(
      checkpoint=_CHECKPOINT_FOR_DOC	,       output_type=__snake_case	,       config_class=_CONFIG_FOR_DOC	,       modality='vision'	,       expected_output=_EXPECTED_OUTPUT_SHAPE	,       )
  def 		A__					(				self	,       __snake_case	,       __snake_case = None	,       __snake_case = None):
   _UpperCamelCase   :		Union[str, Any]			=						(
       output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
   )
   _UpperCamelCase   :		List[str]			=						return_dict if return_dict is not None else self.config.use_return_dict

   _UpperCamelCase   :		Dict			=						self.embedder(__snake_case)

   _UpperCamelCase   :		Any			=						self.encoder(
       __snake_case	,       output_hidden_states=__snake_case	,       return_dict=__snake_case)

   _UpperCamelCase   :		str			=						encoder_outputs[0]

   _UpperCamelCase   :		Union[str, Any]			=						self.pooler(__snake_case)

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

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


  """simple docstring"""




  def __init__(				self	,       __snake_case):
   super().__init__(__snake_case)
   _UpperCamelCase   :		Tuple			=						config.num_labels
   _UpperCamelCase   :		int			=						RegNetModel(__snake_case)
   # classification head
   _UpperCamelCase   :		str			=						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(__snake_case)
  @add_code_sample_docstrings(
      checkpoint=_IMAGE_CLASS_CHECKPOINT	,       output_type=__snake_case	,       config_class=_CONFIG_FOR_DOC	,       expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT	,       )
  def 		A__					(				self	,       __snake_case = None	,       __snake_case = None	,       __snake_case = None	,       __snake_case = None	,       ):
   _UpperCamelCase   :		Any			=						return_dict if return_dict is not None else self.config.use_return_dict

   _UpperCamelCase   :		Optional[Any]			=						self.regnet(__snake_case	,       output_hidden_states=__snake_case	,       return_dict=__snake_case)

   _UpperCamelCase   :		int			=						outputs.pooler_output if return_dict else outputs[1]

   _UpperCamelCase   :		Optional[int]			=						self.classifier(__snake_case)

   _UpperCamelCase   :		Any			=						None

   if labels is not None:
    if self.config.problem_type is None:
     if self.num_labels == 1:
      _UpperCamelCase   :		int			=						'regression'
     elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
      _UpperCamelCase   :		Union[str, Any]			=						'single_label_classification'
     else:
      _UpperCamelCase   :		Dict			=						'multi_label_classification'
    if self.config.problem_type == "regression":
     _UpperCamelCase   :		List[Any]			=						MSELoss()
     if self.num_labels == 1:
      _UpperCamelCase   :		List[Any]			=						loss_fct(logits.squeeze()	,       labels.squeeze())
     else:
      _UpperCamelCase   :		Union[str, Any]			=						loss_fct(__snake_case	,       __snake_case)
    elif self.config.problem_type == "single_label_classification":
     _UpperCamelCase   :		List[Any]			=						CrossEntropyLoss()
     _UpperCamelCase   :		List[Any]			=						loss_fct(logits.view(-1	,       self.num_labels)	,       labels.view(-1))
    elif self.config.problem_type == "multi_label_classification":
     _UpperCamelCase   :		List[str]			=						BCEWithLogitsLoss()
     _UpperCamelCase   :		Any			=						loss_fct(__snake_case	,       __snake_case)

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

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







from ...processing_utils import ProcessorMixin







class 	lowercase    (				_lowercase							):


		"""simple docstring"""




		a__					       =  ["image_processor", "feature_extractor"]
		a__					       =  "TvltImageProcessor"
		a__					       =  "TvltFeatureExtractor"


		def __init__(				self	,       __snake_case	,       __snake_case):
			super().__init__(image_processor=__snake_case	,       feature_extractor=__snake_case)

			_UpperCamelCase   :		List[str]			=						image_processor
			_UpperCamelCase   :		Dict			=						feature_extractor


		def __call__(				self	,       __snake_case=None	,       __snake_case=None	,       __snake_case=None	,       __snake_case=None	,       __snake_case=False	,       __snake_case=False	,       *__snake_case	,       **__snake_case	,       ):

			if images is None and audio is None:
				raise ValueError('You need to specify either an `images` or `audio` input to process.')

			_UpperCamelCase   :		Union[str, Any]			=						None
			if images is not None:
				_UpperCamelCase   :		Tuple			=						self.image_processor(__snake_case	,       mask_pixel=__snake_case	,       *__snake_case	,       **__snake_case)
			if images_mixed is not None:
				_UpperCamelCase   :		Union[str, Any]			=						self.image_processor(__snake_case	,       is_mixed=__snake_case	,       *__snake_case	,       **__snake_case)
			if audio is not None:
				_UpperCamelCase   :		Tuple			=						self.feature_extractor(
				    __snake_case	,       *__snake_case	,       sampling_rate=__snake_case	,       mask_audio=__snake_case	,       **__snake_case)

			_UpperCamelCase   :		Tuple			=						{}
			if audio is not None:
				output_dict.update(__snake_case)
			if images is not None:
				output_dict.update(__snake_case)
			if images_mixed_dict is not None:
				output_dict.update(__snake_case)
			return output_dict






		@property
		def 		A__					(				self):
			_UpperCamelCase   :		List[Any]			=						self.image_processor.model_input_names
			_UpperCamelCase   :		List[Any]			=						self.feature_extractor.model_input_names
			return list(dict.fromkeys(image_processor_input_names + feature_extractor_input_names))




