

import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest
class  _A		(  __magic_name__):
  SCREAMING_SNAKE_CASE    :      Tuple       =  (KDPMaDiscreteScheduler,)
  SCREAMING_SNAKE_CASE    :      List[Any]       =  10

  def  UpperCAmelCase    (				self ,					**_SCREAMING_SNAKE_CASE					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     {
            'num_train_timesteps': 1100,
            'beta_start': 0.0001,
            'beta_end': 0.02,
            'beta_schedule': 'linear',
        }

        config.update(**_SCREAMING_SNAKE_CASE					)
        return config

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







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

  def  UpperCAmelCase    (				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=_SCREAMING_SNAKE_CASE ,					beta_end=_SCREAMING_SNAKE_CASE					)

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







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

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







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

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  List[str]           =     self.scheduler_classes[0]
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     self.get_scheduler_config(prediction_type='v_prediction'					)
        SCREAMING_SNAKE_CASE_				:  Tuple           =     scheduler_class(**_SCREAMING_SNAKE_CASE					)

        scheduler.set_timesteps(self.num_inference_steps					)

        SCREAMING_SNAKE_CASE_				:  List[str]           =     self.dummy_model()
        SCREAMING_SNAKE_CASE_				:  List[str]           =     self.dummy_sample_deter * scheduler.init_noise_sigma
        SCREAMING_SNAKE_CASE_				:  Tuple           =     sample.to(_SCREAMING_SNAKE_CASE					)

        for i, t in enumerate(scheduler.timesteps					):
              SCREAMING_SNAKE_CASE_				:  Any           =     scheduler.scale_model_input(_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE					)

              SCREAMING_SNAKE_CASE_				:  Any           =     model(_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE					)

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

        SCREAMING_SNAKE_CASE_				:  int           =     torch.sum(torch.abs(_SCREAMING_SNAKE_CASE					)					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     torch.mean(torch.abs(_SCREAMING_SNAKE_CASE					)					)

        if torch_device in ["cpu", "mps"]:
              assert abs(result_sum.item() - 4.6_934e-07					) < 1e-2
              assert abs(result_mean.item() - 6.1_112e-10					) < 1e-3
        else:
              # CUDA
              assert abs(result_sum.item() - 4.693_428_650_170_972e-07					) < 1e-2
              assert abs(result_mean.item() - 0.0002					) < 1e-3

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        if torch_device == "mps":
              return
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     self.scheduler_classes[0]
        SCREAMING_SNAKE_CASE_				:  Dict           =     self.get_scheduler_config()
        SCREAMING_SNAKE_CASE_				:  Dict           =     scheduler_class(**_SCREAMING_SNAKE_CASE					)

        scheduler.set_timesteps(self.num_inference_steps					)

        SCREAMING_SNAKE_CASE_				:  Dict           =     self.dummy_model()
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     self.dummy_sample_deter * scheduler.init_noise_sigma
        SCREAMING_SNAKE_CASE_				:  str           =     sample.to(_SCREAMING_SNAKE_CASE					)

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

              SCREAMING_SNAKE_CASE_				:  Optional[int]           =     model(_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE					)

              SCREAMING_SNAKE_CASE_				:  List[Any]           =     scheduler.step(_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE					)
              SCREAMING_SNAKE_CASE_				:  Tuple           =     output.prev_sample

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

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



  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        if torch_device == "mps":
              return
        SCREAMING_SNAKE_CASE_				:  Any           =     self.scheduler_classes[0]
        SCREAMING_SNAKE_CASE_				:  List[str]           =     self.get_scheduler_config()
        SCREAMING_SNAKE_CASE_				:  Tuple           =     scheduler_class(**_SCREAMING_SNAKE_CASE					)

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

        SCREAMING_SNAKE_CASE_				:  Any           =     self.dummy_model()
        SCREAMING_SNAKE_CASE_				:  int           =     self.dummy_sample_deter.to(_SCREAMING_SNAKE_CASE					) * scheduler.init_noise_sigma

        for t in scheduler.timesteps:
              SCREAMING_SNAKE_CASE_				:  Tuple           =     scheduler.scale_model_input(_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE					)

              SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     model(_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE					)

              SCREAMING_SNAKE_CASE_				:  int           =     scheduler.step(_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE ,					_SCREAMING_SNAKE_CASE					)
              SCREAMING_SNAKE_CASE_				:  str           =     output.prev_sample

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

        if str(_SCREAMING_SNAKE_CASE					).startswith('cpu'					):
              # The following sum varies between 148 and 156 on mps. Why?
              assert abs(result_sum.item() - 20.4125					) < 1e-2
              assert abs(result_mean.item() - 0.0266					) < 1e-3
        else:
              # CUDA
              assert abs(result_sum.item() - 20.4125					) < 1e-2
              assert abs(result_mean.item() - 0.0266					) < 1e-3









import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()
class  _A		(  unittest.TestCase):

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        torch.manual_seed(0					)
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     UNetaDModel(
            sample_size=(32, 64) ,					in_channels=1 ,					out_channels=1 ,					layers_per_block=2 ,					block_out_channels=(128, 128) ,					down_block_types=('AttnDownBlock2D', 'DownBlock2D') ,					up_block_types=('UpBlock2D', 'AttnUpBlock2D') ,					)
        return model

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        torch.manual_seed(0					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     UNetaDConditionModel(
            sample_size=(64, 32) ,					in_channels=1 ,					out_channels=1 ,					layers_per_block=2 ,					block_out_channels=(128, 128) ,					down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D') ,					up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D') ,					cross_attention_dim=10 ,					)
        return model

  @property
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        torch.manual_seed(0					)
        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     AutoencoderKL(
            sample_size=(128, 64) ,					in_channels=1 ,					out_channels=1 ,					latent_channels=1 ,					layers_per_block=2 ,					block_out_channels=(128, 128) ,					down_block_types=('DownEncoderBlock2D', 'DownEncoderBlock2D') ,					up_block_types=('UpDecoderBlock2D', 'UpDecoderBlock2D') ,					)
        SCREAMING_SNAKE_CASE_				:  Dict           =     UNetaDModel(
            sample_size=(64, 32) ,					in_channels=1 ,					out_channels=1 ,					layers_per_block=2 ,					block_out_channels=(128, 128) ,					down_block_types=('AttnDownBlock2D', 'DownBlock2D') ,					up_block_types=('UpBlock2D', 'AttnUpBlock2D') ,					)
        return vqvae, unet



  @slow
  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     'cpu'  # ensure determinism for the device-dependent torch.Generator
        SCREAMING_SNAKE_CASE_				:  Any           =     Mel(
            x_res=self.dummy_unet.config.sample_size[1] ,					y_res=self.dummy_unet.config.sample_size[0] ,					)

        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     DDPMScheduler()
        SCREAMING_SNAKE_CASE_				:  str           =     AudioDiffusionPipeline(vqvae=_SCREAMING_SNAKE_CASE ,					unet=self.dummy_unet ,					mel=_SCREAMING_SNAKE_CASE ,					scheduler=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  Dict           =     pipe.to(_SCREAMING_SNAKE_CASE					)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

        SCREAMING_SNAKE_CASE_				:  Dict           =     torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(42					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     pipe(generator=_SCREAMING_SNAKE_CASE ,					steps=4					)
        SCREAMING_SNAKE_CASE_				:  Any           =     output.audios[0]
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     output.images[0]

        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(42					)
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     pipe(generator=_SCREAMING_SNAKE_CASE ,					steps=4 ,					return_dict=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     output[0][0]

        assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
        assert (
            image.height == self.dummy_unet.config.sample_size[0]
            and image.width == self.dummy_unet.config.sample_size[1]
        )
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     np.frombuffer(image.tobytes() ,					dtype='uint8'					)[:10]
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     np.frombuffer(image_from_tuple.tobytes() ,					dtype='uint8'					)[:10]
        SCREAMING_SNAKE_CASE_				:  Any           =     np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]					)

        assert np.abs(image_slice.flatten() - expected_slice					).max() == 0
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice					).max() == 0

        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     Mel(
            x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1] ,					y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0] ,					)

        SCREAMING_SNAKE_CASE_				:  str           =     DDIMScheduler()
        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     self.dummy_vqvae_and_unet
        SCREAMING_SNAKE_CASE_				:  int           =     AudioDiffusionPipeline(
            vqvae=self.dummy_vqvae_and_unet[0] ,					unet=dummy_vqvae_and_unet[1] ,					mel=_SCREAMING_SNAKE_CASE ,					scheduler=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  str           =     pipe.to(_SCREAMING_SNAKE_CASE					)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

        np.random.seed(0					)
        SCREAMING_SNAKE_CASE_				:  Dict           =     np.random.uniform(-1 ,					1 ,					((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)					)
        SCREAMING_SNAKE_CASE_				:  str           =     torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(42					)
        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     pipe(raw_audio=_SCREAMING_SNAKE_CASE ,					generator=_SCREAMING_SNAKE_CASE ,					start_step=5 ,					steps=10					)
        SCREAMING_SNAKE_CASE_				:  int           =     output.images[0]

        assert (
            image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
            and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
        )
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     np.frombuffer(image.tobytes() ,					dtype='uint8'					)[:10]
        SCREAMING_SNAKE_CASE_				:  int           =     np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]					)

        assert np.abs(image_slice.flatten() - expected_slice					).max() == 0

        SCREAMING_SNAKE_CASE_				:  Dict           =     self.dummy_unet_condition
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     AudioDiffusionPipeline(
            vqvae=self.dummy_vqvae_and_unet[0] ,					unet=_SCREAMING_SNAKE_CASE ,					mel=_SCREAMING_SNAKE_CASE ,					scheduler=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  Dict           =     pipe.to(_SCREAMING_SNAKE_CASE					)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

        np.random.seed(0					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     torch.rand((1, 1, 10)					)
        SCREAMING_SNAKE_CASE_				:  Any           =     pipe(generator=_SCREAMING_SNAKE_CASE ,					encoding=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     output.images[0]
        SCREAMING_SNAKE_CASE_				:  Dict           =     np.frombuffer(image.tobytes() ,					dtype='uint8'					)[:10]
        SCREAMING_SNAKE_CASE_				:  Union[str, Any]           =     np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]					)

        assert np.abs(image_slice.flatten() - expected_slice					).max() == 0

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

  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()



  def  UpperCAmelCase    (				self					):





        """simple docstring"""







        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     torch_device

        SCREAMING_SNAKE_CASE_				:  str           =     DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256'					)
        SCREAMING_SNAKE_CASE_				:  Dict           =     pipe.to(_SCREAMING_SNAKE_CASE					)
        pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE					)

        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     torch.Generator(device=_SCREAMING_SNAKE_CASE					).manual_seed(42					)
        SCREAMING_SNAKE_CASE_				:  List[Any]           =     pipe(generator=_SCREAMING_SNAKE_CASE					)
        SCREAMING_SNAKE_CASE_				:  Any           =     output.audios[0]
        SCREAMING_SNAKE_CASE_				:  Optional[Any]           =     output.images[0]

        assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
        assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
        SCREAMING_SNAKE_CASE_				:  Optional[int]           =     np.frombuffer(image.tobytes() ,					dtype='uint8'					)[:10]
        SCREAMING_SNAKE_CASE_				:  Any           =     np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]					)

        assert np.abs(image_slice.flatten() - expected_slice					).max() == 0






