




'''simple docstring'''




import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class 			lowerCamelCase_			(unittest.TestCase					):
      '''simple docstring'''

      def   _A			( self		:       str ):
             # clean up the VRAM after each test
             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()




      @property
      def   _A			( self		:       Tuple ):
             _UpperCAmelCase  :						List[str]							   =					1
             _UpperCAmelCase  :						Dict							   =					3
             _UpperCAmelCase  :						List[str]							   =					(32, 32)

             _UpperCAmelCase  :						Tuple							   =					floats_tensor((batch_size, num_channels) + sizes		,      rng=random.Random(0 ) ).to(A )
             return image




      @property
      def   _A			( self		:       Union[str, Any] ):
             torch.manual_seed(0 )
             _UpperCAmelCase  :						Optional[int]							   =					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		,      )
             return model




      @property
      def   _A			( self		:       Optional[Any] ):
             torch.manual_seed(0 )
             _UpperCAmelCase  :						str							   =					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		,      )
             return model




      @property
      def   _A			( self		:       int ):
             torch.manual_seed(0 )
             _UpperCAmelCase  :						List[Any]							   =					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		,      )
             return CLIPTextModel(A )




      @property
      def   _A			( self		:       List[Any] ):
             def extract(*A		:       Union[str, Any]		,      **A		:       Tuple ):

                    class 			lowerCamelCase_			:
                          '''simple docstring'''

                          def __init__( self		:       Dict ):
                                 _UpperCAmelCase  :						Dict							   =					torch.ones([0] )





                          def   _A			( self		:       Tuple		,      A		:       Dict ):
                                 self.pixel_values.to(A )
                                 return self

                    return Out()

             return extract




      def   _A			( self		:       Dict ):
             _UpperCAmelCase  :						str							   =					"cpu"  # ensure determinism for the device-dependent torch.Generator
             _UpperCAmelCase  :						Tuple							   =					self.dummy_cond_unet
             _UpperCAmelCase  :						Optional[Any]							   =					DDIMScheduler(
                 beta_start=0.00_085		,      beta_end=0.012		,      beta_schedule="scaled_linear"		,      clip_sample=A		,      set_alpha_to_one=A		,      )

             _UpperCAmelCase  :						Any							   =					self.dummy_vae
             _UpperCAmelCase  :						Tuple							   =					self.dummy_text_encoder
             _UpperCAmelCase  :						Optional[int]							   =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

             # make sure here that pndm scheduler skips prk
             _UpperCAmelCase  :						Tuple							   =					StableDiffusionPipeline(
                 unet=A		,      scheduler=A		,      vae=A		,      text_encoder=A		,      tokenizer=A		,      safety_checker=A		,      feature_extractor=self.dummy_extractor		,      )
             _UpperCAmelCase  :						str							   =					sd_pipe.to(A )
             sd_pipe.set_progress_bar_config(disable=A )

             _UpperCAmelCase  :						Tuple							   =					"A painting of a squirrel eating a burger"

             _UpperCAmelCase  :						int							   =					torch.Generator(device=A ).manual_seed(0 )
             _UpperCAmelCase  :						Union[str, Any]							   =					sd_pipe([prompt]		,      generator=A		,      guidance_scale=6.0		,      num_inference_steps=2		,      output_type="np" )
             _UpperCAmelCase  :						str							   =					output.images

             _UpperCAmelCase  :						Union[str, Any]							   =					torch.Generator(device=A ).manual_seed(0 )
             _UpperCAmelCase  :						Tuple							   =					sd_pipe(
                 [prompt]		,      generator=A		,      guidance_scale=6.0		,      num_inference_steps=2		,      output_type="np"		,      return_dict=A		,      )[0]

             _UpperCAmelCase  :						str							   =					image[0, -3:, -3:, -1]
             _UpperCAmelCase  :						int							   =					image_from_tuple[0, -3:, -3:, -1]

             assert image.shape == (1, 64, 64, 3)
             _UpperCAmelCase  :						Optional[int]							   =					np.array([0.5_753, 0.6_114, 0.5_001, 0.5_034, 0.5_470, 0.4_729, 0.4_971, 0.4_867, 0.4_867] )

             assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2
             assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2




      def   _A			( self		:       Tuple ):
             _UpperCAmelCase  :						Dict							   =					"cpu"  # ensure determinism for the device-dependent torch.Generator
             _UpperCAmelCase  :						Tuple							   =					self.dummy_cond_unet
             _UpperCAmelCase  :						Dict							   =					PNDMScheduler(skip_prk_steps=A )
             _UpperCAmelCase  :						List[str]							   =					self.dummy_vae
             _UpperCAmelCase  :						Any							   =					self.dummy_text_encoder
             _UpperCAmelCase  :						Union[str, Any]							   =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

             # make sure here that pndm scheduler skips prk
             _UpperCAmelCase  :						List[str]							   =					StableDiffusionPipeline(
                 unet=A		,      scheduler=A		,      vae=A		,      text_encoder=A		,      tokenizer=A		,      safety_checker=A		,      feature_extractor=self.dummy_extractor		,      )
             _UpperCAmelCase  :						str							   =					sd_pipe.to(A )
             sd_pipe.set_progress_bar_config(disable=A )

             _UpperCAmelCase  :						int							   =					"A painting of a squirrel eating a burger"
             _UpperCAmelCase  :						Optional[Any]							   =					torch.Generator(device=A ).manual_seed(0 )
             _UpperCAmelCase  :						Tuple							   =					sd_pipe([prompt]		,      generator=A		,      guidance_scale=6.0		,      num_inference_steps=2		,      output_type="np" )

             _UpperCAmelCase  :						Union[str, Any]							   =					output.images

             _UpperCAmelCase  :						List[Any]							   =					torch.Generator(device=A ).manual_seed(0 )
             _UpperCAmelCase  :						Tuple							   =					sd_pipe(
                 [prompt]		,      generator=A		,      guidance_scale=6.0		,      num_inference_steps=2		,      output_type="np"		,      return_dict=A		,      )[0]

             _UpperCAmelCase  :						Optional[int]							   =					image[0, -3:, -3:, -1]
             _UpperCAmelCase  :						List[Any]							   =					image_from_tuple[0, -3:, -3:, -1]

             assert image.shape == (1, 64, 64, 3)
             _UpperCAmelCase  :						Union[str, Any]							   =					np.array([0.5_122, 0.5_712, 0.4_825, 0.5_053, 0.5_646, 0.4_769, 0.5_179, 0.4_894, 0.4_994] )

             assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2
             assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1E-2




      def   _A			( self		:       Optional[int] ):
             _UpperCAmelCase  :						Optional[Any]							   =					StableDiffusionPipeline.from_pretrained(
                 "hf-internal-testing/tiny-stable-diffusion-lms-pipe"		,      safety_checker=A )
             assert isinstance(A		,      A )
             assert isinstance(pipe.scheduler		,      A )
             assert pipe.safety_checker is None

             _UpperCAmelCase  :						List[Any]							   =					pipe("example prompt"		,      num_inference_steps=2 ).images[0]
             assert image is not None

             # check that there's no error when saving a pipeline with one of the models being None
             with tempfile.TemporaryDirectory() as tmpdirname:
                    pipe.save_pretrained(A )
                    _UpperCAmelCase  :						Any							   =					StableDiffusionPipeline.from_pretrained(A )

             # sanity check that the pipeline still works
             assert pipe.safety_checker is None
             _UpperCAmelCase  :						Optional[Any]							   =					pipe("example prompt"		,      num_inference_steps=2 ).images[0]
             assert image is not None





      @unittest.skipIf(torch_device != "cuda"		,      "This test requires a GPU" )
      def   _A			( self		:       Any ):
             _UpperCAmelCase  :						Any							   =					self.dummy_cond_unet
             _UpperCAmelCase  :						int							   =					PNDMScheduler(skip_prk_steps=A )
             _UpperCAmelCase  :						Tuple							   =					self.dummy_vae
             _UpperCAmelCase  :						List[str]							   =					self.dummy_text_encoder
             _UpperCAmelCase  :						Any							   =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )

             # put models in fp16
             _UpperCAmelCase  :						List[Any]							   =					unet.half()
             _UpperCAmelCase  :						Optional[int]							   =					vae.half()
             _UpperCAmelCase  :						List[str]							   =					bert.half()

             # make sure here that pndm scheduler skips prk
             _UpperCAmelCase  :						Dict							   =					StableDiffusionPipeline(
                 unet=A		,      scheduler=A		,      vae=A		,      text_encoder=A		,      tokenizer=A		,      safety_checker=A		,      feature_extractor=self.dummy_extractor		,      )
             _UpperCAmelCase  :						int							   =					sd_pipe.to(A )
             sd_pipe.set_progress_bar_config(disable=A )

             _UpperCAmelCase  :						List[str]							   =					"A painting of a squirrel eating a burger"
             _UpperCAmelCase  :						Dict							   =					sd_pipe([prompt]		,      num_inference_steps=2		,      output_type="np" ).images

             assert image.shape == (1, 64, 64, 3)

