
'''simple docstring'''






import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class     __snake_case		(   a__    ,   unittest.TestCase):
					_lowerCAmelCase			     =	DiTPipeline
					_lowerCAmelCase			     =	CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
					_lowerCAmelCase			     =	PipelineTesterMixin.required_optional_params - {
					    '''latents''',
					    '''num_images_per_prompt''',
					    '''callback''',
					    '''callback_steps''',
					}
					_lowerCAmelCase			     =	CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
					_lowerCAmelCase			     =	False
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									torch.manual_seed(0      )
									lowerCamelCase					:			Any									= TransformeraDModel(
									    sample_size=16,    num_layers=2,    patch_size=4,    attention_head_dim=8,    num_attention_heads=2,    in_channels=4,    out_channels=8,    attention_bias=A,    activation_fn='gelu-approximate',    num_embeds_ada_norm=1000,    norm_type='ada_norm_zero',    norm_elementwise_affine=A,    )
									lowerCamelCase					:			List[Any]									= AutoencoderKL()
									lowerCamelCase					:			Tuple									= DDIMScheduler()
									lowerCamelCase					:			List[str]									= {'transformer': transformer.eval(), 'vae': vae.eval(), 'scheduler': scheduler}
									return components
					def     UpperCAmelCase_							(     self,    A,    A=0      ):

									"""simple docstring"""


									if str(A      ).startswith('mps'      ):
													lowerCamelCase					:			Optional[Any]									= torch.manual_seed(A      )
									else:
													lowerCamelCase					:			Union[str, Any]									= torch.Generator(device=A      ).manual_seed(A      )
									lowerCamelCase					:			Tuple									= {
									    'class_labels': [1],
									    'generator': generator,
									    'num_inference_steps': 2,
									    'output_type': 'numpy',
									}
									return inputs
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			int									= 'cpu'

									lowerCamelCase					:			int									= self.get_dummy_components()
									lowerCamelCase					:			Optional[Any]									= self.pipeline_class(**A      )
									pipe.to(A      )
									pipe.set_progress_bar_config(disable=A      )

									lowerCamelCase					:			str									= self.get_dummy_inputs(A      )
									lowerCamelCase					:			str									= pipe(**A      ).images
									lowerCamelCase					:			List[str]									= image[0, -3:, -3:, -1]

									self.assertEqual(image.shape,    (1, 16, 16, 3)      )
									lowerCamelCase					:			str									= np.array([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457]      )
									lowerCamelCase					:			Tuple									= np.abs(image_slice.flatten() - expected_slice      ).max()
									self.assertLessEqual(A,    1e-3      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									self._test_inference_batch_single_identical(relax_max_difference=A,    expected_max_diff=1e-3      )




					@unittest.skipIf(
					    torch_device != 'cuda' or not is_xformers_available(),    reason='XFormers attention is only available with CUDA and `xformers` installed',    )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3      )





@require_torch_gpu
@slow
class     __snake_case		(   unittest.TestCase):
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= torch.manual_seed(0      )

									lowerCamelCase					:			List[Any]									= DiTPipeline.from_pretrained('facebook/DiT-XL-2-256'      )
									pipe.to('cuda'      )

									lowerCamelCase					:			int									= ['vase', 'umbrella', 'white shark', 'white wolf']
									lowerCamelCase					:			Tuple									= pipe.get_label_ids(A      )

									lowerCamelCase					:			List[Any]									= pipe(A,    generator=A,    num_inference_steps=40,    output_type='np'      ).images

									for word, image in zip(A,    A      ):
													lowerCamelCase					:			Tuple									= load_numpy(
													    F'''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy'''      )
													assert np.abs((expected_image - image).max()      ) < 1e-2




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= DiTPipeline.from_pretrained('facebook/DiT-XL-2-512'      )
									lowerCamelCase					:			int									= DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
									pipe.to('cuda'      )

									lowerCamelCase					:			Any									= ['vase', 'umbrella']
									lowerCamelCase					:			Any									= pipe.get_label_ids(A      )

									lowerCamelCase					:			Optional[Any]									= torch.manual_seed(0      )
									lowerCamelCase					:			Union[str, Any]									= pipe(A,    generator=A,    num_inference_steps=25,    output_type='np'      ).images

									for word, image in zip(A,    A      ):
													lowerCamelCase					:			Any									= load_numpy(
													    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
													    F'''/dit/{word}_512.npy'''      )

													assert np.abs((expected_image - image).max()      ) < 1e-1


