diff --git a/docs/user_guide/diffusion/parallelism/cfg_parallel.md b/docs/user_guide/diffusion/parallelism/cfg_parallel.md index 54475440313..1e3b998d8bd 100644 --- a/docs/user_guide/diffusion/parallelism/cfg_parallel.md +++ b/docs/user_guide/diffusion/parallelism/cfg_parallel.md @@ -19,6 +19,12 @@ CFG-Parallel accelerates diffusion models by distributing positive and negative See supported models list in [Supported Models](../../diffusion_features.md#supported-models). +Distilled Cosmos3 checkpoints, including `nvidia/Cosmos3-Super-Image2Video-4Step`, +run without CFG and reject `--cfg-parallel-size` greater than 1 at startup. +Use `--ulysses-degree 2` for two-GPU inference. These checkpoints always use +`guidance_scale=1.0`; other explicitly requested values produce a warning, and +`negative_prompt` does not affect generation. + --- ## Quick Start diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 302409aebaa..9892a72d744 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -9,6 +9,7 @@ from dataclasses import dataclass from types import SimpleNamespace from typing import Any +from unittest.mock import Mock import numpy as np import pytest @@ -376,23 +377,31 @@ def _tokenize( @pytest.mark.parametrize( - ("provided", "value", "default", "is_distilled", "expected"), + ("provided", "value", "default", "is_distilled", "expected", "expected_warning"), [ - (False, 1.0, 7.0, False, 7.0), - (True, 1.0, 7.0, False, 1.0), - (True, 4.5, 7.0, False, 4.5), - (False, 1.0, 7.0, True, 1.0), - (True, 4.5, 7.0, True, 1.0), + (False, 1.0, 7.0, False, 7.0, False), + (True, 1.0, 7.0, False, 1.0, False), + (True, 4.5, 7.0, False, 4.5, False), + (False, 1.0, 7.0, True, 1.0, False), + (True, 1.0, 7.0, True, 1.0, False), + (True, 4.5, 7.0, True, 1.0, True), + (True, 0.0, 7.0, True, 1.0, True), ], ) def test_resolve_guidance_scale( make_cosmos3_pipeline, + monkeypatch: pytest.MonkeyPatch, provided: bool, value: float, default: float, is_distilled: bool, expected: float, + expected_warning: bool, ) -> None: + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + warning_once = Mock() + monkeypatch.setattr(pipeline_cosmos3.logger, "warning_once", warning_once) pipeline = make_cosmos3_pipeline() pipeline.is_distilled_model = is_distilled sp = make_sampling_params( @@ -401,6 +410,12 @@ def test_resolve_guidance_scale( ) assert pipeline._resolve_guidance_scale(sp, default) == expected + if expected_warning: + warning_once.assert_called_once() + assert "overridden to 1.0" in warning_once.call_args.args[0] + assert "negative_prompt does not affect generation" in warning_once.call_args.args[0] + else: + warning_once.assert_not_called() def test_distilled_generation_accepts_t2i_and_i2v(make_cosmos3_pipeline) -> None: @@ -735,6 +750,7 @@ def _make_od_config( custom_pipeline_args={}, model_config=model_config or {}, tf_model_config=tf_model_config, + parallel_config=SimpleNamespace(cfg_parallel_size=1, ulysses_degree=1), ) @@ -770,6 +786,7 @@ def test_pipeline_init_uses_flow_unipc_with_cosmos3_defaults(stub_real_pipeline_ assert pipeline._engine_init_flow_shift == 2.5 +@pytest.mark.parametrize("cfg_parallel_size,ulysses_degree", [(1, 1), (1, 2), (2, 1), (2, 2)]) @pytest.mark.parametrize( ("scheduler_class_name", "expected_distilled"), [ @@ -783,6 +800,8 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file( monkeypatch: pytest.MonkeyPatch, scheduler_class_name: str, expected_distilled: bool, + cfg_parallel_size: int, + ulysses_degree: int, ) -> None: import json @@ -831,6 +850,20 @@ def from_config(cls, config, **kwargs): od_config = _make_od_config(sound_gen=False) od_config.model = str(tmp_path) + od_config.parallel_config.cfg_parallel_size = cfg_parallel_size + od_config.parallel_config.ulysses_degree = ulysses_degree + if expected_distilled and cfg_parallel_size > 1: + monkeypatch.setattr( + pipeline_cosmos3.AutoTokenizer, + "from_pretrained", + lambda *args, **kwargs: pytest.fail("component loading must not start for distilled CFG parallelism"), + ) + with pytest.raises(ValueError, match="Set --cfg-parallel-size 1 and use --ulysses-degree"): + Cosmos3OmniDiffusersPipeline(od_config=od_config) + assert StubFlowMatchScheduler.from_config_calls == [] + assert StubFlowUniPCScheduler.from_config_calls == [] + return + pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config) assert pipeline.is_distilled_model is expected_distilled diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 241eae80577..5977e0eeaac 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -948,6 +948,20 @@ def __init__( model_path = od_config.model local_files_only = os.path.exists(model_path) + # Validate guidance parallelism from checkpoint metadata before loading + # components. Distilled checkpoints have only a conditional branch. + scheduler_config = FlowUniPCMultistepScheduler.load_config( + model_path, + subfolder="scheduler", + local_files_only=local_files_only, + ) + self.is_distilled_model = scheduler_config.get("_class_name") == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS + if self.is_distilled_model and od_config.parallel_config.cfg_parallel_size > 1: + raise ValueError( + "Distilled Cosmos3 checkpoints run without classifier-free guidance. " + "Set --cfg-parallel-size 1 and use --ulysses-degree for multi-GPU inference." + ) + # --- Tokenizer --- self.tokenizer = AutoTokenizer.from_pretrained( model_path, @@ -1010,20 +1024,7 @@ def __init__( logger.info("Cosmos3: session state manager enabled (max_sessions=%d)", mm_max_sessions) # --- Scheduler --- - # Distilled model differs from regular one only by scheduler, - # distilled one uses FlowMatchEulerDiscreteScheduler, while - # regular should use FlowUniPCMultistepScheduler - - scheduler_config = FlowUniPCMultistepScheduler.load_config( - model_path, - subfolder="scheduler", - local_files_only=local_files_only, - ) - - scheduler_class_name = scheduler_config.get("_class_name") - - self.is_distilled_model = False - if scheduler_class_name == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS: + if self.is_distilled_model: fixed_step_config = scheduler_config.get("fixed_step_sampler_config") if not isinstance(fixed_step_config, dict) or fixed_step_config.get("sample_type") != "sde": raise ValueError("Cosmos3 distilled scheduler requires fixed_step_sampler_config.sample_type=sde.") @@ -1035,7 +1036,6 @@ def __init__( stochastic_sampling=True, ) self._scheduler_init_t_list = list(t_list) - self.is_distilled_model = True else: # Preserve compatible solver settings from the checkpoint, but keep # the base shift neutral. The concrete request shift is applied when @@ -1811,6 +1811,11 @@ def _resolve_seed(sp: OmniDiffusionSamplingParams, generator: torch.Generator | def _resolve_guidance_scale(self, sp: OmniDiffusionSamplingParams, default: float) -> float: if self.is_distilled_model: + if sp.guidance_scale_provided and float(sp.guidance_scale) != 1.0: + logger.warning_once( + "Distilled Cosmos3 checkpoints run without classifier-free guidance. " + "The requested guidance_scale is overridden to 1.0; negative_prompt does not affect generation." + ) return 1.0 if sp.guidance_scale_provided: return float(sp.guidance_scale)