@nightly
@require_torch_gpu
class 			lowerCamelCase_			(unittest.TestCase					):
      '''simple docstring'''

      def   _A			( self		:       Any ):
             # clean up the VRAM after each test
             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()




      def   _A			( self		:       Dict ):
             _UpperCAmelCase  :						List[Any]							   =					"cuda"
             _UpperCAmelCase  :						Optional[Any]							   =					StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5" )
             _UpperCAmelCase  :						Tuple							   =					pipe.to(A )
             pipe.set_progress_bar_config(disable=A )

             _UpperCAmelCase  :						int							   =					"a photo of a cat"
             _UpperCAmelCase  :						int							   =					{
                 "editing_prompt": ["sunglasses"],
                 "reverse_editing_direction": [False],
                 "edit_warmup_steps": 10,
                 "edit_guidance_scale": 6,
                 "edit_threshold": 0.95,
                 "edit_momentum_scale": 0.5,
                 "edit_mom_beta": 0.6,
             }

             _UpperCAmelCase  :						Optional[int]							   =					3
             _UpperCAmelCase  :						Dict							   =					7

             # no sega enabled
             _UpperCAmelCase  :						str							   =					torch.Generator(A )
             generator.manual_seed(A )
             _UpperCAmelCase  :						List[Any]							   =					pipe(
                 [prompt]		,      generator=A		,      guidance_scale=A		,      num_inference_steps=50		,      output_type="np"		,      width=512		,      height=512		,      )

             _UpperCAmelCase  :						Optional[Any]							   =					output.images
             _UpperCAmelCase  :						Dict							   =					image[0, -3:, -3:, -1]
             _UpperCAmelCase  :						List[str]							   =					[
                 0.34_673_113,
                 0.38_492_733,
                 0.37_597_352,
                 0.34_086_335,
                 0.35_650_748,
                 0.35_579_205,
                 0.3_384_763,
                 0.34_340_236,
                 0.3_573_271,
             ]

             assert image.shape == (1, 512, 512, 3)

             assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2

             # with sega enabled
             # generator = torch.manual_seed(seed)
             generator.manual_seed(A )
             _UpperCAmelCase  :						Union[str, Any]							   =					pipe(
                 [prompt]		,      generator=A		,      guidance_scale=A		,      num_inference_steps=50		,      output_type="np"		,      width=512		,      height=512		,      **A		,      )

             _UpperCAmelCase  :						int							   =					output.images
             _UpperCAmelCase  :						Optional[Any]							   =					image[0, -3:, -3:, -1]
             _UpperCAmelCase  :						str							   =					[
                 0.41_887_826,
                 0.37_728_766,
                 0.30_138_272,
                 0.41_416_335,
                 0.41_664_985,
                 0.36_283_392,
                 0.36_191_246,
                 0.43_364_465,
                 0.43_001_732,
             ]

             assert image.shape == (1, 512, 512, 3)

             assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2




      def   _A			( self		:       Optional[int] ):
             _UpperCAmelCase  :						Dict							   =					"cuda"
             _UpperCAmelCase  :						Any							   =					StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5" )
             _UpperCAmelCase  :						Union[str, Any]							   =					pipe.to(A )
             pipe.set_progress_bar_config(disable=A )

             _UpperCAmelCase  :						List[Any]							   =					"an image of a crowded boulevard, realistic, 4k"
             _UpperCAmelCase  :						Union[str, Any]							   =					{
                 "editing_prompt": "crowd, crowded, people",
                 "reverse_editing_direction": True,
                 "edit_warmup_steps": 10,
                 "edit_guidance_scale": 8.3,
                 "edit_threshold": 0.9,
                 "edit_momentum_scale": 0.5,
                 "edit_mom_beta": 0.6,
             }

             _UpperCAmelCase  :						Dict							   =					9
             _UpperCAmelCase  :						List[str]							   =					7

             # no sega enabled
             _UpperCAmelCase  :						Optional[int]							   =					torch.Generator(A )
             generator.manual_seed(A )
             _UpperCAmelCase  :						List[Any]							   =					pipe(
                 [prompt]		,      generator=A		,      guidance_scale=A		,      num_inference_steps=50		,      output_type="np"		,      width=512		,      height=512		,      )

             _UpperCAmelCase  :						Tuple							   =					output.images
             _UpperCAmelCase  :						Optional[Any]							   =					image[0, -3:, -3:, -1]
             _UpperCAmelCase  :						List[Any]							   =					[
                 0.43_497_998,
                 0.91_814_065,
                 0.7_540_739,
                 0.55_580_205,
                 0.8_467_265,
                 0.5_389_691,
                 0.62_574_506,
                 0.58_897_763,
                 0.50_926_757,
             ]

             assert image.shape == (1, 512, 512, 3)

             assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2

             # with sega enabled
             # generator = torch.manual_seed(seed)
             generator.manual_seed(A )
             _UpperCAmelCase  :						str							   =					pipe(
                 [prompt]		,      generator=A		,      guidance_scale=A		,      num_inference_steps=50		,      output_type="np"		,      width=512		,      height=512		,      **A		,      )

             _UpperCAmelCase  :						List[str]							   =					output.images
             _UpperCAmelCase  :						List[str]							   =					image[0, -3:, -3:, -1]
             _UpperCAmelCase  :						Tuple							   =					[
                 0.3_089_719,
                 0.30_500_144,
                 0.29_016_042,
                 0.30_630_964,
                 0.325_687,
                 0.29_419_225,
                 0.2_908_091,
                 0.28_723_598,
                 0.27_696_294,
             ]

             assert image.shape == (1, 512, 512, 3)

             assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2




      def   _A			( self		:       str ):
             _UpperCAmelCase  :						List[Any]							   =					"cuda"
             _UpperCAmelCase  :						int							   =					StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5" )
             _UpperCAmelCase  :						List[str]							   =					pipe.to(A )
             pipe.set_progress_bar_config(disable=A )

             _UpperCAmelCase  :						List[str]							   =					"a castle next to a river"
             _UpperCAmelCase  :						Union[str, Any]							   =					{
                 "editing_prompt": ["boat on a river, boat", "monet, impression, sunrise"],
                 "reverse_editing_direction": False,
                 "edit_warmup_steps": [15, 18],
                 "edit_guidance_scale": 6,
                 "edit_threshold": [0.9, 0.8],
                 "edit_momentum_scale": 0.5,
                 "edit_mom_beta": 0.6,
             }

             _UpperCAmelCase  :						List[str]							   =					48
             _UpperCAmelCase  :						Union[str, Any]							   =					7

             # no sega enabled
             _UpperCAmelCase  :						List[Any]							   =					torch.Generator(A )
             generator.manual_seed(A )
             _UpperCAmelCase  :						str							   =					pipe(
                 [prompt]		,      generator=A		,      guidance_scale=A		,      num_inference_steps=50		,      output_type="np"		,      width=512		,      height=512		,      )

             _UpperCAmelCase  :						Dict							   =					output.images
             _UpperCAmelCase  :						int							   =					image[0, -3:, -3:, -1]
             _UpperCAmelCase  :						Tuple							   =					[
                 0.75_163_555,
                 0.76_037_145,
                 0.61_785,
                 0.9_189_673,
                 0.8_627_701,
                 0.85_189_694,
                 0.8_512_813,
                 0.87_012_076,
                 0.8_312_857,
             ]

             assert image.shape == (1, 512, 512, 3)

             assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2

             # with sega enabled
             # generator = torch.manual_seed(seed)
             generator.manual_seed(A )
             _UpperCAmelCase  :						Any							   =					pipe(
                 [prompt]		,      generator=A		,      guidance_scale=A		,      num_inference_steps=50		,      output_type="np"		,      width=512		,      height=512		,      **A		,      )

             _UpperCAmelCase  :						Optional[Any]							   =					output.images
             _UpperCAmelCase  :						Union[str, Any]							   =					image[0, -3:, -3:, -1]
             _UpperCAmelCase  :						Union[str, Any]							   =					[
                 0.73_553_365,
                 0.7_537_271,
                 0.74_341_905,
                 0.66_480_356,
                 0.6_472_925,
                 0.63_039_416,
                 0.64_812_905,
                 0.6_749_717,
                 0.6_517_102,
             ]

             assert image.shape == (1, 512, 512, 3)

             assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2





      def   _A			( self		:       List[str] ):
             _UpperCAmelCase  :						List[str]							   =					"cuda"
             _UpperCAmelCase  :						Any							   =					StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"		,      torch_dtype=torch.floataa )
             _UpperCAmelCase  :						Union[str, Any]							   =					pipe.to(A )
             pipe.set_progress_bar_config(disable=A )

             _UpperCAmelCase  :						List[str]							   =					"a photo of a cat"
             _UpperCAmelCase  :						List[Any]							   =					{
                 "editing_prompt": ["sunglasses"],
                 "reverse_editing_direction": [False],
                 "edit_warmup_steps": 10,
                 "edit_guidance_scale": 6,
                 "edit_threshold": 0.95,
                 "edit_momentum_scale": 0.5,
                 "edit_mom_beta": 0.6,
             }

             _UpperCAmelCase  :						Optional[int]							   =					3
             _UpperCAmelCase  :						Union[str, Any]							   =					7

             # no sega enabled
             _UpperCAmelCase  :						int							   =					torch.Generator(A )
             generator.manual_seed(A )
             _UpperCAmelCase  :						int							   =					pipe(
                 [prompt]		,      generator=A		,      guidance_scale=A		,      num_inference_steps=50		,      output_type="np"		,      width=512		,      height=512		,      )

             _UpperCAmelCase  :						List[Any]							   =					output.images
             _UpperCAmelCase  :						Dict							   =					image[0, -3:, -3:, -1]
             _UpperCAmelCase  :						Dict							   =					[
                 0.34_887_695,
                 0.3_876_953,
                 0.375,
                 0.34_423_828,
                 0.3_581_543,
                 0.35_717_773,
                 0.3_383_789,
                 0.34_570_312,
                 0.359_375,
             ]

             assert image.shape == (1, 512, 512, 3)

             assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2

             # with sega enabled
             # generator = torch.manual_seed(seed)
             generator.manual_seed(A )
             _UpperCAmelCase  :						List[str]							   =					pipe(
                 [prompt]		,      generator=A		,      guidance_scale=A		,      num_inference_steps=50		,      output_type="np"		,      width=512		,      height=512		,      **A		,      )

             _UpperCAmelCase  :						str							   =					output.images
             _UpperCAmelCase  :						Optional[Any]							   =					image[0, -3:, -3:, -1]
             _UpperCAmelCase  :						Optional[Any]							   =					[
                 0.42_285_156,
                 0.36_914_062,
                 0.29_077_148,
                 0.42_041_016,
                 0.41_918_945,
                 0.35_498_047,
                 0.3_618_164,
                 0.4_423_828,
                 0.43_115_234,
             ]

             assert image.shape == (1, 512, 512, 3)

             assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2


