tailot/universal-upscale-distilled

Universal Upscale (Distilled) - native 9x super-resolution in ~5 MB

A compact knowledge-distilled student that upscales any RGB image by 9x in a single forward pass: 1.30 M parameters, 5.2 MB ONNX, no PyTorch required.

This is the distilled variant. The non-distilled model — one arbitrary-scale checkpoint that handles any factor from 2x to 9.2x — lives at tailot/universal-upscale. Use that one when you need a free scale factor; use this one when you want a tiny, fast, fixed 9x model that runs anywhere ONNX Runtime runs.

What the distilled model does

• One forward pass at a fixed 9x factor: H x W in, 9*H x 9*W out. Batch, height and width are dynamic axes, so any input size works.
• Compact trunk: 8 residual blocks x 64 features, PReLU, and a staged 3 x 3 upsample built with DepthToSpace - no transposed convolutions, no custom operators.
• Trained by knowledge distillation from a realesrgan_x4plus teacher plus ground-truth targets, so a 1.30 M-parameter network learns behaviour that would otherwise need a much larger generator.
• Runs on ONNX Runtime: CPU, CUDA, DirectML, or the browser (onnxruntime-web). Graph ops: Add, Conv, DepthToSpace, PRelu.

Files in this repository

File	What it is
model.onnx	Main artifact. Single-file ONNX (opset ai.onnx 17) with the weights embedded: no .onnx.data sidecar, no external data, 5.2 MB.
model_scripted.pt	The same weights as a TorchScript module (forward(x) returns the 9x output) for PyTorch users.
config.json	Architecture, training recipe, input/output contract and verification summary.
upscale.py	Standalone tiled-inference CLI (numpy + OpenCV + ONNX Runtime) that also emits the bicubic baseline.
requirements.txt	Dependencies of upscale.py.
examples/	verification.json (full report), verify_panel.png (DIV2K crops vs bicubic), compare_6x.png (real photo at 6x).

Quick start (ONNX Runtime)

import numpy as np
import onnxruntime as ort
from PIL import Image

session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
input_name = session.get_inputs()[0].name

img = Image.open("photo.jpg").convert("RGB")
lr = np.asarray(img, dtype=np.float32).transpose(2, 0, 1)[None, ...] / 255.0
sr = session.run(None, {input_name: lr})[0][0]          # -> [3, 9*H, 9*W]

arr = (np.clip(sr, 0.0, 1.0).transpose(1, 2, 0) * 255.0).round().astype("uint8")
Image.fromarray(arr).save("photo_9x.png")

Fetch the weights straight from the Hub instead of keeping a local copy:

from huggingface_hub import hf_hub_download

path = hf_hub_download(repo_id="tailot/universal-upscale-distilled", filename="model.onnx")

TorchScript flavour (same weights, forward(x), no ONNX Runtime needed):

import torch

model = torch.jit.load("model_scripted.pt")
model.eval()

CLI demo

pip install -r requirements.txt

python upscale.py photo.jpg --scale 9          # native 9x
python upscale.py photo.jpg --scale 6          # 9x inference, then area downsample
python upscale.py photo.jpg --scale 4 --tile 96 --pad 8

upscale.py runs the network tile by tile in low-resolution space, so a 4K output does not need 4K-sized intermediate buffers, and it writes <output>_bicubic.png next to the result so the comparison is one folder away.

API contract

input : float32 [N, 3, H, W]  RGB, in [0, 1]   (batch, height and width are dynamic axes)
output: float32 [N, 3, 9*H, 9*W]  RGB, in [0, 1]

Native factor is 9x. For a smaller factor, run 9x and area-downsample to the target (sharper than interpolating the input directly); for a larger factor, run 9x and bicubic-upsample the result. upscale.py does both for you, so --scale 6 or --scale 12 work even though the network is 9x.

Quality vs a bicubic baseline

Model	PSNR (dB)	SSIM	MS-SSIM	Detail / ground truth
Distilled student	23.4219	0.5414	0.7677	0.1288
Bicubic baseline	23.2327	0.5568	0.7628	0.0820

At 9x the student reproduces 1.57x the high-frequency energy of the bicubic baseline (+0.189 PSNR, -0.015 SSIM, +0.005 MS-SSIM): it reconstructs detail instead of only interpolating it, and the SSIM gap is the price of that extra sharpness.

Measured with scripts/verify_onnx.py on 6 DIV2K validation crops (384 px HR, bicubic-downscaled 9x, ESRGAN border 4 px, tiled inference). Full numbers, including per-shape sanity checks, are in examples/verification.json.

The notebook reported 28.70 dB validation PSNR for these weights after 4,000 distillation iterations (its own validation crops, not the ones in the table above).

Limitations

• The factor is fixed at 9x. Other factors are obtained by resampling the 9x output, which is what upscale.py does; for a model that renders any factor natively, use tailot/universal-upscale.
• Alpha channels are not upscaled by the network (it is 3-channel in / 3-channel out); upscale.py resamples the alpha separately and re-encodes it.
• Outputs can leave the [0, 1] range by a few thousandths; clamp before quantising.
• These are research weights trained on DIV2K + Flickr2K and distilled from Real-ESRGAN: treat them as research-only, not as a drop-in replacement for a commercial upscaler.

Training recipe

Item	Value
Method	knowledge distillation (student trained on teacher + HR targets)
Teacher	realesrgan_x4plus
Distillation weight (alpha)	0.6
Target source	both
Iterations	4,000
Training time	48.9 min on one GPU
Loss weights	L1 1.0 + perceptual 0.1 + edge 0.0
LR patch / batch	32 px / 8
Optimizer	AdamW lr 0.0002, cosine schedule, EMA 0.999
Notebook	colab/distill_sr9x.py in the source repository

License

• Weights: cc-by-nc-4.0. The training data (DIV2K + Flickr2K) is licensed for non-commercial research use, so these weights are research-only as well.
• Code in this repo: MIT.

Citation

@misc{universal_upscale_distilled,
  title = {Universal Upscale (Distilled): native 9x super-resolution knowledge-distilled student},
  author = {Vincenzo Tilotta},
  year  = {2026},
  publisher = {Hugging Face},
  url = {https://huggingface.co/tailot/universal-upscale-distilled},
  note  = {Distilled variant of universal-upscale},
}