[Feat] Support CFG-Parallel For Sensenova-U1-8B-MoT (#4188)

Signed-off-by: sphinxkkkbc <binchengkang8@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: sphinxkkkbc <binchengkang8@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

docs/user_guide/diffusion_features.md

@@ -133,7 +133,7 @@ The following tables show which models support each feature:
 | **Qwen-Image-Edit**      |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Qwen-Image-Edit-2509** |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |        ✅ (decode)         |          ✅           |       ❌        |        ❌         |
 | **Qwen-Image-Layered**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **SenseNova-U1**         |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **SenseNova-U1**         |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Stable-Diffusion-XL**  |     ❌     |     ❌      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Stable-Diffusion3.5**  |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **Z-Image**              |     ✅     |     ✅      |           ✅           |       ❓        |   ✅ (TP=2 only)   |          ❌          |   ✅    |             ❌             |      ✅ (decode)      |       ✅        |        ❌         |

examples/offline_inference/sensenova_u1/end2end.py

@@ -139,6 +139,12 @@ def parse_args():
         default=1,
         help="Number of GPUs for tensor parallelism.",
     )
+    parser.add_argument(
+        "--cfg-parallel-size",
+        type=int,
+        default=1,
+        help="Number of GPUs for CFG parallelism.",
+    )
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
@@ -194,6 +200,7 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         cache_backend=args.cache_backend,
         enable_cache_dit_summary=args.enable_cache_dit_summary,
