


import gc
import unittest

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

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()







class      _lowercase					(				lowercase__				,    unittest.TestCase):


       """simple docstring"""




       A__   		=     UnCLIPPipeline
       A__   		=     TEXT_TO_IMAGE_PARAMS - {
           "negative_prompt",
           "height",
           "width",
           "negative_prompt_embeds",
           "guidance_scale",
           "prompt_embeds",
           "cross_attention_kwargs",
       }
       A__   		=     TEXT_TO_IMAGE_BATCH_PARAMS
       A__   		=     [
           "generator",
           "return_dict",
           "prior_num_inference_steps",
           "decoder_num_inference_steps",
           "super_res_num_inference_steps",
       ]
       A__   		=     False





       @property
       def 		lowerCAmelCase	(  self				:     List[Any]			):


             '''simple docstring'''



             return 32





       @property
       def 		lowerCAmelCase	(  self				:     List[Any]			):


             '''simple docstring'''



             return 32





       @property
       def 		lowerCAmelCase	(  self				:     Optional[int]			):


             '''simple docstring'''



             return self.time_input_dim





       @property
       def 		lowerCAmelCase	(  self				:     List[Any]			):


             '''simple docstring'''



             return self.time_input_dim * 4





       @property
       def 		lowerCAmelCase	(  self				:     int			):


             '''simple docstring'''



             return 100





       @property
       def 		lowerCAmelCase	(  self				:     int			):


             '''simple docstring'''



             lowerCamelCase__						:							str       =    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"			)
             return tokenizer





       @property
       def 		lowerCAmelCase	(  self				:     int			):


             '''simple docstring'''



             torch.manual_seed(0			)
             lowerCamelCase__						:							Optional[int]       =    CLIPTextConfig(
                 bos_token_id=0   , eos_token_id=2   , hidden_size=self.text_embedder_hidden_size   , projection_dim=self.text_embedder_hidden_size   , intermediate_size=37   , layer_norm_eps=1E-0_5   , num_attention_heads=4   , num_hidden_layers=5   , pad_token_id=1   , vocab_size=1000   , )
             return CLIPTextModelWithProjection(__lowerCamelCase			)





       @property
       def 		lowerCAmelCase	(  self				:     Any			):


             '''simple docstring'''



             torch.manual_seed(0			)

             lowerCamelCase__						:							Union[str, Any]       =    {
                 "num_attention_heads": 2,
                 "attention_head_dim": 12,
                 "embedding_dim": self.text_embedder_hidden_size,
                 "num_layers": 1,
             }

             lowerCamelCase__						:							int       =    PriorTransformer(**__lowerCamelCase			)
             return model





       @property
       def 		lowerCAmelCase	(  self				:     int			):


             '''simple docstring'''



             torch.manual_seed(0			)

             lowerCamelCase__						:							Union[str, Any]       =    {
                 "clip_embeddings_dim": self.text_embedder_hidden_size,
                 "time_embed_dim": self.time_embed_dim,
                 "cross_attention_dim": self.cross_attention_dim,
             }

             lowerCamelCase__						:							List[Any]       =    UnCLIPTextProjModel(**__lowerCamelCase			)
             return model





       @property
       def 		lowerCAmelCase	(  self				:     Union[str, Any]			):


             '''simple docstring'''



             torch.manual_seed(0			)

             lowerCamelCase__						:							Optional[Any]       =    {
                 "sample_size": 32,
                 # RGB in channels
                 "in_channels": 3,
                 # Out channels is double in channels because predicts mean and variance
                 "out_channels": 6,
                 "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
                 "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
                 "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
                 "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
                 "layers_per_block": 1,
                 "cross_attention_dim": self.cross_attention_dim,
                 "attention_head_dim": 4,
                 "resnet_time_scale_shift": "scale_shift",
                 "class_embed_type": "identity",
             }

             lowerCamelCase__						:							Any       =    UNetaDConditionModel(**__lowerCamelCase			)
             return model





       @property
       def 		lowerCAmelCase	(  self				:     Optional[Any]			):


             '''simple docstring'''



             return {
                 "sample_size": 64,
                 "layers_per_block": 1,
                 "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
                 "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
                 "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
                 "in_channels": 6,
                 "out_channels": 3,
             }





       @property
       def 		lowerCAmelCase	(  self				:     List[Any]			):


             '''simple docstring'''



             torch.manual_seed(0			)

             lowerCamelCase__						:							int       =    UNetaDModel(**self.dummy_super_res_kwargs			)
             return model





       @property
       def 		lowerCAmelCase	(  self				:     Tuple			):


             '''simple docstring'''



             torch.manual_seed(1			)

             lowerCamelCase__						:							str       =    UNetaDModel(**self.dummy_super_res_kwargs			)
             return model





       def 		lowerCAmelCase	(  self				:     str			):


             '''simple docstring'''



             lowerCamelCase__						:							List[str]       =    self.dummy_prior
             lowerCamelCase__						:							str       =    self.dummy_decoder
             lowerCamelCase__						:							Tuple       =    self.dummy_text_proj
             lowerCamelCase__						:							List[str]       =    self.dummy_text_encoder
             lowerCamelCase__						:							Dict       =    self.dummy_tokenizer
             lowerCamelCase__						:							Any       =    self.dummy_super_res_first
             lowerCamelCase__						:							Union[str, Any]       =    self.dummy_super_res_last

             lowerCamelCase__						:							Tuple       =    UnCLIPScheduler(
                 variance_type="fixed_small_log"   , prediction_type="sample"   , num_train_timesteps=1000   , clip_sample_range=5.0   , )

             lowerCamelCase__						:							Tuple       =    UnCLIPScheduler(
                 variance_type="learned_range"   , prediction_type="epsilon"   , num_train_timesteps=1000   , )

             lowerCamelCase__						:							Tuple       =    UnCLIPScheduler(
                 variance_type="fixed_small_log"   , prediction_type="epsilon"   , num_train_timesteps=1000   , )

             lowerCamelCase__						:							List[Any]       =    {
                 "prior": prior,
                 "decoder": decoder,
                 "text_proj": text_proj,
                 "text_encoder": text_encoder,
                 "tokenizer": tokenizer,
                 "super_res_first": super_res_first,
                 "super_res_last": super_res_last,
                 "prior_scheduler": prior_scheduler,
                 "decoder_scheduler": decoder_scheduler,
                 "super_res_scheduler": super_res_scheduler,
             }

             return components





       def 		lowerCAmelCase	(  self				:     Any   , __lowerCamelCase				:     Optional[Any]   , __lowerCamelCase				:     Optional[Any]=0			):


             '''simple docstring'''



             if str(__lowerCamelCase			).startswith("mps"			):
                   lowerCamelCase__						:							Tuple       =    torch.manual_seed(__lowerCamelCase			)
             else:
                   lowerCamelCase__						:							str       =    torch.Generator(device=__lowerCamelCase			).manual_seed(__lowerCamelCase			)
             lowerCamelCase__						:							List[str]       =    {
                 "prompt": "horse",
                 "generator": generator,
                 "prior_num_inference_steps": 2,
                 "decoder_num_inference_steps": 2,
                 "super_res_num_inference_steps": 2,
                 "output_type": "numpy",
             }
             return inputs





       def 		lowerCAmelCase	(  self				:     Optional[Any]			):


             '''simple docstring'''



             lowerCamelCase__						:							Optional[Any]       =    "cpu"

             lowerCamelCase__						:							Union[str, Any]       =    self.get_dummy_components()

             lowerCamelCase__						:							Optional[Any]       =    self.pipeline_class(**__lowerCamelCase			)
             lowerCamelCase__						:							Any       =    pipe.to(__lowerCamelCase			)

             pipe.set_progress_bar_config(disable=__lowerCamelCase			)

             lowerCamelCase__						:							Union[str, Any]       =    pipe(**self.get_dummy_inputs(__lowerCamelCase			)			)
             lowerCamelCase__						:							Optional[Any]       =    output.images

             lowerCamelCase__						:							Union[str, Any]       =    pipe(
                 **self.get_dummy_inputs(__lowerCamelCase			)   , return_dict=__lowerCamelCase   , )[0]

             lowerCamelCase__						:							str       =    image[0, -3:, -3:, -1]
             lowerCamelCase__						:							List[str]       =    image_from_tuple[0, -3:, -3:, -1]

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

             lowerCamelCase__						:							int       =    np.array(
                 [
                     0.9_9_9_7,
                     0.9_9_8_8,
                     0.0_0_2_8,
                     0.9_9_9_7,
                     0.9_9_8_4,
                     0.9_9_6_5,
                     0.0_0_2_9,
                     0.9_9_8_6,
                     0.0_0_2_5,
                 ]			)

             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 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             lowerCamelCase__						:							List[str]       =    torch.device("cpu"			)







             class      _lowercase					:


                    """simple docstring"""




                    A__   		=     1



             lowerCamelCase__						:							Tuple       =    self.get_dummy_components()

             lowerCamelCase__						:							Tuple       =    self.pipeline_class(**__lowerCamelCase			)
             lowerCamelCase__						:							Union[str, Any]       =    pipe.to(__lowerCamelCase			)

             lowerCamelCase__						:							Any       =    components["prior"]
             lowerCamelCase__						:							int       =    components["decoder"]
             lowerCamelCase__						:							str       =    components["super_res_first"]
             lowerCamelCase__						:							Dict       =    components["tokenizer"]
             lowerCamelCase__						:							Any       =    components["text_encoder"]

             lowerCamelCase__						:							Optional[int]       =    torch.Generator(device=__lowerCamelCase			).manual_seed(0			)
             lowerCamelCase__						:							Optional[Any]       =    prior.dtype
             lowerCamelCase__						:							Optional[int]       =    1

             lowerCamelCase__						:							Optional[int]       =    (batch_size, prior.config.embedding_dim)
             lowerCamelCase__						:							Dict       =    pipe.prepare_latents(
                 __lowerCamelCase   , dtype=__lowerCamelCase   , device=__lowerCamelCase   , generator=__lowerCamelCase   , latents=__lowerCamelCase   , scheduler=DummyScheduler()			)
             lowerCamelCase__						:							int       =    (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
             lowerCamelCase__						:							Any       =    pipe.prepare_latents(
                 __lowerCamelCase   , dtype=__lowerCamelCase   , device=__lowerCamelCase   , generator=__lowerCamelCase   , latents=__lowerCamelCase   , scheduler=DummyScheduler()			)

             lowerCamelCase__						:							List[Any]       =    (
                 batch_size,
                 super_res_first.config.in_channels // 2,
                 super_res_first.config.sample_size,
                 super_res_first.config.sample_size,
             )
             lowerCamelCase__						:							Tuple       =    pipe.prepare_latents(
                 __lowerCamelCase   , dtype=__lowerCamelCase   , device=__lowerCamelCase   , generator=__lowerCamelCase   , latents=__lowerCamelCase   , scheduler=DummyScheduler()			)

             pipe.set_progress_bar_config(disable=__lowerCamelCase			)

             lowerCamelCase__						:							Optional[Any]       =    "this is a prompt example"

             lowerCamelCase__						:							Dict       =    torch.Generator(device=__lowerCamelCase			).manual_seed(0			)
             lowerCamelCase__						:							Union[str, Any]       =    pipe(
                 [prompt]   , generator=__lowerCamelCase   , prior_num_inference_steps=2   , decoder_num_inference_steps=2   , super_res_num_inference_steps=2   , prior_latents=__lowerCamelCase   , decoder_latents=__lowerCamelCase   , super_res_latents=__lowerCamelCase   , output_type="np"   , )
             lowerCamelCase__						:							Optional[Any]       =    output.images

             lowerCamelCase__						:							List[Any]       =    tokenizer(
                 __lowerCamelCase   , padding="max_length"   , max_length=tokenizer.model_max_length   , return_tensors="pt"   , )
             lowerCamelCase__						:							List[str]       =    text_encoder(text_inputs.input_ids			)
             lowerCamelCase__						:							List[str]       =    text_inputs.attention_mask

             lowerCamelCase__						:							Dict       =    torch.Generator(device=__lowerCamelCase			).manual_seed(0			)
             lowerCamelCase__						:							Dict       =    pipe(
                 generator=__lowerCamelCase   , prior_num_inference_steps=2   , decoder_num_inference_steps=2   , super_res_num_inference_steps=2   , prior_latents=__lowerCamelCase   , decoder_latents=__lowerCamelCase   , super_res_latents=__lowerCamelCase   , text_model_output=__lowerCamelCase   , text_attention_mask=__lowerCamelCase   , output_type="np"   , )[0]

             # make sure passing text embeddings manually is identical
             assert np.abs(image - image_from_text			).max() < 1E-4





       @skip_mps
       def 		lowerCAmelCase	(  self				:     str			):


             '''simple docstring'''



             lowerCamelCase__						:							Optional[int]       =    torch_device == "cpu"

             self._test_attention_slicing_forward_pass(test_max_difference=__lowerCamelCase   , expected_max_diff=0.0_1			)





       @skip_mps
       def 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             lowerCamelCase__						:							Optional[Any]       =    torch_device == "cpu"
             lowerCamelCase__						:							int       =    True
             lowerCamelCase__						:							Dict       =    [
                 "prior_num_inference_steps",
                 "decoder_num_inference_steps",
                 "super_res_num_inference_steps",
             ]

             self._test_inference_batch_single_identical(
                 test_max_difference=__lowerCamelCase   , relax_max_difference=__lowerCamelCase   , additional_params_copy_to_batched_inputs=__lowerCamelCase   , )





       def 		lowerCAmelCase	(  self				:     str			):


             '''simple docstring'''



             lowerCamelCase__						:							List[str]       =    [
                 "prior_num_inference_steps",
                 "decoder_num_inference_steps",
                 "super_res_num_inference_steps",
             ]

             if torch_device == "mps":
                   # TODO: MPS errors with larger batch sizes
                   lowerCamelCase__						:							Dict       =    [2, 3]
                   self._test_inference_batch_consistent(
                       batch_sizes=__lowerCamelCase   , additional_params_copy_to_batched_inputs=__lowerCamelCase   , )
             else:
                   self._test_inference_batch_consistent(
                       additional_params_copy_to_batched_inputs=__lowerCamelCase			)





       @skip_mps
       def 		lowerCAmelCase	(  self				:     List[Any]			):


             '''simple docstring'''



             return super().test_dict_tuple_outputs_equivalent()





       @skip_mps
       def 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             return super().test_save_load_local()





       @skip_mps
       def 		lowerCAmelCase	(  self				:     Optional[int]			):


             '''simple docstring'''



             return super().test_save_load_optional_components()







@nightly
class      _lowercase					(				unittest.TestCase):


       """simple docstring"""




       def 		lowerCAmelCase	(  self				:     Tuple			):


             '''simple docstring'''



             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()





       def 		lowerCAmelCase	(  self				:     str			):


             '''simple docstring'''



             lowerCamelCase__						:							Dict       =    load_numpy(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                 "/unclip/karlo_v1_alpha_horse_cpu.npy"			)

             lowerCamelCase__						:							List[str]       =    UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"			)
             pipeline.set_progress_bar_config(disable=__lowerCamelCase			)

             lowerCamelCase__						:							List[Any]       =    torch.manual_seed(0			)
             lowerCamelCase__						:							Optional[int]       =    pipeline(
                 "horse"   , num_images_per_prompt=1   , generator=__lowerCamelCase   , output_type="np"   , )

             lowerCamelCase__						:							Any       =    output.images[0]

             assert image.shape == (256, 256, 3)
             assert np.abs(expected_image - image			).max() < 1E-1







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


       """simple docstring"""




       def 		lowerCAmelCase	(  self				:     Any			):


             '''simple docstring'''



             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()





       def 		lowerCAmelCase	(  self				:     Union[str, Any]			):


             '''simple docstring'''



             lowerCamelCase__						:							Union[str, Any]       =    load_numpy(
                 "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
                 "/unclip/karlo_v1_alpha_horse_fp16.npy"			)

             lowerCamelCase__						:							Tuple       =    UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"   , torch_dtype=torch.floataa			)
             lowerCamelCase__						:							int       =    pipeline.to(__lowerCamelCase			)
             pipeline.set_progress_bar_config(disable=__lowerCamelCase			)

             lowerCamelCase__						:							Any       =    torch.Generator(device="cpu"			).manual_seed(0			)
             lowerCamelCase__						:							List[str]       =    pipeline(
                 "horse"   , generator=__lowerCamelCase   , output_type="np"   , )

             lowerCamelCase__						:							Optional[int]       =    output.images[0]

             assert image.shape == (256, 256, 3)

             assert_mean_pixel_difference(__lowerCamelCase   , __lowerCamelCase			)





       def 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             torch.cuda.empty_cache()
             torch.cuda.reset_max_memory_allocated()
             torch.cuda.reset_peak_memory_stats()

             lowerCamelCase__						:							Tuple       =    UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"   , torch_dtype=torch.floataa			)
             lowerCamelCase__						:							Dict       =    pipe.to(__lowerCamelCase			)
             pipe.set_progress_bar_config(disable=__lowerCamelCase			)
             pipe.enable_attention_slicing()
             pipe.enable_sequential_cpu_offload()

             lowerCamelCase__						:							int       =    pipe(
                 "horse"   , num_images_per_prompt=1   , prior_num_inference_steps=2   , decoder_num_inference_steps=2   , super_res_num_inference_steps=2   , output_type="np"   , )

             lowerCamelCase__						:							str       =    torch.cuda.max_memory_allocated()
             # make sure that less than 7 GB is allocated
             assert mem_bytes < 7 * 10**9






import inspect
import unittest

from transformers import ResNetConfig, 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.resnet.modeling_flax_resnet import FlaxResNetForImageClassification, FlaxResNetModel

if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor







class      _lowercase					(				unittest.TestCase):


       """simple docstring"""




       def __init__(  self				:     Dict   , __lowerCamelCase				:     int   , __lowerCamelCase				:     List[Any]=3   , __lowerCamelCase				:     Dict=32   , __lowerCamelCase				:     Optional[Any]=3   , __lowerCamelCase				:     List[str]=10   , __lowerCamelCase				:     List[str]=[10, 20, 30, 40]   , __lowerCamelCase				:     List[Any]=[1, 1, 2, 1]   , __lowerCamelCase				:     Any=True   , __lowerCamelCase				:     Optional[int]=True   , __lowerCamelCase				:     int="relu"   , __lowerCamelCase				:     List[str]=3   , __lowerCamelCase				:     Any=None   , ):


             '''simple docstring'''



             lowerCamelCase__						:							int       =    parent
             lowerCamelCase__						:							Optional[int]       =    batch_size
             lowerCamelCase__						:							List[Any]       =    image_size
             lowerCamelCase__						:							Union[str, Any]       =    num_channels
             lowerCamelCase__						:							Any       =    embeddings_size
             lowerCamelCase__						:							Tuple       =    hidden_sizes
             lowerCamelCase__						:							Tuple       =    depths
             lowerCamelCase__						:							List[str]       =    is_training
             lowerCamelCase__						:							Optional[Any]       =    use_labels
             lowerCamelCase__						:							str       =    hidden_act
             lowerCamelCase__						:							Any       =    num_labels
             lowerCamelCase__						:							List[str]       =    scope
             lowerCamelCase__						:							Any       =    len(__lowerCamelCase			)





       def 		lowerCAmelCase	(  self				:     List[str]			):


             '''simple docstring'''



             lowerCamelCase__						:							Any       =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

             lowerCamelCase__						:							Any       =    self.get_config()

             return config, pixel_values





       def 		lowerCAmelCase	(  self				:     Optional[int]			):


             '''simple docstring'''



             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				:     str   , __lowerCamelCase				:     List[Any]   , __lowerCamelCase				:     Tuple			):


             '''simple docstring'''



             lowerCamelCase__						:							Any       =    FlaxResNetModel(config=__lowerCamelCase			)
             lowerCamelCase__						:							str       =    model(__lowerCamelCase			)

             # Output shape (b, c, h, w)
             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				:     int   , __lowerCamelCase				:     Optional[Any]   , __lowerCamelCase				:     List[str]			):


             '''simple docstring'''



             lowerCamelCase__						:							Optional[Any]       =    self.num_labels
             lowerCamelCase__						:							int       =    FlaxResNetForImageClassification(config=__lowerCamelCase			)
             lowerCamelCase__						:							Optional[int]       =    model(__lowerCamelCase			)
             self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.num_labels)			)





       def 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             lowerCamelCase__						:							Tuple       =    self.prepare_config_and_inputs()
             lowerCamelCase__       ,  lowerCamelCase__						:							Tuple       =    config_and_inputs
             lowerCamelCase__						:							Optional[int]       =    {"pixel_values": pixel_values}
             return config, inputs_dict







