asandeistefan/jpeg-restore-meanflows

MeanFlow JPEG Restoration

Model weights and minimal inference code for the paper Efficient JPEG Restoration in the Wavelet Domain via Mean Flows - a 65M Linear DiT, trained with mean flows, that restores JPEG images in one or two steps, at up to 20 img/s on a single RTX 3090.

project repo , arxiv

Overview

• jpeg_model.py: model definition (torch only), plus load_jpeg_model() / restore() helpers
• sample.py: CLI example that loads the checkpoint and restores an image

Input

• image: RGB tensor of shape (B, 3, H, W) with values in [0, 1], the JPEG-degraded image. Any resolution works; sizes are padded internally to a multiple of 4.

Output

• RGB tensor of shape (B, 3, H, W) with values in [0, 1], the restored image

Requirements

pip install torch torchvision huggingface_hub

Any GPU with at least 1GB of VRAM (CPU works too).

Usage

import torch
from jpeg_model import load_jpeg_model, restore

model = load_jpeg_model("meanflow_jpeg_qf5_20.pth", device="cuda")

degraded = torch.rand(1, 3, 720, 1280).cuda()   # degraded JPEG pixels in [0, 1]
restored = restore(model, degraded, sample_steps=2)

Or from the command line:

python sample.py path/to/image.jpg --steps 2

Citation

@article{asandeiandradu2026efficient,
  title   = {Efficient JPEG Restoration in the Wavelet Domain via Mean Flows},
  author  = {Stefan-Alexandru Asandei and Mihai-Alexandru Radu},
  journal = {arXiv preprint arXiv:2608.28730},
  year    = {2026}
}

This project is licensed under the MIT License.