+        cfg_parallel_size=args.cfg_parallel_size,
     )
 
     extra_args = {
@@ -237,6 +244,7 @@ def main():
     print(f"  Seed           : {args.seed}")
     print(f"  Think mode     : {args.think}")
     print(f"  TP size        : {args.tensor_parallel_size}")
+    print(f"  CFG size       : {args.cfg_parallel_size}")
     print(f"  Cache backend  : {args.cache_backend or 'none'}")
     print(f"{'=' * 60}\n")
 

tests/diffusion/distributed/test_cfg_parallel.py

@@ -208,6 +208,7 @@ def _test_cfg_parallel_worker(
             positive_kwargs=positive_kwargs,
             negative_kwargs=negative_kwargs,
             cfg_normalize=test_config["cfg_normalize"],
+            kwargs=test_config["kwargs"],
         )
 
     # CFG parallel returns the combined prediction on every rank.
@@ -284,6 +285,7 @@ def _test_cfg_sequential_worker(
             positive_kwargs=positive_kwargs,
             negative_kwargs=negative_kwargs,
             cfg_normalize=test_config["cfg_normalize"],
+            kwargs=test_config["kwargs"],
         )
 
     # Sequential CFG always returns output
@@ -297,7 +299,10 @@ def _test_cfg_sequential_worker(
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
 @pytest.mark.parametrize("batch_size", [2])
 @pytest.mark.parametrize("cfg_normalize", [False, True])
-def test_predict_noise_maybe_with_cfg(cfg_parallel_size: int, dtype: torch.dtype, batch_size: int, cfg_normalize: bool):
+@pytest.mark.parametrize("kwargs", [None, {"step_i": 1}])
+def test_predict_noise_maybe_with_cfg(
+    cfg_parallel_size: int, dtype: torch.dtype, batch_size: int, cfg_normalize: bool, kwargs: dict | None
+):
     """
     Test that predict_noise_maybe_with_cfg produces identical results
     with and without CFG parallel.
@@ -322,6 +327,7 @@ def test_predict_noise_maybe_with_cfg(cfg_parallel_size: int, dtype: torch.dtype
         "cfg_normalize": cfg_normalize,
         "model_seed": 42,  # Fixed seed for model initialization
         "input_seed": 123,  # Fixed seed for input generation
+        "kwargs": kwargs,  # Additional kwargs to test passing through CFG parallel
     }
 
     mp_context = torch.multiprocessing.get_context("spawn")

tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline
+
+
+def _pipeline_without_init() -> SenseNovaU1Pipeline:
+    return object.__new__(SenseNovaU1Pipeline)
+
+
+@pytest.mark.parametrize("kwargs", [None, {"step_i": 0}])
+def test_combine_cfg_noise_requires_is_it2i(kwargs):
+    pipe = _pipeline_without_init()
+    out_cond = (torch.ones(1, 2, 3),)
+    out_uncond = (torch.zeros(1, 2, 3),)
+
+    with pytest.raises(ValueError, match="is_it2i"):
+        pipe.combine_cfg_noise(
+            out_cond,
+            out_uncond,
+            cfg_scale=4.0,
+            cfg_norm="cfg_zero_star",
+            kwargs=kwargs,
+        )
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [{"is_it2i": False}, {"is_it2i": False, "step_i": None}],
+)
+def test_cfg_zero_star_requires_step_i(kwargs):
+    pipe = _pipeline_without_init()
+    out_cond = (torch.ones(1, 2, 3),)
+    out_uncond = (torch.zeros(1, 2, 3),)
+    with pytest.raises(ValueError, match="step_i"):
+        pipe.combine_cfg_noise(
+            out_cond,
+            out_uncond,
+            cfg_scale=4.0,
+            cfg_norm="cfg_zero_star",
+            kwargs=kwargs,
+        )
+
+
+def test_cfg_zero_star_accepts_step_i():
+    pipe = _pipeline_without_init()
+    out_cond = (torch.ones(1, 2, 3),)
+    out_uncond = (torch.zeros(1, 2, 3),)
+    result = pipe.combine_cfg_noise(
+        out_cond,
+        out_uncond,
+        cfg_scale=4.0,
+        cfg_norm="cfg_zero_star",
+        kwargs={"is_it2i": False, "step_i": 0},
+    )
+
+    assert result.shape == out_cond[0].shape
+    assert torch.equal(result, torch.zeros_like(out_cond[0]))
+    assert torch.isfinite(result).all()

vllm_omni/diffusion/distributed/cfg_parallel.py

@@ -81,6 +81,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
         negative_kwargs: dict[str, Any] | None,
         cfg_normalize: bool = True,
         output_slice: int | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """
         Predict noise with optional classifier-free guidance.
@@ -92,6 +93,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
             negative_kwargs: Kwargs for negative/unconditional prediction
             cfg_normalize: Whether to normalize CFG output (default: True)
             output_slice: If set, slice each output to [:, :output_slice] for image editing
+            kwargs: Optional extra context for custom combine implementations.
 
         Returns:
             Predicted noise tensor or tuple of tensors.
@@ -111,8 +113,12 @@ class CFGParallelMixin(metaclass=ABCMeta):
                 cfg_rank = get_classifier_free_guidance_rank()
 
                 # Each rank computes one branch
-                kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs
-                local_pred = _wrap(self.predict_noise(**kwargs))
+                if cfg_rank == 0:
+                    logger.debug("CFG Parallel: Rank 0 computing positive branch")
+                    local_pred = _wrap(self.predict_noise(**positive_kwargs))
+                else:
+                    logger.debug("CFG Parallel: Rank %d computing negative branch", cfg_rank)
+                    local_pred = _wrap(self.predict_noise(**negative_kwargs))
 
                 if output_slice is not None:
                     local_pred = _slice_pred(local_pred, output_slice)
@@ -128,6 +134,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
                     negative_noise_pred,
                     true_cfg_scale,
                     cfg_normalize,
+                    **({} if kwargs is None else {"kwargs": kwargs}),
                 )
             else:
                 # Sequential CFG: compute both positive and negative
@@ -143,6 +150,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
                     negative_noise_pred,
                     true_cfg_scale,
                     cfg_normalize,
+                    **({} if kwargs is None else {"kwargs": kwargs}),
                 )
         else:
             # No CFG: only compute positive/conditional prediction
