fix(flux2): stop the sampler from advancing reference image tokens

invokeai/backend/flux2/denoise.py

@@ -78,13 +78,22 @@ def denoise(
     """
     total_steps = len(timesteps) - 1
 
-    # Store original sequence length for extracting output later (before concatenating reference images)
+    # Sequence length of the generated latents, i.e. everything that is *not* reference conditioning.
     original_seq_len = img.shape[1]
 
-    # Concatenate reference image conditioning if provided (multi-reference image editing)
-    if img_cond_seq is not None and img_cond_seq_ids is not None:
-        img = torch.cat([img, img_cond_seq], dim=1)
-        img_ids = torch.cat([img_ids, img_cond_seq_ids], dim=1)
+    # Reference image conditioning (multi-reference image editing) is constant context: it is
+    # concatenated onto the latents for every forward pass, but it must NOT be advanced by the
+    # sampler. Concatenating it onto `img` itself would make every sampler step integrate the
+    # reference tokens along the model's velocity field, so the reference drifts away from the
+    # encoded image over the schedule (and drags the generated image with it). Instead, build the
+    # model input per step and slice the prediction back to the generated tokens - same as the
+    # FLUX.1 Kontext path and diffusers' Flux2KleinPipeline.
+    # The position IDs for the concatenated sequence are constant, so precompute them once.
+    if img_cond_seq is not None:
+        assert img_cond_seq_ids is not None, "You need to provide either both or neither of the sequence conditioning"
+        model_img_ids = torch.cat([img_ids, img_cond_seq_ids], dim=1)
+    else:
+        model_img_ids = img_ids
 
     # The transformer forward() requires a guidance tensor even when guidance_embeds=False,
     # because the Flux2TimestepGuidanceEmbeddings forward signature takes it unconditionally.
@@ -132,12 +141,15 @@ def denoise(
             # Track if we're in first or second order step (for Heun)
             in_first_order = scheduler.state_in_first_order if is_heun else True
 
+            # Append the (unchanged) reference conditioning for this forward pass only.
+            img_input = torch.cat([img, img_cond_seq], dim=1) if img_cond_seq is not None else img
+
             # Run the transformer model (matching diffusers: guidance=guidance, return_dict=False)
             output = model(
-                hidden_states=img,
+                hidden_states=img_input,
                 encoder_hidden_states=txt,
                 timestep=t_vec,
-                img_ids=img_ids,
+                img_ids=model_img_ids,
                 txt_ids=txt_ids,
                 guidance=guidance_vec,
                 joint_attention_kwargs=pos_joint_attention_kwargs,
@@ -147,6 +159,10 @@ def denoise(
             # Extract the sample from the output (return_dict=False returns tuple)
             pred = output[0] if isinstance(output, tuple) else output
 
+            # Drop the prediction for the reference tokens - they are context, not sampled state.
+            if img_cond_seq is not None:
+                pred = pred[:, :original_seq_len]
+
             step_cfg_scale = cfg_scale[min(user_step, len(cfg_scale) - 1)]
 
             # Apply CFG if scale is not 1.0
@@ -155,16 +171,18 @@ def denoise(
                     raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
 
                 neg_output = model(
-                    hidden_states=img,
+                    hidden_states=img_input,
                     encoder_hidden_states=neg_txt,
                     timestep=t_vec,
-                    img_ids=img_ids,
+                    img_ids=model_img_ids,
                     txt_ids=neg_txt_ids if neg_txt_ids is not None else txt_ids,
                     guidance=guidance_vec,
                     return_dict=False,
                 )
 
                 neg_pred = neg_output[0] if isinstance(neg_output, tuple) else neg_output
+                if img_cond_seq is not None:
+                    neg_pred = neg_pred[:, :original_seq_len]
                 pred = neg_pred + step_cfg_scale * (pred - neg_pred)
 
             # Use scheduler.step() for the update
@@ -179,15 +197,7 @@ def denoise(
 
             # Apply inpainting merge at each step
             if inpaint_extension is not None:
-                # Separate the generated latents from the reference conditioning
-                gen_img = img[:, :original_seq_len, :]
-                ref_img = img[:, original_seq_len:, :]
-
-                # Merge only the generated part
-                gen_img = inpaint_extension.merge_intermediate_latents_with_init_latents(gen_img, t_prev)
-
-                # Concatenate back together
-                img = torch.cat([gen_img, ref_img], dim=1)
+                img = inpaint_extension.merge_intermediate_latents_with_init_latents(img, t_prev)
 
             # For Heun, only increment user step after second-order step completes
             if is_heun:
@@ -200,17 +210,13 @@ def denoise(
                             preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(
                                 preview_img, 0.0
                             )
-                        # Extract only the generated image portion for preview (exclude reference images)
-                        callback_latents = (
-                            preview_img[:, :original_seq_len, :] if img_cond_seq is not None else preview_img
-                        )
                         step_callback(
                             PipelineIntermediateState(
                                 step=user_step,
                                 order=2,
                                 total_steps=total_steps,
                                 timestep=int(t_curr * 1000),
-                                latents=callback_latents,
+                                latents=preview_img,
                             ),
                         )
             else:
@@ -220,15 +226,13 @@ def denoise(
                     preview_img = img - t_curr * pred
                     if inpaint_extension is not None:
                         preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(preview_img, 0.0)
-                    # Extract only the generated image portion for preview (exclude reference images)
-                    callback_latents = preview_img[:, :original_seq_len, :] if img_cond_seq is not None else preview_img
                     step_callback(
                         PipelineIntermediateState(
                             step=user_step,
                             order=1,
                             total_steps=total_steps,
                             timestep=int(t_curr * 1000),
-                            latents=callback_latents,
+                            latents=preview_img,
                         ),
                     )
 
@@ -241,12 +245,15 @@ def denoise(
         ):
             t_vec = torch.full((img.shape[0],), t_curr, dtype=img.dtype, device=img.device)
 
+            # Append the (unchanged) reference conditioning for this forward pass only.
+            img_input = torch.cat([img, img_cond_seq], dim=1) if img_cond_seq is not None else img
+
             # Run the transformer model (matching diffusers: guidance=guidance, return_dict=False)
             output = model(
-                hidden_states=img,
+                hidden_states=img_input,
                 encoder_hidden_states=txt,
                 timestep=t_vec,
-                img_ids=img_ids,
+                img_ids=model_img_ids,
                 txt_ids=txt_ids,
                 guidance=guidance_vec,
                 joint_attention_kwargs=pos_joint_attention_kwargs,
@@ -256,6 +263,10 @@ def denoise(
             # Extract the sample from the output (return_dict=False returns tuple)
             pred = output[0] if isinstance(output, tuple) else output
 
+            # Drop the prediction for the reference tokens - they are context, not sampled state.
+            if img_cond_seq is not None:
+                pred = pred[:, :original_seq_len]
+
             step_cfg_scale = cfg_scale[step_index]
 
             # Apply CFG if scale is not 1.0
@@ -264,16 +275,18 @@ def denoise(
                     raise ValueError("Negative text conditioning is required when cfg_scale is not 1.0.")
 
                 neg_output = model(
-                    hidden_states=img,
+                    hidden_states=img_input,
                     encoder_hidden_states=neg_txt,
                     timestep=t_vec,
-                    img_ids=img_ids,
+                    img_ids=model_img_ids,
                     txt_ids=neg_txt_ids if neg_txt_ids is not None else txt_ids,
                     guidance=guidance_vec,
                     return_dict=False,
                 )
 
                 neg_pred = neg_output[0] if isinstance(neg_output, tuple) else neg_output
+                if img_cond_seq is not None:
+                    neg_pred = neg_pred[:, :original_seq_len]
                 pred = neg_pred + step_cfg_scale * (pred - neg_pred)
 
             # Euler step
@@ -282,34 +295,17 @@ def denoise(
 
             # Apply inpainting merge at each step
             if inpaint_extension is not None:
-                # Separate the generated latents from the reference conditioning
-                gen_img = img[:, :original_seq_len, :]
-                ref_img = img[:, original_seq_len:, :]
-
-                # Merge only the generated part
-                gen_img = inpaint_extension.merge_intermediate_latents_with_init_latents(gen_img, t_prev)
-
-                # Concatenate back together
-                img = torch.cat([gen_img, ref_img], dim=1)
+                img = inpaint_extension.merge_intermediate_latents_with_init_latents(img, t_prev)
+                preview_img = inpaint_extension.merge_intermediate_latents_with_init_latents(preview_img, 0.0)
 
-                # Handling preview images
-                preview_gen = preview_img[:, :original_seq_len, :]
-                preview_gen = inpaint_extension.merge_intermediate_latents_with_init_latents(preview_gen, 0.0)
-
-            # Extract only the generated image portion for preview (exclude reference images)
-            callback_latents = preview_img[:, :original_seq_len, :] if img_cond_seq is not None else preview_img
             step_callback(
                 PipelineIntermediateState(
                     step=step_index + 1,
                     order=1,
                     total_steps=total_steps,
                     timestep=int(t_curr),
-                    latents=callback_latents,
+                    latents=preview_img,
                 ),
             )
 
-    # Extract only the generated image portion (exclude concatenated reference images)
-    if img_cond_seq is not None:
-        img = img[:, :original_seq_len, :]
-
     return img

tests/backend/flux2/test_denoise_ref_image_conditioning.py

@@ -0,0 +1,160 @@
+"""Regression tests for FLUX.2 reference image conditioning during denoising.
+
+Reference image latents are *context*: they are concatenated onto the latents for every forward
+pass, but the sampler must never advance them. A previous implementation concatenated them onto
+the sampled tensor once before the loop, so every step integrated the reference tokens along the
+model's velocity field - the reference silently drifted away from the encoded image over the
+schedule and dragged the generated image with it (visible as a reproducible spatial shift of the
+edited result).
+"""
+
+from typing import Any
+
+import pytest
+import torch
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+
+from invokeai.backend.flux2.denoise import denoise
+
+BATCH = 1
+GEN_SEQ_LEN = 6
+REF_SEQ_LEN = 4
+CHANNELS = 8
+TXT_SEQ_LEN = 3
+
+
+class RecordingModel:
+    """A stand-in transformer that records its inputs and predicts a constant velocity.
+
+    A constant non-zero prediction is enough to expose sampler updates leaking into the reference
+    tokens: any integration of the reference part changes it away from the encoded latents.
+    """
+
+    def __init__(self) -> None:
+        self.calls: list[dict[str, Any]] = []
+
+    def __call__(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        img_ids: torch.Tensor,
+        txt_ids: torch.Tensor,
+        guidance: torch.Tensor,
+        joint_attention_kwargs: dict[str, Any] | None = None,
+        return_dict: bool = True,
+    ) -> tuple[torch.Tensor]:
+        self.calls.append({"hidden_states": hidden_states.clone(), "img_ids": img_ids.clone()})
+        return (torch.full_like(hidden_states, 0.5),)
+
+
+def _inputs() -> dict[str, Any]:
+    generator = torch.Generator().manual_seed(0)
+    return {
+        "img": torch.randn(BATCH, GEN_SEQ_LEN, CHANNELS, generator=generator),
+        "img_ids": torch.zeros(BATCH, GEN_SEQ_LEN, 4, dtype=torch.long),
+        "txt": torch.randn(BATCH, TXT_SEQ_LEN, CHANNELS, generator=generator),
+        "txt_ids": torch.zeros(BATCH, TXT_SEQ_LEN, 4, dtype=torch.long),
+        "img_cond_seq": torch.randn(BATCH, REF_SEQ_LEN, CHANNELS, generator=generator),
+        "img_cond_seq_ids": torch.full((BATCH, REF_SEQ_LEN, 4), 10, dtype=torch.long),
+    }
+
+
+def _scheduler() -> FlowMatchEulerDiscreteScheduler:
+    return FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=3.0)
+
+
+@pytest.mark.parametrize("use_scheduler", [False, True], ids=["euler", "scheduler"])
+def test_reference_conditioning_is_constant_across_steps(use_scheduler: bool):
+    inputs = _inputs()
+    model = RecordingModel()
+    ref_latents = inputs["img_cond_seq"].clone()
+
+    out = denoise(
+        model=model,  # type: ignore[arg-type]
+        img=inputs["img"],
+        img_ids=inputs["img_ids"],
+        txt=inputs["txt"],
+        txt_ids=inputs["txt_ids"],
+        timesteps=[1.0, 0.75, 0.5, 0.25, 0.0],
+        step_callback=lambda state: None,
+        guidance=1.0,
+        cfg_scale=[1.0] * 4,
+        scheduler=_scheduler() if use_scheduler else None,
+        img_cond_seq=inputs["img_cond_seq"],
+        img_cond_seq_ids=inputs["img_cond_seq_ids"],
+    )
+
+    assert len(model.calls) == 4
+
+    expected_ids = torch.cat([inputs["img_ids"], inputs["img_cond_seq_ids"]], dim=1)
+    for step, call in enumerate(model.calls):
+        hidden_states = call["hidden_states"]
+        assert hidden_states.shape == (BATCH, GEN_SEQ_LEN + REF_SEQ_LEN, CHANNELS)
+        # The reference tokens must be bit-identical to the encoded latents at every step.
+        assert torch.equal(hidden_states[:, GEN_SEQ_LEN:, :], ref_latents), f"reference drifted at step {step}"
+        assert torch.equal(call["img_ids"], expected_ids)
+
+    # The caller's tensor must not be mutated either.
+    assert torch.equal(inputs["img_cond_seq"], ref_latents)
+
+    # Only the generated tokens are returned, and they were actually denoised.
+    assert out.shape == (BATCH, GEN_SEQ_LEN, CHANNELS)
+    assert not torch.equal(out, inputs["img"])
+
+
+@pytest.mark.parametrize("use_scheduler", [False, True], ids=["euler", "scheduler"])
+def test_reference_conditioning_with_cfg(use_scheduler: bool):
+    """With CFG the negative prediction must be sliced to the generated tokens as well."""
+    inputs = _inputs()
+    model = RecordingModel()
+    ref_latents = inputs["img_cond_seq"].clone()
+
+    out = denoise(
+        model=model,  # type: ignore[arg-type]
+        img=inputs["img"],
+        img_ids=inputs["img_ids"],
+        txt=inputs["txt"],
+        txt_ids=inputs["txt_ids"],
+        timesteps=[1.0, 0.5, 0.0],
+        step_callback=lambda state: None,
+        guidance=1.0,
+        cfg_scale=[2.0] * 2,
+        neg_txt=torch.zeros(BATCH, TXT_SEQ_LEN, CHANNELS),
+        neg_txt_ids=torch.zeros(BATCH, TXT_SEQ_LEN, 4, dtype=torch.long),
+        scheduler=_scheduler() if use_scheduler else None,
+        img_cond_seq=inputs["img_cond_seq"],
+        img_cond_seq_ids=inputs["img_cond_seq_ids"],
+    )
+
+    # One positive and one negative forward pass per step.
+    assert len(model.calls) == 4
+    for call in model.calls:
+        assert torch.equal(call["hidden_states"][:, GEN_SEQ_LEN:, :], ref_latents)
+
+    assert out.shape == (BATCH, GEN_SEQ_LEN, CHANNELS)
+
+
+@pytest.mark.parametrize("use_scheduler", [False, True], ids=["euler", "scheduler"])
+def test_without_reference_conditioning(use_scheduler: bool):
+    inputs = _inputs()
+    model = RecordingModel()
+
+    out = denoise(
+        model=model,  # type: ignore[arg-type]
+        img=inputs["img"],
+        img_ids=inputs["img_ids"],
+        txt=inputs["txt"],
+        txt_ids=inputs["txt_ids"],
+        timesteps=[1.0, 0.5, 0.0],
+        step_callback=lambda state: None,
+        guidance=1.0,
+        cfg_scale=[1.0] * 2,
+        scheduler=_scheduler() if use_scheduler else None,
+    )
+
+    for call in model.calls:
+        assert call["hidden_states"].shape == (BATCH, GEN_SEQ_LEN, CHANNELS)
+        assert torch.equal(call["img_ids"], inputs["img_ids"])
+
+    assert out.shape == (BATCH, GEN_SEQ_LEN, CHANNELS)