"""ModelRegistry: load every restoration model ONCE, frozen, on CPU. Built once in the FastAPI lifespan and stashed on ``app.state``. Loading mirrors the verified checkpoint formats (webapp_spec "Model load"): - agent : bare state_dict -> ``AgentNet().load_state_dict(.., weights_only=True)`` - tools : ``load_tools(name, device)`` (reads ``{"model": sd}`` per tool) - monolith: ``{"model": sd, "val_psnr":..}`` -> ``MonoRestoreCNN(14,64)``, ``weights_only=False`` ANY method whose checkpoint is missing is omitted gracefully: ``load()`` never raises for an absent optional checkpoint, and ``available_methods`` reflects only what actually loaded. The agent + the "full" toolbox are treated as required (the demo's headline path); if those are missing we surface a clear error. """ from __future__ import annotations import logging import os from dataclasses import dataclass, field from pathlib import Path import torch from torch import nn from rlrestore.agent.env import load_tools from rlrestore.agent.net import AgentNet from rlrestore.baselines.monolith import MonoRestoreCNN from rlrestore.baselines.unet_restore import UNetRestore from .realesrgan_arch import RRDBNet, load_realesrgan_x4 log = logging.getLogger("rlrestore.backend.models") # Checkpoint root resolution, in priority order: # 1. $RLR_CKPT_ROOT (the Docker/HF-Space sets this to the bundled checkpoints/). # 2. A `checkpoints/` next to hf_space/ (standalone Space layout). # 3. The repo root's checkpoints/ (local dev from inside the research repo). _SPACE_ROOT = Path(__file__).resolve().parents[1] # hf_space/ _REPO_ROOT = Path(__file__).resolve().parents[2] # research repo root (local dev) def _resolve_ckpt_root() -> Path: env = os.environ.get("RLR_CKPT_ROOT") if env: return Path(env) space_local = _SPACE_ROOT / "checkpoints" if space_local.is_dir(): return space_local return _REPO_ROOT / "checkpoints" _CKPT_ROOT = _resolve_ckpt_root() AGENT_CKPT = _CKPT_ROOT / "agent_full_v2" / "agent_net.pt" TOOLS_ROOT = _CKPT_ROOT / "tools" MONOLITH_CKPT = _CKPT_ROOT / "monolith" / "monolith" / "best.pt" UNET_HQ_CKPT = _CKPT_ROOT / "unet_hq" / "best.pt" REALESRGAN_CKPT = _CKPT_ROOT / "realesrgan" / "RealESRGAN_x4plus.pth" MONOLITH_DEPTH = 14 MONOLITH_WIDTH = 64 # Public method ids -> the tools key they roll out against. METHOD_TOOLS_KEY = {"agent": "tools_full", "agent_ft": "tools_full_ft"} @dataclass class ModelRegistry: """Holds frozen models + metadata. Construct then call ``load()``.""" device: torch.device agent_ckpt: Path = AGENT_CKPT tools_root: Path = TOOLS_ROOT monolith_ckpt: Path = MONOLITH_CKPT unet_hq_ckpt: Path = UNET_HQ_CKPT realesrgan_ckpt: Path = REALESRGAN_CKPT agent: AgentNet | None = field(default=None, init=False) tools_full: list[nn.Module] | None = field(default=None, init=False) tools_full_ft: list[nn.Module] | None = field(default=None, init=False) monolith: MonoRestoreCNN | None = field(default=None, init=False) monolith_val_psnr: float | None = field(default=None, init=False) unet_hq: UNetRestore | None = field(default=None, init=False) unet_hq_val_psnr: float | None = field(default=None, init=False) realesrgan: RRDBNet | None = field(default=None, init=False) available_methods: list[str] = field(default_factory=list, init=False) # ------------------------------------------------------------------ load def load(self) -> "ModelRegistry": """Load all present checkpoints. Idempotent enough for a single call.""" self._load_tools_full() self._load_agent() self._load_tools_full_ft() self._load_monolith() self._load_unet_hq() self._load_realesrgan() self._recompute_methods() log.info( "ModelRegistry loaded on %s; methods=%s", self.device, self.available_methods, ) return self # ------------------------------------------------------------ components def _load_tools_full(self) -> None: path = self.tools_root / "full" / "tool00" / "best.pt" if not path.exists(): log.warning("tools/full missing (%s); agent path unavailable", path) return self.tools_full = load_tools("full", self.device, ckpt_root=self.tools_root) log.info("loaded 12 'full' tools") def _load_tools_full_ft(self) -> None: path = self.tools_root / "full_ft" / "tool00" / "best.pt" if not path.exists(): log.info("tools/full_ft absent (%s); agent_ft omitted", path) return self.tools_full_ft = load_tools( "full_ft", self.device, ckpt_root=self.tools_root ) log.info("loaded 12 'full_ft' tools") def _load_agent(self) -> None: if not self.agent_ckpt.exists(): log.warning("agent checkpoint missing (%s)", self.agent_ckpt) return net = AgentNet() sd = torch.load(self.agent_ckpt, map_location=self.device, weights_only=True) net.load_state_dict(sd) net.to(self.device).eval().requires_grad_(False) self.agent = net log.info("loaded agent (%d params)", sum(p.numel() for p in net.parameters())) def _load_monolith(self) -> None: if not self.monolith_ckpt.exists(): log.info("monolith checkpoint absent (%s); monolith omitted", self.monolith_ckpt) return model = MonoRestoreCNN(depth=MONOLITH_DEPTH, width=MONOLITH_WIDTH) ck = torch.load(self.monolith_ckpt, map_location=self.device, weights_only=False) model.load_state_dict(ck["model"]) model.to(self.device).eval().requires_grad_(False) self.monolith = model val = ck.get("val_psnr") self.monolith_val_psnr = float(val) if val is not None else None log.info("loaded monolith (val_psnr=%s)", self.monolith_val_psnr) def _load_unet_hq(self) -> None: if not self.unet_hq_ckpt.exists(): log.info("unet_hq checkpoint absent (%s); HQ Restore omitted", self.unet_hq_ckpt) return ck = torch.load(self.unet_hq_ckpt, map_location=self.device, weights_only=False) model = UNetRestore(width=int(ck.get("width", 48))) model.load_state_dict(ck["model"]) model.to(self.device).eval().requires_grad_(False) self.unet_hq = model val = ck.get("val_psnr") self.unet_hq_val_psnr = float(val) if val is not None else None log.info("loaded unet_hq (val_psnr=%s)", self.unet_hq_val_psnr) def _load_realesrgan(self) -> None: if not self.realesrgan_ckpt.exists(): log.info( "realesrgan checkpoint absent (%s); Enhance finisher omitted", self.realesrgan_ckpt, ) return self.realesrgan = load_realesrgan_x4(self.realesrgan_ckpt, self.device) log.info( "loaded realesrgan x4 (%d params)", sum(p.numel() for p in self.realesrgan.parameters()), ) # --------------------------------------------------------------- methods def _recompute_methods(self) -> None: methods: list[str] = [] if self.agent is not None and self.tools_full is not None: methods.append("agent") if self.agent is not None and self.tools_full_ft is not None: methods.append("agent_ft") if self.monolith is not None: methods.append("monolith") if self.unet_hq is not None: methods.append("unet_hq") self.available_methods = methods # --------------------------------------------------------------- helpers @property def models_loaded(self) -> bool: return bool(self.available_methods) @property def enhance_available(self) -> bool: """The Real-ESRGAN 'Enhance & Upscale' finisher is loaded and usable.""" return self.realesrgan is not None def tools_for(self, tools_key: str) -> list[nn.Module]: tools = getattr(self, tools_key, None) if tools is None: raise KeyError(f"tools '{tools_key}' not loaded") return tools def has_method(self, method: str) -> bool: return method in self.available_methods