






import inspect
import unittest

from transformers import RegNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
		import jax
		import jax.numpy as jnp

		from transformers.models.regnet.modeling_flax_regnet import FlaxRegNetForImageClassification, FlaxRegNetModel

if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor



class   lowerCamelCase      (unittest.TestCase			):







		'''simple docstring'''



		def __init__(							self				,		_UpperCamelCase				,		_UpperCamelCase=3				,		_UpperCamelCase=3_2				,		_UpperCamelCase=3				,		_UpperCamelCase=1_0				,		_UpperCamelCase=[1_0, 2_0, 3_0, 4_0]				,		_UpperCamelCase=[1, 1, 2, 1]				,		_UpperCamelCase=True				,		_UpperCamelCase=True				,		_UpperCamelCase="relu"				,		_UpperCamelCase=3				,		_UpperCamelCase=None				,		)	->							Union[str, Any]:
							UpperCAmelCase_    :      Any            =						parent
							UpperCAmelCase_    :      List[str]            =						batch_size
							UpperCAmelCase_    :      Optional[int]            =						image_size
							UpperCAmelCase_    :      Union[str, Any]            =						num_channels
							UpperCAmelCase_    :      int            =						embeddings_size
							UpperCAmelCase_    :      Optional[int]            =						hidden_sizes
							UpperCAmelCase_    :      Optional[int]            =						depths
							UpperCAmelCase_    :      Optional[Any]            =						is_training
							UpperCAmelCase_    :      Tuple            =						use_labels
							UpperCAmelCase_    :      List[str]            =						hidden_act
							UpperCAmelCase_    :      List[Any]            =						num_labels
							UpperCAmelCase_    :      Optional[Any]            =						scope
							UpperCAmelCase_    :      Any            =						len(_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							List[Any]:
							UpperCAmelCase_    :      Union[str, Any]            =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

							UpperCAmelCase_    :      str            =						self.get_config()

							return config, pixel_values



		def 					__UpperCAmelCase			(							self   )	->							Optional[int]:
							return RegNetConfig(
							    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 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Dict:
							UpperCAmelCase_    :      Any            =						FlaxRegNetModel(config=_UpperCamelCase   )
							UpperCAmelCase_    :      Union[str, Any]            =						model(_UpperCamelCase   )

							# Output shape (b, c, h, w)
							self.parent.assertEqual(
							    result.last_hidden_state.shape				,		(self.batch_size, self.hidden_sizes[-1], self.image_size // 3_2, self.image_size // 3_2)				,		)



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase   )	->							List[Any]:
							UpperCAmelCase_    :      Union[str, Any]            =						self.num_labels
							UpperCAmelCase_    :      List[Any]            =						FlaxRegNetForImageClassification(config=_UpperCamelCase   )
							UpperCAmelCase_    :      Union[str, Any]            =						model(_UpperCamelCase   )
							self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.num_labels)   )



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							UpperCAmelCase_    :      str            =						self.prepare_config_and_inputs()
							UpperCAmelCase_       ,  UpperCAmelCase_    :      Any            =						config_and_inputs
							UpperCAmelCase_    :      Optional[Any]            =						{'pixel_values': pixel_values}
							return config, inputs_dict





@require_flax
class   lowerCamelCase      (_snake_case ,  unittest.TestCase			):







		'''simple docstring'''



		_snake_case    :    List[str]       						=  (FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

		_snake_case    :    int       						=  False
		_snake_case    :    Optional[Any]       						=  False
		_snake_case    :    Optional[int]       						=  False



		def 					__UpperCAmelCase			(							self   )	->							None:
							UpperCAmelCase_    :      List[Any]            =						FlaxRegNetModelTester(self   )
							UpperCAmelCase_    :      Dict            =						ConfigTester(self				,		config_class=_UpperCamelCase				,		has_text_modality=_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Dict:
							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 					__UpperCAmelCase			(							self   )	->							str:
							return



		def 					__UpperCAmelCase			(							self   )	->							Any:
							UpperCAmelCase_    :      Union[str, Any]            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Optional[int]:
							UpperCAmelCase_    :      str            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*_UpperCamelCase   )



		@unittest.skip(reason='RegNet does not use inputs_embeds'   )
		def 					__UpperCAmelCase			(							self   )	->							str:
							pass



		@unittest.skip(reason='RegNet does not support input and output embeddings'   )
		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							pass



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							UpperCAmelCase_       ,  UpperCAmelCase_    :      Tuple            =						self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												UpperCAmelCase_    :      Tuple            =						model_class(_UpperCamelCase   )
												UpperCAmelCase_    :      Any            =						inspect.signature(model.__call__   )
												# signature.parameters is an OrderedDict => so arg_names order is deterministic
												UpperCAmelCase_    :      Dict            =						[*signature.parameters.keys()]

												UpperCAmelCase_    :      List[str]            =						['pixel_values']
												self.assertListEqual(arg_names[:1]				,		_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							str:
							def check_hidden_states_output(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   ):
												UpperCAmelCase_    :      Optional[Any]            =						model_class(_UpperCamelCase   )

												UpperCAmelCase_    :      str            =						model(**self._prepare_for_class(_UpperCamelCase				,		_UpperCamelCase   )   )

												UpperCAmelCase_    :      Optional[int]            =						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

												UpperCAmelCase_    :      Tuple            =						self.model_tester.num_stages
												self.assertEqual(len(_UpperCamelCase   )				,		expected_num_stages + 1   )

							UpperCAmelCase_       ,  UpperCAmelCase_    :      Optional[int]            =						self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												UpperCAmelCase_    :      Tuple            =						True
												check_hidden_states_output(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )

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

												check_hidden_states_output(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Tuple:
							UpperCAmelCase_       ,  UpperCAmelCase_    :      str            =						self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												with self.subTest(model_class.__name__   ):
																	UpperCAmelCase_    :      Dict            =						self._prepare_for_class(_UpperCamelCase				,		_UpperCamelCase   )
																	UpperCAmelCase_    :      Dict            =						model_class(_UpperCamelCase   )

																	@jax.jit
																	def model_jitted(_UpperCamelCase				,		**_UpperCamelCase   ):
																						return model(pixel_values=_UpperCamelCase				,		**_UpperCamelCase   )

																	with self.subTest('JIT Enabled'   ):
																						UpperCAmelCase_    :      Optional[int]            =						model_jitted(**_UpperCamelCase   ).to_tuple()

																	with self.subTest('JIT Disabled'   ):
																						with jax.disable_jit():
																											UpperCAmelCase_    :      int            =						model_jitted(**_UpperCamelCase   ).to_tuple()

																	self.assertEqual(len(_UpperCamelCase   )				,		len(_UpperCamelCase   )   )
																	for jitted_output, output in zip(_UpperCamelCase				,		_UpperCamelCase   ):
																						self.assertEqual(jitted_output.shape				,		output.shape   )





def      lowercase__      (     ):


					'''simple docstring'''


					UpperCAmelCase_    :      List[str]            =						Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'  )
					return image






@require_flax
class   lowerCamelCase      (unittest.TestCase			):







		'''simple docstring'''



		@cached_property
		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							return AutoImageProcessor.from_pretrained('facebook/regnet-y-040'   ) if is_vision_available() else None



		@slow
		def 					__UpperCAmelCase			(							self   )	->							int:
							UpperCAmelCase_    :      Tuple            =						FlaxRegNetForImageClassification.from_pretrained('facebook/regnet-y-040'   )

							UpperCAmelCase_    :      Tuple            =						self.default_image_processor
							UpperCAmelCase_    :      List[Any]            =						prepare_img()
							UpperCAmelCase_    :      Any            =						image_processor(images=_UpperCamelCase				,		return_tensors='np'   )

							UpperCAmelCase_    :      int            =						model(**_UpperCamelCase   )

							# verify the logits
							UpperCAmelCase_    :      Union[str, Any]            =						(1, 1_0_0_0)
							self.assertEqual(outputs.logits.shape				,		_UpperCamelCase   )

							UpperCAmelCase_    :      Tuple            =						jnp.array([-0.41_80, -1.50_51, -3.48_36]   )

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













from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import get_down_block, get_mid_block, get_out_block, get_up_block

@dataclass
class    lowerCAmelCase							(					__UpperCamelCase			):
 UpperCAmelCase__													=		42


class    lowerCAmelCase							(					__UpperCamelCase,  __UpperCamelCase			):
 @register_to_config
 def __init__(      self  :      List[str]			,     UpperCAmelCase  :      int = 65536			,     UpperCAmelCase  :      Optional[int] = None			,     UpperCAmelCase  :      int = 2			,     UpperCAmelCase  :      int = 2			,     UpperCAmelCase  :      int = 0			,     UpperCAmelCase  :      str = "fourier"			,     UpperCAmelCase  :      bool = True			,     UpperCAmelCase  :      bool = False			,     UpperCAmelCase  :      float = 0.0			,     UpperCAmelCase  :      Tuple[str] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")			,     UpperCAmelCase  :      Tuple[str] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")			,     UpperCAmelCase  :      Tuple[str] = "UNetMidBlock1D"			,     UpperCAmelCase  :      str = None			,     UpperCAmelCase  :      Tuple[int] = (32, 32, 64)			,     UpperCAmelCase  :      str = None			,     UpperCAmelCase  :      int = 8			,     UpperCAmelCase  :      int = 1			,     UpperCAmelCase  :      bool = False			,     )			->	List[Any]:
     super().__init__()
     lowerCamelCase__    :      Optional[int]					 =		sample_size

     # time
     if time_embedding_type == "fourier":
         lowerCamelCase__    :      Optional[Any]					 =		GaussianFourierProjection(
             embedding_size=8			,     set_W_to_weight=UpperCAmelCase			,     log=UpperCAmelCase			,     flip_sin_to_cos=UpperCAmelCase	)
         lowerCamelCase__    :      Any					 =		2 * block_out_channels[0]
     elif time_embedding_type == "positional":
         lowerCamelCase__    :      List[Any]					 =		Timesteps(
             block_out_channels[0]			,     flip_sin_to_cos=UpperCAmelCase			,     downscale_freq_shift=UpperCAmelCase	)
         lowerCamelCase__    :      Dict					 =		block_out_channels[0]

     if use_timestep_embedding:
         lowerCamelCase__    :      str					 =		block_out_channels[0] * 4
         lowerCamelCase__    :      List[Any]					 =		TimestepEmbedding(
             in_channels=UpperCAmelCase			,     time_embed_dim=UpperCAmelCase			,     act_fn=UpperCAmelCase			,     out_dim=block_out_channels[0]			,     )

     lowerCamelCase__    :      Any					 =		nn.ModuleList([]	)
     lowerCamelCase__    :      Tuple					 =		None
     lowerCamelCase__    :      List[str]					 =		nn.ModuleList([]	)
     lowerCamelCase__    :      Optional[int]					 =		None

     # down
     lowerCamelCase__    :      Optional[int]					 =		in_channels
     for i, down_block_type in enumerate(UpperCAmelCase	):
         lowerCamelCase__    :      Union[str, Any]					 =		output_channel
         lowerCamelCase__    :      Tuple					 =		block_out_channels[i]

         if i == 0:
             input_channel += extra_in_channels

         lowerCamelCase__    :      Union[str, Any]					 =		i == len(UpperCAmelCase	) - 1

         lowerCamelCase__    :      Optional[int]					 =		get_down_block(
             UpperCAmelCase			,     num_layers=UpperCAmelCase			,     in_channels=UpperCAmelCase			,     out_channels=UpperCAmelCase			,     temb_channels=block_out_channels[0]			,     add_downsample=not is_final_block or downsample_each_block			,     )
         self.down_blocks.append(UpperCAmelCase	)

     # mid
     lowerCamelCase__    :      Optional[int]					 =		get_mid_block(
         UpperCAmelCase			,     in_channels=block_out_channels[-1]			,     mid_channels=block_out_channels[-1]			,     out_channels=block_out_channels[-1]			,     embed_dim=block_out_channels[0]			,     num_layers=UpperCAmelCase			,     add_downsample=UpperCAmelCase			,     )

     # up
     lowerCamelCase__    :      Optional[int]					 =		list(reversed(UpperCAmelCase	)	)
     lowerCamelCase__    :      Optional[int]					 =		reversed_block_out_channels[0]
     if out_block_type is None:
         lowerCamelCase__    :      List[str]					 =		out_channels
     else:
         lowerCamelCase__    :      Any					 =		block_out_channels[0]

     for i, up_block_type in enumerate(UpperCAmelCase	):
         lowerCamelCase__    :      Tuple					 =		output_channel
         lowerCamelCase__    :      Union[str, Any]					 =		(
             reversed_block_out_channels[i + 1] if i < len(UpperCAmelCase	) - 1 else final_upsample_channels
         )

         lowerCamelCase__    :      List[str]					 =		i == len(UpperCAmelCase	) - 1

         lowerCamelCase__    :      Dict					 =		get_up_block(
             UpperCAmelCase			,     num_layers=UpperCAmelCase			,     in_channels=UpperCAmelCase			,     out_channels=UpperCAmelCase			,     temb_channels=block_out_channels[0]			,     add_upsample=not is_final_block			,     )
         self.up_blocks.append(UpperCAmelCase	)
         lowerCamelCase__    :      int					 =		output_channel

     # out
     lowerCamelCase__    :      int					 =		norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4			,     32	)
     lowerCamelCase__    :      List[Any]					 =		get_out_block(
         out_block_type=UpperCAmelCase			,     num_groups_out=UpperCAmelCase			,     embed_dim=block_out_channels[0]			,     out_channels=UpperCAmelCase			,     act_fn=UpperCAmelCase			,     fc_dim=block_out_channels[-1] // 4			,     )







 def      A_	(      self  :      List[Any]			,     UpperCAmelCase  :      torch.FloatTensor			,     UpperCAmelCase  :      Union[torch.Tensor, float, int]			,     UpperCAmelCase  :      bool = True			,     )			->	Union[UNetaDOutput, Tuple]:
     lowerCamelCase__    :      Optional[Any]					 =		timestep
     if not torch.is_tensor(UpperCAmelCase	):
         lowerCamelCase__    :      Optional[int]					 =		torch.tensor([timesteps]			,     dtype=torch.long			,     device=sample.device	)
     elif torch.is_tensor(UpperCAmelCase	) and len(timesteps.shape	) == 0:
         lowerCamelCase__    :      List[str]					 =		timesteps[None].to(sample.device	)

     lowerCamelCase__    :      Optional[int]					 =		self.time_proj(UpperCAmelCase	)
     if self.config.use_timestep_embedding:
         lowerCamelCase__    :      str					 =		self.time_mlp(UpperCAmelCase	)
     else:
         lowerCamelCase__    :      List[str]					 =		timestep_embed[..., None]
         lowerCamelCase__    :      str					 =		timestep_embed.repeat([1, 1, sample.shape[2]]	).to(sample.dtype	)
         lowerCamelCase__    :      str					 =		timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:])	)

     # 2. down
     lowerCamelCase__    :      str					 =		()
     for downsample_block in self.down_blocks:
         lowerCamelCase__ ,							lowerCamelCase__    :      Union[str, Any]					 =		downsample_block(hidden_states=UpperCAmelCase			,     temb=UpperCAmelCase	)
         down_block_res_samples += res_samples

     # 3. mid
     if self.mid_block:
         lowerCamelCase__    :      Optional[Any]					 =		self.mid_block(UpperCAmelCase			,     UpperCAmelCase	)

     # 4. up
     for i, upsample_block in enumerate(self.up_blocks	):
         lowerCamelCase__    :      Dict					 =		down_block_res_samples[-1:]
         lowerCamelCase__    :      Optional[Any]					 =		down_block_res_samples[:-1]
         lowerCamelCase__    :      Any					 =		upsample_block(UpperCAmelCase			,     res_hidden_states_tuple=UpperCAmelCase			,     temb=UpperCAmelCase	)

     # 5. post-process
     if self.out_block:
         lowerCamelCase__    :      Any					 =		self.out_block(UpperCAmelCase			,     UpperCAmelCase	)

     if not return_dict:
         return (sample,)

     return UNetaDOutput(sample=UpperCAmelCase	)