'''simple docstring'''




from typing import Optional, Tuple, Union

import tensorflow as tf

from ...activations_tf import ACTaFN
from ...modeling_tf_outputs import (
    TFBaseModelOutputWithNoAttention,
    TFBaseModelOutputWithPoolingAndNoAttention,
    TFImageClassifierOutputWithNoAttention,
)
from ...modeling_tf_utils import TFPreTrainedModel, TFSequenceClassificationLoss, keras_serializable, unpack_inputs
from ...tf_utils import shape_list
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_resnet import ResNetConfig


__SCREAMING_SNAKE_CASE :  Tuple  =						logging.get_logger(__name__)

# General docstring
__SCREAMING_SNAKE_CASE :  List[str]  =						"""ResNetConfig"""

# Base docstring
__SCREAMING_SNAKE_CASE :  List[str]  =						"""microsoft/resnet-50"""
__SCREAMING_SNAKE_CASE :  Tuple  =						[1, 2_048, 7, 7]

# Image classification docstring
__SCREAMING_SNAKE_CASE :  Optional[Any]  =						"""microsoft/resnet-50"""
__SCREAMING_SNAKE_CASE :  Union[str, Any]  =						"""tiger cat"""

__SCREAMING_SNAKE_CASE :  Dict  =						[
    """microsoft/resnet-50""",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]

