Flax doesn’t support the load_textual_inversion() method, but the textual_inversion_flax.py script saves the learned embeddings as a part of the model after training. This means you can use the model for inference like any other Flax model:   Copied import jax