




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__       (   snake_case__					,							snake_case__					,							snake_case__					,							unittest.TestCase ):
						_a				:       List[str]   =       StableDiffusionImageVariationPipeline
						_a				:       Any   =       IMAGE_VARIATION_PARAMS
						_a				:       Dict   =       IMAGE_VARIATION_BATCH_PARAMS
						_a				:       Optional[int]   =       frozenset([] )
						# TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
						_a				:       Any   =       frozenset([] )






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											torch.manual_seed(0     )
											__lowerCAmelCase							  =			UNetaDConditionModel(
											    block_out_channels=(3_2, 6_4)						,     layers_per_block=2						,     sample_size=3_2						,     in_channels=4						,     out_channels=4						,     down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")						,     up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")						,     cross_attention_dim=3_2						,     )
											__lowerCAmelCase							  =			PNDMScheduler(skip_prk_steps=_A     )
											torch.manual_seed(0     )
											__lowerCAmelCase							  =			AutoencoderKL(
											    block_out_channels=[3_2, 6_4]						,     in_channels=3						,     out_channels=3						,     down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]						,     up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]						,     latent_channels=4						,     )
											torch.manual_seed(0     )
											__lowerCAmelCase							  =			CLIPVisionConfig(
											    hidden_size=3_2						,     projection_dim=3_2						,     intermediate_size=3_7						,     layer_norm_eps=1E-0_5						,     num_attention_heads=4						,     num_hidden_layers=5						,     image_size=3_2						,     patch_size=4						,     )
											__lowerCAmelCase							  =			CLIPVisionModelWithProjection(_A     )
											__lowerCAmelCase							  =			CLIPImageProcessor(crop_size=3_2						,     size=3_2     )

											__lowerCAmelCase							  =			{
											    "unet": unet,
											    "scheduler": scheduler,
											    "vae": vae,
											    "image_encoder": image_encoder,
											    "feature_extractor": feature_extractor,
											    "safety_checker": None,
											}
											return components






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A=0     ):




											"""simple docstring"""



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






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			"cpu"  # ensure determinism for the device-dependent torch.Generator
											__lowerCAmelCase							  =			self.get_dummy_components()
											__lowerCAmelCase							  =			StableDiffusionImageVariationPipeline(**_A     )
											__lowerCAmelCase							  =			sd_pipe.to(_A     )
											sd_pipe.set_progress_bar_config(disable=_A     )

											__lowerCAmelCase							  =			self.get_dummy_inputs(_A     )
											__lowerCAmelCase							  =			sd_pipe(**_A     ).images
											__lowerCAmelCase							  =			image[0, -3:, -3:, -1]

											assert image.shape == (1, 6_4, 6_4, 3)
											__lowerCAmelCase							  =			np.array([0.52_39, 0.57_23, 0.47_96, 0.50_49, 0.55_50, 0.46_85, 0.53_29, 0.48_91, 0.49_21]     )

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






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			"cpu"  # ensure determinism for the device-dependent torch.Generator
											__lowerCAmelCase							  =			self.get_dummy_components()
											__lowerCAmelCase							  =			StableDiffusionImageVariationPipeline(**_A     )
											__lowerCAmelCase							  =			sd_pipe.to(_A     )
											sd_pipe.set_progress_bar_config(disable=_A     )

											__lowerCAmelCase							  =			self.get_dummy_inputs(_A     )
											__lowerCAmelCase							  =			2 * [inputs["image"]]
											__lowerCAmelCase							  =			sd_pipe(**_A     )

											__lowerCAmelCase							  =			output.images

											__lowerCAmelCase							  =			image[-1, -3:, -3:, -1]

											assert image.shape == (2, 6_4, 6_4, 3)
											__lowerCAmelCase							  =			np.array([0.68_92, 0.56_37, 0.58_36, 0.57_71, 0.62_54, 0.64_09, 0.55_80, 0.55_69, 0.52_89]     )

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






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



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






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






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



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






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A="cpu"						,     _A=torch.floataa						,     _A=0     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			torch.Generator(device=_A     ).manual_seed(_A     )
											__lowerCAmelCase							  =			load_image(
											    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
											    "/stable_diffusion_imgvar/input_image_vermeer.png"     )
											__lowerCAmelCase							  =			np.random.RandomState(_A     ).standard_normal((1, 4, 6_4, 6_4)     )
											__lowerCAmelCase							  =			torch.from_numpy(_A     ).to(device=_A						,     dtype=_A     )
											__lowerCAmelCase							  =			{
											    "image": init_image,
											    "latents": latents,
											    "generator": generator,
											    "num_inference_steps": 3,
											    "guidance_scale": 7.5,
											    "output_type": "numpy",
											}
											return inputs






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			StableDiffusionImageVariationPipeline.from_pretrained(
											    "lambdalabs/sd-image-variations-diffusers"						,     safety_checker=_A     )
											__lowerCAmelCase							  =			sd_pipe.to(_A     )
											sd_pipe.set_progress_bar_config(disable=_A     )

											__lowerCAmelCase							  =			self.get_inputs(_A     )
											__lowerCAmelCase							  =			sd_pipe(**_A     ).images
											__lowerCAmelCase							  =			image[0, -3:, -3:, -1].flatten()

											assert image.shape == (1, 5_1_2, 5_1_2, 3)
											__lowerCAmelCase							  =			np.array([0.8_44_91, 0.9_07_89, 0.7_57_08, 0.7_87_34, 0.8_34_85, 0.7_00_99, 0.6_69_38, 0.6_87_27, 0.6_13_79]     )
											assert np.abs(image_slice - expected_slice     ).max() < 6E-3






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			0

											def callback_fn(_A						,     _A						,     _A     ) -> None:
																__lowerCAmelCase							  =			True
																nonlocal number_of_steps
																number_of_steps += 1
																if step == 1:
																					__lowerCAmelCase							  =			latents.detach().cpu().numpy()
																					assert latents.shape == (1, 4, 6_4, 6_4)
																					__lowerCAmelCase							  =			latents[0, -3:, -3:, -1]
																					__lowerCAmelCase							  =			np.array(
																					    [-0.16_21, 0.28_37, -0.79_79, -0.12_21, -1.30_57, 0.76_81, -2.11_91, 0.04_64, 1.63_09]     )

																					assert np.abs(latents_slice.flatten() - expected_slice     ).max() < 5E-2
																elif step == 2:
																					__lowerCAmelCase							  =			latents.detach().cpu().numpy()
																					assert latents.shape == (1, 4, 6_4, 6_4)
																					__lowerCAmelCase							  =			latents[0, -3:, -3:, -1]
																					__lowerCAmelCase							  =			np.array([0.62_99, 1.75_00, 1.19_92, -2.15_82, -1.89_94, 0.73_34, -0.70_90, 1.01_37, 1.52_73]     )

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

											__lowerCAmelCase							  =			False

											__lowerCAmelCase							  =			StableDiffusionImageVariationPipeline.from_pretrained(
											    "fusing/sd-image-variations-diffusers"						,     safety_checker=_A						,     torch_dtype=torch.floataa						,     )
											pipe.to(_A     )
											pipe.set_progress_bar_config(disable=_A     )
											pipe.enable_attention_slicing()

											__lowerCAmelCase							  =			self.get_inputs(_A						,     dtype=torch.floataa     )
											pipe(**_A						,     callback=_A						,     callback_steps=1     )
											assert callback_fn.has_been_called
											assert number_of_steps == inputs["num_inference_steps"]






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



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

											__lowerCAmelCase							  =			"fusing/sd-image-variations-diffusers"
											__lowerCAmelCase							  =			StableDiffusionImageVariationPipeline.from_pretrained(
											    _A						,     safety_checker=_A						,     torch_dtype=torch.floataa     )
											__lowerCAmelCase							  =			pipe.to(_A     )
											pipe.set_progress_bar_config(disable=_A     )
											pipe.enable_attention_slicing(1     )
											pipe.enable_sequential_cpu_offload()

											__lowerCAmelCase							  =			self.get_inputs(_A						,     dtype=torch.floataa     )
											__lowerCAmelCase							  =			pipe(**_A     )

											__lowerCAmelCase							  =			torch.cuda.max_memory_allocated()
											# make sure that less than 2.6 GB is allocated
											assert mem_bytes < 2.6 * 1_0**9






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






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



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






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A="cpu"						,     _A=torch.floataa						,     _A=0     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			torch.Generator(device=_A     ).manual_seed(_A     )
											__lowerCAmelCase							  =			load_image(
											    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
											    "/stable_diffusion_imgvar/input_image_vermeer.png"     )
											__lowerCAmelCase							  =			np.random.RandomState(_A     ).standard_normal((1, 4, 6_4, 6_4)     )
											__lowerCAmelCase							  =			torch.from_numpy(_A     ).to(device=_A						,     dtype=_A     )
											__lowerCAmelCase							  =			{
											    "image": init_image,
											    "latents": latents,
											    "generator": generator,
											    "num_inference_steps": 5_0,
											    "guidance_scale": 7.5,
											    "output_type": "numpy",
											}
											return inputs






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			StableDiffusionImageVariationPipeline.from_pretrained("fusing/sd-image-variations-diffusers"     )
											sd_pipe.to(_A     )
											sd_pipe.set_progress_bar_config(disable=_A     )

											__lowerCAmelCase							  =			self.get_inputs(_A     )
											__lowerCAmelCase							  =			sd_pipe(**_A     ).images[0]

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






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											__lowerCAmelCase							  =			StableDiffusionImageVariationPipeline.from_pretrained("fusing/sd-image-variations-diffusers"     )
											__lowerCAmelCase							  =			DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config     )
											sd_pipe.to(_A     )
											sd_pipe.set_progress_bar_config(disable=_A     )

											__lowerCAmelCase							  =			self.get_inputs(_A     )
											__lowerCAmelCase							  =			2_5
											__lowerCAmelCase							  =			sd_pipe(**_A     ).images[0]

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










