File size: 718 Bytes
e623ad6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | """Load ForgeryNet weights from this Hub repo (requires redrob-verify code)."""
from pathlib import Path
import torch
from safetensors.torch import load_file
# pip/uv install from https://github.com/savagemanage/redrob-verify
from services.forgery.model import ForgeryNet
def load(repo_dir: str | Path = ".") -> ForgeryNet:
cfg = __import__("json").loads(Path(repo_dir, "config.json").read_text())
model = ForgeryNet(imagenet=False, loc_head=True)
state = load_file(Path(repo_dir) / "model.safetensors")
model.load_state_dict(state)
model.eval()
print("image_size", cfg.get("image_size"), "threshold", cfg.get("recommended_threshold"))
return model
if __name__ == "__main__":
load()
|