[object Object]

← back to Exo

Image streaming while generation

4874295b34af0e6e441925e1fe6c74a542f9a417 · 2024-11-20 18:08:54 -0500 · Pranav Veldurthi

Files touched

Diff

commit 4874295b34af0e6e441925e1fe6c74a542f9a417
Author: Pranav Veldurthi <veldurthipranav@gmail.com>
Date:   Wed Nov 20 18:08:54 2024 -0500

    Image streaming while generation
---
 exo/api/chatgpt_api.py                              | 4 ++--
 exo/inference/mlx/models/StableDiffusionPipeline.py | 5 +++--
 exo/orchestration/standard_node.py                  | 2 +-
 exo/tinychat/index.js                               | 6 +++++-
 4 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/exo/api/chatgpt_api.py b/exo/api/chatgpt_api.py
index e53448c0..8f98aa6f 100644
--- a/exo/api/chatgpt_api.py
+++ b/exo/api/chatgpt_api.py
@@ -425,8 +425,8 @@ class ChatGPTAPI:
             full_image_url = base_url + str(image_url)
             
             await response.write(json.dumps({'images': [{'url': str(full_image_url), 'content_type': 'image/png'}]}).encode('utf-8') + b'\n')
-
-            await response.write_eof()
+            if is_finished:
+              await response.write_eof()
               
 
       stream_task = None
diff --git a/exo/inference/mlx/models/StableDiffusionPipeline.py b/exo/inference/mlx/models/StableDiffusionPipeline.py
index 443887a3..2868da95 100644
--- a/exo/inference/mlx/models/StableDiffusionPipeline.py
+++ b/exo/inference/mlx/models/StableDiffusionPipeline.py
@@ -198,6 +198,7 @@ class Model(nn.Module):
         if self.shard_unet.start_layer != -1:
             with tqdm(total=total_steps,initial=step+1) as pbar:
                 if step<total_steps:
+                    x = x_t_prev
                     if self.shard_unet.is_first_layer():
                         x_t_unet = mx.concatenate([x] * 2, axis=0) if cfg_weight> 1 else x
                     else:
@@ -214,7 +215,6 @@ class Model(nn.Module):
                     
         if self.shard_vae.is_last_layer():
             is_step_finished=True
-        if t_prev.item() ==0:
             if self.shard_vae.start_layer != -1:
                 x=self.first_stage_model.decode(x)
             if self.shard_vae.is_last_layer():
@@ -224,7 +224,8 @@ class Model(nn.Module):
                 x = x.reshape(1, B // 1, H, W, C).transpose(0, 2, 1, 3, 4)
                 x = x.reshape(1 * H, B // 1 * W, C)
                 x = (x * 255).astype(mx.uint8)
-                is_finished=True   
+                if t_prev.item() ==0:
+                    is_finished=True   
          
         return x, {'conditioning':conditioning, 'mask':mask,'residual':residual,'x_t_prev':x_t_prev,'is_finished':is_finished,'is_step_finished':is_step_finished, 'step':step, 'total_steps':total_steps}
     
diff --git a/exo/orchestration/standard_node.py b/exo/orchestration/standard_node.py
index 055d3fdb..7a72c325 100644
--- a/exo/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -477,7 +477,7 @@ class StandardNode(Node):
     if inference_state['is_step_finished']:
       inference_state['step']+=1
     progress = [inference_state['step'],inference_state['total_steps']]
-    intermediate_result = progress
+    intermediate_result = result
     if progress[0] == progress[1]:
       intermediate_result = result
     return intermediate_result, inference_state
diff --git a/exo/tinychat/index.js b/exo/tinychat/index.js
index 5e1fea9c..fb767369 100644
--- a/exo/tinychat/index.js
+++ b/exo/tinychat/index.js
@@ -272,9 +272,13 @@ document.addEventListener("alpine:init", () => {
                 this.cstate.messages[this.cstate.messages.length - 1].content = parsed.progress;
               }
               else if (parsed.images) {
+                if (!gottenFirstChunk) {
+                  this.cstate.messages.push({ role: "assistant", content: "" });
+                  gottenFirstChunk = true;
+                }
                 const imageUrl = parsed.images[0].url;
                 console.log(imageUrl)
-                this.cstate.messages[this.cstate.messages.length - 1].content = `![Generated Image](${imageUrl})`;
+                this.cstate.messages[this.cstate.messages.length - 1].content = `![Generated Image](${imageUrl}?t=${Date.now()})`;
               }
             }
           }

← 90765922 added one file  ·  back to Exo  ·  test clean ups 5396f080 →