revert to plain track_tqdm progress (Z-Image pattern), drop generator/callback

app.py

@@ -93,38 +93,9 @@ def _duration(*args, **kwargs):
     return int(120 + int(num_frames) * per_frame)
 
 
-
-class StreamRun:
-    """Run pipe() in a thread; yield live status strings (generator yields DO forward on ZeroGPU)."""
-    def __init__(self, call_pipe, num_steps):
-        self.call_pipe, self.num_steps = call_pipe, num_steps
-        self.state = {"step": 0}
-        self.holder = {}
-    def _cb(self, p, i, t, kw):
-        self.state["step"] = i + 1
-        return {}
-    def _run(self):
-        try:
-            self.holder["out"] = self.call_pipe(self._cb)
-        except Exception as e:
-            self.holder["err"] = e
-    def stream(self):
-        th = threading.Thread(target=self._run); th.start()
-        while th.is_alive():
-            s = self.state["step"]
-            yield (s / self.num_steps if s else 0.0, f"step {s}/{self.num_steps}" if s else "Loading model…")
-            time.sleep(0.4)
-        th.join()
-        if "err" in self.holder:
-            raise self.holder["err"]
-    @property
-    def result(self):
-        return self.holder["out"]
-
-
 @spaces.GPU(duration=_duration)
 def shave(video, prompt, preset, num_frames, seed, randomize,
-          progress=gr.Progress()):
+          progress=gr.Progress(track_tqdm=True)):
     if video is None:
         raise gr.Error("Please upload a video of a bearded subject.")
     if randomize:
@@ -140,8 +111,7 @@ def shave(video, prompt, preset, num_frames, seed, randomize,
     full_prompt = _build_prompt(prompt)
 
 
-    def _call_pipe(_cb):
-        return pipe(
+    video_out, audio_out = pipe(
         prompt=full_prompt, negative_prompt=NEGATIVE,
         reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
         reference_downscale_factor=1,
@@ -149,16 +119,11 @@ def shave(video, prompt, preset, num_frames, seed, randomize,
         num_inference_steps=NUM_STEPS, guidance_scale=GUIDANCE,
         spatio_temporal_guidance_blocks=STG_BLOCKS,
         generator=torch.Generator(device="cuda").manual_seed(seed),
-        output_type="np", return_dict=False, callback_on_step_end=_cb,
+        output_type="np", return_dict=False,
         )
-    runner = StreamRun(_call_pipe, NUM_STEPS)
-    for _frac, _desc in runner.stream():
-        progress(_frac, desc=_desc)
-        yield gr.update(), gr.update()
-    video_out, audio_out = runner.result
     out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
     _export(video_out[0], audio_out, out_path)
-    yield out_path, seed
+    return out_path, seed
 
 
 with gr.Blocks(title="LTX-2.3 Beard Removal") as demo: