






"""simple docstring"""



import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoImageProcessor, ViTImageProcessor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / """utils"""))

from test_module.custom_image_processing import CustomImageProcessor  # noqa E402


_lowerCAmelCase	:	int				=      get_tests_dir("""fixtures""")




class 			lowerCAmelCase__			(    unittest.TestCase   ):
   def       __a							(		self      :	Union[str, Any]							):


        '''simple docstring'''
        # A mock response for an HTTP head request to emulate server down
        UpperCAmelCase__						:    Dict							=							mock.Mock()
        UpperCAmelCase__						:    Union[str, Any]							=							5_0_0
        UpperCAmelCase__						:    Dict							=							{}
        UpperCAmelCase__						:    Tuple							=							HTTPError
        UpperCAmelCase__						:    Dict							=							{}

        # Download this model to make sure it's in the cache.
        UpperCAmelCase__						:    str							=							ViTImageProcessor.from_pretrained("hf-internal-testing/tiny-random-vit"							)
        # Under the mock environment we get a 500 error when trying to reach the model.
        with mock.patch("requests.Session.request"							,			return_value=snake_case__							) as mock_head:
             UpperCAmelCase__						:    List[str]							=							ViTImageProcessor.from_pretrained("hf-internal-testing/tiny-random-vit"							)
             # This check we did call the fake head request
             mock_head.assert_called()
   def       __a							(		self      :	Dict							):


        '''simple docstring'''
        # This test is for deprecated behavior and can be removed in v5
        UpperCAmelCase__						:    Union[str, Any]							=							ViTImageProcessor.from_pretrained(
            "https://huggingface.co/hf-internal-testing/tiny-random-vit/resolve/main/preprocessor_config.json"							)




   def       __a							(		self      :	str							):


        '''simple docstring'''
        with self.assertRaises(snake_case__							):
             # config is in subfolder, the following should not work without specifying the subfolder
             UpperCAmelCase__						:    List[Any]							=							AutoImageProcessor.from_pretrained("hf-internal-testing/stable-diffusion-all-variants"							)

        UpperCAmelCase__						:    Optional[int]							=							AutoImageProcessor.from_pretrained(
            "hf-internal-testing/stable-diffusion-all-variants"							,			subfolder="feature_extractor"							)

        self.assertIsNotNone(snake_case__							)




@is_staging_test
class 			lowerCAmelCase__			(    unittest.TestCase   ):
   @classmethod
   def       __a							(		cls      :	Dict							):


        '''simple docstring'''
        UpperCAmelCase__						:    int							=							TOKEN
        HfFolder.save_token(snake_case__							)
   @classmethod
   def       __a							(		cls      :	int							):


        '''simple docstring'''
        try:
             delete_repo(token=cls._token							,			repo_id="test-image-processor"							)
        except HTTPError:
             pass

        try:
             delete_repo(token=cls._token							,			repo_id="valid_org/test-image-processor-org"							)
        except HTTPError:
             pass

        try:
             delete_repo(token=cls._token							,			repo_id="test-dynamic-image-processor"							)
        except HTTPError:
             pass
   def       __a							(		self      :	int							):


        '''simple docstring'''
        UpperCAmelCase__						:    List[str]							=							ViTImageProcessor.from_pretrained(snake_case__							)
        image_processor.push_to_hub("test-image-processor"							,			use_auth_token=self._token							)

        UpperCAmelCase__						:    Optional[int]							=							ViTImageProcessor.from_pretrained(f'{USER}/test-image-processor'							)
        for k, v in image_processor.__dict__.items():
             self.assertEqual(snake_case__							,			getattr(snake_case__							,			snake_case__							)							)

        # Reset repo
        delete_repo(token=self._token							,			repo_id="test-image-processor"							)

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
             image_processor.save_pretrained(
                 snake_case__							,			repo_id="test-image-processor"							,			push_to_hub=snake_case__							,			use_auth_token=self._token							)

        UpperCAmelCase__						:    Dict							=							ViTImageProcessor.from_pretrained(f'{USER}/test-image-processor'							)
        for k, v in image_processor.__dict__.items():
             self.assertEqual(snake_case__							,			getattr(snake_case__							,			snake_case__							)							)
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        UpperCAmelCase__						:    str							=							ViTImageProcessor.from_pretrained(snake_case__							)
        image_processor.push_to_hub("valid_org/test-image-processor"							,			use_auth_token=self._token							)

        UpperCAmelCase__						:    Tuple							=							ViTImageProcessor.from_pretrained("valid_org/test-image-processor"							)
        for k, v in image_processor.__dict__.items():
             self.assertEqual(snake_case__							,			getattr(snake_case__							,			snake_case__							)							)

        # Reset repo
        delete_repo(token=self._token							,			repo_id="valid_org/test-image-processor"							)

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
             image_processor.save_pretrained(
                 snake_case__							,			repo_id="valid_org/test-image-processor-org"							,			push_to_hub=snake_case__							,			use_auth_token=self._token							)

        UpperCAmelCase__						:    Any							=							ViTImageProcessor.from_pretrained("valid_org/test-image-processor-org"							)
        for k, v in image_processor.__dict__.items():
             self.assertEqual(snake_case__							,			getattr(snake_case__							,			snake_case__							)							)




   def       __a							(		self      :	List[str]							):


        '''simple docstring'''
        CustomImageProcessor.register_for_auto_class()
        UpperCAmelCase__						:    List[str]							=							CustomImageProcessor.from_pretrained(snake_case__							)

        image_processor.push_to_hub("test-dynamic-image-processor"							,			use_auth_token=self._token							)

        # This has added the proper auto_map field to the config
        self.assertDictEqual(
            image_processor.auto_map							,			{"AutoImageProcessor": "custom_image_processing.CustomImageProcessor"}							,			)

        UpperCAmelCase__						:    Union[str, Any]							=							AutoImageProcessor.from_pretrained(
            f'{USER}/test-dynamic-image-processor'							,			trust_remote_code=snake_case__							)
        # Can't make an isinstance check because the new_image_processor is from the CustomImageProcessor class of a dynamic module
        self.assertEqual(new_image_processor.__class__.__name__							,			"CustomImageProcessor"							)










