litert-community/NAFNet-GoPro-width32-LiteRT

NAFNet-GoPro-width32 — LiteRT (on-device image deblur, fully-GPU)

NAFNet (Nonlinear Activation Free Network, ECCV 2022) image restoration, converted to LiteRT and running fully on the CompiledModel GPU (ML Drift) on Android. NAFNet is a U-Net of NAFBlocks with no activation functions at all (SimpleGate = channel-split multiply), so the whole network is a clean CNN on the GPU delegate. This is the GoPro-width32 variant — motion deblur.

On-device (Pixel 8a, Tensor G3 — verified)

nodes on GPU	2179 / 2179 LITERT_CL (full residency)
inference	~42 ms (256×256)
size	38 MB (fp16)
accuracy	device output == PyTorch (corr 1.000000) — re-authoring is numerically exact

image[1,3,256,256] (RGB [0,1]) →[GPU: NAFNet U-Net]→ restored[1,3,256,256]

Minimal usage

Android (Kotlin, CompiledModel GPU)

val model = CompiledModel.create(context.assets, "nafnet_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 restored = outputs[0].readFloat()    // [1,3,256,256] in [0,1]

Python (desktop verification)

import numpy as np
from PIL import Image
from ai_edge_litert.interpreter import Interpreter

img = Image.open("blurry.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_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 repo under compiled_model_api/image_restoration.

How it converts (litert-torch)

NAFNet is fully convolutional (any size that is a multiple of 16; exported here at 256×256). Three numerically-exact GPU re-authorings:

1. LayerNorm2d → fp16-safe channel LayerNorm. 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 numerically exact and fp16-safe.
2. Simplified Channel Attention AdaptiveAvgPool2d(1) → mean(3).mean(2) (two single-axis means).
3. Upsample Conv2d(1×1)+PixelShuffle(2) → Conv2d + depth-to-space ZeroStuffConvT2d.

Result: banned ops NONE, all tensors ≤4D, tflite-vs-torch corr 1.0, device-vs-torch corr 1.0.

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean.

Runtime	Backend	Graph on GPU	Latency
LiteRT CompiledModel (LITERT_CL)	GPU	2179 / 2179	~42 ms
TFLite benchmark_model (TfLiteGpuDelegateV2)	GPU (OpenCL)	2179 / 2179	97.2 ms
TFLite benchmark_model	CPU (XNNPACK, 4 threads)	—	XNNPACK declined the graph

The two GPU rows are different runtimes, not a contradiction. The LITERT_CL figure is the one recorded when this model shipped, taken through LiteRT's own CompiledModel accelerator — the path the Kotlin sample app and the LiteRT API use. The TfLiteGpuDelegateV2 figure is the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. They agree on how much of the graph the GPU takes; they disagree on speed, and the classic delegate is the slower of the two here. Read the TfLiteGpuDelegateV2 row as a reproducible floor, not as this model's speed on LiteRT.

XNNPACK declines these fp16 graphs — it reports failed to delegate DEPTHWISE_CONV_2D and then fails to allocate tensors — so there is no usable CPU number. Disabling XNNPACK falls back to reference kernels, which measured about 20× slower than the GPU on models of this size and would not represent CPU inference anyone would ship.

License

MIT. Upstream: megvii-research/NAFNet. Original weights: NAFNet-GoPro-width32 from the official release.