feat(z-image): add scheduler selection for Z-Image models

Add support for alternative diffusers Flow Matching schedulers for Z-Image:
- Euler (default) - 1st order, optimized for Z-Image-Turbo (8 steps)
- Heun (2nd order) - Better quality, 2x slower
- LCM - Optimized for few-step generation

Backend:
- Extend schedulers.py with Z-Image scheduler types and mapping
- Add scheduler InputField to z_image_denoise invocation (v1.3.0)
- Refactor denoising loop to support diffusers schedulers

Frontend:
- Add zImageScheduler to Redux state in paramsSlice
- Create ParamZImageScheduler component for Linear UI
- Add scheduler to buildZImageGraph for generation

Author: Pfannkuchensack

invokeai/app/invocations/z_image_denoise.py

@@ -1,7 +1,7 @@
-"""Flux scheduler definitions and mapping.
+"""Flow Matching scheduler definitions and mapping.
 
-This module provides the scheduler types and mapping for Flux models,
-supporting multiple Flow Matching schedulers from the diffusers library.
+This module provides the scheduler types and mapping for Flow Matching models
+(Flux and Z-Image), supporting multiple schedulers from the diffusers library.
 """
 
 from typing import Literal, Type
@@ -38,3 +38,25 @@
 
 if _HAS_LCM:
     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.
+ZIMAGE_SCHEDULER_NAME_VALUES = Literal["euler", "heun", "lcm"]
+
+# Human-readable labels for the UI
+ZIMAGE_SCHEDULER_LABELS: dict[str, str] = {
+    "euler": "Euler",
+    "heun": "Heun (2nd order)",
+    "lcm": "LCM",
+}
+
+# Mapping from scheduler names to scheduler classes (same as Flux)
+ZIMAGE_SCHEDULER_MAP: dict[str, Type[SchedulerMixin]] = {
+    "euler": FlowMatchEulerDiscreteScheduler,
+    "heun": FlowMatchHeunDiscreteScheduler,
+}
+
+if _HAS_LCM:
+    ZIMAGE_SCHEDULER_MAP["lcm"] = FlowMatchLCMScheduler

invokeai/backend/flux/schedulers.py

@@ -72,6 +72,9 @@ const slice = createSlice({
     setFluxScheduler: (state, action: PayloadAction<'euler' | 'heun' | 'lcm'>) => {
       state.fluxScheduler = action.payload;
     },
+    setZImageScheduler: (state, action: PayloadAction<'euler' | 'heun' | 'lcm'>) => {
+      state.zImageScheduler = action.payload;
+    },
     setUpscaleScheduler: (state, action: PayloadAction<ParameterScheduler>) => {
       state.upscaleScheduler = action.payload;
     },
@@ -453,6 +456,7 @@ export const {
   setGuidance,
   setScheduler,
   setFluxScheduler,
+  setZImageScheduler,
   setUpscaleScheduler,
   setUpscaleCfgScale,
   setSeed,
@@ -593,6 +597,7 @@ export const selectModelSupportsOptimizedDenoising = createSelector(
 );
 export const selectScheduler = createParamsSelector((params) => params.scheduler);
 export const selectFluxScheduler = createParamsSelector((params) => params.fluxScheduler);
+export const selectZImageScheduler = createParamsSelector((params) => params.zImageScheduler);
 export const selectSeamlessXAxis = createParamsSelector((params) => params.seamlessXAxis);
 export const selectSeamlessYAxis = createParamsSelector((params) => params.seamlessYAxis);
 export const selectSeed = createParamsSelector((params) => params.seed);

invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts

@@ -24,6 +24,7 @@ import {
   zParameterStrength,
   zParameterT5EncoderModel,
   zParameterVAEModel,
+  zParameterZImageScheduler,
 } from 'features/parameters/types/parameterSchemas';
 import type { JsonObject } from 'type-fest';
 import { z } from 'zod';
@@ -598,6 +599,7 @@ export const zParamsState = z.object({
   iterations: z.number(),
   scheduler: zParameterScheduler,
   fluxScheduler: zParameterFluxScheduler,
+  zImageScheduler: zParameterZImageScheduler,
   upscaleScheduler: zParameterScheduler,
   upscaleCfgScale: zParameterCFGScale,
   seed: zParameterSeed,
@@ -653,6 +655,7 @@ export const getInitialParamsState = (): ParamsState => ({
   iterations: 1,
   scheduler: 'dpmpp_3m_k',
   fluxScheduler: 'euler',
+  zImageScheduler: 'euler',
   upscaleScheduler: 'kdpm_2',
   upscaleCfgScale: 2,
   seed: 0,

invokeai/frontend/web/src/features/controlLayers/store/types.ts

@@ -67,6 +67,9 @@ export type SchedulerField = z.infer<typeof zSchedulerField>;
 
 // 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)
+export const zZImageSchedulerField = z.enum(['euler', 'heun', 'lcm']);
 // #endregion
 
 // #region Model-related schemas

invokeai/frontend/web/src/features/nodes/types/common.ts

@@ -53,7 +53,7 @@ export const buildZImageGraph = async (arg: GraphBuilderArg): Promise<GraphBuild
 
   // Z-Image-Turbo uses guidance_scale (stored as cfgScale), defaults to 1.0 for no CFG
   // (1.0 means no CFG effect, matching FLUX convention)
-  const { cfgScale: guidance_scale, steps } = params;
+  const { cfgScale: guidance_scale, steps, zImageScheduler } = params;
 
   const prompts = selectPresetModifiedPrompts(state);
 
@@ -113,6 +113,7 @@ export const buildZImageGraph = async (arg: GraphBuilderArg): Promise<GraphBuild
     id: getPrefixedId('denoise_latents'),
     guidance_scale,
     steps,
+    scheduler: zImageScheduler,
   });
   const l2i = g.addNode({
     type: 'z_image_l2i',

invokeai/frontend/web/src/features/nodes/util/graph/generation/buildZImageGraph.ts

@@ -0,0 +1,44 @@
+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 { 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[] = [
+  { value: 'euler', label: 'Euler' },
+  { value: 'heun', label: 'Heun (2nd order)' },
+  { value: 'lcm', label: 'LCM' },
+];
+
+const ParamZImageScheduler = () => {
+  const dispatch = useAppDispatch();
+  const { t } = useTranslation();
+  const zImageScheduler = useAppSelector(selectZImageScheduler);
+
+  const onChange = useCallback<ComboboxOnChange>(
+    (v) => {
+      if (!isParameterZImageScheduler(v?.value)) {
+        return;
+      }
+      dispatch(setZImageScheduler(v.value));
+    },
+    [dispatch]
+  );
+
+  const value = useMemo(() => ZIMAGE_SCHEDULER_OPTIONS.find((o) => o.value === zImageScheduler), [zImageScheduler]);
+
+  return (
+    <FormControl>
+      <InformationalPopover feature="paramScheduler">
+        <FormLabel>{t('parameters.scheduler')}</FormLabel>
+      </InformationalPopover>
+      <Combobox value={value} options={ZIMAGE_SCHEDULER_OPTIONS} onChange={onChange} />
+    </FormControl>
+  );
+};
+
+export default memo(ParamZImageScheduler);

invokeai/frontend/web/src/features/parameters/components/Core/ParamZImageScheduler.tsx

@@ -1,7 +1,12 @@
 import { NUMPY_RAND_MAX } from 'app/constants';
 import { roundToMultiple } from 'common/util/roundDownToMultiple';
 import { buildZodTypeGuard } from 'common/util/zodUtils';
-import { zFluxSchedulerField, zModelIdentifierField, zSchedulerField } from 'features/nodes/types/common';
+import {
+  zFluxSchedulerField,
+  zModelIdentifierField,
+  zSchedulerField,
+  zZImageSchedulerField,
+} from 'features/nodes/types/common';
 import { z } from 'zod';
 
 /**
@@ -66,6 +71,11 @@ export const [zParameterFluxScheduler, isParameterFluxScheduler] = buildParamete
 export type ParameterFluxScheduler = z.infer<typeof zParameterFluxScheduler>;
 // #endregion
 
+// #region Z-Image Scheduler
+export const [zParameterZImageScheduler, isParameterZImageScheduler] = buildParameter(zZImageSchedulerField);
+export type ParameterZImageScheduler = z.infer<typeof zParameterZImageScheduler>;
+// #endregion
+
 // #region seed
 export const [zParameterSeed, isParameterSeed] = buildParameter(z.number().int().min(0).max(NUMPY_RAND_MAX));
 export type ParameterSeed = z.infer<typeof zParameterSeed>;

invokeai/frontend/web/src/features/parameters/types/parameterSchemas.ts

@@ -12,6 +12,7 @@ import ParamFluxScheduler from 'features/parameters/components/Core/ParamFluxSch
 import ParamGuidance from 'features/parameters/components/Core/ParamGuidance';
 import ParamScheduler from 'features/parameters/components/Core/ParamScheduler';
 import ParamSteps from 'features/parameters/components/Core/ParamSteps';
+import ParamZImageScheduler from 'features/parameters/components/Core/ParamZImageScheduler';
 import { MainModelPicker } from 'features/settingsAccordions/components/GenerationSettingsAccordion/MainModelPicker';
 import { useExpanderToggle } from 'features/settingsAccordions/hooks/useExpanderToggle';
 import { useStandaloneAccordionToggle } from 'features/settingsAccordions/hooks/useStandaloneAccordionToggle';
@@ -70,6 +71,7 @@ export const GenerationSettingsAccordion = memo(() => {
             <FormControlGroup formLabelProps={formLabelProps}>
               {!isFLUX && !isSD3 && !isCogView4 && !isZImage && <ParamScheduler />}
               {isFLUX && <ParamFluxScheduler />}
+              {isZImage && <ParamZImageScheduler />}
               <ParamSteps />
               {isFLUX && modelConfig && !isFluxFillMainModelModelConfig(modelConfig) && <ParamGuidance />}
               {!isFLUX && <ParamCFGScale />}

invokeai/frontend/web/src/features/settingsAccordions/components/GenerationSettingsAccordion/GenerationSettingsAccordion.tsx

@@ -25555,6 +25555,13 @@ export type components = {
              * @default null
              */
             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.
+             * @default euler
+             * @enum {string}
+             */
+            scheduler?: "euler" | "heun" | "lcm";
             /**
              * type
              * @default z_image_denoise