class 			lowerCamelCase_			(tf.keras.layers.Layer					):
      '''simple docstring'''

      def __init__( self		:       List[Any]		,      A		:       int		,      A		:       int = 3		,      A		:       int = 1		,      A		:       str = "relu"		,      **A		:       List[Any] ):
             super().__init__(**A )
             _UpperCAmelCase  :						Dict							   =					kernel_size // 2
             _UpperCAmelCase  :						Dict							   =					tf.keras.layers.ConvaD(
                 A		,      kernel_size=A		,      strides=A		,      padding="valid"		,      use_bias=A		,      name="convolution" )
             # Use same default momentum and epsilon as PyTorch equivalent
             _UpperCAmelCase  :						Optional[Any]							   =					tf.keras.layers.BatchNormalization(epsilon=1E-5		,      momentum=0.9		,      name="normalization" )
             _UpperCAmelCase  :						Tuple							   =					ACTaFN[activation] if activation is not None else tf.keras.layers.Activation("linear" )




      def   _A			( self		:       Optional[Any]		,      A		:       tf.Tensor ):
             # Pad to match that done in the PyTorch Conv2D model
             _UpperCAmelCase  :						Optional[Any]							   =					(self.pad_value, self.pad_value)
             _UpperCAmelCase  :						Any							   =					tf.pad(A		,      [(0, 0), height_pad, width_pad, (0, 0)] )
             _UpperCAmelCase  :						Dict							   =					self.conv(A )
             return hidden_state





      def   _A			( self		:       List[str]		,      A		:       tf.Tensor		,      A		:       bool = False ):
             _UpperCAmelCase  :						int							   =					self.convolution(A )
             _UpperCAmelCase  :						Optional[int]							   =					self.normalization(A		,      training=A )
             _UpperCAmelCase  :						List[str]							   =					self.activation(A )
             return hidden_state

