

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






  """simple docstring"""


  def    lowercase_	(							self		:	int      )	->							Optional[int]:
       SCREAMING_SNAKE_CASE__					=					0






  def    lowercase_	(							self		:	Any      )	->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained('''openai/clip-vit-base-patch32'''      )
       self.assertIsInstance(__lowerCamelCase			,     __lowerCamelCase      )






  def    lowercase_	(							self		:	int      )	->							Optional[Any]:
       with tempfile.TemporaryDirectory() as tmpdirname:
            SCREAMING_SNAKE_CASE__					=					Path(__lowerCamelCase      ) / '''preprocessor_config.json'''
            SCREAMING_SNAKE_CASE__					=					Path(__lowerCamelCase      ) / '''config.json'''
            json.dump(
                {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			,     open(__lowerCamelCase			,     '''w'''      )			,     )
            json.dump({'''model_type''': '''clip'''}			,     open(__lowerCamelCase			,     '''w'''      )      )

            SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained(__lowerCamelCase      )
            self.assertIsInstance(__lowerCamelCase			,     __lowerCamelCase      )






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

            SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained(__lowerCamelCase      )
            self.assertIsInstance(__lowerCamelCase			,     __lowerCamelCase      )






  def    lowercase_	(							self		:	Optional[int]      )	->							Any:
       with tempfile.TemporaryDirectory() as tmpdirname:
            SCREAMING_SNAKE_CASE__					=					CLIPConfig()

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

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

            config_dict.pop('''image_processor_type'''      )
            SCREAMING_SNAKE_CASE__					=					CLIPImageProcessor(**__lowerCamelCase      )

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

            SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained(__lowerCamelCase      )

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

       self.assertIsInstance(__lowerCamelCase			,     __lowerCamelCase      )






  def    lowercase_	(							self		:	str      )	->							Optional[Any]:
       with tempfile.TemporaryDirectory() as tmpdirname:
            SCREAMING_SNAKE_CASE__					=					Path(__lowerCamelCase      ) / '''preprocessor_config.json'''
            json.dump(
                {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			,     open(__lowerCamelCase			,     '''w'''      )			,     )

            SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained(__lowerCamelCase      )
            self.assertIsInstance(__lowerCamelCase			,     __lowerCamelCase      )






  def    lowercase_	(							self		:	int      )	->							Any:
       with self.assertRaisesRegex(
           __lowerCamelCase			,     '''clip-base is not a local folder and is not a valid model identifier'''      ):
            SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained('''clip-base'''      )






  def    lowercase_	(							self		:	Dict      )	->							List[Any]:
       with self.assertRaisesRegex(
           __lowerCamelCase			,     r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''      ):
            SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained(__lowerCamelCase			,     revision='''aaaaaa'''      )






  def    lowercase_	(							self		:	Tuple      )	->							Any:
       with self.assertRaisesRegex(
           __lowerCamelCase			,     '''hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json.'''			,     ):
            SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained('''hf-internal-testing/config-no-model'''      )






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

       SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained(
           '''hf-internal-testing/test_dynamic_image_processor'''			,     trust_remote_code=__lowerCamelCase      )
       self.assertEqual(image_processor.__class__.__name__			,     '''NewImageProcessor'''      )

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






  def    lowercase_	(							self		:	Dict      )	->							List[str]:
       try:
            AutoConfig.register('''custom'''			,     __lowerCamelCase      )
            AutoImageProcessor.register(__lowerCamelCase			,     __lowerCamelCase      )
            # Trying to register something existing in the Transformers library will raise an error
            with self.assertRaises(__lowerCamelCase      ):
                 AutoImageProcessor.register(__lowerCamelCase			,     __lowerCamelCase      )

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

                 SCREAMING_SNAKE_CASE__					=					CustomImageProcessor.from_pretrained(__lowerCamelCase      )

            # 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(__lowerCamelCase      )
                 SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained(__lowerCamelCase      )
                 self.assertIsInstance(__lowerCamelCase			,     __lowerCamelCase      )

       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		:	Optional[Any]      )	->							List[Any]:






       class  UpperCAmelCase__	( A__ ):






         """simple docstring"""


         a     =      True

       try:
            AutoConfig.register('''custom'''			,     __lowerCamelCase      )
            AutoImageProcessor.register(__lowerCamelCase			,     __lowerCamelCase      )
            # If remote code is not set, the default is to use local
            SCREAMING_SNAKE_CASE__					=					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.
            SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained(
                '''hf-internal-testing/test_dynamic_image_processor'''			,     trust_remote_code=__lowerCamelCase      )
            self.assertEqual(image_processor.__class__.__name__			,     '''NewImageProcessor'''      )
            self.assertTrue(image_processor.is_local      )

            # If remote is enabled, we load from the Hub
            SCREAMING_SNAKE_CASE__					=					AutoImageProcessor.from_pretrained(
                '''hf-internal-testing/test_dynamic_image_processor'''			,     trust_remote_code=__lowerCamelCase      )
            self.assertEqual(image_processor.__class__.__name__			,     '''NewImageProcessor'''      )
            self.assertTrue(not hasattr(__lowerCamelCase			,     '''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]





import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, XLMRobertaTokenizer

from diffusers import AltDiffusionPipeline, AutoencoderKL, DDIMScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class  UpperCAmelCase__	( A__		,  A__		,  A__		,  unittest.TestCase ):






  """simple docstring"""


  a     =      AltDiffusionPipeline
  a     =      TEXT_TO_IMAGE_PARAMS
  a     =      TEXT_TO_IMAGE_BATCH_PARAMS
  a     =      TEXT_TO_IMAGE_IMAGE_PARAMS
  a     =      TEXT_TO_IMAGE_IMAGE_PARAMS






  def    lowercase_	(							self		:	Tuple      )	->							Tuple:
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					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			,     )
       SCREAMING_SNAKE_CASE__					=					DDIMScheduler(
           beta_start=0.00085			,     beta_end=0.012			,     beta_schedule='''scaled_linear'''			,     clip_sample=__lowerCamelCase			,     set_alpha_to_one=__lowerCamelCase			,     )
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					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			,     )

       # TODO: address the non-deterministic text encoder (fails for save-load tests)
       # torch.manual_seed(0)
       # text_encoder_config = RobertaSeriesConfig(
       #     hidden_size=32,
       #     project_dim=32,
       #     intermediate_size=37,
       #     layer_norm_eps=1e-05,
       #     num_attention_heads=4,
       #     num_hidden_layers=5,
       #     vocab_size=5002,
       # )
       # text_encoder = RobertaSeriesModelWithTransformation(text_encoder_config)

       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					CLIPTextConfig(
           bos_token_id=0			,     eos_token_id=2			,     hidden_size=32			,     projection_dim=32			,     intermediate_size=37			,     layer_norm_eps=1e-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     pad_token_id=1			,     vocab_size=5002			,     )
       SCREAMING_SNAKE_CASE__					=					CLIPTextModel(__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					XLMRobertaTokenizer.from_pretrained('''hf-internal-testing/tiny-xlm-roberta'''      )
       SCREAMING_SNAKE_CASE__					=					77

       SCREAMING_SNAKE_CASE__					=					{
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }
       return components






  def    lowercase_	(							self		:	Any			,     __lowerCamelCase		:	Tuple			,     __lowerCamelCase		:	Optional[int]=0      )	->							Union[str, Any]:
       if str(__lowerCamelCase      ).startswith('''mps'''      ):
            SCREAMING_SNAKE_CASE__					=					torch.manual_seed(__lowerCamelCase      )
       else:
            SCREAMING_SNAKE_CASE__					=					torch.Generator(device=__lowerCamelCase      ).manual_seed(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					{
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
       }
       return inputs






  def    lowercase_	(							self		:	List[Any]      )	->							str:
       super().test_attention_slicing_forward_pass(expected_max_diff=3e-3      )






  def    lowercase_	(							self		:	List[Any]      )	->							int:
       super().test_inference_batch_single_identical(expected_max_diff=3e-3      )






  def    lowercase_	(							self		:	Dict      )	->							Optional[int]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesConfig(
           hidden_size=32			,     project_dim=32			,     intermediate_size=37			,     layer_norm_eps=1e-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     vocab_size=5002			,     )
       # TODO: remove after fixing the non-deterministic text encoder
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesModelWithTransformation(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					text_encoder

       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_inputs(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					'''A photo of an astronaut'''
       SCREAMING_SNAKE_CASE__					=					alt_pipe(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					output.images
       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       SCREAMING_SNAKE_CASE__					=					np.array(
           [0.5748162, 0.60447145, 0.48821217, 0.50100636, 0.5431185, 0.45763683, 0.49657696, 0.48132733, 0.47573093]      )

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







  def    lowercase_	(							self		:	List[Any]      )	->							Optional[Any]:
       SCREAMING_SNAKE_CASE__					=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_components()
       SCREAMING_SNAKE_CASE__					=					PNDMScheduler(skip_prk_steps=__lowerCamelCase      )
       torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesConfig(
           hidden_size=32			,     project_dim=32			,     intermediate_size=37			,     layer_norm_eps=1e-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     vocab_size=5002			,     )
       # TODO: remove after fixing the non-deterministic text encoder
       SCREAMING_SNAKE_CASE__					=					RobertaSeriesModelWithTransformation(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					text_encoder
       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					self.get_dummy_inputs(__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe(**__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					output.images
       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       SCREAMING_SNAKE_CASE__					=					np.array(
           [0.51605093, 0.5707241, 0.47365507, 0.50578886, 0.5633877, 0.4642503, 0.5182081, 0.48763484, 0.49084237]      )

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






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






  """simple docstring"""


  def    lowercase_	(							self		:	Tuple      )	->							str:
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






  def    lowercase_	(							self		:	str      )	->							Any:
       # make sure here that pndm scheduler skips prk
       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''			,     safety_checker=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					'''A painting of a squirrel eating a burger'''
       SCREAMING_SNAKE_CASE__					=					torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe([prompt]			,     generator=__lowerCamelCase			,     guidance_scale=6.0			,     num_inference_steps=20			,     output_type='''np'''      )

       SCREAMING_SNAKE_CASE__					=					output.images

       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE__					=					np.array([0.1010, 0.0800, 0.0794, 0.0885, 0.0843, 0.0762, 0.0769, 0.0729, 0.0586]      )

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







  def    lowercase_	(							self		:	Optional[int]      )	->							List[str]:
       SCREAMING_SNAKE_CASE__					=					DDIMScheduler.from_pretrained('''BAAI/AltDiffusion'''			,     subfolder='''scheduler'''      )

       SCREAMING_SNAKE_CASE__					=					AltDiffusionPipeline.from_pretrained('''BAAI/AltDiffusion'''			,     scheduler=__lowerCamelCase			,     safety_checker=__lowerCamelCase      )
       SCREAMING_SNAKE_CASE__					=					alt_pipe.to(__lowerCamelCase      )
       alt_pipe.set_progress_bar_config(disable=__lowerCamelCase      )

       SCREAMING_SNAKE_CASE__					=					'''A painting of a squirrel eating a burger'''
       SCREAMING_SNAKE_CASE__					=					torch.manual_seed(0      )

       SCREAMING_SNAKE_CASE__					=					alt_pipe([prompt]			,     generator=__lowerCamelCase			,     num_inference_steps=2			,     output_type='''numpy'''      )
       SCREAMING_SNAKE_CASE__					=					output.images

       SCREAMING_SNAKE_CASE__					=					image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE__					=					np.array([0.4019, 0.4052, 0.3810, 0.4119, 0.3916, 0.3982, 0.4651, 0.4195, 0.5323]      )

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


