



"""simple docstring"""
from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import ResNetConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import TFResNetForImageClassification, TFResNetModel
						from transformers.models.resnet.modeling_tf_resnet import TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import AutoImageProcessor
class 				lowerCamelCase  :






					'''simple docstring'''





					def __init__(       self:							str						,     snake_case:							Dict						,     snake_case:							List[Any]=3						,     snake_case:							Any=32						,     snake_case:							Optional[int]=3						,     snake_case:							List[Any]=10						,     snake_case:							List[str]=[10, 20, 30, 40]						,     snake_case:							Dict=[1, 1, 2, 1]						,     snake_case:							Optional[int]=True						,     snake_case:							Dict=True						,     snake_case:							Union[str, Any]="relu"						,     snake_case:							List[Any]=3						,     snake_case:							Dict=None						,     )   ->	Dict:
											snake_case_ :str						=			parent
											snake_case_ :List[Any]						=			batch_size
											snake_case_ :int						=			image_size
											snake_case_ :Dict						=			num_channels
											snake_case_ :Any						=			embeddings_size
											snake_case_ :str						=			hidden_sizes
											snake_case_ :Tuple						=			depths
											snake_case_ :str						=			is_training
											snake_case_ :int						=			use_labels
											snake_case_ :Optional[int]						=			hidden_act
											snake_case_ :Dict						=			num_labels
											snake_case_ :Tuple						=			scope
											snake_case_ :List[Any]						=			len(snake_case						)






					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	Optional[Any]:
											snake_case_ :Dict						=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

											snake_case_ :Any						=			None
											if self.use_labels:
																	snake_case_ :Any						=			ids_tensor([self.batch_size]						,     self.num_labels						)

											snake_case_ :Union[str, Any]						=			self.get_config()

											return config, pixel_values, labels






					def    lowerCAmelCase_				(       self:							Optional[int]						)   ->	Union[str, Any]:
											return ResNetConfig(
											    num_channels=self.num_channels						,     embeddings_size=self.embeddings_size						,     hidden_sizes=self.hidden_sizes						,     depths=self.depths						,     hidden_act=self.hidden_act						,     num_labels=self.num_labels						,     image_size=self.image_size						,     )






					def    lowerCAmelCase_				(       self:							int						,     snake_case:							List[str]						,     snake_case:							Any						,     snake_case:							Optional[int]						)   ->	Tuple:
											snake_case_ :Dict						=			TFResNetModel(config=snake_case						)
											snake_case_ :Dict						=			model(snake_case						)
											# expected last hidden states: B, C, H // 32, W // 32
											self.parent.assertEqual(
											    result.last_hidden_state.shape						,     (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)						,     )






					def    lowerCAmelCase_				(       self:							str						,     snake_case:							List[str]						,     snake_case:							Union[str, Any]						,     snake_case:							Any						)   ->	int:
											snake_case_ :Any						=			self.num_labels
											snake_case_ :Optional[int]						=			TFResNetForImageClassification(snake_case						)
											snake_case_ :List[Any]						=			model(snake_case						,     labels=snake_case						)
											self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.num_labels)						)






					def    lowerCAmelCase_				(       self:							Tuple						)   ->	Union[str, Any]:
											snake_case_ :Optional[int]						=			self.prepare_config_and_inputs()
											snake_case_,       snake_case_,       snake_case_ :Tuple						=			config_and_inputs
											snake_case_ :Tuple						=			{"""pixel_values""": pixel_values}
											return config, inputs_dict

