longcat: report the real denoise step, over every segment (v0.2.87)

A 304-frame request sat at "generating 0/16 (0%)" for its whole nine-minute
first segment while the service's log showed Denoising: 15/16. Two defects.

The service set _progress["step"] = 0 at the start of every segment and never
again, so the step count existed only inside the vendored pipeline's
tqdm(total=len(timesteps), desc="Denoising") -- which takes no callback.
Nothing reached GET /progress; the poller in video.py faithfully reported the
0 it was handed. Each pipeline module did `from tqdm import tqdm`, so the bar
is instrumented by rebinding the name on the MODULE (patching tqdm.tqdm would
be too late), and only bars labelled "Denoising" are tracked so a
checkpoint-shard bar cannot move a request's progress.

The totals also counted ONE segment: a four-segment render would have run
0->16 four times against a denominator of 16. `step` now carries the completed
segments' offset so it counts the stage end to end, and the bar spans
steps x segments -- with refinement excluded, being a single pass over the
finished video. A request that asked for a duration rather than a segment
count only learns the real count from the service, so the denominator is
corrected on the first report, through a setter that does not restart the bar
and lose the elapsed time the rate is derived from.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.2.86"
+__version__ = "0.2.87"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/api/video.py

@@ -86,6 +86,16 @@ def _vid_progress_done():
     _vid_progress["active"] = False
     _vid_progress["phase"] = "idle"
 
+def _vid_progress_total(total: int):
+    """Correct the denominator mid-render, without resetting the bar.
+
+    Some paths only learn how much work there is after the worker has started -- a
+    LongCat request that asked for a duration rather than a segment count is told the
+    segment count by the service -- and restarting the bar there would throw away the
+    elapsed time the rate is derived from."""
+    _vid_progress["total"] = int(total)
+
+
 def _vid_progress_step(step: int):
     _vid_progress["current"] = step
     elapsed = time.monotonic() - _vid_progress["started_at"]
@@ -4240,10 +4250,22 @@ async def _generate_longcat(request: VideoGenerationRequest, model_name: str,
                     else "source" if _k == "longcat_source" else _k] = str(model_cfg[_k])
 
     # Total steps across the stages we are about to run, so the bar spans the whole
-    # request rather than restarting per stage.
-    _total = sum(lc.stage_params(s, request.num_inference_steps,
-                                 request.guidance_scale)["num_inference_steps"]
-                 for s in stages)
+    # request rather than restarting per stage -- and across SEGMENTS, because a long
+    # video runs the generative stages once per segment. Counting a segment's worth
+    # made a four-segment render read 0/16 four times over instead of 0/64.
+    def _grand_total(segments: int) -> int:
+        segments = max(1, int(segments or 1))
+        return sum(
+            lc.stage_params(s, request.num_inference_steps,
+                            request.guidance_scale)["num_inference_steps"]
+            * (1 if s == "refinement" else segments)
+            for s in stages)
+
+    # Refinement is a single pass over the finished video, so it is not multiplied.
+    # The service owns the real segment count (it derives one from total_frames when
+    # the caller asked for a duration); the poll loop corrects this once it reports.
+    _segments_assumed = max(1, int(payload.get("num_segments") or 1))
+    _total = _grand_total(_segments_assumed)
     _vid_progress_reset(_total)
 
     _stop = threading.Event()
