[Feature][Diffusion]: CFG parallelism support for Boogu-Image (#6786)

Signed-off-by: nagisa <1434936049@qq.com>
Signed-off-by: nagisa-kun <1434936049@qq.com>

tests/diffusion/distributed/test_cfg_parallel.py

@@ -17,7 +17,12 @@ import torch
 
 from tests.helpers.mark import hardware_marks, hardware_test
 from vllm_omni.diffusion.distributed import parallel_state
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _dispatch_branches, _wrap
+from vllm_omni.diffusion.distributed.cfg_parallel import (
+    CFGParallelMixin,
+    _dispatch_branches,
+    _get_cfg_world_size_or_one,
+    _wrap,
+)
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
     get_classifier_free_guidance_rank,
@@ -393,6 +398,68 @@ def _build_pipeline_on_device(
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_world_size_defaults_only_when_group_is_uninitialized(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(parallel_state, "_CFG", None)
+
+    assert not parallel_state.is_cfg_group_initialized()
+    assert _get_cfg_world_size_or_one() == 1
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_world_size_does_not_swallow_initialized_group_errors(monkeypatch: pytest.MonkeyPatch):
+    class BrokenCfgGroup:
+        @property
+        def world_size(self) -> int:
+            raise AssertionError("CFG group is broken")
+
+    monkeypatch.setattr(parallel_state, "_CFG", BrokenCfgGroup())
+
+    assert parallel_state.is_cfg_group_initialized()
+    with pytest.raises(AssertionError, match="CFG group is broken"):
+        _get_cfg_world_size_or_one()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_defaults_to_sequential_without_initialized_group():
+    pipeline = _build_pipeline_on_device(
+        TestCFGPipeline,
+        hidden_dim=32,
+        model_seed=42,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    positive_kwargs, negative_kwargs = _make_two_branch_inputs(
+        batch_size=1,
+        channels=4,
+        height=4,
+        width=4,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        input_seed=123,
+    )
+
+    with torch.no_grad():
+        positive = pipeline.predict_noise(**positive_kwargs)
+        negative = pipeline.predict_noise(**negative_kwargs)
+        actual = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=True,
+            true_cfg_scale=5.0,
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    expected = negative + 5.0 * (positive - negative)
+    torch.testing.assert_close(actual, expected)
+
+
 @pytest.mark.core_model
 @pytest.mark.diffusion
 @pytest.mark.cpu