Fixes

abarankab · Jul 20, 2021 · 78cc338 · 78cc338
1 parent f6e2dd6
commit 78cc338
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/ddpm/diffusion.py b/ddpm/diffusion.py
@@ -5,6 +5,8 @@
 from functools import partial
 from copy import deepcopy
 
+from ddpm.ema import EMA
+from ddpm.utils import extract
 
 class GaussianDiffusion(nn.Module):
     __doc__ = r"""Gaussian Diffusion model. Forwarding through the module returns diffusion reversal scalar loss tensor.

diff --git a/ddpm/utils.py b/ddpm/utils.py
@@ -21,7 +21,7 @@ def preprocess_image(img):
         raise ValueError("incorrect number of channels in an image")
 
 
-def generate_video(samples_list, delay=100, id=0):
+def show_video(samples_list, delay=100, id=0):
     fig = plt.figure(figsize=(8, 8))
     images = []
     for sample in samples_list:
@@ -30,6 +30,7 @@ def generate_video(samples_list, delay=100, id=0):
     ani = animation.ArtistAnimation(fig, images, interval=delay)
     display(HTML(ani.to_html5_video()))
 
+
 def show_images(image_tensor, rows, cols=None, colorbar=False):
     if cols is None:
         cols = (len(image_tensor) + rows - 1) // rows
@@ -61,6 +62,7 @@ def show_images(image_tensor, rows, cols=None, colorbar=False):
 
     plt.show()
 
+
 def show_image(img, colorbar=False):
     if img.shape[0] == 3:
         if colorbar: