generator-based live step status (track_tqdm doesnt forward on zerogpu); fix gen_mask tqdm crash

app.py

@@ -5,6 +5,8 @@ os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
 
 import random
 import tempfile
+import threading
+import time
 
 import cv2
 import numpy as np
@@ -82,8 +84,7 @@ def _sam3_video_masks(frames_pil, text):
     session = sam3_processor.add_text_prompt(inference_session=session, text=text)
     n = len(frames_pil)
     masks = np.zeros((n, H, W), dtype=bool)
-    for mo in tqdm(sam3.propagate_in_video_iterator(inference_session=session, max_frame_num_to_track=n),
-                   total=n, desc="SAM3 tracking"):
+    for mo in sam3.propagate_in_video_iterator(inference_session=session, max_frame_num_to_track=n):
         proc = sam3_processor.postprocess_outputs(session, mo)
         m = proc.get("masks")
         if m is not None and len(m):
@@ -163,9 +164,9 @@ def gen_mask(video, mask_text, preset, progress=gr.Progress(track_tqdm=True)):
     return _write_mask_video(masks)
 
 
-# --- Inference --------------------------------------------------------------
+# --- Inference (generator: yields live step status across the ZeroGPU boundary) ---
 @spaces.GPU(duration=_duration)
-def inpaint(video, mask_video, prompt, preset, num_frames, seed, randomize, progress=gr.Progress(track_tqdm=True)):
+def inpaint(video, mask_video, prompt, preset, num_frames, seed, randomize):
     if video is None:
         raise gr.Error("Please upload a video.")
     if mask_video is None:
@@ -195,16 +196,41 @@ def inpaint(video, mask_video, prompt, preset, num_frames, seed, randomize, prog
     am[masks] = 0.0
     attn_mask = torch.from_numpy(am)[None, None]
 
-    video_out, audio_out = pipe(
-        prompt=prompt.strip(), negative_prompt="",
-        reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
-        conditioning_attention_mask=attn_mask, reference_downscale_factor=1,
-        width=width, height=height, num_frames=num_frames, frame_rate=FPS,
-        num_inference_steps=NUM_STEPS, sigmas=DISTILLED_SIGMA_VALUES,
-        guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
-        generator=torch.Generator(device="cuda").manual_seed(seed),
-        output_type="np", return_dict=False,
-    )
+    state = {"step": 0}
+    holder = {}
+
+    def _cb(p, i, t, kw):
+        state["step"] = i + 1
+        return {}
+
+    def _run():
+        try:
+            holder["out"] = pipe(
+                prompt=prompt.strip(), negative_prompt="",
+                reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
+                conditioning_attention_mask=attn_mask, reference_downscale_factor=1,
+                width=width, height=height, num_frames=num_frames, frame_rate=FPS,
+                num_inference_steps=NUM_STEPS, sigmas=DISTILLED_SIGMA_VALUES,
+                guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
+                generator=torch.Generator(device="cuda").manual_seed(seed),
+                output_type="np", return_dict=False, callback_on_step_end=_cb,
+            )
+        except Exception as e:  # surface in main thread
+            holder["err"] = e
+
+    th = threading.Thread(target=_run)
+    th.start()
+    yield None, gr.update(), "### ⏳ Preparing…"
+    while th.is_alive():
+        s = state["step"]
+        msg = f"### 🪄 Denoising — step {s}/{NUM_STEPS}" if s else "### ⏳ Loading model / encoding…"
+        yield gr.update(), gr.update(), msg
+        time.sleep(0.4)
+    th.join()
+    if "err" in holder:
+        raise holder["err"]
+    yield gr.update(), gr.update(), "### 🎬 Decoding & encoding video…"
+    video_out, audio_out = holder["out"]
 
     gen = (np.clip(video_out[0], 0, 1) * 255).astype(np.uint8)
     orig_arr = np.stack(orig).astype(np.float32)
@@ -216,7 +242,7 @@ def inpaint(video, mask_video, prompt, preset, num_frames, seed, randomize, prog
 
     out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
     _export(out, audio_out, out_path)
-    return out_path, seed
+    yield out_path, seed, "### ✅ Done"
 
 
 # --- UI ---------------------------------------------------------------------
@@ -249,6 +275,7 @@ with gr.Blocks(title="LTX-2.3 Video Inpainting") as demo:
             run = gr.Button("Inpaint", variant="primary")
         with gr.Column():
             video_out = gr.Video(label="Inpainted result")
+            status = gr.Markdown("")
             used_seed = gr.Number(label="Seed used", interactive=False)
 
     # a new input video invalidates any stored base mask
@@ -263,7 +290,7 @@ with gr.Blocks(title="LTX-2.3 Video Inpainting") as demo:
     dilate_px.release(apply_dilation, inputs=[base_mask, mask_video, dilate_px], outputs=[mask_video, base_mask])
 
     run.click(inpaint, inputs=[video_in, mask_video, prompt, preset, num_frames, seed, randomize],
-              outputs=[video_out, used_seed])
+              outputs=[video_out, used_seed, status])
 
     gr.Examples(
         examples=[
@@ -275,7 +302,7 @@ with gr.Blocks(title="LTX-2.3 Video Inpainting") as demo:
              "Fast (768×448)", 49, 42, False],
         ],
         inputs=[video_in, mask_video, prompt, preset, num_frames, seed, randomize],
-        outputs=[video_out, used_seed], fn=inpaint, cache_examples=True, cache_mode="lazy",
+        outputs=[video_out, used_seed, status], fn=inpaint, cache_examples=True, cache_mode="lazy",
     )
 
 if __name__ == "__main__":