class 			lowerCamelCase_			(tf.keras.layers.Layer					):
      '''simple docstring'''

      def __init__( self		:       List[str]		,      A		:       ResNetConfig		,      **A		:       Tuple ):
             super().__init__(**A )
             _UpperCAmelCase  :						str							   =					TFResNetConvLayer(
                 config.embedding_size		,      kernel_size=7		,      stride=2		,      activation=config.hidden_act		,      name="embedder"		,      )
             _UpperCAmelCase  :						Optional[int]							   =					tf.keras.layers.MaxPoolaD(pool_size=3		,      strides=2		,      padding="valid"		,      name="pooler" )
             _UpperCAmelCase  :						List[str]							   =					config.num_channels





      def   _A			( self		:       List[str]		,      A		:       tf.Tensor		,      A		:       bool = False ):
             _UpperCAmelCase   ,						_UpperCAmelCase   ,						_UpperCAmelCase   ,						_UpperCAmelCase  :						Any							   =					shape_list(A )
             if tf.executing_eagerly() and 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  :						Any							   =					pixel_values
             _UpperCAmelCase  :						Union[str, Any]							   =					self.embedder(A )
             _UpperCAmelCase  :						int							   =					tf.pad(A		,      [[0, 0], [1, 1], [1, 1], [0, 0]] )
             _UpperCAmelCase  :						Tuple							   =					self.pooler(A )
             return hidden_state

