diff --git a/fastvideo/configs/pipelines/cosmos3.py b/fastvideo/configs/pipelines/cosmos3.py new file mode 100644 index 000000000..f66a6617f --- /dev/null +++ b/fastvideo/configs/pipelines/cosmos3.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 pipeline configuration. + +Port reference: vllm-omni PR #3454 (https://github.com/vllm-project/vllm-omni/pull/3454) +HEAD ``8536f5b1421f``. Cosmos3 is structurally different from Cosmos 2.5: + +- Dual-stack DiT (UND + GEN) lives entirely inside ``Cosmos3VFMTransformer``. +- No separate text encoder — ``Cosmos3LanguageModel`` is inside the DiT, so + ``text_encoder_configs`` is the empty tuple. The tokenizer is loaded by the + pipeline's ``__init__`` from the checkpoint's ``text_tokenizer/`` subfolder; + that is an instance-level concern, not a config-level concern. +- VAE is ``DistributedAutoencoderKLWan`` (same as Cosmos 2.5), so the + ``Cosmos25VAEConfig`` is reused as-is. +- T2I default ``flow_shift`` is 3.0 (set per-request by ``_set_flow_shift``); + T2V/I2V use the engine-init default of 1.0 baked into this config. +""" +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig +from fastvideo.configs.models.dits.cosmos3 import (Cosmos3ArchConfig, Cosmos3VideoConfig) +from fastvideo.configs.models.encoders import BaseEncoderOutput +from fastvideo.configs.models.vaes import Cosmos25VAEConfig # Cosmos3 reuses the Wan-based VAE +from fastvideo.configs.pipelines.base import PipelineConfig + + +def _identity_preprocess_text(prompt: str) -> str: + return prompt + + +@dataclass +class Cosmos3Config(PipelineConfig): + """Configuration for the Cosmos3 video generation pipeline (T2V/I2V/T2I). + + Tier A scope: architectural skeleton. Real weights from ``nvidia/Cosmos3-Nano`` + are not yet published; this config is shaped to be activatable when they + land without further config changes. + """ + + dit_config: DiTConfig = field(default_factory=lambda: Cosmos3VideoConfig(arch_config=Cosmos3ArchConfig())) + + vae_config: VAEConfig = field(default_factory=Cosmos25VAEConfig) + + # No separate text encoder: Cosmos3LanguageModel lives inside the DiT. + text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=tuple) + + preprocess_text_funcs: tuple[Callable[[str], str], + ...] = field(default_factory=lambda: (_identity_preprocess_text, )) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(default_factory=tuple) + + dit_precision: str = "bf16" + vae_precision: str = "bf16" + text_encoder_precisions: tuple[str, ...] = field(default_factory=tuple) + + embedded_cfg_scale: float = 0.0 + flow_shift: float = 1.0 # T2V/I2V engine-init; T2I uses 3.0 set per-request + + vae_tiling: bool = False + vae_sp: bool = False + + def __post_init__(self): + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 4c7497b1d..3290b3acd 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -20,6 +20,7 @@ from fastvideo.configs.pipelines.cosmos2_5 import ( Cosmos25Config, Cosmos25_14BConfig, ) +from fastvideo.configs.pipelines.cosmos3 import Cosmos3Config from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig from fastvideo.configs.pipelines.gen3c import Gen3CConfig @@ -555,6 +556,22 @@ def _register_configs() -> None: default_preset="gen3c_cosmos_7b", ) + # Cosmos 3 (must register before Cosmos 2.5 and generic Cosmos detectors + # so the cosmos3 path-detection takes precedence) + register_configs( + sampling_param_cls=None, + pipeline_config_cls=Cosmos3Config, + workload_types=(WorkloadType.T2V, ), + hf_model_paths=[ + "nvidia/Cosmos3-Nano", + ], + model_detectors=[ + lambda path: "cosmos3" in path.lower() or "cosmos-3" in path.lower(), + ], + model_family="cosmos3", + default_preset="cosmos3_nano", + ) + # Cosmos 2.5 (2B) register_configs( sampling_param_cls=None,