feat(z-image): add Z-Image Base (undistilled) model variant support (#8799)

* feat(z-image): add Z-Image Base (undistilled) model variant support

- Add ZImageVariantType enum with 'turbo' and 'zbase' variants
- Auto-detect variant on import via scheduler_config.json shift value (3.0=turbo, 6.0=zbase)
- Add database migration to populate variant field for existing Z-Image models
- Re-add LCM scheduler with variant-aware filtering (LCM hidden for zbase)
- Auto-reset scheduler to Euler when switching to zbase model if LCM selected
- Update frontend to show/hide LCM option based on model variant
- Add toast notification when scheduler is auto-reset

Z-Image Base models are undistilled and require more steps (28-50) with higher
guidance (3.0-5.0), while Z-Image Turbo is distilled for ~8 steps with CFG 1.0.
LCM scheduler only works with distilled (Turbo) models.

* Chore ruff format

* Chore fix windows path

* feat(z-image): filter LoRAs by variant compatibility and warn on mismatch

LoRA picker now hides Z-Image LoRAs with incompatible variants (e.g. ZBase
LoRAs when using Turbo model). LoRAs without a variant are always shown.
Backend loaders warn at runtime if a LoRA variant doesn't match the
transformer variant.

* Chore typegen

---------

Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>

Author: Pfannkuchensack

invokeai/app/invocations/z_image_denoise.py

@@ -107,8 +107,8 @@ class ZImageDenoiseInvocation(BaseInvocation):
     # Scheduler selection for the denoising process
     scheduler: ZIMAGE_SCHEDULER_NAME_VALUES = InputField(
         default="euler",
-        description="Scheduler (sampler) for the denoising process. Euler is the default and recommended for "
-        "Z-Image-Turbo. Heun is 2nd-order (better quality, 2x slower). LCM is optimized for few steps.",
+        description="Scheduler (sampler) for the denoising process. Euler is the default and recommended. "
+        "Heun is 2nd-order (better quality, 2x slower). LCM works with Turbo only (not Base).",
         ui_choice_labels=ZIMAGE_SCHEDULER_LABELS,
     )
 
@@ -387,12 +387,11 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                 num_train_timesteps=1000,
                 shift=1.0,
             )
-            # Set timesteps - LCM should use num_inference_steps (it has its own sigma schedule),
+            # Set timesteps - LCM uses its own sigma schedule (num_inference_steps),
             # while other schedulers can use custom sigmas if supported
             is_lcm = self.scheduler == "lcm"
             set_timesteps_sig = inspect.signature(scheduler.set_timesteps)
             if not is_lcm and "sigmas" in set_timesteps_sig.parameters:
-                # Convert sigmas list to tensor for scheduler
                 scheduler.set_timesteps(sigmas=sigmas, device=device)
             else:
                 # LCM or scheduler doesn't support custom sigmas - use num_inference_steps
@@ -644,10 +643,8 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                                     ),
                                 )
                     else:
-                        # For LCM and other first-order schedulers
+                        # For first-order schedulers (Euler, LCM)
                         user_step += 1
-                        # Only call step_callback if we haven't exceeded total_steps
-                        # (LCM scheduler may have more internal steps than user-facing steps)
                         if user_step <= total_steps:
                             pbar.update(1)
                             step_callback(

invokeai/app/invocations/z_image_lora_loader.py

@@ -66,6 +66,18 @@ def invoke(self, context: InvocationContext) -> ZImageLoRALoaderOutput:
         if self.qwen3_encoder and any(lora.lora.key == lora_key for lora in self.qwen3_encoder.loras):
             raise ValueError(f'LoRA "{lora_key}" already applied to Qwen3 encoder.')
 
+        # Warn on variant mismatch between LoRA and transformer.
+        lora_config = context.models.get_config(lora_key)
+        lora_variant = getattr(lora_config, "variant", None)
+        if lora_variant and self.transformer is not None:
+            transformer_config = context.models.get_config(self.transformer.transformer.key)
+            transformer_variant = getattr(transformer_config, "variant", None)
+            if transformer_variant and lora_variant != transformer_variant:
+                context.logger.warning(
+                    f"LoRA variant mismatch: LoRA '{lora_config.name}' is for {lora_variant.value} "
+                    f"but transformer is {transformer_variant.value}. This may cause unexpected results."
+                )
+
         output = ZImageLoRALoaderOutput()
 
         # Attach LoRA layers to the models.
@@ -142,6 +154,18 @@ def invoke(self, context: InvocationContext) -> ZImageLoRALoaderOutput:
                     "not Z-Image models. Ensure you are using a Z-Image compatible LoRA."
                 )
 
