


import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class 		_snake_case  (      _snake_case						,    _snake_case						,    unittest.TestCase							):
 SCREAMING_SNAKE_CASE__      =       StableDiffusionDiffEditPipeline
 SCREAMING_SNAKE_CASE__      =       TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'height', 'width', 'image'} | {'image_latents'}
 SCREAMING_SNAKE_CASE__      =       TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {'image'} | {'image_latents'}
 SCREAMING_SNAKE_CASE__      =       frozenset(
     []							)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 SCREAMING_SNAKE_CASE__      =       frozenset([]							)




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       torch.manual_seed(0    )
       a						:Dict 							=       UNetaDConditionModel(
           block_out_channels=(32, 64)					,  layers_per_block=2					,  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=_lowerCamelCase					,  )
       a						:int 							=       DDIMScheduler(
           beta_start=0.0_0085					,  beta_end=0.012					,  beta_schedule='''scaled_linear'''					,  clip_sample=_lowerCamelCase					,  set_alpha_to_one=_lowerCamelCase					,  )
       a						:Tuple 							=       DDIMInverseScheduler(
           beta_start=0.0_0085					,  beta_end=0.012					,  beta_schedule='''scaled_linear'''					,  clip_sample=_lowerCamelCase					,  set_alpha_to_zero=_lowerCamelCase					,  )
       torch.manual_seed(0    )
       a						:Dict 							=       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						:int 							=       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						:Optional[int] 							=       CLIPTextModel(_lowerCamelCase    )
       a						:str 							=       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )

       a						:Optional[Any] 							=       {
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''inverse_scheduler''': inverse_scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }

       return components




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase					,  _lowerCamelCase=0    ):
       a						:Optional[Any] 							=       floats_tensor((1, 16, 16)					,  rng=random.Random(_lowerCamelCase    )    ).to(_lowerCamelCase    )
       a						:Tuple 							=       floats_tensor((1, 2, 4, 16, 16)					,  rng=random.Random(_lowerCamelCase    )    ).to(_lowerCamelCase    )
       if str(_lowerCamelCase    ).startswith('''mps'''    ):
             a						:int 							=       torch.manual_seed(_lowerCamelCase    )
       else:
             a						:int 							=       torch.Generator(device=_lowerCamelCase    ).manual_seed(_lowerCamelCase    )
       a						:Dict 							=       {
           '''prompt''': '''a dog and a newt''',
           '''mask_image''': mask,
           '''image_latents''': latents,
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''inpaint_strength''': 1.0,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
       }

       return inputs




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase					,  _lowerCamelCase=0    ):
       a						:str 							=       floats_tensor((1, 3, 32, 32)					,  rng=random.Random(_lowerCamelCase    )    ).to(_lowerCamelCase    )
       a						:str 							=       image.cpu().permute(0					,  2					,  3					,  1    )[0]
       a						:Optional[Any] 							=       Image.fromarray(np.uinta(_lowerCamelCase    )    ).convert('''RGB'''    )
       if str(_lowerCamelCase    ).startswith('''mps'''    ):
             a						:Dict 							=       torch.manual_seed(_lowerCamelCase    )
       else:
             a						:Any 							=       torch.Generator(device=_lowerCamelCase    ).manual_seed(_lowerCamelCase    )
       a						:Union[str, Any] 							=       {
           '''image''': image,
           '''source_prompt''': '''a cat and a frog''',
           '''target_prompt''': '''a dog and a newt''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''num_maps_per_mask''': 2,
           '''mask_encode_strength''': 1.0,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
       }

       return inputs




 def        SCREAMING_SNAKE_CASE__				(		self					,  _lowerCamelCase					,  _lowerCamelCase=0    ):
       a						:List[str] 							=       floats_tensor((1, 3, 32, 32)					,  rng=random.Random(_lowerCamelCase    )    ).to(_lowerCamelCase    )
       a						:int 							=       image.cpu().permute(0					,  2					,  3					,  1    )[0]
       a						:Optional[Any] 							=       Image.fromarray(np.uinta(_lowerCamelCase    )    ).convert('''RGB'''    )
       if str(_lowerCamelCase    ).startswith('''mps'''    ):
             a						:Dict 							=       torch.manual_seed(_lowerCamelCase    )
       else:
             a						:Any 							=       torch.Generator(device=_lowerCamelCase    ).manual_seed(_lowerCamelCase    )
       a						:List[Any] 							=       {
           '''image''': image,
           '''prompt''': '''a cat and a frog''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''inpaint_strength''': 1.0,
           '''guidance_scale''': 6.0,
           '''decode_latents''': True,
           '''output_type''': '''numpy''',
       }
       return inputs




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       if not hasattr(self.pipeline_class					,  '''_optional_components'''    ):
             return

       a						:Any 							=       self.get_dummy_components()
       a						:str 							=       self.pipeline_class(**_lowerCamelCase    )
       pipe.to(_lowerCamelCase    )
       pipe.set_progress_bar_config(disable=_lowerCamelCase    )

       # set all optional components to None and update pipeline config accordingly
       for optional_component in pipe._optional_components:
             setattr(_lowerCamelCase					,  _lowerCamelCase					,  _lowerCamelCase    )
       pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}    )

       a						:Union[str, Any] 							=       self.get_dummy_inputs(_lowerCamelCase    )
       a						:List[Any] 							=       pipe(**_lowerCamelCase    )[0]

       with tempfile.TemporaryDirectory() as tmpdir:
             pipe.save_pretrained(_lowerCamelCase    )
             a						:int 							=       self.pipeline_class.from_pretrained(_lowerCamelCase    )
             pipe_loaded.to(_lowerCamelCase    )
             pipe_loaded.set_progress_bar_config(disable=_lowerCamelCase    )

       for optional_component in pipe._optional_components:
             self.assertTrue(
                 getattr(_lowerCamelCase					,  _lowerCamelCase    ) is None					,  F'''`{optional_component}` did not stay set to None after loading.'''					,  )

       a						:Dict 							=       self.get_dummy_inputs(_lowerCamelCase    )
       a						:Tuple 							=       pipe_loaded(**_lowerCamelCase    )[0]

       a						:List[str] 							=       np.abs(output - output_loaded    ).max()
       self.assertLess(_lowerCamelCase					,  1e-4    )




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       a						:Any 							=       '''cpu'''

       a						:Optional[int] 							=       self.get_dummy_components()
       a						:List[str] 							=       self.pipeline_class(**_lowerCamelCase    )
       pipe.to(_lowerCamelCase    )
       pipe.set_progress_bar_config(disable=_lowerCamelCase    )

       a						:Any 							=       self.get_dummy_mask_inputs(_lowerCamelCase    )
       a						:str 							=       pipe.generate_mask(**_lowerCamelCase    )
       a						:List[str] 							=       mask[0, -3:, -3:]

       self.assertEqual(mask.shape					,  (1, 16, 16)    )
       a						:List[Any] 							=       np.array([0] * 9    )
       a						:str 							=       np.abs(mask_slice.flatten() - expected_slice    ).max()
       self.assertLessEqual(_lowerCamelCase					,  1e-3    )
       self.assertEqual(mask[0, -3, -4]					,  0    )




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       a						:Any 							=       '''cpu'''

       a						:List[str] 							=       self.get_dummy_components()
       a						:List[Any] 							=       self.pipeline_class(**_lowerCamelCase    )
       pipe.to(_lowerCamelCase    )
       pipe.set_progress_bar_config(disable=_lowerCamelCase    )

       a						:str 							=       self.get_dummy_inversion_inputs(_lowerCamelCase    )
       a						:Any 							=       pipe.invert(**_lowerCamelCase    ).images
       a						:Union[str, Any] 							=       image[0, -1, -3:, -3:]

       self.assertEqual(image.shape					,  (2, 32, 32, 3)    )
       a						:str 							=       np.array(
           [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.5_1050, 0.5015, 0.4407, 0.4799]					,  )
       a						:str 							=       np.abs(image_slice.flatten() - expected_slice    ).max()
       self.assertLessEqual(_lowerCamelCase					,  1e-3    )




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       super().test_inference_batch_single_identical(expected_max_diff=5e-3    )




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       a						:str 							=       '''cpu'''

       a						:str 							=       self.get_dummy_components()

       a						:Union[str, Any] 							=       {'''beta_start''': 0.0_0085, '''beta_end''': 0.012, '''beta_schedule''': '''scaled_linear'''}
       a						:Dict 							=       DPMSolverMultistepScheduler(**_lowerCamelCase    )
       a						:List[str] 							=       DPMSolverMultistepInverseScheduler(**_lowerCamelCase    )

       a						:int 							=       self.pipeline_class(**_lowerCamelCase    )
       pipe.to(_lowerCamelCase    )
       pipe.set_progress_bar_config(disable=_lowerCamelCase    )

       a						:str 							=       self.get_dummy_inversion_inputs(_lowerCamelCase    )
       a						:Tuple 							=       pipe.invert(**_lowerCamelCase    ).images
       a						:Any 							=       image[0, -1, -3:, -3:]

       self.assertEqual(image.shape					,  (2, 32, 32, 3)    )
       a						:int 							=       np.array(
           [0.5150, 0.5134, 0.5043, 0.5376, 0.4694, 0.5_1050, 0.5015, 0.4407, 0.4799]					,  )
       a						:Tuple 							=       np.abs(image_slice.flatten() - expected_slice    ).max()
       self.assertLessEqual(_lowerCamelCase					,  1e-3    )







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




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




 @classmethod
 def        SCREAMING_SNAKE_CASE__				(		cls    ):
       a						:Tuple 							=       load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'''    )

       a						:Union[str, Any] 							=       raw_image.convert('''RGB'''    ).resize((768, 768)    )

       a						:List[Any] 							=       raw_image




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       a						:Tuple 							=       torch.manual_seed(0    )

       a						:int 							=       StableDiffusionDiffEditPipeline.from_pretrained(
           '''stabilityai/stable-diffusion-2-1'''					,  safety_checker=_lowerCamelCase					,  torch_dtype=torch.floataa    )
       a						:Optional[int] 							=       DDIMScheduler.from_config(pipe.scheduler.config    )
       a						:Optional[int] 							=       DDIMInverseScheduler.from_config(pipe.scheduler.config    )
       pipe.enable_model_cpu_offload()
       pipe.set_progress_bar_config(disable=_lowerCamelCase    )

       a						:Optional[Any] 							=       '''a bowl of fruit'''
       a						:Any 							=       '''a bowl of pears'''

       a						:Any 							=       pipe.generate_mask(
           image=self.raw_image					,  source_prompt=_lowerCamelCase					,  target_prompt=_lowerCamelCase					,  generator=_lowerCamelCase					,  )

       a						:Dict 							=       pipe.invert(
           prompt=_lowerCamelCase					,  image=self.raw_image					,  inpaint_strength=0.7					,  generator=_lowerCamelCase    ).latents

       a						:List[str] 							=       pipe(
           prompt=_lowerCamelCase					,  mask_image=_lowerCamelCase					,  image_latents=_lowerCamelCase					,  generator=_lowerCamelCase					,  negative_prompt=_lowerCamelCase					,  inpaint_strength=0.7					,  output_type='''numpy'''					,  ).images[0]

       a						:List[str] 							=       (
           np.array(
               load_image(
                   '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                   '''/diffedit/pears.png'''    ).resize((768, 768)    )    )
           / 255
       )
       assert np.abs((expected_image - image).max()    ) < 5e-1




 def        SCREAMING_SNAKE_CASE__				(		self    ):
       a						:Optional[int] 							=       torch.manual_seed(0    )

       a						:List[Any] 							=       StableDiffusionDiffEditPipeline.from_pretrained(
           '''stabilityai/stable-diffusion-2-1'''					,  safety_checker=_lowerCamelCase					,  torch_dtype=torch.floataa    )
       a						:Any 							=       DPMSolverMultistepScheduler.from_config(pipe.scheduler.config    )
       a						:Union[str, Any] 							=       DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config    )
       pipe.enable_model_cpu_offload()
       pipe.set_progress_bar_config(disable=_lowerCamelCase    )

       a						:Dict 							=       '''a bowl of fruit'''
       a						:Optional[Any] 							=       '''a bowl of pears'''

       a						:Tuple 							=       pipe.generate_mask(
           image=self.raw_image					,  source_prompt=_lowerCamelCase					,  target_prompt=_lowerCamelCase					,  generator=_lowerCamelCase					,  )

       a						:Dict 							=       pipe.invert(
           prompt=_lowerCamelCase					,  image=self.raw_image					,  inpaint_strength=0.7					,  generator=_lowerCamelCase					,  num_inference_steps=25					,  ).latents

       a						:str 							=       pipe(
           prompt=_lowerCamelCase					,  mask_image=_lowerCamelCase					,  image_latents=_lowerCamelCase					,  generator=_lowerCamelCase					,  negative_prompt=_lowerCamelCase					,  inpaint_strength=0.7					,  num_inference_steps=25					,  output_type='''numpy'''					,  ).images[0]

       a						:List[Any] 							=       (
           np.array(
               load_image(
                   '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                   '''/diffedit/pears.png'''    ).resize((768, 768)    )    )
           / 255
       )
       assert np.abs((expected_image - image).max()    ) < 5e-1






'''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 ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


