Running pre-trained ML models in Godot
Tags: godot, machinelearning, gamedev, ai
Published: 2024-01-18 12:21:04

So I have been developing this GDExtension called [iree.gd](https://github.com/RechieKho/iree.gd). It is mission to embed [IREE](https://github.com/openxla/iree), another cool project that compiles and runs ML models, into Godot. It took me quite a while, but finally It has reached [alpha](https://github.com/RechieKho/iree.gd/releases). 
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](https://github.com/RechieKho/iree.gd/issues/37#issuecomment-1897703268)

## 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](https://www.kaggle.com/models). Then, you'll need to preprocess them with python [via this detailed documentation](https://iree.dev/guides/ml-frameworks/tflite/). You'll need to set the HAL backend following [this table depending on your targeted platform](https://github.com/RechieKho/iree.gd#supported-platforms). 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.

```swift
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.

```swift
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](https://github.com/RechieKho/iree.gd/releases). 

> 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.