render diff

app.py

@@ -46,8 +46,7 @@ def generate_diffusion_intermediates_streaming(label):
         x = x.clamp(-1, 1)
 
         outputs = plot_diff(outputs,x,t,noise_pred)
-
-        if t % 10 == 0:
+        if t % 10 == 0 or t in [499, 399, 299, 199, 99, 0,400, 300, 200, 100, 1]:
             yield tuple(outputs)
             time.sleep(0.06)
 
@@ -90,7 +89,7 @@ def generate_flow_intermediates_streaming(label,noise_type):
         outputs = plot_flow(outputs,i,x,dt,v)
         if i % 2 == 0:
             yield tuple(outputs)
-            time.sleep(0.2) # sleep to render properly in gradio
+            time.sleep(0.15) # sleep to render properly in gradio
         if ENV=="LOCAL":
             time.sleep(TIME_SLEEP)
     yield tuple(outputs)