Spaces:
Sleeping
Sleeping
| """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) | |