VOSR-MLX / README.md

Model card and the original license

0201c51 verified about 9 hours ago

19.3 kB

---
base_model: CSWRY/VOSR
base_model_relation: quantized
license: apache-2.0
pipeline_tag: image-to-image
library_name: mlx
tags:
- mlx
- moonbox
- vosr
- super-resolution
- image-upscaling
- image-restoration
---
# VOSR 2.0 MLX

MLX builds of [VOSR 2.0](https://github.com/cswry/VOSR), 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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/fp16/config.json) | 0.8 kB | `7831ae8dacd5` | PSNR 42.3 dB mean / 31.4 worst, LPIPS <= 0.011 |
| fp16 | [`fp16/mlx_model.py`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/fp16/mlx_model.py) | 23.6 kB | `227a42122e08` | PSNR 42.3 dB mean / 31.4 worst, LPIPS <= 0.011 |
| fp16 | [`fp16/model.safetensors`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/fp16/model.safetensors) | 3.33 GB | `de064c413024` | PSNR 42.3 dB mean / 31.4 worst, LPIPS <= 0.011 |
| 8bit | [`8bit/config.json`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`):

```python
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](https://github.com/cswry/VOSR) 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

<div align="center">

<p align="center">
  <img src="assets/logo.png" alt="VOSR logo" width="72%">
</p>

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

<p align="center"><i>A framework for native generative image restoration.</i></p>

#### &#x1F6A9; Accepted by CVPR 2026

[![Paper](https://img.shields.io/badge/Paper-arXiv-b31b1b.svg)](https://arxiv.org/pdf/2604.03225)
[![VOSR-ModelScope](https://img.shields.io/badge/VOSR-ModelScope-6246EA.svg)](https://modelscope.cn/models/LULALULALU/VOSR_CKPT)
[![VOSR-HuggingFace](https://img.shields.io/badge/VOSR-HuggingFace-FCC624.svg)](https://huggingface.co/CSWRY/VOSR)
[![ScreenSR-ModelScope](https://img.shields.io/badge/ScreenSR-ModelScope-0EA5E9.svg)](https://modelscope.cn/datasets/LULALULALU/ScreenSR)
[![ScreenSR-HuggingFace](https://img.shields.io/badge/ScreenSR-HuggingFace-FFB000.svg)](https://huggingface.co/datasets/CSWRY/ScreenSR)

[Rongyuan Wu](https://scholar.google.com/citations?user=A-U8zE8AAAAJ&hl=zh-CN)<sup>1,2,&#42;</sup> |
[Lingchen Sun](https://scholar.google.com/citations?user=ZCDjTn8AAAAJ&hl=zh-CN)<sup>1,2,&#42;</sup> |
[Zhengqiang Zhang](https://scholar.google.com/citations?user=UX26wSMAAAAJ&hl=en)<sup>1,2</sup> |
[Xiangtao Kong](https://scholar.google.com/citations?user=lueNzSgAAAAJ&hl=zh-CN)<sup>1,2</sup> <br>
[Jixin Zhao](https://scholar.google.com/citations?user=0Z89rfUAAAAJ)<sup>1,2</sup> |
[Shihao Wang](https://scholar.google.com/citations?user=7TWugs4AAAAJ&hl=zh-CN)<sup>1</sup> |
[Lei Zhang](https://www4.comp.polyu.edu.hk/~cslzhang/)<sup>1,2,&dagger;</sup>

<sup>1</sup> The Hong Kong Polytechnic University  
<sup>2</sup> OPPO Research Institute

<sup>&#42;</sup> Equal contribution. <sup>&dagger;</sup> Corresponding author.

</div>

<!-- Optional teaser image -->
<!-- ![teaser](assets/teaser.png) -->

<p align="center">
  <img src="assets/overview.jpg" alt="VOSR overview" width="100%">
</p>
<p align="center"><em>Overview of the VOSR framework, including the overall pipeline and our condition / guidance design.</em></p>


## &#x1F4CC; Quick Links

- [&#x1F4F0; News](#news)
- [&#x1F9F0; Preparation](#preparation)
- [&#x1F3CB;&#xFE0F; Training](#training)
- [&#x1F50D; Inference](#inference)
- [&#x1F4EE; Contact](#contact)
- [&#x1F4DA; Citation](#citation)

<a id="news"></a>
## &#x1F4F0; 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](https://cswry.github.io/vosr2/). Thanks to [ylchen333](https://github.com/ylchen333) for bringing VOSR to ComfyUI with [ComfyUI-VOSR2](https://github.com/ylchen333/ComfyUI-VOSR2).

- 2026.04.10 Public release: training and inference code, [pretrained checkpoints](https://modelscope.cn/models/LULALULALU/VOSR_CKPT), bundled VAE / decoder assets, and the [ScreenSR](https://modelscope.cn/datasets/LULALULALU/ScreenSR) benchmark. Setup and file layout: [Preparation](#preparation); commands: [Training](#training) and [Inference](#inference).

---

<p align="center">
  <img src="assets/vosr_comparison_figure.png" alt="VOSR performance, efficiency, and training cost comparison" width="100%">
</p>
<p align="center"><em>Comparison with prior methods in performance, inference efficiency, and training cost.</em></p>

<p align="center">
  <img src="assets/vis_comp.jpg" alt="Qualitative comparison on natural and text-rich images" width="100%">
</p>
<p align="center"><em>VOSR better preserves fine structures and text readability.</em></p>

<a id="preparation"></a>
## &#x1F9F0; Preparation

### Dependencies and Installation

```bash
## 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
```

### &#x1F4E6; Model Weights

Download all pretrained weights from [ModelScope](https://modelscope.cn/models/LULALULALU/VOSR_CKPT) or [Hugging Face](https://huggingface.co/CSWRY/VOSR), and place them under `preset/ckpts/`. The expected structure:

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

### &#x1F5C2;&#xFE0F; 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`):

```text
/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.

### &#x1F9EA; New Real-World Paired Benchmark

Download the [ScreenSR benchmark from ModelScope](https://modelscope.cn/datasets/LULALULALU/ScreenSR) or [Hugging Face](https://huggingface.co/datasets/CSWRY/ScreenSR), place it wherever you like, and then point `-i` to that folder when running inference.

[ScreenSR](https://modelscope.cn/datasets/LULALULALU/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.

<p align="center">
  <img src="assets/montage_13x10_256px.jpg" alt="ScreenSR benchmark montage" width="100%">
</p>
<p align="center"><em>Thumbnail montage of the ScreenSR benchmark, covering diverse scenes, subjects, and multilingual text.</em></p>

---

<a id="training"></a>
## &#x1F3CB;&#xFE0F; 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

```bash
# 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.

```bash
# 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
```

---

<a id="inference"></a>
## &#x1F50D; 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).

```bash
# 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](https://modelscope.cn/datasets/LULALULALU/ScreenSR).

```bash
# 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).

```bash
# 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`).

```bash
# 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
```

---

<a id="contact"></a>
## &#x1F4EE; Contact

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

<a id="citation"></a>
## &#x1F4DA; Citation

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

```bibtex
@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](https://huggingface.co/Manojb/stable-diffusion-2-1-base), [LightningDiT](https://github.com/hustvl/LightningDiT), [Qwen-Image](https://huggingface.co/Qwen/Qwen-Image), [DINOv2](https://github.com/facebookresearch/dinov2), [BasicSR](https://github.com/XPixelGroup/BasicSR), [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN), [RCGM](https://github.com/LINs-lab/RCGM), and [Shortcut-models](https://github.com/kvfrans/shortcut-models). We thank the authors for their open-source contributions.

## &#x2696;&#xFE0F; License

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

	---
	base_model: CSWRY/VOSR
	base_model_relation: quantized
	license: apache-2.0
	pipeline_tag: image-to-image
	library_name: mlx
	tags:
	- mlx
	- moonbox
	- vosr
	- super-resolution
	- image-upscaling
	- image-restoration
	---

	# VOSR 2.0 MLX

	MLX builds of [VOSR 2.0](https://github.com/cswry/VOSR), 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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/fp16/config.json) \| 0.8 kB \| `7831ae8dacd5` \| PSNR 42.3 dB mean / 31.4 worst, LPIPS <= 0.011 \|
	\| fp16 \| [`fp16/mlx_model.py`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/fp16/mlx_model.py) \| 23.6 kB \| `227a42122e08` \| PSNR 42.3 dB mean / 31.4 worst, LPIPS <= 0.011 \|
	\| fp16 \| [`fp16/model.safetensors`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/fp16/model.safetensors) \| 3.33 GB \| `de064c413024` \| PSNR 42.3 dB mean / 31.4 worst, LPIPS <= 0.011 \|
	\| 8bit \| [`8bit/config.json`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`](https://huggingface.co/moonbox-ai/VOSR-MLX/blob/main/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`):

	```python
	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](https://github.com/cswry/VOSR) 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

	<div align="center">

	<p align="center">
	<img src="assets/logo.png" alt="VOSR logo" width="72%">
	</p>

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

	<p align="center"><i>A framework for native generative image restoration.</i></p>

	#### 🚩 Accepted by CVPR 2026

	[![Paper](https://img.shields.io/badge/Paper-arXiv-b31b1b.svg)](https://arxiv.org/pdf/2604.03225)
	[![VOSR-ModelScope](https://img.shields.io/badge/VOSR-ModelScope-6246EA.svg)](https://modelscope.cn/models/LULALULALU/VOSR_CKPT)
	[![VOSR-HuggingFace](https://img.shields.io/badge/VOSR-HuggingFace-FCC624.svg)](https://huggingface.co/CSWRY/VOSR)
	[![ScreenSR-ModelScope](https://img.shields.io/badge/ScreenSR-ModelScope-0EA5E9.svg)](https://modelscope.cn/datasets/LULALULALU/ScreenSR)
	[![ScreenSR-HuggingFace](https://img.shields.io/badge/ScreenSR-HuggingFace-FFB000.svg)](https://huggingface.co/datasets/CSWRY/ScreenSR)

	[Rongyuan Wu](https://scholar.google.com/citations?user=A-U8zE8AAAAJ&hl=zh-CN)<sup>1,2,*</sup> \|
	[Lingchen Sun](https://scholar.google.com/citations?user=ZCDjTn8AAAAJ&hl=zh-CN)<sup>1,2,*</sup> \|
	[Zhengqiang Zhang](https://scholar.google.com/citations?user=UX26wSMAAAAJ&hl=en)<sup>1,2</sup> \|
	[Xiangtao Kong](https://scholar.google.com/citations?user=lueNzSgAAAAJ&hl=zh-CN)<sup>1,2</sup> <br>
	[Jixin Zhao](https://scholar.google.com/citations?user=0Z89rfUAAAAJ)<sup>1,2</sup> \|
	[Shihao Wang](https://scholar.google.com/citations?user=7TWugs4AAAAJ&hl=zh-CN)<sup>1</sup> \|
	[Lei Zhang](https://www4.comp.polyu.edu.hk/~cslzhang/)<sup>1,2,&dagger;</sup>

	<sup>1</sup> The Hong Kong Polytechnic University
	<sup>2</sup> OPPO Research Institute

	<sup>*</sup> Equal contribution. <sup>&dagger;</sup> Corresponding author.

	</div>

	<!-- Optional teaser image -->
	<!-- ![teaser](assets/teaser.png) -->

	<p align="center">
	<img src="assets/overview.jpg" alt="VOSR overview" width="100%">
	</p>
	<p align="center"><em>Overview of the VOSR framework, including the overall pipeline and our condition / guidance design.</em></p>


	## 📌 Quick Links

	- [📰 News](#news)
	- [🧰 Preparation](#preparation)
	- [🏋️ Training](#training)
	- [🔍 Inference](#inference)
	- [📮 Contact](#contact)
	- [📚 Citation](#citation)

	<a id="news"></a>
	## 📰 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](https://cswry.github.io/vosr2/). Thanks to [ylchen333](https://github.com/ylchen333) for bringing VOSR to ComfyUI with [ComfyUI-VOSR2](https://github.com/ylchen333/ComfyUI-VOSR2).

	- 2026.04.10 Public release: training and inference code, [pretrained checkpoints](https://modelscope.cn/models/LULALULALU/VOSR_CKPT), bundled VAE / decoder assets, and the [ScreenSR](https://modelscope.cn/datasets/LULALULALU/ScreenSR) benchmark. Setup and file layout: [Preparation](#preparation); commands: [Training](#training) and [Inference](#inference).

	---

	<p align="center">
	<img src="assets/vosr_comparison_figure.png" alt="VOSR performance, efficiency, and training cost comparison" width="100%">
	</p>
	<p align="center"><em>Comparison with prior methods in performance, inference efficiency, and training cost.</em></p>

	<p align="center">
	<img src="assets/vis_comp.jpg" alt="Qualitative comparison on natural and text-rich images" width="100%">
	</p>
	<p align="center"><em>VOSR better preserves fine structures and text readability.</em></p>

	<a id="preparation"></a>
	## 🧰 Preparation

	### Dependencies and Installation

	```bash
	## 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](https://modelscope.cn/models/LULALULALU/VOSR_CKPT) or [Hugging Face](https://huggingface.co/CSWRY/VOSR), and place them under `preset/ckpts/`. The expected structure:

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

	```text
	/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](https://modelscope.cn/datasets/LULALULALU/ScreenSR) or [Hugging Face](https://huggingface.co/datasets/CSWRY/ScreenSR), place it wherever you like, and then point `-i` to that folder when running inference.

	[ScreenSR](https://modelscope.cn/datasets/LULALULALU/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.

	<p align="center">
	<img src="assets/montage_13x10_256px.jpg" alt="ScreenSR benchmark montage" width="100%">
	</p>
	<p align="center"><em>Thumbnail montage of the ScreenSR benchmark, covering diverse scenes, subjects, and multilingual text.</em></p>

	---

	<a id="training"></a>
	## 🏋️ 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

	```bash
	# 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.

	```bash
	# 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
	```

	---

	<a id="inference"></a>
	## 🔍 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).

	```bash
	# 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](https://modelscope.cn/datasets/LULALULALU/ScreenSR).

	```bash
	# 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).

	```bash
	# 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`).

	```bash
	# 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
	```

	---

	<a id="contact"></a>
	## 📮 Contact

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

	<a id="citation"></a>
	## 📚 Citation

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

	```bibtex
	@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](https://huggingface.co/Manojb/stable-diffusion-2-1-base), [LightningDiT](https://github.com/hustvl/LightningDiT), [Qwen-Image](https://huggingface.co/Qwen/Qwen-Image), [DINOv2](https://github.com/facebookresearch/dinov2), [BasicSR](https://github.com/XPixelGroup/BasicSR), [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN), [RCGM](https://github.com/LINs-lab/RCGM), and [Shortcut-models](https://github.com/kvfrans/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](LICENSE) for details. Downloadable model weights, benchmark data, and external assets may have separate terms on their hosting pages.