

'''simple docstring'''






import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNetaDModel, VQModel
from diffusers.utils import PIL_INTERPOLATION, floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch


enable_full_determinism()
class 	SCREAMING_SNAKE_CASE__			(				unittest.TestCase      ):




    @property
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       1
          SCREAMING_SNAKE_CASE_      :			Dict     =       3
          SCREAMING_SNAKE_CASE_      :			str     =       (32, 32)

          SCREAMING_SNAKE_CASE_      :			Tuple     =       floats_tensor((batch_size, num_channels) + sizes				,     rng=random.Random(0			)			).to(lowercase__			)
          return image




    @property
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          torch.manual_seed(0			)
          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       UNetaDModel(
              block_out_channels=(32, 64)				,     layers_per_block=2				,     sample_size=32				,     in_channels=6				,     out_channels=3				,     down_block_types=("DownBlock2D", "AttnDownBlock2D")				,     up_block_types=("AttnUpBlock2D", "UpBlock2D")				,     )
          return model




    @property
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          torch.manual_seed(0			)
          SCREAMING_SNAKE_CASE_      :			Dict     =       VQModel(
              block_out_channels=[32, 64]				,     in_channels=3				,     out_channels=3				,     down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]				,     up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]				,     latent_channels=3				,     )
          return model




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       "cpu"
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       self.dummy_uncond_unet
          SCREAMING_SNAKE_CASE_      :			Any     =       DDIMScheduler()
          SCREAMING_SNAKE_CASE_      :			int     =       self.dummy_vq_model

          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       LDMSuperResolutionPipeline(unet=lowercase__				,     vqvae=lowercase__				,     scheduler=lowercase__			)
          ldm.to(lowercase__			)
          ldm.set_progress_bar_config(disable=lowercase__			)

          SCREAMING_SNAKE_CASE_      :			Any     =       self.dummy_image.to(lowercase__			)

          SCREAMING_SNAKE_CASE_      :			str     =       torch.Generator(device=lowercase__			).manual_seed(0			)
          SCREAMING_SNAKE_CASE_      :			str     =       ldm(image=lowercase__				,     generator=lowercase__				,     num_inference_steps=2				,     output_type="numpy"			).images

          SCREAMING_SNAKE_CASE_      :			Tuple     =       image[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       np.array([0.8678, 0.8245, 0.6381, 0.6830, 0.4385, 0.5599, 0.4641, 0.6201, 0.5150]			)

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







    @unittest.skipIf(torch_device != "cuda"				,     "This test requires a GPU"			)
    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       self.dummy_uncond_unet
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       DDIMScheduler()
          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.dummy_vq_model

          # put models in fp16
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       unet.half()
          SCREAMING_SNAKE_CASE_      :			int     =       vqvae.half()

          SCREAMING_SNAKE_CASE_      :			List[str]     =       LDMSuperResolutionPipeline(unet=lowercase__				,     vqvae=lowercase__				,     scheduler=lowercase__			)
          ldm.to(lowercase__			)
          ldm.set_progress_bar_config(disable=lowercase__			)

          SCREAMING_SNAKE_CASE_      :			Tuple     =       self.dummy_image.to(lowercase__			)

          SCREAMING_SNAKE_CASE_      :			List[str]     =       ldm(lowercase__				,     num_inference_steps=2				,     output_type="numpy"			).images

          assert image.shape == (1, 64, 64, 3)






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




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       load_image(
              "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
              "/vq_diffusion/teddy_bear_pool.png"			)
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       init_image.resize((64, 64)				,     resample=PIL_INTERPOLATION["lanczos"]			)

          SCREAMING_SNAKE_CASE_      :			List[Any]     =       LDMSuperResolutionPipeline.from_pretrained("duongna/ldm-super-resolution"				,     device_map="auto"			)
          ldm.set_progress_bar_config(disable=lowercase__			)

          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       torch.manual_seed(0			)
          SCREAMING_SNAKE_CASE_      :			str     =       ldm(image=lowercase__				,     generator=lowercase__				,     num_inference_steps=20				,     output_type="numpy"			).images

          SCREAMING_SNAKE_CASE_      :			int     =       image[0, -3:, -3:, -1]

          assert image.shape == (1, 256, 256, 3)
          SCREAMING_SNAKE_CASE_      :			Tuple     =       np.array([0.7644, 0.7679, 0.7642, 0.7633, 0.7666, 0.7560, 0.7425, 0.7257, 0.6907]			)

          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 	SCREAMING_SNAKE_CASE__			(				_UpperCAmelCase      ):
    _A			    =							(HeunDiscreteScheduler,)
    _A			    =							10




    def   __lowerCamelCase      (					self				,     **lowercase__			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			List[Any]     =       {
              "num_train_timesteps": 1100,
              "beta_start": 0.0001,
              "beta_end": 0.02,
              "beta_schedule": "linear",
          }

          config.update(**lowercase__			)
          return config




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






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




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






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




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






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




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






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




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






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




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Tuple     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.get_scheduler_config()
          SCREAMING_SNAKE_CASE_      :			List[str]     =       scheduler_class(**lowercase__			)

          scheduler.set_timesteps(self.num_inference_steps			)

          SCREAMING_SNAKE_CASE_      :			int     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			Dict     =       self.dummy_sample_deter * scheduler.init_noise_sigma
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       sample.to(lowercase__			)

          for i, t in enumerate(scheduler.timesteps			):
                SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       scheduler.scale_model_input(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Dict     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			str     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__			)
                SCREAMING_SNAKE_CASE_      :			Dict     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			List[Any]     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			List[str]     =       torch.mean(torch.abs(lowercase__			)			)

          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   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       self.get_scheduler_config(prediction_type="v_prediction"			)
          SCREAMING_SNAKE_CASE_      :			Any     =       scheduler_class(**lowercase__			)

          scheduler.set_timesteps(self.num_inference_steps			)

          SCREAMING_SNAKE_CASE_      :			Any     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       self.dummy_sample_deter * scheduler.init_noise_sigma
          SCREAMING_SNAKE_CASE_      :			List[str]     =       sample.to(lowercase__			)

          for i, t in enumerate(scheduler.timesteps			):
                SCREAMING_SNAKE_CASE_      :			Tuple     =       scheduler.scale_model_input(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			int     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__			)
                SCREAMING_SNAKE_CASE_      :			Optional[int]     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			int     =       torch.mean(torch.abs(lowercase__			)			)

          if torch_device in ["cpu", "mps"]:
                assert abs(result_sum.item() - 4.69_34e-07			) < 1e-2
                assert abs(result_mean.item() - 6.11_12e-10			) < 1e-3
          else:
                # CUDA
                assert abs(result_sum.item() - 4.6_93_42_86_50_17_09_72e-07			) < 1e-2
                assert abs(result_mean.item() - 0.0002			) < 1e-3




    def   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       self.get_scheduler_config()
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       scheduler_class(**lowercase__			)

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

          SCREAMING_SNAKE_CASE_      :			Any     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			str     =       self.dummy_sample_deter.to(lowercase__			) * scheduler.init_noise_sigma

          for t in scheduler.timesteps:
                SCREAMING_SNAKE_CASE_      :			Optional[int]     =       scheduler.scale_model_input(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Dict     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Optional[Any]     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__			)
                SCREAMING_SNAKE_CASE_      :			List[Any]     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			str     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			List[Any]     =       torch.mean(torch.abs(lowercase__			)			)

          if str(lowercase__			).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(lowercase__			).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   __lowerCamelCase      (					self			):






          """simple docstring"""






          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       self.scheduler_classes[0]
          SCREAMING_SNAKE_CASE_      :			int     =       self.get_scheduler_config()
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       scheduler_class(**lowercase__				,     use_karras_sigmas=lowercase__			)

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

          SCREAMING_SNAKE_CASE_      :			List[str]     =       self.dummy_model()
          SCREAMING_SNAKE_CASE_      :			Optional[int]     =       self.dummy_sample_deter.to(lowercase__			) * scheduler.init_noise_sigma
          SCREAMING_SNAKE_CASE_      :			Any     =       sample.to(lowercase__			)

          for t in scheduler.timesteps:
                SCREAMING_SNAKE_CASE_      :			Optional[int]     =       scheduler.scale_model_input(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			List[str]     =       model(lowercase__				,     lowercase__			)

                SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       scheduler.step(lowercase__				,     lowercase__				,     lowercase__			)
                SCREAMING_SNAKE_CASE_      :			Tuple     =       output.prev_sample

          SCREAMING_SNAKE_CASE_      :			Any     =       torch.sum(torch.abs(lowercase__			)			)
          SCREAMING_SNAKE_CASE_      :			Union[str, Any]     =       torch.mean(torch.abs(lowercase__			)			)

          assert abs(result_sum.item() - 0.00015			) < 1e-2
          assert abs(result_mean.item() - 1.98_69_55_45_35_03_46_95e-07			) < 1e-2

