"""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()