[Feature][Bagel] Add CFG parallel mode (#1578)

Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

tests/diffusion/models/bagel/test_combine_cfg.py

@@ -0,0 +1,316 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for Bagel._combine_cfg logic."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.bagel.bagel_transformer import Bagel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestCombineCfg:
+    """Tests for the _combine_cfg static method."""
+
+    def _make_tensors(self, shape=(10, 64), seed=42):
+        """Create deterministic test tensors."""
+        gen = torch.Generator().manual_seed(seed)
+        v_t = torch.randn(shape, generator=gen)
+        cfg_text_v_t = torch.randn(shape, generator=gen)
+        cfg_img_v_t = torch.randn(shape, generator=gen)
+        return v_t, cfg_text_v_t, cfg_img_v_t
+
+    def test_text_channel_renorm_preserves_direction(self):
+        """text_channel renorm should change direction but constrain magnitude."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        # Result norm per token should be <= original v_t norm (clamp max=1.0)
+        result_norm = torch.norm(result, dim=-1)
+        v_t_norm = torch.norm(v_t, dim=-1)
+        assert torch.all(result_norm <= v_t_norm + 1e-6), "text_channel renorm should not increase per-token norm"
+
+    def test_scale_1_returns_v_t(self):
+        """cfg_text_scale=1.0 means no CFG: result should equal v_t."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=1.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        # scale=1 → v_t_text_ = cfg_text + 1*(v_t - cfg_text) = v_t
+        # renorm scale = norm(v_t)/norm(v_t) = 1.0, so result = v_t
+        assert torch.allclose(result, v_t, atol=1e-6)
+
+    def test_img_cfg_applied_when_scale_gt_1(self):
+        """When cfg_img_scale > 1.0, result should differ from text-only CFG."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        text_only = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        with_img = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert not torch.allclose(text_only, with_img, atol=1e-6), (
+            "Image CFG should produce different result from text-only CFG"
+        )
+
+    def test_img_cfg_none_ignored(self):
+        """cfg_img_v_t=None should be equivalent to cfg_img_scale <= 1.0."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result_none = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        result_low_scale = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=0.5,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert torch.allclose(result_none, result_low_scale, atol=1e-6), (
+            "cfg_img_v_t=None and cfg_img_scale<=1.0 should give same result"
+        )
+
+    def test_global_renorm(self):
+        """global renorm should produce valid output without error."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="global",
+            cfg_renorm_min=0.0,
+        )
+
+        assert result.shape == v_t.shape
+        assert not torch.any(torch.isnan(result))
+
+    def test_channel_renorm(self):
+        """channel renorm should produce valid output without error."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert result.shape == v_t.shape
+        assert not torch.any(torch.isnan(result))
+
+    def test_invalid_renorm_type_raises(self):
+        """Unknown renorm type should raise NotImplementedError."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        with pytest.raises(NotImplementedError):
+            Bagel._combine_cfg(
+                v_t,
+                cfg_text_v_t,
+                None,
+                cfg_text_scale=4.0,
+                cfg_img_scale=1.0,
+                cfg_renorm_type="unknown",
+                cfg_renorm_min=0.0,
+            )
+
+    def test_renorm_min_clamps_scale(self):
+        """cfg_renorm_min should prevent scale from going too low."""
+        v_t = torch.ones(10, 64)
+        # Make cfg_text_v_t very different so CFG amplifies heavily
+        cfg_text_v_t = torch.zeros(10, 64)
+
+        result_no_min = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=100.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        result_with_min = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=100.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.5,
+        )
+
+        # With higher renorm_min, result magnitude should be larger
+        # (scale is clamped to at least 0.5 instead of going near 0)
+        norm_no_min = torch.norm(result_no_min)
+        norm_with_min = torch.norm(result_with_min)
+        assert norm_with_min >= norm_no_min - 1e-6, "Higher cfg_renorm_min should preserve more magnitude"
+
+    def test_global_renorm_with_img_cfg(self):
+        """global renorm + img CFG should produce valid, different output."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        text_only = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="global",
+            cfg_renorm_min=0.0,
+        )
+
+        with_img = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            cfg_img_v_t.clone(),
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="global",
+            cfg_renorm_min=0.0,
+        )
+
+        assert not torch.allclose(text_only, with_img, atol=1e-6), (
+            "global renorm + img CFG should differ from text-only"
+        )
+        assert not torch.any(torch.isnan(with_img))
+
+    def test_channel_renorm_with_img_cfg(self):
+        """channel renorm + img CFG should produce valid, different output."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        text_only = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="channel",
+            cfg_renorm_min=0.0,
+        )
+
+        with_img = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            cfg_img_v_t.clone(),
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert not torch.allclose(text_only, with_img, atol=1e-6), (
+            "channel renorm + img CFG should differ from text-only"
+        )
+        assert not torch.any(torch.isnan(with_img))
+
+    def test_global_channel_renorm_constrains_norm(self):
+        """global and channel renorm should not increase overall norm."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        for renorm_type in ["global", "channel"]:
+            result = Bagel._combine_cfg(
+                v_t.clone(),
+                cfg_text_v_t.clone(),
+                cfg_img_v_t.clone(),
+                cfg_text_scale=4.0,
+                cfg_img_scale=1.5,
+                cfg_renorm_type=renorm_type,
+                cfg_renorm_min=0.0,
+            )
+            # Global norm of result should be <= global norm of v_t (clamp max=1.0)
+            assert torch.norm(result) <= torch.norm(v_t) + 1e-5, f"{renorm_type} renorm should not increase global norm"
+
+    def test_text_channel_img_cfg_no_second_renorm(self):
+        """text_channel mode: img CFG is applied AFTER renorm, without a second renorm.
+        So the result norm can exceed v_t norm when img_scale > 1."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=2.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        # text_channel renorms after text CFG, then applies img CFG without renorm
+        # So result norm CAN exceed v_t norm — this is expected behavior
+        assert result.shape == v_t.shape
+        assert not torch.any(torch.isnan(result))
+
+    def test_all_renorm_types_consistent_direction(self):
+        """All renorm types should guide in the same general direction."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        results = {}
+        for renorm_type in ["text_channel", "global", "channel"]:
+            results[renorm_type] = Bagel._combine_cfg(
+                v_t.clone(),
+                cfg_text_v_t.clone(),
+                None,
+                cfg_text_scale=4.0,
+                cfg_img_scale=1.0,
+                cfg_renorm_type=renorm_type,
+                cfg_renorm_min=0.0,
+            )
+
+        # All results should have positive cosine similarity with each other
+        for a_name, a in results.items():
+            for b_name, b in results.items():
+                cos_sim = torch.nn.functional.cosine_similarity(a.flatten().unsqueeze(0), b.flatten().unsqueeze(0))
+                assert cos_sim > 0.5, (
+                    f"{a_name} and {b_name} should point in similar direction, "
+                    f"but cosine similarity = {cos_sim.item():.4f}"
+                )