


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


_lowercase							:   Optional[int]						=						logging.get_logger(__name__)

# General docstring
_lowercase							:   Any						=						"RegNetConfig"

# Base docstring
_lowercase							:   Dict						=						"facebook/regnet-y-040"
_lowercase							:   Any						=						[1, 1088, 7, 7]

# Image classification docstring
_lowercase							:   Tuple						=						"facebook/regnet-y-040"
_lowercase							:   Union[str, Any]						=						"tabby, tabby cat"

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






class    _UpperCamelCase     (       nn.Module					):

				"""simple docstring"""







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

				def   _UpperCAmelCase       (    self					, a__ )						->					str:
					A             =					self.convolution(a__ )
					A             =					self.normalization(a__ )
					A             =					self.activation(a__ )
					return hidden_state






class    _UpperCamelCase     (       nn.Module					):

				"""simple docstring"""







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

				def   _UpperCAmelCase       (    self					, a__ )						->					Dict:
					A             =					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.""" )
					A             =					self.embedder(a__ )
					return hidden_state






class    _UpperCamelCase     (       nn.Module					):

				"""simple docstring"""







				def __init__(    self					, a__					, a__					, a__ = 2 )						->					List[Any]:
					super().__init__()
					A             =					nn.Convad(a__					, a__					, kernel_size=1					, stride=a__					, bias=a__ )
					A             =					nn.BatchNormad(a__ )

				def   _UpperCAmelCase       (    self					, a__ )						->					Tensor:
					A             =					self.convolution(a__ )
					A             =					self.normalization(a__ )
					return hidden_state






class    _UpperCamelCase     (       nn.Module					):

				"""simple docstring"""







				def __init__(    self					, a__					, a__ )						->					int:
					super().__init__()

					A             =					nn.AdaptiveAvgPoolad((1, 1) )
					A             =					nn.Sequential(
					    nn.Convad(a__					, a__					, kernel_size=1 )					, nn.ReLU()					, nn.Convad(a__					, a__					, kernel_size=1 )					, nn.Sigmoid()					, )

				def   _UpperCAmelCase       (    self					, a__ )						->					Tuple:
					# b c h w -> b c 1 1
					A             =					self.pooler(a__ )
					A             =					self.attention(a__ )
					A             =					hidden_state * attention
					return hidden_state






class    _UpperCamelCase     (       nn.Module					):

				"""simple docstring"""







				def __init__(    self					, a__					, a__					, a__					, a__ = 1 )						->					Optional[Any]:
					super().__init__()
					A             =					in_channels != out_channels or stride != 1
					A             =					max(1					, out_channels // config.groups_width )
					A             =					(
					    RegNetShortCut(a__					, a__					, stride=a__ ) if should_apply_shortcut else nn.Identity()
					)
					A             =					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__ )					, )
					A             =					ACTaFN[config.hidden_act]

				def   _UpperCAmelCase       (    self					, a__ )						->					Optional[Any]:
					A             =					hidden_state
					A             =					self.layer(a__ )
					A             =					self.shortcut(a__ )
					hidden_state += residual
					A             =					self.activation(a__ )
					return hidden_state






class    _UpperCamelCase     (       nn.Module					):

				"""simple docstring"""







				def __init__(    self					, a__					, a__					, a__					, a__ = 1 )						->					Optional[int]:
					super().__init__()
					A             =					in_channels != out_channels or stride != 1
					A             =					max(1					, out_channels // config.groups_width )
					A             =					(
					    RegNetShortCut(a__					, a__					, stride=a__ ) if should_apply_shortcut else nn.Identity()
					)
					A             =					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__ )					, )
					A             =					ACTaFN[config.hidden_act]

				def   _UpperCAmelCase       (    self					, a__ )						->					Optional[Any]:
					A             =					hidden_state
					A             =					self.layer(a__ )
					A             =					self.shortcut(a__ )
					hidden_state += residual
					A             =					self.activation(a__ )
					return hidden_state






class    _UpperCamelCase     (       nn.Module					):

				"""simple docstring"""







				def __init__(    self					, a__					, a__					, a__					, a__ = 2					, a__ = 2					, )						->					Any:
					super().__init__()

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

					A             =					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   _UpperCAmelCase       (    self					, a__ )						->					Union[str, Any]:
					A             =					self.layers(a__ )
					return hidden_state






class    _UpperCamelCase     (       nn.Module					):

				"""simple docstring"""







				def __init__(    self					, a__ )						->					Union[str, Any]:
					super().__init__()
					A             =					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]					, ) )
					A             =					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   _UpperCAmelCase       (    self					, a__					, a__ = False					, a__ = True )						->					BaseModelOutputWithNoAttention:
					A             =					() if output_hidden_states else None

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

						A             =					stage_module(a__ )

					if output_hidden_states:
						A             =					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    _UpperCamelCase     (       __snake_case					):

				"""simple docstring"""







				lowerCAmelCase					=			RegNetConfig
				lowerCAmelCase					=			'regnet'
				lowerCAmelCase					=			'pixel_values'
				lowerCAmelCase					=			True

				def   _UpperCAmelCase       (    self					, a__ )						->					Optional[Any]:
					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   _UpperCAmelCase       (    self					, a__					, a__=False )						->					str:
					if isinstance(a__					, a__ ):
						A             =					value
_lowercase							:   str						=						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"

_lowercase							:   int						=						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    _UpperCamelCase     (       __snake_case					):

				"""simple docstring"""







				def __init__(    self					, a__ )						->					List[str]:
					super().__init__(a__ )
					A             =					config
					A             =					RegNetEmbeddings(a__ )
					A             =					RegNetEncoder(a__ )
					A             =					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   _UpperCAmelCase       (    self					, a__					, a__ = None					, a__ = None )						->					BaseModelOutputWithPoolingAndNoAttention:
					A             =					(
					    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
					)
					A             =					return_dict if return_dict is not None else self.config.use_return_dict

					A             =					self.embedder(a__ )

					A             =					self.encoder(
					    a__					, output_hidden_states=a__					, return_dict=a__ )

					A             =					encoder_outputs[0]

					A             =					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    '     ,		__snake_case     ,		)
# Copied from transformers.models.resnet.modeling_resnet.ResNetForImageClassification with RESNET->REGNET,ResNet->RegNet,resnet->regnet
class    _UpperCamelCase     (       __snake_case					):

				"""simple docstring"""







				def __init__(    self					, a__ )						->					List[Any]:
					super().__init__(a__ )
					A             =					config.num_labels
					A             =					RegNetModel(a__ )
					# classification head
					A             =					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   _UpperCAmelCase       (    self					, a__ = None					, a__ = None					, a__ = None					, a__ = None					, )						->					ImageClassifierOutputWithNoAttention:
					A             =					return_dict if return_dict is not None else self.config.use_return_dict

					A             =					self.regnet(a__					, output_hidden_states=a__					, return_dict=a__ )

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

					A             =					self.classifier(a__ )

					A             =					None

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

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

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







import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionAttendAndExcitePipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_numpy, skip_mps, slow
from diffusers.utils.testing_utils import require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


_lowercase							:   Optional[int]						=						False






@skip_mps
class    _UpperCamelCase     (       __snake_case     ,		__snake_case     ,		__snake_case     ,		unittest.TestCase					):

				"""simple docstring"""







				lowerCAmelCase					=			StableDiffusionAttendAndExcitePipeline
				lowerCAmelCase					=			False
				lowerCAmelCase					=			TEXT_TO_IMAGE_PARAMS
				lowerCAmelCase					=			TEXT_TO_IMAGE_BATCH_PARAMS.union({'token_indices'}					)
				lowerCAmelCase					=			TEXT_TO_IMAGE_IMAGE_PARAMS
				lowerCAmelCase					=			TEXT_TO_IMAGE_IMAGE_PARAMS

				@classmethod
				def   _UpperCAmelCase       (    cls )						->					List[Any]:
					super().setUpClass()
					torch.use_deterministic_algorithms(a__ )

				@classmethod
				def   _UpperCAmelCase       (    cls )						->					Tuple:
					super().tearDownClass()
					torch.use_deterministic_algorithms(a__ )

				def   _UpperCAmelCase       (    self )						->					Dict:
					torch.manual_seed(0 )
					A             =					UNetaDConditionModel(
					    block_out_channels=(32, 64)					, layers_per_block=1					, sample_size=32					, in_channels=4					, out_channels=4					, down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")					, up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")					, cross_attention_dim=32					, attention_head_dim=(2, 4)					, use_linear_projection=a__					, )
					A             =					DDIMScheduler(
					    beta_start=0.0_00_85					, beta_end=0.0_12					, beta_schedule="""scaled_linear"""					, clip_sample=a__					, set_alpha_to_one=a__					, )
					torch.manual_seed(0 )
					A             =					AutoencoderKL(
					    block_out_channels=[32, 64]					, in_channels=3					, out_channels=3					, down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]					, up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]					, latent_channels=4					, sample_size=128					, )
					torch.manual_seed(0 )
					A             =					CLIPTextConfig(
					    bos_token_id=0					, eos_token_id=2					, hidden_size=32					, intermediate_size=37					, layer_norm_eps=1e-05					, num_attention_heads=4					, num_hidden_layers=5					, pad_token_id=1					, vocab_size=1000					, hidden_act="""gelu"""					, projection_dim=512					, )
					A             =					CLIPTextModel(a__ )
					A             =					CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""" )

					A             =					{
					    """unet""": unet,
					    """scheduler""": scheduler,
					    """vae""": vae,
					    """text_encoder""": text_encoder,
					    """tokenizer""": tokenizer,
					    """safety_checker""": None,
					    """feature_extractor""": None,
					}

					return components

				def   _UpperCAmelCase       (    self					, a__					, a__=0 )						->					Optional[Any]:
					if str(a__ ).startswith("""mps""" ):
						A             =					torch.manual_seed(a__ )
					else:
						A             =					torch.Generator(device=a__ ).manual_seed(a__ )
					A             =					A             =					{
					    """prompt""": """a cat and a frog""",
					    """token_indices""": [2, 5],
					    """generator""": generator,
					    """num_inference_steps""": 1,
					    """guidance_scale""": 6.0,
					    """output_type""": """numpy""",
					    """max_iter_to_alter""": 2,
					    """thresholds""": {0: 0.7},
					}
					return inputs

				def   _UpperCAmelCase       (    self )						->					Union[str, Any]:
					A             =					"""cpu"""

					A             =					self.get_dummy_components()
					A             =					self.pipeline_class(**a__ )
					pipe.to(a__ )
					pipe.set_progress_bar_config(disable=a__ )

					A             =					self.get_dummy_inputs(a__ )
					A             =					pipe(**a__ ).images
					A             =					image[0, -3:, -3:, -1]

					self.assertEqual(image.shape					, (1, 64, 64, 3) )
					A             =					np.array(
					    [0.63_90_53_64, 0.62_89_73_07, 0.48_59_90_17, 0.5_13_36_24, 0.5_55_00_48, 0.45_76_95_16, 0.50_32_69_73, 0.5_02_31_39, 0.45_38_44_96] )
					A             =					np.abs(image_slice.flatten() - expected_slice ).max()
					self.assertLessEqual(a__					, 1e-3 )

				def   _UpperCAmelCase       (    self )						->					List[Any]:
					super().test_cpu_offload_forward_pass(expected_max_diff=5e-4 )

				def   _UpperCAmelCase       (    self )						->					Dict:
					# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
					self._test_inference_batch_consistent(batch_sizes=[1, 2] )

				def   _UpperCAmelCase       (    self )						->					Any:
					self._test_inference_batch_single_identical(batch_size=2					, expected_max_diff=7e-4 )

				def   _UpperCAmelCase       (    self )						->					Optional[Any]:
					super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3 )

				def   _UpperCAmelCase       (    self )						->					str:
					super().test_pt_np_pil_outputs_equivalent(expected_max_diff=5e-4 )

				def   _UpperCAmelCase       (    self )						->					int:
					super().test_save_load_local(expected_max_difference=5e-4 )

				def   _UpperCAmelCase       (    self )						->					Optional[Any]:
					super().test_save_load_optional_components(expected_max_difference=4e-4 )






