Image-Upscaler-150M

Single-model image upscaler: one network covers every standard mode (720p, 1080p, 2K, 4K) plus arbitrary in-between targets by composing 2x/3x/4x hops. Architecture is a bicubic base with a window-attention transformer residual (hop-conditioned via sincos embedding), trained on DIV2K and fine-tuned on mixed degradations (clean, Gaussian blur, JPEG) so it restores blur and compression damage instead of reproducing it. This is the v2 checkpoint (degradation fine-tune); the v1 lanczos-only numbers are included for comparison.

Quick start

pip install -r requirements.txt

# Inspect model specifications
python infer.py --info

# Upscale to any target (long side, pixels)
python infer.py -i photo.jpg -t 2160

# Explicit output path
python infer.py -i photo.jpg -t 1080 -o out.png

Everything required to run the model is contained inside this directory. infer.py imports only PyTorch, Pillow and NumPy, contains the complete model definition, tiled inference with feathered blending, and the hop-chain planner. Standalone inference: python infer.py -i photo.jpg -t 2160 writes photo_2160p.png. python infer.py --info prints the model card. Python API: UpscaleRunner().upscale_to(PIL.Image, target_px).

from infer import UpscaleRunner
from PIL import Image

runner = UpscaleRunner(device="cuda")  # or "cpu"
image = Image.open("photo.jpg")
result = runner.upscale_to(image, 2160)
result.save("photo_2160p.png")

Quality

Held-out DIV2K-valid (PSNR dB / SSIM; higher is better). v2 = this checkpoint:

Robustness challenges (user photo, PSNR dB)

Blur challenge: the source is blurred (Gaussian sigma 2) before upscaling; the score is measured against the sharp ground truth. v1 reproduced the blur (trained on lanczos only); the v2 fine-tune removes it.

Sharp-source challenge: the reference is the original full-resolution photo, so no method can score high — what counts is closing the gap to plain lanczos. v2 recovers most of it:

How it works

1. Hop chain planning. The requested target is decomposed into 2x/3x/4x hops chosen greedily in log space (e.g. 144p to 4K walks 4x then 2x), with a bounded overshoot followed by an area downsample when that lands closer.

2. Bicubic base + learned residual. Each hop upscales the input bicubically, then a window-attention transformer predicts the missing detail as a residual. Hop identity is injected as a sincos embedding with a zero-init per-block bias, and the output head is zero-init, so the network starts as pure bicubic and learns corrections.

3. Tiled feathered inference. Large outputs are produced in overlapping tiles with cosine-feathered alpha blending, so memory stays flat regardless of output size.

4. Degradation-robust training. v2 was fine-tuned from v1 for 3,804 steps with per-sample mixed degradations (blur sigma 0.8-2.2, JPEG quality 35-85, and both combined) while the target stayed sharp — teaching restoration, not just resampling.

Sample Sheets

Comparison sheets generated during evaluation (lanczos / nearest / model columns, labeled per sheet):

• samples/user_cat_blur_720.png: blur source -> 720p
• samples/user_cat_blur_1080.png: blur source -> 1080p
• samples/user_cat_blur_2160.png: blur source -> 4K
• samples/user_cat_blur_face.png: blurred face crop 4x
• samples/user_cat_720.png: sharp source -> 720p
• samples/user_cat_2160.png: sharp source -> 4K
• samples/user_cat_face.png: sharp face crop 4x
• samples/samples_flickr2k_720p.png: Flickr2K 144p -> 720p
• samples/samples_flickr2k_2160p.png: Flickr2K 540p -> 4K
• samples/samples_144to4k_0.png: 144p -> 4K multi-hop demo

Limitations

• Extreme blur at large factors. On the sigma-2 blur challenge at 4.5x total scale (144p-class source to 4K), information loss dominates; this release degrades gracefully to baseline quality rather than hallucinating detail. The 1M release shows the strongest blur restoration because that regime is nearest its training distribution.
• Clean-source PSNR. The degradation fine-tune trades a small amount of clean-input PSNR (at most 0.05 dB on held-out data) for large robustness and sharp-source-detail gains.
• Not a hallucinator. The Charbonnier-regression objective does not invent textures (no GAN/adversarial loss); outputs stay faithful to the input.

Training

DIV2K train set (800 images, box-masked canvas cache), random crops (crop 144 px), mixed hop sampling 0.45/0.27/0.28 for 2x/3x/4x, Charbonnier loss, AdamW lr 3e-3 cosine over a 15-minute wall-clock budget with 200-step warmup, gradient clip 1.0, ModelEMA served at save. v2 initialized from the v1 EMA weights. Steps this round: 3,804.

Files in this Release

• infer.py: Standalone inference CLI and Python interface (model, tiling, hop chains).
• config.json: Model hyperparameters and release metadata.
• requirements.txt: Python package dependencies (torch, pillow, numpy).
• the 150M checkpoint: Trained checkpoint weights (EMA) for the 142,196,659-parameter model.
• charts/: Vector graphics for the benchmark results above.
• samples/: Comparison sheets used during evaluation.

Provenance

Trained with the local training pipeline train.py --upscale on a single RTX 5090; training data DIV2K (HuggingFace bdanko/DIV2K_train_HR); eval sets DIV2K-valid, Flickr2K, Set14, Set5. The v1-to-v2 fine-tune and all measurements are reproducible from the repository.