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-denoiser", filename="model.py")
ckpt_file = hf_hub_download(repo_id="olaverse/prism-denoiser", filename="pytorch_model.pt")
config_file = hf_hub_download(repo_id="olaverse/prism-denoiser", 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.UNet(**config)
model.load_state_dict(torch.load(ckpt_file, map_location="cpu")["net"])
model.eval()
img = Image.open("noisy.jpg").convert("RGB").resize((128, 128))
x = TF.to_tensor(img).unsqueeze(0)
with torch.no_grad():
output = model(x).clamp(0, 1)
TF.to_pil_image(output[0]).save("denoised.jpg")