





'''simple docstring'''






import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    CLIPConfig,
    CLIPImageProcessor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER


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

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_image_processing import CustomImageProcessor  # noqa E402




class 			UpperCAmelCase					(      unittest.TestCase       ):




      def     lowercase__   (      self      :			str  )    ->				Optional[int]:
            _lowerCAmelCase												=							0




      def     lowercase__   (      self      :			int  )    ->				str:
            _lowerCAmelCase												=							AutoImageProcessor.from_pretrained("""openai/clip-vit-base-patch32"""  )
            self.assertIsInstance(__snake_case      ,      __snake_case  )




      def     lowercase__   (      self      :			int  )    ->				Any:
            with tempfile.TemporaryDirectory() as tmpdirname:
                  _lowerCAmelCase												=							Path(__snake_case  ) / """preprocessor_config.json"""
                  _lowerCAmelCase												=							Path(__snake_case  ) / """config.json"""
                  json.dump(
                      {"""image_processor_type""": """CLIPImageProcessor""", """processor_class""": """CLIPProcessor"""}      ,      open(__snake_case      ,      """w"""  )      ,      )
                  json.dump({"""model_type""": """clip"""}      ,      open(__snake_case      ,      """w"""  )  )

                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(__snake_case  )
                  self.assertIsInstance(__snake_case      ,      __snake_case  )




      def     lowercase__   (      self      :			Tuple  )    ->				List[Any]:
            # Ensure we can load the image processor from the feature extractor config
            with tempfile.TemporaryDirectory() as tmpdirname:
                  _lowerCAmelCase												=							Path(__snake_case  ) / """preprocessor_config.json"""
                  _lowerCAmelCase												=							Path(__snake_case  ) / """config.json"""
                  json.dump(
                      {"""feature_extractor_type""": """CLIPFeatureExtractor""", """processor_class""": """CLIPProcessor"""}      ,      open(__snake_case      ,      """w"""  )      ,      )
                  json.dump({"""model_type""": """clip"""}      ,      open(__snake_case      ,      """w"""  )  )

                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(__snake_case  )
                  self.assertIsInstance(__snake_case      ,      __snake_case  )




      def     lowercase__   (      self      :			List[Any]  )    ->				int:
            with tempfile.TemporaryDirectory() as tmpdirname:
                  _lowerCAmelCase												=							CLIPConfig()

                  # Create a dummy config file with image_proceesor_type
                  _lowerCAmelCase												=							Path(__snake_case  ) / """preprocessor_config.json"""
                  _lowerCAmelCase												=							Path(__snake_case  ) / """config.json"""
                  json.dump(
                      {"""image_processor_type""": """CLIPImageProcessor""", """processor_class""": """CLIPProcessor"""}      ,      open(__snake_case      ,      """w"""  )      ,      )
                  json.dump({"""model_type""": """clip"""}      ,      open(__snake_case      ,      """w"""  )  )

                  # remove image_processor_type to make sure config.json alone is enough to load image processor locally
                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(__snake_case  ).to_dict()

                  config_dict.pop("""image_processor_type"""  )
                  _lowerCAmelCase												=							CLIPImageProcessor(**__snake_case  )

                  # save in new folder
                  model_config.save_pretrained(__snake_case  )
                  config.save_pretrained(__snake_case  )

                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(__snake_case  )

                  # make sure private variable is not incorrectly saved
                  _lowerCAmelCase												=							json.loads(config.to_json_string()  )
                  self.assertTrue("""_processor_class""" not in dict_as_saved  )

            self.assertIsInstance(__snake_case      ,      __snake_case  )




      def     lowercase__   (      self      :			List[str]  )    ->				Union[str, Any]:
            with tempfile.TemporaryDirectory() as tmpdirname:
                  _lowerCAmelCase												=							Path(__snake_case  ) / """preprocessor_config.json"""
                  json.dump(
                      {"""image_processor_type""": """CLIPImageProcessor""", """processor_class""": """CLIPProcessor"""}      ,      open(__snake_case      ,      """w"""  )      ,      )

                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(__snake_case  )
                  self.assertIsInstance(__snake_case      ,      __snake_case  )




      def     lowercase__   (      self      :			int  )    ->				Optional[int]:
            with self.assertRaisesRegex(
                __snake_case      ,      """clip-base is not a local folder and is not a valid model identifier"""  ):
                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained("""clip-base"""  )




      def     lowercase__   (      self      :			Any  )    ->				Tuple:
            with self.assertRaisesRegex(
                __snake_case      ,      R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""  ):
                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(__snake_case      ,      revision="""aaaaaa"""  )




      def     lowercase__   (      self      :			Optional[Any]  )    ->				str:
            with self.assertRaisesRegex(
                __snake_case      ,      """hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json."""      ,      ):
                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained("""hf-internal-testing/config-no-model"""  )




      def     lowercase__   (      self      :			Optional[Any]  )    ->				Any:
            # If remote code is not set, we will time out when asking whether to load the model.
            with self.assertRaises(__snake_case  ):
                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained("""hf-internal-testing/test_dynamic_image_processor"""  )
            # If remote code is disabled, we can't load this config.
            with self.assertRaises(__snake_case  ):
                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(
                      """hf-internal-testing/test_dynamic_image_processor"""      ,      trust_remote_code=__snake_case  )

            _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(
                """hf-internal-testing/test_dynamic_image_processor"""      ,      trust_remote_code=__snake_case  )
            self.assertEqual(image_processor.__class__.__name__      ,      """NewImageProcessor"""  )

            # Test image processor can be reloaded.
            with tempfile.TemporaryDirectory() as tmp_dir:
                  image_processor.save_pretrained(__snake_case  )
                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(__snake_case      ,      trust_remote_code=__snake_case  )
            self.assertEqual(reloaded_image_processor.__class__.__name__      ,      """NewImageProcessor"""  )




      def     lowercase__   (      self      :			int  )    ->				Dict:
            try:
                  AutoConfig.register("""custom"""      ,      __snake_case  )
                  AutoImageProcessor.register(__snake_case      ,      __snake_case  )
                  # Trying to register something existing in the Transformers library will raise an error
                  with self.assertRaises(__snake_case  ):
                        AutoImageProcessor.register(__snake_case      ,      __snake_case  )

                  with tempfile.TemporaryDirectory() as tmpdirname:
                        _lowerCAmelCase												=							Path(__snake_case  ) / """preprocessor_config.json"""
                        _lowerCAmelCase												=							Path(__snake_case  ) / """config.json"""
                        json.dump(
                            {"""feature_extractor_type""": """CLIPFeatureExtractor""", """processor_class""": """CLIPProcessor"""}      ,      open(__snake_case      ,      """w"""  )      ,      )
                        json.dump({"""model_type""": """clip"""}      ,      open(__snake_case      ,      """w"""  )  )

                        _lowerCAmelCase												=							CustomImageProcessor.from_pretrained(__snake_case  )

                  # Now that the config is registered, it can be used as any other config with the auto-API
                  with tempfile.TemporaryDirectory() as tmp_dir:
                        image_processor.save_pretrained(__snake_case  )
                        _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(__snake_case  )
                        self.assertIsInstance(__snake_case      ,      __snake_case  )

            finally:
                  if "custom" in CONFIG_MAPPING._extra_content:
                        del CONFIG_MAPPING._extra_content["custom"]
                  if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
                        del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]




      def     lowercase__   (      self      :			Dict  )    ->				List[Any]:




            class 			UpperCAmelCase					(      snake_case_       ):
                  _lowercase:      Dict											=       True

            try:
                  AutoConfig.register("""custom"""      ,      __snake_case  )
                  AutoImageProcessor.register(__snake_case      ,      __snake_case  )
                  # If remote code is not set, the default is to use local
                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained("""hf-internal-testing/test_dynamic_image_processor"""  )
                  self.assertEqual(image_processor.__class__.__name__      ,      """NewImageProcessor"""  )
                  self.assertTrue(image_processor.is_local  )

                  # If remote code is disabled, we load the local one.
                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(
                      """hf-internal-testing/test_dynamic_image_processor"""      ,      trust_remote_code=__snake_case  )
                  self.assertEqual(image_processor.__class__.__name__      ,      """NewImageProcessor"""  )
                  self.assertTrue(image_processor.is_local  )

                  # If remote is enabled, we load from the Hub
                  _lowerCAmelCase												=							AutoImageProcessor.from_pretrained(
                      """hf-internal-testing/test_dynamic_image_processor"""      ,      trust_remote_code=__snake_case  )
                  self.assertEqual(image_processor.__class__.__name__      ,      """NewImageProcessor"""  )
                  self.assertTrue(not hasattr(__snake_case      ,      """is_local"""  )  )

            finally:
                  if "custom" in CONFIG_MAPPING._extra_content:
                        del CONFIG_MAPPING._extra_content["custom"]
                  if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
                        del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]













