






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 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

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





       @property
       def 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

          _lowercase              =1
          _lowercase              =3
          _lowercase              =(32, 32)

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





       @property
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          torch.manual_seed(0)
          _lowercase              =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 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          torch.manual_seed(0)
          _lowercase              =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 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          torch.manual_seed(0)
          _lowercase              =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=5006,      )
          return RobertaSeriesModelWithTransformation(snake_case)





       @property
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          def extract(*snake_case	:Tuple,      **snake_case	:List[Any]):






             class 					SCREAMING_SNAKE_CASE_  :
                    """simple docstring"""


                    def __init__(				self	:List[str]):



                       """simple docstring"""

                       _lowercase              =torch.ones([0])





                    def 				UpperCamelCase__						(				self	:List[Any],      snake_case	:Optional[Any]):



                       """simple docstring"""

                       self.pixel_values.to(snake_case)
                       return self

             return Out()

          return extract





       def 				UpperCamelCase__						(				self	:List[str]):



          """simple docstring"""

          _lowercase              ='cpu'  # ensure determinism for the device-dependent torch.Generator
          _lowercase              =self.dummy_cond_unet
          _lowercase              =PNDMScheduler(skip_prk_steps=snake_case)
          _lowercase              =self.dummy_vae
          _lowercase              =self.dummy_text_encoder
          _lowercase              =XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta')
          _lowercase              =77

          _lowercase              =self.dummy_image.to(snake_case)
          _lowercase              =init_image / 2 + 0.5

          # make sure here that pndm scheduler skips prk
          _lowercase              =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,      )
          _lowercase              =VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor,      do_normalize=snake_case)
          _lowercase              =alt_pipe.to(snake_case)
          alt_pipe.set_progress_bar_config(disable=snake_case)

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

          _lowercase              =output.images

          _lowercase              =torch.Generator(device=snake_case).manual_seed(0)
          _lowercase              =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]

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

          assert image.shape == (1, 32, 32, 3)
          _lowercase              =np.array([0.4_4_2_7, 0.3_7_3_1, 0.4_2_4_9, 0.4_9_4_1, 0.4_5_4_6, 0.4_1_4_8, 0.4_1_9_3, 0.4_6_6_6, 0.4_4_9_9])

          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 				UpperCamelCase__						(				self	:List[str]):



          """simple docstring"""

          _lowercase              =self.dummy_cond_unet
          _lowercase              =PNDMScheduler(skip_prk_steps=snake_case)
          _lowercase              =self.dummy_vae
          _lowercase              =self.dummy_text_encoder
          _lowercase              =XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta')
          _lowercase              =77

          _lowercase              =self.dummy_image.to(snake_case)

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

          # make sure here that pndm scheduler skips prk
          _lowercase              =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,      )
          _lowercase              =VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor,      do_normalize=snake_case)
          _lowercase              =alt_pipe.to(snake_case)
          alt_pipe.set_progress_bar_config(disable=snake_case)

          _lowercase              ='A painting of a squirrel eating a burger'
          _lowercase              =torch.manual_seed(0)
          _lowercase              =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 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

          _lowercase              =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
          _lowercase              =init_image.resize((760, 504))

          _lowercase              ='BAAI/AltDiffusion'
          _lowercase              =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()

          _lowercase              ='A fantasy landscape, trending on artstation'

          _lowercase              =torch.manual_seed(0)
          _lowercase              =pipe(
              prompt=snake_case,      image=snake_case,      strength=0.7_5,      guidance_scale=7.5,      generator=snake_case,      output_type='np',      )
          _lowercase              =output.images[0]

          _lowercase              =image[255:258, 383:386, -1]

          assert image.shape == (504, 760, 3)
          _lowercase              =np.array([0.9_3_5_8, 0.9_3_9_7, 0.9_5_9_9, 0.9_9_0_1, 1.0_0_0_0, 1.0_0_0_0, 0.9_8_8_2, 1.0_0_0_0, 1.0_0_0_0])

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






