From b6f4a8d16616014d14c4676ace3f593fd27e9cc9 Mon Sep 17 00:00:00 2001 From: GokayAI <60583610+gokay-ai@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:22:11 +0100 Subject: [PATCH 1/2] fix(flux): seed LCM step noise for FLUX.1, FLUX.2 and Z-Image FlowMatchLCMScheduler.step re-noises the sample every step, but these loops called it without a generator, so the noise came from the global torch RNG and the same seed gave a different image on each run. Pass a generator seeded from the node's seed, as the ERNIE-Image path does. Fixes #9791 Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com> --- invokeai/app/invocations/flux/flux_denoise.py | 4 ++ .../app/invocations/flux2/flux2_denoise.py | 4 ++ .../invocations/z_image/z_image_denoise.py | 8 +++- invokeai/backend/flux/denoise.py | 7 +++- invokeai/backend/flux2/denoise.py | 7 +++- .../invocations/test_denoise_noise_inputs.py | 42 ++++++++++++++++++- tests/backend/flux/test_denoise.py | 19 +++++++++ 7 files changed, 85 insertions(+), 6 deletions(-) diff --git a/invokeai/app/invocations/flux/flux_denoise.py b/invokeai/app/invocations/flux/flux_denoise.py index 5542aaba5a4..82106a3a422 100644 --- a/invokeai/app/invocations/flux/flux_denoise.py +++ b/invokeai/app/invocations/flux/flux_denoise.py @@ -543,6 +543,10 @@ def _run_diffusion( img_cond_seq_ids=img_cond_seq_ids, dype_extension=dype_extension, scheduler=scheduler, + # LCM re-noises the sample on every step. Without an explicit generator that noise + # comes from the global RNG and the same seed gives a different image each run. + # XOR keeps the step noise decorrelated from the initial noise (as in ERNIE-Image). + generator=torch.Generator(device="cpu").manual_seed(self.seed ^ 0xFFFFFFFF), ) x = unpack(x.float(), self.height, self.width) diff --git a/invokeai/app/invocations/flux2/flux2_denoise.py b/invokeai/app/invocations/flux2/flux2_denoise.py index 19e4fd48738..928d2e1d117 100644 --- a/invokeai/app/invocations/flux2/flux2_denoise.py +++ b/invokeai/app/invocations/flux2/flux2_denoise.py @@ -643,6 +643,10 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: img_cond_seq=img_cond_seq, img_cond_seq_ids=img_cond_seq_ids, pos_joint_attention_kwargs=pos_joint_attention_kwargs, + # LCM re-noises the sample on every step. Without an explicit generator that noise + # comes from the global RNG and the same seed gives a different image each run. + # XOR keeps the step noise decorrelated from the initial noise (as in ERNIE-Image). + generator=torch.Generator(device="cpu").manual_seed(self.seed ^ 0xFFFFFFFF), ) # Apply BN denormalization if BN stats are available diff --git a/invokeai/app/invocations/z_image/z_image_denoise.py b/invokeai/app/invocations/z_image/z_image_denoise.py index a440c3d6c82..f9a6cea6f77 100644 --- a/invokeai/app/invocations/z_image/z_image_denoise.py +++ b/invokeai/app/invocations/z_image/z_image_denoise.py @@ -713,6 +713,10 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # Use diffusers scheduler for stepping # Use tqdm with total_steps (user-facing steps) not num_scheduler_steps (internal steps) # This ensures progress bar shows 1/8, 2/8, etc. even when scheduler uses more internal steps + # LCM re-noises the sample on every step. Without an explicit generator that noise + # comes from the global RNG and the same seed gives a different image each run. + # XOR keeps the step noise decorrelated from the initial noise (as in ERNIE-Image). + step_generator = torch.Generator(device="cpu").manual_seed(self.seed ^ 0xFFFFFFFF) pbar = tqdm(total=total_steps, desc=f"Denoising{TorchDevice.get_session_device_label()}") for step_index in range(num_scheduler_steps): sched_timestep = scheduler.timesteps[step_index] @@ -787,7 +791,9 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: noise_pred = noise_pred_cond # Use scheduler.step() for the update - step_output = scheduler.step(model_output=noise_pred, timestep=sched_timestep, sample=latents) + step_output = scheduler.step( + model_output=noise_pred, timestep=sched_timestep, sample=latents, generator=step_generator + ) latents = step_output.prev_sample # Get sigma_prev for inpainting (next sigma value) diff --git a/invokeai/backend/flux/denoise.py b/invokeai/backend/flux/denoise.py index 49783bc1aa3..db0f990ab68 100644 --- a/invokeai/backend/flux/denoise.py +++ b/invokeai/backend/flux/denoise.py @@ -44,6 +44,8 @@ def denoise( dype_extension: DyPEExtension | None = None, # Defaults to Diffusers Euler when omitted. scheduler: SchedulerMixin | None = None, + # RNG for stochastic schedulers. FlowMatchLCMScheduler.step re-noises the sample every step. + generator: torch.Generator | None = None, ): if len(timesteps) <= 1: return img @@ -217,8 +219,9 @@ def denoise( ) and user_step < total_steps preview_img = img - t_curr * pred if should_emit_preview else None - # Use scheduler.step() for the update - step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img) + # Use scheduler.step() for the update. `generator` matters for LCM, which re-noises every + # step; Euler and Heun only consult it when `s_churn > 0`, which is 0 here. + step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img, generator=generator) img = step_output.prev_sample # Get t_prev for inpainting (next sigma value) diff --git a/invokeai/backend/flux2/denoise.py b/invokeai/backend/flux2/denoise.py index 4dbb99a5b14..a273d2cd77c 100644 --- a/invokeai/backend/flux2/denoise.py +++ b/invokeai/backend/flux2/denoise.py @@ -44,6 +44,8 @@ def denoise( # Optional joint_attention_kwargs (e.g. {"attention_mask": ...}) applied to positive forward # passes only. Negative forwards always run unmasked. pos_joint_attention_kwargs: dict[str, Any] | None = None, + # RNG for stochastic schedulers. FlowMatchLCMScheduler.step re-noises the sample every step. + generator: torch.Generator | None = None, ) -> torch.Tensor: """Denoise latents using a FLUX.2 Klein transformer model. @@ -185,8 +187,9 @@ def denoise( neg_pred = neg_pred[:, :original_seq_len] pred = neg_pred + step_cfg_scale * (pred - neg_pred) - # Use scheduler.step() for the update - step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img) + # Use scheduler.step() for the update. `generator` matters for LCM, which re-noises every + # step; Euler and Heun only consult it when `s_churn > 0`, which is 0 here. + step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img, generator=generator) img = step_output.prev_sample # Get t_prev for inpainting (next sigma value) diff --git a/tests/app/invocations/test_denoise_noise_inputs.py b/tests/app/invocations/test_denoise_noise_inputs.py index 26cca592553..a292f7242d4 100644 --- a/tests/app/invocations/test_denoise_noise_inputs.py +++ b/tests/app/invocations/test_denoise_noise_inputs.py @@ -766,7 +766,13 @@ def __init__(self) -> None: def set_timesteps(self, num_inference_steps: int, device: str | torch.device, mu: float | None = None) -> None: self.received_mu = mu - def step(self, model_output: torch.Tensor, timestep: torch.Tensor, sample: torch.Tensor): + def step( + self, + model_output: torch.Tensor, + timestep: torch.Tensor, + sample: torch.Tensor, + generator: torch.Generator | None = None, + ): return SimpleNamespace(prev_sample=sample) class DummyModel(torch.nn.Module): @@ -801,6 +807,40 @@ def forward( assert scheduler.received_mu == pytest.approx(0.42) +@pytest.mark.skipif(FLUX_SCHEDULER_MAP.get("lcm") is None, reason="FlowMatchLCMScheduler not available") +def test_flux2_lcm_step_noise_follows_the_generator(): + """LCM re-noises the sample on every step, so FLUX.2 must draw that noise from the seeded + generator instead of the global RNG.""" + from invokeai.backend.flux2.denoise import denoise + + class DummyModel(torch.nn.Module): + def forward(self, hidden_states: torch.Tensor, **kwargs: object): + del kwargs + return (0.1 * hidden_states,) + + img = torch.randn(1, 4, 8, generator=torch.Generator().manual_seed(7)) + + def run(seed: int) -> torch.Tensor: + # Advance the global RNG between runs, like a long-lived server process does. + torch.rand(64) + return denoise( + model=DummyModel(), + img=img.clone(), + img_ids=torch.zeros(1, 4, 4, dtype=torch.long), + txt=torch.zeros(1, 4, 8), + txt_ids=torch.zeros(1, 4, 4, dtype=torch.long), + timesteps=[1.0, 0.75, 0.5, 0.25, 0.0], + step_callback=lambda _: None, + guidance=1.0, + cfg_scale=[1.0] * 4, + scheduler=FLUX_SCHEDULER_MAP["lcm"](num_train_timesteps=1000), + generator=torch.Generator().manual_seed(seed), + ) + + assert torch.equal(run(1234), run(1234)) + assert not torch.equal(run(1), run(2)) + + @pytest.mark.parametrize( "scheduler_name", [ diff --git a/tests/backend/flux/test_denoise.py b/tests/backend/flux/test_denoise.py index c699ee610c8..5cccb5fbd40 100644 --- a/tests/backend/flux/test_denoise.py +++ b/tests/backend/flux/test_denoise.py @@ -544,6 +544,25 @@ def test_lcm_preview_uses_the_model_evaluation_sample_and_sigma(monkeypatch): assert torch.allclose(callbacks[0].latents, expected_preview) +@pytest.mark.skipif(not _HAS_LCM, reason="FlowMatchLCMScheduler not available in this Diffusers version") +def test_lcm_step_noise_follows_the_generator(monkeypatch): + """LCM re-noises the sample on every step. With a seeded generator the result must not depend + on the global RNG, otherwise the same seed gives a different image on each run.""" + monkeypatch.setattr("invokeai.backend.flux.denoise.tqdm", _fake_tqdm) + + def run(seed: int) -> torch.Tensor: + # Advance the global RNG between runs, like a long-lived server process does. + torch.rand(64) + scheduler = FLUX_SCHEDULER_MAP["lcm"](num_train_timesteps=1000) + result, _, _, _ = _run_with_scheduler( + scheduler, [1.0, 0.75, 0.5, 0.25, 0.0], generator=torch.Generator().manual_seed(seed) + ) + return result + + assert torch.equal(run(1234), run(1234)) + assert not torch.equal(run(1), run(2)) + + def test_denoise_rejects_an_unrelated_scheduler_with_custom_sigmas_support(monkeypatch): class _UnrelatedScheduler: def __init__(self) -> None: From eb773dec230b814809f8761bdac581e7ef824567 Mon Sep 17 00:00:00 2001 From: GokayAI <60583610+gokay-ai@users.noreply.github.com> Date: Wed, 14 Oct 2026 07:23:04 +0100 Subject: [PATCH 2/2] fix(flux): cover the LCM step generator at the invocation level Address review on #9837: - Add invocation-level tests that drive the real _run_diffusion for FLUX.1, FLUX.2 and Z-Image and assert the step generator is a CPU generator seeded from seed ^ 0xFFFFFFFF. Dropping generator= from either FLUX node, or the Z-Image step_generator, now fails them. - Move the Z-Image step generator and its comment above the scheduler stepping comments, so the tqdm comment sits next to the tqdm call. - Reword the backend comments: Euler and Heun do not use the step noise here (s_churn is 0), rather than not consulting the generator. Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com> --- .../invocations/z_image/z_image_denoise.py | 6 +- invokeai/backend/flux/denoise.py | 2 +- invokeai/backend/flux2/denoise.py | 2 +- .../invocations/test_denoise_noise_inputs.py | 196 +++++++++++++++++- 4 files changed, 200 insertions(+), 6 deletions(-) diff --git a/invokeai/app/invocations/z_image/z_image_denoise.py b/invokeai/app/invocations/z_image/z_image_denoise.py index f9a6cea6f77..2293901d432 100644 --- a/invokeai/app/invocations/z_image/z_image_denoise.py +++ b/invokeai/app/invocations/z_image/z_image_denoise.py @@ -710,13 +710,13 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: user_step = 0 if use_scheduler and scheduler is not None: - # Use diffusers scheduler for stepping - # Use tqdm with total_steps (user-facing steps) not num_scheduler_steps (internal steps) - # This ensures progress bar shows 1/8, 2/8, etc. even when scheduler uses more internal steps # LCM re-noises the sample on every step. Without an explicit generator that noise # comes from the global RNG and the same seed gives a different image each run. # XOR keeps the step noise decorrelated from the initial noise (as in ERNIE-Image). step_generator = torch.Generator(device="cpu").manual_seed(self.seed ^ 0xFFFFFFFF) + # Use diffusers scheduler for stepping + # Use tqdm with total_steps (user-facing steps) not num_scheduler_steps (internal steps) + # This ensures progress bar shows 1/8, 2/8, etc. even when scheduler uses more internal steps pbar = tqdm(total=total_steps, desc=f"Denoising{TorchDevice.get_session_device_label()}") for step_index in range(num_scheduler_steps): sched_timestep = scheduler.timesteps[step_index] diff --git a/invokeai/backend/flux/denoise.py b/invokeai/backend/flux/denoise.py index db0f990ab68..2245d752d67 100644 --- a/invokeai/backend/flux/denoise.py +++ b/invokeai/backend/flux/denoise.py @@ -220,7 +220,7 @@ def denoise( preview_img = img - t_curr * pred if should_emit_preview else None # Use scheduler.step() for the update. `generator` matters for LCM, which re-noises every - # step; Euler and Heun only consult it when `s_churn > 0`, which is 0 here. + # step. Euler and Heun don't use step noise here: with `s_churn` at 0 it is never applied. step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img, generator=generator) img = step_output.prev_sample diff --git a/invokeai/backend/flux2/denoise.py b/invokeai/backend/flux2/denoise.py index a273d2cd77c..3699db70ef3 100644 --- a/invokeai/backend/flux2/denoise.py +++ b/invokeai/backend/flux2/denoise.py @@ -188,7 +188,7 @@ def denoise( pred = neg_pred + step_cfg_scale * (pred - neg_pred) # Use scheduler.step() for the update. `generator` matters for LCM, which re-noises every - # step; Euler and Heun only consult it when `s_churn > 0`, which is 0 here. + # step. Euler and Heun don't use step noise here: with `s_churn` at 0 it is never applied. step_output = scheduler.step(model_output=pred, timestep=timestep, sample=img, generator=generator) img = step_output.prev_sample diff --git a/tests/app/invocations/test_denoise_noise_inputs.py b/tests/app/invocations/test_denoise_noise_inputs.py index a292f7242d4..500d8b94600 100644 --- a/tests/app/invocations/test_denoise_noise_inputs.py +++ b/tests/app/invocations/test_denoise_noise_inputs.py @@ -1,5 +1,5 @@ import inspect -from contextlib import nullcontext +from contextlib import ExitStack, nullcontext from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -841,6 +841,200 @@ def run(seed: int) -> torch.Tensor: assert not torch.equal(run(1), run(2)) +class _StopAtStep(Exception): + """Raised by a recording stub so `_run_diffusion` stops once the step generator has been captured.""" + + +# Mirrors the nodes: the step noise is seeded from `seed`, XORed to stay apart from the initial noise. +_LCM_STEP_SEEDS = [(0, 0xFFFFFFFF), (1234, 1234 ^ 0xFFFFFFFF), (0xFFFFFFFF, 0)] + + +def _assert_step_generator(generator: object, expected_seed: int) -> None: + assert isinstance(generator, torch.Generator) + assert generator.device == torch.device("cpu") + assert generator.initial_seed() == expected_seed + + +def _run_until_backend_denoise(invocation, module: str, context: MagicMock, *extra_patches) -> dict[str, object]: + """Drive the real `_run_diffusion` up to the backend `denoise()` call and return its kwargs.""" + captured: dict[str, object] = {} + + def recording_denoise(**kwargs: object): + captured.update(kwargs) + raise _StopAtStep + + transformer_info = MagicMock() + transformer_info.model_on_device.return_value = nullcontext((None, transformer_info.model)) + context.models.load.return_value = transformer_info + + with ExitStack() as stack: + stack.enter_context(patch(f"{module}.TorchDevice.choose_torch_device", return_value=torch.device("cpu"))) + stack.enter_context(patch(f"{module}.peak_dequant_transient_bytes", return_value=0)) + stack.enter_context(patch(f"{module}.requires_sidecar_patching", return_value=False)) + stack.enter_context(patch(f"{module}.LayerPatcher.apply_smart_model_patches", return_value=nullcontext())) + stack.enter_context(patch(f"{module}.denoise", side_effect=recording_denoise)) + for extra_patch in extra_patches: + stack.enter_context(extra_patch) + with pytest.raises(_StopAtStep): + invocation._run_diffusion(context) + + return captured + + +@pytest.mark.parametrize(("seed", "expected_seed"), _LCM_STEP_SEEDS) +def test_flux_invocation_passes_a_seeded_step_generator_to_denoise(seed: int, expected_seed: int): + """Without `generator=`, FLUX.1 LCM draws its per-step noise from the global RNG again (#9791).""" + module = "invokeai.app.invocations.flux.flux_denoise" + invocation = FluxDenoiseInvocation.model_construct( + latents=None, + add_noise=True, + width=64, + height=64, + num_steps=2, + denoising_start=0.0, + denoising_end=1.0, + positive_text_conditioning=MagicMock(conditioning_name="positive"), + transformer=MagicMock(transformer="transformer", loras=[]), + scheduler="lcm", + seed=seed, + ) + dummy_conditioning = SimpleNamespace( + t5_embeds=torch.zeros(1, 4, 16), + clip_embeds=torch.zeros(1, 768), + to=lambda **_: dummy_conditioning, + ) + context = MagicMock() + context.conditioning.load.return_value = SimpleNamespace(conditionings=[dummy_conditioning]) + context.models.get_config.return_value = SimpleNamespace( + base=BaseModelType.Flux, type=ModelType.Main, variant=FluxVariantType.Dev, format=ModelFormat.Checkpoint + ) + + captured = _run_until_backend_denoise( + invocation, + module, + context, + patch(f"{module}.FLUXConditioningInfo", object), + patch(f"{module}.RegionalPromptingExtension.from_text_conditioning", return_value=MagicMock()), + patch(f"{module}.get_schedule", return_value=[1.0, 0.5, 0.0]), + patch.object(invocation, "_load_redux_conditioning", return_value=[]), + patch.object(invocation, "_prepare_noise_tensor", return_value=torch.zeros(1, 16, 8, 8)), + # `_run_diffusion` asserts the loaded transformer is a `Flux`. + patch(f"{module}.Flux", object), + ) + + _assert_step_generator(captured["generator"], expected_seed) + + +@pytest.mark.parametrize(("seed", "expected_seed"), _LCM_STEP_SEEDS) +def test_flux2_invocation_passes_a_seeded_step_generator_to_denoise(seed: int, expected_seed: int): + """Without `generator=`, FLUX.2 LCM draws its per-step noise from the global RNG again (#9791).""" + module = "invokeai.app.invocations.flux2.flux2_denoise" + invocation = Flux2DenoiseInvocation.model_construct( + latents=MagicMock(latents_name="latents"), + add_noise=False, + width=64, + height=64, + num_steps=2, + denoising_start=0.0, + denoising_end=1.0, + positive_text_conditioning=MagicMock(conditioning_name="positive", mask=None), + transformer=MagicMock(transformer="transformer", loras=[]), + vae=MagicMock(vae="vae"), + scheduler="lcm", + seed=seed, + ) + context = MagicMock() + context.tensors.load.return_value = torch.zeros(1, 32, 8, 8) + context.conditioning.load.return_value = SimpleNamespace( + conditionings=[ + SimpleNamespace( + t5_embeds=torch.zeros(1, 4, 16), to=lambda **_: SimpleNamespace(t5_embeds=torch.zeros(1, 4, 16)) + ) + ] + ) + context.models.get_config.return_value = SimpleNamespace( + base=BaseModelType.Flux2, type=ModelType.Main, format=ModelFormat.Checkpoint, variant=None + ) + + captured = _run_until_backend_denoise( + invocation, + module, + context, + patch(f"{module}.FLUXConditioningInfo", object), + patch.object(invocation, "_get_bn_stats", return_value=None), + ) + + _assert_step_generator(captured["generator"], expected_seed) + + +@pytest.mark.skipif(ZIMAGE_SCHEDULER_MAP.get("lcm") is None, reason="FlowMatchLCMScheduler not available") +@pytest.mark.parametrize(("seed", "expected_seed"), _LCM_STEP_SEEDS) +def test_z_image_invocation_steps_lcm_with_a_seeded_generator(seed: int, expected_seed: int): + """Z-Image steps the scheduler inline, so check the generator that reaches `scheduler.step()`.""" + module = "invokeai.app.invocations.z_image.z_image_denoise" + lcm_class = ZIMAGE_SCHEDULER_MAP["lcm"] + step_generators: list[object] = [] + + class _RecordingLCM(lcm_class): + def step(self, *args: object, **kwargs: object): + step_generators.append(kwargs.get("generator")) + raise _StopAtStep + + class _ZeroTransformer(torch.nn.Module): + def forward(self, x: list[torch.Tensor], t: torch.Tensor, cap_feats: list[torch.Tensor]): + del t, cap_feats + return ([torch.zeros_like(item) for item in x],) + + transformer = _ZeroTransformer() + transformer_info = MagicMock() + transformer_info.model = transformer + transformer_info.model_on_device.return_value = nullcontext((None, transformer)) + + invocation = ZImageDenoiseInvocation.model_construct( + latents=MagicMock(latents_name="latents"), + noise=None, + denoise_mask=None, + add_noise=False, + width=64, + height=64, + steps=2, + denoising_start=0.0, + denoising_end=1.0, + positive_conditioning=SimpleNamespace(conditioning_name="positive", mask=None), + negative_conditioning=None, + guidance_scale=1.0, + transformer=MagicMock(transformer="transformer", loras=[]), + control=None, + shift=None, + scheduler="lcm", + seed=seed, + ) + regional_extension = SimpleNamespace( + regional_text_conditioning=SimpleNamespace(prompt_embeds=torch.zeros(4, 16)), + regional_attn_mask=None, + ) + context = MagicMock() + context.tensors.load.return_value = torch.zeros(1, 16, 8, 8) + context.models.load.return_value = transformer_info + context.models.get_config.return_value = SimpleNamespace( + base=BaseModelType.ZImage, type=ModelType.Main, format=ModelFormat.Checkpoint + ) + + with ( + patch(f"{module}.TorchDevice.choose_torch_device", return_value=torch.device("cpu")), + patch(f"{module}.TorchDevice.choose_bfloat16_safe_dtype", return_value=torch.bfloat16), + patch(f"{module}.ZImageRegionalPromptingExtension.from_text_conditionings", return_value=regional_extension), + patch.object(invocation, "_load_text_conditioning", return_value=[]), + patch.object(invocation, "_estimate_working_memory", return_value=0), + patch.dict(ZIMAGE_SCHEDULER_MAP, {"lcm": _RecordingLCM}), + pytest.raises(_StopAtStep), + ): + invocation._run_diffusion(context) + + assert len(step_generators) == 1 + _assert_step_generator(step_generators[0], expected_seed) + + @pytest.mark.parametrize( "scheduler_name", [