z-image-control-turbo-unified-8-steps-v2

Z-Image Turbo Control Unified V2 (V2.1 8-Steps)

This repository hosts the Z-Image Turbo Control Unified V2 model. This is a specialized architecture that unifies the powerful Z-Image Turbo base transformer with enhanced ControlNet capabilities into a single, cohesive model. This unified pipeline supports multiple generation modes in one place: Text-to-Image, Image-to-Image, ControlNet, and Inpainting.

Unlike traditional pipelines where ControlNet is an external add-on, this model integrates control layers directly into the transformer structure. This enables Unified GGUF Quantization, allowing the entire merged architecture (Base + Control) to be quantized (e.g., Q4_K_M, Q8_0) and run efficiently on consumer hardware with limited VRAM. This version also introduces significant optimizations, architectural improvements, and bug fixes for features like group_offload.

📥 Installation

To set up the environment, simply install the dependencies:

#create virtual env
python -m venv venv

# Activate your venv

#upgrade pip
python.exe -m pip install --upgrade pip

#install requirements
pip install -r requirements.txt

Note: This repository contains a diffusers_local folder with the custom ZImageControlUnifiedPipeline and transformer logic required to run this specific architecture.

🚀 Usage

This repository provides separate, easy-to-use scripts for each generation task.

High-Level Scripts

• infer_t2i.py: For Text-to-Image generation.
• infer_i2i.py: For Image-to-Image generation.
• infer_controlnet.py: For ControlNet-guided generation (Pose, Canny, Depth, etc.).
• infer_inpaint.py: For all inpainting tasks.

Hardware Options

Option 1: Low VRAM (GGUF) - Recommended

Use this version if you have limited VRAM (e.g., 6GB - 8GB). It loads the model from a quantized GGUF file. To use it, set use_gguf = True in the desired inference script and provide the path to the .gguf file.

Key Features:

• Loads the unified transformer from a single 4-bit or 8-bit quantized file.
• Enables aggressive group_offload to fit large models on consumer GPUs.

Option 2: High Precision (Diffusers/BF16)

Use this version if you have ample VRAM (e.g., 24GB+). Set use_gguf = False in the script to load the model using the standard from_pretrained directory structure for full BFloat16 precision.

🎨 Inpainting Guide

The infer_inpaint.py script leverages a powerful, unified inpainting system with multiple modes controlled by the inpaint_mode parameter.

Preparing Your Mask

For best results, especially when removing objects or dealing with complex edges, it's recommended to pre-process your mask. We provide a utility script for this.

prepare_mask.py This script expands the white areas of your mask and applies a feather (blur) to the edges. This helps to completely cover artifacts from the old image and ensures a smooth, seamless blend with the new generated content.

Usage:

python prepare_mask.py <input_mask_path> <output_mask_path> --expand 15 --feather 10

• --expand: Expands the mask to cover "ghosting".
• --feather: Creates a soft gradient for seamless blending.

Inpainting Modes in infer_inpaint.py

You can choose the inpainting method by setting the inpaint_mode variable in the script:

1. inpaint_mode = "default"

   • Uses the standard ControlNet-based inpainting. Good for general-purpose tasks.

2. inpaint_mode = "diff"

   • Uses the "Differential Diffusion" inpainting technique. This method is excellent for preserving the original background texture and lighting perfectly while generating new content in the masked area. It works by composing latents at each step of the diffusion process.

3. inpaint_mode = "diff+inpaint"

   • Combines both methods. It uses the diff mode for background preservation while also feeding the inpainting context to the ControlNet layers. This can be useful for complex scenes where both structural guidance and texture preservation are needed.

🛠️ Model Features & Configuration (V2)

Original Features

• This ControlNet is added on 15 layer blocks and 2 refiner layer blocks.
• The model was trained from scratch for 70,000 steps on a dataset of 1 million high-quality images.
• Multiple Control Conditions Supports Canny, HED, Depth, Pose, and MLSD, which can be used like a standard ControlNet.
• You can adjust controlnet_conditioning_scale for stronger control. For better stability, we highly recommend using a detailed prompt. The optimal range for controlnet_conditioning_scale is from 0.65 to 0.90.
  • Note on Steps: As you increase the control strength, it's recommended to appropriately increase the number of inference steps to achieve better results.

This optmized V2 model introduces several new features and parameters for enhanced control and flexibility:

