





'''simple docstring'''

import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

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,
    assert_mean_pixel_difference,
)


enable_full_determinism()







class __a			( _snake_case, _snake_case, _snake_case, unittest.TestCase      ):
   __UpperCamelCase :			Any						   =						StableUnCLIPPipeline
   __UpperCamelCase :			Tuple						   =						TEXT_TO_IMAGE_PARAMS
   __UpperCamelCase :			List[str]						   =						TEXT_TO_IMAGE_BATCH_PARAMS
   __UpperCamelCase :			Optional[Any]						   =						TEXT_TO_IMAGE_IMAGE_PARAMS
   __UpperCamelCase :			List[str]						   =						TEXT_TO_IMAGE_IMAGE_PARAMS

   # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
   __UpperCamelCase :			List[str]						   =						False





   def   UpperCAmelCase__  (       self					:      int     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							32
      __SCREAMING_SNAKE_CASE   				=							embedder_hidden_size

      # prior components

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							CLIPTextModelWithProjection(
          CLIPTextConfig(
              bos_token_id=0					,eos_token_id=2					,hidden_size=lowerCamelCase					,projection_dim=lowerCamelCase					,intermediate_size=37					,layer_norm_eps=1E-0_5					,num_attention_heads=4					,num_hidden_layers=5					,pad_token_id=1					,vocab_size=1000					,)     )

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							PriorTransformer(
          num_attention_heads=2					,attention_head_dim=12					,embedding_dim=lowerCamelCase					,num_layers=1					,)

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							DDPMScheduler(
          variance_type="""fixed_small_log"""					,prediction_type="""sample"""					,num_train_timesteps=1000					,clip_sample=lowerCamelCase					,clip_sample_range=5.0					,beta_schedule="""squaredcos_cap_v2"""					,)

      # regular denoising components

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							StableUnCLIPImageNormalizer(embedding_dim=lowerCamelCase     )
      __SCREAMING_SNAKE_CASE   				=							DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""     )

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							CLIPTextModel(
          CLIPTextConfig(
              bos_token_id=0					,eos_token_id=2					,hidden_size=lowerCamelCase					,projection_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=1000					,)     )

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							UNetaDConditionModel(
          sample_size=32					,in_channels=4					,out_channels=4					,down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")					,up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")					,block_out_channels=(32, 64)					,attention_head_dim=(2, 4)					,class_embed_type="""projection"""					,projection_class_embeddings_input_dim=embedder_projection_dim * 2					,cross_attention_dim=lowerCamelCase					,layers_per_block=1					,upcast_attention=lowerCamelCase					,use_linear_projection=lowerCamelCase					,)

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							DDIMScheduler(
          beta_schedule="""scaled_linear"""					,beta_start=0.00_085					,beta_end=0.012					,prediction_type="""v_prediction"""					,set_alpha_to_one=lowerCamelCase					,steps_offset=1					,)

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							AutoencoderKL()

      __SCREAMING_SNAKE_CASE   				=							{
          # prior components
          """prior_tokenizer""": prior_tokenizer,
          """prior_text_encoder""": prior_text_encoder,
          """prior""": prior,
          """prior_scheduler""": prior_scheduler,
          # image noising components
          """image_normalizer""": image_normalizer,
          """image_noising_scheduler""": image_noising_scheduler,
          # regular denoising components
          """tokenizer""": tokenizer,
          """text_encoder""": text_encoder,
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
      }

      return components





   def   UpperCAmelCase__  (       self					:      Dict					,lowerCamelCase					:      int					,lowerCamelCase					:      Tuple=0     ):

      '''simple docstring'''

      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,
          """prior_num_inference_steps""": 2,
          """output_type""": """numpy""",
      }
      return inputs





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

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							torch_device == """cpu"""

      self._test_attention_slicing_forward_pass(test_max_difference=lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Optional[int]     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							torch_device in ["""cpu""", """mps"""]

      self._test_inference_batch_single_identical(test_max_difference=lowerCamelCase     )







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





   def   UpperCAmelCase__  (       self					:      List[str]     ):

      '''simple docstring'''

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





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

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""     )

      __SCREAMING_SNAKE_CASE   				=							StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""					,torch_dtype=torch.floataa     )
      pipe.to(lowerCamelCase     )
      pipe.set_progress_bar_config(disable=lowerCamelCase     )
      # stable unclip will oom when integration tests are run on a V100,
      # so turn on memory savings
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      __SCREAMING_SNAKE_CASE   				=							torch.Generator(device="""cpu"""     ).manual_seed(0     )
      __SCREAMING_SNAKE_CASE   				=							pipe("""anime turle"""					,generator=lowerCamelCase					,output_type="""np"""     )

      __SCREAMING_SNAKE_CASE   				=							output.images[0]

      assert image.shape == (768, 768, 3)

      assert_mean_pixel_difference(lowerCamelCase					,lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      List[str]     ):

      '''simple docstring'''

      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()

      __SCREAMING_SNAKE_CASE   				=							StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""					,torch_dtype=torch.floataa     )
      __SCREAMING_SNAKE_CASE   				=							pipe.to(lowerCamelCase     )
      pipe.set_progress_bar_config(disable=lowerCamelCase     )
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      __SCREAMING_SNAKE_CASE   				=							pipe(
          """anime turtle"""					,prior_num_inference_steps=2					,num_inference_steps=2					,output_type="""np"""					,)

      __SCREAMING_SNAKE_CASE   				=							torch.cuda.max_memory_allocated()
      # make sure that less than 7 GB is allocated
      assert mem_bytes < 7 * 10**9












'''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 __a			( unittest.TestCase      ):





   def   UpperCAmelCase__  (       self					:      Any     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							0





   def   UpperCAmelCase__  (       self					:      str     ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE   				=							AutoImageProcessor.from_pretrained("""openai/clip-vit-base-patch32"""     )
      self.assertIsInstance(lowerCamelCase					,lowerCamelCase     )





   def   UpperCAmelCase__  (       self					:      Dict     ):

      '''simple docstring'''

      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   UpperCAmelCase__  (       self					:      Optional[Any]     ):

      '''simple docstring'''

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

      '''simple docstring'''

      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   UpperCAmelCase__  (       self					:      Union[str, Any]     ):

      '''simple docstring'''

      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   UpperCAmelCase__  (       self					:      Dict     ):

      '''simple docstring'''

      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   UpperCAmelCase__  (       self					:      Any     ):

      '''simple docstring'''

      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   UpperCAmelCase__  (       self					:      List[Any]     ):

      '''simple docstring'''

      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   UpperCAmelCase__  (       self					:      Any     ):

      '''simple docstring'''

      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   UpperCAmelCase__  (       self					:      str     ):

      '''simple docstring'''

      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   UpperCAmelCase__  (       self					:      str     ):

      '''simple docstring'''

      class __a			( _snake_case      ):
         __UpperCamelCase :			int						   =						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]





