fix(flux): skip packing noise when FLUX Denoise continues from latents without add_noise

invokeai/app/invocations/flux/flux_denoise.py

@@ -364,7 +364,7 @@ def _run_diffusion(
         # Pack all latent tensors.
         init_latents = pack(init_latents) if init_latents is not None else None
         inpaint_mask = pack(inpaint_mask) if inpaint_mask is not None else None
-        noise = pack(noise)
+        noise = pack(noise) if noise is not None else None
         x = pack(x)
 
         # Now that we have 'packed' the latent tensors, verify that we calculated the image_seq_len, packed_h, and

tests/app/invocations/test_denoise_noise_inputs.py

@@ -1,4 +1,5 @@
 import inspect
+from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
@@ -26,7 +27,7 @@
     set_heun_timesteps_from_sigmas,
 )
 from invokeai.backend.flux2.sampling_utils import get_schedule_flux2
-from invokeai.backend.model_manager.taxonomy import BaseModelType, FluxVariantType, ModelType
+from invokeai.backend.model_manager.taxonomy import BaseModelType, FluxVariantType, ModelFormat, ModelType
 
 
 def test_flux_prepare_noise_uses_external_noise():
@@ -101,6 +102,69 @@ def test_flux_add_noise_false_ignores_connected_noise():
     assert torch.equal(result, init_latents)
 
 
+def test_flux_add_noise_false_denoises_incoming_latents_without_noise():
+    """A second stage that continues from latents with add_noise=False runs its remaining steps."""
+    invocation = FluxDenoiseInvocation.model_construct(
+        latents=MagicMock(latents_name="latents"),
+        add_noise=False,
+        width=64,
+        height=64,
+        num_steps=4,
+        denoising_start=0.5,
+        denoising_end=1.0,
+        positive_text_conditioning=MagicMock(conditioning_name="positive"),
+        transformer=MagicMock(transformer="transformer"),
+        seed=123,
+    )
+    init_latents = (torch.arange(16 * 8 * 8) % 8).float().reshape(1, 16, 8, 8)
+    dummy_conditioning = SimpleNamespace(
+        t5_embeds=torch.zeros(1, 4, 16),
+        clip_embeds=torch.zeros(1, 768),
+        to=lambda **_: dummy_conditioning,
+    )
+    mock_context = MagicMock()
+    mock_context.tensors.load.return_value = init_latents
+    mock_context.conditioning.load.return_value = SimpleNamespace(conditionings=[dummy_conditioning])
+    mock_context.models.get_config.return_value = SimpleNamespace(
+        base=BaseModelType.Flux, type=ModelType.Main, variant=FluxVariantType.Dev, format=ModelFormat.GGUFQuantized
+    )
+    mock_context.models.load.return_value.model_on_device.return_value = nullcontext((None, MagicMock()))
+    denoise_calls: list[dict[str, object]] = []
+
+    def fake_denoise(**kwargs):
+        denoise_calls.append(kwargs)
+        return kwargs["img"]
+
+    with (
+        patch(
+            "invokeai.app.invocations.flux.flux_denoise.TorchDevice.choose_torch_device",
+            return_value=torch.device("cpu"),
+        ),
+        patch("invokeai.app.invocations.flux.flux_denoise.FLUXConditioningInfo", object),
+        patch("invokeai.app.invocations.flux.flux_denoise.Flux", object),
+        patch(
+            "invokeai.app.invocations.flux.flux_denoise.RegionalPromptingExtension.from_text_conditioning",
+            return_value=MagicMock(),
+        ),
+        patch.object(invocation, "_prepare_noise_tensor", side_effect=AssertionError("noise should be ignored")),
+        patch.object(invocation, "_load_redux_conditioning", return_value=[]),
+        patch.object(invocation, "_prep_controlnet_extensions", return_value=[]),
+        patch.object(invocation, "_prep_ip_adapter_extensions", return_value=([], [])),
+        patch("invokeai.app.invocations.flux.flux_denoise.peak_dequant_transient_bytes", return_value=0),
+        patch(
+            "invokeai.app.invocations.flux.flux_denoise.LayerPatcher.apply_smart_model_patches",
+            return_value=nullcontext(),
+        ),
+        patch("invokeai.app.invocations.flux.flux_denoise.denoise", side_effect=fake_denoise),
+    ):
+        result = invocation._run_diffusion(mock_context)
+
+    assert len(denoise_calls) == 1
+    assert len(denoise_calls[0]["timesteps"]) > 1
+    assert denoise_calls[0]["inpaint_extension"] is None
+    assert torch.equal(result, init_latents)
+
+
 @pytest.mark.parametrize(
     ["variant", "expects_shift"],
     [(FluxVariantType.Schnell, False), (FluxVariantType.Dev, True)],