diff --git a/invokeai/backend/flux2/denoise.py b/invokeai/backend/flux2/denoise.py index f5900778064..4dbb99a5b14 100644 --- a/invokeai/backend/flux2/denoise.py +++ b/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 diff --git a/tests/backend/flux2/test_denoise_ref_image_conditioning.py b/tests/backend/flux2/test_denoise_ref_image_conditioning.py new file mode 100644 index 00000000000..b53c5624d72 --- /dev/null +++ b/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)