






import tempfile

import torch

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

from .test_schedulers import SchedulerCommonTest
class 						lowerCamelCase_				(      lowerCamelCase ):
				a__					       =  (UniPCMultistepScheduler,)
				a__					       =  (('''num_inference_steps''', 25),)







				def 			A	(						self       ,					**__lowerCAmelCase     ):
								"""simple docstring"""
								__magic_name__		:Union[str, Any]        =      {
								    '''num_train_timesteps''': 1_0_0_0,
								    '''beta_start''': 0.0001,
								    '''beta_end''': 0.02,
								    '''beta_schedule''': '''linear''',
								    '''solver_order''': 2,
								    '''solver_type''': '''bh2''',
								}

								config.update(**__lowerCAmelCase     )
								return config







				def 			A	(						self       ,					__lowerCAmelCase=0       ,					**__lowerCAmelCase     ):
								"""simple docstring"""
								__magic_name__		:Dict        =      dict(self.forward_default_kwargs     )
								__magic_name__		:Optional[int]        =      kwargs.pop('''num_inference_steps'''       ,					__lowerCAmelCase     )
								__magic_name__		:Optional[Any]        =      self.dummy_sample
								__magic_name__		:Any        =      0.1 * sample
								__magic_name__		:List[str]        =      [residual + 0.2, residual + 0.15, residual + 0.10]

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

												with tempfile.TemporaryDirectory() as tmpdirname:
																scheduler.save_config(__lowerCAmelCase     )
																__magic_name__		:Optional[Any]        =      scheduler_class.from_pretrained(__lowerCAmelCase     )
																new_scheduler.set_timesteps(__lowerCAmelCase     )
																# copy over dummy past residuals
																__magic_name__		:Any        =      dummy_past_residuals[: new_scheduler.config.solver_order]

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

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







				def 			A	(						self       ,					__lowerCAmelCase=0       ,					**__lowerCAmelCase     ):
								"""simple docstring"""
								__magic_name__		:Dict        =      dict(self.forward_default_kwargs     )
								__magic_name__		:List[Any]        =      kwargs.pop('''num_inference_steps'''       ,					__lowerCAmelCase     )
								__magic_name__		:Tuple        =      self.dummy_sample
								__magic_name__		:Optional[int]        =      0.1 * sample
								__magic_name__		:List[str]        =      [residual + 0.2, residual + 0.15, residual + 0.10]

								for scheduler_class in self.scheduler_classes:
												__magic_name__		:List[str]        =      self.get_scheduler_config()
												__magic_name__		:Optional[int]        =      scheduler_class(**__lowerCAmelCase     )
												scheduler.set_timesteps(__lowerCAmelCase     )

												# copy over dummy past residuals (must be after setting timesteps)
												__magic_name__		:Union[str, Any]        =      dummy_past_residuals[: scheduler.config.solver_order]

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

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

												__magic_name__		:Optional[Any]        =      scheduler.step(__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					**__lowerCAmelCase     ).prev_sample
												__magic_name__		:str        =      new_scheduler.step(__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					**__lowerCAmelCase     ).prev_sample

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







				def 			A	(						self       ,					__lowerCAmelCase=None       ,					**__lowerCAmelCase     ):
								"""simple docstring"""
								if scheduler is None:
												__magic_name__		:Tuple        =      self.scheduler_classes[0]
												__magic_name__		:int        =      self.get_scheduler_config(**__lowerCAmelCase     )
												__magic_name__		:Union[str, Any]        =      scheduler_class(**__lowerCAmelCase     )

								__magic_name__		:int        =      self.scheduler_classes[0]
								__magic_name__		:Dict        =      self.get_scheduler_config(**__lowerCAmelCase     )
								__magic_name__		:Union[str, Any]        =      scheduler_class(**__lowerCAmelCase     )

								__magic_name__		:int        =      1_0
								__magic_name__		:Optional[Any]        =      self.dummy_model()
								__magic_name__		:Tuple        =      self.dummy_sample_deter
								scheduler.set_timesteps(__lowerCAmelCase     )

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

								return sample







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:List[str]        =      dict(self.forward_default_kwargs     )

								__magic_name__		:Optional[Any]        =      kwargs.pop('''num_inference_steps'''       ,					__lowerCAmelCase     )

								for scheduler_class in self.scheduler_classes:
												__magic_name__		:Optional[int]        =      self.get_scheduler_config()
												__magic_name__		:Dict        =      scheduler_class(**__lowerCAmelCase     )

												__magic_name__		:Union[str, Any]        =      self.dummy_sample
												__magic_name__		:Tuple        =      0.1 * sample

												if num_inference_steps is not None and hasattr(__lowerCAmelCase       ,					'''set_timesteps'''     ):
																scheduler.set_timesteps(__lowerCAmelCase     )
												elif num_inference_steps is not None and not hasattr(__lowerCAmelCase       ,					'''set_timesteps'''     ):
																__magic_name__		:Tuple        =      num_inference_steps

												# copy over dummy past residuals (must be done after set_timesteps)
												__magic_name__		:List[Any]        =      [residual + 0.2, residual + 0.15, residual + 0.10]
												__magic_name__		:Tuple        =      dummy_past_residuals[: scheduler.config.solver_order]

												__magic_name__		:str        =      scheduler.timesteps[5]
												__magic_name__		:Dict        =      scheduler.timesteps[6]

												__magic_name__		:List[Any]        =      scheduler.step(__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					**__lowerCAmelCase     ).prev_sample
												__magic_name__		:str        =      scheduler.step(__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase       ,					**__lowerCAmelCase     ).prev_sample

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







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:List[Any]        =      UniPCMultistepScheduler(**self.get_scheduler_config()     )
								__magic_name__		:str        =      self.full_loop(scheduler=__lowerCAmelCase     )
								__magic_name__		:Any        =      torch.mean(torch.abs(__lowerCAmelCase     )     )

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

								__magic_name__		:Dict        =      DPMSolverSinglestepScheduler.from_config(scheduler.config     )
								__magic_name__		:str        =      DEISMultistepScheduler.from_config(scheduler.config     )
								__magic_name__		:int        =      DPMSolverMultistepScheduler.from_config(scheduler.config     )
								__magic_name__		:Optional[int]        =      UniPCMultistepScheduler.from_config(scheduler.config     )

								__magic_name__		:str        =      self.full_loop(scheduler=__lowerCAmelCase     )
								__magic_name__		:Any        =      torch.mean(torch.abs(__lowerCAmelCase     )     )

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







				def 			A	(						self     ):
								"""simple docstring"""
								for timesteps in [2_5, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
												self.check_over_configs(num_train_timesteps=__lowerCAmelCase     )







				def 			A	(						self     ):
								"""simple docstring"""
								self.check_over_configs(thresholding=__lowerCAmelCase     )
								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=__lowerCAmelCase       ,					prediction_type=__lowerCAmelCase       ,					sample_max_value=__lowerCAmelCase       ,					solver_order=__lowerCAmelCase       ,					solver_type=__lowerCAmelCase       ,					)







				def 			A	(						self     ):
								"""simple docstring"""
								for prediction_type in ["epsilon", "v_prediction"]:
												self.check_over_configs(prediction_type=__lowerCAmelCase     )







				def 			A	(						self     ):
								"""simple docstring"""
								for solver_type in ["bh1", "bh2"]:
												for order in [1, 2, 3]:
																for prediction_type in ["epsilon", "sample"]:
																				self.check_over_configs(
																				    solver_order=__lowerCAmelCase       ,					solver_type=__lowerCAmelCase       ,					prediction_type=__lowerCAmelCase       ,					)
																				__magic_name__		:List[Any]        =      self.full_loop(
																				    solver_order=__lowerCAmelCase       ,					solver_type=__lowerCAmelCase       ,					prediction_type=__lowerCAmelCase       ,					)
																				assert not torch.isnan(__lowerCAmelCase     ).any(), "Samples have nan numbers"







				def 			A	(						self     ):
								"""simple docstring"""
								self.check_over_configs(lower_order_final=__lowerCAmelCase     )
								self.check_over_configs(lower_order_final=__lowerCAmelCase     )







				def 			A	(						self     ):
								"""simple docstring"""
								for num_inference_steps in [1, 2, 3, 5, 1_0, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
												self.check_over_forward(num_inference_steps=__lowerCAmelCase       ,					time_step=0     )







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Union[str, Any]        =      self.full_loop()
								__magic_name__		:Any        =      torch.mean(torch.abs(__lowerCAmelCase     )     )

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







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Union[str, Any]        =      self.full_loop(prediction_type='''v_prediction'''     )
								__magic_name__		:Dict        =      torch.mean(torch.abs(__lowerCAmelCase     )     )

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







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Optional[int]        =      self.scheduler_classes[0]
								__magic_name__		:int        =      self.get_scheduler_config(thresholding=__lowerCAmelCase       ,					dynamic_thresholding_ratio=0     )
								__magic_name__		:List[str]        =      scheduler_class(**__lowerCAmelCase     )

								__magic_name__		:Any        =      1_0
								__magic_name__		:Tuple        =      self.dummy_model()
								__magic_name__		:Optional[Any]        =      self.dummy_sample_deter.half()
								scheduler.set_timesteps(__lowerCAmelCase     )

								for i, t in enumerate(scheduler.timesteps     ):
												__magic_name__		:Union[str, Any]        =      model(__lowerCAmelCase       ,					__lowerCAmelCase     )
												__magic_name__		:Tuple        =      scheduler.step(__lowerCAmelCase       ,					__lowerCAmelCase       ,					__lowerCAmelCase     ).prev_sample

								assert sample.dtype == torch.floataa







				def 			A	(						self       ,					**__lowerCAmelCase     ):
								"""simple docstring"""
								for scheduler_class in self.scheduler_classes:
												__magic_name__		:int        =      self.get_scheduler_config(**__lowerCAmelCase     )
												__magic_name__		:str        =      scheduler_class(**__lowerCAmelCase     )

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














import unittest

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

from diffusers import DDIMScheduler, LDMPipeline, UNetaDModel, VQModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()
class 						lowerCamelCase_				(      unittest.TestCase ):







				@property
				def 			A	(						self     ):
								"""simple docstring"""
								torch.manual_seed(0     )
								__magic_name__		:Dict        =      UNetaDModel(
								    block_out_channels=(3_2, 6_4)       ,					layers_per_block=2       ,					sample_size=3_2       ,					in_channels=3       ,					out_channels=3       ,					down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')       ,					up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')       ,					)
								return model







				@property
				def 			A	(						self     ):
								"""simple docstring"""
								torch.manual_seed(0     )
								__magic_name__		:List[str]        =      VQModel(
								    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=3       ,					)
								return model







				@property
				def 			A	(						self     ):
								"""simple docstring"""
								torch.manual_seed(0     )
								__magic_name__		:int        =      CLIPTextConfig(
								    bos_token_id=0       ,					eos_token_id=2       ,					hidden_size=3_2       ,					intermediate_size=3_7       ,					layer_norm_eps=1E-05       ,					num_attention_heads=4       ,					num_hidden_layers=5       ,					pad_token_id=1       ,					vocab_size=1_0_0_0       ,					)
								return CLIPTextModel(__lowerCAmelCase     )







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Optional[Any]        =      self.dummy_uncond_unet
								__magic_name__		:Optional[int]        =      DDIMScheduler()
								__magic_name__		:List[str]        =      self.dummy_vq_model

								__magic_name__		:Tuple        =      LDMPipeline(unet=__lowerCAmelCase       ,					vqvae=__lowerCAmelCase       ,					scheduler=__lowerCAmelCase     )
								ldm.to(__lowerCAmelCase     )
								ldm.set_progress_bar_config(disable=__lowerCAmelCase     )

								__magic_name__		:List[Any]        =      torch.manual_seed(0     )
								__magic_name__		:List[str]        =      ldm(generator=__lowerCAmelCase       ,					num_inference_steps=2       ,					output_type='''numpy'''     ).images

								__magic_name__		:List[Any]        =      torch.manual_seed(0     )
								__magic_name__		:Any        =      ldm(generator=__lowerCAmelCase       ,					num_inference_steps=2       ,					output_type='''numpy'''       ,					return_dict=__lowerCAmelCase     )[0]

								__magic_name__		:Any        =      image[0, -3:, -3:, -1]
								__magic_name__		:Dict        =      image_from_tuple[0, -3:, -3:, -1]

								assert image.shape == (1, 6_4, 6_4, 3)
								__magic_name__		:Union[str, Any]        =      np.array([0.8512, 0.818, 0.6411, 0.6808, 0.4465, 0.5618, 0.46, 0.6231, 0.5172]     )
								__magic_name__		:Any        =      1E-2 if torch_device != '''mps''' else 3E-2

								assert np.abs(image_slice.flatten() - expected_slice     ).max() < tolerance
								assert np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < tolerance







@slow
@require_torch
class 						lowerCamelCase_				(      unittest.TestCase ):







				def 			A	(						self     ):
								"""simple docstring"""
								__magic_name__		:Tuple        =      LDMPipeline.from_pretrained('''CompVis/ldm-celebahq-256'''     )
								ldm.to(__lowerCAmelCase     )
								ldm.set_progress_bar_config(disable=__lowerCAmelCase     )

								__magic_name__		:Optional[Any]        =      torch.manual_seed(0     )
								__magic_name__		:Optional[int]        =      ldm(generator=__lowerCAmelCase       ,					num_inference_steps=5       ,					output_type='''numpy'''     ).images

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

								assert image.shape == (1, 2_5_6, 2_5_6, 3)
								__magic_name__		:List[str]        =      np.array([0.4399, 0.44975, 0.46825, 0.474, 0.4359, 0.4581, 0.45095, 0.4341, 0.4447]     )
								__magic_name__		:Tuple        =      1E-2 if torch_device != '''mps''' else 3E-2

								assert np.abs(image_slice.flatten() - expected_slice     ).max() < tolerance