@@ -173,6 +181,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
         negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
         true_cfg_scale: float,
         cfg_normalize: bool = False,
+        kwargs: dict[str, Any] | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """
         Combine conditional and unconditional noise predictions with CFG.
@@ -196,6 +205,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
             negative_noise_pred: Negative/unconditional prediction(s) — Tensor or tuple
             true_cfg_scale: CFG scale factor
             cfg_normalize: Whether to normalize the combined prediction (default: False)
+            kwargs: Optional extra context for custom combine implementations.
 
         Returns:
             Combined noise prediction(s) — same type as inputs
@@ -276,7 +286,7 @@ class CFGParallelMixin(metaclass=ABCMeta):
         branches_kwargs: list[dict[str, Any]],
         n_branches: int,
         cfg_world_size: int,
-        true_cfg_scale: float,
+        true_cfg_scale: float | dict[str, float],
         cfg_normalize: bool,
         output_slice: int | None,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:

vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py

@@ -22,7 +22,7 @@ import math
 import os
 from collections.abc import Iterable
 from types import SimpleNamespace
-from typing import ClassVar
+from typing import Any, ClassVar
 
 import numpy as np
 import torch
@@ -33,6 +33,7 @@ from transformers import AutoTokenizer
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
@@ -489,7 +490,7 @@ def _optimized_scale(positive_flat, negative_flat):
 # ---------------------------------------------------------------------------
 
 
-class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin):
+class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin, CFGParallelMixin):
     """SenseNova-U1 text-to-image and image-to-image pipeline for vllm-omni.
 
     Builds the full model graph internally:
@@ -1055,123 +1056,136 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             timesteps=timesteps,
         )
 