@require_tf
class 				lowerCamelCase  (			_lowerCAmelCase    ,						_lowerCAmelCase    ,						unittest.TestCase				):






					'''simple docstring'''





					_A		:			Union[str, Any]						=     (TFResNetModel, TFResNetForImageClassification) if is_tf_available() else ()
					_A		:			List[str]						=     (
					    {"""feature-extraction""": TFResNetModel, """image-classification""": TFResNetForImageClassification}
					    if is_tf_available()
					    else {}
					)

					_A		:			List[str]						=     False
					_A		:			Any						=     False
					_A		:			int						=     False
					_A		:			List[Any]						=     False
					_A		:			Any						=     False






					def    lowerCAmelCase_				(       self:							str						)   ->	Optional[Any]:
											snake_case_ :List[str]						=			TFResNetModelTester(self						)
											snake_case_ :List[Any]						=			ConfigTester(self						,     config_class=snake_case						,     has_text_modality=snake_case						)






					def    lowerCAmelCase_				(       self:							int						)   ->	List[Any]:
											self.create_and_test_config_common_properties()
											self.config_tester.create_and_test_config_to_json_string()
											self.config_tester.create_and_test_config_to_json_file()
											self.config_tester.create_and_test_config_from_and_save_pretrained()
											self.config_tester.create_and_test_config_with_num_labels()
											self.config_tester.check_config_can_be_init_without_params()
											self.config_tester.check_config_arguments_init()






					def    lowerCAmelCase_				(       self:							List[str]						)   ->	List[str]:
											return






					@unittest.skip(reason="""ResNet does not use inputs_embeds"""						)
					def    lowerCAmelCase_				(       self:							Dict						)   ->	Optional[Any]:
											pass






					@unittest.skip(reason="""ResNet does not support input and output embeddings"""						)
					def    lowerCAmelCase_				(       self:							Dict						)   ->	Optional[Any]:
											pass






					def    lowerCAmelCase_				(       self:							Dict						)   ->	Optional[int]:
											snake_case_,       snake_case_ :Tuple						=			self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	snake_case_ :Tuple						=			model_class(snake_case						)
																	snake_case_ :Optional[int]						=			inspect.signature(model.call						)
																	# signature.parameters is an OrderedDict => so arg_names order is deterministic
																	snake_case_ :Optional[int]						=			[*signature.parameters.keys()]

																	snake_case_ :str						=			["""pixel_values"""]
																	self.assertListEqual(arg_names[:1]						,     snake_case						)






					def    lowerCAmelCase_				(       self:							List[Any]						)   ->	List[str]:
											snake_case_ :Dict						=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*snake_case						)






					def    lowerCAmelCase_				(       self:							str						)   ->	List[Any]:
											def check_hidden_states_output(snake_case:							List[str]						,     snake_case:							List[str]						,     snake_case:							str						):
																	snake_case_ :Optional[Any]						=			model_class(snake_case						)
																	snake_case_ :Optional[int]						=			model(**self._prepare_for_class(snake_case						,     snake_case						)						)

																	snake_case_ :List[str]						=			outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

																	snake_case_ :List[str]						=			self.model_tester.num_stages
																	self.assertEqual(len(snake_case						)						,     expected_num_stages + 1						)

																	# ResNet's feature maps are of shape (batch_size, num_channels, height, width)
																	self.assertListEqual(
																	    list(hidden_states[0].shape[-2:]						)						,     [self.model_tester.image_size // 4, self.model_tester.image_size // 4]						,     )

											snake_case_,       snake_case_ :Optional[Any]						=			self.model_tester.prepare_config_and_inputs_for_common()
											snake_case_ :Tuple						=			["""basic""", """bottleneck"""]
											for model_class in self.all_model_classes:
																	for layer_type in layers_type:
																							snake_case_ :Any						=			layer_type
																							snake_case_ :List[Any]						=			True
																							check_hidden_states_output(snake_case						,     snake_case						,     snake_case						)

																							# check that output_hidden_states also work using config
																							del inputs_dict["output_hidden_states"]
																							snake_case_ :str						=			True

																							check_hidden_states_output(snake_case						,     snake_case						,     snake_case						)






					def    lowerCAmelCase_				(       self:							List[Any]						)   ->	Optional[Any]:
											snake_case_ :str						=			self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*snake_case						)






					@slow
					def    lowerCAmelCase_				(       self:							Dict						)   ->	Optional[int]:
											for model_name in TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	snake_case_ :int						=			TFResNetModel.from_pretrained(snake_case						)
																	self.assertIsNotNone(snake_case						)

def    A_  (			):







						'''simple docstring'''


						snake_case_ :Dict						=			Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""	)
						return image

@require_tf
@require_vision
class 				lowerCamelCase  (			unittest.TestCase				):






					'''simple docstring'''





					@cached_property
					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	Any:
											return (
											    AutoImageProcessor.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]						)
											    if is_vision_available()
											    else None
											)






					@slow
					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	List[str]:
											snake_case_ :Any						=			TFResNetForImageClassification.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]						)

											snake_case_ :int						=			self.default_image_processor
											snake_case_ :Union[str, Any]						=			prepare_img()
											snake_case_ :Optional[Any]						=			image_processor(images=snake_case						,     return_tensors="""tf"""						)

											# forward pass
											snake_case_ :List[str]						=			model(**snake_case						)

											# verify the logits
											snake_case_ :Optional[Any]						=			tf.TensorShape((1, 1_000)						)
											self.assertEqual(outputs.logits.shape						,     snake_case						)

											snake_case_ :Any						=			tf.constant([-1_1.1_0_6_9, -9.7_8_7_7, -8.3_7_7_7]						)

											self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()						,     snake_case						,     atol=1E-4						)						)







from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)





@flax.struct.dataclass
class A_					(			_lowerCamelCase   ):
   lowerCAmelCase__        = 42





