Spaces:
Running on Zero
Running on Zero
File size: 3,292 Bytes
333ff0e 75c39e8 333ff0e 75c39e8 333ff0e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """The stateful engine: resolve the device once, load each model once, hold the optional recogniser.
The UI and the API share one :class:`StudioEngine` for the process lifetime. Models are cached by weights
path so clicking a checkpoint twice doesn't reload it; the cache is capped (default 2, enough for the
single-line A/B compare) and evicts oldest-first with a CUDA cache flush, so switching through many
checkpoints never grows VRAM without bound and never OOMs the co-tenant training run.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from diffu_studio.checkpoints import ArchConfig, LoadedModel, load_model, resolve_weights
from diffu_studio.device import Device, resolve_device
from diffu_studio.line import Recognizer, load_recognizer
if TYPE_CHECKING:
from diffu.config import Config
class StudioEngine:
"""Owns the device + a small model cache + the read-back recogniser for the app's lifetime."""
def __init__(
self,
*,
device: Device | None = None,
ckpt_root: str = "checkpoints",
arch: ArchConfig | None = None,
recognizer_name: str | None = None,
max_cached_models: int = 2,
) -> None:
self.device = device or resolve_device()
self.ckpt_root = ckpt_root
self.arch = arch or ArchConfig()
self.max_cached_models = max_cached_models
self._models: dict[str, LoadedModel] = {}
self._recognizer: Recognizer | None = (
load_recognizer(recognizer_name, self.device.torch_device) if recognizer_name else None
)
@property
def recognizer(self) -> Recognizer | None:
return self._recognizer
def model(self, path: str) -> LoadedModel:
"""Return the model at ``path`` (weights file or ``step_*`` dir), loading + caching it on first use.
Raises :class:`~diffu_studio.checkpoints.CheckpointMismatch` if the checkpoint doesn't fit the code.
"""
key = resolve_weights(path)
cached = self._models.get(key)
if cached is not None:
return cached
loaded = load_model(path, self.arch, self.device)
# ZeroGPU: opt-in AoT compile of the backbone (aokit). OFF by default (set DIFFU_STUDIO_AOT=1) until
# the compiled graph's persistence across ZeroGPU's per-request forks is validated — the compile runs
# in a @spaces.GPU fork whose install would not survive to the parent. Fallback-safe regardless.
import os
if os.environ.get("DIFFU_STUDIO_AOT") == "1":
from diffu_studio.optimization import aot_compile_backbone
aot_compile_backbone(loaded)
self._models[key] = loaded
self._evict_overflow()
return loaded
def config_for(self, path: str) -> Config:
return self.model(path).cfg
def _evict_overflow(self) -> None:
while len(self._models) > self.max_cached_models:
oldest = next(iter(self._models))
del self._models[oldest]
_free_cuda()
def _free_cuda() -> None:
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception: # noqa: BLE001 — freeing cache is best-effort; never fail an eviction on it
pass
|