mlx-community/SCUNet-color-real-psnr-fp32

mlx-community/SCUNet-color-real-psnr-fp32

SCUNet blind real-world denoising, converted to Apple MLX for Apple-Silicon inference via `mlx-scunet-swift`.

Zhang et al., Practical Blind Denoising via Swin-Conv-UNet and Data Synthesis. 17,946,072 parameters (71.8 MB) — config=[4]*7, dim=64.

This is the `MSE` checkpoint. Trained on a fidelity objective. Conservative: it leaves residual noise rather than inventing detail, which is what you want when the output feeds another stage or is measured against a reference.

🔑 The point of this model: there is nothing to configure

SCUNet takes no noise level. Its sibling DRUNet takes σ as a model input and exposes a strength dial; NAFNet, FFTformer and Restormer bake a degradation into the checkpoint. SCUNet takes neither — one forward pass, no σ to estimate, nothing for a caller to get wrong on a real photograph whose noise level nobody has measured.

import SCUNetMLXCore

let model = SCUNet()                      // config=[4]*7 — NOT upstream's [2]*7 default
try model.loadWeights(from: weightsURL)
let clean = model.denoiseTiled(imageNHWC) // NHWC RGB in [0,1]

Or as an MLXEngine imageRestore ModelPackage (MLXSCUNet.SCUNetRestorePackage), which declares supportsStrength: false — the contract distinguishing a blind backer from a dialled one.

Architecture note: window attention tiles almost for free

SCUNet is a Swin-Conv-UNet — ConvTransBlock splits channels between a conv path and a shifted-window attention path and re-fuses them. Because the attention is strictly local (8x8 windows, alternating W / SW), tiling barely perturbs the result. Measured at 512^2, tiled at 256 with 64 overlap versus the full-frame reference:

A ratio of 1.00x means tile boundaries are statistically indistinguishable from ordinary image content. Restormer, whose attention is spatially global, does not get off this lightly.

Tile geometry must be 64-aligned: the forward pass pads to a multiple of 64 and lays the window grid out from the tile's own origin, so an unaligned origin shifts the window phase between neighbouring tiles and leaves a seam feathering cannot remove.

Conversion

MLX NHWC. 540 tensors: 117 Conv2d, 3 `ConvTranspose2d`, 112 Linear (passthrough), 308 passthrough. The two 4-D transposes cannot be told apart by shape:

Exactly m_up{1,2,3}.0.weight are the transposed convs; the converter asserts that count.

Two further traps worth knowing if you port this yourself:

• —`relative_position_params` is stored pre-permuted. The constructor allocates ((2w-1)^2, heads) and then re-assigns the parameter through .view(2w-1, 2w-1, heads).transpose(1,2).transpose(0,1) — so the checkpoint carries (heads, 2w-1, 2w-1). Read the constructor, not the declaration.
• —The QKV head split is not per-head triples. rearrange(qkv, 'b nw np (threeh c) -> threeh b nw np c').chunk(3, dim=0) puts all q heads, then all k, then all v. Splitting it the intuitive way is shape-identical and silently wrong.

Parity

Gated against the PyTorch oracle on the CPU stream, fp32, relative error:

• —key contract — 540 tensors / 17,946,072 params / 0 missing / 0 unused, strict load
• —attention internals — 5/5 at exactly 0.00e+00: the stored bias table, the gathered (heads, 64, 64) bias, and the SW attention mask
• —WMSA end-to-end — W 2.36e-07, SW 1.93e-07
• —blocks — Block and ConvTransBlock, both types, worst 4.57e-07
• —resamplers — the transposed conv is bit-identical (0.00e+00)
• —full model — 64^2 / 128^2 / 100^2, worst 3.42e-06 (100^2 exercises the internal ReplicationPad2d and the crop back)

Measured on real sensor noise

No primary source reports SCUNet's SIDD or DND — the authors deliberately skipped both — so we measured it on NIND (CC0): 5 scenes x 4 ISOs on a locked-off camera with a compensating shutter, 768^2 centre crops, PSNR against the ISO-100 reference. Pairs verified pixel-aligned and brightness-matched first.

real-psnr is the strongest blind denoiser in the set, and its margin over Restormer grows with noise (+0.01 -> +0.17 -> +0.70 dB) — consistent with the randomized-degradation training that is the model's whole thesis. A correctly-tuned DRUNet wins at ISO 1600, but that requires knowing sigma; at a wrong sigma it scores -2.73 dB, worse than leaving the image alone.

real-gan costs 0.88-1.50 dB and is effectively a no-op at ISO 1600 (+0.02 dB) — a perceptual mode, not a default.

NAFNet-SIDD-width64 — the incumbent, and 6.5x larger at 116.0 M params — is last at every ISO, 3.00 / 2.96 / 2.90 dB behind, and at ISO 1600 it scores -1.48 dB, worse than leaving the image alone. That is the failure randomized-degradation training exists to fix: NAFNet trains on SIDD's five smartphone sensors, and NIND is DSLR-class Canon, so being off the training sensors is the whole test.

⚠️ This is a generalization result. It says nothing about NAFNet's in-domain SIDD performance, which is what a phone photo would exercise. NIND is DSLR-class hardware, and PSNR judges the GAN variant on the axis it deliberately trades away.

Code: Apache-2.0 (`cszn/SCUNet`). Weights: MIT, published first-party by the author in the `cszn/KAIR` v1.0 release.