from typing import Any, Callable, Dict, List, Optional, Union

import PIL.Image
import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DiffusionPipeline,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionImgaImgPipeline,
    StableDiffusionInpaintPipelineLegacy,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.configuration_utils import FrozenDict
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.utils import deprecate, logging


UpperCamelCase__            =     logging.get_logger(__name__)  # pylint: disable=invalid-name
class      a__       (   snake_case__ ):
						_a				:       Tuple   =       ["""safety_checker""", """feature_extractor"""]






						def __init__(  self						,     _A						,     _A						,     _A						,     _A						,     _A						,     _A						,     _A						,     _A = True						,     ):




											"""simple docstring"""



											super().__init__()
											if hasattr(scheduler.config						,     "steps_offset"     ) and scheduler.config.steps_offset != 1:
																__lowerCAmelCase							  =			(
																    f"""The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"""
																    f""" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure """
																    "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
																    " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
																    " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
																    " file"
																)
																deprecate("steps_offset!=1"						,     "1.0.0"						,     _A						,     standard_warn=_A     )
																__lowerCAmelCase							  =			dict(scheduler.config     )
																__lowerCAmelCase							  =			1
																__lowerCAmelCase							  =			FrozenDict(_A     )

											self.register_modules(
											    vae=_A						,     text_encoder=_A						,     tokenizer=_A						,     unet=_A						,     scheduler=_A						,     safety_checker=_A						,     feature_extractor=_A						,     )
											self.register_to_config(requires_safety_checker=_A     )






						@property
						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											return {k: getattr(self						,     _A     ) for k in self.config.keys() if not k.startswith("_"     )}






						def 							__SCREAMING_SNAKE_CASE(  self						,     _A = "auto"     ):




											"""simple docstring"""



											if slice_size == "auto":
																# half the attention head size is usually a good trade-off between
																# speed and memory
																__lowerCAmelCase							  =			self.unet.config.attention_head_dim // 2
											self.unet.set_attention_slice(_A     )






						def 							__SCREAMING_SNAKE_CASE(  self     ):




											"""simple docstring"""



											self.enable_attention_slicing(_A     )






						@torch.no_grad()
						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A						,     _A						,     _A = 0.8						,     _A = 5_0						,     _A = 7.5						,     _A = None						,     _A = 1						,     _A = 0.0						,     _A = None						,     _A = "pil"						,     _A = True						,     _A = None						,     _A = 1						,     ):




											"""simple docstring"""



											return StableDiffusionInpaintPipelineLegacy(**self.components     )(
											    prompt=_A						,     image=_A						,     mask_image=_A						,     strength=_A						,     num_inference_steps=_A						,     guidance_scale=_A						,     negative_prompt=_A						,     num_images_per_prompt=_A						,     eta=_A						,     generator=_A						,     output_type=_A						,     return_dict=_A						,     callback=_A						,     )






						@torch.no_grad()
						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A						,     _A = 0.8						,     _A = 5_0						,     _A = 7.5						,     _A = None						,     _A = 1						,     _A = 0.0						,     _A = None						,     _A = "pil"						,     _A = True						,     _A = None						,     _A = 1						,     **_A						,     ):




											"""simple docstring"""



											return StableDiffusionImgaImgPipeline(**self.components     )(
											    prompt=_A						,     image=_A						,     strength=_A						,     num_inference_steps=_A						,     guidance_scale=_A						,     negative_prompt=_A						,     num_images_per_prompt=_A						,     eta=_A						,     generator=_A						,     output_type=_A						,     return_dict=_A						,     callback=_A						,     callback_steps=_A						,     )






						@torch.no_grad()
						def 							__SCREAMING_SNAKE_CASE(  self						,     _A						,     _A = 5_1_2						,     _A = 5_1_2						,     _A = 5_0						,     _A = 7.5						,     _A = None						,     _A = 1						,     _A = 0.0						,     _A = None						,     _A = None						,     _A = "pil"						,     _A = True						,     _A = None						,     _A = 1						,     ):




											"""simple docstring"""



											return StableDiffusionPipeline(**self.components     )(
											    prompt=_A						,     height=_A						,     width=_A						,     num_inference_steps=_A						,     guidance_scale=_A						,     negative_prompt=_A						,     num_images_per_prompt=_A						,     eta=_A						,     generator=_A						,     latents=_A						,     output_type=_A						,     return_dict=_A						,     callback=_A						,     callback_steps=_A						,     )




