Z-Image ComfyUI LoRA Lighting Fix and Auto Prompt Optimization: Complete Guide

Say goodbye to flat lighting and LoRA loading errors: Master advanced Z-Image LoRA workflows in ComfyUI, combined with automatic prompt optimization, to make your training results shine.

Have You Experienced These Problems?

"lora key not found" error when loading Z-Image LoRA in ComfyUI
Generated images have flat, flat lighting with no depth or dimension
LoRA training results look great in AI Toolkit but significantly worse in ComfyUI
Confusion about which LoRA format to use with Z-Image (LoRA vs LoRA-LyCORIS vs LoKR)

If any of these resonate, this article is for you.

The Root Cause: Z-Image's Unique Architecture

Z-Image uses a Single-Stream Diffusion Transformer architecture, fundamentally different from traditional Stable Diffusion:

Feature	Stable Diffusion XL	Z-Image
Architecture	U-Net + Dual Encoder	Single-Stream Diffusion Transformer
Parameters	~6.6B	~6B
LoRA Format	Standard SD LoRA	Needs adaptation
Default LoRA Node	Works directly	May corrupt data

Why Does the Default LoRA Node Fail?

ComfyUI's default LoRA loading node is designed for SDXL/SD1.5, assuming specific weight distributions and layer structures. When directly applied to Z-Image:

Key name mismatch: Z-Image's Transformer layer naming differs from SDXL
Different attention head shapes: The single-stream architecture has a different attention mechanism
Data corruption: Incorrect loading can overwrite or corrupt base model weights

This is the root cause of the "lora key not found" warning and flat lighting effects.

Solution 1: Use the Correct LoRA Node

Recommended: AI Toolkit Native LoRA Node

# Install Ostris AI Toolkit in ComfyUI
git clone https://github.com/ostris/ai-toolkit.git custom_nodes/ai-toolkit
pip install -r custom_nodes/ai-toolkit/requirements.txt

The AI Toolkit LoRA node supports Z-Image's Transformer architecture:

# ComfyUI workflow configuration
{
    "loRA_node": "AIT_LoRALoader",  # AI Toolkit dedicated node
    "model": "Tongyi-MAI/Z-Image-Turbo",
    "lora": "/path/to/your/lora.safetensors",
    "strength_model": 0.8,
    "strength_clip": 0.8
}

NOT Recommended: Standard ComfyUI LoRA Node

❌ LoadLoRA (standard node) — May cause:
   - Partial weight loading failure
   - Flat lighting effect
   - Loss of character features

Solution 2: Train with the Correct Format

Training Format Selection

Format	Z-Image Compatibility	Rating	Notes
Standard LoRA	✅ Fully compatible	⭐⭐⭐⭐⭐	Train with AI Toolkit
LoRA-LyCORIS	⚠️ Partially compatible	⭐⭐⭐	Needs extra config
LoKR	❌ Incompatible	⭐	Avoid
Full Fine-tune	✅ Fully compatible	⭐⭐⭐⭐	Large file size

AI Toolkit Training Config (Recommended)

# config.yaml — AI Toolkit training config
base_model: Tongyi-MAI/Z-Image-Base
precision: bf16
mixed_precision: bf16
optimizer:
  name: prodigy
  lr: 1.0
  weight_decay: 0.01

network:
  type: loha  # Standard LoRA format, best for Z-Image
  dim: 32
  alpha: 16
  modules:
    - .*to_q
    - .*to_k
    - .*to_v
    - .*to_out.0
    - .*ff.net.0
    - .*ff.net.2

dataset:
  directory: ./dataset
  caption_extension: .txt
  shuffle: true
  cache_latents: true

training:
  steps: 2000
  batch_size: 1
  save_every: 500
  save_precision: fp16

Key Training Parameters for Lighting

# Lighting preservation parameters
training:
  # Preserve lighting information
  keep_input_rgb: true
  # Gradient clipping to prevent training instability
  gradient_accumulation_steps: 4
  # Learning rate schedule
  lr_scheduler: cosine_with_restarts
  lr_warmup_steps: 100

Solution 3: Fix Flat Lighting

Problem Symptoms

After using incorrect LoRA loading, generated images show:

Missing or unnatural shadows
Overly uniform facial illumination
Lack of object depth and dimension
Overall grayish color tone

Fix Methods

Method A: Prompt Lighting Enhancement

Add lighting descriptors to your prompt:

# Basic prompt
portrait of a warrior

# Enhanced lighting prompt
cinematic portrait of a warrior, dramatic side lighting,
rim light, volumetric fog, golden hour,
studio lighting setup, soft shadows, depth of field

Method B: ComfyUI Lighting Enhancement Workflow

[Checkpoint] → [AI Toolkit LoRA] → [KSampler] → [VAE Decode]
                                    ↓
                          [CLIP Set Last Layer] (set to 20)
                                    ↓
                          [CLIP Vision Encode] (reference image lighting)

Setting CLIP Set Last Layer to 20 significantly enhances lighting effects — a Z-Image-specific optimization technique.

Method C: ControlNet Lighting Control

# ControlNet lighting reference workflow
from comfyui_nodes import ControlNetApply

controlnet = ControlNetApply(
    controlnet_model="light-controlnet-v1",
    strength=0.6,
    start_percent=0.0,
    end_percent=1.0,
    condition_image=light_reference_image
)

Lighting Keywords Quick Reference

Lighting Type	English Keywords	Use Case
Side light	`side lighting, dramatic shadows`	Portraits, dramatic effect
Rim light	`rim light, backlight`	Subject outline emphasis
Soft light	`soft light, diffused lighting, beauty dish`	Portraits, products
Volumetric	`volumetric lighting, god rays, light rays`	Atmosphere, cinematic
Golden hour	`golden hour, warm sunset light`	Nature, warm tones
Studio	`studio lighting, three-point lighting`	Professional portraits, products

Solution 4: Automatic Prompt Optimization

What is Automatic Prompt Optimization?

Writing high-quality prompts manually is time-consuming and requires expertise. Automatic prompt optimization analyzes your LoRA training data to generate optimal prompt templates.

Option A: ComfyUI Auto Prompt Node

# Install Auto Prompt node
git clone https://github.com/jonstreeter/comfyui-Lora-Tag-Power-Loader.git

How it works:

Analyze training tags in the LoRA file
Generate base prompt templates
Optimize keywords for Z-Image characteristics
Automatically add lighting, style, and quality enhancers

Option B: LLM-Assisted Prompt Generation

import requests

def optimize_prompt(base_prompt, style="photorealistic"):
    """Optimize Z-Image prompt using LLM"""
    system_prompt = """You are a Z-Image prompt optimization expert.
    Given a base prompt, enhance it with:
    1. Lighting descriptors (golden hour, side lighting, etc.)
    2. Quality boosters (8k, ultra-detailed, photorealistic)
    3. Camera settings (35mm, f/1.8, shallow depth of field)
    4. Style modifiers matching the requested style
    
    Keep the enhanced prompt under 75 words."""
    
    response = requests.post(
        "http://localhost:8000/v1/chat/completions",
        json={
            "model": "qwen-7b",
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Style: {style}\nBase prompt: {base_prompt}"}
            ]
        }
    )
    return response.json()["choices"][0]["message"]["content"]

# Example
base = "a woman reading a book in a library"
optimized = optimize_prompt(base, style="cinematic")
# Output: "A beautiful woman reading a vintage book in a grand library,
#          warm golden hour light streaming through tall arched windows,
#          dust particles floating in volumetric light beams,
#          cinematic 35mm shot, f/1.8, ultra-detailed, photorealistic"

Option C: Z-Image Dedicated Prompt Templates

# prompts.yaml — Z-Image prompt template library
portrait:
  template: |
    {quality}, portrait of {subject}, {lighting},
    {camera}, {background}, {mood}
  quality: "8k resolution, ultra-detailed, photorealistic, masterpiece"
  lighting: "cinematic lighting, dramatic side light, rim light"
  camera: "shot on Canon EOS R5, 85mm f/1.4, shallow depth of field"
  background: "bokeh background, elegant studio backdrop"
  mood: "moody atmosphere, golden hour warmth"

product:
  template: |
    {quality}, product photography of {subject}, {lighting},
    {camera}, {background}
  quality: "commercial photography, 8k, ultra-sharp focus"
  lighting: "studio lighting, softbox, clean white background"
  camera: "Canon EOS R3, 100mm macro, f/8"
  background: "pure white seamless background"

