pisa-sr-diffusers

PiSA-SR — diffusers-native repack

A format conversion of the official PiSA-SR (CVPR 2025) weights into standard 🤗 diffusers format, so inference needs no vendored model code and no PEFT/LoRA plumbing. No new training — all credit belongs to the original authors.

Upstream ships a raw PEFT state dict with six sub-adapters whose keys embed adapter names, so load_lora_weights() doesn't work on it. Here the deltas are pre-fused into two ready-to-load UNets.

Contents

VAE, text encoder and tokenizer are unmodified SD 2.1-base — load them from the base model. The upstream checkpoint contains no VAE or text-encoder weights.

Usage

Not a standard diffusion pipeline: no scheduler, no noise, no sampling loop. The LQ image is upsampled first, encoded, and the UNet runs a single pass at t=1 on an empty prompt. Its output is a residual subtracted from the input latent.

example_inference.py in this repo is a runnable version of everything below.

import torch, PIL.Image as Image
import torchvision.transforms.functional as TF
from torchvision import transforms
from diffusers import AutoencoderKL, UNet2DConditionModel
from transformers import AutoTokenizer, CLIPTextModel

BASE, REPO = "stabilityai/stable-diffusion-2-1-base", "ndtran0101/pisa-sr-diffusers"
dev, dt = "cuda", torch.float16

tok  = AutoTokenizer.from_pretrained(BASE, subfolder="tokenizer")
te   = CLIPTextModel.from_pretrained(BASE, subfolder="text_encoder").to(dev, dt).eval()
vae  = AutoencoderKL.from_pretrained(BASE, subfolder="vae").to(dev, dt).eval()
unet = UNet2DConditionModel.from_pretrained(REPO, subfolder="unet_full").to(dev, dt).eval()

img = Image.open("lq.png").convert("RGB")
img = img.resize((img.width * 4, img.height * 4))
img = img.resize((img.width - img.width % 8, img.height - img.height % 8), Image.LANCZOS)

with torch.no_grad():
    x   = TF.to_tensor(img).unsqueeze(0).to(dev, dt) * 2 - 1
    ids = tok("", max_length=tok.model_max_length, padding="max_length",
              truncation=True, return_tensors="pt").input_ids.to(dev)
    emb = te(ids)[0].to(dt)
    t   = torch.tensor([1], device=dev).long()

    z    = vae.encode(x).latent_dist.sample() * vae.config.scaling_factor
    pred = unet(z, t, encoder_hidden_states=emb).sample
    out  = vae.decode((z - pred) / vae.config.scaling_factor).sample.clamp(-1, 1)

transforms.ToPILImage()((out * 0.5 + 0.5).clamp(0, 1)[0].float().cpu()).save("sr.png")

Adjustable mode

Load unet_pix as well and combine the two predictions. Higher lambda_pix removes noise and compression artifacts (too high → over-smoothed); higher lambda_sem adds semantic detail (too high → artifacts). Both are 1.0 in the default mode above.

unet_pix = UNet2DConditionModel.from_pretrained(REPO, subfolder="unet_pix").to(dev, dt).eval()

with torch.no_grad():
    pred_sem = unet(z, t, encoder_hidden_states=emb).sample
    pred_pix = unet_pix(z, t, encoder_hidden_states=emb).sample
    pred = lambda_pix * pred_pix + lambda_sem * (pred_sem - pred_pix)
    out  = vae.decode((z - pred) / vae.config.scaling_factor).sample.clamp(-1, 1)

Colour fix

The upstream pipeline applies an AdaIN colour transfer from the upsampled input to the decoded output. It sits outside the network and skipping it shifts colour noticeably.

def adain(target, source):
    t = TF.to_tensor(target).unsqueeze(0)
    s = TF.to_tensor(source).unsqueeze(0)
    t_mean, t_std = t.mean([2, 3], keepdim=True), t.std([2, 3], keepdim=True)
    s_mean, s_std = s.mean([2, 3], keepdim=True), s.std([2, 3], keepdim=True)
    return transforms.ToPILImage()(
        (((t - t_mean) / (t_std + 1e-5)) * s_std + s_mean).clamp(0, 1)[0])

Also easy to get wrong

• Pre-upsampling — ×4 happens before the UNet, so compute scales with output pixels, not input. A 128² input at ×4 costs the same as a 512² input at ×1.
• Large outputs — no tiling is shipped here. Use vae.enable_tiling() and tile the image yourself above ~768².
• Seeds — inference is deterministic apart from latent_dist.sample(); use .mode() for bit-reproducible output.

Verification

Checked against the official implementation (RealSR crops, RTX 4090, fp16):

Visually and metrically indistinguishable. Upstream metrics reproduce to within 0.39% relative across 27 values (StableSR protocol, ×4, 1 step):

PSNR/SSIM are Y-channel (YCbCr); MANIQA (not shown) needs the PIPAL weights — RGB PSNR or KonIQ MANIQA will not reproduce the paper.

Speed, single RTX 4090 fp16, 1 step: 512² 0.07 s / 5.0 GB · 1024² 0.47 s / 7.4 GB · 2048² 17.9 s / 15.4 GB. Larger outputs work with a reduced VAE-decoder tile (8192² in 8.3 GB) — the limit is time, not VRAM.

License

Two licenses apply:

• unet_pix/, unet_full/ are derivatives of Stable Diffusion 2.1-base → CreativeML Open RAIL++-M, including its use-based restrictions, which you must pass on downstream.
• lora/*.safetensors contain only PiSA-SR-trained parameters → Apache 2.0.

Citation

@inproceedings{sun2025pisasr,
  title     = {Pixel-level and Semantic-level Adjustable Super-resolution: A Dual-LoRA Approach},
  author    = {Sun, Lingchen and Wu, Rongyuan and Ma, Zhiyuan and Liu, Shuaizheng and Yi, Qiaosi and Zhang, Lei},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year      = {2025},
  eprint    = {2412.03017},
  archivePrefix = {arXiv},
  url       = {https://arxiv.org/abs/2412.03017}
}