


'''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    UpperCAmelCase_  (_UpperCAmelCase					,		_UpperCAmelCase					,		_UpperCAmelCase					,		unittest.TestCase  ):





    """simple docstring"""






    lowerCamelCase  :					int							     =			StableUnCLIPPipeline
    lowerCamelCase  :					str							     =			TEXT_TO_IMAGE_PARAMS
    lowerCamelCase  :					int							     =			TEXT_TO_IMAGE_BATCH_PARAMS
    lowerCamelCase  :					List[Any]							     =			TEXT_TO_IMAGE_IMAGE_PARAMS
    lowerCamelCase  :					Optional[int]							     =			TEXT_TO_IMAGE_IMAGE_PARAMS

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







    def       lowercase_				(							self	)				->     int:
         __lowerCamelCase   :	Optional[Any]								=		32
         __lowerCamelCase   :	int								=		embedder_hidden_size

         # prior components

         torch.manual_seed(0	)
         __lowerCamelCase   :	Any								=		CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

         torch.manual_seed(0	)
         __lowerCamelCase   :	str								=		CLIPTextModelWithProjection(
             CLIPTextConfig(
                 bos_token_id=0			,     eos_token_id=2			,     hidden_size=SCREAMING_SNAKE_CASE_			,     projection_dim=SCREAMING_SNAKE_CASE_			,     intermediate_size=37			,     layer_norm_eps=1E-05			,     num_attention_heads=4			,     num_hidden_layers=5			,     pad_token_id=1			,     vocab_size=10_00			,     )	)

         torch.manual_seed(0	)
         __lowerCamelCase   :	Optional[int]								=		PriorTransformer(
             num_attention_heads=2			,     attention_head_dim=12			,     embedding_dim=SCREAMING_SNAKE_CASE_			,     num_layers=1			,     )

         torch.manual_seed(0	)
         __lowerCamelCase   :	Tuple								=		DDPMScheduler(
             variance_type='fixed_small_log'			,     prediction_type='sample'			,     num_train_timesteps=10_00			,     clip_sample=SCREAMING_SNAKE_CASE_			,     clip_sample_range=5.0			,     beta_schedule='squaredcos_cap_v2'			,     )

         # regular denoising components

         torch.manual_seed(0	)
         __lowerCamelCase   :	Optional[Any]								=		StableUnCLIPImageNormalizer(embedding_dim=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		DDPMScheduler(beta_schedule='squaredcos_cap_v2'	)

         torch.manual_seed(0	)
         __lowerCamelCase   :	Any								=		CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'	)

         torch.manual_seed(0	)
         __lowerCamelCase   :	Optional[int]								=		CLIPTextModel(
             CLIPTextConfig(
                 bos_token_id=0			,     eos_token_id=2			,     hidden_size=SCREAMING_SNAKE_CASE_			,     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=10_00			,     )	)

         torch.manual_seed(0	)
         __lowerCamelCase   :	Optional[int]								=		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=SCREAMING_SNAKE_CASE_			,     layers_per_block=1			,     upcast_attention=SCREAMING_SNAKE_CASE_			,     use_linear_projection=SCREAMING_SNAKE_CASE_			,     )

         torch.manual_seed(0	)
         __lowerCamelCase   :	Tuple								=		DDIMScheduler(
             beta_schedule='scaled_linear'			,     beta_start=0.0_0_0_8_5			,     beta_end=0.0_1_2			,     prediction_type='v_prediction'			,     set_alpha_to_one=SCREAMING_SNAKE_CASE_			,     steps_offset=1			,     )

         torch.manual_seed(0	)
         __lowerCamelCase   :	List[str]								=		AutoencoderKL()

         __lowerCamelCase   :	Any								=		{
             # 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       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=0	)				->     str:
         if str(SCREAMING_SNAKE_CASE_	).startswith('mps'	):
              __lowerCamelCase   :	str								=		torch.manual_seed(SCREAMING_SNAKE_CASE_	)
         else:
              __lowerCamelCase   :	Tuple								=		torch.Generator(device=SCREAMING_SNAKE_CASE_	).manual_seed(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Dict								=		{
             '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       lowercase_				(							self	)				->     Dict:
         __lowerCamelCase   :	Optional[int]								=		torch_device == 'cpu'

         self._test_attention_slicing_forward_pass(test_max_difference=SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self	)				->     Optional[Any]:
         __lowerCamelCase   :	Union[str, Any]								=		torch_device in ['cpu', 'mps']

         self._test_inference_batch_single_identical(test_max_difference=SCREAMING_SNAKE_CASE_	)







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





    """simple docstring"""






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







    def       lowercase_				(							self	)				->     str:
         __lowerCamelCase   :	Dict								=		load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy'	)

         __lowerCamelCase   :	Tuple								=		StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'			,     torch_dtype=torch.floataa	)
         pipe.to(SCREAMING_SNAKE_CASE_	)
         pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)
         # 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()

         __lowerCamelCase   :	Any								=		torch.Generator(device='cpu'	).manual_seed(0	)
         __lowerCamelCase   :	Union[str, Any]								=		pipe('anime turle'			,     generator=SCREAMING_SNAKE_CASE_			,     output_type='np'	)

         __lowerCamelCase   :	int								=		output.images[0]

         assert image.shape == (7_68, 7_68, 3)

         assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self	)				->     Any:
         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

         __lowerCamelCase   :	Any								=		StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'			,     torch_dtype=torch.floataa	)
         __lowerCamelCase   :	Optional[int]								=		pipe.to(SCREAMING_SNAKE_CASE_	)
         pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_	)
         pipe.enable_attention_slicing()
         pipe.enable_sequential_cpu_offload()

         __lowerCamelCase   :	Tuple								=		pipe(
             'anime turtle'			,     prior_num_inference_steps=2			,     num_inference_steps=2			,     output_type='np'			,     )

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










'''simple docstring'''



from importlib import import_module

from .logging import get_logger


A__		:					str												=  get_logger(__name__)





class    UpperCAmelCase_  :





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=None	)				->     List[Any]:
         __lowerCamelCase   :	List[str]								=		attrs or []
         if module is not None:
              for key in module.__dict__:
                   if key in attrs or not key.startswith('__'	):
                        setattr(self			,     SCREAMING_SNAKE_CASE_			,     getattr(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)	)
         __lowerCamelCase   :	Optional[int]								=		module._original_module if isinstance(SCREAMING_SNAKE_CASE_			,     _PatchedModuleObj	) else module







class    UpperCAmelCase_  :





    """simple docstring"""






    lowerCamelCase  :					Tuple							     =			[]







    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=None	)				->     Union[str, Any]:
         __lowerCamelCase   :	Optional[int]								=		obj
         __lowerCamelCase   :	List[Any]								=		target
         __lowerCamelCase   :	Union[str, Any]								=		new
         __lowerCamelCase   :	Union[str, Any]								=		target.split('.'	)[0]
         __lowerCamelCase   :	Dict								=		{}
         __lowerCamelCase   :	Dict								=		attrs or []







    def __enter__(							self	)				->     Optional[int]:
         *__lowerCamelCase				,       __lowerCamelCase   :	int								=		self.target.split('.'	)

         # Patch modules:
         # it's used to patch attributes of submodules like "os.path.join";
         # in this case we need to patch "os" and "os.path"

         for i in range(len(SCREAMING_SNAKE_CASE_	)	):
              try:
                   __lowerCamelCase   :	Optional[int]								=		import_module('.'.join(submodules[: i + 1]	)	)
              except ModuleNotFoundError:
                   continue
              # We iterate over all the globals in self.obj in case we find "os" or "os.path"
              for attr in self.obj.__dir__():
                   __lowerCamelCase   :	List[Any]								=		getattr(self.obj			,     SCREAMING_SNAKE_CASE_	)
                   # We don't check for the name of the global, but rather if its value *is* "os" or "os.path".
                   # This allows to patch renamed modules like "from os import path as ospath".
                   if obj_attr is submodule or (
                       (isinstance(SCREAMING_SNAKE_CASE_			,     _PatchedModuleObj	) and obj_attr._original_module is submodule)
                   ):
                        __lowerCamelCase   :	Optional[Any]								=		obj_attr
                        # patch at top level
                        setattr(self.obj			,     SCREAMING_SNAKE_CASE_			,     _PatchedModuleObj(SCREAMING_SNAKE_CASE_			,     attrs=self.attrs	)	)
                        __lowerCamelCase   :	str								=		getattr(self.obj			,     SCREAMING_SNAKE_CASE_	)
                        # construct lower levels patches
                        for key in submodules[i + 1 :]:
                             setattr(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     _PatchedModuleObj(getattr(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)			,     attrs=self.attrs	)	)
                             __lowerCamelCase   :	List[Any]								=		getattr(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
                        # finally set the target attribute
                        setattr(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     self.new	)

        # Patch attribute itself:
        # it's used for builtins like "open",
        # and also to patch "os.path.join" we may also need to patch "join"
        # itself if it was imported as "from os.path import join".

         if submodules:  # if it's an attribute of a submodule like "os.path.join"
              try:
                   __lowerCamelCase   :	Union[str, Any]								=		getattr(import_module('.'.join(SCREAMING_SNAKE_CASE_	)	)			,     SCREAMING_SNAKE_CASE_	)
              except (AttributeError, ModuleNotFoundError):
                   return
              # We iterate over all the globals in self.obj in case we find "os.path.join"
              for attr in self.obj.__dir__():
                   # We don't check for the name of the global, but rather if its value *is* "os.path.join".
                   # This allows to patch renamed attributes like "from os.path import join as pjoin".
                   if getattr(self.obj			,     SCREAMING_SNAKE_CASE_	) is attr_value:
                        __lowerCamelCase   :	Optional[int]								=		getattr(self.obj			,     SCREAMING_SNAKE_CASE_	)
                        setattr(self.obj			,     SCREAMING_SNAKE_CASE_			,     self.new	)
         elif target_attr in globals()["__builtins__"]:  # if it'a s builtin like "open"
              __lowerCamelCase   :	List[Any]								=		globals()['__builtins__'][target_attr]
              setattr(self.obj			,     SCREAMING_SNAKE_CASE_			,     self.new	)
         else:
              raise RuntimeError(f'Tried to patch attribute {target_attr} instead of a submodule.'	)







    def __exit__(							self			,     *SCREAMING_SNAKE_CASE_	)				->     Optional[Any]:
         for attr in list(self.original	):
              setattr(self.obj			,     SCREAMING_SNAKE_CASE_			,     self.original.pop(SCREAMING_SNAKE_CASE_	)	)







    def       lowercase_				(							self	)				->     Optional[int]:
         self.__enter__()
         self._active_patches.append(self	)







    def       lowercase_				(							self	)				->     str:
         try:
              self._active_patches.remove(self	)
         except ValueError:
              # If the patch hasn't been started this will fail
              return None

         return self.__exit__()






