

'''simple docstring'''

import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()





class 							_snake_case (       _a	,  _a	,  _a	,  unittest.TestCase						):
   _A       :     Any       =					StableUnCLIPImgaImgPipeline
   _A       :     List[str]       =					TEXT_GUIDED_IMAGE_VARIATION_PARAMS
   _A       :     Optional[Any]       =					TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
   _A       :     Dict       =					frozenset(
       []						)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
   _A       :     Optional[Any]       =					frozenset([]						)
   def     __UpperCamelCase						( self			:      Union[str, Any]							):
     SCREAMING_SNAKE_CASE:str						      =					32
     SCREAMING_SNAKE_CASE:Optional[int]						      =					embedder_hidden_size

     # image encoding components

     SCREAMING_SNAKE_CASE:List[str]						      =					CLIPImageProcessor(crop_size=32     ,size=32							)

     torch.manual_seed(0							)
     SCREAMING_SNAKE_CASE:str						      =					CLIPVisionModelWithProjection(
         CLIPVisionConfig(
             hidden_size=SCREAMING_SNAKE_CASE__     ,projection_dim=SCREAMING_SNAKE_CASE__     ,num_hidden_layers=5     ,num_attention_heads=4     ,image_size=32     ,intermediate_size=37     ,patch_size=1     ,)							)

     # regular denoising components

     torch.manual_seed(0							)
     SCREAMING_SNAKE_CASE:Tuple						      =					StableUnCLIPImageNormalizer(embedding_dim=SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					DDPMScheduler(beta_schedule="squaredcos_cap_v2"							)

     torch.manual_seed(0							)
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							)

     torch.manual_seed(0							)
     SCREAMING_SNAKE_CASE:List[Any]						      =					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=1_000     ,)							)

     torch.manual_seed(0							)
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					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							)
     SCREAMING_SNAKE_CASE:Tuple						      =					DDIMScheduler(
         beta_schedule="scaled_linear"     ,beta_start=0.00_085     ,beta_end=0.012     ,prediction_type="v_prediction"     ,set_alpha_to_one=SCREAMING_SNAKE_CASE__     ,steps_offset=1     ,)

     torch.manual_seed(0							)
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					AutoencoderKL()

     SCREAMING_SNAKE_CASE:Optional[int]						      =					{
         # image encoding components
         "feature_extractor": feature_extractor,
         "image_encoder": image_encoder.eval(),
         # image noising components
         "image_normalizer": image_normalizer.eval(),
         "image_noising_scheduler": image_noising_scheduler,
         # regular denoising components
         "tokenizer": tokenizer,
         "text_encoder": text_encoder.eval(),
         "unet": unet.eval(),
         "scheduler": scheduler,
         "vae": vae.eval(),
     }

     return components
   def     __UpperCamelCase						( self			:      Optional[int]     ,SCREAMING_SNAKE_CASE__			:      Tuple     ,SCREAMING_SNAKE_CASE__			:      Any=0     ,SCREAMING_SNAKE_CASE__			:      Any=True							):
     if str(SCREAMING_SNAKE_CASE__							).startswith("mps"							):
       SCREAMING_SNAKE_CASE:Dict						      =					torch.manual_seed(SCREAMING_SNAKE_CASE__							)
     else:
       SCREAMING_SNAKE_CASE:Optional[int]						      =					torch.Generator(device=SCREAMING_SNAKE_CASE__							).manual_seed(SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:List[Any]						      =					floats_tensor((1, 3, 32, 32)     ,rng=random.Random(SCREAMING_SNAKE_CASE__							)							).to(SCREAMING_SNAKE_CASE__							)

     if pil_image:
       SCREAMING_SNAKE_CASE:Dict						      =					input_image * 0.5 + 0.5
       SCREAMING_SNAKE_CASE:str						      =					input_image.clamp(0     ,1							)
       SCREAMING_SNAKE_CASE:Tuple						      =					input_image.cpu().permute(0     ,2     ,3     ,1							).float().numpy()
       SCREAMING_SNAKE_CASE:int						      =					DiffusionPipeline.numpy_to_pil(SCREAMING_SNAKE_CASE__							)[0]

     return {
         "prompt": "An anime racoon running a marathon",
         "image": input_image,
         "generator": generator,
         "num_inference_steps": 2,
         "output_type": "np",
     }
   @skip_mps
   def     __UpperCamelCase						( self			:      Tuple							):
     SCREAMING_SNAKE_CASE:Dict						      =					"cpu"  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE:List[str]						      =					self.get_dummy_components()
     SCREAMING_SNAKE_CASE:Dict						      =					StableUnCLIPImgaImgPipeline(**SCREAMING_SNAKE_CASE__							)
     SCREAMING_SNAKE_CASE:str						      =					sd_pipe.to(SCREAMING_SNAKE_CASE__							)
     sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

     SCREAMING_SNAKE_CASE:Any						      =					self.get_dummy_inputs(SCREAMING_SNAKE_CASE__							)
     inputs.update({"image_embeds": None}							)
     SCREAMING_SNAKE_CASE:List[str]						      =					sd_pipe(**SCREAMING_SNAKE_CASE__							).images
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					image[0, -3:, -3:, -1]

     assert image.shape == (1, 32, 32, 3)
     SCREAMING_SNAKE_CASE:str						      =					np.array([0.3_872, 0.7_224, 0.5_601, 0.4_741, 0.6_872, 0.5_814, 0.4_636, 0.3_867, 0.5_078]							)

     assert np.abs(image_slice.flatten() - expected_slice							).max() < 1e-3
   def     __UpperCamelCase						( self			:      str							):
     SCREAMING_SNAKE_CASE:int						      =					torch_device in ["cpu", "mps"]

     self._test_attention_slicing_forward_pass(test_max_difference=SCREAMING_SNAKE_CASE__							)
   def     __UpperCamelCase						( self			:      List[str]							):
     SCREAMING_SNAKE_CASE:int						      =					torch_device in ["cpu", "mps"]

     self._test_inference_batch_single_identical(test_max_difference=SCREAMING_SNAKE_CASE__							)







   @unittest.skipIf(
       torch_device != "cuda" or not is_xformers_available()     ,reason="XFormers attention is only available with CUDA and `xformers` installed"     ,)
   def     __UpperCamelCase						( self			:      List[str]							):
     self._test_xformers_attention_forwardGenerator_pass(test_max_difference=SCREAMING_SNAKE_CASE__							)





@slow
@require_torch_gpu
class 							_snake_case (       unittest.TestCase						):
   def     __UpperCamelCase						( self			:      List[str]							):
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()
   def     __UpperCamelCase						( self			:      Union[str, Any]							):
     SCREAMING_SNAKE_CASE:Any						      =					load_image(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"							)

     SCREAMING_SNAKE_CASE:Optional[Any]						      =					load_numpy(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy"							)

     SCREAMING_SNAKE_CASE:List[Any]						      =					StableUnCLIPImgaImgPipeline.from_pretrained(
         "fusing/stable-unclip-2-1-l-img2img"     ,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()

     SCREAMING_SNAKE_CASE:Dict						      =					torch.Generator(device="cpu"							).manual_seed(0							)
     SCREAMING_SNAKE_CASE:str						      =					pipe(SCREAMING_SNAKE_CASE__     ,"anime turle"     ,generator=SCREAMING_SNAKE_CASE__     ,output_type="np"							)

     SCREAMING_SNAKE_CASE:List[Any]						      =					output.images[0]

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

     assert_mean_pixel_difference(SCREAMING_SNAKE_CASE__     ,SCREAMING_SNAKE_CASE__							)
   def     __UpperCamelCase						( self			:      List[Any]							):
     SCREAMING_SNAKE_CASE:List[str]						      =					load_image(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"							)

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

     SCREAMING_SNAKE_CASE:int						      =					StableUnCLIPImgaImgPipeline.from_pretrained(
         "fusing/stable-unclip-2-1-h-img2img"     ,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()

     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					torch.Generator(device="cpu"							).manual_seed(0							)
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					pipe(SCREAMING_SNAKE_CASE__     ,"anime turle"     ,generator=SCREAMING_SNAKE_CASE__     ,output_type="np"							)

     SCREAMING_SNAKE_CASE:int						      =					output.images[0]

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

     assert_mean_pixel_difference(SCREAMING_SNAKE_CASE__     ,SCREAMING_SNAKE_CASE__							)







   def     __UpperCamelCase						( self			:      List[str]							):
     SCREAMING_SNAKE_CASE:Any						      =					load_image(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"							)

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

     SCREAMING_SNAKE_CASE:Dict						      =					StableUnCLIPImgaImgPipeline.from_pretrained(
         "fusing/stable-unclip-2-1-h-img2img"     ,torch_dtype=torch.floataa							)
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					pipe.to(SCREAMING_SNAKE_CASE__							)
     pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)
     pipe.enable_attention_slicing()
     pipe.enable_sequential_cpu_offload()

     SCREAMING_SNAKE_CASE:int						      =					pipe(
         SCREAMING_SNAKE_CASE__     ,"anime turtle"     ,num_inference_steps=2     ,output_type="np"     ,)

     SCREAMING_SNAKE_CASE:Optional[int]						      =					torch.cuda.max_memory_allocated()
     # make sure that less than 7 GB is allocated
     assert mem_bytes < 7 * 10**9









'''simple docstring'''

import warnings
from functools import wraps
from typing import Callable





def    A_  ( snake_case ):

  @wraps(snake_case )
  def _inner_fn(*snake_case      ,       **snake_case ):
    warnings.warn(
        (F'''\'{fn.__name__}\' is experimental and might be subject to breaking changes in the future.''')      ,       snake_case      ,       )
    return fn(*snake_case      ,       **snake_case )

  return _inner_fn






