fix(flux2): img2img posterization - noise scaling and sigma schedule

invokeai/app/invocations/flux2_denoise.py

@@ -36,6 +36,7 @@
     get_noise_flux2,
     get_schedule_flux2,
     pack_flux2,
+    time_shift_flux2,
     unpack_flux2,
 )
 from invokeai.backend.flux2.text_conditioning import Flux2TextConditioning
@@ -234,6 +235,46 @@ def _bn_denormalize(
         bn_std = bn_std.to(x.device, x.dtype)
         return x * bn_std + bn_mean
 
+    def _prepare_normalized_start_latents(
+        self,
+        init_latents_packed: torch.Tensor,
+        noise_packed: Optional[torch.Tensor],
+        t_0: float,
+        bn_mean: Optional[torch.Tensor],
+        bn_std: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Build the img2img/inpainting start latents in the transformer's BN-normalized space.
+
+        Only the init latents are normalized. The noise must NOT be: it is already N(0, 1),
+        exactly the distribution the transformer expects in normalized space. The rectified-flow
+        preblend is therefore computed *from* the normalized operands rather than by normalizing
+        the raw mixture -- normalizing the mixture would divide its noise term by bn_std (~1.77 for
+        the FLUX.2 VAE) as well, so the sample would carry only ~57% of the noise implied by the
+        timestep it is handed to the transformer with. The model then over-denoises and flattens
+        fine detail into posterized patches, progressively worse at higher denoise strengths
+        (see #8964).
+
+        Args:
+            init_latents_packed: Packed, un-normalized init latents of shape (B, seq, 128).
+            noise_packed: Packed N(0, 1) noise of shape (B, seq, 128). Required if add_noise.
+            t_0: First timestep of the clipped schedule.
+            bn_mean: BN running mean of shape (128,), or None if the VAE exposes no BN stats.
+            bn_std: BN running std of shape (128,), or None if the VAE exposes no BN stats.
+
+        Returns:
+            Tuple of (start latents, normalized init latents), both of shape (B, seq, 128).
+        """
+        if bn_mean is not None and bn_std is not None:
+            init_latents_packed = self._bn_normalize(init_latents_packed, bn_mean, bn_std)
+
+        if self.add_noise:
+            assert noise_packed is not None
+            x = t_0 * noise_packed + (1.0 - t_0) * init_latents_packed
+        else:
+            x = init_latents_packed
+
+        return x, init_latents_packed
+
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> LatentsOutput:
         latents = self._run_diffusion(context)
@@ -324,6 +365,14 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
         # Compute mu for dynamic schedule shifting (used by FlowMatchEulerDiscreteScheduler)
         mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=self.num_steps)
 
+        # img2img and inpainting step this schedule manually (see the scheduler setup below), while
+        # txt2img hands it to the scheduler, which applies the exponential shift from mu itself. Apply
+        # the same shift here so both paths follow one schedule -- and apply it before clipping, so
+        # denoising_start/end select a fraction of the schedule the model is actually run on.
+        uses_manual_euler = self.denoise_mask is not None or self.denoising_start > 1e-5
+        if uses_manual_euler:
+            timesteps = time_shift_flux2(timesteps, mu)
+
         # Clip the timesteps schedule based on denoising_start and denoising_end
         timesteps = clip_timestep_schedule_fractional(timesteps, self.denoising_start, self.denoising_end)
 
@@ -366,30 +415,19 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
         noise_packed = pack_flux2(noise) if noise is not None else None
         x = pack_flux2(x)
 
-        # BN normalization for img2img/inpainting:
-        # - The init_latents from VAE encode are NOT BN-normalized
-        # - The transformer operates in BN-normalized space
-        # - We must normalize x, init_latents, AND noise for InpaintExtension
-        # - Output MUST be denormalized after denoising before VAE decode
-        #
-        # This ensures that:
-        # 1. x starts in the correct normalized space for the transformer
-        # 2. When InpaintExtension merges intermediate_latents with noised_init_latents,
-        #    both are in the same scale/space (noise and init_latents must be in same space
-        #    for the linear interpolation: noised = noise * t + init * (1-t))
-        if bn_mean is not None and bn_std is not None:
-            if init_latents_packed is not None:
-                init_latents_packed = self._bn_normalize(init_latents_packed, bn_mean, bn_std)
-                # Also normalize noise for InpaintExtension - it's used to compute
-                # noised_init_latents = noise * t + init_latents * (1-t)
-                # Both operands must be in the same normalized space
-                if noise_packed is not None:
-                    noise_packed = self._bn_normalize(noise_packed, bn_mean, bn_std)
-            # For img2img/inpainting, x is computed from init_latents and must also be normalized
-            # For txt2img, x is pure noise (already N(0,1)) - normalizing it would be incorrect
-            # We detect img2img by checking if init_latents was provided
-            if init_latents is not None:
-                x = self._bn_normalize(x, bn_mean, bn_std)
+        # The init latents from VAE encode are NOT BN-normalized, but the transformer operates in
+        # BN-normalized space, so the img2img/inpainting start latents are rebuilt there. The noise
+        # is left untouched -- see _prepare_normalized_start_latents. pack_flux2 is a pure rearrange,
+        # so blending before or after packing is equivalent. Output is denormalized again below,
+        # after denoising, before it goes to the VAE.
+        if init_latents_packed is not None:
+            x, init_latents_packed = self._prepare_normalized_start_latents(
+                init_latents_packed=init_latents_packed,
+                noise_packed=noise_packed,
+                t_0=timesteps[0],
+                bn_mean=bn_mean,
+                bn_std=bn_std,
+            )
 
         # Verify packed dimensions
         assert packed_h * packed_w == x.shape[1]
@@ -409,9 +447,6 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
         num_steps = len(timesteps) - 1
         cfg_scale_list = [self.cfg_scale] * num_steps
 
-        # Check if we're doing inpainting (have a mask or a clipped schedule)
-        is_inpainting = self.denoise_mask is not None or self.denoising_start > 1e-5
-
         # Create scheduler with FLUX.2 Klein configuration
         # For inpainting/img2img, use manual Euler stepping to preserve the exact
         # clipped timestep schedule used for the initial latent/noise preblend.
@@ -421,7 +456,7 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
         # change the first effective timestep/sigma and break parity with the
         # preblend computed above.
         scheduler = None
-        if self.scheduler in FLUX_SCHEDULER_MAP and not is_inpainting:
+        if self.scheduler in FLUX_SCHEDULER_MAP and not uses_manual_euler:
             # Only use scheduler for txt2img - use manual Euler for inpainting to preserve exact timesteps
             scheduler_class = FLUX_SCHEDULER_MAP[self.scheduler]
             # FlowMatchHeunDiscreteScheduler only supports num_train_timesteps and shift parameters

invokeai/backend/flux2/sampling_utils.py

@@ -162,6 +162,36 @@ def get_schedule_flux2(
     return sigmas_list
 
 
+def time_shift_flux2(sigmas: list[float], mu: float) -> list[float]:
+    """Apply the exponential schedule shift that the txt2img scheduler applies internally.
+
+    ``get_schedule_flux2()`` returns an unshifted linear schedule because the
+    FlowMatchEulerDiscreteScheduler used for txt2img shifts it itself from ``mu``. Code paths that
+    step the schedule manually (img2img, inpainting) must apply the shift themselves, otherwise they
+    run the model on a completely different sigma trajectory than txt2img does: with 9 steps the
+    shifted schedule bottoms out at sigma 0.485, the linear one at 0.111. A distilled model such as
+    FLUX.2 Klein is never trained at the low end of the linear schedule and leaves a grainy residue
+    there.
+
+    Mirrors diffusers' ``FlowMatchEulerDiscreteScheduler._time_shift_exponential`` with an exponent
+    of 1.0, which is how the scheduler invokes it::
+
+        sigma' = exp(mu) / (exp(mu) + (1 / sigma - 1))
+
+    Args:
+        sigmas: Unshifted sigma schedule, descending from 1.0 to 0.0.
+        mu: Shift parameter, as returned by ``compute_empirical_mu()``.
+
+    Returns:
+        The shifted schedule. The 1.0 and 0.0 endpoints are fixed points of the transform and are
+        passed through directly to avoid a division by zero.
+    """
+    exp_mu = math.exp(mu)
+    return [
+        1.0 if sigma >= 1.0 else 0.0 if sigma <= 0.0 else exp_mu / (exp_mu + (1.0 / sigma - 1.0)) for sigma in sigmas
+    ]
+
+
 def generate_img_ids_flux2(h: int, w: int, batch_size: int, device: torch.device) -> torch.Tensor:
     """Generate tensor of image position ids for FLUX.2 with RoPE scaling.
 

tests/app/invocations/test_flux2_denoise_img2img_normalization.py

@@ -0,0 +1,124 @@
+"""Regression tests for the FLUX.2 img2img/inpainting latent normalization (see #8964).
+
+The FLUX.2 transformer operates on BN-normalized latents, while the VAE encode node emits raw
+latents. The img2img preblend used to be computed in raw space and normalized afterwards, which
+divided the noise term by bn_std (~1.77 for the FLUX.2 VAE) as well. The start latents then carried
+only ~57% of the noise implied by their timestep, the model over-denoised, and fine detail collapsed
+into posterized patches -- progressively worse the higher the denoise strength.
+"""
+
+import pytest
+import torch
+
+from invokeai.app.invocations.flux2_denoise import Flux2DenoiseInvocation
+
+# Measured on the BFL FLUX.2 VAE (bn.running_mean / bn.running_var): mean ~0, var ~3.13.
+BN_STD_VALUE = 1.7676
+PACKED_CHANNELS = 128
+
+
+def _bn_stats() -> tuple[torch.Tensor, torch.Tensor]:
+    bn_mean = torch.full((PACKED_CHANNELS,), 0.05)
+    bn_std = torch.full((PACKED_CHANNELS,), BN_STD_VALUE)
+    return bn_mean, bn_std
+
+
+def _packed(seed: int) -> torch.Tensor:
+    generator = torch.Generator().manual_seed(seed)
+    return torch.randn(1, 64, PACKED_CHANNELS, generator=generator)
+
+
+@pytest.mark.parametrize("t_0", [0.2, 0.5, 0.85, 1.0])
+def test_noise_keeps_unit_scale_in_normalized_start_latents(t_0: float) -> None:
+    """The noise term must survive the preblend at full N(0, 1) scale, for every denoise strength."""
+    bn_mean, bn_std = _bn_stats()
+    init_latents = _packed(1)
+    noise = _packed(2)
+    invocation = Flux2DenoiseInvocation.model_construct(add_noise=True)
+
+    x, normalized_init = invocation._prepare_normalized_start_latents(
+        init_latents_packed=init_latents,
+        noise_packed=noise,
+        t_0=t_0,
+        bn_mean=bn_mean,
+        bn_std=bn_std,
+    )
+
+    # x == t_0 * noise + (1 - t_0) * normalize(init), so peeling off the init term must recover the
+    # noise unscaled -- not noise / bn_std, which is what normalizing the raw mixture produced.
+    recovered_noise = (x - (1.0 - t_0) * normalized_init) / t_0
+    assert torch.allclose(recovered_noise, noise, atol=1e-5)
+
+
+def test_start_latents_differ_from_normalizing_the_raw_mixture() -> None:
+    """Guard against a regression back to normalizing the blended tensor as a whole."""
+    bn_mean, bn_std = _bn_stats()
+    init_latents = _packed(3)
+    noise = _packed(4)
+    t_0 = 0.6
+    invocation = Flux2DenoiseInvocation.model_construct(add_noise=True)
+
+    x, _ = invocation._prepare_normalized_start_latents(
+        init_latents_packed=init_latents,
+        noise_packed=noise,
+        t_0=t_0,
+        bn_mean=bn_mean,
+        bn_std=bn_std,
+    )
+
+    buggy = invocation._bn_normalize(t_0 * noise + (1.0 - t_0) * init_latents, bn_mean, bn_std)
+    # The buggy variant attenuates the noise term by 1 / bn_std.
+    assert not torch.allclose(x, buggy, atol=1e-3)
+    assert torch.allclose(x - buggy, t_0 * (noise - (noise - bn_mean) / bn_std), atol=1e-5)
+
+
+def test_init_latents_are_normalized() -> None:
+    bn_mean, bn_std = _bn_stats()
+    init_latents = _packed(5)
+    invocation = Flux2DenoiseInvocation.model_construct(add_noise=True)
+
+    _, normalized_init = invocation._prepare_normalized_start_latents(
+        init_latents_packed=init_latents,
+        noise_packed=_packed(6),
+        t_0=0.4,
+        bn_mean=bn_mean,
+        bn_std=bn_std,
+    )
+
+    assert torch.allclose(normalized_init, (init_latents - bn_mean) / bn_std, atol=1e-6)
+
+
+def test_without_add_noise_start_latents_are_the_normalized_init_latents() -> None:
+    bn_mean, bn_std = _bn_stats()
+    init_latents = _packed(7)
+    invocation = Flux2DenoiseInvocation.model_construct(add_noise=False)
+
+    x, normalized_init = invocation._prepare_normalized_start_latents(
+        init_latents_packed=init_latents,
+        noise_packed=_packed(8),
+        t_0=0.4,
+        bn_mean=bn_mean,
+        bn_std=bn_std,
+    )
+
+    assert torch.allclose(x, normalized_init, atol=1e-6)
+    assert torch.allclose(x, (init_latents - bn_mean) / bn_std, atol=1e-6)
+
+
+def test_without_bn_stats_the_raw_preblend_is_preserved() -> None:
+    """VAE formats that expose no BN stats keep the previous raw-space behaviour."""
+    init_latents = _packed(9)
+    noise = _packed(10)
+    t_0 = 0.3
+    invocation = Flux2DenoiseInvocation.model_construct(add_noise=True)
+
+    x, normalized_init = invocation._prepare_normalized_start_latents(
+        init_latents_packed=init_latents,
+        noise_packed=noise,
+        t_0=t_0,
+        bn_mean=None,
+        bn_std=None,
+    )
+
+    assert torch.allclose(normalized_init, init_latents, atol=1e-6)
+    assert torch.allclose(x, t_0 * noise + (1.0 - t_0) * init_latents, atol=1e-6)

tests/backend/flux2/test_flux2_schedule_shift.py

@@ -0,0 +1,70 @@
+"""Regression tests for the FLUX.2 schedule shift shared by the txt2img and img2img paths.
+
+``get_schedule_flux2()`` returns an unshifted linear schedule because the txt2img scheduler applies
+the exponential shift from ``mu`` itself. img2img and inpainting step the schedule manually and must
+apply the same shift, otherwise they run the model on a sigma trajectory txt2img never visits -- with
+9 steps the shifted schedule bottoms out at 0.485 while the linear one runs down to 0.111, and a
+distilled model like FLUX.2 Klein leaves a grainy residue down there.
+"""
+
+import numpy as np
+import pytest
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+from invokeai.backend.flux2.sampling_utils import compute_empirical_mu, get_schedule_flux2, time_shift_flux2
+
+
+def _txt2img_scheduler() -> FlowMatchEulerDiscreteScheduler:
+    """The scheduler exactly as Flux2DenoiseInvocation builds it for txt2img."""
+    return FlowMatchEulerDiscreteScheduler(
+        num_train_timesteps=1000,
+        shift=3.0,
+        use_dynamic_shifting=True,
+        base_shift=0.5,
+        max_shift=1.15,
+        base_image_seq_len=256,
+        max_image_seq_len=4096,
+        time_shift_type="exponential",
+    )
+
+
+@pytest.mark.parametrize("num_steps", [4, 9, 20, 30])
+def test_shift_matches_the_txt2img_scheduler(num_steps: int) -> None:
+    """The manual shift must reproduce the sigmas the txt2img scheduler produces."""
+    image_seq_len = 64 * 64
+    timesteps = get_schedule_flux2(num_steps=num_steps, image_seq_len=image_seq_len)
+    mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=num_steps)
+
+    scheduler = _txt2img_scheduler()
+    scheduler.set_timesteps(sigmas=timesteps[:-1], mu=mu)
+
+    shifted = time_shift_flux2(timesteps, mu)
+    np.testing.assert_allclose(shifted, [float(s) for s in scheduler.sigmas], rtol=0, atol=1e-6)
+
+
+def test_shift_is_a_no_op_at_the_endpoints() -> None:
+    """1.0 and 0.0 are fixed points, and 0.0 must not divide by zero."""
+    shifted = time_shift_flux2([1.0, 0.5, 0.0], mu=2.02)
+    assert shifted[0] == 1.0
+    assert shifted[-1] == 0.0
+
+
+def test_shift_is_strictly_decreasing_and_bounded() -> None:
+    timesteps = get_schedule_flux2(num_steps=30, image_seq_len=64 * 64)
+    shifted = time_shift_flux2(timesteps, mu=compute_empirical_mu(image_seq_len=64 * 64, num_steps=30))
+
+    assert all(0.0 <= s <= 1.0 for s in shifted)
+    assert all(a > b for a, b in zip(shifted[:-1], shifted[1:], strict=True))
+
+
+def test_shift_raises_the_schedule_floor_above_the_linear_one() -> None:
+    """The property that matters: the model is never asked for the low sigmas of the linear schedule."""
+    num_steps = 9
+    linear = get_schedule_flux2(num_steps=num_steps, image_seq_len=64 * 64)
+    shifted = time_shift_flux2(linear, mu=compute_empirical_mu(image_seq_len=64 * 64, num_steps=num_steps))
+
+    # Lowest sigma the model is actually evaluated at (the final 0.0 entry is the step target, not a
+    # timestep the model is called with).
+    assert linear[-2] == pytest.approx(1 / num_steps)
+    assert shifted[-2] > 0.45
+    assert all(s >= lin for s, lin in zip(shifted, linear, strict=True))