+            # Warn on variant mismatch between LoRA and transformer.
+            lora_config = context.models.get_config(lora.lora.key)
+            lora_variant = getattr(lora_config, "variant", None)
+            if lora_variant and self.transformer is not None:
+                transformer_config = context.models.get_config(self.transformer.transformer.key)
+                transformer_variant = getattr(transformer_config, "variant", None)
+                if transformer_variant and lora_variant != transformer_variant:
+                    context.logger.warning(
+                        f"LoRA variant mismatch: LoRA '{lora_config.name}' is for {lora_variant.value} "
+                        f"but transformer is {transformer_variant.value}. This may cause unexpected results."
+                    )
+
             added_loras.append(lora.lora.key)
 
             if self.transformer is not None and output.transformer is not None:

invokeai/app/services/model_records/model_records_base.py

@@ -27,6 +27,7 @@
     ModelVariantType,
     Qwen3VariantType,
     SchedulerPredictionType,
+    ZImageVariantType,
 )
 
 
@@ -92,9 +93,9 @@ class ModelRecordChanges(BaseModelExcludeNull):
 
     # Checkpoint-specific changes
     # TODO(MM2): Should we expose these? Feels footgun-y...
-    variant: Optional[ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | Qwen3VariantType] = (
-        Field(description="The variant of the model.", default=None)
-    )
+    variant: Optional[
+        ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | ZImageVariantType | Qwen3VariantType
+    ] = Field(description="The variant of the model.", default=None)
     prediction_type: Optional[SchedulerPredictionType] = Field(
         description="The prediction type of the model.", default=None
     )

invokeai/app/services/shared/sqlite/sqlite_util.py

@@ -28,6 +28,7 @@
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_23 import build_migration_23
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_24 import build_migration_24
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_25 import build_migration_25
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_26 import build_migration_26
 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -73,6 +74,7 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto
     migrator.register_migration(build_migration_23(app_config=config, logger=logger))
     migrator.register_migration(build_migration_24(app_config=config, logger=logger))
     migrator.register_migration(build_migration_25(app_config=config, logger=logger))
+    migrator.register_migration(build_migration_26(app_config=config, logger=logger))
     migrator.run_migrations()
 
     return db

invokeai/app/services/shared/sqlite_migrator/migrations/migration_26.py

@@ -0,0 +1,115 @@
+import json
+import sqlite3
+from logging import Logger
+from pathlib import Path
+from typing import Any
+
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, ZImageVariantType
+
+
+class Migration26Callback:
+    def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
+        self._app_config = app_config
+        self._logger = logger
+
+    def _detect_variant_from_scheduler(self, model_path: Path) -> ZImageVariantType:
+        """Detect Z-Image variant from scheduler config for Diffusers models.
+
+        Z-Image variants are distinguished by the scheduler shift value:
+        - Turbo (distilled): shift = 3.0
+        - Base (undistilled): shift = 6.0
+        """
+        scheduler_config_path = model_path / "scheduler" / "scheduler_config.json"
+
+        if not scheduler_config_path.exists():
+            return ZImageVariantType.Turbo
+
+        try:
+            with open(scheduler_config_path, "r", encoding="utf-8") as f:
+                scheduler_config = json.load(f)
+
+            shift = scheduler_config.get("shift", 3.0)
+
+            # ZBase (undistilled) uses shift = 6.0, Turbo uses shift = 3.0
+            if shift >= 5.0:
+                return ZImageVariantType.ZBase
+            else:
+                return ZImageVariantType.Turbo
+        except (json.JSONDecodeError, OSError) as e:
+            self._logger.warning(f"Could not read scheduler config: {e}, defaulting to Turbo")
+            return ZImageVariantType.Turbo
+
+    def __call__(self, cursor: sqlite3.Cursor) -> None:
+        cursor.execute("SELECT id, config FROM models;")
+        rows = cursor.fetchall()
+
+        migrated_turbo = 0
+        migrated_base = 0
+
+        for model_id, config_json in rows:
+            try:
+                config_dict: dict[str, Any] = json.loads(config_json)
+
+                # Only migrate Z-Image main models
+                if config_dict.get("base") != BaseModelType.ZImage.value:
+                    continue
+
+                if config_dict.get("type") != ModelType.Main.value:
+                    continue
+
+                # Skip if variant already set
+                if "variant" in config_dict:
+                    continue
+
+                # Determine variant based on format
+                model_format = config_dict.get("format")
+                model_path = config_dict.get("path")
+
+                if model_format == ModelFormat.Diffusers.value and model_path:
+                    # For Diffusers models, detect from scheduler config
+                    variant = self._detect_variant_from_scheduler(Path(model_path))
+                else:
+                    # For Checkpoint/GGUF, default to Turbo (Base only available as Diffusers)
+                    variant = ZImageVariantType.Turbo
+
+                config_dict["variant"] = variant.value
+
+                cursor.execute(
+                    "UPDATE models SET config = ? WHERE id = ?;",
+                    (json.dumps(config_dict), model_id),
+                )
+
+                if variant == ZImageVariantType.ZBase:
+                    migrated_base += 1
+                else:
+                    migrated_turbo += 1
+
+            except json.JSONDecodeError as e:
+                self._logger.error("Invalid config JSON for model %s: %s", model_id, e)
+                raise
+
+        total = migrated_turbo + migrated_base
+        if total > 0:
+            self._logger.info(
+                f"Migration complete: {total} Z-Image model configs updated "
+                f"({migrated_turbo} Turbo, {migrated_base} Base)"
+            )
+        else:
+            self._logger.info("Migration complete: no Z-Image model configs needed migration")
+
+
+def build_migration_26(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
+    """Builds the migration object for migrating from version 25 to version 26.
+
+    This migration adds the variant field to existing Z-Image main models.
+    Models installed before the variant field was added will default to Turbo
+    (the only variant available before Z-Image Base support was added).
+    """
+
+    return Migration(
+        from_version=25,
+        to_version=26,
+        callback=Migration26Callback(app_config=app_config, logger=logger),
+    )

invokeai/backend/flux/schedulers.py

@@ -40,9 +40,9 @@
     FLUX_SCHEDULER_MAP["lcm"] = FlowMatchLCMScheduler
 
 
-# Z-Image scheduler types (same schedulers as Flux, both use Flow Matching)
-# Note: Z-Image-Turbo is optimized for ~8 steps with Euler, but other schedulers
-# can be used for experimentation.
+# Z-Image scheduler types (Flow Matching schedulers)
+# Note: Z-Image-Turbo is optimized for ~8 steps with Euler, LCM can also work.
+# Z-Image Base (undistilled) should only use Euler or Heun (LCM not supported for undistilled models).
 ZIMAGE_SCHEDULER_NAME_VALUES = Literal["euler", "heun", "lcm"]
 
 # Human-readable labels for the UI
@@ -52,7 +52,7 @@
     "lcm": "LCM",
 }
 
