Feature: Add Z-Image-Turbo model support (#8671)

Add comprehensive support for Z-Image-Turbo (S3-DiT) models including:

Backend:
- New BaseModelType.ZImage in taxonomy
- Z-Image model config classes (ZImageTransformerConfig,
Qwen3TextEncoderConfig)
- Model loader for Z-Image transformer and Qwen3 text encoder
- Z-Image conditioning data structures
- Step callback support for Z-Image with FLUX latent RGB factors

Invocations:
- z_image_model_loader: Load Z-Image transformer and Qwen3 encoder
- z_image_text_encoder: Encode prompts using Qwen3 with chat template
- z_image_denoise: Flow matching denoising with time-shifted sigmas
- z_image_image_to_latents: Encode images to 16-channel latents
- z_image_latents_to_image: Decode latents using FLUX VAE

Frontend:
- Z-Image graph builder for text-to-image generation
- Model picker and validation updates for z-image base type
- CFG scale now allows 0 (required for Z-Image-Turbo)
- Clip skip disabled for Z-Image (uses Qwen3, not CLIP)
- Optimal dimension settings for Z-Image (1024x1024)

Technical details:
- Uses Qwen3 text encoder (not CLIP/T5)
- 16 latent channels with FLUX-compatible VAE
- Flow matching scheduler with dynamic time shift
- 8 inference steps recommended for Turbo variant
- bfloat16 inference dtype

## Summary

<!--A description of the changes in this PR. Include the kind of change
(fix, feature, docs, etc), the "why" and the "how". Screenshots or
videos are useful for frontend changes.-->

## Related Issues / Discussions

<!--WHEN APPLICABLE: List any related issues or discussions on github or
discord. If this PR closes an issue, please use the "Closes #1234"
format, so that the issue will be automatically closed when the PR
merges.-->

## QA Instructions

- Install a Z-Image-Turbo model (e.g., from HuggingFace)
- Select the model in the Model Picker
- Generate a text-to-image with:
- CFG Scale: 0
- Steps: 8
- Resolution: 1024x1024
- Verify the generated image is coherent (not noise)

## Merge Plan

Standard merge, no special considerations needed.

## Checklist

- [x] _The PR has a short but descriptive title, suitable for a
changelog_
- [ ] _Tests added / updated (if applicable)_
- [ ] _❗Changes to a redux slice have a corresponding migration_
- [ ] _Documentation added / updated (if applicable)_
- [ ] _Updated `What's New` copy (if doing a release after this PR)_

Author: blessedcoolant

invokeai/app/api/dependencies.py

@@ -49,6 +49,7 @@
     FLUXConditioningInfo,
     SD3ConditioningInfo,
     SDXLConditioningInfo,
+    ZImageConditioningInfo,
 )
 from invokeai.backend.util.logging import InvokeAILogger
 from invokeai.version.invokeai_version import __version__
@@ -129,6 +130,7 @@ def initialize(
                     FLUXConditioningInfo,
                     SD3ConditioningInfo,
                     CogView4ConditioningInfo,
+                    ZImageConditioningInfo,
                 ],
                 ephemeral=True,
             ),

invokeai/app/invocations/fields.py

@@ -154,6 +154,7 @@ class FieldDescriptions:
     clip = "CLIP (tokenizer, text encoder, LoRAs) and skipped layer count"
     t5_encoder = "T5 tokenizer and text encoder"
     glm_encoder = "GLM (THUDM) tokenizer and text encoder"
+    qwen3_encoder = "Qwen3 tokenizer and text encoder"
     clip_embed_model = "CLIP Embed loader"
     clip_g_model = "CLIP-G Embed loader"
     unet = "UNet (scheduler, LoRAs)"
@@ -169,6 +170,7 @@ class FieldDescriptions:
     flux_model = "Flux model (Transformer) to load"
     sd3_model = "SD3 model (MMDiTX) to load"
     cogview4_model = "CogView4 model (Transformer) to load"
+    z_image_model = "Z-Image model (Transformer) to load"
     sdxl_main_model = "SDXL Main model (UNet, VAE, CLIP1, CLIP2) to load"
     sdxl_refiner_model = "SDXL Refiner Main Modde (UNet, VAE, CLIP2) to load"
     onnx_main_model = "ONNX Main model (UNet, VAE, CLIP) to load"
@@ -321,6 +323,12 @@ class CogView4ConditioningField(BaseModel):
     conditioning_name: str = Field(description="The name of conditioning tensor")
 
 
+class ZImageConditioningField(BaseModel):
+    """A Z-Image conditioning tensor primitive value"""
+
+    conditioning_name: str = Field(description="The name of conditioning tensor")
+
+
 class ConditioningField(BaseModel):
     """A conditioning tensor primitive value"""
 

invokeai/app/invocations/latents_to_image.py

