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

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)

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

invokeai/app/invocations/z_image/z_image_denoise.py

@@ -710,6 +710,10 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
             user_step = 0
 
             if use_scheduler and scheduler is not None:
+                # 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
@@ -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)

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 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
 
             # Get t_prev for inpainting (next sigma value)

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 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
 
             # Get t_prev for inpainting (next sigma value)

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
 
@@ -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,234 @@ 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))
+
+
+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",
     [

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: