



"""simple docstring"""





import unittest

import numpy as np
import torch

from diffusers import DDPMPipeline, DDPMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device


enable_full_determinism()


class       __lowerCamelCase	(   unittest.TestCase    ):







   @property
   def 							A__						(self				):




      '''simple docstring'''




      torch.manual_seed(0				)
      _lowerCAmelCase					     =		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""")       ,	)
      return model







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		"""cpu"""
      _lowerCAmelCase					     =		self.dummy_uncond_unet
      _lowerCAmelCase					     =		DDPMScheduler()

      _lowerCAmelCase					     =		DDPMPipeline(unet=lowerCamelCase       ,	scheduler=lowerCamelCase				)
      ddpm.to(lowerCamelCase				)
      ddpm.set_progress_bar_config(disable=lowerCamelCase				)

      _lowerCAmelCase					     =		torch.Generator(device=lowerCamelCase				).manual_seed(0				)
      _lowerCAmelCase					     =		ddpm(generator=lowerCamelCase       ,	num_inference_steps=2       ,	output_type="""numpy"""				).images

      _lowerCAmelCase					     =		torch.Generator(device=lowerCamelCase				).manual_seed(0				)
      _lowerCAmelCase					     =		ddpm(generator=lowerCamelCase       ,	num_inference_steps=2       ,	output_type="""numpy"""       ,	return_dict=lowerCamelCase				)[0]

      _lowerCAmelCase					     =		image[0, -3:, -3:, -1]
      _lowerCAmelCase					     =		image_from_tuple[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      _lowerCAmelCase					     =		np.array(
          [9.956e-01, 5.785e-01, 4.675e-01, 9.930e-01, 0.0, 1.000, 1.199e-03, 2.648e-04, 5.101e-04]				)

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







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.dummy_uncond_unet
      _lowerCAmelCase					     =		DDPMScheduler(prediction_type="""sample"""				)

      _lowerCAmelCase					     =		DDPMPipeline(unet=lowerCamelCase       ,	scheduler=lowerCamelCase				)
      ddpm.to(lowerCamelCase				)
      ddpm.set_progress_bar_config(disable=lowerCamelCase				)

      _lowerCAmelCase					     =		torch.manual_seed(0				)
      _lowerCAmelCase					     =		ddpm(generator=lowerCamelCase       ,	num_inference_steps=2       ,	output_type="""numpy"""				).images

      _lowerCAmelCase					     =		torch.manual_seed(0				)
      _lowerCAmelCase					     =		ddpm(generator=lowerCamelCase       ,	num_inference_steps=2       ,	output_type="""numpy"""				)[0]

      _lowerCAmelCase					     =		image[0, -3:, -3:, -1]
      _lowerCAmelCase					     =		image_eps[0, -3:, -3:, -1]

      assert image.shape == (1, 32, 32, 3)
      _lowerCAmelCase					     =		1e-2 if torch_device != """mps""" else 3e-2
      assert np.abs(image_slice.flatten() - image_eps_slice.flatten()				).max() < tolerance





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







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		"""google/ddpm-cifar10-32"""

      _lowerCAmelCase					     =		UNetaDModel.from_pretrained(lowerCamelCase				)
      _lowerCAmelCase					     =		DDPMScheduler.from_pretrained(lowerCamelCase				)

      _lowerCAmelCase					     =		DDPMPipeline(unet=lowerCamelCase       ,	scheduler=lowerCamelCase				)
      ddpm.to(lowerCamelCase				)
      ddpm.set_progress_bar_config(disable=lowerCamelCase				)

      _lowerCAmelCase					     =		torch.manual_seed(0				)
      _lowerCAmelCase					     =		ddpm(generator=lowerCamelCase       ,	output_type="""numpy"""				).images

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

      assert image.shape == (1, 32, 32, 3)
      _lowerCAmelCase					     =		np.array([0.4200, 0.3588, 0.1939, 0.3847, 0.3382, 0.2647, 0.4155, 0.3582, 0.3385]				)
      assert np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2




"""simple docstring"""





import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest


class       __lowerCamelCase	(   __lowercase    ):
   __UpperCamelCase									=   (HeunDiscreteScheduler,)
   __UpperCamelCase									=   10







   def 							A__						(self       ,	**lowerCamelCase				):




      '''simple docstring'''




      _lowerCAmelCase					     =		{
          """num_train_timesteps""": 1_100,
          """beta_start""": 0.0001,
          """beta_end""": 0.02,
          """beta_schedule""": """linear""",
      }

      config.update(**lowerCamelCase				)
      return config







   def 							A__						(self				):




      '''simple docstring'''




      for timesteps in [10, 50, 100, 1_000]:
         self.check_over_configs(num_train_timesteps=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]       ,	[0.0002, 0.002, 0.02]				):
         self.check_over_configs(beta_start=lowerCamelCase       ,	beta_end=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for schedule in ["linear", "scaled_linear", "exp"]:
         self.check_over_configs(beta_schedule=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




      for clip_sample_range in [1.0, 2.0, 3.0]:
         self.check_over_configs(clip_sample_range=lowerCamelCase       ,	clip_sample=lowerCamelCase				)







   def 							A__						(self				):




      '''simple docstring'''




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







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _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				)

      for i, t in enumerate(scheduler.timesteps				):
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      if torch_device in ["cpu", "mps"]:
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3
      else:
         # CUDA
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _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				)

      for i, t in enumerate(scheduler.timesteps				):
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      if torch_device in ["cpu", "mps"]:
         assert abs(result_sum.item() - 4.6934e-07				) < 1e-2
         assert abs(result_mean.item() - 6.1112e-10				) < 1e-3
      else:
         # CUDA
         assert abs(result_sum.item() - 4.693428650170972e-07				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase				)

      scheduler.set_timesteps(self.num_inference_steps       ,	device=lowerCamelCase				)

      _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				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      if str(lowerCamelCase				).startswith("""cpu"""				):
         # The following sum varies between 148 and 156 on mps. Why?
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3
      elif str(lowerCamelCase				).startswith("""mps"""				):
         # Larger tolerance on mps
         assert abs(result_mean.item() - 0.0002				) < 1e-2
      else:
         # CUDA
         assert abs(result_sum.item() - 0.1233				) < 1e-2
         assert abs(result_mean.item() - 0.0002				) < 1e-3







   def 							A__						(self				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.scheduler_classes[0]
      _lowerCAmelCase					     =		self.get_scheduler_config()
      _lowerCAmelCase					     =		scheduler_class(**lowerCamelCase       ,	use_karras_sigmas=lowerCamelCase				)

      scheduler.set_timesteps(self.num_inference_steps       ,	device=lowerCamelCase				)

      _lowerCAmelCase					     =		self.dummy_model()
      _lowerCAmelCase					     =		self.dummy_sample_deter.to(lowerCamelCase				) * scheduler.init_noise_sigma
      _lowerCAmelCase					     =		sample.to(lowerCamelCase				)

      for t in scheduler.timesteps:
         _lowerCAmelCase					     =		scheduler.scale_model_input(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		model(lowerCamelCase       ,	lowerCamelCase				)

         _lowerCAmelCase					     =		scheduler.step(lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase				)
         _lowerCAmelCase					     =		output.prev_sample

      _lowerCAmelCase					     =		torch.sum(torch.abs(lowerCamelCase				)				)
      _lowerCAmelCase					     =		torch.mean(torch.abs(lowerCamelCase				)				)

      assert abs(result_sum.item() - 0.0_0015				) < 1e-2
      assert abs(result_mean.item() - 1.9869554535034695e-07				) < 1e-2