feat: add configurable shift parameter for Z-Image (#9004)

* feat: add configurable shift parameter for Z-Image sigma schedule

Add a shift (mu) override to the Z-Image denoise invocation and expose
it in the UI. When left blank, shift is auto-calculated from image
dimensions (existing behavior). Users can override to fine-tune the
timestep schedule, with an inline X button to reset back to auto.

* refactor: switch Z-Image sigma schedule from exponential to linear time shift

Use shift directly as a linear multiplier instead of exp(mu), giving
more predictable and uniform control over the timestep schedule.
Auto-calculated values are converted via exp(mu) to preserve identical
default behavior.

* feat: recall Z-Image shift parameter from metadata

Write z_image_shift into graph metadata and add a ZImageShift recall
handler so the shift override can be restored from previously generated
images. Auto-mode (null) is omitted from metadata to avoid persisting a
stale value.

---------

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

Author: Pfannkuchensack

invokeai/app/invocations/z_image_denoise.py

@@ -50,7 +50,7 @@
     title="Denoise - Z-Image",
     tags=["image", "z-image"],
     category="image",
-    version="1.4.0",
+    version="1.5.0",
     classification=Classification.Prototype,
 )
 class ZImageDenoiseInvocation(BaseInvocation):
@@ -104,6 +104,15 @@ class ZImageDenoiseInvocation(BaseInvocation):
         description=FieldDescriptions.vae + " Required for control conditioning.",
         input=Input.Connection,
     )
