Running pre-trained ML models in Godot

So I have been developing this GDExtension called iree.gd. It is mission to embed IREE, another cool project that compiles and runs ML models, into Godot. It took me quite a while, but finally It has reached alpha.
Hope you guys could check it out the sample.

I can't upload video to here, but you still can watch a small clips of iree.gd doing its magic here

Overview

Currently, it is still having a very thin abstraction layer, and some manual stuff need to be done in prior.

Preparing the ML model

Before doing anything, you'll need to download your trained ML model from the internet. You could find some from Kaggle. Then, you'll need to preprocess them with python via this detailed documentation. You'll need to set the HAL backend following this table depending on your targeted platform. As you generate the .vmfb byte code, you are ready to proceed to the next step.

Use the ML model

.vmfb files are automatically imported into the project as you drag and drop 'em into it. You could then use them in your script by loading them.

var model : IREEModule = preload("res://your_model.vmfb")
@export var model_2 : IREEModule # Drag and drop in editor

Then, you'll need to prepare your inputs with IREETensor, then execute the ML model.

var input := IREETensor.from_bytes(image.get_data(), [1, 50, 50, 3]) # Remember to consider the input type.
var outputs := module.run_module("module.main", [input])
for output in outputs:
    pass # Do something with the `output`.

The input and function name vary models to models. You could inspect the .vmfb file using iree-dump-module from the IREE packages you installed during preparing the ML model.

Here is one section of the dump of the ML model used:

...
Exported Functions:
  [  0] main(!vm.ref<?>) -> (!vm.ref<?>)
        iree.abi.declaration: sync func @main(%input0: tensor<1x50x50x3xf32> {ml_program.identifier = "input_0"}) -> (%output0: tensor<1x200x200x3xf32> {ml_program.identifier = "Identity"})
  [  1] __init() -> ()
...

You could kind of guess it, it has main function that takes a tensor of float32 (1x50x50x3xf32) and outputs another tensor of float32 (1x200x200x3xf32).

Sample

There is a readily available sample of esrgan for you to try in the release page.

The sample only work on windows and linux.

After download and open iree-gd-sample-*.zip with Godot 4.2+, you will need to do these:

1. set res://scenes/main/main.tscn to be loaded on start.
2. Drag and drop appropriate byte code in res://bytecodes/esrgan/esrgan.*.vmfb into UI/TextureRect's module export.
3. Run it.

Afterword

Hopefully, in the future, Godot developer could utilise this piece of technology and create something extraordinary and magical. Thank you for reading this.

Comments

[Techie Kho]

Do note that it is still a very young project depending on another not very mature project. There are plenty of challenges. Here to name a few.

• It is a port from IREE, another young but capable library maintained by Google, as it is still not very mature, some of the things can't be polished in GDExtension as well, the good thing is that it is GPU accelerated.

• Because ML models usually only stays in python, less concern is put onto documenting the format of input and output data. Each model has its own arrangement of bits and bytes. That one I couldn't do much, as from IREE side of view, we only get to know the shape but not the arrangement of the data.

• IREE uses compiler to compile multiple type of ML models (tflite, tensorflow, PyTorch) to get a unified byte code format to be executed in its runtime. So it is inevitable that the programmer to do this on their own. Yes, it is troublesome and I wholeheartedly admit it. But with my current machine, I don't have the spec to embed the compiler as you literally need to compile the whole llvm project 😞. We do have dream to support direct compilation in Godot to ease the workflow.

There is more problem to face...

[FJones]

Funny, I just spent a few weeks writing a neural network implementation with serialization for Godot. Still testing some intricacies, so it's not published, but there's quite a bit of room for good ML implementations in Godot!

[Techie Kho]

Cool, good luck on your implementation!