@@ -4254,7 +4276,7 @@ async def _generate_longcat(request: VideoGenerationRequest, model_name: str,
     # claimed in its own docstring to "carry cancellation across the process
     # boundary" and never did, so a LongCat render — the longest-running video job
     # there is — was the only one nothing could stop.
-    _lc_steps = sum(lc.STAGE_DEFAULTS.get(_s, (0, 0))[0] for _s in stages) or 1
+    _lc_steps = _total or 1
     _tid = task_registry.register(
         "video", title=(payload.get("prompt") or mode or "")[:80],
         model=model_name or "", total=_lc_steps)
@@ -4270,6 +4292,10 @@ async def _generate_longcat(request: VideoGenerationRequest, model_name: str,
         returns with an error. That is also exactly what eviction does to it, so the
         path is already exercised."""
         _seen_stage = ""
+        # What the stage now running reports as its own total, segments included, so a
+        # finished stage contributes what it actually ran rather than a default.
+        _stage_total = 0
+        _live = {"segments": _segments_assumed, "total": _total}
         while not _stop.wait(2.0):
             try:
                 if not _cancelled["flag"]:
@@ -4292,12 +4318,21 @@ async def _generate_longcat(request: VideoGenerationRequest, model_name: str,
                 st = str(p.get("stage") or "")
                 if st and st != _seen_stage:
                     if _seen_stage:
-                        _done_steps["base"] += lc.STAGE_DEFAULTS.get(
+                        _done_steps["base"] += _stage_total or lc.STAGE_DEFAULTS.get(
                             _seen_stage, (0, 0))[0]
                     _seen_stage = st
+                    _stage_total = 0
+                _stage_total = int(p.get("steps") or 0) or _stage_total
+                _segs_now = int(p.get("segments") or 0)
+                if _segs_now and _segs_now != _live["segments"]:
+                    _live["segments"] = _segs_now
+                    _live["total"] = _grand_total(_segs_now)
+                    _vid_progress_total(_live["total"])
+                # `step` counts the whole stage, every segment of it: the service adds
+                # the completed segments' offset before reporting the bar's position.
                 _step = _done_steps["base"] + int(p.get("step") or 0)
                 _vid_progress_step(_step)
-                task_registry.step(_tid, _step)
+                task_registry.step(_tid, _step, total=_live["total"])
             except Exception:
                 pass
 

tools/longcat_service.py

@@ -70,6 +70,13 @@ _state = {
 }
 _progress = {"active": False, "stage": "", "segment": 0, "segments": 0,
              "step": 0, "steps": 0}
+# Denoise-step bookkeeping for /progress. The vendored pipelines render their step
+# count onto a tqdm bar and nothing else, so until this was wired the bar the caller
+# polls reported 0/N for an entire render. ``offset`` is how many steps earlier
+# SEGMENTS of the current stage already contributed, so ``_progress["step"]`` counts
+# the stage end to end instead of restarting at every segment boundary;
+# ``pass_steps`` is one segment's step count, learnt from the first bar's total.
+_bar = {"offset": 0, "pass_steps": 0}
 # Set by POST /yield: stop at the next SEGMENT boundary and return what has been
 # generated, so another model can have the card without the request being lost.
 _yield_flag = threading.Event()
@@ -79,6 +86,58 @@ def log(msg):
     print(f"[longcat] {msg}", flush=True)
 
 
+def _install_progress_bars() -> None:
+    """Teach the pipelines' denoising bars to publish their step count.
+
+    Upstream's generate_* methods drive a bare ``tqdm(total=len(timesteps),
+    desc="Denoising")`` and take no callback, so the only way to observe a step
+    without forking the vendored pipeline is to replace the ``tqdm`` the pipeline
+    module resolved at import time. Each module did ``from tqdm import tqdm``, so the
+    name has to be rebound on the MODULE -- patching ``tqdm.tqdm`` would be too late.
+
+    Only bars labelled "Denoising" are tracked: the pipelines use tqdm for nothing
+    else today, and a checkpoint-shard bar must not move the request's progress."""
+    import tqdm as _tqdm_mod
+
+    base = _tqdm_mod.tqdm
+
+    class _ProgressTqdm(base):
+        # Checked instead of class identity: this function builds a fresh subclass on
+        # every call, so identity would re-wrap an already-wrapped bar.
+        _lc_progress_bar = True
+
+        def __init__(self, *a, **kw):
+            super().__init__(*a, **kw)
+            self._lc_track = str(kw.get("desc") or "") == "Denoising"
+            if not self._lc_track:
+                return
+            with _lock:
+                total = int(self.total or 0)
+                if total:
+                    _bar["pass_steps"] = total
+                segments = max(1, int(_progress.get("segments") or 1))
+                _progress["steps"] = total * segments
+                _progress["step"] = int(_bar["offset"])
+
+        def update(self, n=1):
+            ret = super().update(n)
+            if getattr(self, "_lc_track", False):
+                with _lock:
+                    _progress["step"] = int(_bar["offset"]) + int(self.n or 0)
+            return ret
+
+    patched = []
+    for name in ("longcat_video.pipeline_longcat_video",
+                 "longcat_video.pipeline_longcat_video_avatar"):
+        mod = sys.modules.get(name)
+        if mod is not None and not getattr(
+                getattr(mod, "tqdm", None), "_lc_progress_bar", False):
+            mod.tqdm = _ProgressTqdm
+            patched.append(name.rsplit(".", 1)[-1])
+    if patched:
+        log(f"progress: denoising bars instrumented ({', '.join(patched)})")
+
+
 # ── loading ───────────────────────────────────────────────────────────────────
 
 def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
@@ -226,6 +285,9 @@ def load_pipeline(checkpoint_dir: str, source_dir: str, dtype: str = "bfloat16",
                 raise RuntimeError(
                     "use_distill needs the DiT's enable_loras() — this build of "
                     "longcat_video does not expose it")
+        # Both pipeline modules are imported by now (the avatar one only for an
+        # avatar family), so this is the first point where the bars can be rebound.
+        _install_progress_bars()
         _place_pipeline(pipe, offload)
         _state.update(pipe=pipe, model=root, source=src, dtype=dtype)
         log("ready")
@@ -1097,10 +1159,15 @@ def _run_segments(pipe, task: str, stage: str, ctx: dict, log=print):
 
     for index in range(segments):
         with _lock:
-            _progress.update(segment=index + 1, segments=segments, step=0,
-                             steps=LC.stage_params(
-                                 stage, ctx.get("num_inference_steps"),
-                                 ctx.get("guidance_scale"))["num_inference_steps"])
+            # One segment's step count comes from the first bar's own total, which is
+            # authoritative (the avatar distilled pass is 8 steps, the base one 16).
+            # Before that bar exists, the stage default is the estimate.
+            per_pass = int(_bar["pass_steps"]) or int(LC.stage_params(
+                stage, ctx.get("num_inference_steps"),
+                ctx.get("guidance_scale"))["num_inference_steps"])
+            _bar["offset"] = index * per_pass
+            _progress.update(segment=index + 1, segments=segments,
+                             step=index * per_pass, steps=per_pass * segments)
         out = _one_pass(pipe, task, stage, ctx,
                         cond=cur if (index or first_is_continuation) else None, log=log)
         new = _frames_from_output(out)
@@ -1354,6 +1421,8 @@ def generate(body: dict) -> dict:
     frames = None
     yielded = 0
     with _lock:
+        _bar["offset"] = 0
+        _bar["pass_steps"] = 0
         _progress.update(active=True, stage="", segment=0, segments=segments,
                          step=0, steps=0)
     try:
@@ -1377,6 +1446,12 @@ def generate(body: dict) -> dict:
                 if image is not None:
                     kw["image"] = image
                     kw["num_cond_frames"] = 1
+                with _lock:
+                    # Refinement is one pass over the whole video, so the segment
+                    # offset from the stage before it must not carry over.
+                    _bar["offset"] = 0
+                    _progress.update(segment=1, segments=1, step=0,
+                                     steps=params["num_inference_steps"])
                 log(f"stage 'refinement': refining {len(frames)} frames")
                 frames = _frames_from_output(pipe.generate_refine(**kw)[0])
             else:
@@ -1386,6 +1461,7 @@ def generate(body: dict) -> dict:
     finally:
         _yield_flag.clear()
         with _lock:
+            _bar["offset"] = 0
             _progress.update(active=False, stage="", step=0)
 
     if stages[-1] == "refinement":