fix(flux): make LCM reproducible with a fixed seed on FLUX.1, FLUX.2 and Z-Image

[gokay-ai]

Summary

With the LCM scheduler, FLUX.1, FLUX.2 and Z-Image give a different image each run for the same seed. FlowMatchLCMScheduler.step() re-noises the sample every step, and these three loops called it without a generator, so that noise came from the global torch RNG.

The loops now pass a CPU generator seeded from the node's seed, the same way the ERNIE-Image path already does (seed ^ 0xFFFFFFFF, so the step noise isn't correlated with the initial noise). The backend denoise() functions for FLUX.1 and FLUX.2 take an optional generator that defaults to None, so other callers are unchanged. Euler and Heun accept the argument too and only use it when s_churn > 0, which is 0 on these paths, so their output doesn't change.

Related Issues / Discussions

Fixes #9791

QA Instructions

New tests:

• tests/backend/flux/test_denoise.py::test_lcm_step_noise_follows_the_generator (FLUX.1)
• tests/app/invocations/test_denoise_noise_inputs.py::test_flux2_lcm_step_noise_follows_the_generator (FLUX.2)

Each runs the real LCM scheduler twice with the same seeded generator, advancing the global RNG in between, and checks the outputs match. Different seeds must give different outputs, so the test can't pass by ignoring the generator. Without a generator, the same FLUX.2 setup differs between runs by up to ~1.43. The dummy scheduler in test_flux2_lcm_scheduler_setup_passes_mu now accepts generator, like the real Diffusers schedulers do.

• pytest tests/backend/flux tests/backend/flux2 tests/backend/ernie_image tests/app/invocations -k "flux or z_image or denoise": 995 passed, 5 skipped, 5 xfailed.
• ruff check / ruff format --check pass on the changed files.

The Z-Image loop lives inside the invocation, so there's no unit test for it; it gets the same change. I tested on CPU with mocks only and didn't run real LCM generations on a GPU.

Merge Plan

N/A

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)

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Pfannkuchensack]

Verified on GPU (RTX 4090): same seed, three runs each, invocation cache disabled. On main, LCM results differed between runs by a mean of 12.5/255 (FLUX.1 dev), 51.8/255 (FLUX.2 Klein 4B) and 8.7/255 (Z-Image Turbo). With this PR, repeated runs are bit-identical and different seeds still give different images. Euler and Heun outputs are unchanged with the generator passed.

A few changes before merging:

1. Test gap: both new tests call the backend denoise() with an explicit generator. Nothing checks that the invocations actually pass one, so removing generator= from flux_denoise.py / flux2_denoise.py, or the Z-Image step_generator, would keep the suite green. Please add coverage at the invocation level, e.g. assert that denoise receives a generator seeded from seed, like the existing _run_diffusion tests that patch denoise. Ideally cover Z-Image too.
2. Z-Image comment placement: in z_image_denoise.py the new comment and the step_generator line sit between the existing tqdm comment and the pbar = tqdm(...) line it describes. Please move the generator and its comment above the "Use diffusers scheduler for stepping" block.
3. Comment wording in backend/flux/denoise.py and backend/flux2/denoise.py: "Euler and Heun only consult it when s_churn > 0" isn't quite right. FlowMatchHeunDiscreteScheduler.step always draws the noise and only applies it when gamma > 0. The output is unchanged (verified), but the comment should say that Euler and Heun don't use the noise here, rather than that they don't consult the generator.

Otherwise this looks good to me.

[gokay-ai]

Thanks for checking this on GPU. All three are addressed in 175ce72:

1. Added invocation-level tests in tests/app/invocations/test_denoise_noise_inputs.py that drive the real _run_diffusion for FLUX.1, FLUX.2 and Z-Image. The FLUX ones patch denoise and assert it receives a CPU generator seeded with seed ^ 0xFFFFFFFF. Z-Image steps inline, so its test swaps in a recording LCM scheduler and checks the generator that reaches scheduler.step(). Each runs for a few seeds, including 0 and 0xFFFFFFFF. I checked that removing generator= from either FLUX node, or the Z-Image step_generator, makes all of them fail.
2. Moved the Z-Image generator and its comment above the "Use diffusers scheduler for stepping" block, so the tqdm comment is next to the tqdm(...) call again.
3. Reworded the backend comments: Euler and Heun don't use the step noise here, since with s_churn at 0 it is never applied.

[Pfannkuchensack]

Heads-up on merge order: this PR and #9836 both change the same import line in tests/app/invocations/test_denoise_noise_inputs.py (#9836 adds from contextlib import nullcontext, this one adds from contextlib import ExitStack, nullcontext). Whichever merges second will hit a small conflict there. The correct resolution is to keep from contextlib import ExitStack, nullcontext; the rest merges cleanly (checked locally with both branches at their current heads).