@@ -2,12 +2,6 @@
 
 import torch
 from diffusers.image_processor import VaeImageProcessor
-from diffusers.models.attention_processor import (
-    AttnProcessor2_0,
-    LoRAAttnProcessor2_0,
-    LoRAXFormersAttnProcessor,
-    XFormersAttnProcessor,
-)
 from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
 from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny
 
@@ -77,26 +71,9 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             assert isinstance(vae, (AutoencoderKL, AutoencoderTiny))
             latents = latents.to(TorchDevice.choose_torch_device())
             if self.fp32:
+                # FP32 mode: convert everything to float32 for maximum precision
                 vae.to(dtype=torch.float32)
-
-                use_torch_2_0_or_xformers = hasattr(vae.decoder, "mid_block") and isinstance(
-                    vae.decoder.mid_block.attentions[0].processor,
-                    (
-                        AttnProcessor2_0,
-                        XFormersAttnProcessor,
-                        LoRAXFormersAttnProcessor,
-                        LoRAAttnProcessor2_0,
-                    ),
-                )
-                # if xformers or torch_2_0 is used attention block does not need
-                # to be in float32 which can save lots of memory
-                if use_torch_2_0_or_xformers:
-                    vae.post_quant_conv.to(latents.dtype)
-                    vae.decoder.conv_in.to(latents.dtype)
-                    vae.decoder.mid_block.to(latents.dtype)
-                else:
-                    latents = latents.float()
-
+                latents = latents.float()
             else:
                 vae.to(dtype=torch.float16)
                 latents = latents.half()

invokeai/app/invocations/metadata.py

@@ -158,6 +158,10 @@ def invoke(self, context: InvocationContext) -> MetadataOutput:
     "cogview4_img2img",
     "cogview4_inpaint",
     "cogview4_outpaint",
+    "z_image_txt2img",
+    "z_image_img2img",
+    "z_image_inpaint",
+    "z_image_outpaint",
 ]
 
 
@@ -166,7 +170,7 @@ def invoke(self, context: InvocationContext) -> MetadataOutput:
     title="Core Metadata",
     tags=["metadata"],
     category="metadata",
-    version="2.0.0",
+    version="2.1.0",
     classification=Classification.Internal,
 )
 class CoreMetadataInvocation(BaseInvocation):
@@ -217,6 +221,10 @@ class CoreMetadataInvocation(BaseInvocation):
         default=None,
         description="The VAE used for decoding, if the main model's default was not used",
     )
+    qwen3_encoder: Optional[ModelIdentifierField] = InputField(
+        default=None,
+        description="The Qwen3 text encoder model used for Z-Image inference",
+    )
 
     # High resolution fix metadata.
     hrf_enabled: Optional[bool] = InputField(

invokeai/app/invocations/model.py

@@ -72,6 +72,14 @@ class GlmEncoderField(BaseModel):
     text_encoder: ModelIdentifierField = Field(description="Info to load text_encoder submodel")
 
 
+class Qwen3EncoderField(BaseModel):
+    """Field for Qwen3 text encoder used by Z-Image models."""
+
+    tokenizer: ModelIdentifierField = Field(description="Info to load tokenizer submodel")
+    text_encoder: ModelIdentifierField = Field(description="Info to load text_encoder submodel")
+    loras: List[LoRAField] = Field(default_factory=list, description="LoRAs to apply on model loading")
+
+
 class VAEField(BaseModel):
     vae: ModelIdentifierField = Field(description="Info to load vae submodel")
     seamless_axes: List[str] = Field(default_factory=list, description='Axes("x" and "y") to which apply seamless')

invokeai/app/invocations/primitives.py

@@ -27,6 +27,7 @@
     SD3ConditioningField,
     TensorField,
     UIComponent,
+    ZImageConditioningField,
 )
 from invokeai.app.services.images.images_common import ImageDTO
 from invokeai.app.services.shared.invocation_context import InvocationContext
@@ -461,6 +462,17 @@ def build(cls, conditioning_name: str) -> "CogView4ConditioningOutput":
         return cls(conditioning=CogView4ConditioningField(conditioning_name=conditioning_name))
 
 
+@invocation_output("z_image_conditioning_output")
+class ZImageConditioningOutput(BaseInvocationOutput):
+    """Base class for nodes that output a Z-Image text conditioning tensor."""
+
+    conditioning: ZImageConditioningField = OutputField(description=FieldDescriptions.cond)
+
+    @classmethod
+    def build(cls, conditioning_name: str) -> "ZImageConditioningOutput":
+        return cls(conditioning=ZImageConditioningField(conditioning_name=conditioning_name))
+
+
 @invocation_output("conditioning_output")
 class ConditioningOutput(BaseInvocationOutput):
     """Base class for nodes that output a single conditioning tensor"""