"""simple docstring"""


import torch

from diffusers import KDPMaAncestralDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest






class        __UpperCAmelCase						(     __lowerCAmelCase     ):
				A__				:       str  							=	(KDPMaAncestralDiscreteScheduler,)
				A__				:       Optional[int]  							=	1_0

				def 		_a     (				self     ,			**_lowerCamelCase				):
										lowerCamelCase__ ={
										    "num_train_timesteps": 1100,
										    "beta_start": 0.0_0_0_1,
										    "beta_end": 0.0_2,
										    "beta_schedule": "linear",
										}

										config.update(**_lowerCamelCase				)
										return config

				def 		_a     (				self				):
										for timesteps in [10, 50, 100, 1000]:
																self.check_over_configs(num_train_timesteps=_lowerCamelCase				)

				def 		_a     (				self				):
										for beta_start, beta_end in zip([0.0_0_0_0_1, 0.0_0_0_1, 0.0_0_1]     ,			[0.0_0_0_2, 0.0_0_2, 0.0_2]				):
																self.check_over_configs(beta_start=_lowerCamelCase     ,			beta_end=_lowerCamelCase				)

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

				def 		_a     (				self				):
										if torch_device == "mps":
																return
										lowerCamelCase__ =self.scheduler_classes[0]
										lowerCamelCase__ =self.get_scheduler_config()
										lowerCamelCase__ =scheduler_class(**_lowerCamelCase				)

										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(_lowerCamelCase				)

										for i, t in enumerate(scheduler.timesteps				):
																lowerCamelCase__ =scheduler.scale_model_input(_lowerCamelCase     ,			_lowerCamelCase				)

																lowerCamelCase__ =model(_lowerCamelCase     ,			_lowerCamelCase				)

																lowerCamelCase__ =scheduler.step(_lowerCamelCase     ,			_lowerCamelCase     ,			_lowerCamelCase     ,			generator=_lowerCamelCase				)
																lowerCamelCase__ =output.prev_sample

										lowerCamelCase__ =torch.sum(torch.abs(_lowerCamelCase				)				)
										lowerCamelCase__ =torch.mean(torch.abs(_lowerCamelCase				)				)

										assert abs(result_sum.item() - 1_3_8_4_9.3_8_7_7				) < 1E-2
										assert abs(result_mean.item() - 1_8.0_3_3_1				) < 5E-3

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

				def 		_a     (				self				):
										if torch_device == "mps":
																return
										lowerCamelCase__ =self.scheduler_classes[0]
										lowerCamelCase__ =self.get_scheduler_config(prediction_type="v_prediction"				)
										lowerCamelCase__ =scheduler_class(**_lowerCamelCase				)

										scheduler.set_timesteps(self.num_inference_steps				)

										lowerCamelCase__ =self.dummy_model()
										lowerCamelCase__ =self.dummy_sample_deter * scheduler.init_noise_sigma
										lowerCamelCase__ =sample.to(_lowerCamelCase				)

										lowerCamelCase__ =torch.manual_seed(0				)

										for i, t in enumerate(scheduler.timesteps				):
																lowerCamelCase__ =scheduler.scale_model_input(_lowerCamelCase     ,			_lowerCamelCase				)

																lowerCamelCase__ =model(_lowerCamelCase     ,			_lowerCamelCase				)

																lowerCamelCase__ =scheduler.step(_lowerCamelCase     ,			_lowerCamelCase     ,			_lowerCamelCase     ,			generator=_lowerCamelCase				)
																lowerCamelCase__ =output.prev_sample

										lowerCamelCase__ =torch.sum(torch.abs(_lowerCamelCase				)				)
										lowerCamelCase__ =torch.mean(torch.abs(_lowerCamelCase				)				)

										assert abs(result_sum.item() - 3_2_8.9_9_7_0				) < 1E-2
										assert abs(result_mean.item() - 0.4_2_8_4				) < 1E-3






				def 		_a     (				self				):
										if torch_device == "mps":
																return
										lowerCamelCase__ =self.scheduler_classes[0]
										lowerCamelCase__ =self.get_scheduler_config()
										lowerCamelCase__ =scheduler_class(**_lowerCamelCase				)

										scheduler.set_timesteps(self.num_inference_steps     ,			device=_lowerCamelCase				)
										lowerCamelCase__ =torch.manual_seed(0				)

										lowerCamelCase__ =self.dummy_model()
										lowerCamelCase__ =self.dummy_sample_deter.to(_lowerCamelCase				) * scheduler.init_noise_sigma

										for t in scheduler.timesteps:
																lowerCamelCase__ =scheduler.scale_model_input(_lowerCamelCase     ,			_lowerCamelCase				)

																lowerCamelCase__ =model(_lowerCamelCase     ,			_lowerCamelCase				)

																lowerCamelCase__ =scheduler.step(_lowerCamelCase     ,			_lowerCamelCase     ,			_lowerCamelCase     ,			generator=_lowerCamelCase				)
																lowerCamelCase__ =output.prev_sample

										lowerCamelCase__ =torch.sum(torch.abs(_lowerCamelCase				)				)
										lowerCamelCase__ =torch.mean(torch.abs(_lowerCamelCase				)				)

										assert abs(result_sum.item() - 1_3_8_4_9.3_8_1_8				) < 1E-1
										assert abs(result_mean.item() - 1_8.0_3_3_1				) < 1E-3



"""simple docstring"""


import random
import unittest

import torch

from diffusers import IFSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin






@skip_mps
class        __UpperCAmelCase						(     __lowerCAmelCase				,				__lowerCAmelCase				,				unittest.TestCase     ):
				A__				:       Any  							=	IFSuperResolutionPipeline
				A__				:       Tuple  							=	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''width''', '''height'''}
				A__				:       Optional[int]  							=	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
				A__				:       List[Any]  							=	PipelineTesterMixin.required_optional_params - {'''latents'''}

				def 		_a     (				self				):
										return self._get_superresolution_dummy_components()

				def 		_a     (				self     ,			_lowerCamelCase     ,			_lowerCamelCase=0				):
										if str(_lowerCamelCase				).startswith("mps"				):
																lowerCamelCase__ =torch.manual_seed(_lowerCamelCase				)
										else:
																lowerCamelCase__ =torch.Generator(device=_lowerCamelCase				).manual_seed(_lowerCamelCase				)

										lowerCamelCase__ =floats_tensor((1, 3, 32, 32)     ,			rng=random.Random(_lowerCamelCase				)				).to(_lowerCamelCase				)

										lowerCamelCase__ ={
										    "prompt": "A painting of a squirrel eating a burger",
										    "image": image,
										    "generator": generator,
										    "num_inference_steps": 2,
										    "output_type": "numpy",
										}

										return inputs

				@unittest.skipIf(
				    torch_device != "cuda" or not is_xformers_available()     ,			reason="XFormers attention is only available with CUDA and `xformers` installed"     ,			)
				def 		_a     (				self				):
										self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3				)

				def 		_a     (				self				):
										self._test_save_load_optional_components()

				@unittest.skipIf(torch_device != "cuda"     ,			reason="float16 requires CUDA"				)
				def 		_a     (				self				):
										# Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
										super().test_save_load_floataa(expected_max_diff=1E-1				)

				def 		_a     (				self				):
										self._test_attention_slicing_forward_pass(expected_max_diff=1E-2				)

				def 		_a     (				self				):
										self._test_save_load_local()






				def 		_a     (				self				):
										self._test_inference_batch_single_identical(
										    expected_max_diff=1E-2     ,			)


