feat: stream generation + assistant progress over SSE

Convert generate_image / prompt_assistant into progress-yielding generators
so each yield surfaces as an SSE `generating` frame on /gradio_api/call. A
hidden JSON progress output (index 3 for generate_image, index 1 for
prompt_assistant) carries {stage,p,step,total,label} for downstream orchestration.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

app.py

@@ -1,7 +1,10 @@
+import inspect
 import io
 import os
+import queue
 import random
 import re
+import threading
 import time
 
 import numpy as np
@@ -16,7 +19,24 @@ from diffusers import (
 )
 from compel import Compel, ReturnedEmbeddingsType
 from huggingface_hub import hf_hub_download
-from transformers import AutoProcessor, AutoModelForImageTextToText
+from transformers import AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer
+
+# Structured progress contract (also see wan2.2 / LTX2.3 Spaces and the generator
+# orchestrator). Every streaming endpoint yields a hidden JSON "progress" output
+# in addition to its real outputs. Each yielded progress dict carries:
+#   {"stage": <phase id>, "p": <0..1 fraction>, "step": int, "total": int,
+#    "label": <human text>}
+# Because the function is a generator, Gradio surfaces every yield as an
+# `event: generating` frame on the /gradio_api/call SSE stream, so a downstream
+# consumer reading the JSON at the progress index gets live progress over SSE.
+def _progress(stage, p, step=0, total=0, label=""):
+    return {
+        "stage": stage,
+        "p": max(0.0, min(1.0, float(p))),
+        "step": int(step),
+        "total": int(total),
+        "label": label,
+    }
 
 import r2_uploader
 
