keep video decode + trend rendering off the GPU; tighten ZeroGPU duration

app.py

@@ -202,7 +202,7 @@ def _render_trend_frames(values, sampled_indices, fps, size, task_title):
     actual = (x[-1] - x) / float(fps)
 
     ax1.plot(x, y, color="tab:blue", linewidth=1.0, alpha=0.25)
-    (line,) = ax1.plot([], [], color="tab:blue", linewidth=2.0, label="predicted remaining")
+    (line,) = ax1.plot([], [], color="tab:blue", linewidth=2.0, label="predicted (left axis)")
     pt1 = ax1.scatter([x[0]], [y[0]], color="red", s=32, zorder=3, label="current prediction")
     ax1.set_xlabel("Frame", fontsize=9)
     ax1.set_ylabel("Predicted value (s)", color="tab:blue", fontsize=9)
@@ -219,7 +219,9 @@ def _render_trend_frames(values, sampled_indices, fps, size, task_title):
     ax1.set_ylim(y_min - margin, y_max + margin)
 
     ax2 = ax1.twinx()
-    ax2.plot(x, actual, color="green", linestyle="--", linewidth=2.0, label="actual remaining")
+    ax2.plot(
+        x, actual, color="green", linestyle="--", linewidth=2.0, label="actual (right axis)"
+    )
     pt2 = ax2.scatter([x[0]], [actual[0]], color="green", s=26, zorder=3)
     ax2.set_ylabel("Actual remaining (s)", color="green", fontsize=9)
     ax2.tick_params(axis="y", labelcolor="green", labelsize=8)
