Break down upscale's non-model overhead into named profiling spans (#61)

UpscaleProfiler previously only measured model() calls, leaving the ~74%
of upscale wall time spent outside them as one opaque "non_model" number.
Adds a span() context manager and wraps job_grouping, tile_merge,
frames_to_tensor, pre_pad, host_transfer, and pil_convert so the next
UPSCALE_PROFILE=1 run on the dev Space localizes where that time actually
goes.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

postprocess/upscale/profiling.py

@@ -6,6 +6,11 @@ only marks a dim dynamic after seeing more than one value for it, so the very fi
 given (shape, batch_size) pair is the one that risks eating a recompile, not just the first call
 overall. Also tracks peak CUDA memory for the run. See issue #48.
 
+Also buckets the non-model spans identified in issue #61 (frame/tensor prep, tile-shape grouping,
+tile merge, the post-tile host transfer, and per-frame PIL conversion) via `span()`, so a dev-Space
+run localizes exactly where the ~74%-of-wall-time "non_model" cost goes instead of leaving it as
+one opaque number.
+
 Off by default — the `synchronize()` calls this needs for accurate per-call timing would
 otherwise skew real request latency by killing async kernel overlap.
 """
@@ -28,6 +33,8 @@ class UpscaleProfiler:
         self._first_call_time = 0.0
         self._steady_call_count = 0
         self._steady_call_time = 0.0
+        self._span_times: dict[str, float] = {}
+        self._span_counts: dict[str, int] = {}
         self._wall_t0 = _time.perf_counter()
         if self.enabled and device.type == "cuda":
             torch.cuda.reset_peak_memory_stats(device)
@@ -55,6 +62,23 @@ class UpscaleProfiler:
             self._first_call_count += 1
             self._first_call_time += elapsed
 
+    @contextmanager
+    def span(self, name: str) -> Iterator[None]:
+        """Wrap an arbitrary non-model span (e.g. tile merge, host transfer, PIL conversion),
+        accumulating elapsed time under `name` for the non_model breakdown in `report()`."""
+        if not self.enabled:
+            yield
+            return
+        if self._device.type == "cuda":
+            torch.cuda.synchronize()
+        t0 = _time.perf_counter()
+        yield
+        if self._device.type == "cuda":
+            torch.cuda.synchronize()
+        elapsed = _time.perf_counter() - t0
+        self._span_times[name] = self._span_times.get(name, 0.0) + elapsed
+        self._span_counts[name] = self._span_counts.get(name, 0) + 1
+
     def report(self, frame_count: int, tile_size: int, frame_batch_size: int, max_tile_batch: int) -> None:
         if not self.enabled:
             return
@@ -70,6 +94,18 @@ class UpscaleProfiler:
             f"| steady calls={self._steady_call_count} total={self._steady_call_time:.2f}s avg={steady_avg:.3f}s/call",
             flush=True,
         )
+        if self._span_times:
+            non_model = wall_elapsed - model_time
+            span_total = sum(self._span_times.values())
+            breakdown = " ".join(
+                f"{name}={t:.2f}s(n={self._span_counts[name]})"
+                for name, t in sorted(self._span_times.items(), key=lambda kv: -kv[1])
+            )
+            print(
+                f"[upscale] profile: non_model breakdown: {breakdown} | "
+                f"tracked_total={span_total:.2f}s unaccounted={non_model - span_total:.2f}s",
+                flush=True,
+            )
         if self._device.type == "cuda":
             peak_allocated = torch.cuda.max_memory_allocated(self._device) / 1024**3
             peak_reserved = torch.cuda.max_memory_reserved(self._device) / 1024**3

postprocess/upscale/upscale.py

@@ -130,17 +130,18 @@ def _tile_process(model: torch.nn.Module, img: torch.Tensor, profiler: UpscalePr
     # patch shape so each model() call is a single batched forward pass over same-shaped patches
     # (chunked to MAX_TILE_BATCH to bound memory) instead of one tile at a time.
     jobs_by_shape: dict[tuple[int, int], list[tuple[int, int, int, int, int, int, int, int, int]]] = {}
-    for b in range(batch):
-        for y in range(tiles_y):
-            for x in range(tiles_x):
-                in_x0, in_x1 = x * TILE_SIZE, min((x + 1) * TILE_SIZE, width)
-                in_y0, in_y1 = y * TILE_SIZE, min((y + 1) * TILE_SIZE, height)
-                pad_x0, pad_x1 = max(in_x0 - TILE_PAD, 0), min(in_x1 + TILE_PAD, width)
-                pad_y0, pad_y1 = max(in_y0 - TILE_PAD, 0), min(in_y1 + TILE_PAD, height)
-                shape = (pad_y1 - pad_y0, pad_x1 - pad_x0)
-                jobs_by_shape.setdefault(shape, []).append(
-                    (b, in_x0, in_y0, in_x1, in_y1, pad_x0, pad_y0, pad_x1, pad_y1)
-                )
+    with profiler.span("job_grouping"):
+        for b in range(batch):
+            for y in range(tiles_y):
+                for x in range(tiles_x):
+                    in_x0, in_x1 = x * TILE_SIZE, min((x + 1) * TILE_SIZE, width)
+                    in_y0, in_y1 = y * TILE_SIZE, min((y + 1) * TILE_SIZE, height)
+                    pad_x0, pad_x1 = max(in_x0 - TILE_PAD, 0), min(in_x1 + TILE_PAD, width)
+                    pad_y0, pad_y1 = max(in_y0 - TILE_PAD, 0), min(in_y1 + TILE_PAD, height)
+                    shape = (pad_y1 - pad_y0, pad_x1 - pad_x0)
+                    jobs_by_shape.setdefault(shape, []).append(
+                        (b, in_x0, in_y0, in_x1, in_y1, pad_x0, pad_y0, pad_x1, pad_y1)
+                    )
 
     for shape, jobs in jobs_by_shape.items():
         for chunk_start in range(0, len(jobs), MAX_TILE_BATCH):
@@ -152,13 +153,14 @@ def _tile_process(model: torch.nn.Module, img: torch.Tensor, profiler: UpscalePr
             with profiler.timed(shape, patch.shape[0]):
                 tile_out = model(patch)
 
-            for i, (b, in_x0, in_y0, in_x1, in_y1, pad_x0, pad_y0, pad_x1, pad_y1) in enumerate(chunk):
-                trim_x0, trim_y0 = (in_x0 - pad_x0) * SCALE, (in_y0 - pad_y0) * SCALE
-                trim_x1 = trim_x0 + (in_x1 - in_x0) * SCALE
-                trim_y1 = trim_y0 + (in_y1 - in_y0) * SCALE
-                output[b:b + 1, :, in_y0 * SCALE:in_y1 * SCALE, in_x0 * SCALE:in_x1 * SCALE] = (
-                    tile_out[i:i + 1, :, trim_y0:trim_y1, trim_x0:trim_x1]
-                )
+            with profiler.span("tile_merge"):
+                for i, (b, in_x0, in_y0, in_x1, in_y1, pad_x0, pad_y0, pad_x1, pad_y1) in enumerate(chunk):
+                    trim_x0, trim_y0 = (in_x0 - pad_x0) * SCALE, (in_y0 - pad_y0) * SCALE
+                    trim_x1 = trim_x0 + (in_x1 - in_x0) * SCALE
+                    trim_y1 = trim_y0 + (in_y1 - in_y0) * SCALE
+                    output[b:b + 1, :, in_y0 * SCALE:in_y1 * SCALE, in_x0 * SCALE:in_x1 * SCALE] = (
+                        tile_out[i:i + 1, :, trim_y0:trim_y1, trim_x0:trim_x1]
+                    )
     return output
 
 
@@ -207,14 +209,18 @@ def upscale_frames(
             batch.append(frames[i])
             i += 1
 
-        tensor = _frames_to_tensor(batch)
-        padded = _pre_pad(tensor, TILE_PAD)
+        with profiler.span("frames_to_tensor"):
+            tensor = _frames_to_tensor(batch)
+        with profiler.span("pre_pad"):
+            padded = _pre_pad(tensor, TILE_PAD)
         upscaled = _tile_process(model, padded, profiler)
         # crop the pre-pad border (scaled) back off
         upscaled = upscaled[:, :, : h * SCALE, : w * SCALE]
-        results = upscaled.clamp(0, 1).permute(0, 2, 3, 1).float().cpu().numpy()
+        with profiler.span("host_transfer"):
+            results = upscaled.clamp(0, 1).permute(0, 2, 3, 1).float().cpu().numpy()
         for result in results:
-            out_frames.append(Image.fromarray((result * 255.0).round().astype(np.uint8)))
+            with profiler.span("pil_convert"):
+                out_frames.append(Image.fromarray((result * 255.0).round().astype(np.uint8)))
             if progress_callback is not None:
                 progress_callback(len(out_frames), len(frames))