"""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): the first call at a new shape is where one-off costs land (CUDA/cuDNN lazy init today; per-shape recompiles back when the model was torch.compile'd — dropped per issue #50). Also tracks peak CUDA memory for the run. See issue #48. The memory line also reports what was already allocated when the run started (e.g. the diffusion pipeline's weights, resident in the same process — `reset_peak_memory_stats` resets the peak, not the baseline), the upscale run's own delta above that, and the device's total/free VRAM, so `TILE_SIZE`/`FRAME_BATCH_SIZE`/`MAX_TILE_BATCH` can be tuned against real headroom. See issue #51. 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. """ 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._span_times: dict[str, float] = {} self._span_counts: dict[str, int] = {} self._baseline_allocated = 0 self._baseline_reserved = 0 self._wall_t0 = _time.perf_counter() if self.enabled and device.type == "cuda": torch.cuda.reset_peak_memory_stats(device) self._baseline_allocated = torch.cuda.memory_allocated(device) self._baseline_reserved = torch.cuda.memory_reserved(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 @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 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._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": gib = 1024**3 peak_allocated = torch.cuda.max_memory_allocated(self._device) peak_reserved = torch.cuda.max_memory_reserved(self._device) # mem_get_info is driver-level: free/total across every process on the device, not just # this caching allocator — "free" is what's actually left right now, after the run. free, total = torch.cuda.mem_get_info(self._device) print( f"[upscale] profile: peak CUDA memory allocated={peak_allocated / gib:.2f} GiB " f"reserved={peak_reserved / gib:.2f} GiB (TILE_SIZE={tile_size}, " f"FRAME_BATCH_SIZE={frame_batch_size}, MAX_TILE_BATCH={max_tile_batch})", flush=True, ) print( f"[upscale] profile: baseline at start allocated={self._baseline_allocated / gib:.2f} GiB " f"reserved={self._baseline_reserved / gib:.2f} GiB | upscale delta " f"allocated=+{(peak_allocated - self._baseline_allocated) / gib:.2f} GiB " f"reserved=+{(peak_reserved - self._baseline_reserved) / gib:.2f} GiB | device " f"total={total / gib:.2f} GiB free_now={free / gib:.2f} GiB " f"headroom_vs_peak_reserved={(total - peak_reserved) / gib:.2f} GiB", flush=True, )