| """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) |
| |
| |
| 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, |
| ) |
|
|
