



"""simple docstring"""
import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest
class 				lowerCamelCase  (			_lowerCAmelCase				):






					'''simple docstring'''





					_A		:			Tuple						=     (UniPCMultistepScheduler,)
					_A		:			Any						=     (("""num_inference_steps""", 2_5),)






					def    lowerCAmelCase_				(       self:							Any						,     **snake_case:							Optional[Any]						)   ->	Optional[int]:
											snake_case_ :Tuple						=			{
											    """num_train_timesteps""": 1_000,
											    """beta_start""": 0.0_0_0_1,
											    """beta_end""": 0.0_2,
											    """beta_schedule""": """linear""",
											    """solver_order""": 2,
											    """solver_type""": """bh2""",
											}

											config.update(**snake_case						)
											return config






					def    lowerCAmelCase_				(       self:							List[Any]						,     snake_case:							Dict=0						,     **snake_case:							List[str]						)   ->	Union[str, Any]:
											snake_case_ :Any						=			dict(self.forward_default_kwargs						)
											snake_case_ :str						=			kwargs.pop("""num_inference_steps"""						,     snake_case						)
											snake_case_ :Optional[int]						=			self.dummy_sample
											snake_case_ :Tuple						=			0.1 * sample
											snake_case_ :str						=			[residual + 0.2, residual + 0.1_5, residual + 0.1_0]

											for scheduler_class in self.scheduler_classes:
																	snake_case_ :Union[str, Any]						=			self.get_scheduler_config(**snake_case						)
																	snake_case_ :Tuple						=			scheduler_class(**snake_case						)
																	scheduler.set_timesteps(snake_case						)
																	# copy over dummy past residuals
																	snake_case_ :Optional[int]						=			dummy_past_residuals[: scheduler.config.solver_order]

																	with tempfile.TemporaryDirectory() as tmpdirname:
																							scheduler.save_config(snake_case						)
																							snake_case_ :int						=			scheduler_class.from_pretrained(snake_case						)
																							new_scheduler.set_timesteps(snake_case						)
																							# copy over dummy past residuals
																							snake_case_ :Union[str, Any]						=			dummy_past_residuals[: new_scheduler.config.solver_order]

																	snake_case_,       snake_case_ :Union[str, Any]						=			sample, sample
																	for t in range(snake_case						,     time_step + scheduler.config.solver_order + 1						):
																							snake_case_ :Any						=			scheduler.step(snake_case						,     snake_case						,     snake_case						,     **snake_case						).prev_sample
																							snake_case_ :Dict						=			new_scheduler.step(snake_case						,     snake_case						,     snake_case						,     **snake_case						).prev_sample

																							assert torch.sum(torch.abs(output - new_output						)						) < 1E-5, "Scheduler outputs are not identical"






					def    lowerCAmelCase_				(       self:							Tuple						,     snake_case:							List[Any]=0						,     **snake_case:							Union[str, Any]						)   ->	Union[str, Any]:
											snake_case_ :Dict						=			dict(self.forward_default_kwargs						)
											snake_case_ :List[Any]						=			kwargs.pop("""num_inference_steps"""						,     snake_case						)
											snake_case_ :Optional[int]						=			self.dummy_sample
											snake_case_ :Any						=			0.1 * sample
											snake_case_ :Union[str, Any]						=			[residual + 0.2, residual + 0.1_5, residual + 0.1_0]

											for scheduler_class in self.scheduler_classes:
																	snake_case_ :Tuple						=			self.get_scheduler_config()
																	snake_case_ :Dict						=			scheduler_class(**snake_case						)
																	scheduler.set_timesteps(snake_case						)

																	# copy over dummy past residuals (must be after setting timesteps)
																	snake_case_ :Dict						=			dummy_past_residuals[: scheduler.config.solver_order]

																	with tempfile.TemporaryDirectory() as tmpdirname:
																							scheduler.save_config(snake_case						)
																							snake_case_ :Tuple						=			scheduler_class.from_pretrained(snake_case						)
																							# copy over dummy past residuals
																							new_scheduler.set_timesteps(snake_case						)

																							# copy over dummy past residual (must be after setting timesteps)
																							snake_case_ :List[Any]						=			dummy_past_residuals[: new_scheduler.config.solver_order]

																	snake_case_ :List[Any]						=			scheduler.step(snake_case						,     snake_case						,     snake_case						,     **snake_case						).prev_sample
																	snake_case_ :List[Any]						=			new_scheduler.step(snake_case						,     snake_case						,     snake_case						,     **snake_case						).prev_sample

																	assert torch.sum(torch.abs(output - new_output						)						) < 1E-5, "Scheduler outputs are not identical"






					def    lowerCAmelCase_				(       self:							Any						,     snake_case:							List[str]=None						,     **snake_case:							int						)   ->	Any:
											if scheduler is None:
																	snake_case_ :Optional[int]						=			self.scheduler_classes[0]
																	snake_case_ :Optional[int]						=			self.get_scheduler_config(**snake_case						)
																	snake_case_ :int						=			scheduler_class(**snake_case						)

											snake_case_ :Optional[int]						=			self.scheduler_classes[0]
											snake_case_ :List[str]						=			self.get_scheduler_config(**snake_case						)
											snake_case_ :Optional[int]						=			scheduler_class(**snake_case						)

											snake_case_ :Dict						=			10
											snake_case_ :List[str]						=			self.dummy_model()
											snake_case_ :str						=			self.dummy_sample_deter
											scheduler.set_timesteps(snake_case						)

											for i, t in enumerate(scheduler.timesteps						):
																	snake_case_ :Dict						=			model(snake_case						,     snake_case						)
																	snake_case_ :int						=			scheduler.step(snake_case						,     snake_case						,     snake_case						).prev_sample

											return sample






					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	Optional[int]:
											snake_case_ :Dict						=			dict(self.forward_default_kwargs						)

											snake_case_ :List[Any]						=			kwargs.pop("""num_inference_steps"""						,     snake_case						)

											for scheduler_class in self.scheduler_classes:
																	snake_case_ :Union[str, Any]						=			self.get_scheduler_config()
																	snake_case_ :List[str]						=			scheduler_class(**snake_case						)

																	snake_case_ :int						=			self.dummy_sample
																	snake_case_ :Tuple						=			0.1 * sample

																	if num_inference_steps is not None and hasattr(snake_case						,     """set_timesteps"""						):
																							scheduler.set_timesteps(snake_case						)
																	elif num_inference_steps is not None and not hasattr(snake_case						,     """set_timesteps"""						):
																							snake_case_ :Optional[int]						=			num_inference_steps

																	# copy over dummy past residuals (must be done after set_timesteps)
																	snake_case_ :Optional[int]						=			[residual + 0.2, residual + 0.1_5, residual + 0.1_0]
																	snake_case_ :int						=			dummy_past_residuals[: scheduler.config.solver_order]

																	snake_case_ :int						=			scheduler.timesteps[5]
																	snake_case_ :List[str]						=			scheduler.timesteps[6]

																	snake_case_ :Dict						=			scheduler.step(snake_case						,     snake_case						,     snake_case						,     **snake_case						).prev_sample
																	snake_case_ :Any						=			scheduler.step(snake_case						,     snake_case						,     snake_case						,     **snake_case						).prev_sample

																	self.assertEqual(output_a.shape						,     sample.shape						)
																	self.assertEqual(output_a.shape						,     output_a.shape						)






					def    lowerCAmelCase_				(       self:							int						)   ->	List[Any]:
											# make sure that iterating over schedulers with same config names gives same results
											# for defaults
											snake_case_ :str						=			UniPCMultistepScheduler(**self.get_scheduler_config()						)
											snake_case_ :List[Any]						=			self.full_loop(scheduler=snake_case						)
											snake_case_ :Optional[int]						=			torch.mean(torch.abs(snake_case						)						)

											assert abs(result_mean.item() - 0.2_4_6_4						) < 1E-3

											snake_case_ :Optional[int]						=			DPMSolverSinglestepScheduler.from_config(scheduler.config						)
											snake_case_ :Union[str, Any]						=			DEISMultistepScheduler.from_config(scheduler.config						)
											snake_case_ :int						=			DPMSolverMultistepScheduler.from_config(scheduler.config						)
											snake_case_ :Optional[Any]						=			UniPCMultistepScheduler.from_config(scheduler.config						)

											snake_case_ :str						=			self.full_loop(scheduler=snake_case						)
											snake_case_ :Dict						=			torch.mean(torch.abs(snake_case						)						)

											assert abs(result_mean.item() - 0.2_4_6_4						) < 1E-3






					def    lowerCAmelCase_				(       self:							Any						)   ->	List[str]:
											for timesteps in [25, 50, 100, 999, 1_000]:
																	self.check_over_configs(num_train_timesteps=snake_case						)






					def    lowerCAmelCase_				(       self:							Tuple						)   ->	Tuple:
											self.check_over_configs(thresholding=snake_case						)
											for order in [1, 2, 3]:
																	for solver_type in ["bh1", "bh2"]:
																							for threshold in [0.5, 1.0, 2.0]:
																													for prediction_type in ["epsilon", "sample"]:
																																			self.check_over_configs(
																																			    thresholding=snake_case						,     prediction_type=snake_case						,     sample_max_value=snake_case						,     solver_order=snake_case						,     solver_type=snake_case						,     )






					def    lowerCAmelCase_				(       self:							Dict						)   ->	List[str]:
											for prediction_type in ["epsilon", "v_prediction"]:
																	self.check_over_configs(prediction_type=snake_case						)






					def    lowerCAmelCase_				(       self:							str						)   ->	Optional[int]:
											for solver_type in ["bh1", "bh2"]:
																	for order in [1, 2, 3]:
																							for prediction_type in ["epsilon", "sample"]:
																													self.check_over_configs(
																													    solver_order=snake_case						,     solver_type=snake_case						,     prediction_type=snake_case						,     )
																													snake_case_ :str						=			self.full_loop(
																													    solver_order=snake_case						,     solver_type=snake_case						,     prediction_type=snake_case						,     )
																													assert not torch.isnan(snake_case						).any(), "Samples have nan numbers"






					def    lowerCAmelCase_				(       self:							List[Any]						)   ->	Union[str, Any]:
											self.check_over_configs(lower_order_final=snake_case						)
											self.check_over_configs(lower_order_final=snake_case						)






					def    lowerCAmelCase_				(       self:							List[str]						)   ->	int:
											for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1_000]:
																	self.check_over_forward(num_inference_steps=snake_case						,     time_step=0						)






					def    lowerCAmelCase_				(       self:							str						)   ->	Optional[int]:
											snake_case_ :Tuple						=			self.full_loop()
											snake_case_ :List[Any]						=			torch.mean(torch.abs(snake_case						)						)

											assert abs(result_mean.item() - 0.2_4_6_4						) < 1E-3






					def    lowerCAmelCase_				(       self:							Tuple						)   ->	Optional[int]:
											snake_case_ :Dict						=			self.full_loop(prediction_type="""v_prediction"""						)
											snake_case_ :Any						=			torch.mean(torch.abs(snake_case						)						)

											assert abs(result_mean.item() - 0.1_0_1_4						) < 1E-3






					def    lowerCAmelCase_				(       self:							Any						)   ->	List[str]:
											snake_case_ :Dict						=			self.scheduler_classes[0]
											snake_case_ :Tuple						=			self.get_scheduler_config(thresholding=snake_case						,     dynamic_thresholding_ratio=0						)
											snake_case_ :List[Any]						=			scheduler_class(**snake_case						)

											snake_case_ :Any						=			10
											snake_case_ :int						=			self.dummy_model()
											snake_case_ :Union[str, Any]						=			self.dummy_sample_deter.half()
											scheduler.set_timesteps(snake_case						)

											for i, t in enumerate(scheduler.timesteps						):
																	snake_case_ :int						=			model(snake_case						,     snake_case						)
																	snake_case_ :List[Any]						=			scheduler.step(snake_case						,     snake_case						,     snake_case						).prev_sample

											assert sample.dtype == torch.floataa






					def    lowerCAmelCase_				(       self:							Tuple						,     **snake_case:							List[Any]						)   ->	Dict:
											for scheduler_class in self.scheduler_classes:
																	snake_case_ :int						=			self.get_scheduler_config(**snake_case						)
																	snake_case_ :Tuple						=			scheduler_class(**snake_case						)

																	scheduler.set_timesteps(scheduler.config.num_train_timesteps						)
																	assert len(scheduler.timesteps.unique()						) == scheduler.num_inference_steps





