Larger batches with OOM fallback; size GPU duration from measurements

app.py

@@ -356,17 +356,29 @@ def _badge(flag):
 # Inference
 # ----------------------------------------------------------------------------
 def _gpu_duration(*args, **kwargs):
+    """Size the ZeroGPU reservation from the measured cost of one run.
+
+    Reference points measured on this Space (448 px, 24 frames/prefix,
+    batch 8): 32 prefixes over a 429-frame video = ~30 s wall clock end to end,
+    including decode and rendering. Cost is dominated by the value pass, which
+    scales with ``num_steps × num_frames`` and roughly with the square of the
+    image side (the eager attention is O(L²)).
+    """
     num_steps = kwargs.get("num_steps", DEFAULT_NUM_STEPS)
     num_frames = kwargs.get("num_frames", DEFAULT_NUM_FRAMES)
+    side = kwargs.get("max_image_side", DEFAULT_MAX_SIDE)
     if len(args) > 4:
         num_steps = args[4]
     if len(args) > 5:
         num_frames = args[5]
+    if len(args) > 6:
+        side = args[6]
     try:
-        work = int(num_steps) * int(num_frames)
+        work = min(int(num_steps) * int(num_frames), WORK_BUDGET)
+        factor = (float(side) / DEFAULT_MAX_SIDE) ** 2.5
     except Exception:
-        work = DEFAULT_NUM_STEPS * DEFAULT_NUM_FRAMES
-    return int(min(180, 45 + work * 0.035))
+        work, factor = DEFAULT_NUM_STEPS * DEFAULT_NUM_FRAMES, 1.0
+    return int(min(180, max(30, 18 + work * 0.045 * factor)))
 
 
 @spaces.GPU(duration=_gpu_duration)
@@ -448,7 +460,10 @@ def analyze(
     t_prep = time.perf_counter() - t1
 
     seq_len = int(samples[0]["input_ids"].shape[-1])
-    batch_size = 4 if seq_len <= 3500 else (2 if seq_len <= 5200 else 1)
+    # `pred_slot_isolated_eager` materialises a full B×32×L×L attention matrix,
+    # so the batch size is memory- rather than compute-bound. Start optimistic
+    # and halve on OOM (see the loop below).
+    batch_size = 8 if seq_len <= 3600 else (4 if seq_len <= 5400 else 2)
 
     def run_batch(batch):
         kwargs = dict(
@@ -465,12 +480,23 @@ def analyze(
 
     t2 = time.perf_counter()
     values = []
-    for start in range(0, len(samples), batch_size):
-        chunk = samples[start : start + batch_size]
-        values.extend(run_batch(chunk))
+    while len(values) < len(samples):
+        chunk = samples[len(values) : len(values) + batch_size]
+        try:
+            values.extend(run_batch(chunk))
+        except torch.cuda.OutOfMemoryError:
+            torch.cuda.empty_cache()
+            if batch_size == 1:
+                raise gr.Error(
+                    "Ran out of GPU memory. Try a smaller 'Max image side' or fewer "
+                    "'Frames per prefix' in Advanced settings."
+                )
+            batch_size = max(1, batch_size // 2)
+            print(f"[oom] falling back to batch_size={batch_size}", flush=True)
+            continue
         progress(
-            0.2 + 0.55 * (start + len(chunk)) / len(samples),
-            desc=f"Scoring prefix {start + len(chunk)}/{len(samples)}…",
+            0.2 + 0.55 * len(values) / len(samples),
+            desc=f"Scoring prefix {len(values)}/{len(samples)}…",
         )
     t_value = time.perf_counter() - t2