





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				,			__a				,			__a=3				,			__a=32				,			__a=3				,			__a=10				,			__a=[10, 20, 30, 40]				,			__a=[1, 1, 2, 1]				,			__a=True				,			__a=True				,			__a="relu"				,			__a=3				,			__a=None				,			):






           '''simple docstring'''



           lowerCamelCase								=		parent
           lowerCamelCase								=		batch_size
           lowerCamelCase								=		image_size
           lowerCamelCase								=		num_channels
           lowerCamelCase								=		embeddings_size
           lowerCamelCase								=		hidden_sizes
           lowerCamelCase								=		depths
           lowerCamelCase								=		is_training
           lowerCamelCase								=		use_labels
           lowerCamelCase								=		hidden_act
           lowerCamelCase								=		num_labels
           lowerCamelCase								=		scope
           lowerCamelCase								=		len(SCREAMING_SNAKE_CASE_    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

           lowerCamelCase								=		self.get_config()

           return config, pixel_values







     def 						_a      (self    ):






           '''simple docstring'''



           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 						_a      (self				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		FlaxRegNetModel(config=SCREAMING_SNAKE_CASE_    )
           lowerCamelCase								=		model(SCREAMING_SNAKE_CASE_    )

           # 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 						_a      (self				,			__a				,			__a    ):






           '''simple docstring'''



           lowerCamelCase								=		self.num_labels
           lowerCamelCase								=		FlaxRegNetForImageClassification(config=SCREAMING_SNAKE_CASE_    )
           lowerCamelCase								=		model(SCREAMING_SNAKE_CASE_    )
           self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.num_labels)    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.prepare_config_and_inputs()
           lowerCamelCase			,							lowerCamelCase								=		config_and_inputs
           lowerCamelCase								=		{"pixel_values": pixel_values}
           return config, inputs_dict





@require_flax
class    lowerCamelCase__  (						__lowerCamelCase       ,    unittest.TestCase):
     """simple docstring"""



     _A       =			(FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

     _A       =			False
     _A       =			False
     _A       =			False







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		FlaxRegNetModelTester(self    )
           lowerCamelCase								=		ConfigTester(self				,			config_class=SCREAMING_SNAKE_CASE_				,			has_text_modality=SCREAMING_SNAKE_CASE_    )







     def 						_a      (self    ):






           '''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 						_a      (self    ):






           '''simple docstring'''



           return







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE_    )







     @unittest.skip(reason="RegNet does not use inputs_embeds"    )
     def 						_a      (self    ):






           '''simple docstring'''



           pass







     @unittest.skip(reason="RegNet does not support input and output embeddings"    )
     def 						_a      (self    ):






           '''simple docstring'''



           pass







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase			,							lowerCamelCase								=		self.model_tester.prepare_config_and_inputs_for_common()

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

                 lowerCamelCase								=		["pixel_values"]
                 self.assertListEqual(arg_names[:1]				,			SCREAMING_SNAKE_CASE_    )







     def 						_a      (self    ):






           '''simple docstring'''



           def check_hidden_states_output(__a				,			__a				,			__a    ):
                 lowerCamelCase								=		model_class(SCREAMING_SNAKE_CASE_    )

                 lowerCamelCase								=		model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_				,			SCREAMING_SNAKE_CASE_    )    )

                 lowerCamelCase								=		outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                 lowerCamelCase								=		self.model_tester.num_stages
                 self.assertEqual(len(SCREAMING_SNAKE_CASE_    )				,			expected_num_stages + 1    )

           lowerCamelCase			,							lowerCamelCase								=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 lowerCamelCase								=		True
                 check_hidden_states_output(SCREAMING_SNAKE_CASE_				,			SCREAMING_SNAKE_CASE_				,			SCREAMING_SNAKE_CASE_    )

                 # check that output_hidden_states also work using config
                 del inputs_dict["output_hidden_states"]
                 lowerCamelCase								=		True

                 check_hidden_states_output(SCREAMING_SNAKE_CASE_				,			SCREAMING_SNAKE_CASE_				,			SCREAMING_SNAKE_CASE_    )







     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase			,							lowerCamelCase								=		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                 with self.subTest(model_class.__name__    ):
                       lowerCamelCase								=		self._prepare_for_class(SCREAMING_SNAKE_CASE_				,			SCREAMING_SNAKE_CASE_    )
                       lowerCamelCase								=		model_class(SCREAMING_SNAKE_CASE_    )

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

                       with self.subTest("JIT Enabled"    ):
                             lowerCamelCase								=		model_jitted(**SCREAMING_SNAKE_CASE_    ).to_tuple()

                       with self.subTest("JIT Disabled"    ):
                             with jax.disable_jit():
                                   lowerCamelCase								=		model_jitted(**SCREAMING_SNAKE_CASE_    ).to_tuple()

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



def     __lowercase(	):


      """simple docstring"""



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






@require_flax
class    lowerCamelCase__  (						unittest.TestCase):
     """simple docstring"""



     @cached_property
     def 						_a      (self    ):






           '''simple docstring'''



           return AutoImageProcessor.from_pretrained("facebook/regnet-y-040"    ) if is_vision_available() else None







     @slow
     def 						_a      (self    ):






           '''simple docstring'''



           lowerCamelCase								=		FlaxRegNetForImageClassification.from_pretrained("facebook/regnet-y-040"    )

           lowerCamelCase								=		self.default_image_processor
           lowerCamelCase								=		prepare_img()
           lowerCamelCase								=		image_processor(images=SCREAMING_SNAKE_CASE_				,			return_tensors="np"    )

           lowerCamelCase								=		model(**SCREAMING_SNAKE_CASE_    )

           # verify the logits
           lowerCamelCase								=		(1, 10_00)
           self.assertEqual(outputs.logits.shape				,			SCREAMING_SNAKE_CASE_    )

           lowerCamelCase								=		jnp.array([-0.4180, -1.5051, -3.4836]    )

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




import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel

from ...utils import logging


a_						:     Tuple	     =		logging.get_logger(__name__)





class    lowerCamelCase__  (						UpperCAmelCase_):
     """simple docstring"""



     _A       =			CLIPConfig

     _A       =			['CLIPEncoderLayer']







     def __init__(self				,			__a    ):






           '''simple docstring'''



           super().__init__(__a    )

           lowerCamelCase								=		CLIPVisionModelWithProjection(config.vision_config    )

           lowerCamelCase								=		nn.Linear(config.vision_config.projection_dim				,			1    )
           lowerCamelCase								=		nn.Linear(config.vision_config.projection_dim				,			1    )







     @torch.no_grad()
     def 						_a      (self				,			__a				,			__a				,			__a=0.5				,			__a=0.5    ):






           '''simple docstring'''



           lowerCamelCase								=		self.vision_model(__a    )[0]

           lowerCamelCase								=		self.p_head(__a    )
           lowerCamelCase								=		nsfw_detected.flatten()
           lowerCamelCase								=		nsfw_detected > p_threshold
           lowerCamelCase								=		nsfw_detected.tolist()

           if any(__a    ):
                 logger.warning(
                     "Potential NSFW content was detected in one or more images. A black image will be returned instead."
                     " Try again with a different prompt and/or seed."    )

           for idx, nsfw_detected_ in enumerate(__a    ):
                 if nsfw_detected_:
                       lowerCamelCase								=		np.zeros(images[idx].shape    )

           lowerCamelCase								=		self.w_head(__a    )
           lowerCamelCase								=		watermark_detected.flatten()
           lowerCamelCase								=		watermark_detected > w_threshold
           lowerCamelCase								=		watermark_detected.tolist()

           if any(__a    ):
                 logger.warning(
                     "Potential watermarked content was detected in one or more images. A black image will be returned instead."
                     " Try again with a different prompt and/or seed."    )

           for idx, watermark_detected_ in enumerate(__a    ):
                 if watermark_detected_:
                       lowerCamelCase								=		np.zeros(images[idx].shape    )

           return images, nsfw_detected, watermark_detected