Skip to content

Commit f82ebb9

Browse files
patrickvonplatenpatrickvonplaten
committed
fix some model tests
1 parent 63c68d9 commit f82ebb9

5 files changed

Lines changed: 89 additions & 10 deletions

File tree

‎src/diffusers/pipelines/ddpm/pipeline_ddpm.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def __call__(self, batch_size=1, generator=None, torch_device=None):
3535

3636
# Sample gaussian noise to begin loop
3737
image = torch.randn(
38-
(batch_size, self.unet.in_channels, self.unet.resolution, self.unet.resolution),
38+
(batch_size, self.unet.in_channels, self.unet.image_size, self.unet.image_size),
3939
generator=generator,
4040
)
4141
image = image.to(torch_device)

‎src/diffusers/schedulers/scheduling_ddim.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,9 @@ def __init__(
9090
self.num_inference_steps = None
9191
self.timesteps = np.arange(0, num_train_timesteps)[::-1].copy()
9292

93+
self.tensor_format = tensor_format
94+
self.set_format(tensor_format=tensor_format)
95+
9396
def _get_variance(self, timestep, prev_timestep):
9497
alpha_prod_t = self.alphas_cumprod[timestep]
9598
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.one
@@ -102,9 +105,9 @@ def _get_variance(self, timestep, prev_timestep):
102105

103106
def set_timesteps(self, num_inference_steps):
104107
self.num_inference_steps = num_inference_steps
105-
self.timesteps = np.arange(0, self.config.timesteps, self.config.timesteps // self.num_inference_steps)[
106-
::-1
107-
].copy()
108+
self.timesteps = np.arange(
109+
0, self.config.num_train_timesteps, self.config.num_train_timesteps // self.num_inference_steps
110+
)[::-1].copy()
108111
self.set_format(tensor_format=self.tensor_format)
109112

110113
def step(

‎src/diffusers/schedulers/scheduling_sde_ve.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
# DISCLAIMER: This file is strongly influenced by https://github.com/yang-song/score_sde_pytorch
1616

1717
# TODO(Patrick, Anton, Suraj) - make scheduler framework indepedent and clean-up a bit
18-
import pdb
1918

2019
import numpy as np
2120
import torch
@@ -110,10 +109,10 @@ def step_pred(self, score, x, t):
110109
Predict the sample at the previous timestep by reversing the SDE.
111110
"""
112111
# TODO(Patrick) better comments + non-PyTorch
113-
t = self.repeat_scalar(t, x.shape[0])
114-
timesteps = self.long((t * (len(self.timesteps) - 1)))
112+
t = self.repeat_scalar(t, x.shape[0]).to(x.device)
113+
timesteps = self.long((t * (len(self.timesteps) - 1))).to(x.device)
115114

116-
sigma = self.discrete_sigmas[timesteps]
115+
sigma = self.discrete_sigmas[timesteps].to(x.device)
117116
adjacent_sigma = self.get_adjacent_sigma(timesteps, t)
118117
drift = self.zeros_like(x)
119118
diffusion = (sigma**2 - adjacent_sigma**2) ** 0.5

‎tests/test_modeling_utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -911,7 +911,7 @@ def test_from_pretrained_save_pretrained(self):
911911
down_blocks=("UNetResDownBlock2D", "UNetResAttnDownBlock2D"),
912912
up_blocks=("UNetResAttnUpBlock2D", "UNetResUpBlock2D"),
913913
)
914-
schedular = DDPMScheduler(timesteps=10)
914+
schedular = DDPMScheduler(num_train_timesteps=10)
915915

916916
ddpm = DDPMPipeline(model, schedular)
917917

‎tests/test_scheduler.py‎

Lines changed: 78 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
15-
import pdb
1615
import tempfile
1716
import unittest
1817

@@ -618,3 +617,81 @@ def test_full_loop_no_noise(self):
618617

619618
assert abs(result_sum.item() - 10629923278.7104) < 1e-2
620619
assert abs(result_mean.item() - 13841045.9358) < 1e-3
620+
621+
def test_from_pretrained_save_pretrained(self):
622+
kwargs = dict(self.forward_default_kwargs)
623+
624+
num_inference_steps = kwargs.pop("num_inference_steps", None)
625+
626+
for scheduler_class in self.scheduler_classes:
627+
sample = self.dummy_sample
628+
residual = 0.1 * sample
629+
630+
scheduler_config = self.get_scheduler_config()
631+
scheduler = scheduler_class(**scheduler_config)
632+
633+
with tempfile.TemporaryDirectory() as tmpdirname:
634+
scheduler.save_config(tmpdirname)
635+
new_scheduler = scheduler_class.from_config(tmpdirname)
636+
637+
if num_inference_steps is not None and hasattr(scheduler, "set_timesteps"):
638+
scheduler.set_timesteps(num_inference_steps)
639+
new_scheduler.set_timesteps(num_inference_steps)
640+
elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"):
641+
kwargs["num_inference_steps"] = num_inference_steps
642+
643+
output = scheduler.step_pred(residual, 1, sample, **kwargs)["prev_sample"]
644+
new_output = new_scheduler.step_pred(residual, 1, sample, **kwargs)["prev_sample"]
645+
646+
assert np.sum(np.abs(output - new_output)) < 1e-5, "Scheduler outputs are not identical"
647+
648+
def test_step_shape(self):
649+
kwargs = dict(self.forward_default_kwargs)
650+
651+
num_inference_steps = kwargs.pop("num_inference_steps", None)
652+
653+
for scheduler_class in self.scheduler_classes:
654+
scheduler_config = self.get_scheduler_config()
655+
scheduler = scheduler_class(**scheduler_config)
656+
657+
sample = self.dummy_sample
658+
residual = 0.1 * sample
659+
660+
if num_inference_steps is not None and hasattr(scheduler, "set_timesteps"):
661+
scheduler.set_timesteps(num_inference_steps)
662+
elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"):
663+
kwargs["num_inference_steps"] = num_inference_steps
664+
665+
output_0 = scheduler.step_pred(residual, 0, sample, **kwargs)["prev_sample"]
666+
output_1 = scheduler.step_pred(residual, 1, sample, **kwargs)["prev_sample"]
667+
668+
self.assertEqual(output_0.shape, sample.shape)
669+
self.assertEqual(output_0.shape, output_1.shape)
670+
671+
def test_pytorch_equal_numpy(self):
672+
kwargs = dict(self.forward_default_kwargs)
673+
674+
num_inference_steps = kwargs.pop("num_inference_steps", None)
675+
676+
for scheduler_class in self.scheduler_classes:
677+
sample = self.dummy_sample
678+
residual = 0.1 * sample
679+
680+
sample_pt = torch.tensor(sample)
681+
residual_pt = 0.1 * sample_pt
682+
683+
scheduler_config = self.get_scheduler_config()
684+
scheduler = scheduler_class(**scheduler_config)
685+
686+
scheduler_pt = scheduler_class(tensor_format="pt", **scheduler_config)
687+
688+
if num_inference_steps is not None and hasattr(scheduler, "set_timesteps"):
689+
scheduler.set_timesteps(num_inference_steps)
690+
scheduler_pt.set_timesteps(num_inference_steps)
691+
elif num_inference_steps is not None and not hasattr(scheduler, "set_timesteps"):
692+
kwargs["num_inference_steps"] = num_inference_steps
693+
694+
output = scheduler.step_pred(residual, 1, sample, **kwargs)["prev_sample"]
695+
output_pt = scheduler_pt.step_pred(residual_pt, 1, sample_pt, **kwargs)["prev_sample"]
696+
697+
assert np.sum(np.abs(output - output_pt.numpy())) < 1e-4, "Scheduler outputs are not identical"

0 commit comments

Comments
 (0)