landscape:
  template: |
    {quality}, {scene}, {lighting}, {camera}, {mood}
  quality: "national geographic quality, ultra HD, ultra-detailed"
  scene: "breathtaking landscape with dramatic sky"
  lighting: "golden hour, volumetric fog, dramatic clouds"
  camera: "wide angle shot, tilt-shift effect"
  mood: "epic, majestic, awe-inspiring"

Complete Workflow: Training to Generation

Step 1: Train the LoRA

# Train Z-Image-compatible LoRA with AI Toolkit
python train.py --config config.yaml --base_model Tongyi-MAI/Z-Image-Base

Step 2: Load in ComfyUI

1. Use AI Toolkit LoRA node (NOT standard LoRA node)
2. Set strength_model = 0.6-0.8 (too high damages base model)
3. Set CLIP Set Last Layer = 20 (enhance lighting)
4. Add lighting keywords to prompt

Step 3: Verify and Tune

1. Generate test images (same prompt, different LoRA strengths)
2. Check if lighting effects are natural
3. Adjust strength_model and strength_clip
4. If needed, fine-tune lighting keywords in prompt

FAQ

Q: Why does my LoRA work well in AI Toolkit but poorly in ComfyUI?

A: Most likely, you're using the standard ComfyUI LoRA node instead of the AI Toolkit dedicated node. Switch to AIT_LoRALoader and set CLIP Set Last Layer = 20 — this usually fixes the issue.

Q: Can't I use LoKR format?

A: Z-Image currently doesn't support LoKR format (ComfyUI Issue #10973). When training LoRAs, use standard LoRA or LoHA format.

Q: Recommended training steps for LoRA?

Character/style LoRA: 1,000-3,000 steps
Object/scene LoRA: 500-1,500 steps
With Prodigy optimizer, set learning rate to 1.0

Q: How to fix "lora key not found" error?

Ensure LoRA was trained on Z-Image Base model
Use AI Toolkit LoRA node
Check that LoRA is in standard format (not LoKR)
Update ComfyUI to latest version

Summary

Z-Image LoRA workflows in ComfyUI have specific considerations, but following these principles delivers the best results:

✅ Use AI Toolkit LoRA node (not standard node)
✅ Train with standard LoRA/LoHA format (not LoKR)
✅ Set CLIP Last Layer = 20 for enhanced lighting
✅ Add lighting keywords to prompts
✅ Use auto prompt optimization tools to reduce manual effort

These techniques ensure your Z-Image LoRA training results shine in ComfyUI — with natural lighting, rich details, and consistent quality.

Z-Image ComfyUI LoRA Lighting Fix and Auto Prompt Optimization: Complete Guide

Table of Contents

Z-Image ComfyUI LoRA Lighting Fix and Auto Prompt Optimization: Complete Guide

Have You Experienced These Problems?

The Root Cause: Z-Image's Unique Architecture

Why Does the Default LoRA Node Fail?

Solution 1: Use the Correct LoRA Node

Recommended: AI Toolkit Native LoRA Node

NOT Recommended: Standard ComfyUI LoRA Node

Solution 2: Train with the Correct Format

Training Format Selection

AI Toolkit Training Config (Recommended)

Key Training Parameters for Lighting

Solution 3: Fix Flat Lighting

Problem Symptoms

Fix Methods

Method A: Prompt Lighting Enhancement

Method B: ComfyUI Lighting Enhancement Workflow

Method C: ControlNet Lighting Control

Lighting Keywords Quick Reference

Solution 4: Automatic Prompt Optimization

What is Automatic Prompt Optimization?

Option A: ComfyUI Auto Prompt Node

Option B: LLM-Assisted Prompt Generation

Option C: Z-Image Dedicated Prompt Templates

Complete Workflow: Training to Generation

Step 1: Train the LoRA

Step 2: Load in ComfyUI

Step 3: Verify and Tune

FAQ

Q: Why does my LoRA work well in AI Toolkit but poorly in ComfyUI?

Q: Can't I use LoKR format?

Q: Recommended training steps for LoRA?

Q: How to fix "lora key not found" error?

Summary