xocialize's picture
Upload README.md with huggingface_hub
df2b6f8 verified
|
Raw
History Blame Contribute Delete
4.58 kB
---
library_name: mlx
license: mit
license_link: https://github.com/kkkls/FFTformer/blob/main/license
base_model: kkkls/FFTformer
pipeline_tag: image-to-image
tags:
- mlx
- deblurring
- image-deblurring
- image-restoration
- fftformer
---
# mlx-community/FFTformer-GoPro-fp32
[FFTformer](https://github.com/kkkls/FFTformer) motion deblurring (CVPR 2023, `fftformer_GoPro.pth`),
converted to **Apple MLX** for Apple-Silicon inference via the
[`mlx-fftformer-swift`](https://github.com/xocialize/mlx-fftformer-swift) Swift package.
Kong et al., *Efficient Frequency Domain-based Transformers for High-Quality Image Deblurring*.
16,560,474 parameters. Blurry image in β†’ deblurred image at the same resolution out.
The architecture replaces the `QKα΅€` matmul with an element-wise product of two spectra inside each
8Γ—8 patch (`FSAS`), and gates a learned per-frequency mask in the feed-forward path (`DFFN`) β€” O(N)
in pixels rather than O(NΒ²).
## Use with mlx-fftformer-swift (Swift / MLX)
```swift
// Package.swift β†’ .package(url: "https://github.com/xocialize/mlx-fftformer-swift", from: "0.1.0")
import FFTformerMLXCore
let model = FFTformer()
try model.loadWeights(from: weightsURL) // model.safetensors from this repo
let deblurred = model.restoreTiled(imageNHWC) // NHWC RGB in [0,1]
```
Or as an MLXEngine `imageRestore` ModelPackage (`MLXFFTformer.FFTformerRestorePackage`), which
resolves this repo automatically via the Hub.
## Two things to know before you use it
**1. Tile. Do not run this model full-frame.** Level 1 runs at full resolution and both blocks expand
channels 6Γ— (48 β†’ 288), and the patch decomposition copies that again. A single 1080p frame peaks at
about **40 GB** of MLX allocation (measured; predicted within 10% by arithmetic). `restoreTiled`
defaults to a 256-px tile, which holds 1080p to ~8 GB.
⚠️ **Tile geometry must be 32-aligned.** Every FFT block decomposes from the *tile origin*, and a
level-3 patch spans 32 full-res pixels, so a tile origin that is not ≑ 0 (mod 32) shifts the patch
grid out of phase. Measured cost of getting this wrong: ~20.6 dB vs ~26 dB. Both tile size and
overlap must be multiples of 32.
Tiling also matches how the model was trained β€” GoPro used **128Γ—128 crops**, so it has never seen a
full frame.
**2. This build is fp32, not fp16 β€” measured, against the usual "small conv nets ship fp16" rule.**
End-to-end vs the fp32 reference on a 256Β² input:
| dtype | cosine | PSNR vs fp32 |
|---|---|---|
| fp32 | 0.99999994 | 108.99 dB |
| fp16 | 0.99994284 | 50.13 dB |
| bf16 | 0.98789388 | **30.09 dB** |
bf16 is disqualified: 30 dB of *dtype* error is the same order as the model's entire task signal
(GoPro 34.21 dB). fp16 is viable, but at only 66 MB the saving is not worth leaving the dtype error
~16 dB below the signal.
Note that **fp16 beats bf16 here by 20 dB**, inverting the usual default. Peak activation is ~2955,
comfortably inside fp16's range, so nothing overflows β€” this is mantissa precision, where fp16's 10
bits beat bf16's 7. If you re-quantize, keep the 54 learned `.fft` masks in fp32; the FFT interior is
fp32 by construction upstream (`rfft2(x.float())`).
## Conversion
MLX **NHWC** layout: conv weights `(O,I,kH,kW) β†’ (O,kH,kW,I)` (257 tensors), learned FFT masks
`(C,1,1,8,5) β†’ (8,5,C)` for trailing-dim broadcast (54 tensors), and the `Downsample`/`Upsample`
`nn.Sequential` convs renamed `.body.1. β†’ .conv.` (4 tensors). 535 tensors total, loading
`strict=True` / `verify: .all` clean.
## Parity
Gated against the PyTorch oracle on the CPU stream, fp32, judged on **relative** error:
- **key contract** β€” 535 tensors, 16,560,474 params, 0 missing / 0 unused / 0 shape mismatches
- **primitives** β€” 7/7; `patchify`, `unpatchify` and the FFT gate are **bit-identical** (0.00e+00),
i.e. MLX's `rfft2`/`irfft2` reproduce PyTorch exactly on these 8Γ—8 patches
- **blocks** β€” 5/5, worst relative 5.4e-07
- **full model** β€” 64Β² / 128Β² / 256Β² all at cosine **1.00000000**; structured 256Β² tile 3.6e-05
## Honest limitation
GoPro rank is a weak predictor of real-world value. *"Deblurring in the Wild"* (2026) found that on
real smartphone blur **every** GoPro-trained SOTA method scored *below the blurry input* β€” baseline
32.38 dB, FFTformer 32.25, the best of a bad set. GoPro's blur is frame-averaged camera shake and so
globally correlated, which is not what handheld capture produces. Validate on your own footage before
relying on it.
Weights: MIT (kkkls/FFTformer, committed in-repo). Port code: MIT.