+    # Shift override for the sigma schedule. If None, shift is auto-calculated from image dimensions.
+    shift: Optional[float] = InputField(
+        default=None,
+        ge=0.0,
+        description="Override the timestep shift (mu) for the sigma schedule. "
+        "Leave blank to auto-calculate based on image dimensions (recommended). "
+        "Lower values (~0.5) produce less noise shifting, higher values (~1.15) produce more.",
+        title="Shift",
+    )
     # Scheduler selection for the denoising process
     scheduler: ZIMAGE_SCHEDULER_NAME_VALUES = InputField(
         default="euler",
@@ -225,34 +234,36 @@ def _calculate_shift(
         """Calculate timestep shift based on image sequence length.
 
         Based on diffusers ZImagePipeline.calculate_shift method.
+        Returns a linear shift value (exp(mu) from the original formula).
         """
+        import math
+
         m = (max_shift - base_shift) / (max_image_seq_len - base_image_seq_len)
         b = base_shift - m * base_image_seq_len
         mu = image_seq_len * m + b
-        return mu
+        # Convert from exponential mu to linear shift value
+        return math.exp(mu)
 
-    def _get_sigmas(self, mu: float, num_steps: int) -> list[float]:
-        """Generate sigma schedule with time shift.
+    def _get_sigmas(self, shift: float, num_steps: int) -> list[float]:
+        """Generate sigma schedule with linear time shift.
 
-        Based on FlowMatchEulerDiscreteScheduler with shift.
+        Uses linear time shift: shift / (shift + (1/t - 1)).
+        The shift value is used directly as a multiplier.
         Generates num_steps + 1 sigma values (including terminal 0.0).
         """
-        import math
 
-        def time_shift(mu: float, sigma: float, t: float) -> float:
-            """Apply time shift to a single timestep value."""
+        def time_shift(shift: float, t: float) -> float:
+            """Apply linear time shift to a single timestep value."""
             if t <= 0:
                 return 0.0
             if t >= 1:
                 return 1.0
-            return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+            return shift / (shift + (1 / t - 1))
 
-        # Generate linearly spaced values from 1 to 0 (excluding endpoints for safety)
-        # then apply time shift
         sigmas = []
         for i in range(num_steps + 1):
             t = 1.0 - i / num_steps  # Goes from 1.0 to 0.0
-            sigma = time_shift(mu, 1.0, t)
+            sigma = time_shift(shift, t)
             sigmas.append(sigma)
 
         return sigmas
@@ -313,11 +324,14 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
             # Concatenate all negative embeddings
             neg_prompt_embeds = torch.cat([tc.prompt_embeds for tc in neg_text_conditionings], dim=0)
 
-        # Calculate shift based on image sequence length
-        mu = self._calculate_shift(img_seq_len)
+        # Calculate shift based on image sequence length, or use override
+        if self.shift is not None:
+            shift = self.shift
+        else:
+            shift = self._calculate_shift(img_seq_len)
 
         # Generate sigma schedule with time shift
-        sigmas = self._get_sigmas(mu, self.steps)
+        sigmas = self._get_sigmas(shift, self.steps)
 
         # Apply denoising_start and denoising_end clipping
         if self.denoising_start > 0 or self.denoising_end < 1:

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

@@ -1001,6 +1001,7 @@
         "seedVarianceEnabled": "Seed Variance Enabled",
         "seedVarianceStrength": "Seed Variance Strength",
         "seedVarianceRandomizePercent": "Seed Variance Randomize %",
+        "zImageShift": "Z-Image Shift",
         "seed": "Seed",
         "steps": "Steps",
         "strength": "Image to image strength",

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

@@ -86,6 +86,9 @@ const slice = createSlice({
     setZImageScheduler: (state, action: PayloadAction<'euler' | 'heun' | 'lcm'>) => {
       state.zImageScheduler = action.payload;
     },
+    setZImageShift: (state, action: PayloadAction<number | null>) => {
+      state.zImageShift = action.payload;
+    },
     setZImageSeedVarianceEnabled: (state, action: PayloadAction<boolean>) => {
       state.zImageSeedVarianceEnabled = action.payload;
     },
@@ -535,6 +538,7 @@ export const {
   setFluxDypeScale,
   setFluxDypeExponent,
   setZImageScheduler,
+  setZImageShift,
   setZImageSeedVarianceEnabled,
   setZImageSeedVarianceStrength,
   setZImageSeedVarianceRandomizePercent,
@@ -696,6 +700,7 @@ export const selectFluxDypePreset = createParamsSelector((params) => params.flux
 export const selectFluxDypeScale = createParamsSelector((params) => params.fluxDypeScale);
 export const selectFluxDypeExponent = createParamsSelector((params) => params.fluxDypeExponent);
 export const selectZImageScheduler = createParamsSelector((params) => params.zImageScheduler);
+export const selectZImageShift = createParamsSelector((params) => params.zImageShift);
 export const selectZImageSeedVarianceEnabled = createParamsSelector((params) => params.zImageSeedVarianceEnabled);
 export const selectZImageSeedVarianceStrength = createParamsSelector((params) => params.zImageSeedVarianceStrength);
 export const selectZImageSeedVarianceRandomizePercent = createParamsSelector(

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

@@ -717,6 +717,7 @@ export const zParamsState = z.object({
   fluxDypeScale: zParameterFluxDypeScale,
   fluxDypeExponent: zParameterFluxDypeExponent,
   zImageScheduler: zParameterZImageScheduler,
+  zImageShift: z.number().min(0).max(3).nullable(),
   upscaleScheduler: zParameterScheduler,
   upscaleCfgScale: zParameterCFGScale,
   seed: zParameterSeed,
@@ -788,6 +789,7 @@ export const getInitialParamsState = (): ParamsState => ({
   fluxDypeScale: 2.0,
   fluxDypeExponent: 2.0,
   zImageScheduler: 'euler',
+  zImageShift: null,
   upscaleScheduler: 'kdpm_2',
   upscaleCfgScale: 2,
   seed: 0,

invokeai/frontend/web/src/features/metadata/parsing.tsx

@@ -43,6 +43,7 @@ import {
   setZImageSeedVarianceEnabled,
   setZImageSeedVarianceRandomizePercent,
   setZImageSeedVarianceStrength,
+  setZImageShift,
   vaeSelected,
   widthChanged,
   zImageQwen3EncoderModelSelected,
@@ -686,6 +687,24 @@ const ZImageSeedVarianceRandomizePercent: SingleMetadataHandler<number> = {
 };
 //#endregion ZImageSeedVarianceRandomizePercent
 
+//#region ZImageShift
+const ZImageShift: SingleMetadataHandler<number> = {
+  [SingleMetadataKey]: true,
+  type: 'ZImageShift',
+  parse: (metadata, _store) => {
+    const raw = getProperty(metadata, 'z_image_shift');
+    const parsed = z.number().min(0).max(3).parse(raw);
+    return Promise.resolve(parsed);
+  },
+  recall: (value, store) => {
+    store.dispatch(setZImageShift(value));
+  },
+  i18nKey: 'metadata.zImageShift',
+  LabelComponent: MetadataLabel,
+  ValueComponent: ({ value }: SingleMetadataValueProps<number>) => <MetadataPrimitiveValue value={value} />,
+};
+//#endregion ZImageShift
+
 //#region RefinerModel
 const RefinerModel: SingleMetadataHandler<ParameterSDXLRefinerModel> = {
   [SingleMetadataKey]: true,
@@ -1314,6 +1333,7 @@ export const ImageMetadataHandlers = {
   ZImageSeedVarianceEnabled,
   ZImageSeedVarianceStrength,
   ZImageSeedVarianceRandomizePercent,
+  ZImageShift,
   LoRAs,
   CanvasLayers,
   RefImages,

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

@@ -8,6 +8,7 @@ import {
   selectZImageSeedVarianceEnabled,
   selectZImageSeedVarianceRandomizePercent,
   selectZImageSeedVarianceStrength,
+  selectZImageShift,
   selectZImageVaeModel,
 } from 'features/controlLayers/store/paramsSlice';
 import { selectCanvasMetadata, selectCanvasSlice } from 'features/controlLayers/store/selectors';
@@ -58,6 +59,9 @@ export const buildZImageGraph = async (arg: GraphBuilderArg): Promise<GraphBuild
   // (1.0 means no CFG effect, matching FLUX convention)
   const { cfgScale: guidance_scale, steps, zImageScheduler } = params;
 
+  // Shift override (null = auto-calculate from image dimensions)
+  const zImageShift = selectZImageShift(state);
+
   // Seed Variance Enhancer settings
   const seedVarianceEnabled = selectZImageSeedVarianceEnabled(state);
   const seedVarianceStrength = selectZImageSeedVarianceStrength(state);
@@ -122,6 +126,7 @@ export const buildZImageGraph = async (arg: GraphBuilderArg): Promise<GraphBuild
     guidance_scale,
     steps,
     scheduler: zImageScheduler,
+    shift: zImageShift ?? undefined,
   });
   const l2i = g.addNode({
     type: 'z_image_l2i',
@@ -216,6 +221,7 @@ export const buildZImageGraph = async (arg: GraphBuilderArg): Promise<GraphBuild
     z_image_seed_variance_enabled: seedVarianceEnabled,
     z_image_seed_variance_strength: seedVarianceStrength,
     z_image_seed_variance_randomize_percent: seedVarianceRandomizePercent,
+    z_image_shift: zImageShift ?? undefined,
   });
   g.addEdgeToMetadata(seed, 'value', 'seed');
   g.addEdgeToMetadata(positivePrompt, 'value', 'positive_prompt');

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

@@ -0,0 +1,73 @@
+import { CompositeNumberInput, CompositeSlider, FormControl, FormLabel, Text } from '@invoke-ai/ui-library';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { selectZImageShift, setZImageShift } from 'features/controlLayers/store/paramsSlice';
+import type React from 'react';
+import { memo, useCallback } from 'react';
+import { PiXBold } from 'react-icons/pi';
+
+const CONSTRAINTS = {
+  initial: 3,
+  sliderMin: 1,
+  sliderMax: 7,
+  numberInputMin: 0,
+  numberInputMax: 10,
+  fineStep: 0.1,
+  coarseStep: 0.5,
+};
+
+const MARKS = [1, 2, 3, 4, 5, 6, 7];
+
+const ParamZImageShift = () => {
+  const shift = useAppSelector(selectZImageShift);
+  const dispatch = useAppDispatch();
+
+  const onChange = useCallback((v: number) => dispatch(setZImageShift(v)), [dispatch]);
+  const onReset = useCallback(
+    (e: React.MouseEvent) => {
+      e.preventDefault();
+      e.stopPropagation();
+      dispatch(setZImageShift(null));
+    },
+    [dispatch]
+  );
+
+  const displayValue = shift ?? CONSTRAINTS.initial;
+
+  return (
+    <FormControl>
+      <FormLabel>
+        Shift{' '}
+        {shift !== null ? (
+          <Text as="span" cursor="pointer" onClick={onReset} display="inline-flex" verticalAlign="middle">
+            <PiXBold />
+          </Text>
+        ) : (
+          <Text as="span" opacity={0.5} fontWeight="normal" fontSize="xs">
+            (auto)
+          </Text>
+        )}
+      </FormLabel>
+      <CompositeSlider
+        value={displayValue}
+        defaultValue={CONSTRAINTS.initial}
+        min={CONSTRAINTS.sliderMin}
+        max={CONSTRAINTS.sliderMax}
+        step={CONSTRAINTS.coarseStep}
+        fineStep={CONSTRAINTS.fineStep}
+        onChange={onChange}
+        marks={MARKS}
+      />
+      <CompositeNumberInput
+        value={displayValue}
+        defaultValue={CONSTRAINTS.initial}
+        min={CONSTRAINTS.numberInputMin}
+        max={CONSTRAINTS.numberInputMax}
+        step={CONSTRAINTS.coarseStep}
+        fineStep={CONSTRAINTS.fineStep}
+        onChange={onChange}
+      />
+    </FormControl>
+  );
+};
+
+export default memo(ParamZImageShift);

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

@@ -25,6 +25,7 @@ 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 ParamZImageShift from 'features/parameters/components/Core/ParamZImageShift';
 import ParamZImageSeedVarianceSettings from 'features/parameters/components/SeedVariance/ParamZImageSeedVarianceSettings';
 import { MainModelPicker } from 'features/settingsAccordions/components/GenerationSettingsAccordion/MainModelPicker';
 import { useExpanderToggle } from 'features/settingsAccordions/hooks/useExpanderToggle';
@@ -92,6 +93,7 @@ export const GenerationSettingsAccordion = memo(() => {
               <ParamSteps />
               {(isFLUX || isFlux2) && modelConfig && !isFluxFillMainModelModelConfig(modelConfig) && <ParamGuidance />}
               {!isFLUX && !isFlux2 && <ParamCFGScale />}
+              {isZImage && <ParamZImageShift />}
               {isFLUX && <ParamFluxDypePreset />}
               {isFLUX && fluxDypePreset === 'manual' && <ParamFluxDypeScale />}
               {isFLUX && fluxDypePreset === 'manual' && <ParamFluxDypeExponent />}

invokeai/frontend/web/src/services/api/schema.ts

@@ -29072,6 +29072,12 @@ export type components = {
              * @default null
              */
             vae?: components["schemas"]["VAEField"] | null;
+            /**
+             * Shift
+             * @description Override the timestep shift (mu) for the sigma schedule. Leave blank to auto-calculate based on image dimensions (recommended). Lower values (~0.5) produce less noise shifting, higher values (~1.15) produce more.
+             * @default null
+             */
+            shift?: number | null;
             /**
              * Scheduler
              * @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).
@@ -29199,6 +29205,12 @@ export type components = {
              * @default null
              */
             vae?: components["schemas"]["VAEField"] | null;
+            /**
+             * Shift
+             * @description Override the timestep shift (mu) for the sigma schedule. Leave blank to auto-calculate based on image dimensions (recommended). Lower values (~0.5) produce less noise shifting, higher values (~1.15) produce more.
+             * @default null
+             */
+            shift?: number | null;
             /**
              * Scheduler
              * @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).