@slow
@require_torch_gpu
class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          _lowercase              =load_image(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/img2img/sketch-mountains-input.jpg')
          _lowercase              =init_image.resize((768, 512))
          _lowercase              =load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy')

          _lowercase              ='BAAI/AltDiffusion'
          _lowercase              =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()

          _lowercase              ='A fantasy landscape, trending on artstation'

          _lowercase              =torch.manual_seed(0)
          _lowercase              =pipe(
              prompt=snake_case,      image=snake_case,      strength=0.7_5,      guidance_scale=7.5,      generator=snake_case,      output_type='np',      )
          _lowercase              =output.images[0]

          assert image.shape == (512, 768, 3)
          # img2img is flaky across GPUs even in fp32, so using MAE here
          assert np.abs(expected_image - image).max() < 1e-2










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 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          _lowercase              =0





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          _lowercase              =AutoImageProcessor.from_pretrained('openai/clip-vit-base-patch32')
          self.assertIsInstance(snake_case,      snake_case)





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



          """simple docstring"""

          with tempfile.TemporaryDirectory() as tmpdirname:
             _lowercase              =Path(snake_case) / 'preprocessor_config.json'
             _lowercase              =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'))

             _lowercase              =AutoImageProcessor.from_pretrained(snake_case)
             self.assertIsInstance(snake_case,      snake_case)





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



          """simple docstring"""

          with tempfile.TemporaryDirectory() as tmpdirname:
             _lowercase              =Path(snake_case) / 'preprocessor_config.json'
             _lowercase              =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'))

             _lowercase              =AutoImageProcessor.from_pretrained(snake_case)
             self.assertIsInstance(snake_case,      snake_case)





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



          """simple docstring"""

          with tempfile.TemporaryDirectory() as tmpdirname:
             _lowercase              =CLIPConfig()

             # Create a dummy config file with image_proceesor_type
             _lowercase              =Path(snake_case) / 'preprocessor_config.json'
             _lowercase              =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
             _lowercase              =AutoImageProcessor.from_pretrained(snake_case).to_dict()

             config_dict.pop('image_processor_type')
             _lowercase              =CLIPImageProcessor(**snake_case)

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

             _lowercase              =AutoImageProcessor.from_pretrained(snake_case)

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

          self.assertIsInstance(snake_case,      snake_case)





       def 				UpperCamelCase__						(				self	:Optional[int]):



          """simple docstring"""

          with tempfile.TemporaryDirectory() as tmpdirname:
             _lowercase              =Path(snake_case) / 'preprocessor_config.json'
             json.dump(
                 {'image_processor_type': 'CLIPImageProcessor', 'processor_class': 'CLIPProcessor'},      open(snake_case,      'w'),      )

             _lowercase              =AutoImageProcessor.from_pretrained(snake_case)
             self.assertIsInstance(snake_case,      snake_case)





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

          with self.assertRaisesRegex(
              snake_case,      'clip-base is not a local folder and is not a valid model identifier'):
             _lowercase              =AutoImageProcessor.from_pretrained('clip-base')





       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          with self.assertRaisesRegex(
              snake_case,      r'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'):
             _lowercase              =AutoImageProcessor.from_pretrained(snake_case,      revision='aaaaaa')





       def 				UpperCamelCase__						(				self	:List[Any]):



          """simple docstring"""

          with self.assertRaisesRegex(
              snake_case,      'hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json.',      ):
             _lowercase              =AutoImageProcessor.from_pretrained('hf-internal-testing/config-no-model')





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          with self.assertRaises(snake_case):
             _lowercase              =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):
             _lowercase              =AutoImageProcessor.from_pretrained(
                 'hf-internal-testing/test_dynamic_image_processor',      trust_remote_code=snake_case)

          _lowercase              =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)
             _lowercase              =AutoImageProcessor.from_pretrained(snake_case,      trust_remote_code=snake_case)
          self.assertEqual(reloaded_image_processor.__class__.__name__,      'NewImageProcessor')





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          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:
                _lowercase              =Path(snake_case) / 'preprocessor_config.json'
                _lowercase              =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'))

                _lowercase              =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)
                _lowercase              =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 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          class 					SCREAMING_SNAKE_CASE_  (       _a ):
                 """simple docstring"""


                 __lowerCAmelCase    :		Union[str, Any]					=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
             _lowercase              =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.
             _lowercase              =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
             _lowercase              =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]