-    def _denoise_step(self, image_prediction, ns, t, z, image_embeds, caches, p, step_i):
-        """Execute one denoising step with the appropriate CFG strategy."""
-        kv_cond = caches["cond"]
-        idx_cond = caches["idx_cond"]
-        mask_cond = caches["mask_cond"]
-        is_it2i = "img_cond" in caches
-        if is_it2i:
-            use_cfg = (t > p.cfg_interval[0] and t < p.cfg_interval[1]) or p.cfg_interval[0] == 0
-            needs_cfg = not (p.cfg_scale == 1 and p.img_cfg_scale == 1)
-            has_cached_partner = use_cfg and needs_cfg
-        else:
-            has_cached_partner = t >= p.cfg_interval[0] and t <= p.cfg_interval[1] and p.cfg_scale > 1
-
-        out_cond = self._t2i_predict_v(
-            image_embeds,
-            idx_cond,
-            mask_cond,
-            kv_cond,
-            t,
-            z,
+    def _get_cfg_kwargs(self, caches: dict, image_embeds, t, z, ns, p, branch: str, cache_dit_skip: bool = False):
+        required = (branch, f"idx_{branch}", f"mask_{branch}")
+        missing = [key for key in required if key not in caches]
+        if missing:
+            raise KeyError(f"Missing {branch} CFG cache keys: {missing}")
+        kwargs = dict(
+            input_embeds=image_embeds,
+            past_key_values=caches[branch],
+            indexes_image=caches[f"idx_{branch}"],
+            attn_mask=caches[f"mask_{branch}"],
+            t=t,
+            z=z,
             image_token_num=ns.token_h * ns.token_w,
             image_size=p.image_size,
-            cache_dit_skip=not has_cached_partner,
         )
+        if cache_dit_skip:
+            kwargs["cache_dit_skip"] = True
+        return kwargs
 
-        if is_it2i:
-            return self._denoise_step_it2i(out_cond, image_embeds, ns, t, z, caches, p, step_i)
-        return self._denoise_step_t2i(out_cond, image_embeds, ns, t, z, caches, p, step_i)
-
-    def _denoise_step_t2i(self, out_cond, image_embeds, ns, t, z, caches, p, step_i):
-        """T2I: two-way CFG (condition vs uncondition), supports cfg_zero_star."""
-        if t >= p.cfg_interval[0] and t <= p.cfg_interval[1] and p.cfg_scale > 1:
-            out_uncond = self._t2i_predict_v(
-                image_embeds,
-                caches["idx_uncond"],
-                caches["mask_uncond"],
-                caches["uncond"],
-                t,
-                z,
-                image_token_num=ns.token_h * ns.token_w,
-                image_size=p.image_size,
+    def _denoise(self, image_prediction, ns, t, z, image_embeds, caches, p, step_i, is_it2i):
+        if not is_it2i:
+            has_cached_partner = t >= p.cfg_interval[0] and t <= p.cfg_interval[1] and p.cfg_scale > 1
+            cond_kwargs = self._get_cfg_kwargs(
+                caches, image_embeds, t, z, ns, p, branch="cond", cache_dit_skip=not has_cached_partner
             )
-            if p.cfg_norm == "cfg_zero_star":
-                pos_flat = out_cond.view(p.batch_size, -1)
-                neg_flat = out_uncond.view(p.batch_size, -1)
-                alpha = _optimized_scale(pos_flat, neg_flat)
-                alpha = alpha.view(p.batch_size, *([1] * (len(out_cond.shape) - 1))).to(pos_flat.dtype)
-                v_pred = (
-                    out_cond * 0.0
-                    if step_i <= 0
-                    else (out_uncond * alpha + p.cfg_scale * (out_cond - out_uncond * alpha))
+
+            in_interval = t >= p.cfg_interval[0] and t <= p.cfg_interval[1]
+            if not (in_interval and p.cfg_scale > 1):
+                return self.predict_noise(**cond_kwargs)
+
+            uncond_kwargs = self._get_cfg_kwargs(caches, image_embeds, t, z, ns, p, branch="uncond")
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=True,
+                true_cfg_scale=p.cfg_scale,
+                positive_kwargs=cond_kwargs,
+                negative_kwargs=uncond_kwargs,
+                cfg_normalize=p.cfg_norm,
+                kwargs={"step_i": step_i, "is_it2i": is_it2i},
+            )
+            return noise_pred
+        else:
+            use_cfg = (t > p.cfg_interval[0] and t < p.cfg_interval[1]) or p.cfg_interval[0] == 0
+            needs_cfg = p.cfg_scale != 1 or p.img_cfg_scale != 1
+            has_cached_partner = use_cfg and needs_cfg
+            cond_kwargs = self._get_cfg_kwargs(
+                caches, image_embeds, t, z, ns, p, branch="cond", cache_dit_skip=not has_cached_partner
+            )
+
+            if not use_cfg or not needs_cfg:
+                return self.predict_noise(**cond_kwargs)
+
+            cfg_norm = p.cfg_norm if (p.cfg_scale > 1 or p.img_cfg_scale > 1) else None
+            if p.img_cfg_scale == 1:
+                image_cond_kwargs = self._get_cfg_kwargs(caches, image_embeds, t, z, ns, p, branch="img_cond")
+                noise_pred = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg=True,
+                    true_cfg_scale=p.cfg_scale,
+                    positive_kwargs=cond_kwargs,
+                    negative_kwargs=image_cond_kwargs,
+                    cfg_normalize=cfg_norm,
+                    kwargs={"is_it2i": is_it2i},
+                )
+            elif p.cfg_scale == p.img_cfg_scale:
+                uncond_kwargs = self._get_cfg_kwargs(caches, image_embeds, t, z, ns, p, branch="uncond")
+                noise_pred = self.predict_noise_maybe_with_cfg(
+                    do_true_cfg=True,
+                    true_cfg_scale=p.cfg_scale,
+                    positive_kwargs=cond_kwargs,
+                    negative_kwargs=uncond_kwargs,
+                    cfg_normalize=cfg_norm,
+                    kwargs={"is_it2i": is_it2i},
                 )
             else:
-                v_pred = out_uncond + p.cfg_scale * (out_cond - out_uncond)
-                v_pred = self._apply_cfg_norm(v_pred, out_cond, p.cfg_norm)
+                image_cond_kwargs = self._get_cfg_kwargs(
+                    caches, image_embeds, t, z, ns, p, branch="img_cond", cache_dit_skip=True
+                )
+                uncond_kwargs = self._get_cfg_kwargs(caches, image_embeds, t, z, ns, p, branch="uncond")
+                noise_pred = self.predict_noise_with_multi_branch_cfg(
+                    do_true_cfg=True,
+                    true_cfg_scale={
+                        "cfg_scale": p.cfg_scale,
+                        "img_cfg_scale": p.img_cfg_scale,
+                    },
+                    branches_kwargs=[cond_kwargs, image_cond_kwargs, uncond_kwargs],
+                    cfg_normalize=cfg_norm,
+                )
+            return noise_pred
+
+    def predict_noise(self, **kwargs):
+        return self._t2i_predict_v(**kwargs)
+
+    def combine_cfg_noise(self, out_cond, out_uncond, cfg_scale, cfg_norm, kwargs: dict[str, Any] | None = None):
+        if not isinstance(kwargs, dict) or "is_it2i" not in kwargs:
+            raise ValueError("SenseNovaU1 CFG requires kwargs['is_it2i']")
+        if isinstance(cfg_scale, dict):
+            raise NotImplementedError("dict cfg_scale requires multi-branch CFG")
+
+        out_cond = out_cond[0]
+        out_uncond = out_uncond[0]
+        batch_size = out_cond.shape[0]
+
+        is_it2i = kwargs.get("is_it2i")
+        if not is_it2i and cfg_norm == "cfg_zero_star":
+            if not isinstance(kwargs, dict) or not isinstance(kwargs.get("step_i"), int):
+                raise ValueError("SenseNovaU1 T2I CFG requires kwargs['step_i']")
+            step_i = kwargs.get("step_i")
+            pos_flat = out_cond.view(batch_size, -1)
+            neg_flat = out_uncond.view(batch_size, -1)
+            alpha = _optimized_scale(pos_flat, neg_flat)
+            alpha = alpha.view(batch_size, *([1] * (len(out_cond.shape) - 1))).to(pos_flat.dtype)
+            v_pred = (
+                out_cond * 0.0 if step_i <= 0 else (out_uncond * alpha + cfg_scale * (out_cond - out_uncond * alpha))
+            )
         else:
-            v_pred = out_cond
+            if cfg_norm == "cfg_zero_star":
+                logger.warning("cfg_zero_star is T2I-only; please use 'global' or 'channel' instead.")
+            # cfg_zero_star is T2I-only; _apply_cfg_norm treats it as a no-op for IT2I.
+            v_pred = out_uncond + cfg_scale * (out_cond - out_uncond)
+            v_pred = self._apply_cfg_norm(v_pred, out_cond, cfg_norm)
         return v_pred
 
-    def _denoise_step_it2i(self, out_cond, image_embeds, ns, t, z, caches, p, step_i):
-        """IT2I: dual CFG with ``cfg_scale`` (text) and ``img_cfg_scale`` (image)."""
-        use_cfg = (t > p.cfg_interval[0] and t < p.cfg_interval[1]) or p.cfg_interval[0] == 0
-        needs_cfg = not (p.cfg_scale == 1 and p.img_cfg_scale == 1)
-
-        if not use_cfg or not needs_cfg:
-            return out_cond
-
-        predict_kw = dict(image_token_num=ns.token_h * ns.token_w, image_size=p.image_size)
-
-        if p.img_cfg_scale == 1:
-            out_img_cond = self._t2i_predict_v(
-                image_embeds,
-                caches["idx_img_cond"],
-                caches["mask_img_cond"],
-                caches["img_cond"],
-                t,
-                z,
-                **predict_kw,
-            )
-            v_pred = out_img_cond + p.cfg_scale * (out_cond - out_img_cond)
-        elif p.cfg_scale == p.img_cfg_scale:
-            out_uncond = self._t2i_predict_v(
-                image_embeds,
-                caches["idx_uncond"],
-                caches["mask_uncond"],
-                caches["uncond"],
-                t,
-                z,
-                **predict_kw,
-            )
-            v_pred = out_uncond + p.cfg_scale * (out_cond - out_uncond)
-        else:
-            out_img_cond = self._t2i_predict_v(
-                image_embeds,
-                caches["idx_img_cond"],
-                caches["mask_img_cond"],
-                caches["img_cond"],
-                t,
-                z,
-                cache_dit_skip=True,
-                **predict_kw,
-            )
-            out_uncond = self._t2i_predict_v(
-                image_embeds,
-                caches["idx_uncond"],
-                caches["mask_uncond"],
-                caches["uncond"],
-                t,
-                z,
-                **predict_kw,
-            )
-            v_pred = (
-                out_uncond + p.cfg_scale * (out_cond - out_img_cond) + p.img_cfg_scale * (out_img_cond - out_uncond)
-            )
-
-        if p.cfg_scale > 1 or p.img_cfg_scale > 1:
-            v_pred = self._apply_cfg_norm(v_pred, out_cond, p.cfg_norm)
+    def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize):
+        out_cond = predictions[0]
+        out_img_cond = predictions[1]
+        out_uncond = predictions[2]
+        cfg_scale = true_cfg_scale["cfg_scale"]
+        img_cfg_scale = true_cfg_scale["img_cfg_scale"]
+        v_pred = out_uncond + cfg_scale * (out_cond - out_img_cond) + img_cfg_scale * (out_img_cond - out_uncond)
+        if cfg_scale > 1 or img_cfg_scale > 1:
+            v_pred = self._apply_cfg_norm(v_pred, out_cond, cfg_normalize)
         return v_pred
 
     @staticmethod
