moonbox-ai/VOSR-MLX

VOSR 2.0 MLX

MLX builds of VOSR 2.0, made by Moonbox for Apple Silicon Macs. Each precision is a folder; keep a folder's files together. Only the parts listed below are converted: the rest of the pipeline comes from the original repository unless the section below says otherwise.

Files

The last column is how far each build is from the original PyTorch model on the same inputs (see the section below for what was measured).

Precision	File	Size	sha256	Checked against the original
bf16	bf16/config.json	0.8 kB	7a65073099dd	below threshold (experimental): PSNR 33.8 dB mean / 23.2 worst, LPIPS <= 0.045
bf16	bf16/mlx_model.py	23.6 kB	227a42122e08	below threshold (experimental): PSNR 33.8 dB mean / 23.2 worst, LPIPS <= 0.045
bf16	bf16/model.safetensors	3.33 GB	d1e436b94a3d	below threshold (experimental): PSNR 33.8 dB mean / 23.2 worst, LPIPS <= 0.045
fp16	fp16/config.json	0.8 kB	7831ae8dacd5	PSNR 42.3 dB mean / 31.4 worst, LPIPS <= 0.011
fp16	fp16/mlx_model.py	23.6 kB	227a42122e08	PSNR 42.3 dB mean / 31.4 worst, LPIPS <= 0.011
fp16	fp16/model.safetensors	3.33 GB	de064c413024	PSNR 42.3 dB mean / 31.4 worst, LPIPS <= 0.011
8bit	8bit/config.json	0.8 kB	21f182efebad	below threshold (experimental): PSNR 33.7 dB mean / 23.1 worst, LPIPS <= 0.047
8bit	8bit/mlx_model.py	23.6 kB	227a42122e08	below threshold (experimental): PSNR 33.7 dB mean / 23.1 worst, LPIPS <= 0.047
8bit	8bit/model.safetensors	1.82 GB	c5e2274356c5	below threshold (experimental): PSNR 33.7 dB mean / 23.1 worst, LPIPS <= 0.047
6bit	6bit/config.json	0.8 kB	42928eb73593	below threshold (experimental): PSNR 33.1 dB mean / 22.8 worst, LPIPS <= 0.048
6bit	6bit/mlx_model.py	23.6 kB	227a42122e08	below threshold (experimental): PSNR 33.1 dB mean / 22.8 worst, LPIPS <= 0.048
6bit	6bit/model.safetensors	1.41 GB	13e8fee0b43e	below threshold (experimental): PSNR 33.1 dB mean / 22.8 worst, LPIPS <= 0.048
4bit	4bit/config.json	0.8 kB	70a18af94fd5	below threshold (experimental): PSNR 29.1 dB mean / 20.5 worst, LPIPS <= 0.094
4bit	4bit/mlx_model.py	23.6 kB	227a42122e08	below threshold (experimental): PSNR 29.1 dB mean / 20.5 worst, LPIPS <= 0.094
4bit	4bit/model.safetensors	1.01 GB	1e2f7b023f9c	below threshold (experimental): PSNR 29.1 dB mean / 20.5 worst, LPIPS <= 0.094

How to run

VOSR 2.0 (the one-step 1.4B model) ported to MLX: the whole pipeline of inference_vosr_onestep.py (bicubic upscale, the 2D Qwen-Image VAE, DINOv2-L features, one LightningDiT step, the wavelet color fix). Each folder holds model.safetensors, config.json and mlx_model.py, the port (mlx, numpy and Pillow):

import numpy as np, mlx_model
from PIL import Image
model = mlx_model.load("fp16")
out = mlx_model.upscale(model, np.array(Image.open("in.png").convert("RGB")), scale=4, seed=42)
Image.fromarray(out).save("out.png")

