From 1c86808392dad9018a83bd8be275585d1ecdeec6 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 28 Jan 2026 21:38:45 +0100 Subject: [PATCH 1/3] 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. --- invokeai/app/invocations/z_image_denoise.py | 11 +- .../model_records/model_records_base.py | 7 +- .../app/services/shared/sqlite/sqlite_util.py | 2 + .../migrations/migration_26.py | 115 ++++++++++++++++++ invokeai/backend/flux/schedulers.py | 8 +- .../backend/model_manager/configs/main.py | 48 +++++++- invokeai/backend/model_manager/taxonomy.py | 16 ++- invokeai/frontend/web/public/locales/en.json | 2 + .../listeners/modelSelected.ts | 19 +++ .../web/src/features/metadata/parsing.tsx | 2 +- .../web/src/features/modelManagerV2/models.ts | 2 + .../src/features/nodes/types/common.test-d.ts | 2 + .../web/src/features/nodes/types/common.ts | 5 +- .../components/Core/ParamZImageScheduler.tsx | 27 +++- .../frontend/web/src/services/api/schema.ts | 25 ++-- 15 files changed, 255 insertions(+), 36 deletions(-) create mode 100644 invokeai/app/services/shared/sqlite_migrator/migrations/migration_26.py diff --git a/invokeai/app/invocations/z_image_denoise.py b/invokeai/app/invocations/z_image_denoise.py index 44fa82dbeb0..6ab43d66574 100644 --- a/invokeai/app/invocations/z_image_denoise.py +++ b/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( diff --git a/invokeai/app/services/model_records/model_records_base.py b/invokeai/app/services/model_records/model_records_base.py index 546790cb3e5..7a31d1a34a7 100644 --- a/invokeai/app/services/model_records/model_records_base.py +++ b/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 ) diff --git a/invokeai/app/services/shared/sqlite/sqlite_util.py b/invokeai/app/services/shared/sqlite/sqlite_util.py index f6cb70b9df0..6621170a976 100644 --- a/invokeai/app/services/shared/sqlite/sqlite_util.py +++ b/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 diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_26.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_26.py new file mode 100644 index 00000000000..d392d284139 --- /dev/null +++ b/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), + ) diff --git a/invokeai/backend/flux/schedulers.py b/invokeai/backend/flux/schedulers.py index 0593d727476..e5a8a7137c2 100644 --- a/invokeai/backend/flux/schedulers.py +++ b/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, diff --git a/invokeai/backend/model_manager/configs/main.py b/invokeai/backend/model_manager/configs/main.py index 9f983795ea4..73970749997 100644 --- a/invokeai/backend/model_manager/configs/main.py +++ b/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: diff --git a/invokeai/backend/model_manager/taxonomy.py b/invokeai/backend/model_manager/taxonomy.py index aa0660152a5..023788958de 100644 --- a/invokeai/backend/model_manager/taxonomy.py +++ b/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,7 @@ 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) diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index f819ae10cea..272e8baf5ce 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -1495,6 +1495,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", diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts index cdafc2e8d9d..e3ca6c09199 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts @@ -7,6 +7,7 @@ import { kleinQwen3EncoderModelSelected, kleinVaeModelSelected, modelChanged, + setZImageScheduler, syncedToOptimalDimension, vaeSelected, zImageQwen3EncoderModelSelected, @@ -275,6 +276,24 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) = } } + // Handle Z-Image scheduler when switching to Z-Image Base (zbase) model + // LCM is not supported for undistilled models, so reset to euler + if (newBase === 'z-image' && state.params.zImageScheduler === 'lcm') { + const modelConfigsResult = selectModelConfigsQuery(state); + if (modelConfigsResult.data) { + const newModelConfig = modelConfigsAdapterSelectors.selectById(modelConfigsResult.data, newModel.key); + if (newModelConfig && 'variant' in newModelConfig && newModelConfig.variant === 'zbase') { + dispatch(setZImageScheduler('euler')); + toast({ + id: 'ZIMAGE_SCHEDULER_RESET', + title: t('toast.schedulerReset'), + description: t('toast.schedulerResetZImageBase'), + status: 'info', + }); + } + } + } + dispatch(modelChanged({ model: newModel, previousModel: state.params.model })); const modelBase = selectBboxModelBase(state); diff --git a/invokeai/frontend/web/src/features/metadata/parsing.tsx b/invokeai/frontend/web/src/features/metadata/parsing.tsx index 90eef870793..dfb754c4721 100644 --- a/invokeai/frontend/web/src/features/metadata/parsing.tsx +++ b/invokeai/frontend/web/src/features/metadata/parsing.tsx @@ -411,7 +411,7 @@ const Scheduler: SingleMetadataHandler = { store.dispatch(setFluxScheduler(value)); } } else if (base === 'z-image') { - // Z-Image only supports euler, heun, lcm + // Z-Image supports euler, heun, lcm (but LCM only works well with Turbo, not Base) if (value === 'euler' || value === 'heun' || value === 'lcm') { store.dispatch(setZImageScheduler(value)); } diff --git a/invokeai/frontend/web/src/features/modelManagerV2/models.ts b/invokeai/frontend/web/src/features/modelManagerV2/models.ts index cd83315d48c..239fdca119d 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/models.ts +++ b/invokeai/frontend/web/src/features/modelManagerV2/models.ts @@ -214,6 +214,8 @@ export const MODEL_VARIANT_TO_LONG_NAME: Record = { klein_4b: 'FLUX.2 Klein 4B', klein_9b: 'FLUX.2 Klein 9B', klein_9b_base: 'FLUX.2 Klein 9B Base', + turbo: 'Z-Image Turbo', + zbase: 'Z-Image Base', large: 'CLIP L', gigantic: 'CLIP G', qwen3_4b: 'Qwen3 4B', diff --git a/invokeai/frontend/web/src/features/nodes/types/common.test-d.ts b/invokeai/frontend/web/src/features/nodes/types/common.test-d.ts index 25ba48c293f..04e0fea2cc9 100644 --- a/invokeai/frontend/web/src/features/nodes/types/common.test-d.ts +++ b/invokeai/frontend/web/src/features/nodes/types/common.test-d.ts @@ -18,6 +18,7 @@ import type { zModelVariantType, zQwen3VariantType, zSubModelType, + zZImageVariantType, } from 'features/nodes/types/common'; import type { Invocation, S } from 'services/api/types'; import type { Equals, Extends } from 'tsafe'; @@ -50,6 +51,7 @@ describe('Common types', () => { test('ModelVariantType', () => assert, S['ModelVariantType']>>()); test('FluxVariantType', () => assert, S['FluxVariantType']>>()); test('Flux2VariantType', () => assert, S['Flux2VariantType']>>()); + test('ZImageVariantType', () => assert, S['ZImageVariantType']>>()); test('Qwen3VariantType', () => assert, S['Qwen3VariantType']>>()); test('ModelFormat', () => assert, S['ModelFormat']>>()); diff --git a/invokeai/frontend/web/src/features/nodes/types/common.ts b/invokeai/frontend/web/src/features/nodes/types/common.ts index af91ebe1e18..4ddaff644c5 100644 --- a/invokeai/frontend/web/src/features/nodes/types/common.ts +++ b/invokeai/frontend/web/src/features/nodes/types/common.ts @@ -68,7 +68,8 @@ export type SchedulerField = z.infer; // Flux-specific scheduler options (Flow Matching schedulers) export const zFluxSchedulerField = z.enum(['euler', 'heun', 'lcm']); -// Z-Image scheduler options (Flow Matching schedulers, same as Flux) +// Z-Image scheduler options (Flow Matching schedulers) +// Note: LCM is only supported for Z-Image Turbo, not for Z-Image Base (undistilled) export const zZImageSchedulerField = z.enum(['euler', 'heun', 'lcm']); // Flux DyPE (Dynamic Position Extrapolation) preset options for high-resolution generation @@ -134,12 +135,14 @@ export const zClipVariantType = z.enum(['large', 'gigantic']); export const zModelVariantType = z.enum(['normal', 'inpaint', 'depth']); export const zFluxVariantType = z.enum(['dev', 'dev_fill', 'schnell']); export const zFlux2VariantType = z.enum(['klein_4b', 'klein_9b', 'klein_9b_base']); +export const zZImageVariantType = z.enum(['turbo', 'zbase']); export const zQwen3VariantType = z.enum(['qwen3_4b', 'qwen3_8b']); export const zAnyModelVariant = z.union([ zModelVariantType, zClipVariantType, zFluxVariantType, zFlux2VariantType, + zZImageVariantType, zQwen3VariantType, ]); export type AnyModelVariant = z.infer; diff --git a/invokeai/frontend/web/src/features/parameters/components/Core/ParamZImageScheduler.tsx b/invokeai/frontend/web/src/features/parameters/components/Core/ParamZImageScheduler.tsx index 4353811a73a..7953b14e61e 100644 --- a/invokeai/frontend/web/src/features/parameters/components/Core/ParamZImageScheduler.tsx +++ b/invokeai/frontend/web/src/features/parameters/components/Core/ParamZImageScheduler.tsx @@ -2,22 +2,39 @@ import type { ComboboxOnChange, ComboboxOption } from '@invoke-ai/ui-library'; import { Combobox, FormControl, FormLabel } from '@invoke-ai/ui-library'; import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; import { InformationalPopover } from 'common/components/InformationalPopover/InformationalPopover'; -import { selectZImageScheduler, setZImageScheduler } from 'features/controlLayers/store/paramsSlice'; +import { + selectMainModelConfig, + selectZImageScheduler, + setZImageScheduler, +} from 'features/controlLayers/store/paramsSlice'; import { isParameterZImageScheduler } from 'features/parameters/types/parameterSchemas'; import { memo, useCallback, useMemo } from 'react'; import { useTranslation } from 'react-i18next'; -// Z-Image scheduler options (Flow Matching schedulers, same as Flux) -const ZIMAGE_SCHEDULER_OPTIONS: ComboboxOption[] = [ +// All Z-Image scheduler options +const ZIMAGE_SCHEDULER_OPTIONS_ALL: ComboboxOption[] = [ { value: 'euler', label: 'Euler' }, { value: 'heun', label: 'Heun (2nd order)' }, { value: 'lcm', label: 'LCM' }, ]; +// Z-Image Base (zbase) scheduler options - LCM not supported for undistilled models +const ZIMAGE_SCHEDULER_OPTIONS_BASE: ComboboxOption[] = [ + { value: 'euler', label: 'Euler' }, + { value: 'heun', label: 'Heun (2nd order)' }, +]; + const ParamZImageScheduler = () => { const dispatch = useAppDispatch(); const { t } = useTranslation(); const zImageScheduler = useAppSelector(selectZImageScheduler); + const modelConfig = useAppSelector(selectMainModelConfig); + + // Determine if the selected model is Z-Image Base (zbase variant) + // LCM is not supported for undistilled models + // Need to check base first to narrow the type, since only z-image models have this variant + const isZImageBase = modelConfig?.base === 'z-image' && modelConfig.variant === 'zbase'; + const options = isZImageBase ? ZIMAGE_SCHEDULER_OPTIONS_BASE : ZIMAGE_SCHEDULER_OPTIONS_ALL; const onChange = useCallback( (v) => { @@ -29,14 +46,14 @@ const ParamZImageScheduler = () => { [dispatch] ); - const value = useMemo(() => ZIMAGE_SCHEDULER_OPTIONS.find((o) => o.value === zImageScheduler), [zImageScheduler]); + const value = useMemo(() => options.find((o) => o.value === zImageScheduler), [options, zImageScheduler]); return ( {t('parameters.scheduler')} - + ); }; diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index a99021bc9d0..bdb69b224e3 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -13570,14 +13570,14 @@ export type components = { * Convert Cache Dir * Format: path * @description Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions). - * @default models/.convert_cache + * @default models\.convert_cache */ convert_cache_dir?: string; /** * Download Cache Dir * Format: path * @description Path to the directory that contains dynamically downloaded models. - * @default models/.download_cache + * @default models\.download_cache */ download_cache_dir?: string; /** @@ -16926,6 +16926,8 @@ export type components = { * @constant */ format: "checkpoint"; + /** @default turbo */ + variant: components["schemas"]["ZImageVariantType"]; }; /** Main_Diffusers_CogView4_Config */ Main_Diffusers_CogView4_Config: { @@ -17568,7 +17570,7 @@ export type components = { }; /** * Main_Diffusers_ZImage_Config - * @description Model config for Z-Image diffusers models (Z-Image-Turbo, Z-Image-Base, Z-Image-Edit). + * @description Model config for Z-Image diffusers models (Z-Image-Turbo, Z-Image-Base). */ Main_Diffusers_ZImage_Config: { /** @@ -17645,6 +17647,7 @@ export type components = { * @constant */ base: "z-image"; + variant: components["schemas"]["ZImageVariantType"]; }; /** * Main_GGUF_FLUX_Config @@ -17896,6 +17899,8 @@ export type components = { * @constant */ format: "gguf_quantized"; + /** @default turbo */ + variant: components["schemas"]["ZImageVariantType"]; }; /** * Combine Masks @@ -20102,7 +20107,7 @@ export type components = { * Variant * @description The variant of the model. */ - variant?: components["schemas"]["ModelVariantType"] | components["schemas"]["ClipVariantType"] | components["schemas"]["FluxVariantType"] | components["schemas"]["Flux2VariantType"] | components["schemas"]["Qwen3VariantType"] | null; + variant?: components["schemas"]["ModelVariantType"] | components["schemas"]["ClipVariantType"] | components["schemas"]["FluxVariantType"] | components["schemas"]["Flux2VariantType"] | components["schemas"]["ZImageVariantType"] | components["schemas"]["Qwen3VariantType"] | null; /** @description The prediction type of the model. */ prediction_type?: components["schemas"]["SchedulerPredictionType"] | null; /** @@ -23943,7 +23948,7 @@ export type components = { path_or_prefix: string; model_type: components["schemas"]["ModelType"]; /** Variant */ - variant?: components["schemas"]["ModelVariantType"] | components["schemas"]["ClipVariantType"] | components["schemas"]["FluxVariantType"] | components["schemas"]["Flux2VariantType"] | components["schemas"]["Qwen3VariantType"] | null; + variant?: components["schemas"]["ModelVariantType"] | components["schemas"]["ClipVariantType"] | components["schemas"]["FluxVariantType"] | components["schemas"]["Flux2VariantType"] | components["schemas"]["ZImageVariantType"] | components["schemas"]["Qwen3VariantType"] | null; }; /** * Subtract Integers @@ -26567,7 +26572,7 @@ export type components = { vae?: components["schemas"]["VAEField"] | null; /** * Scheduler - * @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). * @default euler * @enum {string} */ @@ -26694,7 +26699,7 @@ export type components = { vae?: components["schemas"]["VAEField"] | null; /** * Scheduler - * @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). * @default euler * @enum {string} */ @@ -27121,6 +27126,12 @@ export type components = { */ type: "z_image_text_encoder"; }; + /** + * ZImageVariantType + * @description Z-Image model variants. + * @enum {string} + */ + ZImageVariantType: "turbo" | "zbase"; }; responses: never; parameters: never; From c044e81519782da727d988f34b29827909f296d6 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 28 Jan 2026 21:42:36 +0100 Subject: [PATCH 2/3] Chore ruff format --- invokeai/backend/model_manager/taxonomy.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/invokeai/backend/model_manager/taxonomy.py b/invokeai/backend/model_manager/taxonomy.py index 023788958de..eec3ac14a48 100644 --- a/invokeai/backend/model_manager/taxonomy.py +++ b/invokeai/backend/model_manager/taxonomy.py @@ -211,7 +211,9 @@ class FluxLoRAFormat(str, Enum): XLabs = "flux.xlabs" -AnyVariant: TypeAlias = Union[ModelVariantType, ClipVariantType, FluxVariantType, Flux2VariantType, ZImageVariantType, Qwen3VariantType] +AnyVariant: TypeAlias = Union[ + ModelVariantType, ClipVariantType, FluxVariantType, Flux2VariantType, ZImageVariantType, Qwen3VariantType +] variant_type_adapter = TypeAdapter[ ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | ZImageVariantType | Qwen3VariantType ](ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | ZImageVariantType | Qwen3VariantType) From b05f9b48488b656985b6d3e53cd0ea7ca466fd62 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 28 Jan 2026 22:34:11 +0100 Subject: [PATCH 3/3] Chore fix windows path --- invokeai/frontend/web/src/services/api/schema.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index bdb69b224e3..eaae5326dac 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -13570,14 +13570,14 @@ export type components = { * Convert Cache Dir * Format: path * @description Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions). - * @default models\.convert_cache + * @default models/.convert_cache */ convert_cache_dir?: string; /** * Download Cache Dir * Format: path * @description Path to the directory that contains dynamically downloaded models. - * @default models\.download_cache + * @default models/.download_cache */ download_cache_dir?: string; /**