Update core/models/ddim/ddim_vd.py

core/models/ddim/ddim_vd.py

@@ -6,6 +6,7 @@ import torch
 import numpy as np
 from tqdm import tqdm
 from functools import partial
+import streamlit as st
 
 from .diffusion_utils import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
 
@@ -27,7 +28,8 @@ class DDIMSampler_VD(DDIMSampler):
                mix_weight=None,
                noise_dropout=0.,
                verbose=True,
-               log_every_t=100, ):
+               log_every_t=100,
+               progress_bar=False, ):
 
         self.make_schedule(ddim_num_steps=steps, ddim_eta=eta, verbose=verbose)
         print(f'Data shape for DDIM sampling is {shape}, eta {eta}')
@@ -42,7 +44,8 @@ class DDIMSampler_VD(DDIMSampler):
             noise_dropout=noise_dropout,
             temperature=temperature,
             log_every_t=log_every_t,
-            mix_weight=mix_weight, )
+            mix_weight=mix_weight,
+            progress_bar=progress_bar, )
         return samples, intermediates
 
     @torch.no_grad()
@@ -58,7 +61,8 @@ class DDIMSampler_VD(DDIMSampler):
                       noise_dropout=0.,
                       temperature=1.,
                       mix_weight=None,
-                      log_every_t=100, ):
+                      log_every_t=100,
+                      progress_bar=False,):
 
         device = self.model.device
         dtype = condition[0][0].dtype
@@ -86,7 +90,12 @@ class DDIMSampler_VD(DDIMSampler):
 
         pred_xt = xt
         iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+        if progress_bar is not None:
+            progress_bar.progress(0)
+            progress_bar.text("Generating samples...")
         for i, step in enumerate(iterator):
+            if progress_bar is not None:
+                progress_bar.progress(i/total_steps)
             index = total_steps - i - 1
             ts = torch.full((bs,), step, device=device, dtype=torch.long)
 
@@ -107,6 +116,9 @@ class DDIMSampler_VD(DDIMSampler):
                 intermediates['pred_xt'].append(pred_xt)
                 intermediates['pred_x0'].append(pred_x0)
 
+        if progress_bar is not None:
+            progress_bar.success("Sampling complete.")
+
         return pred_xt, intermediates
 
     @torch.no_grad()
@@ -172,4 +184,4 @@ class DDIMSampler_VD(DDIMSampler):
             x_prev_i = a_prev.sqrt() * pred_x0_i + dir_xt + noise
             x_prev.append(x_prev_i)
             pred_x0.append(pred_x0_i)
-        return x_prev, pred_x0
+        return x_prev, pred_x0