@@ -335,67 +337,45 @@ def save_video_with_trend(frames, values, sampled_indices, fps, output_path, tas
 
 
 def _gpu_duration(
-    video=None,
     instruction="",
+    model_images=(),
     robot_description=DEFAULT_ROBOT,
     camera_description=DEFAULT_CAMERA,
-    num_frames=DEFAULT_NUM_FRAMES,
-    max_image_side=DEFAULT_MAX_IMAGE_SIDE,
     max_new_tokens=DEFAULT_MAX_NEW_TOKENS,
+    max_image_side=DEFAULT_MAX_IMAGE_SIDE,
     *args,
     **kwargs,
 ):
     try:
-        nf = int(num_frames)
+        nf = len(model_images)
         side = int(max_image_side)
     except (TypeError, ValueError):
         nf, side = DEFAULT_NUM_FRAMES, DEFAULT_MAX_IMAGE_SIDE
     # Sequence length scales with frames x tokens-per-frame, and the eager
     # value-isolation attention makes the forward grow ~quadratically in it.
-    # Calibrated against a measured 32-frame @384px run.
-    n = nf * (side / 384.0) ** 2
-    return int(min(240, 20 + 0.6 * n + 0.012 * n**2))
+    # Calibrated against measured 48-frame @384px runs (~8 s on GPU).
+    n = max(nf, 1) * (side / 384.0) ** 2
+    return int(min(120, 6 + 0.10 * n + 0.0016 * n**2))
 
 
 @spaces.GPU(duration=_gpu_duration)
-def analyze(
-    video,
+def _run_model(
     instruction,
-    robot_description=DEFAULT_ROBOT,
-    camera_description=DEFAULT_CAMERA,
-    num_frames=DEFAULT_NUM_FRAMES,
+    model_images,
+    robot_description,
+    camera_description,
+    max_new_tokens,
     max_image_side=DEFAULT_MAX_IMAGE_SIDE,
-    max_new_tokens=DEFAULT_MAX_NEW_TOKENS,
-    progress=gr.Progress(),
 ):
-    started = time.time()
-    instruction = (instruction or "").strip()
-    if not instruction:
-        raise gr.Error("Please describe the task the agent is supposed to perform.")
-
-    robot_description = (robot_description or "").strip() or None
-    camera_description = (camera_description or "").strip() or None
-    if robot_description is None and camera_description is None:
-        # The checkpoint was trained with use_meta=True; the conversation
-        # builder refuses to build a prompt without at least one meta field.
-        robot_description = DEFAULT_ROBOT
-
-    num_frames = int(num_frames)
-    max_image_side = int(max_image_side)
-    max_new_tokens = int(max_new_tokens)
+    """GPU half: build the episode prompt, read the value heads, write the Analysis.
 
-    progress(0.03, desc="Decoding video...")
-    frames, fps = load_video_frames(video)
-    total = len(frames)
-    model_images = resize_frames(frames, max_image_side)
-
-    num_frames = max(2, min(num_frames, total))
-    sampled_indices = np.linspace(0, total - 1, num_frames, dtype=int).tolist()
-
-    progress(0.18, desc="Building episode prompt...")
+    Video decoding and the trend rendering deliberately stay on the CPU side so
+    they don't consume ZeroGPU time.
+    """
+    started = time.time()
     sample = processor.process_episode(
         instruction=instruction,
-        images=[model_images[j] for j in sampled_indices],
+        images=list(model_images),
         robot_description=robot_description,
         camera_description=camera_description,
     )
@@ -404,7 +384,6 @@ def analyze(
     pixel_values = sample["pixel_values"].flatten(0, 1).to("cuda")
     image_grid_thw = sample["image_grid_thw"].flatten(0, 1).to("cuda").long()
 
-    progress(0.28, desc="Predicting temporal distance...")
     with torch.inference_mode():
         outputs = model(
             input_ids=input_ids,
@@ -412,22 +391,19 @@ def analyze(
             pixel_values=pixel_values,
             image_grid_thw=image_grid_thw,
         )
-    # (num_value_heads, num_frames) -> one remaining-time prediction per frame
+    # decode_from_bins yields (num_value_heads, num_slots); one slot per frame.
     pred_value = outputs.value.pred_value.float()
-    if pred_value.dim() > 1:
-        pred_value = pred_value.reshape(-1, len(sampled_indices)).mean(dim=0)
-    pred_value = pred_value.reshape(-1).tolist()
-    # relative head: signed time delta between consecutive sampled frames
+    pred_value = pred_value.reshape(-1, len(model_images)).mean(dim=0).tolist()
+    # relative head: signed time delta between consecutive presented frames
     rel_value = outputs.relative.pred_value.float().reshape(-1).tolist()
 
-    progress(0.6, desc="Generating analysis...")
     with torch.inference_mode():
         gen_out = model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
             pixel_values=pixel_values,
             image_grid_thw=image_grid_thw,
-            max_new_tokens=max_new_tokens,
+            max_new_tokens=int(max_new_tokens),
             do_sample=False,
             num_beams=1,
             eos_token_id=EOS_TOKEN_ID,
@@ -435,10 +411,54 @@ def analyze(
             use_cache=True,
         )
     analysis_text = tokenizer.decode(gen_out[0, input_ids.shape[1] :], skip_special_tokens=True)
+    return pred_value, rel_value, analysis_text, time.time() - started
+
+
+def analyze(
+    video,
+    instruction,
+    robot_description=DEFAULT_ROBOT,
+    camera_description=DEFAULT_CAMERA,
+    num_frames=DEFAULT_NUM_FRAMES,
+    max_image_side=DEFAULT_MAX_IMAGE_SIDE,
+    max_new_tokens=DEFAULT_MAX_NEW_TOKENS,
+    progress=gr.Progress(),
+):
+    instruction = (instruction or "").strip()
+    if not instruction:
+        raise gr.Error("Please describe the task the agent is supposed to perform.")
+
+    robot_description = (robot_description or "").strip() or None
+    camera_description = (camera_description or "").strip() or None
+    if robot_description is None and camera_description is None:
+        # The checkpoint was trained with use_meta=True; the conversation
+        # builder refuses to build a prompt without at least one meta field.
+        robot_description = DEFAULT_ROBOT
+
+    num_frames = int(num_frames)
+    max_image_side = int(max_image_side)
+    max_new_tokens = int(max_new_tokens)
+
+    progress(0.03, desc="Decoding video...")
+    frames, fps = load_video_frames(video)
+    total = len(frames)
+    model_images = resize_frames(frames, max_image_side)
+
+    num_frames = max(2, min(num_frames, total))
+    sampled_indices = np.linspace(0, total - 1, num_frames, dtype=int).tolist()
+
+    progress(0.2, desc="Predicting temporal distance...")
+    pred_value, rel_value, analysis_text, gpu_seconds = _run_model(
+        instruction,
+        [model_images[j] for j in sampled_indices],
+        robot_description,
+        camera_description,
+        max_new_tokens,
+        max_image_side,
+    )
     analysis = parse_analysis(analysis_text)
-    gpu_seconds = time.time() - started
 
-    progress(0.85, desc="Rendering trend video...")
+    progress(0.75, desc="Rendering trend video...")
     out_path = os.path.join(tempfile.mkdtemp(), "rynnvalue_trend.mp4")
     save_video_with_trend(frames, pred_value, sampled_indices, fps, out_path, instruction)
 
@@ -457,11 +477,15 @@ def analyze(
             f"**Matches the instruction?** {_verdict(analysis['match'])}   ·   "
             f"**Task completed?** {_verdict(analysis['success'])}",
             "",
-            f"**Temporal distance** — {pred_value[0]:.2f} s remaining at the first sampled frame → "
-            f"{pred_value[-1]:.2f} s at the last. The clip itself is {clip_seconds:.1f} s long.",
+            f"**Temporal distance (absolute head)** — cost-to-go falls from "
+            f"**{max(pred_value):.2f} s** to **{min(pred_value):.2f} s**, ending at "
+            f"**{pred_value[-1]:.2f} s** on the final frame. Values are the model's learned "
+            f"estimate of time-to-completion, not a rescaling of the clip's own "
+            f"{clip_seconds:.1f} s length.",
             "",
-            f"**Relative head** — the predicted time deltas between the {num_frames} sampled "
-            f"frames sum to {predicted_span:.1f} s (actual clip length {clip_seconds:.1f} s).",
+            f"**Relative head** — the signed Δt it predicts between the {num_frames} presented "
+            f"frames sums to {predicted_span:.1f} s, against {clip_seconds:.1f} s of real elapsed "
+            "time.",
             "",
             f"<sub>{num_frames} frames · longest side {max_image_side}px · "
             f"{gpu_seconds:.1f} s on GPU</sub>",