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__						( A__   ,		A__   ,		A__   ,		unittest.TestCase ):
					A     	=							StableDiffusionImageVariationPipeline
					A     	=							IMAGE_VARIATION_PARAMS
					A     	=							IMAGE_VARIATION_BATCH_PARAMS
					A     	=							frozenset([] )
					# TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
					A     	=							frozenset([] )

					def     __UpperCamelCase       (   self				:	Tuple   ):







							"""simple docstring"""






							torch.manual_seed(0   )
							SCREAMING_SNAKE_CASE_  :		Optional[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,)
							SCREAMING_SNAKE_CASE_  :		Optional[int]    				=   PNDMScheduler(skip_prk_steps=_A   )
							torch.manual_seed(0   )
							SCREAMING_SNAKE_CASE_  :		Any    				=   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   )
							SCREAMING_SNAKE_CASE_  :		Tuple    				=   CLIPVisionConfig(
							    hidden_size=32,projection_dim=32,intermediate_size=37,layer_norm_eps=1E-05,num_attention_heads=4,num_hidden_layers=5,image_size=32,patch_size=4,)
							SCREAMING_SNAKE_CASE_  :		Any    				=   CLIPVisionModelWithProjection(_A   )
							SCREAMING_SNAKE_CASE_  :		Optional[Any]    				=   CLIPImageProcessor(crop_size=32,size=32   )

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

					def     __UpperCamelCase       (   self				:	Optional[int],_A				:	str,_A				:	Dict=0   ):







							"""simple docstring"""






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

					def     __UpperCamelCase       (   self				:	List[str]   ):







							"""simple docstring"""






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

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

							assert image.shape == (1, 64, 64, 3)
							SCREAMING_SNAKE_CASE_  :		Optional[int]    				=   np.array([0.5239, 0.5723, 0.4796, 0.5049, 0.5550, 0.4685, 0.5329, 0.4891, 0.4921]   )

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

					def     __UpperCamelCase       (   self				:	List[Any]   ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		List[Any]    				=   "cpu"  # ensure determinism for the device-dependent torch.Generator
							SCREAMING_SNAKE_CASE_  :		Union[str, Any]    				=   self.get_dummy_components()
							SCREAMING_SNAKE_CASE_  :		int    				=   StableDiffusionImageVariationPipeline(**_A   )
							SCREAMING_SNAKE_CASE_  :		Optional[int]    				=   sd_pipe.to(_A   )
							sd_pipe.set_progress_bar_config(disable=_A   )

							SCREAMING_SNAKE_CASE_  :		List[str]    				=   self.get_dummy_inputs(_A   )
							SCREAMING_SNAKE_CASE_  :		List[Any]    				=   2 * [inputs["image"]]
							SCREAMING_SNAKE_CASE_  :		Union[str, Any]    				=   sd_pipe(**_A   )

							SCREAMING_SNAKE_CASE_  :		List[Any]    				=   output.images

							SCREAMING_SNAKE_CASE_  :		Dict    				=   image[-1, -3:, -3:, -1]

							assert image.shape == (2, 64, 64, 3)
							SCREAMING_SNAKE_CASE_  :		List[Any]    				=   np.array([0.6892, 0.5637, 0.5836, 0.5771, 0.6254, 0.6409, 0.5580, 0.5569, 0.5289]   )

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






					def     __UpperCamelCase       (   self				:	Any   ):







							"""simple docstring"""






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







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

					def     __UpperCamelCase       (   self				:	List[Any]   ):







							"""simple docstring"""






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

					def     __UpperCamelCase       (   self				:	str,_A				:	str,_A				:	str="cpu",_A				:	Any=torch.floataa,_A				:	Optional[Any]=0   ):







							"""simple docstring"""






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

					def     __UpperCamelCase       (   self				:	List[str]   ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		List[str]    				=   StableDiffusionImageVariationPipeline.from_pretrained(
							    "lambdalabs/sd-image-variations-diffusers",safety_checker=_A   )
							SCREAMING_SNAKE_CASE_  :		Dict    				=   sd_pipe.to(_A   )
							sd_pipe.set_progress_bar_config(disable=_A   )

							SCREAMING_SNAKE_CASE_  :		Dict    				=   self.get_inputs(_A   )
							SCREAMING_SNAKE_CASE_  :		Union[str, Any]    				=   sd_pipe(**_A   ).images
							SCREAMING_SNAKE_CASE_  :		List[Any]    				=   image[0, -3:, -3:, -1].flatten()

							assert image.shape == (1, 512, 512, 3)
							SCREAMING_SNAKE_CASE_  :		Any    				=   np.array([0.84491, 0.90789, 0.75708, 0.78734, 0.83485, 0.70099, 0.66938, 0.68727, 0.61379]   )
							assert np.abs(image_slice - expected_slice   ).max() < 6E-3

					def     __UpperCamelCase       (   self				:	str   ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		Dict    				=   0

							def callback_fn(_A				:	int,_A				:	int,_A				:	torch.FloatTensor   ) -> None:
									SCREAMING_SNAKE_CASE_  :		List[str]    				=   True
									nonlocal number_of_steps
									number_of_steps += 1
									if step == 1:
											SCREAMING_SNAKE_CASE_  :		Tuple    				=   latents.detach().cpu().numpy()
											assert latents.shape == (1, 4, 64, 64)
											SCREAMING_SNAKE_CASE_  :		List[Any]    				=   latents[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE_  :		int    				=   np.array(
											    [-0.1621, 0.2837, -0.7979, -0.1221, -1.3057, 0.7681, -2.1191, 0.0464, 1.6309]   )

											assert np.abs(latents_slice.flatten() - expected_slice   ).max() < 5E-2
									elif step == 2:
											SCREAMING_SNAKE_CASE_  :		Any    				=   latents.detach().cpu().numpy()
											assert latents.shape == (1, 4, 64, 64)
											SCREAMING_SNAKE_CASE_  :		Optional[int]    				=   latents[0, -3:, -3:, -1]
											SCREAMING_SNAKE_CASE_  :		Tuple    				=   np.array([0.6299, 1.7500, 1.1992, -2.1582, -1.8994, 0.7334, -0.7090, 1.0137, 1.5273]   )

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

							SCREAMING_SNAKE_CASE_  :		Tuple    				=   False

							SCREAMING_SNAKE_CASE_  :		Union[str, Any]    				=   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()

							SCREAMING_SNAKE_CASE_  :		Optional[int]    				=   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     __UpperCamelCase       (   self				:	List[str]   ):







							"""simple docstring"""






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

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

							SCREAMING_SNAKE_CASE_  :		List[str]    				=   self.get_inputs(_A,dtype=torch.floataa   )
							SCREAMING_SNAKE_CASE_  :		Optional[Any]    				=   pipe(**_A   )

							SCREAMING_SNAKE_CASE_  :		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     __UpperCamelCase       (   self				:	Optional[int]   ):







							"""simple docstring"""






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

					def     __UpperCamelCase       (   self				:	Tuple,_A				:	int,_A				:	Tuple="cpu",_A				:	Dict=torch.floataa,_A				:	Dict=0   ):







							"""simple docstring"""






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

					def     __UpperCamelCase       (   self				:	Dict   ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		Optional[Any]    				=   StableDiffusionImageVariationPipeline.from_pretrained("fusing/sd-image-variations-diffusers"   )
							sd_pipe.to(_A   )
							sd_pipe.set_progress_bar_config(disable=_A   )

							SCREAMING_SNAKE_CASE_  :		Union[str, Any]    				=   self.get_inputs(_A   )
							SCREAMING_SNAKE_CASE_  :		List[Any]    				=   sd_pipe(**_A   ).images[0]

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






					def     __UpperCamelCase       (   self				:	Optional[int]   ):







							"""simple docstring"""






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

							SCREAMING_SNAKE_CASE_  :		List[str]    				=   self.get_inputs(_A   )
							SCREAMING_SNAKE_CASE_  :		Dict    				=   25
							SCREAMING_SNAKE_CASE_  :		str    				=   sd_pipe(**_A   ).images[0]

							SCREAMING_SNAKE_CASE_  :		Optional[int]    				=   load_numpy(
							    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"
							    "/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy"   )
							SCREAMING_SNAKE_CASE_  :		Union[str, Any]    				=   np.abs(expected_image - image   ).max()
							assert max_diff < 1E-3

import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionDualGuidedPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


__lowerCamelCase :     List[Any]          =  False






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

					def     __UpperCamelCase       (   self				:	Optional[int]   ):







							"""simple docstring"""






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

					def     __UpperCamelCase       (   self				:	Optional[Any]   ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		int    				=   VersatileDiffusionDualGuidedPipeline.from_pretrained("shi-labs/versatile-diffusion"   )
							# remove text_unet
							pipe.remove_unused_weights()
							pipe.to(_A   )
							pipe.set_progress_bar_config(disable=_A   )

							SCREAMING_SNAKE_CASE_  :		Optional[int]    				=   load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg"   )

							SCREAMING_SNAKE_CASE_  :		Optional[Any]    				=   torch.manual_seed(0   )
							SCREAMING_SNAKE_CASE_  :		int    				=   pipe(
							    prompt="first prompt",image=_A,text_to_image_strength=0.75,generator=_A,guidance_scale=7.5,num_inference_steps=2,output_type="numpy",).images

							with tempfile.TemporaryDirectory() as tmpdirname:
									pipe.save_pretrained(_A   )
									SCREAMING_SNAKE_CASE_  :		Dict    				=   VersatileDiffusionDualGuidedPipeline.from_pretrained(_A   )

							pipe.to(_A   )
							pipe.set_progress_bar_config(disable=_A   )

							SCREAMING_SNAKE_CASE_  :		List[str]    				=   generator.manual_seed(0   )
							SCREAMING_SNAKE_CASE_  :		str    				=   pipe(
							    prompt="first prompt",image=_A,text_to_image_strength=0.75,generator=_A,guidance_scale=7.5,num_inference_steps=2,output_type="numpy",).images

							assert np.abs(image - new_image   ).sum() < 1E-5, "Models don't have the same forward pass"






					def     __UpperCamelCase       (   self				:	Any   ):







							"""simple docstring"""






							SCREAMING_SNAKE_CASE_  :		Any    				=   VersatileDiffusionDualGuidedPipeline.from_pretrained("shi-labs/versatile-diffusion"   )
							pipe.remove_unused_weights()
							pipe.to(_A   )
							pipe.set_progress_bar_config(disable=_A   )

							SCREAMING_SNAKE_CASE_  :		Union[str, Any]    				=   "cyberpunk 2077"
							SCREAMING_SNAKE_CASE_  :		Dict    				=   load_image(
							    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg"   )
							SCREAMING_SNAKE_CASE_  :		Tuple    				=   torch.manual_seed(0   )
							SCREAMING_SNAKE_CASE_  :		Dict    				=   pipe(
							    prompt=_A,image=_A,text_to_image_strength=0.75,generator=_A,guidance_scale=7.5,num_inference_steps=50,output_type="numpy",).images

							SCREAMING_SNAKE_CASE_  :		Optional[int]    				=   image[0, 253:256, 253:256, -1]

							assert image.shape == (1, 512, 512, 3)
							SCREAMING_SNAKE_CASE_  :		Optional[Any]    				=   np.array([0.0787, 0.0849, 0.0826, 0.0812, 0.0807, 0.0795, 0.0818, 0.0798, 0.0779]   )

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