@require_flax
class      _lowercase					(				lowercase__				,    unittest.TestCase):


       """simple docstring"""




       A__   		=     (FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

       A__   		=     False
       A__   		=     False
       A__   		=     False





       def 		lowerCAmelCase	(  self				:     List[str]			):


             '''simple docstring'''



             lowerCamelCase__						:							List[str]       =    FlaxResNetModelTester(self			)
             lowerCamelCase__						:							Dict       =    ConfigTester(self   , config_class=__lowerCamelCase   , has_text_modality=__lowerCamelCase			)





       def 		lowerCAmelCase	(  self				:     Optional[Any]			):


             '''simple docstring'''



             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				:     Dict			):


             '''simple docstring'''



             return





       def 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             lowerCamelCase__						:							str       =    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*__lowerCamelCase			)





       def 		lowerCAmelCase	(  self				:     List[Any]			):


             '''simple docstring'''



             lowerCamelCase__						:							str       =    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase			)





       @unittest.skip(reason="ResNet does not use inputs_embeds"			)
       def 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             pass





       @unittest.skip(reason="ResNet does not support input and output embeddings"			)
       def 		lowerCAmelCase	(  self				:     int			):


             '''simple docstring'''



             pass





       def 		lowerCAmelCase	(  self				:     Any			):


             '''simple docstring'''



             lowerCamelCase__       ,  lowerCamelCase__						:							Tuple       =    self.model_tester.prepare_config_and_inputs_for_common()

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

                   lowerCamelCase__						:							List[Any]       =    ["pixel_values"]
                   self.assertListEqual(arg_names[:1]   , __lowerCamelCase			)





       def 		lowerCAmelCase	(  self				:     Any			):


             '''simple docstring'''



             def check_hidden_states_output(__lowerCamelCase				:     List[str]   , __lowerCamelCase				:     Optional[int]   , __lowerCamelCase				:     str			):
                   lowerCamelCase__						:							Any       =    model_class(__lowerCamelCase			)

                   lowerCamelCase__						:							Dict       =    model(**self._prepare_for_class(__lowerCamelCase   , __lowerCamelCase			)			)

                   lowerCamelCase__						:							str       =    outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                   lowerCamelCase__						:							Tuple       =    self.model_tester.num_stages
                   self.assertEqual(len(__lowerCamelCase			)   , expected_num_stages + 1			)





       @unittest.skip(reason="ResNet does not use feedforward chunking"			)
       def 		lowerCAmelCase	(  self				:     Any			):


             '''simple docstring'''



             pass





       def 		lowerCAmelCase	(  self				:     int			):


             '''simple docstring'''



             lowerCamelCase__       ,  lowerCamelCase__						:							Dict       =    self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   with self.subTest(model_class.__name__			):
                         lowerCamelCase__						:							Any       =    self._prepare_for_class(__lowerCamelCase   , __lowerCamelCase			)
                         lowerCamelCase__						:							Optional[int]       =    model_class(__lowerCamelCase			)

                         @jax.jit
                         def model_jitted(__lowerCamelCase				:     Optional[Any]   , **__lowerCamelCase				:     List[Any]			):
                               return model(pixel_values=__lowerCamelCase   , **__lowerCamelCase			)

                         with self.subTest("JIT Enabled"			):
                               lowerCamelCase__						:							Any       =    model_jitted(**__lowerCamelCase			).to_tuple()

                         with self.subTest("JIT Disabled"			):
                               with jax.disable_jit():
                                     lowerCamelCase__						:							Union[str, Any]       =    model_jitted(**__lowerCamelCase			).to_tuple()

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



