





import unittest

import numpy as np
import torch

from diffusers import DDIMPipeline, DDIMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()

class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				unittest.TestCase  ):
      A_						:    str							=   DDIMPipeline
      A_						:    str							=   UNCONDITIONAL_IMAGE_GENERATION_PARAMS
      A_						:    Any							=   PipelineTesterMixin.required_optional_params - {
          'num_images_per_prompt',
          'latents',
          'callback',
          'callback_steps',
      }
      A_						:    int							=   UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
      A_						:    List[str]							=   False


      def        a			(self	:				Optional[int]				):





         """simple docstring"""



         torch.manual_seed(0				)
         __snake_case								=       UNetaDModel(
             block_out_channels=(32, 64)    ,      layers_per_block=2    ,      sample_size=32    ,      in_channels=3    ,      out_channels=3    ,      down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')    ,      up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')    ,      )
         __snake_case								=       DDIMScheduler()
         __snake_case								=       {'''unet''': unet, '''scheduler''': scheduler}
         return components


      def        a			(self	:				Union[str, Any]    ,      a__	:				Tuple    ,      a__	:				Tuple=0				):





         """simple docstring"""



         if str(a__				).startswith('''mps'''				):
            __snake_case								=       torch.manual_seed(a__				)
         else:
            __snake_case								=       torch.Generator(device=a__				).manual_seed(a__				)
         __snake_case								=       {
             '''batch_size''': 1,
             '''generator''': generator,
             '''num_inference_steps''': 2,
             '''output_type''': '''numpy''',
         }
         return inputs


      def        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case								=       '''cpu'''

         __snake_case								=       self.get_dummy_components()
         __snake_case								=       self.pipeline_class(**a__				)
         pipe.to(a__				)
         pipe.set_progress_bar_config(disable=a__				)

         __snake_case								=       self.get_dummy_inputs(a__				)
         __snake_case								=       pipe(**a__				).images
         __snake_case								=       image[0, -3:, -3:, -1]

         self.assertEqual(image.shape    ,      (1, 32, 32, 3)				)
         __snake_case								=       np.array(
             [1.000E00, 5.717E-01, 4.717E-01, 1.000E00, 0.000E00, 1.000E00, 3.000E-04, 0.000E00, 9.000E-04]				)
         __snake_case								=       np.abs(image_slice.flatten() - expected_slice				).max()
         self.assertLessEqual(a__    ,      1E-3				)


      def        a			(self	:				int				):





         """simple docstring"""



         super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3				)


      def        a			(self	:				Any				):





         """simple docstring"""



         super().test_save_load_local(expected_max_difference=3E-3				)


      def        a			(self	:				Union[str, Any]				):





         """simple docstring"""



         super().test_save_load_optional_components(expected_max_difference=3E-3				)


      def        a			(self	:				Tuple				):





         """simple docstring"""



         super().test_inference_batch_single_identical(expected_max_diff=3E-3				)






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


      def        a			(self	:				Any				):





         """simple docstring"""



         __snake_case								=       '''google/ddpm-cifar10-32'''

         __snake_case								=       UNetaDModel.from_pretrained(a__				)
         __snake_case								=       DDIMScheduler()

         __snake_case								=       DDIMPipeline(unet=a__    ,      scheduler=a__				)
         ddim.to(a__				)
         ddim.set_progress_bar_config(disable=a__				)

         __snake_case								=       torch.manual_seed(0				)
         __snake_case								=       ddim(generator=a__    ,      eta=0.0    ,      output_type='''numpy'''				).images

         __snake_case								=       image[0, -3:, -3:, -1]

         assert image.shape == (1, 32, 32, 3)
         __snake_case								=       np.array([0.1_7_2_3, 0.1_6_1_7, 0.1_6_0_0, 0.1_6_2_6, 0.1_4_9_7, 0.1_5_1_3, 0.1_5_0_5, 0.1_4_4_2, 0.1_4_5_3]				)

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


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       '''google/ddpm-ema-bedroom-256'''

         __snake_case								=       UNetaDModel.from_pretrained(a__				)
         __snake_case								=       DDIMScheduler.from_pretrained(a__				)

         __snake_case								=       DDIMPipeline(unet=a__    ,      scheduler=a__				)
         ddpm.to(a__				)
         ddpm.set_progress_bar_config(disable=a__				)

         __snake_case								=       torch.manual_seed(0				)
         __snake_case								=       ddpm(generator=a__    ,      output_type='''numpy'''				).images

         __snake_case								=       image[0, -3:, -3:, -1]

         assert image.shape == (1, 256, 256, 3)
         __snake_case								=       np.array([0.0_0_6_0, 0.0_2_0_1, 0.0_3_4_4, 0.0_0_2_4, 0.0_0_1_8, 0.0_0_0_2, 0.0_0_2_2, 0.0_0_0_0, 0.0_0_6_9]				)

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








