






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()





class     __a						(					UpperCamelCase__ ,     UpperCamelCase__ ,     UpperCamelCase__ ,     unittest.TestCase ):
 _lowerCAmelCase : Optional[Any] 				=				StableDiffusionImageVariationPipeline
 _lowerCAmelCase : Optional[int] 				=				IMAGE_VARIATION_PARAMS
 _lowerCAmelCase : Optional[int] 				=				IMAGE_VARIATION_BATCH_PARAMS
 _lowerCAmelCase : Optional[int] 				=				frozenset([] )
 # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 _lowerCAmelCase : str 				=				frozenset([] )







 def       __lowercase      (	self      :     List[str]		):


      '''simple docstring'''






      torch.manual_seed(0		)
      UpperCamelCase__    :     List[Any]        =						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							,     )
      UpperCamelCase__    :     Union[str, Any]        =						PNDMScheduler(skip_prk_steps=lowerCAmelCase__		)
      torch.manual_seed(0		)
      UpperCamelCase__    :     int        =						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							,     )
      torch.manual_seed(0		)
      UpperCamelCase__    :     List[str]        =						CLIPVisionConfig(
          hidden_size=32							,     projection_dim=32							,     intermediate_size=37							,     layer_norm_eps=1e-0_5							,     num_attention_heads=4							,     num_hidden_layers=5							,     image_size=32							,     patch_size=4							,     )
      UpperCamelCase__    :     Tuple        =						CLIPVisionModelWithProjection(lowerCAmelCase__		)
      UpperCamelCase__    :     int        =						CLIPImageProcessor(crop_size=32							,     size=32		)

      UpperCamelCase__    :     Optional[int]        =						{
          "unet": unet,
          "scheduler": scheduler,
          "vae": vae,
          "image_encoder": image_encoder,
          "feature_extractor": feature_extractor,
          "safety_checker": None,
      }
      return components







 def       __lowercase      (	self      :     Dict							,     SCREAMING_SNAKE_CASE      :     Tuple							,     SCREAMING_SNAKE_CASE      :     Optional[Any]=0		):


      '''simple docstring'''






      UpperCamelCase__    :     Optional[Any]        =						floats_tensor((1, 3, 32, 32)							,     rng=random.Random(lowerCAmelCase__		)		)
      UpperCamelCase__    :     Optional[Any]        =						image.cpu().permute(0							,     2							,     3							,     1		)[0]
      UpperCamelCase__    :     int        =						Image.fromarray(np.uinta(lowerCAmelCase__		)		).convert("RGB"		).resize((32, 32)		)
      if str(lowerCAmelCase__		).startswith("mps"		):
           UpperCamelCase__    :     str        =						torch.manual_seed(lowerCAmelCase__		)
      else:
           UpperCamelCase__    :     str        =						torch.Generator(device=lowerCAmelCase__		).manual_seed(lowerCAmelCase__		)
      UpperCamelCase__    :     str        =						{
          "image": image,
          "generator": generator,
          "num_inference_steps": 2,
          "guidance_scale": 6.0,
          "output_type": "numpy",
      }
      return inputs







 def       __lowercase      (	self      :     Optional[Any]		):


      '''simple docstring'''






      UpperCamelCase__    :     Tuple        =						"cpu"  # ensure determinism for the device-dependent torch.Generator
      UpperCamelCase__    :     Optional[Any]        =						self.get_dummy_components()
      UpperCamelCase__    :     List[str]        =						StableDiffusionImageVariationPipeline(**lowerCAmelCase__		)
      UpperCamelCase__    :     Optional[Any]        =						sd_pipe.to(lowerCAmelCase__		)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

      UpperCamelCase__    :     Any        =						self.get_dummy_inputs(lowerCAmelCase__		)
      UpperCamelCase__    :     Optional[Any]        =						sd_pipe(**lowerCAmelCase__		).images
      UpperCamelCase__    :     int        =						image[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)
      UpperCamelCase__    :     int        =						np.array([0.5_2_3_9, 0.5_7_2_3, 0.4_7_9_6, 0.5_0_4_9, 0.5_5_5_0, 0.4_6_8_5, 0.5_3_2_9, 0.4_8_9_1, 0.4_9_2_1]		)

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







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


      '''simple docstring'''






      UpperCamelCase__    :     int        =						"cpu"  # ensure determinism for the device-dependent torch.Generator
      UpperCamelCase__    :     Any        =						self.get_dummy_components()
      UpperCamelCase__    :     Tuple        =						StableDiffusionImageVariationPipeline(**lowerCAmelCase__		)
      UpperCamelCase__    :     Dict        =						sd_pipe.to(lowerCAmelCase__		)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

      UpperCamelCase__    :     Union[str, Any]        =						self.get_dummy_inputs(lowerCAmelCase__		)
      UpperCamelCase__    :     Optional[int]        =						2 * [inputs["image"]]
      UpperCamelCase__    :     Optional[Any]        =						sd_pipe(**lowerCAmelCase__		)

      UpperCamelCase__    :     Union[str, Any]        =						output.images

      UpperCamelCase__    :     Any        =						image[-1, -3:, -3:, -1]

      assert image.shape == (2, 64, 64, 3)
      UpperCamelCase__    :     List[str]        =						np.array([0.6_8_9_2, 0.5_6_3_7, 0.5_8_3_6, 0.5_7_7_1, 0.6_2_5_4, 0.6_4_0_9, 0.5_5_8_0, 0.5_5_6_9, 0.5_2_8_9]		)

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







 def       __lowercase      (	self      :     Optional[int]		):


      '''simple docstring'''






      super().test_inference_batch_single_identical(expected_max_diff=3e-3		)







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







 def       __lowercase      (	self      :     Optional[Any]		):


      '''simple docstring'''






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







 def       __lowercase      (	self      :     Optional[int]							,     SCREAMING_SNAKE_CASE      :     Union[str, Any]							,     SCREAMING_SNAKE_CASE      :     int="cpu"							,     SCREAMING_SNAKE_CASE      :     List[str]=torch.floataa							,     SCREAMING_SNAKE_CASE      :     int=0		):


      '''simple docstring'''






      UpperCamelCase__    :     int        =						torch.Generator(device=lowerCAmelCase__		).manual_seed(lowerCAmelCase__		)
      UpperCamelCase__    :     List[Any]        =						load_image(
          "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
          "/stable_diffusion_imgvar/input_image_vermeer.png"		)
      UpperCamelCase__    :     int        =						np.random.RandomState(lowerCAmelCase__		).standard_normal((1, 4, 64, 64)		)
      UpperCamelCase__    :     Any        =						torch.from_numpy(lowerCAmelCase__		).to(device=lowerCAmelCase__							,     dtype=lowerCAmelCase__		)
      UpperCamelCase__    :     str        =						{
          "image": init_image,
          "latents": latents,
          "generator": generator,
          "num_inference_steps": 3,
          "guidance_scale": 7.5,
          "output_type": "numpy",
      }
      return inputs







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


      '''simple docstring'''






      UpperCamelCase__    :     int        =						StableDiffusionImageVariationPipeline.from_pretrained(
          "lambdalabs/sd-image-variations-diffusers"							,     safety_checker=lowerCAmelCase__		)
      UpperCamelCase__    :     Union[str, Any]        =						sd_pipe.to(lowerCAmelCase__		)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

      UpperCamelCase__    :     Dict        =						self.get_inputs(lowerCAmelCase__		)
      UpperCamelCase__    :     Any        =						sd_pipe(**lowerCAmelCase__		).images
      UpperCamelCase__    :     List[Any]        =						image[0, -3:, -3:, -1].flatten()

      assert image.shape == (1, 5_12, 5_12, 3)
      UpperCamelCase__    :     Any        =						np.array([0.8_4_4_9_1, 0.9_0_7_8_9, 0.7_5_7_0_8, 0.7_8_7_3_4, 0.8_3_4_8_5, 0.7_0_0_9_9, 0.6_6_9_3_8, 0.6_8_7_2_7, 0.6_1_3_7_9]		)
      assert np.abs(image_slice - expected_slice		).max() < 6e-3







 def       __lowercase      (	self      :     Optional[Any]		):


      '''simple docstring'''






      UpperCamelCase__    :     Tuple        =						0

      def callback_fn(SCREAMING_SNAKE_CASE      :     Any							,     SCREAMING_SNAKE_CASE      :     str							,     SCREAMING_SNAKE_CASE      :     int		) -> None:
           UpperCamelCase__    :     int        =						True
           nonlocal number_of_steps
           number_of_steps += 1
           if step == 1:
                UpperCamelCase__    :     Union[str, Any]        =						latents.detach().cpu().numpy()
                assert latents.shape == (1, 4, 64, 64)
                UpperCamelCase__    :     Any        =						latents[0, -3:, -3:, -1]
                UpperCamelCase__    :     Any        =						np.array(
                    [-0.1_6_2_1, 0.2_8_3_7, -0.7_9_7_9, -0.1_2_2_1, -1.3_0_5_7, 0.7_6_8_1, -2.1_1_9_1, 0.0_4_6_4, 1.6_3_0_9]		)

                assert np.abs(latents_slice.flatten() - expected_slice		).max() < 5e-2
           elif step == 2:
                UpperCamelCase__    :     List[Any]        =						latents.detach().cpu().numpy()
                assert latents.shape == (1, 4, 64, 64)
                UpperCamelCase__    :     Union[str, Any]        =						latents[0, -3:, -3:, -1]
                UpperCamelCase__    :     Dict        =						np.array([0.6_2_9_9, 1.7_5_0_0, 1.1_9_9_2, -2.1_5_8_2, -1.8_9_9_4, 0.7_3_3_4, -0.7_0_9_0, 1.0_1_3_7, 1.5_2_7_3]		)

                assert np.abs(latents_slice.flatten() - expected_slice		).max() < 5e-2

      UpperCamelCase__    :     str        =						False

      UpperCamelCase__    :     List[Any]        =						StableDiffusionImageVariationPipeline.from_pretrained(
          "fusing/sd-image-variations-diffusers"							,     safety_checker=lowerCAmelCase__							,     torch_dtype=torch.floataa							,     )
      pipe.to(lowerCAmelCase__		)
      pipe.set_progress_bar_config(disable=lowerCAmelCase__		)
      pipe.enable_attention_slicing()

      UpperCamelCase__    :     Optional[int]        =						self.get_inputs(lowerCAmelCase__							,     dtype=torch.floataa		)
      pipe(**lowerCAmelCase__							,     callback=lowerCAmelCase__							,     callback_steps=1		)
      assert callback_fn.has_been_called
      assert number_of_steps == inputs["num_inference_steps"]







 def       __lowercase      (	self      :     List[str]		):


      '''simple docstring'''






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

      UpperCamelCase__    :     Dict        =						"fusing/sd-image-variations-diffusers"
      UpperCamelCase__    :     Optional[int]        =						StableDiffusionImageVariationPipeline.from_pretrained(
          lowerCAmelCase__							,     safety_checker=lowerCAmelCase__							,     torch_dtype=torch.floataa		)
      UpperCamelCase__    :     Optional[Any]        =						pipe.to(lowerCAmelCase__		)
      pipe.set_progress_bar_config(disable=lowerCAmelCase__		)
      pipe.enable_attention_slicing(1		)
      pipe.enable_sequential_cpu_offload()

      UpperCamelCase__    :     Any        =						self.get_inputs(lowerCAmelCase__							,     dtype=torch.floataa		)
      UpperCamelCase__    :     str        =						pipe(**lowerCAmelCase__		)

      UpperCamelCase__    :     Tuple        =						torch.cuda.max_memory_allocated()
      # make sure that less than 2.6 GB is allocated
      assert mem_bytes < 2.6 * 10**9







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







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


      '''simple docstring'''






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







 def       __lowercase      (	self      :     Union[str, Any]							,     SCREAMING_SNAKE_CASE      :     str							,     SCREAMING_SNAKE_CASE      :     Optional[int]="cpu"							,     SCREAMING_SNAKE_CASE      :     Optional[Any]=torch.floataa							,     SCREAMING_SNAKE_CASE      :     Tuple=0		):


      '''simple docstring'''






      UpperCamelCase__    :     Tuple        =						torch.Generator(device=lowerCAmelCase__		).manual_seed(lowerCAmelCase__		)
      UpperCamelCase__    :     List[str]        =						load_image(
          "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
          "/stable_diffusion_imgvar/input_image_vermeer.png"		)
      UpperCamelCase__    :     Tuple        =						np.random.RandomState(lowerCAmelCase__		).standard_normal((1, 4, 64, 64)		)
      UpperCamelCase__    :     Any        =						torch.from_numpy(lowerCAmelCase__		).to(device=lowerCAmelCase__							,     dtype=lowerCAmelCase__		)
      UpperCamelCase__    :     str        =						{
          "image": init_image,
          "latents": latents,
          "generator": generator,
          "num_inference_steps": 50,
          "guidance_scale": 7.5,
          "output_type": "numpy",
      }
      return inputs







 def       __lowercase      (	self      :     List[str]		):


      '''simple docstring'''






      UpperCamelCase__    :     Optional[int]        =						StableDiffusionImageVariationPipeline.from_pretrained("fusing/sd-image-variations-diffusers"		)
      sd_pipe.to(lowerCAmelCase__		)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

      UpperCamelCase__    :     Union[str, Any]        =						self.get_inputs(lowerCAmelCase__		)
      UpperCamelCase__    :     Dict        =						sd_pipe(**lowerCAmelCase__		).images[0]

      UpperCamelCase__    :     Any        =						load_numpy(
          "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
          "/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy"		)
      UpperCamelCase__    :     Tuple        =						np.abs(expected_image - image		).max()
      assert max_diff < 1e-3







 def       __lowercase      (	self      :     List[Any]		):


      '''simple docstring'''






      UpperCamelCase__    :     List[Any]        =						StableDiffusionImageVariationPipeline.from_pretrained("fusing/sd-image-variations-diffusers"		)
      UpperCamelCase__    :     Union[str, Any]        =						DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config		)
      sd_pipe.to(lowerCAmelCase__		)
      sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__		)

      UpperCamelCase__    :     Any        =						self.get_inputs(lowerCAmelCase__		)
      UpperCamelCase__    :     Union[str, Any]        =						25
      UpperCamelCase__    :     List[str]        =						sd_pipe(**lowerCAmelCase__		).images[0]

      UpperCamelCase__    :     int        =						load_numpy(
          "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
          "/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy"		)
      UpperCamelCase__    :     Any        =						np.abs(expected_image - image		).max()
      assert max_diff < 1e-3







def 	SCREAMING_SNAKE_CASE							(      __lowerCAmelCase    )				->			str:
     return "".join(chr(ord(__lowerCAmelCase    ) + 32    ) if "A" <= char <= "Z" else char for char in word    )


if __name__ == "__main__":
 from doctest import testmod

 testmod()