class 			lowerCamelCase_			(tf.keras.layers.Layer					):
      '''simple docstring'''

      def __init__( self		:       str		,      A		:       int		,      A		:       int = 2		,      **A		:       Union[str, Any] ):
             super().__init__(**A )
             _UpperCAmelCase  :						Dict							   =					tf.keras.layers.ConvaD(
                 A		,      kernel_size=1		,      strides=A		,      use_bias=A		,      name="convolution" )
             # Use same default momentum and epsilon as PyTorch equivalent
             _UpperCAmelCase  :						Tuple							   =					tf.keras.layers.BatchNormalization(epsilon=1E-5		,      momentum=0.9		,      name="normalization" )





      def   _A			( self		:       List[Any]		,      A		:       tf.Tensor		,      A		:       bool = False ):
             _UpperCAmelCase  :						List[Any]							   =					x
             _UpperCAmelCase  :						Tuple							   =					self.convolution(A )
             _UpperCAmelCase  :						List[str]							   =					self.normalization(A		,      training=A )
             return hidden_state

class 			lowerCamelCase_			(tf.keras.layers.Layer					):
      '''simple docstring'''

      def __init__( self		:       Optional[Any]		,      A		:       int		,      A		:       int		,      A		:       int = 1		,      A		:       str = "relu"		,      **A		:       Dict ):
             super().__init__(**A )
             _UpperCAmelCase  :						List[Any]							   =					in_channels != out_channels or stride != 1
             _UpperCAmelCase  :						Dict							   =					TFResNetConvLayer(A		,      stride=A		,      name="layer.0" )
             _UpperCAmelCase  :						int							   =					TFResNetConvLayer(A		,      activation=A		,      name="layer.1" )
             _UpperCAmelCase  :						List[Any]							   =					(
                 TFResNetShortCut(A		,      stride=A		,      name="shortcut" )
                 if should_apply_shortcut
                 else tf.keras.layers.Activation("linear"		,      name="shortcut" )
             )
             _UpperCAmelCase  :						Optional[int]							   =					ACTaFN[activation]





      def   _A			( self		:       Any		,      A		:       tf.Tensor		,      A		:       bool = False ):
             _UpperCAmelCase  :						Dict							   =					hidden_state
             _UpperCAmelCase  :						Optional[Any]							   =					self.conva(A		,      training=A )
             _UpperCAmelCase  :						List[str]							   =					self.conva(A		,      training=A )
             _UpperCAmelCase  :						Dict							   =					self.shortcut(A		,      training=A )
             hidden_state += residual
             _UpperCAmelCase  :						List[str]							   =					self.activation(A )
             return hidden_state

class 			lowerCamelCase_			(tf.keras.layers.Layer					):
      '''simple docstring'''

      def __init__( self		:       List[str]		,      A		:       int		,      A		:       int		,      A		:       int = 1		,      A		:       str = "relu"		,      A		:       int = 4		,      **A		:       str		,      ):
             super().__init__(**A )
             _UpperCAmelCase  :						List[Any]							   =					in_channels != out_channels or stride != 1
             _UpperCAmelCase  :						List[str]							   =					out_channels // reduction
             _UpperCAmelCase  :						List[Any]							   =					TFResNetConvLayer(A		,      kernel_size=1		,      name="layer.0" )
             _UpperCAmelCase  :						Dict							   =					TFResNetConvLayer(A		,      stride=A		,      name="layer.1" )
             _UpperCAmelCase  :						str							   =					TFResNetConvLayer(A		,      kernel_size=1		,      activation=A		,      name="layer.2" )
             _UpperCAmelCase  :						Optional[Any]							   =					(
                 TFResNetShortCut(A		,      stride=A		,      name="shortcut" )
                 if should_apply_shortcut
                 else tf.keras.layers.Activation("linear"		,      name="shortcut" )
             )
             _UpperCAmelCase  :						Union[str, Any]							   =					ACTaFN[activation]





      def   _A			( self		:       Optional[int]		,      A		:       tf.Tensor		,      A		:       bool = False ):
             _UpperCAmelCase  :						Optional[Any]							   =					hidden_state
             _UpperCAmelCase  :						int							   =					self.conva(A		,      training=A )
             _UpperCAmelCase  :						Any							   =					self.conva(A		,      training=A )
             _UpperCAmelCase  :						List[Any]							   =					self.conva(A		,      training=A )
             _UpperCAmelCase  :						List[str]							   =					self.shortcut(A		,      training=A )
             hidden_state += residual
             _UpperCAmelCase  :						Tuple							   =					self.activation(A )
             return hidden_state

