[ASASR] Coloring the Noise: Adversarial Sobolev Alignment for Faithful Image Super-Resolution

July 17, 2026 Β· View on GitHub

arXiv ICML 2026 Python 3.10 πŸ€— Demo πŸ€— Model PwC

Official implementation of the ICML 2026 paper:

Coloring the Noise: Adversarial Sobolev Alignment for Faithful Image Super-Resolution Hongbo Wang, Huaibo Huang, Pin Wang, Jinhua Hao, Chao Zhou, Ran He International Conference on Machine Learning (ICML), 2026 [arXiv:2605.23264]

News

  • 2026-07-16 β€” πŸ† ASASR is now indexed on Papers with Code: ranked SOTA on DIV2K-Val (Real-ESRGAN Γ—4) and #3 on RealSR!
  • 2026-07-07 β€” πŸ“Œ Poster session at ICML 2026, Seoul: Tue, Jul 7, 10:30 AM–12:15 PM KST, Hall A #1100. I'll be there in person β€” come by to chat!
  • 2026-07-04 β€” πŸ€— Online demo released: try ASASR on HuggingFace Spaces.
  • 2026-06-30 β€” Code and pretrained weights released.

Qualitative comparison against state-of-the-art methods

Highlights

ASASR pipeline overview

  • 🎨 Colored-Noise Flow β€” replaces isotropic Gaussian noise with a spectrally shaped kernel aligned to the natural image manifold.
  • πŸ“ Sobolev-Induced Geometry β€” reformulates the generative flow under a Riemannian metric that respects high-frequency structure (Sobolev frequency-weighted DPO, "S-DPO").
  • βš”οΈ Adversarial Manifold Guidance (AMG) β€” a Riesz-representation-based adversary network produces worst-case, spatially-aligned hard negatives for preference optimization ("AS-DPO").
  • πŸ–ΌοΈ Faithful Super-Resolution β€” built on FLUX.1-dev with dual-LoRA inference (SR prior + DPO refinement); strong improvements in spectral consistency, structural fidelity, and artifact suppression.

Online Demo

Try ASASR directly in your browser β€” no setup required: πŸ€— ASASR Super-Resolution on HuggingFace Spaces

Upload a low-resolution image (or pick an example), and the demo runs the FLUX.1-dev dual-LoRA pipeline to produce a faithful Γ—4 (512Γ—512) reconstruction. It uses the same weights hosted at wafer-bob/ASASR.

The demo was kindly built and hosted by the open-source team at Hugging Face β€” we thank them for their help. The Space is maintained independently of this repository.

Table of Contents

  1. Installation
  2. Pretrained Weights
  3. Data Preparation
  4. Training (full AS-DPO, 2 stages)
  5. Inference
  6. Evaluation
  7. Quick Smoke Test
  8. Scope & Notes

Repository Layout

ASASR/
β”œβ”€β”€ src/flux/                  # FLUX transformer / generate / condition / LoRA controller (OminiControl-style)
β”œβ”€β”€ train_adversary.py         # Stage-1: SFT the AMG adversary LoRA
β”œβ”€β”€ train.py                   # Stage-2: FLUX + Adversarial Sobolev-DPO (AS-DPO) trainer
β”œβ”€β”€ inference.py               # FLUX + dual-LoRA inference (SR + DPO)
β”œβ”€β”€ tools/
β”‚   β”œβ”€β”€ make_paired_data.py    # Real-ESRGAN-style degradation β†’ (HQ, LQ) pairs
β”‚   β”œβ”€β”€ build_dataset.py       # Pack (HQ, LQ) folders into a HuggingFace Dataset
β”‚   └── color_fix.py           # AdaIN / wavelet color correction (post-processing)
β”œβ”€β”€ eval/eval_pyiqa.py         # PSNR/SSIM/LPIPS/DISTS/MANIQA/MUSIQ/CLIPIQA+ via pyiqa
β”œβ”€β”€ scripts/
β”‚   β”œβ”€β”€ train_adversary.sh     # Stage-1 adversary training (multi-GPU)
β”‚   β”œβ”€β”€ train.sh               # Stage-2 AS-DPO training (multi-GPU)
β”‚   β”œβ”€β”€ infer.sh               # dual-LoRA inference
β”‚   └── eval.sh                # pyiqa evaluation
β”œβ”€β”€ examples/                  # sample (lr, gt) pairs for smoke tests
└── checkpoints/               # weights go here (downloaded separately, see Β§2)

1. Installation

