GPU acceleration support
So I noticed in the repo notes that the current execution falls back to CPU due to the classic TFLite GPU delegate hardcoding tensors to 4D (crashing/failing on 5D spatiotemporal tensors).
But recently, google open sourced their new GPU delegate called "ML Drift", which explicitly adds support for 5D tensors, according to google: "5D Tensor Support: The TFLite GPU delegate is structurally hardcoded to 4D tensors, forcing developers to use layout hacks for complex models requiring 5D tensors. ML Drift reflected this long overdue functionality in the new LiteRT ML Drift GPU accelerator to enable 5D tensor support out-of-the-box. With this, LiteRT executes complex workloads like 3D convolutional networks for volumetric/spatial AI and spatiotemporal models (e.g., YOLO 11n, MobileViT v2, and Swin Transformer v2) directly on edge GPUs." So this could potentially make the model able to run on GPUs via the LiteRT runtime and signficantly speed up the generation process compared to the CPU.