From 318eb35ea0d3f9aa30762b4dfa2368ae1ebbb0f3 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lincoln.stein@gmail.com>
Date: Sun, 9 Oct 2022 15:29:04 -0400
Subject: [PATCH] run make_schedule() if it hasn't already been called

- fixes #984
---
 ldm/models/diffusion/ksampler.py | 8 ++++++++
 ldm/models/diffusion/sampler.py  | 9 +++++++++
 2 files changed, 17 insertions(+)

diff --git a/ldm/models/diffusion/ksampler.py b/ldm/models/diffusion/ksampler.py
index a947296f9e..1e44fe1fdf 100644
--- a/ldm/models/diffusion/ksampler.py
+++ b/ldm/models/diffusion/ksampler.py
@@ -171,6 +171,14 @@ class KSampler(Sampler):
             if img_callback is not None:
                 img_callback(k_callback_values['x'],k_callback_values['i'])
 
+        # if make_schedule() hasn't been called, we do it now
+        if self.sigmas is None:
+            self.make_schedule(
+                ddim_num_steps=S,
+                ddim_eta = eta,
+                verbose = False,
+            )
+
         # sigmas are set up in make_schedule - we take the last steps items
         total_steps = len(self.sigmas)
         sigmas = self.sigmas[-S-1:]
diff --git a/ldm/models/diffusion/sampler.py b/ldm/models/diffusion/sampler.py
index 8b19f894e5..88cdc01974 100644
--- a/ldm/models/diffusion/sampler.py
+++ b/ldm/models/diffusion/sampler.py
@@ -20,6 +20,7 @@ from ldm.modules.diffusionmodules.util import (
 class Sampler(object):
     def __init__(self, model, schedule='linear', steps=None, device=None, **kwargs):
         self.model = model
+        self.ddim_timesteps = None
         self.ddpm_num_timesteps = steps
         self.schedule = schedule
         self.device   = device or choose_torch_device()
@@ -157,6 +158,14 @@ class Sampler(object):
         **kwargs,
     ):
 
+        # check to see if make_schedule() has run, and if not, run it
+        if self.ddim_timesteps is None:
+            self.make_schedule(
+                ddim_num_steps=S,
+                ddim_eta = eta,
+                verbose = False,
+            )
+
         ts = self.get_timesteps(S)
 
         # sampling