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

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/services/model_records/model_records_base.py

@@ -27,6 +27,7 @@
     ModelVariantType,
     Qwen3VariantType,
     SchedulerPredictionType,
+    ZImageVariantType,
 )
 
 
@@ -91,9 +92,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/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
@@ -56,7 +57,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:
@@ -66,7 +67,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:
@@ -1076,9 +1084,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:
@@ -1094,19 +1103,42 @@ 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)
+    # Default to Turbo for checkpoint models (Base is currently only available in diffusers format)
+    variant: ZImageVariantType = Field(default=ZImageVariantType.Turbo)
 
     @classmethod
     def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
@@ -1118,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:
@@ -1138,6 +1172,8 @@ 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)
+    # Default to Turbo for GGUF models (Base is currently only available in diffusers format)
+    variant: ZImageVariantType = Field(default=ZImageVariantType.Turbo)
 
     @classmethod
     def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
@@ -1149,7 +1185,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:

invokeai/backend/model_manager/taxonomy.py

@@ -133,6 +133,16 @@ class Flux2VariantType(str, Enum):
     """Flux2 Klein 9B Base variant - undistilled foundation model using Qwen3 8B text encoder."""
 
 
+class ZImageVariantType(str, Enum):
+    """Z-Image model variants."""
+
+    Turbo = "turbo"
+    """Z-Image Turbo - distilled model optimized for 8 steps, no CFG support."""
+
+    ZBase = "zbase"
+    """Z-Image Base - undistilled foundation model with full CFG and negative prompt support."""
+
+
 class Qwen3VariantType(str, Enum):
     """Qwen3 text encoder variants based on model size."""
 
@@ -201,7 +211,9 @@ class FluxLoRAFormat(str, Enum):
     XLabs = "flux.xlabs"
 
 
-AnyVariant: TypeAlias = Union[ModelVariantType, ClipVariantType, FluxVariantType, Flux2VariantType, Qwen3VariantType]
+AnyVariant: TypeAlias = Union[
+    ModelVariantType, ClipVariantType, FluxVariantType, Flux2VariantType, ZImageVariantType, Qwen3VariantType
+]
 variant_type_adapter = TypeAdapter[
-    ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | Qwen3VariantType
-](ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | Qwen3VariantType)
+    ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | ZImageVariantType | Qwen3VariantType
+](ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | ZImageVariantType | Qwen3VariantType)

invokeai/frontend/web/public/locales/en.json

@@ -1501,6 +1501,8 @@
         "baseModelChangedCleared_other": "Updated, cleared or disabled {{count}} incompatible submodels",
         "kleinEncoderCleared": "Qwen3 Encoder Cleared",
         "kleinEncoderClearedDescription": "Please select a compatible Qwen3 encoder for the new Klein model variant",
+        "schedulerReset": "Scheduler Reset",
+        "schedulerResetZImageBase": "LCM scheduler is not compatible with Z-Image Base models. Reset to Euler.",
         "canceled": "Processing Canceled",
         "connected": "Connected to Server",
         "imageCopied": "Image Copied",