class 			lowerCamelCase_			(tf.keras.layers.Layer					):
      '''simple docstring'''

      def __init__( self		:       Dict		,      A		:       ResNetConfig		,      A		:       int		,      A		:       int		,      A		:       int = 2		,      A		:       int = 2		,      **A		:       Dict ):
             super().__init__(**A )

             _UpperCAmelCase  :						Optional[int]							   =					TFResNetBottleNeckLayer if config.layer_type == "bottleneck" else TFResNetBasicLayer

             _UpperCAmelCase  :						Dict							   =					[layer(A		,      A		,      stride=A		,      activation=config.hidden_act		,      name="layers.0" )]
             layers += [
                 layer(A		,      A		,      activation=config.hidden_act		,      name=F"""layers.{i + 1}""" )
                 for i in range(depth - 1 )
             ]
             _UpperCAmelCase  :						int							   =					layers





      def   _A			( self		:       List[Any]		,      A		:       tf.Tensor		,      A		:       bool = False ):
             for layer in self.stage_layers:
                    _UpperCAmelCase  :						int							   =					layer(A		,      training=A )
             return hidden_state

class 			lowerCamelCase_			(tf.keras.layers.Layer					):
      '''simple docstring'''

      def __init__( self		:       Optional[int]		,      A		:       ResNetConfig		,      **A		:       Optional[Any] ):
             super().__init__(**A )
             # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
             _UpperCAmelCase  :						Optional[Any]							   =					[
                 TFResNetStage(
                     A		,      config.embedding_size		,      config.hidden_sizes[0]		,      stride=2 if config.downsample_in_first_stage else 1		,      depth=config.depths[0]		,      name="stages.0"		,      )
             ]
             for i, (in_channels, out_channels, depth) in enumerate(
                 zip(config.hidden_sizes		,      config.hidden_sizes[1:]		,      config.depths[1:] ) ):
                    self.stages.append(TFResNetStage(A		,      A		,      A		,      depth=A		,      name=F"""stages.{i + 1}""" ) )





      def   _A			( self		:       Tuple		,      A		:       tf.Tensor		,      A		:       bool = False		,      A		:       bool = True		,      A		:       bool = False		,      ):
             _UpperCAmelCase  :						List[str]							   =					() if output_hidden_states else None

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

                    _UpperCAmelCase  :						str							   =					stage_module(A		,      training=A )

             if output_hidden_states:
                    _UpperCAmelCase  :						Union[str, Any]							   =					hidden_states + (hidden_state,)

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

             return TFBaseModelOutputWithNoAttention(last_hidden_state=A		,      hidden_states=A )

class 			lowerCamelCase_			(snake_case__					):
      '''simple docstring'''

      __UpperCamelCase:						str      			=		ResNetConfig
      __UpperCamelCase:						List[str]      			=		"resnet"
      __UpperCamelCase:						int      			=		"pixel_values"




      @property
      def   _A			( self		:       List[Any] ):
             return {"pixel_values": tf.TensorSpec(shape=(None, self.config.num_channels, 224, 224)		,      dtype=tf.floataa )}

__SCREAMING_SNAKE_CASE :  Dict  =						R"""
    This model is a TensorFlow
    [tf.keras.layers.Layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer) sub-class. Use it as a
    regular TensorFlow Module and refer to the TensorFlow 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 [`~TFPreTrainedModel.from_pretrained`] method to load the model weights.
"""


__SCREAMING_SNAKE_CASE :  int  =						R"""
    Args:
        pixel_values (`tf.Tensor` 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.
"""

