Train Your Own Z-Image Turbo LoRA on cloud GPUs

Source: Dev.to

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) with captions
• ~2 hours for setup and training

Set Up the EC2 Instance

1. Choose an instance – e.g., g6e.2xlarge (48 GB VRAM).
2. Configure the security group – open inbound ports 8675, 8888, and 22.
3. Select an AMI – Deep Learning OSS Nvidia Driver AMI GPU PyTorch 2.9 (Ubuntu 22.04), which includes CUDA and drivers.
4. Allocate storage – at least 100 GB for models, datasets, and checkpoints.

SSH into the instance

ssh -i your-key.pem ubuntu@YOUR_PUBLIC_IP

Install System Packages and the AI Toolkit

# Update system
sudo apt update && sudo apt upgrade -y

# Install basics
sudo apt install -y git build-essential python3-pip python3-venv

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

# Set up a virtual environment
python3 -m venv venv
source venv/bin/activate

# Install PyTorch (CUDA 12.6) and other requirements
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

Install Node.js (for the UI)

# Install nvm
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"

# Install Node.js 23 (or any 18+ version)
nvm install 23

# Verify
node -v   # should show v23.x.x
npm -v

Run the Toolkit UI

cd ~/ai-toolkit
source venv/bin/activate
cd ui
npm run build_and_start

The UI runs on port 8675. Create an SSH tunnel from your local machine:

ssh -i your-key.pem -L 8675:localhost:8675 -N ubuntu@YOUR_PUBLIC_IP

Then open http://localhost:8675 in a browser.

Set Up Jupyter Notebook

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

Edit ~/.jupyter/jupyter_notebook_config.py and add:

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'

Start Jupyter and tunnel the port:

jupyter notebook --no-browser --port=8888

On your local machine:

ssh -i your-key.pem -L 8888:localhost:8888 -N ubuntu@YOUR_PUBLIC_IP

Access Jupyter at http://localhost:8888 and log in with the password you set.

Create a Dataset in the UI

1. In the AI Toolkit UI, go to Datasets → New Dataset.
2. Name it (e.g., MySubject).
3. Upload the 6–15 images.
4. Add captions that include a unique trigger word (e.g., samt).
   • Caption format: [TRIGGER_WORD], a man with wavy hair …

Tips

• Use a non‑dictionary trigger word (e.g., skt, samt).
• Keep image quality high; vary pose and lighting but keep the subject consistent.

Create a Training Job

Test the Trained LoRA

Run the following code in a Jupyter notebook (replace paths and trigger word as needed):

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 an image
prompt = "YOUR_TRIGGER_WORD, posing in front of the Eiffel Tower"
generator = torch.Generator(device).manual_seed(42)

print("Generating image...")
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    num_inference_steps=9,   # 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)

Replace YOUR_TRIGGER_WORD with the word you used in the captions. After running the notebook, you should obtain a custom Z‑Image Turbo LoRA that generates lifelike images of your subject in under a second.