Stream WanGP events and cache uploads

app.py

@@ -14,6 +14,7 @@ import shutil
 import tempfile
 import threading
 import json
+import time
 from pathlib import Path
 
 import gradio as gr
@@ -25,6 +26,7 @@ WAN2GP_ROOT  = _PERSISTENT / "Wan2GP"
 CKPTS_DIR    = WAN2GP_ROOT / "ckpts"
 LORAS_DIR    = WAN2GP_ROOT / "loras" / "ltx2"
 FINETUNES_DIR = WAN2GP_ROOT / "finetunes"
+UPLOADS_DIR  = _PERSISTENT / "sulphur_uploads"
 
 SULPHUR_ASSETS = [
     ("SulphurAI/Sulphur-2-base", "sulphur_distil_bf16.safetensors", CKPTS_DIR),
@@ -239,6 +241,19 @@ def _wangp_cli_args():
 setup()
 
 
+def cache_uploaded_image(image):
+    if not image:
+        return None
+    src = Path(str(image)).resolve()
+    if not src.exists():
+        raise gr.Error("Uploaded image is no longer available. Please upload it again.")
+    UPLOADS_DIR.mkdir(parents=True, exist_ok=True)
+    suffix = src.suffix or ".png"
+    dest = UPLOADS_DIR / f"input_{time.time_ns()}{suffix}"
+    shutil.copy2(src, dest)
+    return str(dest)
+
+
 def _ensure_session():
     """Create WanGPSession on first call; reuse on all subsequent calls."""
     global _session
@@ -290,14 +305,16 @@ class _QueueWriter(io.TextIOBase):
 
 
 @spaces.GPU
-def generate_video(image, prompt, resolution, steps, guidance_scale, seconds, seed):
-    if image is None:
+def generate_video(image_path_cached, prompt, resolution, steps, guidance_scale, seconds, seed):
+    if image_path_cached is None:
         raise gr.Error("Please upload an image.")
     if not prompt.strip():
         raise gr.Error("Please enter a prompt.")
 
     frames     = int(seconds * 25 // 8) * 8 + 1
-    image_path = str(Path(image.strip()).resolve())
+    image_path = str(Path(str(image_path_cached)).resolve())
+    if not Path(image_path).exists():
+        raise gr.Error("Uploaded image is no longer available. Please upload it again.")
     out_dir    = Path(tempfile.mkdtemp())
     out_file   = out_dir / "output.mp4"
 
@@ -317,56 +334,67 @@ def generate_video(image, prompt, resolution, steps, guidance_scale, seconds, se
         "loras_multipliers":    ["0.5"],
     }
 
-    log_q      = queue.Queue()
-    result_box = [None]
-    error_box  = [None]
-    real_out   = sys.stdout
-    real_err   = sys.stderr
-
-    def _run():
-        writer = _QueueWriter(log_q, real_out)
-        sys.stdout = writer
-        sys.stderr = writer  # tqdm and progress bars write to stderr
-        try:
-            print("[startup] Loading WanGP session...")
-            session = _ensure_session()
-            print("[startup] WanGP session ready. Starting generation...")
-            result_box[0] = session.run_task(task)
-        except Exception as exc:
-            error_box[0] = exc
-        finally:
-            sys.stdout = real_out
-            sys.stderr = real_err
-            log_q.put(None)  # sentinel
-
-    t = threading.Thread(target=_run, daemon=True)
-    t.start()
-
     log_lines = []
-    while True:
-        try:
-            line = log_q.get(timeout=180)
-        except queue.Empty:
-            yield None, "\n".join(log_lines) + "\n\n[TIMEOUT]"
-            return
-        if line is None:
-            break
-        stripped = line.strip()
+    result = None
+
+    def _append_log(line):
+        stripped = str(line or "").strip()
         if not stripped:
-            continue
+            return
         if log_lines and ("%" in stripped or "it/s" in stripped or "step" in stripped.lower()):
             log_lines[-1] = stripped
         else:
             log_lines.append(stripped)
+
+    _append_log("[startup] Loading WanGP session...")
+    yield None, "\n".join(log_lines[-30:])
+
+    try:
+        session = _ensure_session()
+        _append_log("[startup] WanGP session ready. Starting generation...")
         yield None, "\n".join(log_lines[-30:])
 
-    t.join(timeout=10)
+        job = session.submit_task(task)
+        last_event_at = time.time()
+        while True:
+            event = job.events.get(timeout=0.5)
+            if event is None:
+                if job.done:
+                    break
+                if time.time() - last_event_at > 30:
+                    _append_log("[waiting] WanGP is still running...")
+                    last_event_at = time.time()
+                    yield None, "\n".join(log_lines[-30:])
+                continue
+
+            last_event_at = time.time()
+            if event.kind == "stream":
+                data = event.data
+                _append_log(getattr(data, "text", data))
+            elif event.kind == "progress":
+                progress = event.data
+                current = getattr(progress, "current_step", None)
+                total = getattr(progress, "total_steps", None)
+                phase = getattr(progress, "phase", "progress")
+                status = getattr(progress, "status", "")
+                if current is not None and total is not None:
+                    _append_log(f"[{phase}] step {current}/{total} {status}")
+                else:
+                    _append_log(f"[{phase}] {status}")
+            elif event.kind == "error":
+                _append_log(f"[ERROR] {event.data}")
+            elif event.kind == "completed":
+                result = event.data
+                break
 
-    if error_box[0]:
-        yield None, "\n".join(log_lines) + f"\n\n[ERROR] {error_box[0]}"
+            yield None, "\n".join(log_lines[-30:])
+
+        if result is None:
+            result = job.result(timeout=1)
+    except Exception as exc:
+        yield None, "\n".join(log_lines) + f"\n\n[ERROR] {exc}"
         return
 
-    result      = result_box[0]
     output_file = None
 
     if result and result.artifacts:
@@ -393,6 +421,7 @@ def generate_video(image, prompt, resolution, steps, guidance_scale, seconds, se
 
 with gr.Blocks(title="Sulphur — Image to Video") as demo:
     gr.Markdown("# Sulphur — Image to Video\nUsing New 'Experimental' Lora v1")
+    image_state = gr.State(None)
     with gr.Row():
         with gr.Column(scale=1):
             image_in  = gr.Image(type="filepath", label="Input Image")
@@ -408,9 +437,15 @@ with gr.Blocks(title="Sulphur — Image to Video") as demo:
             video_out = gr.Video(label="Output Video")
             log_out   = gr.Textbox(label="Log", lines=10, interactive=False)
 
+    image_in.change(
+        fn=cache_uploaded_image,
+        inputs=[image_in],
+        outputs=[image_state],
+    )
+
     run_btn.click(
         fn=generate_video,
-        inputs=[image_in, prompt_in, resolution_dd, steps_sl, guidance_sl, seconds_sl, seed_num],
+        inputs=[image_state, prompt_in, resolution_dd, steps_sl, guidance_sl, seconds_sl, seed_num],
         outputs=[video_out, log_out],
     )