diff --git a/invokeai/app/invocations/flux/flux_denoise.py b/invokeai/app/invocations/flux/flux_denoise.py index d65640ad1f3..7fc5dee62ae 100644 --- a/invokeai/app/invocations/flux/flux_denoise.py +++ b/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 diff --git a/tests/app/invocations/test_denoise_noise_inputs.py b/tests/app/invocations/test_denoise_noise_inputs.py index c9d7a17cb4c..26cca592553 100644 --- a/tests/app/invocations/test_denoise_noise_inputs.py +++ b/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)],