README.md

---
license: other
license_name: creativeml-openrail-m-plus-plus
license_link: https://huggingface.co/stabilityai/stable-diffusion-2-1-base/blob/main/LICENSE
base_model: stabilityai/stable-diffusion-2-1-base
tags:
  - super-resolution
  - image-to-image
  - diffusers
  - one-step-diffusion
  - real-isr
library_name: diffusers
pipeline_tag: image-to-image
---

# PiSA-SR — diffusers-native repack

A format conversion of the official [PiSA-SR](https://github.com/csslc/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

| Path | What | Size |
|---|---|---|
| `unet_full/` | pix + sem fused — default 1-step mode | 1.7 GB |
| `unet_pix/` | pix only — second model for adjustable mode | 1.7 GB |
| `lora/pisa_{pix,sem}.safetensors` | adapter-only deltas (+ config json) | 16 MB ea |
| `example_inference.py` | runnable default / adjustable inference | — |

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.

```python
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.

```python
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.

```python
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):

| | PSNR | mean abs err (0–255) |
|---|---|---|
| Reference vs itself (noise floor) | 63.1 dB | 0.03 |
| This repack vs reference | **56.7 dB** | 0.13 |

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

| Dataset | PSNR(Y) | SSIM(Y) | LPIPS | DISTS | FID | MUSIQ | CLIPIQA |
|---|---|---|---|---|---|---|---|
| RealSR (100) | 25.50 | 0.7418 | 0.2672 | 0.2044 | 124.13 | 70.15 | 0.6697 |
| DRealSR (93) | 28.32 | 0.7804 | 0.2960 | 0.2169 | 130.45 | 66.11 | 0.6971 |
| DIV2K-Val (3000) | 23.87 | 0.6058 | 0.2823 | 0.1934 | 25.09 | 69.68 | 0.6928 |

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](https://huggingface.co/stabilityai/stable-diffusion-2-1-base/blob/main/LICENSE),
  including its use-based restrictions, which you must pass on downstream.
- `lora/*.safetensors` contain only PiSA-SR-trained parameters →
  [Apache 2.0](https://github.com/csslc/PiSA-SR/blob/main/LICENSE).

## Citation

```bibtex
@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}
}
```