File size: 3,712 Bytes
780c44a 849d4cb 780c44a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | ---
license: mit
library_name: LiteRT
pipeline_tag: image-to-image
tags:
- litert
- tflite
- on-device
- android
- gpu
- image-restoration
- denoising
- nafnet
base_model: megvii-research/NAFNet
---
# NAFNet-SIDD-width32 — LiteRT (on-device image denoising, fully-GPU)
[NAFNet](https://github.com/megvii-research/NAFNet) (Nonlinear Activation Free Network, ECCV 2022) image
restoration, converted to **LiteRT** and running **fully on the `CompiledModel` GPU** (ML Drift) on Android.
This is the **SIDD-width32** variant — real-image **denoising**. NAFNet is a U-Net of NAFBlocks with **no
activation functions** (SimpleGate = channel-split multiply), so the whole network is a clean CNN on the GPU.
## On-device (Pixel 8a, Tensor G3 — verified)
| | |
|---|---|
| nodes on GPU | **2179 / 2179** LITERT_CL (full residency) |
| inference | **~46 ms** (256×256) |
| size | 62.5 MB (fp16) |
| accuracy | device output **== PyTorch (corr 0.999999)** — re-authoring is numerically exact |
```
image[1,3,256,256] (RGB [0,1]) →[GPU: NAFNet U-Net]→ denoised[1,3,256,256]
```
## Minimal usage
**Android (Kotlin, CompiledModel GPU)**
```kotlin
val model = CompiledModel.create(context.assets, "nafnet_sidd_width32_fp16.tflite",
CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()
inputs[0].writeFloat(chw) // [1,3,256,256] RGB in [0,1], NCHW
model.run(inputs, outputs)
val denoised = outputs[0].readFloat() // [1,3,256,256] in [0,1]
```
**Python (desktop verification)**
```python
import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter
img = Image.open("noisy.jpg").convert("RGB").resize((256, 256))
x = (np.asarray(img, np.float32) / 255.0).transpose(2, 0, 1)[None] # [1,3,256,256]
it = Interpreter(model_path="nafnet_sidd_width32_fp16.tflite"); it.allocate_tensors()
it.set_tensor(it.get_input_details()[0]["index"], x); it.invoke()
y = it.get_tensor(it.get_output_details()[0]["index"])[0] # [3,256,256], [0,1]
Image.fromarray((y.transpose(1, 2, 0).clip(0, 1) * 255).astype(np.uint8)).save("restored.png")
```
A complete Android sample (image picker + before/after) is in the official
[google-ai-edge/litert-samples](https://github.com/google-ai-edge/litert-samples) repo under
`compiled_model_api/image_restoration`.
## How it converts (litert-torch)
Pure CNN (no activations). Three numerically-exact re-authorings, the headline being **SafeLayerNorm**:
NAFNet's residual stream grows large (|x|≈175 at the bottleneck), so the LayerNorm channel reductions
`Σ_c x` and `Σ_c (x−μ)²` (~15M) **overflow fp16 (max 65504)** on the Mali delegate (which computes in fp16
regardless of the model dtype) → a grid artifact. Doing the reductions in a down-scaled `x/S` domain (S=128)
and rescaling is exact and fp16-safe. Plus the Simplified Channel Attention `AdaptiveAvgPool2d(1)` →
`mean(3).mean(2)`, and the upsample `Conv2d(1×1)+PixelShuffle(2)` → depth-to-space `ZeroStuffConvT2d`.
Result: banned ops NONE, all tensors ≤4D, tflite-vs-torch corr **1.0**, device-vs-torch corr **1.0**.
A complete Android sample (image picker + before/after) is in the official
[google-ai-edge/litert-samples](https://github.com/google-ai-edge/litert-samples) repo under
`compiled_model_api/image_restoration` (push this `.tflite` in place of the deblur model).
## License
[MIT](https://github.com/megvii-research/NAFNet/blob/main/LICENSE). Upstream:
[megvii-research/NAFNet](https://github.com/megvii-research/NAFNet); weights NAFNet-SIDD-width32.
|