@flax_register_to_config
class A_					(			nn.Module ,      _lowerCamelCase ,      _lowerCamelCase   ):
   lowerCAmelCase__        = 32
   lowerCAmelCase__        = 4
   lowerCAmelCase__        = 4
   lowerCAmelCase__        = (
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "DownBlock2D",
   )
   lowerCAmelCase__        = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
   lowerCAmelCase__        = False
   lowerCAmelCase__        = (320, 640, 1280, 1280)
   lowerCAmelCase__        = 2
   lowerCAmelCase__        = 8
   lowerCAmelCase__        = None
   lowerCAmelCase__        = 1280
   lowerCAmelCase__        = 0.0
   lowerCAmelCase__        = False
   lowerCAmelCase__        = jnp.floataa
   lowerCAmelCase__        = True
   lowerCAmelCase__        = 0
   lowerCAmelCase__        = False




   def _lowerCAmelCase      (self	:Tuple    ,   _UpperCamelCase	:jax.random.KeyArray )->					FrozenDict:
         # init input tensors
         __A          =				(1, self.in_channels, self.sample_size, self.sample_size)
         __A          =				jnp.zeros(_UpperCamelCase    ,   dtype=jnp.floataa )
         __A          =				jnp.ones((1,)    ,   dtype=jnp.intaa )
         __A          =				jnp.zeros((1, 1, self.cross_attention_dim)    ,   dtype=jnp.floataa )

         __A   ,						__A          =				jax.random.split(_UpperCamelCase )
         __A          =				{'''params''': params_rng, '''dropout''': dropout_rng}

         return self.init(_UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase )["params"]




   def _lowerCAmelCase      (self	:Tuple )->					Optional[int]:
         __A          =				self.block_out_channels
         __A          =				block_out_channels[0] * 4

         if self.num_attention_heads is not None:
               raise ValueError(
                   '''At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.''' )

         # If `num_attention_heads` is not defined (which is the case for most models)
         # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
         # The reason for this behavior is to correct for incorrectly named variables that were introduced
         # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
         # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
         # which is why we correct for the naming here.
         __A          =				self.num_attention_heads or self.attention_head_dim

         # input
         __A          =				nn.Conv(
             block_out_channels[0]    ,   kernel_size=(3, 3)    ,   strides=(1, 1)    ,   padding=((1, 1), (1, 1))    ,   dtype=self.dtype    ,   )

         # time
         __A          =				FlaxTimesteps(
             block_out_channels[0]    ,   flip_sin_to_cos=self.flip_sin_to_cos    ,   freq_shift=self.config.freq_shift )
         __A          =				FlaxTimestepEmbedding(_UpperCamelCase    ,   dtype=self.dtype )

         __A          =				self.only_cross_attention
         if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
               __A          =				(only_cross_attention,) * len(self.down_block_types )

         if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
               __A          =				(num_attention_heads,) * len(self.down_block_types )

         # down
         __A          =				[]
         __A          =				block_out_channels[0]
         for i, down_block_type in enumerate(self.down_block_types ):
               __A          =				output_channel
               __A          =				block_out_channels[i]
               __A          =				i == len(_UpperCamelCase ) - 1

               if down_block_type == "CrossAttnDownBlock2D":
                     __A          =				FlaxCrossAttnDownBlockaD(
                         in_channels=_UpperCamelCase    ,   out_channels=_UpperCamelCase    ,   dropout=self.dropout    ,   num_layers=self.layers_per_block    ,   num_attention_heads=num_attention_heads[i]    ,   add_downsample=not is_final_block    ,   use_linear_projection=self.use_linear_projection    ,   only_cross_attention=only_cross_attention[i]    ,   use_memory_efficient_attention=self.use_memory_efficient_attention    ,   dtype=self.dtype    ,   )
               else:
                     __A          =				FlaxDownBlockaD(
                         in_channels=_UpperCamelCase    ,   out_channels=_UpperCamelCase    ,   dropout=self.dropout    ,   num_layers=self.layers_per_block    ,   add_downsample=not is_final_block    ,   dtype=self.dtype    ,   )

               down_blocks.append(_UpperCamelCase )
         __A          =				down_blocks

         # mid
         __A          =				FlaxUNetMidBlockaDCrossAttn(
             in_channels=block_out_channels[-1]    ,   dropout=self.dropout    ,   num_attention_heads=num_attention_heads[-1]    ,   use_linear_projection=self.use_linear_projection    ,   use_memory_efficient_attention=self.use_memory_efficient_attention    ,   dtype=self.dtype    ,   )

         # up
         __A          =				[]
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				list(reversed(_UpperCamelCase ) )
         __A          =				reversed_block_out_channels[0]
         for i, up_block_type in enumerate(self.up_block_types ):
               __A          =				output_channel
               __A          =				reversed_block_out_channels[i]
               __A          =				reversed_block_out_channels[min(i + 1    ,   len(_UpperCamelCase ) - 1 )]

               __A          =				i == len(_UpperCamelCase ) - 1

               if up_block_type == "CrossAttnUpBlock2D":
                     __A          =				FlaxCrossAttnUpBlockaD(
                         in_channels=_UpperCamelCase    ,   out_channels=_UpperCamelCase    ,   prev_output_channel=_UpperCamelCase    ,   num_layers=self.layers_per_block + 1    ,   num_attention_heads=reversed_num_attention_heads[i]    ,   add_upsample=not is_final_block    ,   dropout=self.dropout    ,   use_linear_projection=self.use_linear_projection    ,   only_cross_attention=only_cross_attention[i]    ,   use_memory_efficient_attention=self.use_memory_efficient_attention    ,   dtype=self.dtype    ,   )
               else:
                     __A          =				FlaxUpBlockaD(
                         in_channels=_UpperCamelCase    ,   out_channels=_UpperCamelCase    ,   prev_output_channel=_UpperCamelCase    ,   num_layers=self.layers_per_block + 1    ,   add_upsample=not is_final_block    ,   dropout=self.dropout    ,   dtype=self.dtype    ,   )

               up_blocks.append(_UpperCamelCase )
               __A          =				output_channel
         __A          =				up_blocks

         # out
         __A          =				nn.GroupNorm(num_groups=32    ,   epsilon=1e-5 )
         __A          =				nn.Conv(
             self.out_channels    ,   kernel_size=(3, 3)    ,   strides=(1, 1)    ,   padding=((1, 1), (1, 1))    ,   dtype=self.dtype    ,   )






   def __call__(self	:Union[str, Any]    ,   _UpperCamelCase	:Any    ,   _UpperCamelCase	:List[str]    ,   _UpperCamelCase	:str    ,   _UpperCamelCase	:int=None    ,   _UpperCamelCase	:Optional[int]=None    ,   _UpperCamelCase	:bool = True    ,   _UpperCamelCase	:bool = False    ,   )->					Union[FlaxUNetaDConditionOutput, Tuple]:
         # 1. time
         if not isinstance(_UpperCamelCase    ,   jnp.ndarray ):
               __A          =				jnp.array([timesteps]    ,   dtype=jnp.intaa )
         elif isinstance(_UpperCamelCase    ,   jnp.ndarray ) and len(timesteps.shape ) == 0:
               __A          =				timesteps.astype(dtype=jnp.floataa )
               __A          =				jnp.expand_dims(_UpperCamelCase    ,   0 )

         __A          =				self.time_proj(_UpperCamelCase )
         __A          =				self.time_embedding(_UpperCamelCase )

         # 2. pre-process
         __A          =				jnp.transpose(_UpperCamelCase    ,   (0, 2, 3, 1) )
         __A          =				self.conv_in(_UpperCamelCase )

         # 3. down
         __A          =				(sample,)
         for down_block in self.down_blocks:
               if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
                     __A   ,						__A          =				down_block(_UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase    ,   deterministic=not train )
               else:
                     __A   ,						__A          =				down_block(_UpperCamelCase    ,   _UpperCamelCase    ,   deterministic=not train )
               down_block_res_samples += res_samples

         if down_block_additional_residuals is not None:
               __A          =				()

               for down_block_res_sample, down_block_additional_residual in zip(
                   _UpperCamelCase    ,   _UpperCamelCase ):
                     down_block_res_sample += down_block_additional_residual
                     new_down_block_res_samples += (down_block_res_sample,)

               __A          =				new_down_block_res_samples

         # 4. mid
         __A          =				self.mid_block(_UpperCamelCase    ,   _UpperCamelCase    ,   _UpperCamelCase    ,   deterministic=not train )

         if mid_block_additional_residual is not None:
               sample += mid_block_additional_residual

         # 5. up
         for up_block in self.up_blocks:
               __A          =				down_block_res_samples[-(self.layers_per_block + 1) :]
               __A          =				down_block_res_samples[: -(self.layers_per_block + 1)]
               if isinstance(_UpperCamelCase    ,   _UpperCamelCase ):
                     __A          =				up_block(
                         _UpperCamelCase    ,   temb=_UpperCamelCase    ,   encoder_hidden_states=_UpperCamelCase    ,   res_hidden_states_tuple=_UpperCamelCase    ,   deterministic=not train    ,   )
               else:
                     __A          =				up_block(_UpperCamelCase    ,   temb=_UpperCamelCase    ,   res_hidden_states_tuple=_UpperCamelCase    ,   deterministic=not train )

        # 6. post-process
         __A          =				self.conv_norm_out(_UpperCamelCase )
         __A          =				nn.silu(_UpperCamelCase )
         __A          =				self.conv_out(_UpperCamelCase )
         __A          =				jnp.transpose(_UpperCamelCase    ,   (0, 3, 1, 2) )

         if not return_dict:
               return (sample,)

         return FlaxUNetaDConditionOutput(sample=_UpperCamelCase )