@@ -1252,7 +1266,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             "idx_uncond": indexes_image_uncond,
             "mask_uncond": {"full_attention": None},
         }
-        return self._run_denoising_loop(ns, caches, p, think_text)
+        return self._run_denoising_loop(ns, caches, p, think_text, is_it2i=False)
 
     def _forward_it2i(self, p, input_images: list[Image.Image]) -> DiffusionOutput:
         """Image-to-image (editing) generation path with dual CFG."""
@@ -1368,9 +1382,9 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             if key in caches and not isinstance(caches[key], dict):
                 self._expand_and_prepare_kv(caches[key], ns.token_h * ns.token_w, p.batch_size)
 
-        return self._run_denoising_loop(ns, caches, p, think_text)
+        return self._run_denoising_loop(ns, caches, p, think_text, is_it2i=True)
 
-    def _run_denoising_loop(self, ns, caches, p, think_text="") -> DiffusionOutput:
+    def _run_denoising_loop(self, ns, caches, p, think_text="", is_it2i: bool = False) -> DiffusionOutput:
         """Shared denoising loop for both T2I and IT2I."""
         merge_size = self.merge_size
         image_prediction = ns.image_prediction
@@ -1403,7 +1417,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
                 timestep_embeddings = timestep_embeddings + ns_emb
             image_embeds = image_embeds + timestep_embeddings
 
-            v_pred = self._denoise_step(image_prediction, ns, t, z, image_embeds, caches, p, step_i)
+            v_pred = self._denoise(image_prediction, ns, t, z, image_embeds, caches, p, step_i, is_it2i)
             z = z + (t_next - t) * v_pred
             image_prediction = _unpatchify(z, self.patch_size * merge_size, p.image_size[1], p.image_size[0])