• Unified Pipeline: A single pipeline now handles Text-to-Image, Image-to-Image, ControlNet, and and multiple Inpainting modes.
• Refiner Scale (controlnet_refiner_conditioning_scale): It provides fine-grained control over the influence of the initial refining layers, allowing for isolated adjustments without the influence of the controlnet's conditioning force.
• Optional Refiner (add_control_noise_refiner=False): You can now disable the control noise refiner layers when loading the model to save memory or for different stylistic results.
• Inpainting Blur (mask_blur_radius): A parameter to soften the edges of the inpainting mask for smoother transitions.
• Backward Compatibility: The model supports running weights from V1.
• Group Offload Fixes: The underlying code includes crucial fixes to ensure diffusers group_offload works correctly with use_stream=True, enabling efficient memory management without errors.

🏞️ V2 Examples: Refiner Scale Test

The new controlnet_refiner_conditioning_scale parameter allows for fine-tuning the control signal. Here is a comparison showing its effect while keeping the main control scale fixed.

Prompt: "Photorealistic portrait of a beautiful young East Asian woman with long, vibrant purple hair and a black bow. She is wearing a flowing white summer dress, standing on a sunny beach with a sparkling ocean and clear blue sky in the background. Bright natural sunlight, sharp focus, ultra-detailed." Control Image: Pose.

---

New Tests with this pipeline

<table border="0" style="width: 100%; text-align: left; margin-top: 20px;"> <tr> <td>Pose + Inpaint Image</td> <td>Inpaint Mask</td> <td>Model Inpaint</td> <td>Diff Inpaint</td> <td>Diff + Model Inpaint</td> </tr> <tr> <td><img src="assets/pose.jpg" width="100%" /><img src="assets/inpaint.jpg" width="100%" /></td> <td><img src="assets/inpaint_mask.jpg" width="100%" /></td> <td><img src="results/new_tests/result_inpaint_default.png" width="100%" /></td> <td><img src="results/new_tests/result_inpaint_diff.png" width="100%" /></td> <td><img src="results/new_tests/result_inpaint_diffinpaint.png" width="100%" /></td> </tr> </table> <table border="0" style="width: 100%; text-align: left; margin-top: 20px;"> <tr> <td>Pose</td> <td>Output</td> </tr> <tr> <td><img src="assets/pose.jpg" width="100%" /></td> <td><img src="results/new_tests/result_control_pose.png" width="100%" /></td> </tr> </table> <table border="0" style="width: 100%; text-align: left; margin-top: 20px;"> <tr> <td>Canny</td> <td>Output</td> </tr> <tr> <td><img src="assets/canny.jpg" width="100%" /></td> <td><img src="results/new_tests/result_control_canny.png" width="100%" /></td> </tr> </table> <table border="0" style="width: 100%; text-align: left; margin-top: 20px;"> <tr> <td>HED</td> <td>Output</td> </tr> <tr> <td><img src="assets/man_hed.png" width="100%" /></td> <td><img src="results/new_tests/result_control_hed.png" width="100%" /></td> </tr> </table> <table border="0" style="width: 100%; text-align: left; margin-top: 20px;"> <tr> <td>Depth</td> <td>Output</td> </tr> <tr> <td><img src="assets/depth_cat.png" width="100%" /></td> <td><img src="results/new_tests/result_control_depth.png" width="100%" /></td> </tr> </table> <table border="0" style="width: 100%; text-align: left; margin-top: 20px;"> <tr> <td>MLSD</td> <td>Output</td> </tr> <tr> <td><img src="assets/room_mlsd.png" width="100%" /></td> <td><img src="results/new_tests/result_control_mlsd.png" width="100%" /></td> </tr> </table>

T2I and I2I Results

<table border="0" style="width: 100%; text-align: left; margin-top: 20px;"> <tr> <td>T2I</td> </tr> <tr> <td><img src="results/new_tests/result_t2i.png" width="100%" /></td> </tr> </table> <table border="0" style="width: 100%; text-align: left; margin-top: 20px;"> <tr> <td>I2I</td> <td>Output</td> </tr> <tr> <td><img src="assets/bottle.jpg" width="100%" /></td> <td><img src="results/new_tests/result_i2i.png" width="100%" /></td> </tr> </table>

Original Scale Test Results

The table below shows the generation results under different combinations of Diffusion steps and Control Scale strength from the original model:

---

📂 Repository Structure

• ./transformer/: Directory for model weights (GGUF or standard).
• infer_controlnet.py: Script for ControlNet inference.
• infer_inpaint.py: Script for inpainting inference.
• infer_t2i.py: Script for Text-to-Image inference.
• infer_i2i.py: Script for Image-to-Image inference.
• prepare_mask.py: Utility script to process masks for inpainting.
• diffusers_local/: Custom pipeline code.
• requirements.txt: Python dependencies.