-# Mapping from scheduler names to scheduler classes (same as Flux)
+# Mapping from scheduler names to scheduler classes
 ZIMAGE_SCHEDULER_MAP: dict[str, Type[SchedulerMixin]] = {
     "euler": FlowMatchEulerDiscreteScheduler,
     "heun": FlowMatchHeunDiscreteScheduler,

invokeai/backend/model_manager/configs/base.py

@@ -128,7 +128,7 @@ def get_tag(cls) -> Tag:
                 # defaults. variant does not require a default, but if it has one, we need to add it to the tag. We can
                 # check for the presence of a default by seeing if it's not PydanticUndefined, a sentinel value used by
                 # pydantic to indicate that no default was provided.
-                if field.default is not PydanticUndefined:
+                if field.default is not PydanticUndefined and field.default is not None:
                     # We expect each of these fields has an Enum for its default; we want the value of the enum.
                     tag_strings.append(field.default.value)
         return Tag(".".join(tag_strings))

invokeai/backend/model_manager/configs/lora.py

@@ -27,6 +27,7 @@
     FluxLoRAFormat,
     ModelFormat,
     ModelType,
+    ZImageVariantType,
 )
 from invokeai.backend.model_manager.util.model_util import lora_token_vector_length
 from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import is_state_dict_likely_flux_control
@@ -226,6 +227,7 @@ class LoRA_LyCORIS_ZImage_Config(LoRA_LyCORIS_Config_Base, Config_Base):
     """Model config for Z-Image LoRA models in LyCORIS format."""
 
     base: Literal[BaseModelType.ZImage] = Field(default=BaseModelType.ZImage)
+    variant: ZImageVariantType | None = Field(default=None)
 
     @classmethod
     def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None:
@@ -398,3 +400,4 @@ class LoRA_Diffusers_ZImage_Config(LoRA_Diffusers_Config_Base, Config_Base):
     """Model config for Z-Image LoRA models in Diffusers format."""
 
     base: Literal[BaseModelType.ZImage] = Field(default=BaseModelType.ZImage)
+    variant: ZImageVariantType | None = Field(default=None)

invokeai/backend/model_manager/configs/main.py

@@ -30,6 +30,7 @@
     ModelVariantType,
     SchedulerPredictionType,
     SubModelType,
+    ZImageVariantType,
 )
 from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
 from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