"""simple docstring"""
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class 				lowerCamelCase  (			_lowerCAmelCase    ,						_lowerCAmelCase    ,						unittest.TestCase				):






					'''simple docstring'''





					_A		:			str						=     StableDiffusionSAGPipeline
					_A		:			Optional[Any]						=     TEXT_TO_IMAGE_PARAMS
					_A		:			Any						=     TEXT_TO_IMAGE_BATCH_PARAMS
					_A		:			Tuple						=     TEXT_TO_IMAGE_IMAGE_PARAMS
					_A		:			Tuple						=     TEXT_TO_IMAGE_IMAGE_PARAMS
					_A		:			List[str]						=     False






					def    lowerCAmelCase_				(       self:							Optional[Any]						)   ->	str:
											torch.manual_seed(0						)
											snake_case_ :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						,     )
											snake_case_ :Any						=			DDIMScheduler(
											    beta_start=0.0_0_0_8_5						,     beta_end=0.0_1_2						,     beta_schedule="""scaled_linear"""						,     clip_sample=snake_case						,     set_alpha_to_one=snake_case						,     )
											torch.manual_seed(0						)
											snake_case_ :Optional[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						)
											snake_case_ :Union[str, Any]						=			CLIPTextConfig(
											    bos_token_id=0						,     eos_token_id=2						,     hidden_size=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						,     )
											snake_case_ :Tuple						=			CLIPTextModel(snake_case						)
											snake_case_ :str						=			CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

											snake_case_ :Dict						=			{
											    """unet""": unet,
											    """scheduler""": scheduler,
											    """vae""": vae,
											    """text_encoder""": text_encoder,
											    """tokenizer""": tokenizer,
											    """safety_checker""": None,
											    """feature_extractor""": None,
											}
											return components






					def    lowerCAmelCase_				(       self:							List[str]						,     snake_case:							Tuple						,     snake_case:							List[str]=0						)   ->	str:
											if str(snake_case						).startswith("""mps"""						):
																	snake_case_ :Tuple						=			torch.manual_seed(snake_case						)
											else:
																	snake_case_ :Optional[int]						=			torch.Generator(device=snake_case						).manual_seed(snake_case						)
											snake_case_ :Any						=			{
											    """prompt""": """.""",
											    """generator""": generator,
											    """num_inference_steps""": 2,
											    """guidance_scale""": 1.0,
											    """sag_scale""": 1.0,
											    """output_type""": """numpy""",
											}
											return inputs






					def    lowerCAmelCase_				(       self:							Optional[int]						)   ->	str:
											super().test_inference_batch_single_identical(expected_max_diff=3E-3						)

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






					'''simple docstring'''





					def    lowerCAmelCase_				(       self:							int						)   ->	str:
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()






					def    lowerCAmelCase_				(       self:							int						)   ->	List[str]:
											snake_case_ :Any						=			StableDiffusionSAGPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""						)
											snake_case_ :int						=			sag_pipe.to(snake_case						)
											sag_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_ :Union[str, Any]						=			"""."""
											snake_case_ :str						=			torch.manual_seed(0						)
											snake_case_ :str						=			sag_pipe(
											    [prompt]						,     generator=snake_case						,     guidance_scale=7.5						,     sag_scale=1.0						,     num_inference_steps=20						,     output_type="""np"""						)

											snake_case_ :List[Any]						=			output.images

											snake_case_ :Tuple						=			image[0, -3:, -3:, -1]

											assert image.shape == (1, 512, 512, 3)
											snake_case_ :List[Any]						=			np.array([0.1_5_6_8, 0.1_7_3_8, 0.1_6_9_5, 0.1_6_9_3, 0.1_5_0_7, 0.1_7_0_5, 0.1_5_4_7, 0.1_7_5_1, 0.1_9_4_9]						)

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






					def    lowerCAmelCase_				(       self:							Dict						)   ->	str:
											snake_case_ :Tuple						=			StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""						)
											snake_case_ :Optional[int]						=			sag_pipe.to(snake_case						)
											sag_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_ :Tuple						=			"""."""
											snake_case_ :Union[str, Any]						=			torch.manual_seed(0						)
											snake_case_ :Tuple						=			sag_pipe(
											    [prompt]						,     generator=snake_case						,     guidance_scale=7.5						,     sag_scale=1.0						,     num_inference_steps=20						,     output_type="""np"""						)

											snake_case_ :Optional[int]						=			output.images

											snake_case_ :Optional[Any]						=			image[0, -3:, -3:, -1]

											assert image.shape == (1, 512, 512, 3)
											snake_case_ :Tuple						=			np.array([0.3_4_5_9, 0.2_8_7_6, 0.2_5_3_7, 0.3_0_0_2, 0.2_6_7_1, 0.2_1_6_0, 0.3_0_2_6, 0.2_2_6_2, 0.2_3_7_1]						)

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






					def    lowerCAmelCase_				(       self:							List[str]						)   ->	List[str]:
											snake_case_ :Optional[int]						=			StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base"""						)
											snake_case_ :int						=			sag_pipe.to(snake_case						)
											sag_pipe.set_progress_bar_config(disable=snake_case						)

											snake_case_ :Tuple						=			"""."""
											snake_case_ :Optional[int]						=			torch.manual_seed(0						)
											snake_case_ :List[str]						=			sag_pipe(
											    [prompt]						,     width=768						,     height=512						,     generator=snake_case						,     guidance_scale=7.5						,     sag_scale=1.0						,     num_inference_steps=20						,     output_type="""np"""						,     )

											snake_case_ :Optional[Any]						=			output.images

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