continous streaming generation

app.py

@@ -9,7 +9,7 @@ from src.model import ConditionalUNet
 from huggingface_hub import hf_hub_download
 import time
 device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
-device = 'cpu'
+#device = 'cpu'
 img_shape = (1, 28, 28)
 
 
@@ -73,13 +73,20 @@ def generate_diffusion_intermediates_streaming(label):
             outputs[step_idx] = resize(vel_colored)
             yield tuple(outputs)
 
+        outputs[12] = resize(((x + 1) / 2.0)[0, 0].cpu().numpy(),(300,300))
+
         if t in [400, 300, 200, 100, 1, 0]:
             step_idx = {400: 1, 300: 2, 200: 3, 100: 4, 1: 5, 0 :12}[t]
             if t==0:
-                outputs[step_idx] = resize(((x + 1) / 2.0)[0, 0].cpu().numpy(),(400,400))
+                outputs[step_idx] = resize(((x + 1) / 2.0)[0, 0].cpu().numpy(),(300,300))
             else:
                 outputs[step_idx] = resize(((x + 1) / 2.0)[0, 0].cpu().numpy())
             yield tuple(outputs)
+        if t % 10 == 0:
+            yield tuple(outputs)
+            time.sleep(0.05)
+        #time.sleep(0.1)
+    yield tuple(outputs)
 
 
 def generate_localized_noise(shape, radius=5):
@@ -127,9 +134,9 @@ def generate_flow_intermediates_streaming(label):
         v = model_flow(x, t, y)
         x = x + v * dt
 
-        if i in [10,20,30,40,48,49]:
-            #images.append((x + 1) / 2.0)
-            step_idx = {10: 1, 20: 2, 30: 3, 40: 4, 48: 5, 49:12}[i]
+        outputs[12] =  resize(((x + 1) / 2.0)[0, 0].clamp(0, 1).cpu().numpy(),(300,300))
+        if i in [10,20,30,40,48,49]: #
+            step_idx = {10: 1, 20: 2, 30: 3, 40: 4, 48: 5,49:12}[i] #, 
             if i==49:
                 outputs[step_idx] = resize(((x + 1) / 2.0)[0, 0].clamp(0, 1).cpu().numpy(),(300,300))
             else:
@@ -146,6 +153,12 @@ def generate_flow_intermediates_streaming(label):
             step_idx = {0: 6, 11: 7, 21: 8, 31: 9, 41: 10, 49:11}[i]
             outputs[step_idx] = resize(vel_colored)
             yield tuple(outputs)
+        if t % 10 == 0:
+            yield tuple(outputs)
+            time.sleep(0.05)
+
+        #time.sleep(0.1)
+    yield tuple(outputs)
 
 
 with gr.Blocks() as demo:
@@ -207,5 +220,5 @@ with gr.Blocks() as demo:
             ]
         btn_f.click(fn=generate_flow_intermediates_streaming, inputs=label_f, outputs=outs_f+flow_vel_imgs+flow_result_imgs)
 
-#demo.launch()
-demo.launch(share=False, server_port=9071)
+demo.launch()
+#demo.launch(share=False, server_port=9071)