litert-community/Fast-Neural-Style-LiteRT

Fast Neural Style Transfer — LiteRT (on-device, fully-GPU, 4 styles)

Fast neural style transfer (PyTorch examples TransformerNet, Johnson et al.), converted to LiteRT and running fully on the CompiledModel GPU (ML Drift) on Android. Applies an artistic style to a photo — 4 styles (candy / mosaic / rain_princess / udnie), each a 3.5 MB fp16 graph.

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

nodes on GPU	350 / 350 LITERT_CL (full residency)
inference	~9 ms (256×256)
size	3.5 MB per style (fp16)
accuracy	device-vs-PyTorch corr 0.9998–0.9999 (all 4 styles)

image[1,3,256,256] (RGB 0-255) →[GPU: TransformerNet]→ stylized[1,3,256,256] (RGB 0-255)

Minimal usage

Android (Kotlin, CompiledModel GPU)

val model = CompiledModel.create(context.assets, "style_candy_fp16.tflite",
    CompiledModel.Options(Accelerator.GPU), null)
val inputs = model.createInputBuffers()
val outputs = model.createOutputBuffers()
inputs[0].writeFloat(chw)            // [1,3,256,256] RGB 0-255, NCHW
model.run(inputs, outputs)
val stylized = outputs[0].readFloat()    // [1,3,256,256] RGB 0-255

Python (desktop verification)

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

img = Image.open("photo.jpg").convert("RGB")
w, h = img.size; s = min(w, h)
img = img.crop(((w-s)//2, (h-s)//2, (w+s)//2, (h+s)//2)).resize((256, 256))
x = np.asarray(img, np.float32).transpose(2, 0, 1)[None]     # 0-255, no normalization

# candy / mosaic / rain_princess / udnie
it = Interpreter(model_path="style_candy_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] RGB 0-255
Image.fromarray(y.transpose(1, 2, 0).clip(0, 255).astype(np.uint8)).save("stylized.png")

How it converts (litert-torch) — three numerically-exact re-authorings

1. ReflectionPad2d → zero-pad (GATHER_ND → PAD; border-only difference).
2. Large conv activations → conv-weight scaling. The conv outputs reach ≈ |5000|, where the Mali delegate's fp16 conv accumulation loses precision → garbage (device corr 0.34 at full residency — residency ≠ correctness). Each conv is followed by an InstanceNorm (which is scale-invariant), so scaling those conv weights down so the output is ≈ |10| is exact (IN output unchanged) and keeps the fp16 accumulation precise → corr 1.0.
3. InstanceNorm → SafeInstanceNorm (down-scaled-domain spatial reduction, fp16-safe; SafeLayerNorm class).

Upsample is interpolate(nearest) (no transposed conv → no ZeroStuff). Result: banned ops NONE, ≤4D, tflite-vs-torch corr 1.0, device-vs-torch corr 0.9999.

Preprocessing

Center-crop to square, resize to 256×256, RGB 0–255 (no normalization), NCHW. Output is 0–255 RGB (clamp).

License

BSD-3-Clause. Upstream: pytorch/examples.