'''simple docstring'''






import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class 			UpperCAmelCase					(      unittest.TestCase       ):




      def     lowercase__   (      self      :			Optional[int]  )    ->				Dict:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      @property
      def     lowercase__   (      self      :			List[Any]  )    ->				List[str]:
            _lowerCAmelCase												=							1
            _lowerCAmelCase												=							3
            _lowerCAmelCase												=							(32, 32)

            _lowerCAmelCase												=							floats_tensor((batch_size, num_channels) + sizes      ,      rng=random.Random(0  )  ).to(__snake_case  )
            return image




      @property
      def     lowercase__   (      self      :			int  )    ->				Union[str, Any]:
            torch.manual_seed(0  )
            _lowerCAmelCase												=							UNetaDConditionModel(
                block_out_channels=(32, 64)      ,      layers_per_block=2      ,      sample_size=32      ,      in_channels=4      ,      out_channels=4      ,      down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")      ,      up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")      ,      cross_attention_dim=32      ,      )
            return model




      @property
      def     lowercase__   (      self      :			Optional[int]  )    ->				List[str]:
            torch.manual_seed(0  )
            _lowerCAmelCase												=							AutoencoderKL(
                block_out_channels=[32, 64]      ,      in_channels=3      ,      out_channels=3      ,      down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]      ,      up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]      ,      latent_channels=4      ,      )
            return model




      @property
      def     lowercase__   (      self      :			Dict  )    ->				Optional[Any]:
            torch.manual_seed(0  )
            _lowerCAmelCase												=							RobertaSeriesConfig(
                hidden_size=32      ,      project_dim=32      ,      intermediate_size=37      ,      layer_norm_eps=1E-0_5      ,      num_attention_heads=4      ,      num_hidden_layers=5      ,      pad_token_id=1      ,      vocab_size=50_06      ,      )
            return RobertaSeriesModelWithTransformation(__snake_case  )




      @property
      def     lowercase__   (      self      :			Union[str, Any]  )    ->				str:
            def extract(*__snake_case      :			List[Any]      ,      **__snake_case      :			Any  ):




                  class 			UpperCAmelCase					:




                        def __init__(      self      :			Any  )    ->				Any:
                              _lowerCAmelCase												=							torch.ones([0]  )




                        def     lowercase__   (      self      :			Optional[Any]      ,      __snake_case      :			Tuple  )    ->				Dict:
                              self.pixel_values.to(__snake_case  )
                              return self




                  return Out()

            return extract




      def     lowercase__   (      self      :			List[str]  )    ->				Optional[int]:
            _lowerCAmelCase												=							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            _lowerCAmelCase												=							self.dummy_cond_unet
            _lowerCAmelCase												=							PNDMScheduler(skip_prk_steps=__snake_case  )
            _lowerCAmelCase												=							self.dummy_vae
            _lowerCAmelCase												=							self.dummy_text_encoder
            _lowerCAmelCase												=							XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""  )
            _lowerCAmelCase												=							77

            _lowerCAmelCase												=							self.dummy_image.to(__snake_case  )
            _lowerCAmelCase												=							init_image / 2 + 0.5

            # make sure here that pndm scheduler skips prk
            _lowerCAmelCase												=							AltDiffusionImgaImgPipeline(
                unet=__snake_case      ,      scheduler=__snake_case      ,      vae=__snake_case      ,      text_encoder=__snake_case      ,      tokenizer=__snake_case      ,      safety_checker=__snake_case      ,      feature_extractor=self.dummy_extractor      ,      )
            _lowerCAmelCase												=							VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor      ,      do_normalize=__snake_case  )
            _lowerCAmelCase												=							alt_pipe.to(__snake_case  )
            alt_pipe.set_progress_bar_config(disable=__snake_case  )

            _lowerCAmelCase												=							"""A painting of a squirrel eating a burger"""
            _lowerCAmelCase												=							torch.Generator(device=__snake_case  ).manual_seed(0  )
            _lowerCAmelCase												=							alt_pipe(
                [prompt]      ,      generator=__snake_case      ,      guidance_scale=6.0      ,      num_inference_steps=2      ,      output_type="""np"""      ,      image=__snake_case      ,      )

            _lowerCAmelCase												=							output.images

            _lowerCAmelCase												=							torch.Generator(device=__snake_case  ).manual_seed(0  )
            _lowerCAmelCase												=							alt_pipe(
                [prompt]      ,      generator=__snake_case      ,      guidance_scale=6.0      ,      num_inference_steps=2      ,      output_type="""np"""      ,      image=__snake_case      ,      return_dict=__snake_case      ,      )[0]

            _lowerCAmelCase												=							image[0, -3:, -3:, -1]
            _lowerCAmelCase												=							image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 32, 32, 3)
            _lowerCAmelCase												=							np.array([0.44_27, 0.37_31, 0.42_49, 0.49_41, 0.45_46, 0.41_48, 0.41_93, 0.46_66, 0.44_99]  )

            assert np.abs(image_slice.flatten() - expected_slice  ).max() < 5E-3
            assert np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 5E-3




      @unittest.skipIf(torch_device != """cuda"""      ,      """This test requires a GPU"""  )
      def     lowercase__   (      self      :			Tuple  )    ->				str:
            _lowerCAmelCase												=							self.dummy_cond_unet
            _lowerCAmelCase												=							PNDMScheduler(skip_prk_steps=__snake_case  )
            _lowerCAmelCase												=							self.dummy_vae
            _lowerCAmelCase												=							self.dummy_text_encoder
            _lowerCAmelCase												=							XLMRobertaTokenizer.from_pretrained("""hf-internal-testing/tiny-xlm-roberta"""  )
            _lowerCAmelCase												=							77

            _lowerCAmelCase												=							self.dummy_image.to(__snake_case  )

            # put models in fp16
            _lowerCAmelCase												=							unet.half()
            _lowerCAmelCase												=							vae.half()
            _lowerCAmelCase												=							bert.half()

            # make sure here that pndm scheduler skips prk
            _lowerCAmelCase												=							AltDiffusionImgaImgPipeline(
                unet=__snake_case      ,      scheduler=__snake_case      ,      vae=__snake_case      ,      text_encoder=__snake_case      ,      tokenizer=__snake_case      ,      safety_checker=__snake_case      ,      feature_extractor=self.dummy_extractor      ,      )
            _lowerCAmelCase												=							VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor      ,      do_normalize=__snake_case  )
            _lowerCAmelCase												=							alt_pipe.to(__snake_case  )
            alt_pipe.set_progress_bar_config(disable=__snake_case  )

            _lowerCAmelCase												=							"""A painting of a squirrel eating a burger"""
            _lowerCAmelCase												=							torch.manual_seed(0  )
            _lowerCAmelCase												=							alt_pipe(
                [prompt]      ,      generator=__snake_case      ,      num_inference_steps=2      ,      output_type="""np"""      ,      image=__snake_case      ,      ).images

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




      @unittest.skipIf(torch_device != """cuda"""      ,      """This test requires a GPU"""  )
      def     lowercase__   (      self      :			Optional[Any]  )    ->				Optional[Any]:
            _lowerCAmelCase												=							load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                """/img2img/sketch-mountains-input.jpg"""  )
            # resize to resolution that is divisible by 8 but not 16 or 32
            _lowerCAmelCase												=							init_image.resize((7_60, 5_04)  )

            _lowerCAmelCase												=							"""BAAI/AltDiffusion"""
            _lowerCAmelCase												=							AltDiffusionImgaImgPipeline.from_pretrained(
                __snake_case      ,      safety_checker=__snake_case      ,      )
            pipe.to(__snake_case  )
            pipe.set_progress_bar_config(disable=__snake_case  )
            pipe.enable_attention_slicing()

            _lowerCAmelCase												=							"""A fantasy landscape, trending on artstation"""

            _lowerCAmelCase												=							torch.manual_seed(0  )
            _lowerCAmelCase												=							pipe(
                prompt=__snake_case      ,      image=__snake_case      ,      strength=0.75      ,      guidance_scale=7.5      ,      generator=__snake_case      ,      output_type="""np"""      ,      )
            _lowerCAmelCase												=							output.images[0]

            _lowerCAmelCase												=							image[2_55:2_58, 3_83:3_86, -1]

            assert image.shape == (5_04, 7_60, 3)
            _lowerCAmelCase												=							np.array([0.93_58, 0.93_97, 0.95_99, 0.99_01, 1.00_00, 1.00_00, 0.98_82, 1.00_00, 1.00_00]  )

            assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2




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




      def     lowercase__   (      self      :			int  )    ->				Any:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      def     lowercase__   (      self      :			Dict  )    ->				Tuple:
            _lowerCAmelCase												=							load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                """/img2img/sketch-mountains-input.jpg"""  )
            _lowerCAmelCase												=							init_image.resize((7_68, 5_12)  )
            _lowerCAmelCase												=							load_numpy(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy"""  )

            _lowerCAmelCase												=							"""BAAI/AltDiffusion"""
            _lowerCAmelCase												=							AltDiffusionImgaImgPipeline.from_pretrained(
                __snake_case      ,      safety_checker=__snake_case      ,      )
            pipe.to(__snake_case  )
            pipe.set_progress_bar_config(disable=__snake_case  )
            pipe.enable_attention_slicing()

            _lowerCAmelCase												=							"""A fantasy landscape, trending on artstation"""

            _lowerCAmelCase												=							torch.manual_seed(0  )
            _lowerCAmelCase												=							pipe(
                prompt=__snake_case      ,      image=__snake_case      ,      strength=0.75      ,      guidance_scale=7.5      ,      generator=__snake_case      ,      output_type="""np"""      ,      )
            _lowerCAmelCase												=							output.images[0]

            assert image.shape == (5_12, 7_68, 3)
            # img2img is flaky across GPUs even in fp32, so using MAE here
            assert np.abs(expected_image - image  ).max() < 1E-2






