330 lines
11 KiB
Python
Executable File
330 lines
11 KiB
Python
Executable File
# coding=utf-8
|
|
# Copyright 2022 HuggingFace Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
|
|
import tempfile
|
|
import unittest
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from diffusers import DDIMScheduler, GaussianDDPMScheduler
|
|
|
|
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
|
|
|
|
class SchedulerCommonTest(unittest.TestCase):
|
|
scheduler_classes = ()
|
|
forward_default_kwargs = ()
|
|
|
|
@property
|
|
def dummy_image(self):
|
|
batch_size = 4
|
|
num_channels = 3
|
|
height = 8
|
|
width = 8
|
|
|
|
image = np.random.rand(batch_size, num_channels, height, width)
|
|
|
|
return image
|
|
|
|
@property
|
|
def dummy_image_deter(self):
|
|
batch_size = 4
|
|
num_channels = 3
|
|
height = 8
|
|
width = 8
|
|
|
|
num_elems = batch_size * num_channels * height * width
|
|
image = np.arange(num_elems)
|
|
image = image.reshape(num_channels, height, width, batch_size)
|
|
image = image / num_elems
|
|
image = image.transpose(3, 0, 1, 2)
|
|
|
|
return image
|
|
|
|
def get_scheduler_config(self):
|
|
raise NotImplementedError
|
|
|
|
def dummy_model(self):
|
|
def model(image, t, *args):
|
|
return image * t / (t + 1)
|
|
|
|
return model
|
|
|
|
def check_over_configs(self, time_step=0, **config):
|
|
kwargs = dict(self.forward_default_kwargs)
|
|
|
|
for scheduler_class in self.scheduler_classes:
|
|
scheduler_class = self.scheduler_classes[0]
|
|
image = self.dummy_image
|
|
residual = 0.1 * image
|
|
|
|
scheduler_config = self.get_scheduler_config(**config)
|
|
scheduler = scheduler_class(**scheduler_config)
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
scheduler.save_config(tmpdirname)
|
|
new_scheduler = scheduler_class.from_config(tmpdirname)
|
|
|
|
output = scheduler.step(residual, image, time_step, **kwargs)
|
|
new_output = new_scheduler.step(residual, image, time_step, **kwargs)
|
|
|
|
assert np.sum(np.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical"
|
|
|
|
def check_over_forward(self, time_step=0, **forward_kwargs):
|
|
kwargs = dict(self.forward_default_kwargs)
|
|
kwargs.update(forward_kwargs)
|
|
|
|
for scheduler_class in self.scheduler_classes:
|
|
scheduler_class = self.scheduler_classes[0]
|
|
image = self.dummy_image
|
|
residual = 0.1 * image
|
|
|
|
scheduler_config = self.get_scheduler_config()
|
|
scheduler = scheduler_class(**scheduler_config)
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
scheduler.save_config(tmpdirname)
|
|
new_scheduler = scheduler_class.from_config(tmpdirname)
|
|
|
|
output = scheduler.step(residual, image, time_step, **kwargs)
|
|
new_output = new_scheduler.step(residual, image, time_step, **kwargs)
|
|
|
|
assert np.sum(np.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical"
|
|
|
|
def test_from_pretrained_save_pretrained(self):
|
|
kwargs = dict(self.forward_default_kwargs)
|
|
|
|
for scheduler_class in self.scheduler_classes:
|
|
image = self.dummy_image
|
|
residual = 0.1 * image
|
|
|
|
scheduler_config = self.get_scheduler_config()
|
|
scheduler = scheduler_class(**scheduler_config)
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
scheduler.save_config(tmpdirname)
|
|
new_scheduler = scheduler_class.from_config(tmpdirname)
|
|
|
|
output = scheduler.step(residual, image, 1, **kwargs)
|
|
new_output = new_scheduler.step(residual, image, 1, **kwargs)
|
|
|
|
assert np.sum(np.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical"
|
|
|
|
def test_step_shape(self):
|
|
kwargs = dict(self.forward_default_kwargs)
|
|
|
|
for scheduler_class in self.scheduler_classes:
|
|
scheduler_config = self.get_scheduler_config()
|
|
scheduler = scheduler_class(**scheduler_config)
|
|
|
|
image = self.dummy_image
|
|
residual = 0.1 * image
|
|
|
|
output_0 = scheduler.step(residual, image, 0, **kwargs)
|
|
output_1 = scheduler.step(residual, image, 1, **kwargs)
|
|
|
|
self.assertEqual(output_0.shape, image.shape)
|
|
self.assertEqual(output_0.shape, output_1.shape)
|
|
|
|
def test_pytorch_equal_numpy(self):
|
|
kwargs = dict(self.forward_default_kwargs)
|
|
|
|
for scheduler_class in self.scheduler_classes:
|
|
image = self.dummy_image
|
|
residual = 0.1 * image
|
|
|
|
image_pt = torch.tensor(image)
|
|
residual_pt = 0.1 * image_pt
|
|
|
|
scheduler_config = self.get_scheduler_config()
|
|
scheduler = scheduler_class(**scheduler_config)
|
|
|
|
scheduler_pt = scheduler_class(tensor_format="pt", **scheduler_config)
|
|
|
|
output = scheduler.step(residual, image, 1, **kwargs)
|
|
output_pt = scheduler_pt.step(residual_pt, image_pt, 1, **kwargs)
|
|
|
|
assert np.sum(np.abs(output - output_pt.numpy())) < 1e-5, "Scheduler outputs are not identical"
|
|
|
|
|
|
class DDPMSchedulerTest(SchedulerCommonTest):
|
|
scheduler_classes = (GaussianDDPMScheduler,)
|
|
|
|
def get_scheduler_config(self, **kwargs):
|
|
config = {
|
|
"timesteps": 1000,
|
|
"beta_start": 0.0001,
|
|
"beta_end": 0.02,
|
|
"beta_schedule": "linear",
|
|
"variance_type": "fixed_small",
|
|
"clip_predicted_image": True,
|
|
}
|
|
|
|
config.update(**kwargs)
|
|
return config
|
|
|
|
def test_timesteps(self):
|
|
for timesteps in [1, 5, 100, 1000]:
|
|
self.check_over_configs(timesteps=timesteps)
|
|
|
|
def test_betas(self):
|
|
for beta_start, beta_end in zip([0.0001, 0.001, 0.01, 0.1], [0.002, 0.02, 0.2, 2]):
|
|
self.check_over_configs(beta_start=beta_start, beta_end=beta_end)
|
|
|
|
def test_schedules(self):
|
|
for schedule in ["linear", "squaredcos_cap_v2"]:
|
|
self.check_over_configs(beta_schedule=schedule)
|
|
|
|
def test_variance_type(self):
|
|
for variance in ["fixed_small", "fixed_large", "other"]:
|
|
self.check_over_configs(variance_type=variance)
|
|
|
|
def test_clip_image(self):
|
|
for clip_predicted_image in [True, False]:
|
|
self.check_over_configs(clip_predicted_image=clip_predicted_image)
|
|
|
|
def test_time_indices(self):
|
|
for t in [0, 500, 999]:
|
|
self.check_over_forward(time_step=t)
|
|
|
|
def test_variance(self):
|
|
scheduler_class = self.scheduler_classes[0]
|
|
scheduler_config = self.get_scheduler_config()
|
|
scheduler = scheduler_class(**scheduler_config)
|
|
|
|
assert np.sum(np.abs(scheduler.get_variance(0) - 0.0)) < 1e-5
|
|
assert np.sum(np.abs(scheduler.get_variance(487) - 0.00979)) < 1e-5
|
|
assert np.sum(np.abs(scheduler.get_variance(999) - 0.02)) < 1e-5
|
|
|
|
def test_full_loop_no_noise(self):
|
|
scheduler_class = self.scheduler_classes[0]
|
|
scheduler_config = self.get_scheduler_config()
|
|
scheduler = scheduler_class(**scheduler_config)
|
|
|
|
num_trained_timesteps = len(scheduler)
|
|
|
|
model = self.dummy_model()
|
|
image = self.dummy_image_deter
|
|
|
|
for t in reversed(range(num_trained_timesteps)):
|
|
# 1. predict noise residual
|
|
residual = model(image, t)
|
|
|
|
# 2. predict previous mean of image x_t-1
|
|
pred_prev_image = scheduler.step(residual, image, t)
|
|
|
|
if t > 0:
|
|
noise = self.dummy_image_deter
|
|
variance = scheduler.get_variance(t) ** (0.5) * noise
|
|
|
|
image = pred_prev_image + variance
|
|
|
|
result_sum = np.sum(np.abs(image))
|
|
result_mean = np.mean(np.abs(image))
|
|
|
|
assert result_sum.item() - 732.9947 < 1e-3
|
|
assert result_mean.item() - 0.9544 < 1e-3
|
|
|
|
|
|
class DDIMSchedulerTest(SchedulerCommonTest):
|
|
scheduler_classes = (DDIMScheduler,)
|
|
forward_default_kwargs = (("num_inference_steps", 50), ("eta", 0.0))
|
|
|
|
def get_scheduler_config(self, **kwargs):
|
|
config = {
|
|
"timesteps": 1000,
|
|
"beta_start": 0.0001,
|
|
"beta_end": 0.02,
|
|
"beta_schedule": "linear",
|
|
"clip_predicted_image": True,
|
|
}
|
|
|
|
config.update(**kwargs)
|
|
return config
|
|
|
|
def test_timesteps(self):
|
|
for timesteps in [1, 5, 100, 1000]:
|
|
self.check_over_configs(timesteps=timesteps)
|
|
|
|
def test_betas(self):
|
|
for beta_start, beta_end in zip([0.0001, 0.001, 0.01, 0.1], [0.002, 0.02, 0.2, 2]):
|
|
self.check_over_configs(beta_start=beta_start, beta_end=beta_end)
|
|
|
|
def test_schedules(self):
|
|
for schedule in ["linear", "squaredcos_cap_v2"]:
|
|
self.check_over_configs(beta_schedule=schedule)
|
|
|
|
def test_clip_image(self):
|
|
for clip_predicted_image in [True, False]:
|
|
self.check_over_configs(clip_predicted_image=clip_predicted_image)
|
|
|
|
def test_time_indices(self):
|
|
for t in [1, 10, 49]:
|
|
self.check_over_forward(time_step=t)
|
|
|
|
def test_inference_steps(self):
|
|
for t, num_inference_steps in zip([1, 10, 50], [10, 50, 500]):
|
|
self.check_over_forward(time_step=t, num_inference_steps=num_inference_steps)
|
|
|
|
def test_eta(self):
|
|
for t, eta in zip([1, 10, 49], [0.0, 0.5, 1.0]):
|
|
self.check_over_forward(time_step=t, eta=eta)
|
|
|
|
def test_variance(self):
|
|
scheduler_class = self.scheduler_classes[0]
|
|
scheduler_config = self.get_scheduler_config()
|
|
scheduler = scheduler_class(**scheduler_config)
|
|
|
|
assert np.sum(np.abs(scheduler.get_variance(0, 50) - 0.0)) < 1e-5
|
|
assert np.sum(np.abs(scheduler.get_variance(21, 50) - 0.14771)) < 1e-5
|
|
assert np.sum(np.abs(scheduler.get_variance(49, 50) - 0.32460)) < 1e-5
|
|
assert np.sum(np.abs(scheduler.get_variance(0, 1000) - 0.0)) < 1e-5
|
|
assert np.sum(np.abs(scheduler.get_variance(487, 1000) - 0.00979)) < 1e-5
|
|
assert np.sum(np.abs(scheduler.get_variance(999, 1000) - 0.02)) < 1e-5
|
|
|
|
def test_full_loop_no_noise(self):
|
|
scheduler_class = self.scheduler_classes[0]
|
|
scheduler_config = self.get_scheduler_config()
|
|
scheduler = scheduler_class(**scheduler_config)
|
|
|
|
num_inference_steps, eta = 10, 0.1
|
|
num_trained_timesteps = len(scheduler)
|
|
|
|
inference_step_times = range(0, num_trained_timesteps, num_trained_timesteps // num_inference_steps)
|
|
|
|
model = self.dummy_model()
|
|
image = self.dummy_image_deter
|
|
|
|
for t in reversed(range(num_inference_steps)):
|
|
residual = model(image, inference_step_times[t])
|
|
|
|
pred_prev_image = scheduler.step(residual, image, t, num_inference_steps, eta)
|
|
|
|
variance = 0
|
|
if eta > 0:
|
|
noise = self.dummy_image_deter
|
|
variance = scheduler.get_variance(t, num_inference_steps) ** (0.5) * eta * noise
|
|
|
|
image = pred_prev_image + variance
|
|
|
|
result_sum = np.sum(np.abs(image))
|
|
result_mean = np.mean(np.abs(image))
|
|
|
|
assert result_sum.item() - 270.6214 < 1e-3
|
|
assert result_mean.item() - 0.3524 < 1e-3
|