fix(z-image): hold the schedule shift at the end of its fit

[Pfannkuchensack]

Summary

ZImageDenoiseInvocation._calculate_shift builds a straight line through two points — 256 image tokens to mu 0.5, 4096 tokens (a 1024×1024 frame) to mu 1.15 — and evaluated it past the second one. Extrapolated, the shift reaches 25 at 2048px, 810 at 3072px and 103 777 at 4096px, and _get_sigmas flattens with it.

Because denoising_start is clipped here by index (start_idx = int(denoising_start * (total_sigmas - 1))), the step count survives and the entry sigma moves instead. A canvas img2img pass at strength 0.5 entered at:

frame	shift before	entry sigma before	entry sigma after
1024px	3.16	0.760	0.760
1536px	7.5	0.883	0.760
2048px	25.3	0.962	0.760
4096px	103 777	1.0000	0.760

At 4096px the blend s · noise + (1 − s) · init kept nothing of the input, so a pass asked to refine an image generated a new one instead. This clamps the fit at its upper anchor. Below one megapixel nothing changes, the fit stays live there, and an explicit shift on the invocation still overrides it.

This is the same defect and the same fix as #359 for FLUX.1, which cited the identical two-point fit.

Related Issues / Discussions

Follows #359 (the FLUX.1 clamp). The underlying reason this was reachable at all — Z-Image and Anima clip denoising_start by index where FLUX.1, FLUX.2, SD3 and CogView4 use the shared value-based clip_timestep_schedule_fractional — is a separate change and is not in this PR.

QA Instructions

Reproduce on main: a Z-Image img2img pass at 2048×2048 with denoising_start 0.5 returns an image unrelated to its input. With this change it refines the input.

Checks run:

• pytest tests/app/invocations/test_z_image_schedule_shift.py — 14 passed
• pytest tests/app/invocations -n 4 — 2046 passed, 3 skipped, 6 xfailed
• pytest -k "z_image or zimage" — 192 passed
• pytest tests/app tests/backend/flux -n 4 — 6752 passed, 9 failed; all 9 are pre-existing local download/install failures caused by an HF_ENDPOINT proxy in this environment and pass under env -u HF_ENDPOINT (45 passed)
• ruff check and ruff format --check on the touched paths — clean
• make frontend-openapi and make frontend-typegen — regenerated, diff limited to the two intended changes

The new cells drive the real _run_diffusion as far as the img2img blend and read the entry sigma back out of the blended tensor rather than recomputing the schedule, so they fail if the node stops honouring the clamp, the override, or the clipping. Each was run against a deliberately broken copy of the node; seven mutations, all caught, including deletion of the shift override.

Review

The first version of the override cell recomputed the schedule itself and stayed green when the override branch was deleted — it was replaced with the harness described above, which also closes the previously untested index-clipping path. Two further items were found and fixed: the shift field description documented mu units for a field consumed as exp(mu), and a function-local import math shadowed the module import.

Remaining limitation: this corrects where the pass starts, not the fact that Z-Image measures denoising_start on a different scale from FLUX.1. At the same slider value Z-Image still enters at more noise than FLUX.1 does. That is the value-based clipping change, deliberately kept separate because it alters the meaning of every saved Z-Image setting.

Compatibility / Rollout

Invocation version 1.6.0 → 1.6.1; openapi.json and the generated schema.ts regenerated.

Existing workflows keep working and are unaffected at or below one megapixel. Above it, a partial denoise becomes more conservative — which is the point of the fix — so a saved workflow that relied on the runaway shift to produce near-txt2img output at high resolutions will now refine its input instead. shift can be set explicitly to restore the previous value.

Full denoises (denoising_start = 0) still enter at sigma 1.0 at every resolution, but the intermediate sigmas above one megapixel do change, so a txt2img generation there will differ from 1.6.0 at the same seed. That is the intended consequence of holding the fit; it is the same trade #359 made for FLUX.1.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)