Merge origin/develop, resolving upscale.py profiling conflict

Both branches touched upscale_frames(): develop added opt-in profiling
(UpscaleProfiler) for #48, this branch added the mypy --strict type
annotation on out_frames. Keep both.

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

postprocess/upscale/profiling.py

@@ -0,0 +1,81 @@
+"""Optional profiling for `upscale_frames()`, gated by `UPSCALE_PROFILE=1` (see `upscale.PROFILE`).
+
+Buckets each tiled `model()` call into "first call at this (tile-shape, batch-size)" vs. "steady
+call" (a repeat of a shape/batch-size combo already seen this run): `torch.compile(dynamic=True)`
+only marks a dim dynamic after seeing more than one value for it, so the very first call at a
+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.
+
+Off by default — the `synchronize()` calls this needs for accurate per-call timing would
+otherwise skew real request latency by killing async kernel overlap.
+"""
+
+from __future__ import annotations
+
+import time as _time
+from contextlib import contextmanager
+from typing import Iterator
+
+import torch
+
+
+class UpscaleProfiler:
+    def __init__(self, enabled: bool, device: torch.device) -> None:
+        self.enabled = enabled
+        self._device = device
+        self._seen_keys: set[tuple[tuple[int, int], int]] = set()
+        self._first_call_count = 0
+        self._first_call_time = 0.0
+        self._steady_call_count = 0
+        self._steady_call_time = 0.0
+        self._wall_t0 = _time.perf_counter()
+        if self.enabled and device.type == "cuda":
+            torch.cuda.reset_peak_memory_stats(device)
+
+    @contextmanager
+    def timed(self, shape: tuple[int, int], batch_size: int) -> Iterator[None]:
+        """Wrap a single tiled model() call, bucketing its elapsed time by (shape, batch_size)."""
+        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
+
+        key = (shape, batch_size)
+        if key in self._seen_keys:
+            self._steady_call_count += 1
+            self._steady_call_time += elapsed
+        else:
+            self._seen_keys.add(key)
+            self._first_call_count += 1
+            self._first_call_time += elapsed
+
+    def report(self, frame_count: int, tile_size: int, frame_batch_size: int, max_tile_batch: int) -> None:
+        if not self.enabled:
+            return
+        wall_elapsed = _time.perf_counter() - self._wall_t0
+        model_time = self._first_call_time + self._steady_call_time
+        steady_avg = (
+            self._steady_call_time / self._steady_call_count if self._steady_call_count else float("nan")
+        )
+        print(
+            f"[upscale] profile: frames={frame_count} wall={wall_elapsed:.2f}s model_time={model_time:.2f}s "
+            f"(non_model={wall_elapsed - model_time:.2f}s) | first-at-shape calls={self._first_call_count} "
+            f"total={self._first_call_time:.2f}s avg={self._first_call_time / max(self._first_call_count, 1):.3f}s/call "
+            f"| steady calls={self._steady_call_count} total={self._steady_call_time:.2f}s avg={steady_avg:.3f}s/call",
+            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
+            print(
+                f"[upscale] profile: peak CUDA memory allocated={peak_allocated:.2f} GiB "
+                f"reserved={peak_reserved:.2f} GiB (TILE_SIZE={tile_size}, "
+                f"FRAME_BATCH_SIZE={frame_batch_size}, MAX_TILE_BATCH={max_tile_batch})",
+                flush=True,
+            )

postprocess/upscale/upscale.py

@@ -26,6 +26,7 @@ fetch out of the metered GPU allocation on its first upscale request.
 from __future__ import annotations
 
 import math
+import os
 from functools import lru_cache
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, cast
@@ -37,6 +38,7 @@ from PIL import Image
 from safetensors.torch import load_file
 from torch.hub import download_url_to_file, get_dir
 
+from .profiling import UpscaleProfiler
 from .srvgg_arch import SRVGGNetCompact
 
 if TYPE_CHECKING:
@@ -67,6 +69,10 @@ MAX_TILE_BATCH = 16
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dtype = torch.float16 if device.type == "cuda" else torch.float32
 
+# Set UPSCALE_PROFILE=1 on the dev Space to print per-model()-call timing and peak CUDA memory
+# for each upscale_frames() invocation — see profiling.py. Off by default.
+PROFILE = os.environ.get("UPSCALE_PROFILE", "") not in ("", "0")
+
 ProgressCallback = Callable[[int, int], None]
 
 
@@ -112,7 +118,7 @@ def _pre_pad(tensor: torch.Tensor, pad: int) -> torch.Tensor:
 
 
 @torch.no_grad()
-def _tile_process(model: torch.nn.Module, img: torch.Tensor) -> torch.Tensor:
+def _tile_process(model: torch.nn.Module, img: torch.Tensor, profiler: UpscaleProfiler) -> torch.Tensor:
     batch, channel, height, width = img.shape
     output = img.new_zeros((batch, channel, height * SCALE, width * SCALE))
     tiles_x = math.ceil(width / TILE_SIZE)
@@ -136,14 +142,15 @@ def _tile_process(model: torch.nn.Module, img: torch.Tensor) -> torch.Tensor:
                     (b, in_x0, in_y0, in_x1, in_y1, pad_x0, pad_y0, pad_x1, pad_y1)
                 )
 
-    for jobs in jobs_by_shape.values():
+    for shape, jobs in jobs_by_shape.items():
         for chunk_start in range(0, len(jobs), MAX_TILE_BATCH):
             chunk = jobs[chunk_start:chunk_start + MAX_TILE_BATCH]
             patch = torch.cat(
                 [img[b:b + 1, :, pad_y0:pad_y1, pad_x0:pad_x1] for b, _, _, _, _, pad_x0, pad_y0, pad_x1, pad_y1 in chunk],
                 dim=0,
             )
-            tile_out = model(patch)
+            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
@@ -189,6 +196,7 @@ def upscale_frames(
         return []
     model = _load_model()
     out_frames: list[Image.Image] = []
+    profiler = UpscaleProfiler(PROFILE, device)
 
     i = 0
     while i < len(frames):
@@ -201,7 +209,7 @@ def upscale_frames(
 
         tensor = _frames_to_tensor(batch)
         padded = _pre_pad(tensor, TILE_PAD)
-        upscaled = _tile_process(model, padded)
+        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()
@@ -209,4 +217,6 @@ def upscale_frames(
             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))
+
+    profiler.report(len(frames), TILE_SIZE, FRAME_BATCH_SIZE, MAX_TILE_BATCH)
     return out_frames