def  lowercase_						(						):

      """simple docstring"""




      lowerCamelCase__						:							Optional[Any]       =    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"		)
      return image







@require_flax
class      _lowercase					(				unittest.TestCase):


       """simple docstring"""




       @cached_property
       def 		lowerCAmelCase	(  self				:     Tuple			):


             '''simple docstring'''



             return AutoImageProcessor.from_pretrained("microsoft/resnet-50"			) if is_vision_available() else None





       @slow
       def 		lowerCAmelCase	(  self				:     List[str]			):


             '''simple docstring'''



             lowerCamelCase__						:							List[Any]       =    FlaxResNetForImageClassification.from_pretrained("microsoft/resnet-50"			)

             lowerCamelCase__						:							List[str]       =    self.default_image_processor
             lowerCamelCase__						:							Optional[Any]       =    prepare_img()
             lowerCamelCase__						:							str       =    image_processor(images=__lowerCamelCase   , return_tensors="np"			)

             lowerCamelCase__						:							str       =    model(**__lowerCamelCase			)

             # verify the logits
             lowerCamelCase__						:							int       =    (1, 1000)
             self.assertEqual(outputs.logits.shape   , __lowerCamelCase			)

             lowerCamelCase__						:							Optional[int]       =    jnp.array([-1_1.1_0_6_9, -9.7_8_7_7, -8.3_7_7_7]			)

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


