--- license: bsd-3-clause library_name: LiteRT pipeline_tag: image-to-image tags: [litert, tflite, on-device, android, gpu, style-transfer, neural-style, fast-neural-style] base_model: pytorch/examples --- # Fast Neural Style Transfer — LiteRT (on-device, fully-GPU, 4 styles) Fast neural **style transfer** ([PyTorch examples](https://github.com/pytorch/examples/tree/main/fast_neural_style) `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. ![Fast Neural Style — content + candy / mosaic / rain / udnie (on-device LiteRT GPU)](samples/sample.png) ## 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)** ```kotlin 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)** ```python 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). ## Performance Measured on a **Pixel 8a** (Tensor G3, Android 16) with the standard TFLite [`benchmark_model`](https://ai.google.dev/edge/litert/models/measurement) tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean. | Runtime | Backend | Graph on GPU | Latency | |---|---|---|---| | TFLite `benchmark_model` (`TfLiteGpuDelegateV2`) — `style_udnie_fp16.tflite` | GPU (OpenCL) | 350 / 350 | 37.4 ms | | TFLite `benchmark_model` (`TfLiteGpuDelegateV2`) — `style_mosaic_fp16.tflite` | GPU (OpenCL) | 350 / 350 | 37.5 ms | | TFLite `benchmark_model` (`TfLiteGpuDelegateV2`) — `style_rain_princess_fp16.tflite` | GPU (OpenCL) | 350 / 350 | 37.7 ms | | TFLite `benchmark_model` (`TfLiteGpuDelegateV2`) — `style_candy_fp16.tflite` | GPU (OpenCL) | 350 / 350 | 37.7 ms | | TFLite `benchmark_model` — `style_udnie_fp16.tflite` | CPU (XNNPACK, 4 threads) | — | 406.0 ms | | TFLite `benchmark_model` — `style_mosaic_fp16.tflite` | CPU (XNNPACK, 4 threads) | — | 404.3 ms | | TFLite `benchmark_model` — `style_rain_princess_fp16.tflite` | CPU (XNNPACK, 4 threads) | — | 404.5 ms | | TFLite `benchmark_model` — `style_candy_fp16.tflite` | CPU (XNNPACK, 4 threads) | — | 402.5 ms | **Any on-device figure recorded when this model shipped came from a different runtime.** It was taken through LiteRT's own `CompiledModel` accelerator (logcat reports it as `LITERT_CL`), which is the path the Kotlin sample app and the LiteRT API use, and it appears elsewhere on this card. The rows above are the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. The two are not comparable, so read the rows above as a reproducible floor rather than as this model's speed on LiteRT. ## License [BSD-3-Clause](https://github.com/pytorch/examples/blob/main/LICENSE). Upstream: [pytorch/examples](https://github.com/pytorch/examples/tree/main/fast_neural_style).