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-max", filename="model.py")
ckpt_file = hf_hub_download(repo_id="olaverse/prism-upscaler-max", filename="pytorch_model.pt")
config_file = hf_hub_download(repo_id="olaverse/prism-upscaler-max", 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.LIIF(**config)
model.load_state_dict(torch.load(ckpt_file, map_location="cpu"))
model.eval()
img = Image.open("input.jpg").convert("RGB")
lr_tensor = TF.to_tensor(img)
out_h, out_w = 1024, 1024 # any target resolution you want
with torch.no_grad():
feat = model.gen_feat(lr_tensor.unsqueeze(0))
coord = model_module.make_coord((out_h, out_w), device="cpu").view(1, -1, 2)
cell = torch.tensor([2.0 / out_h, 2.0 / out_w]).view(1, 1, 2).repeat(1, coord.shape[1], 1)
pred = model.query_rgb(feat, coord, cell) # chunk this loop for very large outputs
output = pred.view(1, out_h, out_w, 3).permute(0, 3, 1, 2).clamp(0, 1)
TF.to_pil_image(output[0]).save("output.jpg")