# 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}" )