DEV Community

ishaan gupta
ishaan gupta

Posted on • Originally published at Medium

Train Your Own Z-Image Turbo LoRA on cloud GPUs

Z-Image is Alibaba’s 6B parameter image generation model that produces stunning images in just 8 inference steps. In this guide, you’ll learn how to train a custom LoRA on your own images using AWS EC2 GPU and the Ostris AI Toolkit.

By the end, you’ll have a working LoRA that generates images of your subject in seconds.

What You’ll Need

  • AWS Account with GPU instance access
  • SSH Key Pair for EC2 access
  • 6–15 high-quality images of your subject (1024×1024 recommended) along with captions of em
  • ~2 hours for setup and training

Well, let’s begin then.

Step 1: Launch Your EC2 Instance

We will go with g6e.2xlarge which is enough for both training and inferencing our model. It supports a 48 gigs of VRAM, which is more than enough.

Configure your security group. Open these inbound ports in your EC2 security group: 8675, 8888 & 22.

We will be using the Deep Learning OSS Nvidia Driver AMI GPU PyTorch 2.9 (Ubuntu 22.04) since it comes with CUDA and drivers pre-installed.

Allocate at least 100GB for models, datasets, and checkpoints.

Step 2: Connect and Install Dependencies

SSH into your instance:

ssh -i your-key.pem ubuntu@YOUR_PUBLIC_IP
Enter fullscreen mode Exit fullscreen mode

Update system and install basics:

sudo apt update && sudo apt upgrade -y
sudo apt install -y git build-essential python3-pip python3-venv
Enter fullscreen mode Exit fullscreen mode

Step 3: Install the AI Toolkit

Clone the repo and set up the python env:

cd ~
git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
git submodule update --init --recursive

python3 -m venv venv
source venv/bin/activate

pip3 install --no-cache-dir torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu126
pip3 install -r requirements.txt
Enter fullscreen mode Exit fullscreen mode

Step 4: Install Node.js for the Web UI

The AI Toolkit UI requires Node.js 18+:

cd ~
curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.2/install.sh | bash

export NVM_DIR="$HOME/.nvm"
[ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"

nvm install 23

node -v  # Verify: v23.x.x
npm -v
Enter fullscreen mode Exit fullscreen mode

Step 5: Launch the AI Toolkit UI

cd ~/ai-toolkit
source venv/bin/activate
cd ui
npm run build_and_start
Enter fullscreen mode Exit fullscreen mode

The toolkit UI runs on port 8675. On your local machine terminal, create an SSH tunnel using below command to access it:

ssh -i your-key.pem -L 8675:localhost:8675 -N ubuntu@YOUR_PUBLIC_IP
Enter fullscreen mode Exit fullscreen mode

Now open http://localhost:8675 in your browser. It should look like this

Step 6: Set Up Jupyter Notebook

For testing your trained LoRA:

cd ~/ai-toolkit
source venv/bin/activate

pip install jupyter jupyterlab ipykernel
python -m ipykernel install --user --name=ai-toolkit --display-name="AI Toolkit (Python)"

jupyter notebook --generate-config
jupyter notebook password  # Set a password and remember it for later

nano ~/.jupyter/jupyter_notebook_config.py
Enter fullscreen mode Exit fullscreen mode

Add these lines in the file

c.NotebookApp.ip = 'localhost'
c.NotebookApp.port = 8888
c.NotebookApp.open_browser = False
c.NotebookApp.allow_remote_access = True
c.NotebookApp.notebook_dir = '/home/ubuntu/ai-toolkit
Enter fullscreen mode Exit fullscreen mode

Start Jupyter and ssh into it via ur local machine

jupyter notebook --no-browser --port=8888
Enter fullscreen mode Exit fullscreen mode

Run below in your local terminal

ssh -i your-key.pem -L 8888:localhost:8888 -N ubuntu@YOUR_PUBLIC_IP
Enter fullscreen mode Exit fullscreen mode

Access Jupyter at http://localhost:8888 in ur browser . enter the password u set before.

Step 7: Prepare Your Dataset

  1. - In the AI Toolkit UI, go to Datasets → New Dataset
  2. - Name it (e.g., IAM)
  3. - Upload 6–15 high-quality images of your subject
  4. - Add captions with your trigger word (e.g., IAMD etc) Tips:
  • Use a unique, non-dictionary trigger word like sks or samt
  • Keep image quality high so the model learns what you feed it
  • Vary poses and lighting, but keep the subject consistent
  • captions format should be like this, [TRIGGER_WORD], a man with wavy hair……….

Step 8: Configure Your Training Job

Create a new job with below settings:

Job Settings

  • Training Name: your_training_name
  • Trigger Word: your_trigger_word
  • Model Architecture: Z-Image Turbo (w/ Training Adapter)
  • Low VRAM: Disable
  • Transformer: NONE
  • Cache Text Embeddings: Enable

Under advance section, enable the “Do differential Guidance” and set it to 3

Sample Prompts
Rewrite prompts according to your need, click on create jobs

Step 9: Test Your LoRA

Great, now time to test your LoRA. Use this Python code in jupyter to test your trained LoRA.

import torch
from diffusers import ZImagePipeline
device = "cuda"
dtype = torch.bfloat16  # Must be bfloat16, not float16!
print("Loading Z-Image Turbo pipeline...")
pipe = ZImagePipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=dtype,
    low_cpu_mem_usage=False,
)
pipe.to(device)
print("Pipeline loaded!")
# Load your trained LoRA
lora_path = "/home/ubuntu/ai-toolkit/output/my_first_lora_v1/my_first_lora_v1.safetensors"
print(f"Loading LoRA from: {lora_path}")
pipe.load_lora_weights(lora_path)
print("LoRA loaded!")

# Generate image with your trigger word
prompt = "ADAM, posing infront of Eiffel tower"

generator = torch.Generator(device).manual_seed(42)
print("Generating image...")
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    num_inference_steps=9,  # Results in 8 DiT forwards
    guidance_scale=0.0,     # Must be 0.0 for Turbo inference
    generator=generator,
).images[0]
image.save("output.png")
print("Done!")

# Display in Jupyter
from IPython.display import display
display(image)
Enter fullscreen mode Exit fullscreen mode

make sure to replace your token word in above prompt and run this code in jupyter notebook

Final Thoughts

You now have a custom Z-Image Turbo LoRA that generates real lifelike images of your subject in under a second. You can also download your LoRA and run it locally with the model. Would love to know how it worked for you guys, do share ur reviews or if any errors in the comments!

Top comments (0)