--- 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

VOSR logo

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

A framework for native generative image restoration.

#### 🚩 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)1,2,* | [Lingchen Sun](https://scholar.google.com/citations?user=ZCDjTn8AAAAJ&hl=zh-CN)1,2,* | [Zhengqiang Zhang](https://scholar.google.com/citations?user=UX26wSMAAAAJ&hl=en)1,2 | [Xiangtao Kong](https://scholar.google.com/citations?user=lueNzSgAAAAJ&hl=zh-CN)1,2
[Jixin Zhao](https://scholar.google.com/citations?user=0Z89rfUAAAAJ)1,2 | [Shihao Wang](https://scholar.google.com/citations?user=7TWugs4AAAAJ&hl=zh-CN)1 | [Lei Zhang](https://www4.comp.polyu.edu.hk/~cslzhang/)1,2,† 1 The Hong Kong Polytechnic University 2 OPPO Research Institute * Equal contribution. † Corresponding author.

VOSR overview

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

## 📌 Quick Links - [📰 News](#news) - [🧰 Preparation](#preparation) - [🏋️ Training](#training) - [🔍 Inference](#inference) - [📮 Contact](#contact) - [📚 Citation](#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](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). ---

VOSR performance, efficiency, and training cost comparison

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

Qualitative comparison on natural and text-rich images

VOSR better preserves fine structures and text readability.

## 🧰 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: `, `. 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.

ScreenSR benchmark montage

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 ```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 ``` --- ## 🔍 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 ``` --- ## 📮 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: ```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.