It runs the full picture at once (the original's default, no tiling). Output sides are rounded up to a multiple of 16 internally and cropped back.

Which one: fp16/ follows the original closely (42 dB mean against PyTorch fp32, LPIPS <= 0.011). VOSR makes up detail from noise in one step, so small numeric differences show up as different fine texture: bf16/, 8bit/ and 6bit/ look the same as each other but drift from the original on some pictures (LPIPS up to 0.048, below the lab's upscaling threshold of 0.02, so marked experimental); 4bit/ drifts further.

Precision	Size	PSNR vs PyTorch fp32 (mean / worst)	LPIPS (worst)	Time, 128 -> 512 on an M5
fp16/	3.33 GB	42.3 / 31.4 dB	0.011	1.0 s
bf16/ (experimental)	3.33 GB	33.8 / 23.2 dB	0.045	1.0 s
8bit/ (experimental)	1.82 GB	33.7 / 23.1 dB	0.047	1.1 s
6bit/ (experimental)	1.41 GB	33.1 / 22.8 dB	0.048	1.1 s
4bit/ (experimental)	1.01 GB	29.1 / 20.5 dB	0.094	1.0 s

Quantized builds quantize the DiT's and DINOv2's linear layers (group 64); the VAE is all convolutions and stays 16-bit.

How these were made

From VOSR 2.0 at revision 516f292b99: the original weights loaded in MLX and stored as bf16 or fp16, or quantized with MLX's affine quantization (8bit, 6bit, 4bit, group size 64; layers MLX can't quantize, such as convolutions and norms, stay bf16). Every build was checked against the original before upload; the sha256 column above is what the files must hash to.

License

Apache-2.0. These files are a derivative of the original model and carry its license: it is copied in this repository, and it governs your use of these files. The original model card follows unchanged.

---

Original model card

VOSR: A Vision-Only Generative Model for Image Super-Resolution

A framework for native generative image restoration.

🚩 Accepted by CVPR 2026

Rongyuan Wu^1,2,* | Lingchen Sun^1,2,* | Zhengqiang Zhang^1,2 | Xiangtao Kong^1,2
Jixin Zhao^1,2 | Shihao Wang¹ | Lei Zhang^1,2,†

¹ The Hong Kong Polytechnic University
² OPPO Research Institute

^* Equal contribution. ^† Corresponding author.

Overview of the VOSR framework, including the overall pipeline and our condition / guidance design.

📌 Quick Links

• 📰 News
• 🧰 Preparation
• 🏋️ Training
• 🔍 Inference
• 📮 Contact
• 📚 Citation

📰 News

• 2026.08.22 VOSR 2.0 released: a one-step 1.4B DiT, vision-only model with improved fine-structure restoration. See visual comparison. Thanks to ylchen333 for bringing VOSR to ComfyUI with ComfyUI-VOSR2.

• 2026.04.10 Public release: training and inference code, pretrained checkpoints, bundled VAE / decoder assets, and the ScreenSR benchmark. Setup and file layout: Preparation; commands: Training and Inference.

---

Comparison with prior methods in performance, inference efficiency, and training cost.

VOSR better preserves fine structures and text readability.

🧰 Preparation

Dependencies and Installation

## clone this repository
git clone https://github.com/cswry/VOSR.git
cd VOSR

# create an environment with python >= 3.8
conda create -n vosr python=3.8
conda activate vosr
pip install -r requirements.txt

📦 Model Weights

Download all pretrained weights from ModelScope or Hugging Face, and place them under preset/ckpts/. The expected structure:

preset/ckpts/
|-- Qwen-Image-vae-2d/          # Qwen-Image VAE (2D, for 1.4B models)
|-- stable-diffusion-2-1-base/  # SD2.1 VAE (for 0.5B models)
|-- sd21_lwdecoder.pth          # Lightweight decoder for SD2.1 VAE
|-- torch_cache/                # DINOv2 pretrained weights
|-- VOSR2/                     # VOSR 2.0 one-step (1.4B) model
|-- VOSR_0.5B_ms/               # 0.5B multi-step model
|-- VOSR_0.5B_os/               # 0.5B one-step (distilled) model
|-- VOSR_1.4B_ms/               # 1.4B multi-step model
`-- VOSR_1.4B_os/               # 1.4B one-step (distilled) model

VAE and decoder

For VOSR-0.5B, we provide sd21_lwdecoder.pth, a lightweight replacement for the original SD2.1 VAE decoder. It achieves comparable overall visual quality in our evaluation, while performing slightly better on text-rich and document-like images.

When scaling to VOSR-1.4B, we adopt the 16-channel Qwen-Image VAE to better preserve input fidelity. Although Qwen-Image is designed for T2I generation, it is released in a video-VAE form, which is unnecessarily slow for image SR inference. We therefore provide Qwen-Image-vae-2d, an image-only 2D variant extracted from the original model to remove the overhead of the full 3D design.

🗂️ Training Data

We support two data loading modes, configured via dataset_type in the YAML config:

• txt - Each folder contains individual image files. A txt config lists folders with sampling weights.
• webdataset - Each folder contains .tar shards. Same txt config format, loaded via WebDataset.

Create a dataset config file (e.g., configs/train_txt/train_dataset_txt.txt):

/path/to/dataset_A, 2
/path/to/dataset_B, 1
/path/to/dataset_C, 1

Each line: <folder_path>, <sampling_weight>. Higher weight = more frequent sampling.

For txt mode, each folder should contain HQ images (.png / .jpg). For webdataset mode, each folder should contain .tar shards with images inside.

🧪 New Real-World Paired Benchmark

Download the ScreenSR benchmark from ModelScope or Hugging Face, place it wherever you like, and then point -i to that folder when running inference.

ScreenSR is a real-world paired benchmark for generative SR, built with a screen re-photography pipeline. It provides cleaner references, more diverse content, and broader variation in scenes and scales than existing real-world paired SR benchmarks.

Thumbnail montage of the ScreenSR benchmark, covering diverse scenes, subjects, and multilingual text.

---

🏋️ Training

The provided training configs disable experiment tracking by default (report_to: none). To enable Weights & Biases logging, run wandb login with your own account and set report_to: wandb in the YAML config.

Multi-step Training

# VOSR-0.5B
torchrun --nproc_per_node=8 train_vosr.py --config configs/train_yml/multi_step/VOSR_0.5B.yml

# VOSR-1.4B
torchrun --nproc_per_node=8 train_vosr.py --config configs/train_yml/multi_step/VOSR_1.4B.yml

One-step Distillation

Requires a trained multi-step teacher checkpoint. Set teacher_ckpt and pretrained_ckpt in the YAML config.

# VOSR-0.5B one-step
torchrun --nproc_per_node=8 train_vosr_distill.py --config configs/train_yml/one_step/VOSR_0.5B.yml

# VOSR-1.4B one-step
torchrun --nproc_per_node=8 train_vosr_distill.py --config configs/train_yml/one_step/VOSR_1.4B.yml

---

🔍 Inference

Single-GPU inference.

VOSR 2.0 one-step (default)

VOSR 2.0 is the default one-step model (1.4B DiT, vision-only, with improved fine-structure restoration).

# VOSR2-1.4B one-step
python inference_vosr_onestep.py \
    -c preset/ckpts/VOSR2 \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4

Multi-step models

Multi-step sampling defaults to 25 steps (--infer_steps, default 25). Override if you need fewer or more function evaluations.

Two knobs mainly affect the trade-off between faithfulness to the LR input and generative detail (both can be set via CLI to override args.json):

• --cfg_scale - Higher values tend to emphasize fidelity to the condition; lower values give more generative freedom. The sweet spot depends on input degradation strength. In our experiments, roughly -0.5 to 2 is a usable range; 0.5 is a practical default.
• --weak_cond_strength_aelq - During training this is sampled uniformly in [0.05, 0.25] so the checkpoint supports a wide range at inference via the same flag (smaller -> more generative, larger -> more faithful). Default 0.1.

Benchmark presets (multi-step): for apples-to-apples evaluation we use --cfg_scale -0.5 on RealSR and --cfg_scale 0.5 on ScreenSR.

# Inputs under preset/datasets/inp_data follow RealSR-style evaluation; use --cfg_scale -0.5 (see benchmark presets above).

# VOSR-0.5B multi-step (25 steps)
python inference_vosr.py \
    -c preset/ckpts/VOSR_0.5B_ms \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4

# VOSR-1.4B multi-step (25 steps)
python inference_vosr.py \
    -c preset/ckpts/VOSR_1.4B_ms \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4

One-step models

One-step models use --infer_steps with default 1 (typical for distilled checkpoints).

# VOSR-0.5B one-step
python inference_vosr_onestep.py \
    -c preset/ckpts/VOSR_0.5B_os \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4

# VOSR-1.4B one-step
python inference_vosr_onestep.py \
    -c preset/ckpts/VOSR_1.4B_os \
    -i preset/datasets/inp_data \
    -o preset/results \
    -u 4

Common arguments

-c checkpoint path, -i input image or folder, -o output directory, -u upscale factor. Outputs are saved as PNG directly into the -o folder; color alignment defaults to wavelet (--align_method wavelet|adain|nofix).

• Multi-step (inference_vosr.py): --infer_steps (default 25), --cfg_scale, --weak_cond_strength_aelq (see above).
• One-step (inference_vosr_onestep.py): --infer_steps (default 1).
• Both: --posterior_mode / --no-posterior_mode — deterministic VAE latent (latent_dist.mode()) is the default, for reproducible results and seam-free tiled inference.

Tiled inference for large images

Large images are tiled in latent space to bound GPU memory:

• --tile_size 512 — DiT tile size in pixels (default 0 = disabled). Use for large images; overlapping tiles are Gaussian-blended.
• --vae_tile_size 1024 — VAE encoder/decoder tile size in pixels (default 0 = full-image VAE). The Qwen VAE attention is O(N)-memory, so full-image encode/decode already fits up to ~8K inputs; VAE tiling is only needed for extreme sizes.
• --vae_tile_overlap — VAE tile overlap (default: follows --tile_overlap).

# VOSR-1.4B one-step on a large image with DiT + VAE tiling
python inference_vosr_onestep.py \
    -c preset/ckpts/VOSR_1.4B_os \
    -i large_input.png \
    -o preset/results \
    -u 4 --tile_size 512 --vae_tile_size 1024

---

📮 Contact

If you have any questions, please feel free to contact: rong-yuan.wu@connect.polyu.hk

📚 Citation

If VOSR is useful for your research, please consider citing:

@inproceedings{wu2026vosr,
  title   = {VOSR: A Vision-Only Generative Model for Image Super-Resolution},
  author  = {Wu, Rongyuan and Sun, Lingchen and Zhang, Zhengqiang and Kong, Xiangtao and Zhao, Jixin and Wang, Shihao and Zhang, Lei},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  year    = {2026}
}

Acknowledgements

This project benefits from stable-diffusion-2-1-base, LightningDiT, Qwen-Image, DINOv2, BasicSR, Real-ESRGAN, RCGM, and Shortcut-models. We thank the authors for their open-source contributions.

⚖️ License

This project is released under the Apache License 2.0 unless otherwise noted. See LICENSE for details. Downloadable model weights, benchmark data, and external assets may have separate terms on their hosting pages.