Running pre-trained ML models in Godot

rechie_kho

Techie Kho

Posted on January 18, 2024

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
Enter fullscreen mode Exit fullscreen mode

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`.
Enter fullscreen mode Exit fullscreen mode

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() -> ()
...
Enter fullscreen mode Exit fullscreen mode

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.

💖 💪 🙅 🚩
rechie_kho
Techie Kho

Posted on January 18, 2024

Join Our Newsletter. No Spam, Only the good stuff.

Sign up to receive the latest update from our blog.

Related