from huggingface_hub import hf_hub_download
import torch, importlib.util, json
from PIL import Image
import torchvision.transforms.functional as TF
model_file = hf_hub_download(repo_id="olaverse/prism-upscaler-4x", filename="model.py")
ckpt_file = hf_hub_download(repo_id="olaverse/prism-upscaler-4x", filename="pytorch_model.pt")
config_file = hf_hub_download(repo_id="olaverse/prism-upscaler-4x", filename="config.json")
spec = importlib.util.spec_from_file_location("model", model_file)
model_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(model_module)
config = json.load(open(config_file))
model = model_module.FSRCNN(**config)
model.load_state_dict(torch.load(ckpt_file, map_location="cpu"))
model.eval()
img = Image.open("input.jpg").convert("RGB")
with torch.no_grad():
output = model(TF.to_tensor(img).unsqueeze(0)).clamp(0, 1)
TF.to_pil_image(output[0]).save("output.jpg")