'''simple docstring'''






import tempfile

import torch

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

from .test_schedulers import SchedulerCommonTest




class     __snake_case		(   a__):
					_lowerCAmelCase			     =	(DPMSolverSinglestepScheduler,)
					_lowerCAmelCase			     =	(('''num_inference_steps''', 25),)
					def     UpperCAmelCase_							(     self,    **A      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= {
									    'num_train_timesteps': 1000,
									    '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',
									    'lambda_min_clipped': -float('inf'      ),
									    'variance_type': None,
									}

									config.update(**A      )
									return config
					def     UpperCAmelCase_							(     self,    A=0,    **A      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= dict(self.forward_default_kwargs      )
									lowerCamelCase					:			Optional[Any]									= kwargs.pop('num_inference_steps',    A      )
									lowerCamelCase					:			Union[str, Any]									= self.dummy_sample
									lowerCamelCase					:			Dict									= 0.1 * sample
									lowerCamelCase					:			Dict									= [residual + 0.2, residual + 0.15, residual + 0.10]

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

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

													lowerCamelCase			, lowerCamelCase					:			Optional[int]									= sample, sample
													for t in range(A,    time_step + scheduler.config.solver_order + 1      ):
																	lowerCamelCase					:			Dict									= scheduler.step(A,    A,    A,    **A      ).prev_sample
																	lowerCamelCase					:			Optional[int]									= new_scheduler.step(A,    A,    A,    **A      ).prev_sample

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

									"""simple docstring"""


									pass
					def     UpperCAmelCase_							(     self,    A=0,    **A      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= dict(self.forward_default_kwargs      )
									lowerCamelCase					:			str									= kwargs.pop('num_inference_steps',    A      )
									lowerCamelCase					:			Union[str, Any]									= self.dummy_sample
									lowerCamelCase					:			List[str]									= 0.1 * sample
									lowerCamelCase					:			List[str]									= [residual + 0.2, residual + 0.15, residual + 0.10]

									for scheduler_class in self.scheduler_classes:
													lowerCamelCase					:			Tuple									= self.get_scheduler_config()
													lowerCamelCase					:			Optional[Any]									= scheduler_class(**A      )
													scheduler.set_timesteps(A      )

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

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

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

													lowerCamelCase					:			int									= scheduler.step(A,    A,    A,    **A      ).prev_sample
													lowerCamelCase					:			Dict									= new_scheduler.step(A,    A,    A,    **A      ).prev_sample

													assert torch.sum(torch.abs(output - new_output      )      ) < 1e-5, "Scheduler outputs are not identical"
					def     UpperCAmelCase_							(     self,    A=None,    **A      ):

									"""simple docstring"""


									if scheduler is None:
													lowerCamelCase					:			Any									= self.scheduler_classes[0]
													lowerCamelCase					:			Optional[Any]									= self.get_scheduler_config(**A      )
													lowerCamelCase					:			Optional[int]									= scheduler_class(**A      )

									lowerCamelCase					:			List[Any]									= self.scheduler_classes[0]
									lowerCamelCase					:			Optional[Any]									= self.get_scheduler_config(**A      )
									lowerCamelCase					:			Optional[int]									= scheduler_class(**A      )

									lowerCamelCase					:			Any									= 10
									lowerCamelCase					:			Dict									= self.dummy_model()
									lowerCamelCase					:			Any									= self.dummy_sample_deter
									scheduler.set_timesteps(A      )

									for i, t in enumerate(scheduler.timesteps      ):
													lowerCamelCase					:			Dict									= model(A,    A      )
													lowerCamelCase					:			List[str]									= scheduler.step(A,    A,    A      ).prev_sample

									return sample
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= DPMSolverSinglestepScheduler(**self.get_scheduler_config()      )
									lowerCamelCase					:			Dict									= 50
									lowerCamelCase					:			Tuple									= self.dummy_model()
									lowerCamelCase					:			Optional[int]									= self.dummy_sample_deter
									scheduler.set_timesteps(A      )

									# make sure that the first t is uneven
									for i, t in enumerate(scheduler.timesteps[3:]      ):
													lowerCamelCase					:			Any									= model(A,    A      )
													lowerCamelCase					:			Optional[int]									= scheduler.step(A,    A,    A      ).prev_sample

									lowerCamelCase					:			Any									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.2574      ) < 1e-3
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									for timesteps in [25, 50, 100, 999, 1000]:
													self.check_over_configs(num_train_timesteps=A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= DPMSolverSinglestepScheduler(**self.get_scheduler_config()      )
									lowerCamelCase					:			str									= self.full_loop(scheduler=A      )
									lowerCamelCase					:			Optional[int]									= torch.mean(torch.abs(A      )      )

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

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

									lowerCamelCase					:			str									= self.full_loop(scheduler=A      )
									lowerCamelCase					:			Optional[int]									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.2791      ) < 1e-3
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									self.check_over_configs(thresholding=A      )
									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=A,    prediction_type=A,    sample_max_value=A,    algorithm_type='dpmsolver++',    solver_order=A,    solver_type=A,    )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


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

									"""simple docstring"""


									for algorithm_type in ["dpmsolver", "dpmsolver++"]:
													for solver_type in ["midpoint", "heun"]:
																	for order in [1, 2, 3]:
																					for prediction_type in ["epsilon", "sample"]:
																									self.check_over_configs(
																									    solver_order=A,    solver_type=A,    prediction_type=A,    algorithm_type=A,    )
																									lowerCamelCase					:			Optional[Any]									= self.full_loop(
																									    solver_order=A,    solver_type=A,    prediction_type=A,    algorithm_type=A,    )
																									assert not torch.isnan(A      ).any(), "Samples have nan numbers"
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									self.check_over_configs(lower_order_final=A      )
									self.check_over_configs(lower_order_final=A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									self.check_over_configs(lambda_min_clipped=-float('inf'      )      )
									self.check_over_configs(lambda_min_clipped=-5.1      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									self.check_over_configs(variance_type=A      )
									self.check_over_configs(variance_type='learned_range'      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									for num_inference_steps in [1, 2, 3, 5, 10, 50, 100, 999, 1000]:
													self.check_over_forward(num_inference_steps=A,    time_step=0      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= self.full_loop()
									lowerCamelCase					:			str									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.2791      ) < 1e-3
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= self.full_loop(use_karras_sigmas=A      )
									lowerCamelCase					:			Tuple									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.2248      ) < 1e-3
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= self.full_loop(prediction_type='v_prediction'      )
									lowerCamelCase					:			Dict									= torch.mean(torch.abs(A      )      )

									assert abs(result_mean.item() - 0.1453      ) < 1e-3
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= self.full_loop(prediction_type='v_prediction',    use_karras_sigmas=A      )
									lowerCamelCase					:			Optional[Any]									= torch.mean(torch.abs(A      )      )

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




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[Any]									= self.scheduler_classes[0]
									lowerCamelCase					:			Dict									= self.get_scheduler_config(thresholding=A,    dynamic_thresholding_ratio=0      )
									lowerCamelCase					:			str									= scheduler_class(**A      )

									lowerCamelCase					:			List[Any]									= 10
									lowerCamelCase					:			List[str]									= self.dummy_model()
									lowerCamelCase					:			int									= self.dummy_sample_deter.half()
									scheduler.set_timesteps(A      )

									for i, t in enumerate(scheduler.timesteps      ):
													lowerCamelCase					:			str									= model(A,    A      )
													lowerCamelCase					:			Tuple									= scheduler.step(A,    A,    A      ).prev_sample

									assert sample.dtype == torch.floataa
