diff --git a/invokeai/app/invocations/flux2_denoise.py b/invokeai/app/invocations/flux2_denoise.py index 7a5ce158c98..47234bea909 100644 --- a/invokeai/app/invocations/flux2_denoise.py +++ b/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 diff --git a/invokeai/backend/flux2/sampling_utils.py b/invokeai/backend/flux2/sampling_utils.py index 3981e912756..50cd2b98f71 100644 --- a/invokeai/backend/flux2/sampling_utils.py +++ b/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. diff --git a/tests/app/invocations/test_flux2_denoise_img2img_normalization.py b/tests/app/invocations/test_flux2_denoise_img2img_normalization.py new file mode 100644 index 00000000000..70fc3006d65 --- /dev/null +++ b/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) diff --git a/tests/backend/flux2/test_flux2_schedule_shift.py b/tests/backend/flux2/test_flux2_schedule_shift.py new file mode 100644 index 00000000000..d19fc165408 --- /dev/null +++ b/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))