@require_torch_gpu
@slow
class    _UpperCamelCase     (       unittest.TestCase					):

				"""simple docstring"""







				@classmethod
				def   _UpperCAmelCase       (    cls )						->					Tuple:
					super().setUpClass()
					torch.use_deterministic_algorithms(a__ )

				@classmethod
				def   _UpperCAmelCase       (    cls )						->					Dict:
					super().tearDownClass()
					torch.use_deterministic_algorithms(a__ )

				def   _UpperCAmelCase       (    self )						->					Optional[Any]:
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()

				def   _UpperCAmelCase       (    self )						->					int:
					A             =					torch.manual_seed(51 )

					A             =					StableDiffusionAttendAndExcitePipeline.from_pretrained(
					    """CompVis/stable-diffusion-v1-4"""					, safety_checker=a__					, torch_dtype=torch.floataa )
					pipe.to("""cuda""" )

					A             =					"""a painting of an elephant with glasses"""
					A             =					[5, 7]

					A             =					pipe(
					    prompt=a__					, token_indices=a__					, guidance_scale=7.5					, generator=a__					, num_inference_steps=5					, max_iter_to_alter=5					, output_type="""numpy"""					, ).images[0]

					A             =					load_numpy(
					    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/attend-and-excite/elephant_glasses.npy""" )
					assert np.abs((expected_image - image).max() ) < 5e-1



