



import tempfile

import torch

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

from .test_schedulers import SchedulerCommonTest





class 		__SCREAMING_SNAKE_CASE      (				lowercase__      ):






				"""simple docstring"""

				SCREAMING_SNAKE_CASE_ 				=  (DPMSolverMultistepScheduler,)
				SCREAMING_SNAKE_CASE_ 				=  (('num_inference_steps', 2_5),)







				def 		__lowerCamelCase(				self						,	**SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							int					     =   {
					    """num_train_timesteps""": 10_00,
					    """beta_start""": 0.0001,
					    """beta_end""": 0.02,
					    """beta_schedule""": """linear""",
					    """solver_order""": 2,
					    """prediction_type""": """epsilon""",
					    """thresholding""": False,
					    """sample_max_value""": 1.0,
					    """algorithm_type""": """dpmsolver++""",
					    """solver_type""": """midpoint""",
					    """lower_order_final""": False,
					    """lambda_min_clipped""": -float("""inf"""      ),
					    """variance_type""": None,
					}

					config.update(**SCREAMING_SNAKE_CASE__      )
					return config







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__=0						,	**SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							List[str]					     =   dict(self.forward_default_kwargs      )
					_snake_case					:							Optional[int]					     =   kwargs.pop("""num_inference_steps"""						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Any					     =   self.dummy_sample
					_snake_case					:							str					     =   0.1 * sample
					_snake_case					:							List[Any]					     =   [residual + 0.2, residual + 0.15, residual + 0.10]

					for scheduler_class in self.scheduler_classes:
						_snake_case					:							List[Any]					     =   self.get_scheduler_config(**SCREAMING_SNAKE_CASE__      )
						_snake_case					:							List[Any]					     =   scheduler_class(**SCREAMING_SNAKE_CASE__      )
						scheduler.set_timesteps(SCREAMING_SNAKE_CASE__      )
						# copy over dummy past residuals
						_snake_case					:							Any					     =   dummy_past_residuals[: scheduler.config.solver_order]

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

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

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







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					pass







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__=0						,	**SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							Dict					     =   dict(self.forward_default_kwargs      )
					_snake_case					:							int					     =   kwargs.pop("""num_inference_steps"""						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Any					     =   self.dummy_sample
					_snake_case					:							Dict					     =   0.1 * sample
					_snake_case					:							List[Any]					     =   [residual + 0.2, residual + 0.15, residual + 0.10]

					for scheduler_class in self.scheduler_classes:
						_snake_case					:							int					     =   self.get_scheduler_config()
						_snake_case					:							Dict					     =   scheduler_class(**SCREAMING_SNAKE_CASE__      )
						scheduler.set_timesteps(SCREAMING_SNAKE_CASE__      )

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

						with tempfile.TemporaryDirectory() as tmpdirname:
							scheduler.save_config(SCREAMING_SNAKE_CASE__      )
							_snake_case					:							str					     =   scheduler_class.from_pretrained(SCREAMING_SNAKE_CASE__      )
							# copy over dummy past residuals
							new_scheduler.set_timesteps(SCREAMING_SNAKE_CASE__      )

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

						_snake_case					:							str					     =   scheduler.step(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	**SCREAMING_SNAKE_CASE__      ).prev_sample
						_snake_case					:							int					     =   new_scheduler.step(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	**SCREAMING_SNAKE_CASE__      ).prev_sample

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







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__=None						,	**SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					if scheduler is None:
						_snake_case					:							List[str]					     =   self.scheduler_classes[0]
						_snake_case					:							int					     =   self.get_scheduler_config(**SCREAMING_SNAKE_CASE__      )
						_snake_case					:							List[Any]					     =   scheduler_class(**SCREAMING_SNAKE_CASE__      )

					_snake_case					:							Union[str, Any]					     =   10
					_snake_case					:							Optional[Any]					     =   self.dummy_model()
					_snake_case					:							List[str]					     =   self.dummy_sample_deter
					scheduler.set_timesteps(SCREAMING_SNAKE_CASE__      )

					for i, t in enumerate(scheduler.timesteps      ):
						_snake_case					:							Union[str, Any]					     =   model(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )
						_snake_case					:							int					     =   scheduler.step(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      ).prev_sample

					return sample







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					_snake_case					:							int					     =   dict(self.forward_default_kwargs      )

					_snake_case					:							str					     =   kwargs.pop("""num_inference_steps"""						,	SCREAMING_SNAKE_CASE__      )

					for scheduler_class in self.scheduler_classes:
						_snake_case					:							Dict					     =   self.get_scheduler_config()
						_snake_case					:							Union[str, Any]					     =   scheduler_class(**SCREAMING_SNAKE_CASE__      )

						_snake_case					:							List[Any]					     =   self.dummy_sample
						_snake_case					:							str					     =   0.1 * sample

						if num_inference_steps is not None and hasattr(SCREAMING_SNAKE_CASE__						,	"""set_timesteps"""      ):
							scheduler.set_timesteps(SCREAMING_SNAKE_CASE__      )
						elif num_inference_steps is not None and not hasattr(SCREAMING_SNAKE_CASE__						,	"""set_timesteps"""      ):
							_snake_case					:							Union[str, Any]					     =   num_inference_steps

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

						_snake_case					:							Union[str, Any]					     =   scheduler.timesteps[5]
						_snake_case					:							List[str]					     =   scheduler.timesteps[6]

						_snake_case					:							Any					     =   scheduler.step(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	**SCREAMING_SNAKE_CASE__      ).prev_sample
						_snake_case					:							Dict					     =   scheduler.step(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	**SCREAMING_SNAKE_CASE__      ).prev_sample

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







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					for timesteps in [25, 50, 1_00, 9_99, 10_00]:
						self.check_over_configs(num_train_timesteps=SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					self.check_over_configs(thresholding=SCREAMING_SNAKE_CASE__      )
					for order in [1, 2, 3]:
						for solver_type in ["midpoint", "heun"]:
							for threshold in [0.5, 1.0, 2.0]:
								for prediction_type in ["epsilon", "sample"]:
									self.check_over_configs(
									    thresholding=SCREAMING_SNAKE_CASE__						,	prediction_type=SCREAMING_SNAKE_CASE__						,	sample_max_value=SCREAMING_SNAKE_CASE__						,	algorithm_type="""dpmsolver++"""						,	solver_order=SCREAMING_SNAKE_CASE__						,	solver_type=SCREAMING_SNAKE_CASE__						,	)







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

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







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					for algorithm_type in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]:
						for solver_type in ["midpoint", "heun"]:
							for order in [1, 2, 3]:
								for prediction_type in ["epsilon", "sample"]:
									if algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
										if order == 3:
											continue
									else:
										self.check_over_configs(
										    solver_order=SCREAMING_SNAKE_CASE__						,	solver_type=SCREAMING_SNAKE_CASE__						,	prediction_type=SCREAMING_SNAKE_CASE__						,	algorithm_type=SCREAMING_SNAKE_CASE__						,	)
									_snake_case					:							int					     =   self.full_loop(
									    solver_order=SCREAMING_SNAKE_CASE__						,	solver_type=SCREAMING_SNAKE_CASE__						,	prediction_type=SCREAMING_SNAKE_CASE__						,	algorithm_type=SCREAMING_SNAKE_CASE__						,	)
									assert not torch.isnan(SCREAMING_SNAKE_CASE__      ).any(), "Samples have nan numbers"







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					self.check_over_configs(lower_order_final=SCREAMING_SNAKE_CASE__      )
					self.check_over_configs(lower_order_final=SCREAMING_SNAKE_CASE__      )







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					self.check_over_configs(lambda_min_clipped=-float("""inf"""      )      )
					self.check_over_configs(lambda_min_clipped=-5.1      )







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					self.check_over_configs(variance_type=SCREAMING_SNAKE_CASE__      )
					self.check_over_configs(variance_type="""learned_range"""      )







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					for num_inference_steps in [1, 2, 3, 5, 10, 50, 1_00, 9_99, 10_00]:
						self.check_over_forward(num_inference_steps=SCREAMING_SNAKE_CASE__						,	time_step=0      )







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					_snake_case					:							Tuple					     =   self.full_loop()
					_snake_case					:							Any					     =   torch.mean(torch.abs(SCREAMING_SNAKE_CASE__      )      )

					assert abs(result_mean.item() - 0.3301      ) < 1e-3







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					_snake_case					:							int					     =   self.full_loop(thresholding=SCREAMING_SNAKE_CASE__						,	dynamic_thresholding_ratio=0.87						,	sample_max_value=0.5      )
					_snake_case					:							List[Any]					     =   torch.mean(torch.abs(SCREAMING_SNAKE_CASE__      )      )

					assert abs(result_mean.item() - 1.1364      ) < 1e-3







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					_snake_case					:							str					     =   self.full_loop(prediction_type="""v_prediction"""      )
					_snake_case					:							Dict					     =   torch.mean(torch.abs(SCREAMING_SNAKE_CASE__      )      )

					assert abs(result_mean.item() - 0.2251      ) < 1e-3







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					_snake_case					:							Optional[int]					     =   self.full_loop(prediction_type="""v_prediction"""						,	use_karras_sigmas=SCREAMING_SNAKE_CASE__      )
					_snake_case					:							int					     =   torch.mean(torch.abs(SCREAMING_SNAKE_CASE__      )      )

					assert abs(result_mean.item() - 0.2096      ) < 1e-3







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					_snake_case					:							Dict					     =   DPMSolverMultistepScheduler(**self.get_scheduler_config()      )
					_snake_case					:							Union[str, Any]					     =   self.full_loop(scheduler=SCREAMING_SNAKE_CASE__      )
					_snake_case					:							List[Any]					     =   torch.mean(torch.abs(SCREAMING_SNAKE_CASE__      )      )

					assert abs(result_mean.item() - 0.3301      ) < 1e-3

					_snake_case					:							int					     =   DPMSolverSinglestepScheduler.from_config(scheduler.config      )
					_snake_case					:							int					     =   UniPCMultistepScheduler.from_config(scheduler.config      )
					_snake_case					:							List[str]					     =   DEISMultistepScheduler.from_config(scheduler.config      )
					_snake_case					:							Optional[Any]					     =   DPMSolverMultistepScheduler.from_config(scheduler.config      )

					_snake_case					:							int					     =   self.full_loop(scheduler=SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Union[str, Any]					     =   torch.mean(torch.abs(SCREAMING_SNAKE_CASE__      )      )

					assert abs(result_mean.item() - 0.3301      ) < 1e-3







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					_snake_case					:							Optional[Any]					     =   self.scheduler_classes[0]
					_snake_case					:							str					     =   self.get_scheduler_config(thresholding=SCREAMING_SNAKE_CASE__						,	dynamic_thresholding_ratio=0      )
					_snake_case					:							str					     =   scheduler_class(**SCREAMING_SNAKE_CASE__      )

					_snake_case					:							Tuple					     =   10
					_snake_case					:							Optional[Any]					     =   self.dummy_model()
					_snake_case					:							Tuple					     =   self.dummy_sample_deter.half()
					scheduler.set_timesteps(SCREAMING_SNAKE_CASE__      )

					for i, t in enumerate(scheduler.timesteps      ):
						_snake_case					:							List[str]					     =   model(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )
						_snake_case					:							List[Any]					     =   scheduler.step(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      ).prev_sample

					assert sample.dtype == torch.floataa







				def 		__lowerCamelCase(				self						,	**SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					for scheduler_class in self.scheduler_classes:
						_snake_case					:							Optional[int]					     =   self.get_scheduler_config(**SCREAMING_SNAKE_CASE__      )
						_snake_case					:							List[Any]					     =   scheduler_class(**SCREAMING_SNAKE_CASE__      )

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










import tempfile
import unittest

import numpy as np
import torch

from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor





class 		__SCREAMING_SNAKE_CASE      (				unittest.TestCase      ):






				"""simple docstring"""

				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__ = False      ):






					"""simple docstring"""

					_snake_case					:							Any					     =   10

					if only_cross_attention:
						_snake_case					:							Tuple					     =   12
					else:
						# when only cross attention is not set, the cross attention dim must be the same as the query dim
						_snake_case					:							Any					     =   query_dim

					return {
					    "query_dim": query_dim,
					    "cross_attention_dim": cross_attention_dim,
					    "heads": 2,
					    "dim_head": 4,
					    "added_kv_proj_dim": 6,
					    "norm_num_groups": 1,
					    "only_cross_attention": only_cross_attention,
					    "processor": AttnAddedKVProcessor(),
					}







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

					_snake_case					:							Any					     =   2

					_snake_case					:							List[str]					     =   torch.rand(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	3						,	2      )
					_snake_case					:							List[str]					     =   torch.rand(SCREAMING_SNAKE_CASE__						,	4						,	SCREAMING_SNAKE_CASE__      )
					_snake_case					:							Optional[int]					     =   None

					return {
					    "hidden_states": hidden_states,
					    "encoder_hidden_states": encoder_hidden_states,
					    "attention_mask": attention_mask,
					}







				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					torch.manual_seed(0      )

					_snake_case					:							Optional[int]					     =   self.get_constructor_arguments(only_cross_attention=SCREAMING_SNAKE_CASE__      )
					_snake_case					:							List[str]					     =   Attention(**SCREAMING_SNAKE_CASE__      )

					self.assertTrue(attn.to_k is not None      )
					self.assertTrue(attn.to_v is not None      )

					_snake_case					:							Dict					     =   self.get_forward_arguments(
					    query_dim=constructor_args["""query_dim"""]						,	added_kv_proj_dim=constructor_args["""added_kv_proj_dim"""]      )

					_snake_case					:							Any					     =   attn(**SCREAMING_SNAKE_CASE__      )

					# only self attention

					torch.manual_seed(0      )

					_snake_case					:							Any					     =   self.get_constructor_arguments(only_cross_attention=SCREAMING_SNAKE_CASE__      )
					_snake_case					:							List[str]					     =   Attention(**SCREAMING_SNAKE_CASE__      )

					self.assertTrue(attn.to_k is None      )
					self.assertTrue(attn.to_v is None      )

					_snake_case					:							str					     =   self.get_forward_arguments(
					    query_dim=constructor_args["""query_dim"""]						,	added_kv_proj_dim=constructor_args["""added_kv_proj_dim"""]      )

					_snake_case					:							Any					     =   attn(**SCREAMING_SNAKE_CASE__      )

					self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()      )





class 		__SCREAMING_SNAKE_CASE      (				unittest.TestCase      ):






				"""simple docstring"""

				def 		__lowerCamelCase(				self      ):






					"""simple docstring"""

					_snake_case					:							str					     =   DiffusionPipeline.from_pretrained("""hf-internal-testing/tiny-stable-diffusion-pipe"""						,	safety_checker=SCREAMING_SNAKE_CASE__      )

					_snake_case					:							Tuple					     =   pipe(
					    """foo"""						,	num_inference_steps=2						,	generator=torch.Generator("""cpu"""      ).manual_seed(0      )						,	output_type="""np"""						,	).images

					# the initial conversion succeeds
					_snake_case					:							List[str]					     =   DiffusionPipeline.from_pretrained(
					    """hf-internal-testing/tiny-stable-diffusion-pipe"""						,	device_map="""sequential"""						,	safety_checker=SCREAMING_SNAKE_CASE__      )

					_snake_case					:							Union[str, Any]					     =   pipe(
					    """foo"""						,	num_inference_steps=2						,	generator=torch.Generator("""cpu"""      ).manual_seed(0      )						,	output_type="""np"""						,	).images

					with tempfile.TemporaryDirectory() as tmpdir:
						# save the converted model
						pipe.save_pretrained(SCREAMING_SNAKE_CASE__      )

						# can also load the converted weights
						_snake_case					:							str					     =   DiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE__						,	device_map="""sequential"""						,	safety_checker=SCREAMING_SNAKE_CASE__      )

					_snake_case					:							int					     =   pipe(
					    """foo"""						,	num_inference_steps=2						,	generator=torch.Generator("""cpu"""      ).manual_seed(0      )						,	output_type="""np"""						,	).images

					self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	atol=1e-5      )      )
					self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	atol=1e-5      )      )





