
"""simple docstring"""






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






			def 				UpperCAmelCase       (self     :  int )					->  Union[str, Any]:





				"""simple docstring"""


				torch.manual_seed(0 )
				lowerCAmelCase				       =     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			,)
				lowerCAmelCase				       =     PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE_ )
				torch.manual_seed(0 )
				lowerCAmelCase				       =     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 )
				lowerCAmelCase				       =     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			,)
				lowerCAmelCase				       =     CLIPVisionModelWithProjection(SCREAMING_SNAKE_CASE_ )
				lowerCAmelCase				       =     CLIPImageProcessor(crop_size=32			,size=32 )

				lowerCAmelCase				       =     {
				    '''unet''': unet,
				    '''scheduler''': scheduler,
				    '''vae''': vae,
				    '''image_encoder''': image_encoder,
				    '''feature_extractor''': feature_extractor,
				    '''safety_checker''': None,
				}
				return components






			def 				UpperCAmelCase       (self     :  Tuple			,SCREAMING_SNAKE_CASE_     :  Tuple			,SCREAMING_SNAKE_CASE_     :  Union[str, Any]=0 )					->  str:





				"""simple docstring"""


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






			def 				UpperCAmelCase       (self     :  Any )					->  Any:





				"""simple docstring"""


				lowerCAmelCase				       =     '''cpu'''  # ensure determinism for the device-dependent torch.Generator
				lowerCAmelCase				       =     self.get_dummy_components()
				lowerCAmelCase				       =     StableDiffusionImageVariationPipeline(**SCREAMING_SNAKE_CASE_ )
				lowerCAmelCase				       =     sd_pipe.to(SCREAMING_SNAKE_CASE_ )
				sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

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

				assert image.shape == (1, 64, 64, 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 				UpperCAmelCase       (self     :  str )					->  Union[str, Any]:





				"""simple docstring"""


				lowerCAmelCase				       =     '''cpu'''  # ensure determinism for the device-dependent torch.Generator
				lowerCAmelCase				       =     self.get_dummy_components()
				lowerCAmelCase				       =     StableDiffusionImageVariationPipeline(**SCREAMING_SNAKE_CASE_ )
				lowerCAmelCase				       =     sd_pipe.to(SCREAMING_SNAKE_CASE_ )
				sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

				lowerCAmelCase				       =     self.get_dummy_inputs(SCREAMING_SNAKE_CASE_ )
				lowerCAmelCase				       =     2 * [inputs['''image''']]
				lowerCAmelCase				       =     sd_pipe(**SCREAMING_SNAKE_CASE_ )

				lowerCAmelCase				       =     output.images

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

				assert image.shape == (2, 64, 64, 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 				UpperCAmelCase       (self     :  Union[str, Any] )					->  Optional[Any]:





				"""simple docstring"""


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





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






			def 				UpperCAmelCase       (self     :  List[str] )					->  Union[str, Any]:





				"""simple docstring"""


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






			def 				UpperCAmelCase       (self     :  Any			,SCREAMING_SNAKE_CASE_     :  int			,SCREAMING_SNAKE_CASE_     :  Optional[Any]="cpu"			,SCREAMING_SNAKE_CASE_     :  Tuple=torch.floataa			,SCREAMING_SNAKE_CASE_     :  List[Any]=0 )					->  str:





				"""simple docstring"""


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






			def 				UpperCAmelCase       (self     :  Tuple )					->  List[str]:





				"""simple docstring"""


				lowerCAmelCase				       =     StableDiffusionImageVariationPipeline.from_pretrained(
				    '''lambdalabs/sd-image-variations-diffusers'''			,safety_checker=SCREAMING_SNAKE_CASE_ )
				lowerCAmelCase				       =     sd_pipe.to(SCREAMING_SNAKE_CASE_ )
				sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

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

				assert image.shape == (1, 512, 512, 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 				UpperCAmelCase       (self     :  Tuple )					->  List[Any]:





				"""simple docstring"""


				lowerCAmelCase				       =     0

				def callback_fn(SCREAMING_SNAKE_CASE_     :  int			,SCREAMING_SNAKE_CASE_     :  int			,SCREAMING_SNAKE_CASE_     :  torch.FloatTensor ) -> None:
					lowerCAmelCase				       =     True
					nonlocal number_of_steps
					number_of_steps += 1
					if step == 1:
						lowerCAmelCase				       =     latents.detach().cpu().numpy()
						assert latents.shape == (1, 4, 64, 64)
						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, 64, 64)
						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=SCREAMING_SNAKE_CASE_			,torch_dtype=torch.floataa			,)
				pipe.to(SCREAMING_SNAKE_CASE_ )
				pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )
				pipe.enable_attention_slicing()

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






			def 				UpperCAmelCase       (self     :  List[str] )					->  Optional[int]:





				"""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(
				    SCREAMING_SNAKE_CASE_			,safety_checker=SCREAMING_SNAKE_CASE_			,torch_dtype=torch.floataa )
				lowerCAmelCase				       =     pipe.to(SCREAMING_SNAKE_CASE_ )
				pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )
				pipe.enable_attention_slicing(1 )
				pipe.enable_sequential_cpu_offload()

				lowerCAmelCase				       =     self.get_inputs(SCREAMING_SNAKE_CASE_			,dtype=torch.floataa )
				lowerCAmelCase				       =     pipe(**SCREAMING_SNAKE_CASE_ )

				lowerCAmelCase				       =     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  lowercase				(   unittest.TestCase		):






			def 				UpperCAmelCase       (self     :  Optional[Any] )					->  int:





				"""simple docstring"""


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






			def 				UpperCAmelCase       (self     :  int			,SCREAMING_SNAKE_CASE_     :  Union[str, Any]			,SCREAMING_SNAKE_CASE_     :  str="cpu"			,SCREAMING_SNAKE_CASE_     :  Dict=torch.floataa			,SCREAMING_SNAKE_CASE_     :  List[str]=0 )					->  int:





				"""simple docstring"""


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






			def 				UpperCAmelCase       (self     :  int )					->  Tuple:





				"""simple docstring"""


				lowerCAmelCase				       =     StableDiffusionImageVariationPipeline.from_pretrained('''fusing/sd-image-variations-diffusers''' )
				sd_pipe.to(SCREAMING_SNAKE_CASE_ )
				sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

				lowerCAmelCase				       =     self.get_inputs(SCREAMING_SNAKE_CASE_ )
				lowerCAmelCase				       =     sd_pipe(**SCREAMING_SNAKE_CASE_ ).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 				UpperCAmelCase       (self     :  List[Any] )					->  Optional[Any]:





				"""simple docstring"""


				lowerCAmelCase				       =     StableDiffusionImageVariationPipeline.from_pretrained('''fusing/sd-image-variations-diffusers''' )
				lowerCAmelCase				       =     DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config )
				sd_pipe.to(SCREAMING_SNAKE_CASE_ )
				sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_ )

				lowerCAmelCase				       =     self.get_inputs(SCREAMING_SNAKE_CASE_ )
				lowerCAmelCase				       =     25
				lowerCAmelCase				       =     sd_pipe(**SCREAMING_SNAKE_CASE_ ).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




"""simple docstring"""






import torch

from diffusers import KDPMaAncestralDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest



class  lowercase				(   lowercase__		):
			lowercase			     =    (KDPMaAncestralDiscreteScheduler,)
			lowercase			     =    10






			def 				UpperCAmelCase       (self     :  Optional[int]			,**SCREAMING_SNAKE_CASE_     :  Optional[int] )					->  Optional[Any]:





				"""simple docstring"""


				lowerCAmelCase				       =     {
				    '''num_train_timesteps''': 1_100,
				    '''beta_start''': 0.00_01,
				    '''beta_end''': 0.02,
				    '''beta_schedule''': '''linear''',
				}

				config.update(**SCREAMING_SNAKE_CASE_ )
				return config






			def 				UpperCAmelCase       (self     :  Any )					->  str:





				"""simple docstring"""


				for timesteps in [10, 50, 100, 1_000]:
					self.check_over_configs(num_train_timesteps=SCREAMING_SNAKE_CASE_ )






			def 				UpperCAmelCase       (self     :  Optional[Any] )					->  int:





				"""simple docstring"""


				for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]			,[0.00_02, 0.0_02, 0.02] ):
					self.check_over_configs(beta_start=SCREAMING_SNAKE_CASE_			,beta_end=SCREAMING_SNAKE_CASE_ )






			def 				UpperCAmelCase       (self     :  Optional[Any] )					->  Optional[int]:





				"""simple docstring"""


				for schedule in ["linear", "scaled_linear"]:
					self.check_over_configs(beta_schedule=SCREAMING_SNAKE_CASE_ )






			def 				UpperCAmelCase       (self     :  Union[str, Any] )					->  Dict:





				"""simple docstring"""


				if torch_device == "mps":
					return
				lowerCAmelCase				       =     self.scheduler_classes[0]
				lowerCAmelCase				       =     self.get_scheduler_config()
				lowerCAmelCase				       =     scheduler_class(**SCREAMING_SNAKE_CASE_ )

				scheduler.set_timesteps(self.num_inference_steps )

				lowerCAmelCase				       =     torch.manual_seed(0 )

				lowerCAmelCase				       =     self.dummy_model()
				lowerCAmelCase				       =     self.dummy_sample_deter * scheduler.init_noise_sigma
				lowerCAmelCase				       =     sample.to(SCREAMING_SNAKE_CASE_ )

				for i, t in enumerate(scheduler.timesteps ):
					lowerCAmelCase				       =     scheduler.scale_model_input(SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_ )

					lowerCAmelCase				       =     model(SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_ )

					lowerCAmelCase				       =     scheduler.step(SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_			,generator=SCREAMING_SNAKE_CASE_ )
					lowerCAmelCase				       =     output.prev_sample

				lowerCAmelCase				       =     torch.sum(torch.abs(SCREAMING_SNAKE_CASE_ ) )
				lowerCAmelCase				       =     torch.mean(torch.abs(SCREAMING_SNAKE_CASE_ ) )

				assert abs(result_sum.item() - 1_38_49.38_77 ) < 1e-2
				assert abs(result_mean.item() - 18.03_31 ) < 5e-3






			def 				UpperCAmelCase       (self     :  str )					->  List[Any]:





				"""simple docstring"""


				for prediction_type in ["epsilon", "v_prediction"]:
					self.check_over_configs(prediction_type=SCREAMING_SNAKE_CASE_ )






			def 				UpperCAmelCase       (self     :  List[Any] )					->  List[str]:





				"""simple docstring"""


				if torch_device == "mps":
					return
				lowerCAmelCase				       =     self.scheduler_classes[0]
				lowerCAmelCase				       =     self.get_scheduler_config(prediction_type='''v_prediction''' )
				lowerCAmelCase				       =     scheduler_class(**SCREAMING_SNAKE_CASE_ )

				scheduler.set_timesteps(self.num_inference_steps )

				lowerCAmelCase				       =     self.dummy_model()
				lowerCAmelCase				       =     self.dummy_sample_deter * scheduler.init_noise_sigma
				lowerCAmelCase				       =     sample.to(SCREAMING_SNAKE_CASE_ )

				lowerCAmelCase				       =     torch.manual_seed(0 )

				for i, t in enumerate(scheduler.timesteps ):
					lowerCAmelCase				       =     scheduler.scale_model_input(SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_ )

					lowerCAmelCase				       =     model(SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_ )

					lowerCAmelCase				       =     scheduler.step(SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_			,generator=SCREAMING_SNAKE_CASE_ )
					lowerCAmelCase				       =     output.prev_sample

				lowerCAmelCase				       =     torch.sum(torch.abs(SCREAMING_SNAKE_CASE_ ) )
				lowerCAmelCase				       =     torch.mean(torch.abs(SCREAMING_SNAKE_CASE_ ) )

				assert abs(result_sum.item() - 3_28.99_70 ) < 1e-2
				assert abs(result_mean.item() - 0.42_84 ) < 1e-3






			def 				UpperCAmelCase       (self     :  Dict )					->  Union[str, Any]:





				"""simple docstring"""


				if torch_device == "mps":
					return
				lowerCAmelCase				       =     self.scheduler_classes[0]
				lowerCAmelCase				       =     self.get_scheduler_config()
				lowerCAmelCase				       =     scheduler_class(**SCREAMING_SNAKE_CASE_ )

				scheduler.set_timesteps(self.num_inference_steps			,device=SCREAMING_SNAKE_CASE_ )
				lowerCAmelCase				       =     torch.manual_seed(0 )

				lowerCAmelCase				       =     self.dummy_model()
				lowerCAmelCase				       =     self.dummy_sample_deter.to(SCREAMING_SNAKE_CASE_ ) * scheduler.init_noise_sigma

				for t in scheduler.timesteps:
					lowerCAmelCase				       =     scheduler.scale_model_input(SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_ )

					lowerCAmelCase				       =     model(SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_ )

					lowerCAmelCase				       =     scheduler.step(SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_			,generator=SCREAMING_SNAKE_CASE_ )
					lowerCAmelCase				       =     output.prev_sample

				lowerCAmelCase				       =     torch.sum(torch.abs(SCREAMING_SNAKE_CASE_ ) )
				lowerCAmelCase				       =     torch.mean(torch.abs(SCREAMING_SNAKE_CASE_ ) )

				assert abs(result_sum.item() - 1_38_49.38_18 ) < 1e-1
				assert abs(result_mean.item() - 18.03_31 ) < 1e-3