import tempfile

import torch

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

from .test_schedulers import SchedulerCommonTest

class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase  ):
      A_						:    Union[str, Any]							=   (DPMSolverSinglestepScheduler,)
      A_						:    Union[str, Any]							=   (('num_inference_steps', 25),)


      def        a			(self	:				Dict    ,      **a__	:				Tuple				):





         """simple docstring"""



         __snake_case								=       {
             '''num_train_timesteps''': 1000,
             '''beta_start''': 0.0_0_0_1,
             '''beta_end''': 0.0_2,
             '''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        a			(self	:				str    ,      a__	:				Any=0    ,      **a__	:				Tuple				):





         """simple docstring"""



         __snake_case								=       dict(self.forward_default_kwargs				)
         __snake_case								=       kwargs.pop('''num_inference_steps'''    ,      a__				)
         __snake_case								=       self.dummy_sample
         __snake_case								=       0.1 * sample
         __snake_case								=       [residual + 0.2, residual + 0.1_5, residual + 0.1_0]

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

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

            __snake_case	,     __snake_case								=       sample, sample
            for t in range(a__    ,      time_step + scheduler.config.solver_order + 1				):
               __snake_case								=       scheduler.step(a__    ,      a__    ,      a__    ,      **a__				).prev_sample
               __snake_case								=       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        a			(self	:				Union[str, Any]				):





         """simple docstring"""



         pass


      def        a			(self	:				List[Any]    ,      a__	:				Dict=0    ,      **a__	:				List[str]				):





         """simple docstring"""



         __snake_case								=       dict(self.forward_default_kwargs				)
         __snake_case								=       kwargs.pop('''num_inference_steps'''    ,      a__				)
         __snake_case								=       self.dummy_sample
         __snake_case								=       0.1 * sample
         __snake_case								=       [residual + 0.2, residual + 0.1_5, residual + 0.1_0]

         for scheduler_class in self.scheduler_classes:
            __snake_case								=       self.get_scheduler_config()
            __snake_case								=       scheduler_class(**a__				)
            scheduler.set_timesteps(a__				)

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

            with tempfile.TemporaryDirectory() as tmpdirname:
               scheduler.save_config(a__				)
               __snake_case								=       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)
               __snake_case								=       dummy_past_residuals[: new_scheduler.config.solver_order]

            __snake_case								=       scheduler.step(a__    ,      a__    ,      a__    ,      **a__				).prev_sample
            __snake_case								=       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        a			(self	:				int    ,      a__	:				Tuple=None    ,      **a__	:				List[str]				):





         """simple docstring"""



         if scheduler is None:
            __snake_case								=       self.scheduler_classes[0]
            __snake_case								=       self.get_scheduler_config(**a__				)
            __snake_case								=       scheduler_class(**a__				)

         __snake_case								=       self.scheduler_classes[0]
         __snake_case								=       self.get_scheduler_config(**a__				)
         __snake_case								=       scheduler_class(**a__				)

         __snake_case								=       10
         __snake_case								=       self.dummy_model()
         __snake_case								=       self.dummy_sample_deter
         scheduler.set_timesteps(a__				)

         for i, t in enumerate(scheduler.timesteps				):
            __snake_case								=       model(a__    ,      a__				)
            __snake_case								=       scheduler.step(a__    ,      a__    ,      a__				).prev_sample

         return sample


      def        a			(self	:				str				):





         """simple docstring"""



         __snake_case								=       DPMSolverSinglestepScheduler(**self.get_scheduler_config()				)
         __snake_case								=       50
         __snake_case								=       self.dummy_model()
         __snake_case								=       self.dummy_sample_deter
         scheduler.set_timesteps(a__				)

         # make sure that the first t is uneven
         for i, t in enumerate(scheduler.timesteps[3:]				):
            __snake_case								=       model(a__    ,      a__				)
            __snake_case								=       scheduler.step(a__    ,      a__    ,      a__				).prev_sample

         __snake_case								=       torch.mean(torch.abs(a__				)				)

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


      def        a			(self	:				int				):





         """simple docstring"""



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


      def        a			(self	:				List[str]				):





         """simple docstring"""



         __snake_case								=       DPMSolverSinglestepScheduler(**self.get_scheduler_config()				)
         __snake_case								=       self.full_loop(scheduler=a__				)
         __snake_case								=       torch.mean(torch.abs(a__				)				)

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

         __snake_case								=       DEISMultistepScheduler.from_config(scheduler.config				)
         __snake_case								=       DPMSolverMultistepScheduler.from_config(scheduler.config				)
         __snake_case								=       UniPCMultistepScheduler.from_config(scheduler.config				)
         __snake_case								=       DPMSolverSinglestepScheduler.from_config(scheduler.config				)

         __snake_case								=       self.full_loop(scheduler=a__				)
         __snake_case								=       torch.mean(torch.abs(a__				)				)

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


      def        a			(self	:				List[str]				):





         """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        a			(self	:				Union[str, Any]				):





         """simple docstring"""



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


      def        a			(self	:				Union[str, Any]				):





         """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__    ,      )
                     __snake_case								=       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        a			(self	:				List[str]				):





         """simple docstring"""



         self.check_over_configs(lower_order_final=a__				)
         self.check_over_configs(lower_order_final=a__				)


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



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


      def        a			(self	:				Tuple				):





         """simple docstring"""



         self.check_over_configs(variance_type=a__				)
         self.check_over_configs(variance_type='''learned_range'''				)


      def        a			(self	:				int				):





         """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        a			(self	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       self.full_loop()
         __snake_case								=       torch.mean(torch.abs(a__				)				)

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


      def        a			(self	:				int				):





         """simple docstring"""



         __snake_case								=       self.full_loop(use_karras_sigmas=a__				)
         __snake_case								=       torch.mean(torch.abs(a__				)				)

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


      def        a			(self	:				Tuple				):





         """simple docstring"""



         __snake_case								=       self.full_loop(prediction_type='''v_prediction'''				)
         __snake_case								=       torch.mean(torch.abs(a__				)				)

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


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       self.full_loop(prediction_type='''v_prediction'''    ,      use_karras_sigmas=a__				)
         __snake_case								=       torch.mean(torch.abs(a__				)				)

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


      def        a			(self	:				int				):





         """simple docstring"""



         __snake_case								=       self.scheduler_classes[0]
         __snake_case								=       self.get_scheduler_config(thresholding=a__    ,      dynamic_thresholding_ratio=0				)
         __snake_case								=       scheduler_class(**a__				)

         __snake_case								=       10
         __snake_case								=       self.dummy_model()
         __snake_case								=       self.dummy_sample_deter.half()
         scheduler.set_timesteps(a__				)

         for i, t in enumerate(scheduler.timesteps				):
            __snake_case								=       model(a__    ,      a__				)
            __snake_case								=       scheduler.step(a__    ,      a__    ,      a__				).prev_sample

         assert sample.dtype == torch.floataa

