From dd3a4d0a29756429f703244db5cac13c2d082f33 Mon Sep 17 00:00:00 2001 From: boatman <109857087+sphinxkkkbc@users.noreply.github.com> Date: Tue, 30 Jun 2026 16:03:26 +0800 Subject: [PATCH] [Feat] Support CFG-Parallel For Sensenova-U1-8B-MoT (#4188) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: sphinxkkkbc Signed-off-by: 汪志鹏 Co-authored-by: sphinxkkkbc Co-authored-by: 汪志鹏 --- docs/user_guide/diffusion_features.md | 2 +- .../offline_inference/sensenova_u1/end2end.py | 8 + .../distributed/test_cfg_parallel.py | 8 +- .../sensenova_u1/test_sensenova_u1_cfg.py | 62 +++++ .../diffusion/distributed/cfg_parallel.py | 16 +- .../sensenova_u1/pipeline_sensenova_u1.py | 240 +++++++++--------- 6 files changed, 218 insertions(+), 118 deletions(-) create mode 100644 tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 3ac70a754..6b923a793 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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) | ✅ | ❌ | diff --git a/examples/offline_inference/sensenova_u1/end2end.py b/examples/offline_inference/sensenova_u1/end2end.py index 0716c74d5..05f037dd7 100644 --- a/examples/offline_inference/sensenova_u1/end2end.py +++ b/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") diff --git a/tests/diffusion/distributed/test_cfg_parallel.py b/tests/diffusion/distributed/test_cfg_parallel.py index bf709618d..af81f00ef 100644 --- a/tests/diffusion/distributed/test_cfg_parallel.py +++ b/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") diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py new file mode 100644 index 000000000..818d3d7c0 --- /dev/null +++ b/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() diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py index 6810a447b..83ce73b16 100644 --- a/vllm_omni/diffusion/distributed/cfg_parallel.py +++ b/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, ...]: diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index 080ed10e6..73951fb35 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/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])