"""simple docstring"""



import gc
import unittest

from diffusers import FlaxStableDiffusionInpaintPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
 import jax
 import jax.numpy as jnp
 from flax.jax_utils import replicate
 from flax.training.common_utils import shard




@slow
@require_flax
class 			lowerCAmelCase__			(    unittest.TestCase   ):
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()




   def       __a							(		self      :	List[str]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Optional[int]							=							load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/sd2-inpaint/init_image.png"							)
        UpperCAmelCase__						:    Any							=							load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-inpaint/mask.png"							)

        UpperCAmelCase__						:    str							=							"xvjiarui/stable-diffusion-2-inpainting"
        UpperCAmelCase__							,				UpperCAmelCase__						:    List[Any]							=							FlaxStableDiffusionInpaintPipeline.from_pretrained(snake_case__							,			safety_checker=snake_case__							)

        UpperCAmelCase__						:    List[str]							=							"Face of a yellow cat, high resolution, sitting on a park bench"

        UpperCAmelCase__						:    List[str]							=							jax.random.PRNGKey(0							)
        UpperCAmelCase__						:    Optional[int]							=							5_0

        UpperCAmelCase__						:    List[Any]							=							jax.device_count()
        UpperCAmelCase__						:    int							=							num_samples * [prompt]
        UpperCAmelCase__						:    str							=							num_samples * [init_image]
        UpperCAmelCase__						:    Optional[Any]							=							num_samples * [mask_image]
        UpperCAmelCase__							,				UpperCAmelCase__							,				UpperCAmelCase__						:    Optional[int]							=							pipeline.prepare_inputs(snake_case__							,			snake_case__							,			snake_case__							)

        # shard inputs and rng
        UpperCAmelCase__						:    Dict							=							replicate(snake_case__							)
        UpperCAmelCase__						:    Dict							=							jax.random.split(snake_case__							,			jax.device_count()							)
        UpperCAmelCase__						:    List[str]							=							shard(snake_case__							)
        UpperCAmelCase__						:    str							=							shard(snake_case__							)
        UpperCAmelCase__						:    Optional[int]							=							shard(snake_case__							)

        UpperCAmelCase__						:    Dict							=							pipeline(
            snake_case__							,			snake_case__							,			snake_case__							,			snake_case__							,			snake_case__							,			snake_case__							,			jit=snake_case__							)

        UpperCAmelCase__						:    List[str]							=							output.images.reshape(snake_case__							,			5_1_2							,			5_1_2							,			3							)

        UpperCAmelCase__						:    Tuple							=							images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

        UpperCAmelCase__						:    str							=							jnp.asarray(jax.device_get(image_slice.flatten()							)							)
        UpperCAmelCase__						:    int							=							jnp.array(
            [0.361_1307, 0.3764_9736, 0.375_7408, 0.3821_3953, 0.3929_5167, 0.384_1631, 0.4155_4978, 0.413_7475, 0.421_7084]							)
        print(f'output_slice: {output_slice}'							)
        assert jnp.abs(output_slice - expected_slice							).max() < 1e-2