@@ -334,6 +354,14 @@ def get_embed_new(prompt, pipeline, compel, only_convert_string=False, compel_pr
 # =============================================================================
 # Generation
 # =============================================================================
+def _supports_step_callback(pipe):
+    """True if this diffusers pipeline's __call__ accepts callback_on_step_end."""
+    try:
+        return "callback_on_step_end" in inspect.signature(pipe.__call__).parameters
+    except (TypeError, ValueError):
+        return False
+
+
 @spaces.GPU
 def generate_image(
     model_name,
@@ -348,19 +376,53 @@ def generate_image(
     randomize_seed,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Generate an image from the given prompt using the selected model."""
+    """Generate an image, streaming per-step progress.
+
+    This is a generator (so ZeroGPU streams its yields back over SSE). It yields
+    ``("progress", step, total)`` tuples during sampling and a final
+    ``("image", image, seed)`` tuple. The sampler runs in a worker thread feeding
+    a queue so the main thread can yield progress as each diffusion step lands.
+    """
     _gpu_start = time.time()
+    total_steps = int(num_inference_steps)
+    q = queue.Queue()
+    result = {}
+
+    def _step_cb(_pipe, step, _timestep, callback_kwargs):
+        # diffusers calls this after each step; `step` is the 0-based index.
+        q.put(step + 1)
+        return callback_kwargs
+
+    def _run():
+        try:
+            result["image"], result["seed"] = _generate_image_inner(
+                model_name, prompt, negative_prompt, use_negative_prompt,
+                height, width, total_steps, guidance_scale, seed, randomize_seed,
+                callback=_step_cb,
+            )
+        except Exception as exc:  # noqa: BLE001 - surfaced to the main thread
+            result["error"] = exc
+        finally:
+            q.put(None)  # sentinel: generation finished (ok or error)
+
+    thread = threading.Thread(target=_run, daemon=True)
+    thread.start()
     try:
-        return _generate_image_inner(
-            model_name, prompt, negative_prompt, use_negative_prompt,
-            height, width, num_inference_steps, guidance_scale, seed, randomize_seed,
-        )
+        while True:
+            step = q.get()
+            if step is None:
+                break
+            yield ("progress", step, total_steps)
     finally:
+        thread.join()
         print(
             f"[ImageStudio] GPU time consumed: {time.time() - _gpu_start:.2f}s "
             f"(model={model_name}, steps={num_inference_steps}, {int(width)}x{int(height)})",
             flush=True,
         )
+    if "error" in result:
+        raise result["error"]
+    yield ("image", result["image"], result["seed"])
 
 
 def _generate_image_inner(
@@ -374,6 +436,7 @@ def _generate_image_inner(
     guidance_scale,
     seed,
     randomize_seed,
+    callback=None,
 ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -405,7 +468,7 @@ def _generate_image_inner(
             [cond_prompt, cond_negative], precomputed_padding=empty_padding
         )
 
-        image = noobxl_pipe(
+        kwargs = dict(
             prompt_embeds=cond_prompt,
             pooled_prompt_embeds=pooled_prompt,
             negative_prompt_embeds=cond_negative,
@@ -416,19 +479,25 @@ def _generate_image_inner(
             num_inference_steps=int(num_inference_steps),
             generator=generator,
             use_resolution_binning=True,
-        ).images[0]
+        )
+        if callback is not None and _supports_step_callback(noobxl_pipe):
+            kwargs["callback_on_step_end"] = callback
+        image = noobxl_pipe(**kwargs).images[0]
         return image, seed
 
     # Default: Z-Image-Turbo (guidance-free distilled model)
     generator = torch.Generator("cuda").manual_seed(seed)
-    image = zimage_pipe(
+    kwargs = dict(
         prompt=prompt,
         height=int(height),
         width=int(width),
         num_inference_steps=int(num_inference_steps),
         guidance_scale=0.0,
         generator=generator,
-    ).images[0]
+    )
+    if callback is not None and _supports_step_callback(zimage_pipe):
+        kwargs["callback_on_step_end"] = callback
+    image = zimage_pipe(**kwargs).images[0]
     return image, seed
 
 
@@ -439,52 +508,82 @@ def _generate_image_inner(
 def vlm_chat(message, image, reasoning, max_new_tokens, progress=gr.Progress(track_tqdm=True)):
     """Answer a single user message, optionally grounded on an uploaded image.
 
+    Generator: yields ``("progress", produced, budget)`` as tokens stream in and a
+    final ``("text", answer)`` tuple. Token streaming (TextIteratorStreamer + a
+    worker thread) is the canonical ZeroGPU pattern and lets the downstream
+    orchestrator track this node's progress over SSE.
+
     ``reasoning`` ("On"/"Off") drives Qwen's ``enable_thinking`` switch: Off skips
     the <think> trace for a direct answer (best for prompt rewriting); On lets the
     model reason step-by-step first (slower, needs more max_new_tokens).
     """
     message = (message or "").strip()
     if not message and image is None:
-        return "Please enter a question (and optionally attach an image)."
+        yield ("text", "Please enter a question (and optionally attach an image).")
+        return
 
     enable_thinking = (reasoning == "On")
+    budget = int(max_new_tokens)
     _gpu_start = time.time()
+
+    content = []
+    if image is not None:
+        content.append({"type": "image", "image": image})
+    content.append({"type": "text", "text": message or "Describe this image."})
+    messages = [{"role": "user", "content": content}]
+
+    inputs = vlm_processor.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_dict=True,
+        return_tensors="pt",
+        enable_thinking=enable_thinking,
+    ).to(vlm_model.device)
+
+    tokenizer = getattr(vlm_processor, "tokenizer", vlm_processor)
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
+    result = {}
+
+    def _run():
+        try:
+            with torch.inference_mode():
+                vlm_model.generate(
+                    **inputs,
+                    max_new_tokens=budget,
+                    do_sample=False,
+                    streamer=streamer,
+                )
+        except Exception as exc:  # noqa: BLE001 - surfaced to the main thread
+            result["error"] = exc
+            streamer.end()
+
+    thread = threading.Thread(target=_run, daemon=True)
+    thread.start()
     try:
-        content = []
-        if image is not None:
-            content.append({"type": "image", "image": image})
-        content.append({"type": "text", "text": message or "Describe this image."})
-        messages = [{"role": "user", "content": content}]
-
-        inputs = vlm_processor.apply_chat_template(
-            messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_dict=True,
-            return_tensors="pt",
-            enable_thinking=enable_thinking,
-        ).to(vlm_model.device)
-
-        with torch.inference_mode():
-            generated = vlm_model.generate(
-                **inputs,
-                max_new_tokens=int(max_new_tokens),
-                do_sample=False,
-            )
-        # Drop the prompt tokens so only the freshly generated answer is decoded.
-        trimmed = generated[0][inputs["input_ids"].shape[1]:]
-        text = vlm_processor.decode(trimmed, skip_special_tokens=True).strip()
-        # With reasoning off, drop any stray <think>…</think> block so the answer
-        # stays clean; with it on, keep the trace so the user can see it.
-        if not enable_thinking and "</think>" in text:
-            text = text.split("</think>")[-1].strip()
-        return text
+        text = ""
+        produced = 0
+        for chunk in streamer:
+            text += chunk
+            produced += 1
+            yield ("progress", produced, budget, text)
     finally:
+        thread.join()
         print(
             f"[ImageStudio] Assistant GPU time: {time.time() - _gpu_start:.2f}s "
-            f"(has_image={image is not None}, reasoning={reasoning}, max_new_tokens={int(max_new_tokens)})",
+            f"(has_image={image is not None}, reasoning={reasoning}, max_new_tokens={budget})",
             flush=True,
         )
+    if "error" in result:
+        raise result["error"]
+    text = text.strip()
+    # With reasoning off, drop any stray <think>…</think> block so the answer
+    # stays clean; with it on, keep the trace so the user can see it.
+    if not enable_thinking and "</think>" in text:
+        text = text.split("</think>")[-1].strip()
+    yield ("text", text)
 
 
 def generate_and_upload(
@@ -503,17 +602,30 @@ def generate_and_upload(
 ):
     """Generate, then upload the result to R2 outside the GPU window.
 
-    Returns ``(image, seed, r2_status)``. The image is always the original
+    Generator yielding ``(image, seed, r2_status, progress)``. Streams per-step
+    sampling progress (image still None) and finishes with the real image, seed
+    and R2 status once the upload completes. The image is always the original
     HF-generated asset; ``r2_status`` reports the uploaded filekey on success or
     the error on failure. The caller's unique id (``uid`` cookie) is recorded in
     the uploaded object's metadata.
     """
-    image, used = generate_image(
+    image, used = None, None
+    for ev in generate_image(
         model_name, prompt, negative_prompt, use_negative_prompt,
         height, width, num_inference_steps, guidance_scale, seed, randomize_seed,
         progress=progress,
-    )
+    ):
+        if ev[0] == "progress":
+            _, step, total = ev
+            # Reserve the last 5% of this node for the R2 upload that follows.
+            frac = (step / max(total, 1)) * 0.95
+            yield None, None, None, _progress(
+                "image", frac, step, total, f"Sampling {step}/{total}"
+            )
+        else:
+            _, image, used = ev
 
+    yield None, used, None, _progress("image", 0.97, label="Uploading")
     uid = r2_uploader.uid_from_request(request)
     buf = io.BytesIO()
     image.save(buf, format="PNG")
@@ -541,7 +653,27 @@ def generate_and_upload(
         status = {"r2_filekey": result["filekey"], "r2_bucket": result["bucket"]}
     else:
         status = {"r2_error": result.get("error", "unknown error")}
-    return image, used, status
+    yield image, used, status, _progress("done", 1.0, label="Done")
+
+
+def assistant_chat(
+    message, image, reasoning, max_new_tokens,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Gradio-facing wrapper around ``vlm_chat``.
+
+    Yields ``(answer, progress)`` so the UI streams the text live and a downstream
+    consumer reading the progress index sees this node advance over SSE. The final
+    ``complete`` frame carries the clean answer at index 0.
+    """
+    for ev in vlm_chat(message, image, reasoning, max_new_tokens, progress=progress):
+        if ev[0] == "progress":
+            _, produced, budget, partial = ev
+            frac = min(0.99, produced / max(budget, 1))
+            yield partial, _progress("prompt", frac, produced, budget, "Writing prompt")
+        else:
+            _, text = ev
+            yield text, _progress("done", 1.0, label="Done")
 
 
 # Recommended defaults per model: (steps, guidance, height, width)
@@ -724,6 +856,9 @@ with gr.Blocks(fill_height=True) as demo:
                         label="🎲 Seed Used", interactive=False, container=True,
                     )
                     r2_status = gr.JSON(label="☁️ R2 Upload")
+                    # Hidden structured-progress channel (index 3 of generate_image
+                    # outputs). Surfaces every yield as an SSE `generating` frame.
+                    gen_progress = gr.JSON(label="progress", visible=False)
 
         with gr.Tab("💬 Prompt Assistant"):
             gr.Markdown(
@@ -764,6 +899,9 @@ with gr.Blocks(fill_height=True) as demo:
                         label="🤖 Answer",
                         lines=20,
                     )
+                    # Hidden structured-progress channel (index 1 of prompt_assistant
+                    # outputs).
+                    vlm_progress = gr.JSON(label="progress", visible=False)
 
     gr.Markdown(
         """
@@ -795,20 +933,23 @@ with gr.Blocks(fill_height=True) as demo:
     # resolving even though the click now runs the upload wrapper.
     generate_btn.click(
         fn=generate_and_upload, inputs=gen_inputs,
-        outputs=[output_image, used_seed, r2_status], api_name="generate_image",
+        outputs=[output_image, used_seed, r2_status, gen_progress],
+        api_name="generate_image",
     )
     prompt.submit(
         fn=generate_and_upload, inputs=gen_inputs,
-        outputs=[output_image, used_seed, r2_status],
+        outputs=[output_image, used_seed, r2_status, gen_progress],
     )
 
     # Prompt Assistant (Qwen3.5-4B) — single-turn, optional image
     vlm_inputs = [vlm_prompt, vlm_image, vlm_reasoning, vlm_max_tokens]
     vlm_btn.click(
-        fn=vlm_chat, inputs=vlm_inputs, outputs=[vlm_output],
+        fn=assistant_chat, inputs=vlm_inputs, outputs=[vlm_output, vlm_progress],
         api_name="prompt_assistant",
     )
-    vlm_prompt.submit(fn=vlm_chat, inputs=vlm_inputs, outputs=[vlm_output])
+    vlm_prompt.submit(
+        fn=assistant_chat, inputs=vlm_inputs, outputs=[vlm_output, vlm_progress],
+    )
 
 if __name__ == "__main__":
     demo.launch(