# Conda (recommended)
conda env create -f environment.yml
conda activate asasr

# Or pip into an existing env
pip install -r requirements.txt

Tested on Python 3.10 + PyTorch 2.x (CUDA 11.8/12.x) + diffusers β‰₯ 0.30. For other CUDA versions, install the matching torch / torchvision wheels first, then pip install -r requirements.txt.

2. Pretrained Weights

The weights are not stored in this Git repository (each file exceeds GitHub's 100 MB limit). They are hosted on the HuggingFace Hub at wafer-bob/ASASR. Download them into checkpoints/ in one command:

pip install -U "huggingface_hub[cli]"
huggingface-cli download wafer-bob/ASASR --local-dir ./checkpoints

This yields the layout below (see also checkpoints/README.md):

WeightPathSizeNeeded for
ASASR DPO LoRAcheckpoints/dpo_lora/adapter_model.safetensors~111 MBinference
Base SR LoRAcheckpoints/sr_lora/pytorch_lora_weights_v2.safetensors~885 MBinference
AMG adversary LoRAcheckpoints/adv_lora/adapter_model.safetensors~111 MBtraining only (Stage-2)
FLUX.1-devlocal dir or HF Hub~24 GBboth
  • The base SR LoRA is trained on the OminiControl FLUX dual-stream condition framework.
  • For FLUX.1-dev, set FLUX_MODEL_PATH=/path/to/FLUX.1-dev to use a local checkout; otherwise it is pulled from the HuggingFace Hub (gated β€” export HF_TOKEN=... once). HF_HUB_OFFLINE=1 forces fully-offline use.

3. Data Preparation

ASASR trains on paired (HQ, LQ) crops packed as a HuggingFace dataset with two columns: jpg_0 (HQ / GT) and jpg_1 (LQ). We use DIV2K + LSDIR as HR sources and the higher-order Real-ESRGAN degradation pipeline (hyper-parameters aligned with SeeSR / DreamClear), at a scale factor of Γ—4 (HQ 512Γ—512 / LQ 128Γ—128).

3.1 Generate paired (HQ, LQ) data

python tools/make_paired_data.py \
    --gt_path /path/to/HR_images \
    --save_dir ./data/paired_train \
    --batch_size 192 --epoch 1

Output layout (./data/paired_train/{gt,lr}/*.png, HR 512Γ—512 / LR 128Γ—128). The degradation knobs default to the full SeeSR-aligned strength; use --deg_* to weaken them.

3.2 Pack into a HuggingFace dataset

python tools/build_dataset.py --paired_dir ./data/paired_train --output_dir ./data/dataset

This produces a datasets.Dataset saved with save_to_disk, ready to pass as --dataset_name ./data/dataset to the trainers.

4. Training (full AS-DPO, 2 stages)

The full paper method is Adversarial Sobolev-DPO (AS-DPO): a Sobolev frequency-weighted DPO objective whose hard negatives are synthesized on-the-fly by an adversary network (AMG). Reproduction is two stages.

Stage 1 β€” train the AMG adversary

The adversary is SFT-trained to mimic the reconstruction artifacts of baseline SR methods (e.g. Real-ESRGAN / SeeSR / SUPSR), so that during DPO it produces realistic, spatially-aligned hard negatives. Build a dataset whose jpg_0 = baseline-SR output (artifact proxy) and jpg_1 = the matching LQ (use tools/build_dataset.py, putting the baseline outputs in gt/):

ADV_DATASET=./data/adv_dataset bash scripts/train_adversary.sh
# β†’ outputs/adv_*/final_adv_lora/adapter_model.safetensors

Copy the result to checkpoints/adv_lora/adapter_model.safetensors (or pass it via ADV_LORA).

VarDefaultMeaning
ADV_DATASET./data/adv_datasetdataset with jpg_0=baseline output, jpg_1=LQ
LR5e-5AdamW learning rate (paper: adversary uses 5e-5)
RANK16adversary LoRA rank/Ξ± (paper main setting (16,16))
STEPS1000optimizer steps

Stage 2 β€” Adversarial Sobolev-DPO

bash scripts/train.sh                 # full AS-DPO (uses checkpoints/adv_lora)
VarDefaultMeaning
FLUX_MODEL_PATHblack-forest-labs/FLUX.1-devFLUX checkpoint dir / HF id
DATASET./data/datasetpacked (HQ, LQ) dataset
SR_LORA./checkpoints/sr_lora/...base SR LoRA (frozen reference)
ADV_LORA./checkpoints/adv_lora/adapter_model.safetensorsAMG adversary LoRA from Stage-1
ADV_STRENGTH0.1Ξ», adversarial perturbation strength
BETA4000raw DPO Ξ² (effective Ξ² = 0.5Β·BETA = 2000, the paper value)
SOBOLEV_S1.5Sobolev exponent s
NUM_GPUS / BATCH / GRAD_ACCUM8 / 8 / 8global batch = 512
LR1e-5AdamW learning rate
STEPS1000optimizer steps
RESOLUTION512HQ size (Γ—4 SR, LQ = 128)

Implementation details: BF16 mixed precision + gradient checkpointing, 200-step linear warmup, empty text prompts, a frozen-SR-LoRA + zero-init-DPO-LoRA dual-LoRA reference strategy, and a 28-step sampler at inference. The Sobolev spectral operator is realized with a 2-D FFT (rfft2) on the latent velocity residual.

Outputs: outputs/train_*/checkpoint-*/lora_dpo/adapter_model.safetensors, each plug-compatible with inference as DPO_LORA. Training was run on 8 GPUs.

5. Inference

bash scripts/infer.sh
# β†’ outputs/inference/*.png  (+ *_pair.png: input | output side-by-side)

Manual invocation:

python inference.py \
    --input_dir examples/lr \
    --output_dir outputs/inference \
    --sr_lora_path checkpoints/sr_lora/pytorch_lora_weights_v2.safetensors \
    --dpo_lora_path checkpoints/dpo_lora/adapter_model.safetensors \
    --sr_scale 1.0 --dpo_scale 1.0 \
    --resolution 512 --num_gpus 1 --save_pair
  • FLUX_MODEL_PATH=/path/to/FLUX.1-dev β€” use a local FLUX checkpoint.
  • HF_HUB_OFFLINE=1 β€” fully offline.
  • Multi-GPU via --num_gpus N (one process per GPU, images sharded round-robin).

6. Evaluation

eval/eval_pyiqa.py reports the paper's full metric set via pyiqa: PSNR / SSIM (Y channel of YCbCr), LPIPS / DISTS (full-reference perceptual), and MANIQA / MUSIQ / CLIPIQA+ (no-reference). SR/HR pairs are matched by filename stem; metric weights download on first run and cache under ~/.cache/torch/hub/pyiqa/.

bash scripts/eval.sh
# or, compare any number of methods:
python eval/eval_pyiqa.py --gt_dir examples/gt \
    --sr_dirs asasr:outputs/inference seesr:/path/to/seesr dreamclear:/path/to/dc

7. Quick End-to-End Smoke Test

After installation and placing sr_lora + dpo_lora under checkpoints/:

# (1) imports
python -c "from src.flux import transformer, condition, generate; \
           from tools.color_fix import adain_color_fix; print('ok')"

# (2) inference on the bundled examples (~1 min on one A100)
bash scripts/infer.sh

# (3) evaluate the run
bash scripts/eval.sh

8. Scope & Notes

  • Resolution. This release operates at Γ—4 SR with HQ 512 / LQ 128 (the configuration the released weights were trained for). All scripts default to 512.
  • Reproducibility. The released dpo_lora is produced by the Stage-2 AS-DPO pipeline above. Exact numbers depend on the full DIV2K+LSDIR training set and the 8-GPU global batch of 512; the bundled examples are for smoke-testing, not for matching paper tables.
  • In this repo: SR training (S-DPO / AS-DPO), the AMG adversary trainer, dual-LoRA inference, data synthesis, and the 7-metric IQA evaluation.
  • Not in this repo: the downstream-task evaluations (COCO detection/segmentation, ADE20K, OCR), the user study, and the spectral (LSD) analysis reported in the paper β€” these reuse standard external toolkits and are out of scope for this code release.

Citation

@inproceedings{wang2026asasr,
  title     = {Coloring the Noise: Adversarial Sobolev Alignment for Faithful Image Super-Resolution},
  author    = {Wang, Hongbo and Huang, Huaibo and Wang, Pin and Hao, Jinhua and Zhou, Chao and He, Ran},
  booktitle = {International Conference on Machine Learning (ICML)},
  year      = {2026}
}

Acknowledgements

ASASR is built on the FLUX.1-dev backbone, uses OminiControl for conditional LoRA control, and follows the BasicSR / SeeSR pipeline for degradation synthesis. We thank the authors of these projects.

License

This project is released under CC-BY-NC-4.0 for non-commercial research use only.

Copyright (c) 2026 The Authors and Kuaishou Technology.