





"""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 ):


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




						def     snake_case_   (self   ,							__a		) ->   Optional[Any]:
								UpperCamelCase   =	self.convolution(__a		)
								UpperCamelCase   =	self.normalization(__a		)
								UpperCamelCase   =	self.activation(__a		)
								return hidden_state





class        _lowerCamelCase      (       nn.Module ):


						def __init__(self   ,							__a		) ->   List[Any]:
								super().__init__()
								UpperCamelCase   =	RegNetConvLayer(
								    config.num_channels   ,							config.embedding_size   ,							kernel_size=3   ,							stride=2   ,							activation=config.hidden_act		)
								UpperCamelCase   =	config.num_channels




						def     snake_case_   (self   ,							__a		) ->   List[Any]:
								UpperCamelCase   =	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   =	self.embedder(__a		)
								return hidden_state





class        _lowerCamelCase      (       nn.Module ):


						def __init__(self   ,							__a   ,							__a   ,							__a = 2		) ->   Optional[Any]:
								super().__init__()
								UpperCamelCase   =	nn.Convad(__a   ,							__a   ,							kernel_size=1   ,							stride=__a   ,							bias=__a		)
								UpperCamelCase   =	nn.BatchNormad(__a		)




						def     snake_case_   (self   ,							__a		) ->   Tensor:
								UpperCamelCase   =	self.convolution(__a		)
								UpperCamelCase   =	self.normalization(__a		)
								return hidden_state





class        _lowerCamelCase      (       nn.Module ):


						def __init__(self   ,							__a   ,							__a		) ->   str:
								super().__init__()

								UpperCamelCase   =	nn.AdaptiveAvgPoolad((1, 1)		)
								UpperCamelCase   =	nn.Sequential(
								    nn.Convad(__a   ,							__a   ,							kernel_size=1		)   ,							nn.ReLU()   ,							nn.Convad(__a   ,							__a   ,							kernel_size=1		)   ,							nn.Sigmoid()   ,							)




						def     snake_case_   (self   ,							__a		) ->   int:
								# b c h w -> b c 1 1
								UpperCamelCase   =	self.pooler(__a		)
								UpperCamelCase   =	self.attention(__a		)
								UpperCamelCase   =	hidden_state * attention
								return hidden_state





class        _lowerCamelCase      (       nn.Module ):


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




						def     snake_case_   (self   ,							__a		) ->   int:
								UpperCamelCase   =	hidden_state
								UpperCamelCase   =	self.layer(__a		)
								UpperCamelCase   =	self.shortcut(__a		)
								hidden_state += residual
								UpperCamelCase   =	self.activation(__a		)
								return hidden_state





class        _lowerCamelCase      (       nn.Module ):


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




						def     snake_case_   (self   ,							__a		) ->   Any:
								UpperCamelCase   =	hidden_state
								UpperCamelCase   =	self.layer(__a		)
								UpperCamelCase   =	self.shortcut(__a		)
								hidden_state += residual
								UpperCamelCase   =	self.activation(__a		)
								return hidden_state





class        _lowerCamelCase      (       nn.Module ):


						def __init__(self   ,							__a   ,							__a   ,							__a   ,							__a = 2   ,							__a = 2   ,							) ->   Union[str, Any]:
								super().__init__()

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

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




						def     snake_case_   (self   ,							__a		) ->   Any:
								UpperCamelCase   =	self.layers(__a		)
								return hidden_state





class        _lowerCamelCase      (       nn.Module ):


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




						def     snake_case_   (self   ,							__a   ,							__a = False   ,							__a = True		) ->   BaseModelOutputWithNoAttention:
								UpperCamelCase   =	() if output_hidden_states else None

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

										UpperCamelCase   =	stage_module(__a		)

								if output_hidden_states:
										UpperCamelCase   =	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=__a   ,							hidden_states=__a		)





class        _lowerCamelCase      (       _lowercase ):
						UpperCAmelCase_ =			RegNetConfig
						UpperCAmelCase_ =			"regnet"
						UpperCAmelCase_ =			"pixel_values"
						UpperCAmelCase_ =			True


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




						def     snake_case_   (self   ,							__a   ,							__a=False		) ->   Dict:
								if isinstance(__a   ,							__a		):
										UpperCamelCase   =	value



lowerCAmelCase__     			=			r'''
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`RegNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
'''

lowerCAmelCase__     			=			r'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
'''





@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        _lowerCamelCase      (       _lowercase ):


						def __init__(self   ,							__a		) ->   str:
								super().__init__(__a		)
								UpperCamelCase   =	config
								UpperCamelCase   =	RegNetEmbeddings(__a		)
								UpperCamelCase   =	RegNetEncoder(__a		)
								UpperCamelCase   =	nn.AdaptiveAvgPoolad((1, 1)		)
								# Initialize weights and apply final processing
								self.post_init()




						@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     snake_case_   (self   ,							__a   ,							__a = None   ,							__a = None		) ->   BaseModelOutputWithPoolingAndNoAttention:
								UpperCamelCase   =	(
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								UpperCamelCase   =	return_dict if return_dict is not None else self.config.use_return_dict

								UpperCamelCase   =	self.embedder(__a		)

								UpperCamelCase   =	self.encoder(
								    __a   ,							output_hidden_states=__a   ,							return_dict=__a		)

								UpperCamelCase   =	encoder_outputs[0]

								UpperCamelCase   =	self.pooler(__a		)

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

								return BaseModelOutputWithPoolingAndNoAttention(
								    last_hidden_state=__a   ,							pooler_output=__a   ,							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        _lowerCamelCase      (       _lowercase ):


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

								UpperCamelCase   =	self.regnet(__a   ,							output_hidden_states=__a   ,							return_dict=__a		)

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

								UpperCamelCase   =	self.classifier(__a		)

								UpperCamelCase   =	None

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

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

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













"""simple docstring"""





import os

from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE, hf_cache_home


lowerCAmelCase__     			=			HUGGINGFACE_HUB_CACHE


lowerCAmelCase__     			=			'''config.json'''
lowerCAmelCase__     			=			'''diffusion_pytorch_model.bin'''
lowerCAmelCase__     			=			'''diffusion_flax_model.msgpack'''
lowerCAmelCase__     			=			'''model.onnx'''
lowerCAmelCase__     			=			'''diffusion_pytorch_model.safetensors'''
lowerCAmelCase__     			=			'''weights.pb'''
lowerCAmelCase__     			=			'''https://huggingface.co'''
lowerCAmelCase__     			=			default_cache_path
lowerCAmelCase__     			=			'''diffusers_modules'''
lowerCAmelCase__     			=			os.getenv('''HF_MODULES_CACHE''', os.path.join(hf_cache_home, '''modules'''))
lowerCAmelCase__     			=			['''fp16''', '''non-ema''']
lowerCAmelCase__     			=			'''.self_attn'''