_A			:      Optional[Any]					      =				logging.get_logger(__name__)

# General docstring
_A			:      Optional[Any]					      =				'''ResNetConfig'''

# Base docstring
_A			:      Tuple					      =				'''microsoft/resnet-50'''
_A			:      List[str]					      =				[1, 2048, 7, 7]

# Image classification docstring
_A			:      str					      =				'''microsoft/resnet-50'''
_A			:      Dict					      =				'''tiger cat'''

_A			:      List[Any]					      =				[
    '''microsoft/resnet-50''',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






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





 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor							)  ->    Tensor:
  __lowerCAmelCase											=						self.convolution(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.normalization(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.activation(SCREAMING_SNAKE_CASE__							)
  return hidden_state







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	ResNetConfig							)  ->    List[str]:
  super().__init__()
  __lowerCAmelCase											=						ResNetConvLayer(
      config.num_channels      ,       config.embedding_size      ,       kernel_size=7      ,       stride=2      ,       activation=config.hidden_act							)
  __lowerCAmelCase											=						nn.MaxPoolad(kernel_size=3      ,       stride=2      ,       padding=1							)
  __lowerCAmelCase											=						config.num_channels





 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor							)  ->    Tensor:
  __lowerCAmelCase											=						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."""							)
  __lowerCAmelCase											=						self.embedder(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.pooler(SCREAMING_SNAKE_CASE__							)
  return embedding







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






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





 def a					(      self   :	Optional[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor							)  ->    Tensor:
  __lowerCAmelCase											=						self.convolution(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.normalization(SCREAMING_SNAKE_CASE__							)
  return hidden_state







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	int      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int = 1      ,       SCREAMING_SNAKE_CASE__   :	str = "relu"							)  ->    Dict:
  super().__init__()
  __lowerCAmelCase											=						in_channels != out_channels or stride != 1
  __lowerCAmelCase											=						(
      ResNetShortCut(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       stride=SCREAMING_SNAKE_CASE__							) if should_apply_shortcut else nn.Identity()
  )
  __lowerCAmelCase											=						nn.Sequential(
      ResNetConvLayer(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       stride=SCREAMING_SNAKE_CASE__							)      ,       ResNetConvLayer(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       activation=SCREAMING_SNAKE_CASE__							)      ,       )
  __lowerCAmelCase											=						ACTaFN[activation]





 def a					(      self   :	Optional[Any]      ,       SCREAMING_SNAKE_CASE__   :	List[Any]							)  ->    int:
  __lowerCAmelCase											=						hidden_state
  __lowerCAmelCase											=						self.layer(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.shortcut(SCREAMING_SNAKE_CASE__							)
  hidden_state += residual
  __lowerCAmelCase											=						self.activation(SCREAMING_SNAKE_CASE__							)
  return hidden_state







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int = 1      ,       SCREAMING_SNAKE_CASE__   :	str = "relu"      ,       SCREAMING_SNAKE_CASE__   :	int = 4							)  ->    int:
  super().__init__()
  __lowerCAmelCase											=						in_channels != out_channels or stride != 1
  __lowerCAmelCase											=						out_channels // reduction
  __lowerCAmelCase											=						(
      ResNetShortCut(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       stride=SCREAMING_SNAKE_CASE__							) if should_apply_shortcut else nn.Identity()
  )
  __lowerCAmelCase											=						nn.Sequential(
      ResNetConvLayer(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       kernel_size=1							)      ,       ResNetConvLayer(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       stride=SCREAMING_SNAKE_CASE__							)      ,       ResNetConvLayer(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       kernel_size=1      ,       activation=SCREAMING_SNAKE_CASE__							)      ,       )
  __lowerCAmelCase											=						ACTaFN[activation]





 def a					(      self   :	Any      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]							)  ->    Tuple:
  __lowerCAmelCase											=						hidden_state
  __lowerCAmelCase											=						self.layer(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.shortcut(SCREAMING_SNAKE_CASE__							)
  hidden_state += residual
  __lowerCAmelCase											=						self.activation(SCREAMING_SNAKE_CASE__							)
  return hidden_state







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	ResNetConfig      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int = 2      ,       SCREAMING_SNAKE_CASE__   :	int = 2      ,       )  ->    int:
  super().__init__()

  __lowerCAmelCase											=						ResNetBottleNeckLayer if config.layer_type == """bottleneck""" else ResNetBasicLayer

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





 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor							)  ->    Tensor:
  __lowerCAmelCase											=						input
  for layer in self.layers:
   __lowerCAmelCase											=						layer(SCREAMING_SNAKE_CASE__							)
  return hidden_state







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






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





 def a					(      self   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	Tensor      ,       SCREAMING_SNAKE_CASE__   :	bool = False      ,       SCREAMING_SNAKE_CASE__   :	bool = True							)  ->    BaseModelOutputWithNoAttention:
  __lowerCAmelCase											=						() if output_hidden_states else None

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

   __lowerCAmelCase											=						stage_module(SCREAMING_SNAKE_CASE__							)

  if output_hidden_states:
   __lowerCAmelCase											=						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=SCREAMING_SNAKE_CASE__      ,       hidden_states=SCREAMING_SNAKE_CASE__      ,       )







class 				_lowercase     (				UpperCAmelCase__			):


 '''simple docstring'''






 _SCREAMING_SNAKE_CASE     :       int      	=					ResNetConfig
 _SCREAMING_SNAKE_CASE     :       Union[str, Any]      	=					"""resnet"""
 _SCREAMING_SNAKE_CASE     :       Union[str, Any]      	=					"""pixel_values"""
 _SCREAMING_SNAKE_CASE     :       Union[str, Any]      	=					True
 def a					(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]							)  ->    str:
  if isinstance(SCREAMING_SNAKE_CASE__      ,       nn.Convad							):
   nn.init.kaiming_normal_(module.weight      ,       mode="""fan_out"""      ,       nonlinearity="""relu"""							)
  elif isinstance(SCREAMING_SNAKE_CASE__      ,       (nn.BatchNormad, nn.GroupNorm)							):
   nn.init.constant_(module.weight      ,       1							)
   nn.init.constant_(module.bias      ,       0							)





 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]=False							)  ->    int:
  if isinstance(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__							):
   __lowerCAmelCase											=						value







_A			:      Dict					      =				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 ([`ResNetConfig`]): 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.
'''

_A			:      Optional[int]					      =				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 [`~utils.ModelOutput`] instead of a plain tuple.
'''







@add_start_docstrings(
    """The bare ResNet model outputting raw features without any specific head on top."""  ,							UpperCAmelCase__  ,							)
class 				_lowercase     (				UpperCAmelCase__			):


 '''simple docstring'''






 def __init__(      self   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]							)  ->    List[str]:
  super().__init__(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						config
  __lowerCAmelCase											=						ResNetEmbeddings(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						ResNetEncoder(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						nn.AdaptiveAvgPoolad((1, 1)							)
  # Initialize weights and apply final processing
  self.post_init()





 @add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE__							)
 @add_code_sample_docstrings(
     checkpoint=_CHECKPOINT_FOR_DOC      ,       output_type=SCREAMING_SNAKE_CASE__      ,       config_class=_CONFIG_FOR_DOC      ,       modality="""vision"""      ,       expected_output=_EXPECTED_OUTPUT_SHAPE      ,       )
 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None							)  ->    BaseModelOutputWithPoolingAndNoAttention:
  __lowerCAmelCase											=						(
      output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
  )
  __lowerCAmelCase											=						return_dict if return_dict is not None else self.config.use_return_dict

  __lowerCAmelCase											=						self.embedder(SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						self.encoder(
      SCREAMING_SNAKE_CASE__      ,       output_hidden_states=SCREAMING_SNAKE_CASE__      ,       return_dict=SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						encoder_outputs[0]

  __lowerCAmelCase											=						self.pooler(SCREAMING_SNAKE_CASE__							)

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

  return BaseModelOutputWithPoolingAndNoAttention(
      last_hidden_state=SCREAMING_SNAKE_CASE__      ,       pooler_output=SCREAMING_SNAKE_CASE__      ,       hidden_states=encoder_outputs.hidden_states      ,       )







@add_start_docstrings(
    """
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """  ,							UpperCAmelCase__  ,							)
class 				_lowercase     (				UpperCAmelCase__			):


 '''simple docstring'''






 def __init__(      self   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	Tuple							)  ->    Any:
  super().__init__(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						config.num_labels
  __lowerCAmelCase											=						ResNetModel(SCREAMING_SNAKE_CASE__							)
  # classification head
  __lowerCAmelCase											=						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(SCREAMING_SNAKE_CASE__							)
 @add_code_sample_docstrings(
     checkpoint=_IMAGE_CLASS_CHECKPOINT      ,       output_type=SCREAMING_SNAKE_CASE__      ,       config_class=_CONFIG_FOR_DOC      ,       expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,       )
 def a					(      self   :	int      ,       SCREAMING_SNAKE_CASE__   :	Optional[torch.FloatTensor] = None      ,       SCREAMING_SNAKE_CASE__   :	Optional[torch.LongTensor] = None      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None      ,       )  ->    ImageClassifierOutputWithNoAttention:
  __lowerCAmelCase											=						return_dict if return_dict is not None else self.config.use_return_dict

  __lowerCAmelCase											=						self.resnet(SCREAMING_SNAKE_CASE__      ,       output_hidden_states=SCREAMING_SNAKE_CASE__      ,       return_dict=SCREAMING_SNAKE_CASE__							)

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

  __lowerCAmelCase											=						self.classifier(SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						None

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

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

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







@add_start_docstrings(
    """
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    """  ,							UpperCAmelCase__  ,							)
class 				_lowercase     (				UpperCAmelCase__  ,							UpperCAmelCase__			):


 '''simple docstring'''






 def __init__(      self   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	List[Any]							)  ->    Tuple:
  super().__init__(SCREAMING_SNAKE_CASE__							)
  super()._init_backbone(SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						[config.embedding_size] + config.hidden_sizes
  __lowerCAmelCase											=						ResNetEmbeddings(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						ResNetEncoder(SCREAMING_SNAKE_CASE__							)

  # initialize weights and apply final processing
  self.post_init()





 @add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE__							)
 @replace_return_docstrings(output_type=SCREAMING_SNAKE_CASE__      ,       config_class=_CONFIG_FOR_DOC							)
 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None							)  ->    BackboneOutput:
  __lowerCAmelCase											=						return_dict if return_dict is not None else self.config.use_return_dict
  __lowerCAmelCase											=						(
      output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
  )

  __lowerCAmelCase											=						self.embedder(SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						self.encoder(SCREAMING_SNAKE_CASE__      ,       output_hidden_states=SCREAMING_SNAKE_CASE__      ,       return_dict=SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						outputs.hidden_states

  __lowerCAmelCase											=						()
  for idx, stage in enumerate(self.stage_names							):
   if stage in self.out_features:
    feature_maps += (hidden_states[idx],)

  if not return_dict:
   __lowerCAmelCase											=						(feature_maps,)
   if output_hidden_states:
    output += (outputs.hidden_states,)
   return output

  return BackboneOutput(
      feature_maps=SCREAMING_SNAKE_CASE__      ,       hidden_states=outputs.hidden_states if output_hidden_states else None      ,       attentions=SCREAMING_SNAKE_CASE__      ,       )