@keras_serializable
class 			lowerCamelCase_			(tf.keras.layers.Layer					):
      '''simple docstring'''

      __UpperCamelCase:						Any      			=		ResNetConfig




      def __init__( self		:       int		,      A		:       ResNetConfig		,      **A		:       str ):
             super().__init__(**A )
             _UpperCAmelCase  :						int							   =					config
             _UpperCAmelCase  :						str							   =					TFResNetEmbeddings(A		,      name="embedder" )
             _UpperCAmelCase  :						List[str]							   =					TFResNetEncoder(A		,      name="encoder" )
             _UpperCAmelCase  :						List[Any]							   =					tf.keras.layers.GlobalAveragePoolingaD(keepdims=A )





      @unpack_inputs
      def   _A			( self		:       str		,      A		:       tf.Tensor		,      A		:       Optional[bool] = None		,      A		:       Optional[bool] = None		,      A		:       bool = False		,      ):
             _UpperCAmelCase  :						Union[str, Any]							   =					(
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )
             _UpperCAmelCase  :						int							   =					return_dict if return_dict is not None else self.config.use_return_dict

             # TF 2.0 image layers can't use NCHW format when running on CPU.
             # We transpose to NHWC format and then transpose back after the full forward pass.
             # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
             _UpperCAmelCase  :						Optional[Any]							   =					tf.transpose(A		,      perm=[0, 2, 3, 1] )
             _UpperCAmelCase  :						int							   =					self.embedder(A		,      training=A )

             _UpperCAmelCase  :						List[Any]							   =					self.encoder(
                 A		,      output_hidden_states=A		,      return_dict=A		,      training=A )

             _UpperCAmelCase  :						Any							   =					encoder_outputs[0]

             _UpperCAmelCase  :						List[Any]							   =					self.pooler(A )

             # Transpose all the outputs to the NCHW format
             # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
             _UpperCAmelCase  :						str							   =					tf.transpose(A		,      (0, 3, 1, 2) )
             _UpperCAmelCase  :						Optional[int]							   =					tf.transpose(A		,      (0, 3, 1, 2) )
             _UpperCAmelCase  :						Dict							   =					()
             for hidden_state in encoder_outputs[1:]:
                    _UpperCAmelCase  :						Optional[Any]							   =					hidden_states + tuple(tf.transpose(A		,      (0, 3, 1, 2) ) for h in hidden_state )

             if not return_dict:
                    return (last_hidden_state, pooled_output) + hidden_states

             _UpperCAmelCase  :						List[Any]							   =					hidden_states if output_hidden_states else None

             return TFBaseModelOutputWithPoolingAndNoAttention(
                 last_hidden_state=A		,      pooler_output=A		,      hidden_states=A		,      )

@add_start_docstrings(
    "The bare ResNet model outputting raw features without any specific head on top."							,				snake_case__							,				)
class 			lowerCamelCase_			(snake_case__					):
      '''simple docstring'''

      def __init__( self		:       str		,      A		:       ResNetConfig		,      **A		:       Optional[int] ):
             super().__init__(A		,      **A )
             _UpperCAmelCase  :						List[str]							   =					TFResNetMainLayer(config=A		,      name="resnet" )





      @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		,      )
      @unpack_inputs
      def   _A			( self		:       Optional[int]		,      A		:       tf.Tensor		,      A		:       Optional[bool] = None		,      A		:       Optional[bool] = None		,      A		:       bool = False		,      ):
             _UpperCAmelCase  :						List[str]							   =					(
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )
             _UpperCAmelCase  :						Tuple							   =					return_dict if return_dict is not None else self.config.use_return_dict

             _UpperCAmelCase  :						str							   =					self.resnet(
                 pixel_values=A		,      output_hidden_states=A		,      return_dict=A		,      training=A		,      )
             return resnet_outputs

@add_start_docstrings(
    "\n    ResNet 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__							,				)
class 			lowerCamelCase_			(snake_case__							,				snake_case__					):
      '''simple docstring'''

      def __init__( self		:       Tuple		,      A		:       ResNetConfig		,      **A		:       Any ):
             super().__init__(A		,      **A )
             _UpperCAmelCase  :						List[Any]							   =					config.num_labels
             _UpperCAmelCase  :						int							   =					TFResNetMainLayer(A		,      name="resnet" )
             # classification head
             _UpperCAmelCase  :						Optional[int]							   =					(
                 tf.keras.layers.Dense(config.num_labels		,      name="classifier.1" )
                 if config.num_labels > 0
                 else tf.keras.layers.Activation("linear"		,      name="classifier.1" )
             )




      def   _A			( self		:       Optional[int]		,      A		:       tf.Tensor ):
             _UpperCAmelCase  :						List[Any]							   =					tf.keras.layers.Flatten()(A )
             _UpperCAmelCase  :						List[str]							   =					self.classifier_layer(A )
             return logits





      @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		,      )
      @unpack_inputs
      def   _A			( self		:       List[Any]		,      A		:       tf.Tensor = None		,      A		:       tf.Tensor = None		,      A		:       bool = None		,      A		:       bool = None		,      A		:       bool = False		,      ):
             _UpperCAmelCase  :						Any							   =					return_dict if return_dict is not None else self.config.use_return_dict

             _UpperCAmelCase  :						Optional[Any]							   =					self.resnet(
                 A		,      output_hidden_states=A		,      return_dict=A		,      training=A )

             _UpperCAmelCase  :						List[Any]							   =					outputs.pooler_output if return_dict else outputs[1]

             _UpperCAmelCase  :						Optional[int]							   =					self.classifier(A )

             _UpperCAmelCase  :						Dict							   =					None if labels is None else self.hf_compute_loss(A		,      A )

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

             return TFImageClassifierOutputWithNoAttention(loss=A		,      logits=A		,      hidden_states=outputs.hidden_states )

