runtime.py

"""Eager ImprovedSRV5 registration for local CPU and ZeroGPU runtimes."""

from __future__ import annotations

import spaces

import os
from pathlib import Path

import torch

from inference.checkpoint import (
    CHECKPOINT_FILENAME,
    load_verified_model,
    resolve_checkpoint,
)
from inference.sr import InferenceSettings


MODEL_REPO_ID = "lckly/improved-srv5-x4"
MODEL_REVISION = "39e7df52b2a9716ea7a2434ca0281e2348698b3f"
GPU_DURATION = 60
SETTINGS = InferenceSettings()
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

_local_override = os.environ.get("SR_V5_CHECKPOINT_PATH")
_checkpoint_path = resolve_checkpoint(
    Path(_local_override) if _local_override else None,
    MODEL_REPO_ID,
    MODEL_REVISION,
    CHECKPOINT_FILENAME,
)
MODEL = load_verified_model(_checkpoint_path, torch.device("cpu"))
MODEL.to(DEVICE)
MODEL.eval()
MODEL.requires_grad_(False)