@@ -57,7 +58,7 @@ class MainModelDefaultSettings(BaseModel):
     def from_base(
         cls,
         base: BaseModelType,
-        variant: Flux2VariantType | FluxVariantType | ModelVariantType | None = None,
+        variant: Flux2VariantType | FluxVariantType | ModelVariantType | ZImageVariantType | None = None,
     ) -> Self | None:
         match base:
             case BaseModelType.StableDiffusion1:
@@ -67,7 +68,14 @@ def from_base(
             case BaseModelType.StableDiffusionXL:
                 return cls(width=1024, height=1024)
             case BaseModelType.ZImage:
-                return cls(steps=9, cfg_scale=1.0, width=1024, height=1024)
+                # Different defaults based on variant
+                if variant == ZImageVariantType.ZBase:
+                    # Undistilled base model needs more steps and supports CFG
+                    # Recommended: steps=28-50, cfg_scale=3.0-5.0
+                    return cls(steps=50, cfg_scale=4.0, width=1024, height=1024)
+                else:
+                    # Turbo (distilled) uses fewer steps, no CFG
+                    return cls(steps=9, cfg_scale=1.0, width=1024, height=1024)
             case BaseModelType.Flux2:
                 # Different defaults based on variant
                 if variant == Flux2VariantType.Klein9BBase:
@@ -1077,9 +1085,10 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -
 
 
 class Main_Diffusers_ZImage_Config(Diffusers_Config_Base, Main_Config_Base, Config_Base):
-    """Model config for Z-Image diffusers models (Z-Image-Turbo, Z-Image-Base, Z-Image-Edit)."""
+    """Model config for Z-Image diffusers models (Z-Image-Turbo, Z-Image-Base)."""
 
     base: Literal[BaseModelType.ZImage] = Field(BaseModelType.ZImage)
+    variant: ZImageVariantType = Field()
 
     @classmethod
     def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
@@ -1095,19 +1104,41 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -
             },
         )
 
+        variant = override_fields.get("variant") or cls._get_variant_or_raise(mod)
+
         repo_variant = override_fields.get("repo_variant") or cls._get_repo_variant_or_raise(mod)
 
         return cls(
             **override_fields,
+            variant=variant,
             repo_variant=repo_variant,
         )
 
+    @classmethod
+    def _get_variant_or_raise(cls, mod: ModelOnDisk) -> ZImageVariantType:
+        """Determine Z-Image variant from the scheduler config.
+
+        Z-Image variants are distinguished by the scheduler shift value:
+        - Turbo (distilled): shift = 3.0
+        - Base (undistilled): shift = 6.0
+        """
+        scheduler_config = get_config_dict_or_raise(mod.path / "scheduler" / "scheduler_config.json")
+
+        shift = scheduler_config.get("shift", 3.0)
+
+        # ZBase (undistilled) uses shift = 6.0, Turbo uses shift = 3.0
+        if shift >= 5.0:
+            return ZImageVariantType.ZBase
+        else:
+            return ZImageVariantType.Turbo
+
 
 class Main_Checkpoint_ZImage_Config(Checkpoint_Config_Base, Main_Config_Base, Config_Base):
     """Model config for Z-Image single-file checkpoint models (safetensors, etc)."""
 
     base: Literal[BaseModelType.ZImage] = Field(default=BaseModelType.ZImage)
     format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+    variant: ZImageVariantType = Field()
 
     @classmethod
     def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
@@ -1119,7 +1150,9 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -
 
         cls._validate_does_not_look_like_gguf_quantized(mod)
 
-        return cls(**override_fields)
+        variant = override_fields.get("variant", ZImageVariantType.Turbo)
+
+        return cls(**override_fields, variant=variant)
 
     @classmethod
     def _validate_looks_like_z_image_model(cls, mod: ModelOnDisk) -> None:
@@ -1139,6 +1172,7 @@ class Main_GGUF_ZImage_Config(Checkpoint_Config_Base, Main_Config_Base, Config_B
 
     base: Literal[BaseModelType.ZImage] = Field(default=BaseModelType.ZImage)
     format: Literal[ModelFormat.GGUFQuantized] = Field(default=ModelFormat.GGUFQuantized)
+    variant: ZImageVariantType = Field()
 
     @classmethod
     def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
@@ -1150,7 +1184,9 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -
 
         cls._validate_looks_like_gguf_quantized(mod)
 
-        return cls(**override_fields)
+        variant = override_fields.get("variant", ZImageVariantType.Turbo)
+
+        return cls(**override_fields, variant=variant)
 
     @classmethod
     def _validate_looks_like_z_image_model(cls, mod: ModelOnDisk) -> None: