UltraEdit/diffusers/tests/schedulers/test_scheduler_consistency_model.py

import torch

from diffusers import CMStochasticIterativeScheduler

from .test_schedulers import SchedulerCommonTest


class CMStochasticIterativeSchedulerTest(SchedulerCommonTest):
    scheduler_classes = (CMStochasticIterativeScheduler,)
    num_inference_steps = 10

    def get_scheduler_config(self, **kwargs):
        config = {
            "num_train_timesteps": 201,
            "sigma_min": 0.002,
            "sigma_max": 80.0,
        }

        config.update(**kwargs)
        return config

    # Override test_step_shape to add CMStochasticIterativeScheduler-specific logic regarding timesteps
    # Problem is that we don't know two timesteps that will always be in the timestep schedule from only the scheduler
    # config; scaled sigma_max is always in the timestep schedule, but sigma_min is in the sigma schedule while scaled
    # sigma_min is not in the timestep schedule
    def test_step_shape(self):
        num_inference_steps = 10

        scheduler_config = self.get_scheduler_config()
        scheduler = self.scheduler_classes[0](**scheduler_config)

        scheduler.set_timesteps(num_inference_steps)

        timestep_0 = scheduler.timesteps[0]
        timestep_1 = scheduler.timesteps[1]

        sample = self.dummy_sample
        residual = 0.1 * sample

        output_0 = scheduler.step(residual, timestep_0, sample).prev_sample
        output_1 = scheduler.step(residual, timestep_1, sample).prev_sample

        self.assertEqual(output_0.shape, sample.shape)
        self.assertEqual(output_0.shape, output_1.shape)

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

    def test_clip_denoised(self):
        for clip_denoised in [True, False]:
            self.check_over_configs(clip_denoised=clip_denoised)

    def test_full_loop_no_noise_onestep(self):
        scheduler_class = self.scheduler_classes[0]
        scheduler_config = self.get_scheduler_config()
        scheduler = scheduler_class(**scheduler_config)

        num_inference_steps = 1
        scheduler.set_timesteps(num_inference_steps)
        timesteps = scheduler.timesteps

        generator = torch.manual_seed(0)

        model = self.dummy_model()
        sample = self.dummy_sample_deter * scheduler.init_noise_sigma

        for i, t in enumerate(timesteps):
            # 1. scale model input
            scaled_sample = scheduler.scale_model_input(sample, t)

            # 2. predict noise residual
            residual = model(scaled_sample, t)

            # 3. predict previous sample x_t-1
            pred_prev_sample = scheduler.step(residual, t, sample, generator=generator).prev_sample

            sample = pred_prev_sample

        result_sum = torch.sum(torch.abs(sample))
        result_mean = torch.mean(torch.abs(sample))

        assert abs(result_sum.item() - 192.7614) < 1e-2
        assert abs(result_mean.item() - 0.2510) < 1e-3

    def test_full_loop_no_noise_multistep(self):
        scheduler_class = self.scheduler_classes[0]
        scheduler_config = self.get_scheduler_config()
        scheduler = scheduler_class(**scheduler_config)

        timesteps = [106, 0]
        scheduler.set_timesteps(timesteps=timesteps)
        timesteps = scheduler.timesteps

        generator = torch.manual_seed(0)

        model = self.dummy_model()
        sample = self.dummy_sample_deter * scheduler.init_noise_sigma

        for t in timesteps:
            # 1. scale model input
            scaled_sample = scheduler.scale_model_input(sample, t)

            # 2. predict noise residual
            residual = model(scaled_sample, t)

            # 3. predict previous sample x_t-1
            pred_prev_sample = scheduler.step(residual, t, sample, generator=generator).prev_sample

            sample = pred_prev_sample

        result_sum = torch.sum(torch.abs(sample))
        result_mean = torch.mean(torch.abs(sample))

        assert abs(result_sum.item() - 347.6357) < 1e-2
        assert abs(result_mean.item() - 0.4527) < 1e-3

    def test_full_loop_with_noise(self):
        scheduler_class = self.scheduler_classes[0]
        scheduler_config = self.get_scheduler_config()
        scheduler = scheduler_class(**scheduler_config)

        num_inference_steps = 10
        t_start = 8

        scheduler.set_timesteps(num_inference_steps)
        timesteps = scheduler.timesteps

        generator = torch.manual_seed(0)

        model = self.dummy_model()
        sample = self.dummy_sample_deter * scheduler.init_noise_sigma

        noise = self.dummy_noise_deter
        timesteps = scheduler.timesteps[t_start * scheduler.order :]

        sample = scheduler.add_noise(sample, noise, timesteps[:1])

        for t in timesteps:
            # 1. scale model input
            scaled_sample = scheduler.scale_model_input(sample, t)

            # 2. predict noise residual
            residual = model(scaled_sample, t)

            # 3. predict previous sample x_t-1
            pred_prev_sample = scheduler.step(residual, t, sample, generator=generator).prev_sample

            sample = pred_prev_sample

        result_sum = torch.sum(torch.abs(sample))
        result_mean = torch.mean(torch.abs(sample))

        assert abs(result_sum.item() - 763.9186) < 1e-2, f" expected result sum 763.9186, but get {result_sum}"
        assert abs(result_mean.item() - 0.9947) < 1e-3, f" expected result mean 0.9947, but get {result_mean}"

    def test_custom_timesteps_increasing_order(self):
        scheduler_class = self.scheduler_classes[0]
        scheduler_config = self.get_scheduler_config()
        scheduler = scheduler_class(**scheduler_config)

        timesteps = [39, 30, 12, 15, 0]

        with self.assertRaises(ValueError, msg="`timesteps` must be in descending order."):
            scheduler.set_timesteps(timesteps=timesteps)

    def test_custom_timesteps_passing_both_num_inference_steps_and_timesteps(self):
        scheduler_class = self.scheduler_classes[0]
        scheduler_config = self.get_scheduler_config()
        scheduler = scheduler_class(**scheduler_config)

        timesteps = [39, 30, 12, 1, 0]
        num_inference_steps = len(timesteps)

        with self.assertRaises(ValueError, msg="Can only pass one of `num_inference_steps` or `timesteps`."):
            scheduler.set_timesteps(num_inference_steps=num_inference_steps, timesteps=timesteps)

    def test_custom_timesteps_too_large(self):
        scheduler_class = self.scheduler_classes[0]
        scheduler_config = self.get_scheduler_config()
        scheduler = scheduler_class(**scheduler_config)

        timesteps = [scheduler.config.num_train_timesteps]

        with self.assertRaises(
            ValueError,
            msg="`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}",
        ):
            scheduler.set_timesteps(timesteps=timesteps)