




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 			lowerCamelCase_				( _SCREAMING_SNAKE_CASE			,  _SCREAMING_SNAKE_CASE			,  _SCREAMING_SNAKE_CASE			,  unittest.TestCase   ):


 '''simple docstring'''







 lowercase_							       =       StableUnCLIPPipeline
 lowercase_							       =       TEXT_TO_IMAGE_PARAMS
 lowercase_							       =       TEXT_TO_IMAGE_BATCH_PARAMS
 lowercase_							       =       TEXT_TO_IMAGE_IMAGE_PARAMS
 lowercase_							       =       TEXT_TO_IMAGE_IMAGE_PARAMS

 # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
 lowercase_							       =       False



 def  lowerCAmelCase_    (				self    :						str		):
       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-05							,							num_attention_heads=4							,							num_hidden_layers=5							,							pad_token_id=1							,							vocab_size=1_000							,							)		)

       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=1_000							,							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-05							,							num_attention_heads=4							,							num_hidden_layers=5							,							pad_token_id=1							,							vocab_size=1_000							,							)		)

       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.0_0085							,							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  lowerCAmelCase_    (				self    :						str							,							_lowerCAmelCase    :						Tuple							,							_lowerCAmelCase    :						List[Any]=0		):
       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  lowerCAmelCase_    (				self    :						Optional[int]		):
       SCREAMING_SNAKE_CASE_      		=     torch_device == 'cpu'

       self._test_attention_slicing_forward_pass(test_max_difference=_lowerCAmelCase		)




 def  lowerCAmelCase_    (				self    :						Tuple		):
       SCREAMING_SNAKE_CASE_      		=     torch_device in ['cpu', 'mps']

       self._test_inference_batch_single_identical(test_max_difference=_lowerCAmelCase		)





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


 '''simple docstring'''







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



 def  lowerCAmelCase_    (				self    :						int		):
       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  lowerCAmelCase_    (				self    :						Any		):
       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
import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


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

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


_SCREAMING_SNAKE_CASE 		=      get_tests_dir("""fixtures""")
_SCREAMING_SNAKE_CASE 		=      get_tests_dir("""fixtures/dummy_feature_extractor_config.json""")
_SCREAMING_SNAKE_CASE 		=      get_tests_dir("""fixtures/dummy-config.json""")
class SCREAMING_SNAKE_CASE_		(      unittest.TestCase		):







   def 						lowerCamelCase_	(     self						:						Tuple				):





    """simple docstring"""


    UpperCamelCase			  =      0







   def 						lowerCamelCase_	(     self						:						int				):





    """simple docstring"""


    UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained("""facebook/wav2vec2-base-960h"""				)
    self.assertIsInstance(lowerCamelCase_   ,   lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						Optional[Any]				):





    """simple docstring"""


    UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_				)
    self.assertIsInstance(lowerCamelCase_   ,   lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						Tuple				):





    """simple docstring"""


    with tempfile.TemporaryDirectory() as tmpdirname:
     UpperCamelCase			  =      WavaVecaConfig()

     # remove feature_extractor_type to make sure config.json alone is enough to load feature processor locally
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_				).to_dict()

     config_dict.pop("""feature_extractor_type"""				)
     UpperCamelCase			  =      WavaVecaFeatureExtractor(**lowerCamelCase_				)

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

     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_				)

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

    self.assertIsInstance(lowerCamelCase_   ,   lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						List[str]				):





    """simple docstring"""


    UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_				)
    self.assertIsInstance(lowerCamelCase_   ,   lowerCamelCase_				)







   def 						lowerCamelCase_	(     self						:						int				):





    """simple docstring"""


    with self.assertRaisesRegex(
        lowerCamelCase_   ,   """bert-base is not a local folder and is not a valid model identifier"""				):
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained("""bert-base"""				)







   def 						lowerCamelCase_	(     self						:						str				):





    """simple docstring"""


    with self.assertRaisesRegex(
        lowerCamelCase_   ,   R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""				):
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_   ,   revision="""aaaaaa"""				)







   def 						lowerCamelCase_	(     self						:						Any				):





    """simple docstring"""


    with self.assertRaisesRegex(
        lowerCamelCase_   ,   """hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json."""   ,   ):
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained("""hf-internal-testing/config-no-model"""				)







   def 						lowerCamelCase_	(     self						:						List[Any]				):





    """simple docstring"""


    with self.assertRaises(lowerCamelCase_				):
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
         """hf-internal-testing/test_dynamic_feature_extractor"""				)
    # If remote code is disabled, we can't load this config.
    with self.assertRaises(lowerCamelCase_				):
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
         """hf-internal-testing/test_dynamic_feature_extractor"""   ,   trust_remote_code=lowerCamelCase_				)

    UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
        """hf-internal-testing/test_dynamic_feature_extractor"""   ,   trust_remote_code=lowerCamelCase_				)
    self.assertEqual(feature_extractor.__class__.__name__   ,   """NewFeatureExtractor"""				)

    # Test feature extractor can be reloaded.
    with tempfile.TemporaryDirectory() as tmp_dir:
     feature_extractor.save_pretrained(lowerCamelCase_				)
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_   ,   trust_remote_code=lowerCamelCase_				)
    self.assertEqual(reloaded_feature_extractor.__class__.__name__   ,   """NewFeatureExtractor"""				)







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





    """simple docstring"""


    try:
     AutoConfig.register("""custom"""   ,   lowerCamelCase_				)
     AutoFeatureExtractor.register(lowerCamelCase_   ,   lowerCamelCase_				)
     # Trying to register something existing in the Transformers library will raise an error
     with self.assertRaises(lowerCamelCase_				):
      AutoFeatureExtractor.register(lowerCamelCase_   ,   lowerCamelCase_				)

     # Now that the config is registered, it can be used as any other config with the auto-API
     UpperCamelCase			  =      CustomFeatureExtractor.from_pretrained(lowerCamelCase_				)
     with tempfile.TemporaryDirectory() as tmp_dir:
      feature_extractor.save_pretrained(lowerCamelCase_				)
      UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(lowerCamelCase_				)
      self.assertIsInstance(lowerCamelCase_   ,   lowerCamelCase_				)

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







   def 						lowerCamelCase_	(     self						:						Dict				):





    """simple docstring"""


    class SCREAMING_SNAKE_CASE_		(      __lowerCAmelCase		):
       __lowerCAmelCase           =				True

    try:
     AutoConfig.register("""custom"""   ,   lowerCamelCase_				)
     AutoFeatureExtractor.register(lowerCamelCase_   ,   lowerCamelCase_				)
     # If remote code is not set, the default is to use local
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
         """hf-internal-testing/test_dynamic_feature_extractor"""				)
     self.assertEqual(feature_extractor.__class__.__name__   ,   """NewFeatureExtractor"""				)
     self.assertTrue(feature_extractor.is_local				)

     # If remote code is disabled, we load the local one.
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
         """hf-internal-testing/test_dynamic_feature_extractor"""   ,   trust_remote_code=lowerCamelCase_				)
     self.assertEqual(feature_extractor.__class__.__name__   ,   """NewFeatureExtractor"""				)
     self.assertTrue(feature_extractor.is_local				)

     # If remote is enabled, we load from the Hub
     UpperCamelCase			  =      AutoFeatureExtractor.from_pretrained(
         """hf-internal-testing/test_dynamic_feature_extractor"""   ,   trust_remote_code=lowerCamelCase_				)
     self.assertEqual(feature_extractor.__class__.__name__   ,   """NewFeatureExtractor"""				)
     self.assertTrue(not hasattr(lowerCamelCase_   ,   """is_local"""				)				)

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