From edd91cbdd5fe32c69150d6d5e99ba0d5e6e00c2b Mon Sep 17 00:00:00 2001 From: Dawid Majchrowski Date: Fri, 10 Jul 2026 04:48:04 +0200 Subject: [PATCH] [diffusion] feat: support action output for cosmos3 (#27168) Co-authored-by: Kedi Wu Co-authored-by: Mick --- .../configs/models/dits/cosmos3video.py | 29 +- .../configs/pipeline_configs/base.py | 8 + .../configs/pipeline_configs/cosmos3.py | 11 +- .../configs/pipeline_configs/wan.py | 6 + .../multimodal_gen/configs/sample/cosmos3.py | 40 +- .../configs/sample/sampling_params.py | 69 +- .../entrypoints/diffusion_generator.py | 1 + .../runtime/entrypoints/openai/protocol.py | 3 + .../runtime/entrypoints/openai/utils.py | 11 + .../runtime/entrypoints/openai/video_api.py | 266 +++++-- .../runtime/entrypoints/utils.py | 1 + .../layers/quantization/modelopt_fp8.py | 14 +- .../layers/quantization/modelopt_quant.py | 9 +- .../component_loaders/scheduler_loader.py | 13 +- .../sound_tokenizer_loader.py | 76 ++ .../runtime/loader/transformer_load_utils.py | 2 + .../runtime/managers/scheduler.py | 3 + .../runtime/models/dits/cosmos3video.py | 401 ++++++++-- .../runtime/models/vaes/cosmos3_avae.py | 222 ++++++ .../runtime/pipelines/cosmos3_pipeline.py | 17 +- .../runtime/pipelines_core/schedule_batch.py | 8 + .../stages/model_specific_stages/cosmos3.py | 699 +++++++++++++++--- .../model_specific_stages/cosmos3_action.py | 201 +++++ .../runtime/utils/quantization_utils.py | 12 +- .../test/server/testcase_configs.py | 1 - .../sglang/multimodal_gen/test/test_utils.py | 2 +- .../multimodal_gen/test/unit/test_cosmos3.py | 493 +++++++++++- .../test/unit/test_transformer_quant.py | 32 + .../test/unit/test_wan_pipeline_config.py | 19 + 29 files changed, 2443 insertions(+), 226 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py create mode 100644 python/sglang/multimodal_gen/runtime/models/vaes/cosmos3_avae.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_action.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_wan_pipeline_config.py diff --git a/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py index da7476135..beec2db50 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py @@ -25,23 +25,26 @@ def _build_cosmos3_param_names_mapping() -> dict: norm_moe_gen.weight -> norm_moe_gen.weight time_embedder.linear_{1,2}.weight -> (pass-through) proj_in.weight, proj_out.weight -> (pass-through) + vae2llm.weight -> proj_in.weight (FP8 ckpt alias) + llm2vae.weight -> proj_out.weight (FP8 ckpt alias) GEN patterns (`*_moe_gen`, `add_*`, `to_add_out`, `norm_added_*`) must precede the UND catch-all so the catch-all can't claim GEN keys. `norm.weight` and `lm_head.weight` are inherited from Qwen3-VL - pretraining and not used at inference; audio/action keys are reserved - for a future modality extension — all skipped via empty-string replacement. + pretraining and not used at inference; both are skipped. + Audio and action keys (``audio_proj_*``, ``action_proj_*``, modality + embeds) pass through unchanged. """ return { # Inherited from Qwen3-VL pretraining; unused at diffusion inference. r"^lm_head\.weight$": "", r"^norm\.weight$": "", - # Audio / action modalities — not yet wired; skip to avoid load warnings. - r"^audio_.*$": "", - r"^action_.*$": "", # Top-level norms / embeddings. r"^norm_moe_gen\.(.*)$": r"norm_moe_gen.\1", r"^embed_tokens\.(.*)$": r"language_model.embed_tokens.\1", + # FP8 checkpoint aliases for the latent projection layers. + r"^vae2llm\.(.*)$": r"proj_in.\1", + r"^llm2vae\.(.*)$": r"proj_out.\1", # GEN pathway: per-layer (must run before the UND catch-all below). # Q/K/V merge into MergedColumnParallelLinear to_qkv (concat order: Q, K, V). r"^layers\.(\d+)\.self_attn\.add_q_proj\.(.*)$": ( @@ -151,6 +154,17 @@ class Cosmos3VideoArchConfig(DiTArchConfig): temporal_compression_factor: int = 4 unified_3d_mrope_temporal_modality_margin: int = 15000 + # Audio (sound) modality + sound_gen: bool = False + sound_dim: int = 64 + sound_latent_fps: float = 25.0 + temporal_compression_factor_sound: int = 1 + + # Action modality + action_gen: bool = False + action_dim: int = 64 + num_embodiment_domains: int = 32 + # Timestep embedding timestep_scale: float = 0.001 frequency_embedding_size: int = 256 @@ -167,6 +181,11 @@ class Cosmos3VideoArchConfig(DiTArchConfig): ) reverse_param_names_mapping: dict = field(default_factory=dict) lora_param_names_mapping: dict = field(default_factory=dict) + # FP8 checkpoint quantization_config.ignore uses checkpoint module names; + # translate them to model names so is_layer_excluded matches correctly. + quant_ignore_remap: dict = field( + default_factory=lambda: {"vae2llm": "proj_in", "llm2vae": "proj_out"} + ) def __post_init__(self): super().__post_init__() diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 2b0051f04..9610e8b5f 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -195,6 +195,7 @@ class PipelineConfig: cfg_policy: CFGPolicy = field(default_factory=CFGPolicy) generator_device: str | None = None flow_shift: float | None = None + scheduler_class_override: str | None = None disable_autocast: bool = False # Model configuration @@ -711,6 +712,13 @@ class PipelineConfig: default=PipelineConfig.flow_shift, help="Flow shift parameter", ) + parser.add_argument( + f"--{prefix_with_dot}scheduler-class-override", + type=str, + dest=f"{prefix_with_dot.replace('-', '_')}scheduler_class_override", + default=PipelineConfig.scheduler_class_override, + help="Override the scheduler class from scheduler_config.json.", + ) parser.add_argument( f"--{prefix_with_dot}resolution", type=int, diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py index 41071bf7f..947d6388e 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py @@ -40,7 +40,11 @@ class Cosmos3Config(PipelineConfig): vae_tiling: bool = False vae_sp: bool = False - # Sourced from scheduler_config.json in the checkpoint. + # Cosmos3 reference inference uses FlowUniPC even when the checkpoint + # scheduler_config.json advertises a different scheduler class. + scheduler_class_override: str | None = "FlowUniPCMultistepScheduler" + + # Per-request mode defaults are applied in Cosmos3TimestepPreparationStage. flow_shift: float | None = None precision: str = "bf16" @@ -51,6 +55,11 @@ class Cosmos3Config(PipelineConfig): use_duration_template: bool = True use_system_prompt: bool = False + # Filesystem path to dataset-derived action stats (JSON) for action + # (de)normalization. Set at server launch rather than per request, since it + # names a server-side file. ``None`` disables normalization. + action_stats_path: str | None = None + def __post_init__(self): self.vae_config.arch_config.z_dim = 48 # Encoder is needed for I2V; T2V/T2I never invoke it. diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py index bda7ef9d9..971a7b781 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/wan.py @@ -97,6 +97,12 @@ class WanT2V480PConfig(PipelineConfig): auto_dit_layerwise_offload=True, ) + def get_pos_prompt_embeds(self, batch): + return batch.prompt_embeds[0] + + def get_neg_prompt_embeds(self, batch): + return batch.negative_prompt_embeds[0] + @dataclass class TurboWanT2V480PConfig(WanT2V480PConfig): diff --git a/python/sglang/multimodal_gen/configs/sample/cosmos3.py b/python/sglang/multimodal_gen/configs/sample/cosmos3.py index c67b7e5b9..b1de5d02e 100644 --- a/python/sglang/multimodal_gen/configs/sample/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/sample/cosmos3.py @@ -1,13 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 """Cosmos3 sampling parameters. -A single ``SamplingParams`` class serves T2V, I2V, and T2I — the per-request -mode is dispatched in the pipeline from ``num_frames`` (``== 1`` → T2I) and -``image_path`` (set → I2V). For ``num_frames == 1`` the output ``data_type`` -flips to ``IMAGE`` so the file extension and decode path agree. +A single ``SamplingParams`` class serves T2V, I2V, V2V, T2I, and +action-conditioned variants. Per-request mode is dispatched in the pipeline +from ``num_frames`` (``== 1`` → T2I), ``image_path`` (set → I2V), +``video_path`` (set → V2V), and ``action_mode`` (set → action-conditioned). +For ``num_frames == 1`` the output ``data_type`` flips to ``IMAGE`` +so the file extension and decode path agree. """ from dataclasses import dataclass, field +from typing import Any from sglang.multimodal_gen.configs.sample.sampling_params import ( DataType, @@ -17,7 +20,7 @@ from sglang.multimodal_gen.configs.sample.sampling_params import ( @dataclass class Cosmos3SamplingParams(SamplingParams): - """Cosmos3 sampling parameters (T2V defaults; also used for I2V / T2I).""" + """Cosmos3 sampling parameters (T2V defaults; also used for I2V / V2V / T2I).""" height: int = 720 width: int = 1280 @@ -30,9 +33,16 @@ class Cosmos3SamplingParams(SamplingParams): negative_prompt: str = "" # Optional CFG window — T2I requests typically pass e.g. ``(400, 1000)`` to - # skip guidance at low noise levels. T2V / I2V leave it unset. + # skip guidance at low noise levels. T2V / I2V / V2V leave it unset. guidance_interval: tuple[float, float] | None = None + # V2V conditioning: which latent-frame indices stay locked to the input + # video. ``None`` resolves to ``[0]`` for I2V (single frame) and ``[0, 1]`` + # for V2V. ``condition_video_keep`` controls whether the first or last + # source frames are used when the input video is longer than needed. + condition_frame_indexes: list[int] | None = None + condition_video_keep: str = "first" + supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ (1280, 720), @@ -43,6 +53,24 @@ class Cosmos3SamplingParams(SamplingParams): ] ) + # Action modality (requires action_gen=True in the model checkpoint) + # action_mode: "forward_dynamics" | "policy" | "inverse_dynamics" + action_mode: str | None = None + domain_id: int | None = None + domain_name: str | None = None + raw_action_dim: int | None = None + action_fps: float | None = None + # Action data for forward_dynamics: [T, D] nested list (API) or JSON string + # (CLI via --action). Ignored by the other action modes. + action: Any = None + # Viewpoint phrasing for the structured action caption. + action_view_point: str = "ego_view" + # Optional dataset-derived action stats (JSON) for (de)normalization. When + # set, input actions are normalized and predicted actions de-normalized + # into physical units with ``action_normalization``. + action_stats_path: str | None = None + action_normalization: str = "quantile" + def _set_output_file_name(self) -> None: # The pipeline config's ``task_type=TI2V`` drives ``data_type`` to # VIDEO, but a single-frame request is a T2I and must pick the IMAGE diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 92ce1bd67..3dc24ebc5 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -103,6 +103,9 @@ class SamplingParams: # Image inputs image_path: str | list[str] | None = None + # Video inputs (video-to-video conditioning) + video_path: str | list[str] | None = None + # Text inputs prompt: str | list[str] | None = field( default=None, metadata={"batch_sig_exclude": True} @@ -158,6 +161,9 @@ class SamplingParams: default=None, metadata={"batch_sig_exclude": True} ) # None means all resolutions allowed + # Output audio duration in seconds (models without an audio modality ignore this). + sound_duration: float = 0.0 + # Denoising parameters num_inference_steps: int = None guidance_scale: float = 1.0 @@ -841,6 +847,11 @@ class SamplingParams: type=int, help="Number of frames to generate", ) + add_argument( + "--sound-duration", + type=float, + help="Duration of generated audio in seconds; 0 disables audio output (audio-capable models only)", + ) add_argument( "--height", type=int, @@ -952,9 +963,10 @@ class SamplingParams: "--action", type=str, help=( - "SANA-WM WASD/IJKL action DSL, e.g. " - "'w-80,jw-40,w-40,lw-60,w-100'. Model-specific fields are " - "ignored by other pipelines." + "Action input. SANA-WM uses a WASD/IJKL DSL, e.g. " + "'w-80,jw-40,w-40,lw-60,w-100'. Cosmos3 uses a JSON array of " + "shape [T, D], e.g. '[[0.1, 0.2, ...], ...]'. Model-specific " + "fields are ignored by other pipelines." ), ) add_argument( @@ -978,6 +990,57 @@ class SamplingParams: dest="pitch_limit_deg", help="SANA-WM action DSL absolute pitch clamp in degrees.", ) + add_argument( + "--video-path", + type=str, + nargs="+", + help=( + "Path(s) to input video(s) for video-to-video generation. " + "The first/last frames of the video become the conditioning " + "frames for the generated output." + ), + ) + add_argument( + "--action-mode", + type=str, + dest="action_mode", + help=( + "Cosmos3 action mode: 'forward_dynamics' (predict next frame " + "from action), 'policy' (predict action from frame), or " + "'inverse_dynamics' (predict action from two frames)." + ), + ) + add_argument( + "--domain-id", + type=int, + dest="domain_id", + help="Action embodiment domain ID (integer). Overrides --domain-name.", + ) + add_argument( + "--domain-name", + type=str, + dest="domain_name", + help="Action embodiment domain name (e.g. 'av', 'camera_pose', 'umi').", + ) + add_argument( + "--raw-action-dim", + type=int, + dest="raw_action_dim", + help=( + "Number of active action dimensions to predict; remaining " + "dimensions are zero-padded. Required for 'policy' and " + "'inverse_dynamics' modes." + ), + ) + add_argument( + "--action-fps", + type=float, + dest="action_fps", + help=( + "Frame rate used for action token temporal mRoPE positions. " + "Defaults to the video fps when not set." + ), + ) add_argument( "--moba-config-path", type=str, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index c7885b814..2eaae8706 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -436,6 +436,7 @@ class DiffGenerator: generation_time=generation_time, peak_memory_mb=output_batch.peak_memory_mb, metrics=metrics.to_dict() if metrics else {}, + action=output_batch.action_pred, trajectory_latents=output_batch.trajectory_latents, trajectory_timesteps=output_batch.trajectory_timesteps, rollout_trajectory_data=output_batch.rollout_trajectory_data, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py index 694a9fdc4..c5bbc77da 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py @@ -86,6 +86,7 @@ class VideoResponse(BaseModel): num_outputs: Optional[int] = None peak_memory_mb: Optional[float] = None inference_time_s: Optional[float] = None + action: Optional[Dict[str, Any]] = None class VideoGenerationsRequest(BaseModel): @@ -94,6 +95,8 @@ class VideoGenerationsRequest(BaseModel): prompt: str input_reference: Optional[str] = None reference_url: Optional[str] = None + video_path: Optional[str] = None + video_url: Optional[str] = None model: Optional[str] = None n: Optional[int] = 1 num_outputs_per_prompt: Optional[int] = None diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index 7f6020af1..23ea98e41 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -432,6 +432,17 @@ def add_common_data_to_response( response["id"] = request_id + if result.action_pred is not None: + t = result.action_pred + response["action"] = { + "data": t.tolist(), + "shape": list(t.shape), + "dtype": str(t.dtype).replace("torch.", ""), + "raw_action_dim": result.action_raw_action_dim, + "action_mode": result.action_mode, + "domain_id": result.action_domain_id, + } + return response diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index 4ee93723b..eb5d1bcd4 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -50,11 +50,30 @@ from sglang.srt.observability.trace import extract_trace_headers logger = init_logger(__name__) router = APIRouter(prefix="/v1/videos", tags=["videos"]) +_VIDEO_EXTENSIONS = { + ".avi", + ".gif", + ".m4v", + ".mkv", + ".mov", + ".mp4", + ".mpeg", + ".mpg", + ".webm", +} + def _extra_value(request: VideoGenerationsRequest, name: str) -> Any: return (request.model_extra or {}).get(name) +def _request_value(request: VideoGenerationsRequest, name: str) -> Any: + value = getattr(request, name, None) + if value is not None: + return value + return _extra_value(request, name) + + def _parse_form_extra_value(value: Any) -> Any: if not isinstance(value, str): return value @@ -64,15 +83,146 @@ def _parse_form_extra_value(value: Any) -> Any: return value +def _is_probably_video_source(source: Any) -> bool: + content_type = (getattr(source, "content_type", "") or "").lower() + if content_type.startswith("video/"): + return True + + if isinstance(source, str): + if source.lower().startswith("data:video"): + return True + source_name = source + else: + source_name = getattr(source, "filename", None) + + if not source_name: + return False + source_name = str(source_name).split("?", 1)[0].split("#", 1)[0] + return os.path.splitext(source_name)[1].lower() in _VIDEO_EXTENSIONS + + +def _is_cosmos3_server(server_args) -> bool: + from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config + + return isinstance(server_args.pipeline_config, Cosmos3Config) + + +def _normalize_optional_string(value: Any) -> Any: + if isinstance(value, str) and not value.strip(): + return None + return value + + +def _coerce_optional_int_list(value: Any) -> list[int] | None: + value = _parse_form_extra_value(value) + if value is None: + return None + if isinstance(value, str) and not value.strip(): + return None + if isinstance(value, (list, tuple)): + return [int(item) for item in value] + return [int(value)] + + +def _resolve_video_path(req: VideoGenerationsRequest) -> str | None: + video_path = _request_value(req, "video_path") or _request_value(req, "video_url") + if video_path: + return str(video_path) + + input_reference = _request_value(req, "input_reference") + if _is_probably_video_source(input_reference): + return str(input_reference) + + reference_url = _request_value(req, "reference_url") + if _is_probably_video_source(reference_url): + return str(reference_url) + + return None + + +def _resolve_image_path( + req: VideoGenerationsRequest, video_path: str | None +) -> str | None: + image_path = _request_value(req, "input_reference") + if video_path and image_path == video_path: + return None + if _is_probably_video_source(image_path): + return None + return image_path + + +def _resolve_sound_duration( + req: VideoGenerationsRequest, *, num_frames: int, fps: int +) -> float | None: + generate_sound = _request_value(req, "generate_sound") + sound_duration = _request_value(req, "sound_duration") + + if generate_sound is False: + return 0.0 + if sound_duration is not None: + return float(sound_duration) + if generate_sound is True: + return float(num_frames) / float(fps) + return None + + +def _cosmos3_sampling_param_kwargs( + req: VideoGenerationsRequest, *, num_frames: int, fps: int +) -> Dict[str, Any]: + """Map HTTP/API aliases to Cosmos3SamplingParams field names.""" + kwargs: Dict[str, Any] = {} + + sound_duration = _resolve_sound_duration(req, num_frames=num_frames, fps=fps) + if sound_duration is not None: + kwargs["sound_duration"] = sound_duration + + condition_frame_indexes = _request_value(req, "condition_frame_indexes") + if condition_frame_indexes is None: + condition_frame_indexes = _request_value(req, "condition_frame_indexes_vision") + condition_frame_indexes = _coerce_optional_int_list(condition_frame_indexes) + if condition_frame_indexes is not None: + kwargs["condition_frame_indexes"] = condition_frame_indexes + + for name in ( + "condition_video_keep", + "action_mode", + "domain_id", + "domain_name", + "raw_action_dim", + "action_fps", + "action", + "action_view_point", + "action_normalization", + ): + value = _parse_form_extra_value(_request_value(req, name)) + value = _normalize_optional_string(value) + if value is not None: + kwargs[name] = value + + return kwargs + + def _build_video_sampling_params(request_id: str, request: VideoGenerationsRequest): """Resolve video-specific defaults (fps, seconds → num_frames) then delegate to the shared build_sampling_params.""" + server_args = get_global_server_args() seconds = request.seconds if request.seconds is not None else DEFAULT_VIDEO_SECONDS fps = request.fps if request.fps is not None else DEFAULT_FPS num_frames = request.num_frames if request.num_frames is not None else fps * seconds num_outputs = request.num_outputs_per_prompt if num_outputs is None: num_outputs = request.n or 1 + video_path = _resolve_video_path(request) + image_path = _resolve_image_path(request, video_path) + cosmos3_kwargs = {} + if _is_cosmos3_server(server_args): + cosmos3_kwargs = _cosmos3_sampling_param_kwargs( + request, num_frames=num_frames, fps=fps + ) + if server_args.pipeline_config.action_stats_path is not None: + cosmos3_kwargs["action_stats_path"] = ( + server_args.pipeline_config.action_stats_path + ) return build_sampling_params( request_id, @@ -83,7 +233,8 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque height=request.height, num_frames=num_frames, fps=fps, - image_path=request.input_reference, + image_path=image_path, + video_path=video_path, output_file_name=request_id, seed=request.seed, generator_device=request.generator_device, @@ -110,33 +261,10 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque output_quality=request.output_quality, perf_dump_path=request.perf_dump_path, diffusers_kwargs=request.diffusers_kwargs, + **cosmos3_kwargs, ) -def _reject_unsupported_cosmos3_modes( - req: VideoGenerationsRequest, model_path: str | None -) -> None: - if "cosmos3" not in (model_path or "").lower(): - return - - extra = req.model_extra or {} - if extra.get("generate_sound"): - raise HTTPException( - status_code=400, - detail="Cosmos3 video-with-sound is not supported by SGLang yet; omit generate_sound for video-only generation.", - ) - if extra.get("action_mode"): - raise HTTPException( - status_code=400, - detail="Cosmos3 action generation is not supported by SGLang yet.", - ) - if extra.get("condition_frame_indexes_vision") or extra.get("condition_video_keep"): - raise HTTPException( - status_code=400, - detail="Cosmos3 video-to-video conditioning is not supported by SGLang yet.", - ) - - # extract metadata which http_server needs to know def _video_job_from_sampling( request_id: str, req: VideoGenerationsRequest, sampling: SamplingParams @@ -234,6 +362,9 @@ async def create_video( prompt: Optional[str] = Form(None), input_reference: Optional[UploadFile] = File(None), reference_url: Optional[str] = Form(None), + video_reference: Optional[UploadFile] = File(None), + video_url: Optional[str] = Form(None), + video_path: Optional[str] = Form(None), model: Optional[str] = Form(None), n: Optional[int] = Form(1), num_outputs_per_prompt: Optional[int] = Form(None), @@ -287,24 +418,49 @@ async def create_video( if "multipart/form-data" in content_type: if not prompt: raise HTTPException(status_code=400, detail="prompt is required") - # Validate image input based on model task type + + video_input_path = None image_sources = merge_image_input_list(input_reference, reference_url) + if video_reference is not None: + video_input_path = await _save_first_input_image( + video_reference, + request_id, + uploads_dir, + prefer_remote_source=server_args.input_save_path is None, + ) + elif video_path or video_url: + video_input_path = video_path or video_url + elif input_reference is not None and _is_probably_video_source(input_reference): + video_input_path = await _save_first_input_image( + input_reference, + request_id, + uploads_dir, + prefer_remote_source=server_args.input_save_path is None, + ) + image_sources = merge_image_input_list(reference_url) + elif reference_url and _is_probably_video_source(reference_url): + video_input_path = reference_url + image_sources = merge_image_input_list(input_reference) + + # Validate image input based on model task type if task_type.requires_image_input() and not image_sources: raise HTTPException( status_code=400, detail="input_reference or reference_url is required for image-to-video generation", ) - try: - input_path = await _save_first_input_image( - image_sources, - request_id, - uploads_dir, - prefer_remote_source=server_args.input_save_path is None, - ) - except Exception as e: - raise HTTPException( - status_code=400, detail=f"Failed to process image source: {str(e)}" - ) + input_path = None + if image_sources: + try: + input_path = await _save_first_input_image( + image_sources, + request_id, + uploads_dir, + prefer_remote_source=server_args.input_save_path is None, + ) + except Exception as e: + raise HTTPException( + status_code=400, detail=f"Failed to process image source: {str(e)}" + ) # Parse extra_body JSON (if provided in multipart form) to get fps/num_frames overrides extra_from_form: Dict[str, Any] = {} @@ -322,7 +478,8 @@ async def create_video( pass def form_value(name: str, value: Any) -> Any: - return value if value is not None else extra_from_form.get(name) + selected = value if value is not None else extra_from_form.get(name) + return _parse_form_extra_value(selected) raw_form = await request.form() for key in ( @@ -331,9 +488,19 @@ async def create_video( "use_system_prompt", "use_guardrails", "guardrails", + "video_path", + "video_url", "generate_sound", "sound_duration", + "condition_frame_indexes", "action_mode", + "domain_id", + "domain_name", + "raw_action_dim", + "action_fps", + "action", + "action_view_point", + "action_normalization", "condition_frame_indexes_vision", "condition_video_keep", ): @@ -353,6 +520,8 @@ async def create_video( req = VideoGenerationsRequest( prompt=prompt, input_reference=input_path, + video_path=form_value("video_path", video_input_path), + video_url=form_value("video_url", video_url), model=form_value("model", model), n=form_value("n", n), num_outputs_per_prompt=form_value( @@ -414,8 +583,19 @@ async def create_video( payload.update(flatten_extra_params(extra_json)) flatten_extra_params(payload) # Validate image input based on model task type - has_image_input = payload.get("reference_url") or payload.get( - "input_reference" + if payload.get("video_url") and not payload.get("video_path"): + payload["video_path"] = payload["video_url"] + if _is_probably_video_source(payload.get("reference_url")): + payload.setdefault("video_path", payload.get("reference_url")) + if _is_probably_video_source(payload.get("input_reference")): + payload.setdefault("video_path", payload.get("input_reference")) + + has_image_input = ( + payload.get("reference_url") + and not _is_probably_video_source(payload.get("reference_url")) + ) or ( + payload.get("input_reference") + and not _is_probably_video_source(payload.get("input_reference")) ) if task_type.requires_image_input() and not has_image_input: raise HTTPException( @@ -423,7 +603,9 @@ async def create_video( detail="input_reference or reference_url is required for image-to-video generation", ) # for non-multipart/form-data type - if payload.get("reference_url"): + if payload.get("reference_url") and not _is_probably_video_source( + payload.get("reference_url") + ): try: input_path = await _save_first_input_image( payload.get("reference_url"), @@ -454,8 +636,6 @@ async def create_video( logger.debug(f"Server received from create_video endpoint: req={req}") - _reject_unsupported_cosmos3_modes(req, server_args.model_path) - try: sampling_params = _build_video_sampling_params(request_id, req) except (ValueError, TypeError) as e: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 44a52d4f0..bca5c394c 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -117,6 +117,7 @@ class GenerationResult: samples: Any = None frames: Any = None audio: Any = None + action: Any = None # [T, raw_action_dim] predicted action (policy/inverse_dynamics) prompt: str | None = None size: tuple | None = None # (height, width, num_frames) generation_time: float = 0.0 diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py index ab016215f..e03d7ce63 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py @@ -74,7 +74,11 @@ class ModelOptFp8Config(QuantizationConfig): return [] @classmethod - def from_config(cls, config: Dict[str, Any]) -> ModelOptFp8Config: + def from_config( + cls, + config: Dict[str, Any], + ignore_remap: Optional[Dict[str, str]] = None, + ) -> ModelOptFp8Config: quant_algo = config.get("quant_algo") if quant_algo is None: raise ValueError( @@ -85,6 +89,8 @@ class ModelOptFp8Config(QuantizationConfig): f"ModelOptFp8Config only supports FP8, got quant_algo={quant_algo!r}." ) ignore = config.get("ignore", []) + if ignore_remap and ignore: + ignore = [ignore_remap.get(pattern, pattern) for pattern in ignore] return cls(is_checkpoint_fp8_serialized=True, ignore=ignore) def _is_layer_ignored(self, prefix: str) -> bool: @@ -122,7 +128,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase): """Linear method for ModelOpt static per-tensor FP8 quantization. Uses ``torch._scaled_mm`` (or CUTLASS FP8 GEMM when available) for - the FP8 matrix multiply — the same kernels used by the LLM runtime. + the FP8 matrix multiply - the same kernels used by the LLM runtime. """ def __init__(self, quant_config: ModelOptFp8Config): @@ -177,8 +183,8 @@ class ModelOptFp8LinearMethod(LinearMethodBase): max_w_scale = layer.weight_scale.max() # Transpose weight to [in, out] column-major layout for - # apply_fp8_linear / CUTLASS fp8_scaled_mm. Do NOT call - # .contiguous() — the kernel requires column-major stride. + # apply_fp8_linear / CUTLASS fp8_scaled_mm. Do not call .contiguous(); + # the kernel requires column-major stride. layer.weight = torch.nn.Parameter(layer.weight.data.t(), requires_grad=False) if self.cutlass_fp8_supported: diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py index a8392c976..f539c1ed5 100755 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py @@ -187,7 +187,11 @@ class ModelOptFp8Config(ModelOptQuantConfig): return 89 @classmethod - def from_config(cls, config: Dict[str, Any]) -> ModelOptFp8Config: + def from_config( + cls, + config: Dict[str, Any], + ignore_remap: Optional[Dict[str, str]] = None, + ) -> ModelOptFp8Config: quant_method = config.get("quant_algo") exclude_modules = config.get("ignore") if quant_method is None: @@ -205,6 +209,9 @@ class ModelOptFp8Config(ModelOptQuantConfig): "ModelOptFp8Config only supports static FP8 quantization in SGLang diffusion." ) + if ignore_remap and exclude_modules: + exclude_modules = [ignore_remap.get(p, p) for p in exclude_modules] + return cls( is_checkpoint_fp8_serialized=True, exclude_modules=exclude_modules, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py index a82c1047c..e0de44a5b 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py @@ -23,11 +23,22 @@ class SchedulerLoader(ComponentLoader): """Load the scheduler based on the model path, and inference args.""" config = get_diffusers_component_config(component_path=component_model_path) - class_name = config.pop("_class_name") + checkpoint_class_name = config.pop("_class_name", None) + class_name = ( + getattr(server_args.pipeline_config, "scheduler_class_override", None) + or checkpoint_class_name + ) assert ( class_name is not None ), "Model config does not contain a _class_name attribute. Only diffusers format is supported." + if checkpoint_class_name is not None and class_name != checkpoint_class_name: + logger.info( + "Overriding scheduler class from %s to %s", + checkpoint_class_name, + class_name, + ) + scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name) scheduler = scheduler_cls(**config) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py new file mode 100644 index 000000000..8db15950d --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/sound_tokenizer_loader.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +from safetensors.torch import load_file as safetensors_load_file + +from sglang.multimodal_gen.configs.models import ModelConfig +from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + ComponentLoader, +) +from sglang.multimodal_gen.runtime.loader.utils import ( + _list_safetensors_files, + set_default_torch_dtype, + skip_init_modules, +) +from sglang.multimodal_gen.runtime.models.registry import ModelRegistry +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + get_diffusers_component_config, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + + +class SoundTokenizerLoader(ComponentLoader): + component_names = ["sound_tokenizer"] + expected_library = "diffusers" + + def should_offload( + self, server_args: ServerArgs, model_config: ModelConfig | None = None + ) -> bool: + return server_args.vae_cpu_offload + + def load_customized( + self, component_model_path: str, server_args: ServerArgs, component_name: str + ): + config = get_diffusers_component_config(component_path=component_model_path) + class_name = config.pop("_class_name", None) or self.component_architecture + assert ( + class_name is not None + ), "Sound tokenizer class name must be available from component config." + + server_args.model_paths[component_name] = component_model_path + + try: + precision = server_args.pipeline_config.vae_precision + except AttributeError: + precision = "bf16" + dtype = PRECISION_TO_TYPE[precision] + target_device = self.target_device(self.should_offload(server_args)) + + with set_default_torch_dtype(dtype), skip_init_modules(): + model_cls, _ = ModelRegistry.resolve_model_cls(class_name) + model = model_cls(config).to(target_device) + + safetensors_list = _list_safetensors_files(component_model_path) + assert ( + len(safetensors_list) == 1 + ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}" + loaded = safetensors_load_file(safetensors_list[0]) + incompatible = model.load_state_dict(loaded, strict=False) + missing = getattr(incompatible, "missing_keys", []) + # The tokenizer is decoder-only; the checkpoint's encoder weights are + # expected leftovers, so they're excluded from the load warning. + unexpected = [ + k + for k in getattr(incompatible, "unexpected_keys", []) + if not k.startswith("encoder.") + ] + if missing or unexpected: + logger.warning( + "Loaded sound_tokenizer with missing_keys=%d unexpected_keys=%d", + len(missing), + len(unexpected), + ) + model.eval() + return model diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 5b0f82808..d501cd5ec 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -615,10 +615,12 @@ def _resolve_quant_config( reverse_param_names_mapping_dict = getattr( arch_config, "reverse_param_names_mapping", None ) + quant_ignore_remap_dict = getattr(arch_config, "quant_ignore_remap", None) quant_config = get_quant_config( hf_config, component_model_path, reverse_param_names_mapping=reverse_param_names_mapping_dict, + quant_ignore_remap=quant_ignore_remap_dict, ) quant_config_name = _get_quant_config_name(quant_config) inferred_nvfp4_config = None diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 3562ecd1a..9f631d3f3 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -732,6 +732,9 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag output_batch.audio, start, end, total_items ), audio_sample_rate=output_batch.audio_sample_rate, + action_pred=self._slice_batched_value( + output_batch.action_pred, start, end, total_items + ), trajectory_timesteps=self._slice_batched_value( output_batch.trajectory_timesteps, start, end, total_items ), diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index 9cb0fa060..10c688517 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py @@ -45,6 +45,9 @@ from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping +from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + LayerwiseOffloadableModuleMixin, +) from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.utils import add_prefix @@ -84,6 +87,8 @@ def compute_mrope_position_ids_vision( fps: float | None = None, base_fps: float = 24.0, temporal_compression_factor: int = 4, + base_temporal_compression_factor: int | None = None, + start_frame_offset: int = 0, ) -> tuple[torch.Tensor, int | float]: """Generate 3D mRoPE position IDs for vision tokens. @@ -91,6 +96,12 @@ def compute_mrope_position_ids_vision( per vision segment (Qwen3VL-style). Flattened in T-major order. + When the token rate (``temporal_compression_factor``) differs from the + base-fps rate (``base_temporal_compression_factor``), the two no longer + cancel: action tokens run at frame rate (factor 1) while their temporal + positions are scaled by the video factor. ``base_temporal_compression_factor`` + defaults to ``temporal_compression_factor`` so vision/sound are unchanged. + Returns: (position_ids [3, grid_t * grid_h * grid_w], next_temporal_offset) """ @@ -98,18 +109,28 @@ def compute_mrope_position_ids_vision( if fps_modulation: tps = fps / temporal_compression_factor - base_tps = base_fps / temporal_compression_factor + effective_base_tcf = ( + base_temporal_compression_factor + if base_temporal_compression_factor is not None + else temporal_compression_factor + ) + base_tps = base_fps / effective_base_tcf frame_indices = torch.arange(grid_t, dtype=torch.float32, device=device) t_index = ( - (frame_indices / tps * base_tps + temporal_offset) + ((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset) .view(-1, 1) .expand(-1, grid_h * grid_w) .flatten() ) else: - t_index = torch.arange(grid_t, dtype=torch.long, device=device).view( - -1, 1 - ).expand(-1, grid_h * grid_w).flatten() + int(temporal_offset) + t_index = ( + torch.arange(grid_t, dtype=torch.long, device=device) + .view(-1, 1) + .expand(-1, grid_h * grid_w) + .flatten() + + int(temporal_offset) + + start_frame_offset + ) h_index = ( torch.arange(grid_h, dtype=torch.long, device=device) @@ -135,6 +156,58 @@ def compute_mrope_position_ids_vision( return mrope_ids, next_offset +def compute_mrope_position_ids_sound( + grid_t: int, + temporal_offset: int | float, + sound_latent_fps: float, + device: torch.device, + base_fps: float = 24.0, + temporal_compression_factor_sound: int = 1, +) -> tuple[torch.Tensor, int | float]: + """mRoPE position IDs for sound tokens: a (T, 1, 1) grid.""" + return compute_mrope_position_ids_vision( + grid_t=grid_t, + grid_h=1, + grid_w=1, + temporal_offset=temporal_offset, + device=device, + fps=sound_latent_fps, + base_fps=base_fps, + temporal_compression_factor=temporal_compression_factor_sound, + ) + + +def compute_mrope_position_ids_action( + grid_t: int, + temporal_offset: int | float, + action_fps: float | None, + device: torch.device, + base_fps: float = 24.0, + base_temporal_compression_factor: int = 4, + start_frame_offset: int = 1, +) -> tuple[torch.Tensor, int | float]: + """mRoPE position IDs for action tokens: a (T, 1, 1) grid. + + Action tokens run at frame rate (``temporal_compression_factor=1``) but + their positions are scaled by the video's ``base_temporal_compression_factor`` + so they share the video's temporal coordinate frame. ``start_frame_offset=1`` + (default) shifts them one frame ahead so they align with the video frame + they condition on. + """ + return compute_mrope_position_ids_vision( + grid_t=grid_t, + grid_h=1, + grid_w=1, + temporal_offset=temporal_offset, + device=device, + fps=action_fps, + base_fps=base_fps, + temporal_compression_factor=1, + base_temporal_compression_factor=base_temporal_compression_factor, + start_frame_offset=start_frame_offset, + ) + + # ----------------------------------------------------------------------------- # Qwen3-style RoPE functions # ----------------------------------------------------------------------------- @@ -195,6 +268,44 @@ def _apply_qwen3_qk_norm_rope_split( return _apply_qwen3_rope_from_cache(q, k, cos_sin_cache) +# ----------------------------------------------------------------------------- +# Action domain-aware projection +# ----------------------------------------------------------------------------- + + +class DomainAwareLinear(nn.Module): + """Per-domain linear projection for action conditioning. + + Maintains one weight matrix and bias per embodiment domain via embedding + tables, enabling multi-domain robot action generation from a shared + backbone. + """ + + def __init__(self, input_size: int, output_size: int, num_domains: int) -> None: + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.num_domains = num_domains + self.fc = nn.Embedding(num_domains, output_size * input_size) + self.bias = nn.Embedding(num_domains, output_size) + nn.init.xavier_uniform_( + self.fc.weight.view(num_domains, output_size, input_size) + ) + nn.init.zeros_(self.bias.weight) + + def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: + if domain_id.ndim == 0: + domain_id = domain_id.unsqueeze(0) + domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1) + weight = self.fc(domain_id).view( + domain_id.shape[0], self.input_size, self.output_size + ) + bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size) + if x.ndim == 2: + return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias + return torch.bmm(x, weight) + bias.unsqueeze(1) + + # ----------------------------------------------------------------------------- # Cosmos3 Timestep Embedder # ----------------------------------------------------------------------------- @@ -797,7 +908,7 @@ class Cosmos3LanguageModel(nn.Module): # ----------------------------------------------------------------------------- -class Cosmos3OmniTransformer(CachableDiT): +class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): """Cosmos3 Omni transformer. Dual-pathway architecture: @@ -837,6 +948,8 @@ class Cosmos3OmniTransformer(CachableDiT): self.base_fps = arch.base_fps self.temporal_compression_factor = arch.temporal_compression_factor self.temporal_margin = arch.unified_3d_mrope_temporal_modality_margin + self.sound_latent_fps = arch.sound_latent_fps + self.temporal_compression_factor_sound = arch.temporal_compression_factor_sound self.rms_norm_eps = arch.rms_norm_eps # Ulysses sequence parallelism. When CFG-parallel is also enabled @@ -881,6 +994,40 @@ class Cosmos3OmniTransformer(CachableDiT): prefix="proj_out", ) + self.sound_gen = arch.sound_gen + self.sound_dim = arch.sound_dim + if arch.sound_gen: + self.audio_proj_in = ReplicatedLinear( + self.sound_dim, + self.hidden_size, + bias=True, + quant_config=quant_config, + prefix="audio_proj_in", + ) + self.audio_proj_out = ReplicatedLinear( + self.hidden_size, + self.sound_dim, + bias=True, + quant_config=quant_config, + prefix="audio_proj_out", + ) + self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size)) + + if arch.action_gen: + self.action_dim = arch.action_dim + self.num_embodiment_domains = arch.num_embodiment_domains + self.action_proj_in = DomainAwareLinear( + self.action_dim, + self.hidden_size, + self.num_embodiment_domains, + ) + self.action_proj_out = DomainAwareLinear( + self.hidden_size, + self.action_dim, + self.num_embodiment_domains, + ) + self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size)) + # Timestep embedder self.time_embedder = Cosmos3TimestepEmbedder( hidden_size=self.hidden_size, @@ -920,6 +1067,8 @@ class Cosmos3OmniTransformer(CachableDiT): self.__post_init__() + self.layer_names = ["gen_layers", "language_model.layers"] + def _pad_to_patch_size(self, H: int, W: int) -> tuple[int, int, int, int]: """Compute padded spatial dims aligned to patch_size.""" p = self.latent_patch_size @@ -964,8 +1113,12 @@ class Cosmos3OmniTransformer(CachableDiT): Wp: int, fps: float | None, device: torch.device, + sound_frames: int = 0, + action_frames: int = 0, + action_fps: float | None = None, + action_start_frame_offset: int = 1, ) -> tuple[torch.Tensor, torch.Tensor]: - """Compute mRoPE position IDs for UND text and GEN visual tokens.""" + """Compute mRoPE position IDs for UND text and GEN visual + action + sound tokens.""" B = text_mask.shape[0] S_text = text_mask.shape[1] text_lengths = text_mask.sum(dim=1).long() @@ -978,16 +1131,40 @@ class Cosmos3OmniTransformer(CachableDiT): t_pos, t_offset = compute_mrope_position_ids_text( real_len, temporal_offset=0, device=device ) + media_offset = t_offset + self.temporal_margin v_pos, _ = compute_mrope_position_ids_vision( T, Hp, Wp, - temporal_offset=t_offset + self.temporal_margin, + temporal_offset=media_offset, device=device, fps=effective_fps, base_fps=self.base_fps, temporal_compression_factor=self.temporal_compression_factor, ) + if action_frames > 0: + a_pos, _ = compute_mrope_position_ids_action( + action_frames, + temporal_offset=media_offset, + action_fps=action_fps, + device=device, + base_fps=self.base_fps, + base_temporal_compression_factor=self.temporal_compression_factor, + start_frame_offset=action_start_frame_offset, + ) + pos_dtype = torch.promote_types(v_pos.dtype, a_pos.dtype) + v_pos = torch.cat([v_pos.to(pos_dtype), a_pos.to(pos_dtype)], dim=1) + if sound_frames > 0: + s_pos, _ = compute_mrope_position_ids_sound( + sound_frames, + temporal_offset=media_offset, + sound_latent_fps=self.sound_latent_fps, + device=device, + base_fps=self.base_fps, + temporal_compression_factor_sound=self.temporal_compression_factor_sound, + ) + pos_dtype = torch.promote_types(v_pos.dtype, s_pos.dtype) + v_pos = torch.cat([v_pos.to(pos_dtype), s_pos.to(pos_dtype)], dim=1) if real_len < S_text: t_pos = torch.cat( [ @@ -1002,7 +1179,7 @@ class Cosmos3OmniTransformer(CachableDiT): vis_pos_list.append(v_pos) text_pos_ids = torch.stack(text_pos_list, dim=1).to(device) # [3, B, S_text] - vis_pos_ids = torch.stack(vis_pos_list, dim=1).to(device) # [3, B, S_vis] + vis_pos_ids = torch.stack(vis_pos_list, dim=1).to(device) # [3, B, S_gen] return text_pos_ids, vis_pos_ids @@ -1044,8 +1221,14 @@ class Cosmos3OmniTransformer(CachableDiT): cache_key: str = "default", noisy_frame_mask: torch.Tensor | None = None, max_text_seq_len: int | None = None, + sound_latents: torch.Tensor | None = None, + action_latents: torch.Tensor | None = None, + action_domain_ids: torch.Tensor | None = None, + action_noisy_mask: torch.Tensor | None = None, + action_fps: float | None = None, + action_start_frame_offset: int = 1, **kwargs, - ) -> torch.Tensor: + ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Forward pass for denoising. Args: @@ -1064,9 +1247,20 @@ class Cosmos3OmniTransformer(CachableDiT): ``None`` means every frame is noisy (T2V / T2I). max_text_seq_len: Real text length already computed during tokenization. When omitted it is derived from ``text_mask``. + action_latents: Optional [B, T_action, D_action] noisy action + latents for action generation. + action_domain_ids: [B] embodiment domain IDs (0=no-action default). + action_noisy_mask: [B, T_action, 1] where 1=noisy, 0=conditioned; + controls which action tokens receive the timestep embedding. + ``None`` means all tokens are noisy. + action_fps: Frame rate for action token temporal mRoPE scaling. + Defaults to the video fps when None. + action_start_frame_offset: Temporal offset applied to action + position IDs relative to the video's media_offset (default 1). Returns: - [B, C, T, H, W] velocity prediction + [B, C, T, H, W] velocity prediction, or a tuple + (video_pred, ...) with extra tensors when action/sound are active. """ if text_ids is None or text_mask is None: raise ValueError("Cosmos3 requires text_ids and text_mask to be passed") @@ -1079,9 +1273,31 @@ class Cosmos3OmniTransformer(CachableDiT): text_ids = text_ids[:, :max_text_seq_len] text_mask = text_mask[:, :max_text_seq_len] - # Check if sequence parallelism is enabled + sound_frames = sound_latents.shape[-1] if sound_latents is not None else 0 + + action_frames = 0 + if action_latents is not None: + if self.sp_size > 1: + raise NotImplementedError( + "Cosmos3 action generation does not support sequence parallelism yet" + ) + action_frames = action_latents.shape[1] + if action_domain_ids is None: + action_domain_ids = torch.zeros( + action_latents.shape[0], + dtype=torch.long, + device=action_latents.device, + ) + + extra_frames = action_frames + sound_frames sequence_shard_enabled = self.sp_size > 1 + # Add timestep embedding (computed in float32 for numerical stability, then cast back) + time_embed = self.time_embedder(timestep.float()) + time_embed = time_embed.to( + hidden_states.dtype + ) # Cast to match hidden_gen dtype + # Patchify and project to hidden dim hidden_gen, _ = self.proj_in(self.patchify(hidden_states, T, H, W)) seq_len_orig = hidden_gen.shape[1] @@ -1099,42 +1315,90 @@ class Cosmos3OmniTransformer(CachableDiT): .to(hidden_gen.dtype) ) - # Shard sequence across GPUs if SP enabled - if sequence_shard_enabled: - if seq_len_orig % self.sp_size != 0: - seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size) - pad = torch.zeros( - (batch_size, seq_shard_pad, hidden_gen.shape[2]), - dtype=hidden_gen.dtype, - device=hidden_gen.device, - ) - hidden_gen = torch.cat([hidden_gen, pad], dim=1) - if token_noisy_mask is not None: - mask_pad = torch.zeros( - (batch_size, seq_shard_pad, 1), - dtype=token_noisy_mask.dtype, - device=token_noisy_mask.device, + if extra_frames == 0: + # Video-only: shard the visual tokens, then add the timestep + # embedding on the local shard. + if sequence_shard_enabled: + if seq_len_orig % self.sp_size != 0: + seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size) + pad = torch.zeros( + (batch_size, seq_shard_pad, hidden_gen.shape[2]), + dtype=hidden_gen.dtype, + device=hidden_gen.device, ) - token_noisy_mask = torch.cat([token_noisy_mask, mask_pad], dim=1) - local_seq_len = hidden_gen.shape[1] // self.sp_size - hidden_gen = hidden_gen.view( - batch_size, self.sp_size, local_seq_len, hidden_gen.shape[2] - ) - hidden_gen = hidden_gen[:, self.sp_rank, :, :] + hidden_gen = torch.cat([hidden_gen, pad], dim=1) + if token_noisy_mask is not None: + mask_pad = torch.zeros( + (batch_size, seq_shard_pad, 1), + dtype=token_noisy_mask.dtype, + device=token_noisy_mask.device, + ) + token_noisy_mask = torch.cat( + [token_noisy_mask, mask_pad], dim=1 + ) + local_seq_len = hidden_gen.shape[1] // self.sp_size + hidden_gen = hidden_gen.view( + batch_size, self.sp_size, local_seq_len, hidden_gen.shape[2] + ) + hidden_gen = hidden_gen[:, self.sp_rank, :, :] + if token_noisy_mask is not None: + token_noisy_mask = token_noisy_mask.view( + batch_size, self.sp_size, local_seq_len, 1 + )[:, self.sp_rank, :, :] if token_noisy_mask is not None: - token_noisy_mask = token_noisy_mask.view( - batch_size, self.sp_size, local_seq_len, 1 - )[:, self.sp_rank, :, :] - - # Add timestep embedding (computed in float32 for numerical stability, then cast back) - time_embed = self.time_embedder(timestep.float()) - time_embed = time_embed.to( - hidden_states.dtype - ) # Cast to match hidden_gen dtype - if token_noisy_mask is not None: - hidden_gen = hidden_gen + time_embed.unsqueeze(1) * token_noisy_mask + hidden_gen = hidden_gen + time_embed.unsqueeze(1) * token_noisy_mask + else: + hidden_gen = hidden_gen + time_embed.unsqueeze(1) else: - hidden_gen = hidden_gen + time_embed.unsqueeze(1) + # Multi-modal: assemble the full GEN sequence + # (video[, action][, sound]) with timestep embeddings, then shard + # the combined stream so sequence parallelism splits every modality + # evenly. The per-modality output heads run after the post-loop + # all-gather reassembles the sequence. + if token_noisy_mask is not None: + hidden_gen = hidden_gen + time_embed.unsqueeze(1) * token_noisy_mask + else: + hidden_gen = hidden_gen + time_embed.unsqueeze(1) + + if action_latents is not None: + hidden_action = self.action_proj_in( + action_latents.to(hidden_gen.dtype), action_domain_ids + ) + hidden_action = hidden_action + self.action_modality_embed.to( + hidden_action.dtype + ) + if action_noisy_mask is None: + hidden_action = hidden_action + time_embed.unsqueeze(1) + else: + hidden_action = hidden_action + time_embed.unsqueeze( + 1 + ) * action_noisy_mask.to(hidden_action.dtype) + hidden_gen = torch.cat([hidden_gen, hidden_action], dim=1) + + if sound_latents is not None: + packed_sound = sound_latents.permute(0, 2, 1).to(hidden_gen.dtype) + hidden_sound, _ = self.audio_proj_in(packed_sound) + hidden_sound = hidden_sound + self.audio_modality_embed.to( + hidden_sound.dtype + ) + hidden_sound = hidden_sound + time_embed.unsqueeze(1) + hidden_gen = torch.cat([hidden_gen, hidden_sound], dim=1) + + seq_len_orig = hidden_gen.shape[1] + if sequence_shard_enabled: + if seq_len_orig % self.sp_size != 0: + seq_shard_pad = self.sp_size - (seq_len_orig % self.sp_size) + pad = torch.zeros( + (batch_size, seq_shard_pad, hidden_gen.shape[2]), + dtype=hidden_gen.dtype, + device=hidden_gen.device, + ) + hidden_gen = torch.cat([hidden_gen, pad], dim=1) + local_seq_len = hidden_gen.shape[1] // self.sp_size + hidden_gen = hidden_gen.view( + batch_size, self.sp_size, local_seq_len, hidden_gen.shape[2] + ) + hidden_gen = hidden_gen[:, self.sp_rank, :, :] self._ensure_cache_dicts() @@ -1145,7 +1409,16 @@ class Cosmos3OmniTransformer(CachableDiT): or cache_key not in self.cached_gen_rope_inputs ): text_pos_ids, vis_pos_ids = self._compute_rope_position_ids( - text_mask, T, Hp, Wp, fps, hidden_states.device + text_mask, + T, + Hp, + Wp, + fps, + hidden_states.device, + sound_frames=sound_frames, + action_frames=action_frames, + action_fps=action_fps if action_fps is not None else fps, + action_start_frame_offset=action_start_frame_offset, ) # UND K/V cache is kept FULL on all ranks (not sharded). Text # sequence is short, so memory impact is minimal, and the GEN @@ -1193,14 +1466,42 @@ class Cosmos3OmniTransformer(CachableDiT): # this cuts the post-loop SP collective bandwidth ~21x. hidden_gen = hidden_gen + residual hidden_gen = self.norm_moe_gen(hidden_gen) - output, _ = self.proj_out(hidden_gen) + if extra_frames == 0: + # Video-only: project on the local shard and gather the (much + # smaller) patch-space output. With patch_latent_dim ~= + # hidden_size / 21 for cosmos3, this cuts the post-loop SP + # collective bandwidth ~21x. + output, _ = self.proj_out(hidden_gen) + if sequence_shard_enabled: + output = sequence_model_parallel_all_gather(output, dim=1) + if seq_shard_pad > 0: + output = output[:, :seq_len_orig, :] + return self.unpatchify(output, T, H, W) + + # Multi-modal: gather the full GEN hidden and drop shard padding, then + # split per modality so each output head sees its own contiguous tokens. if sequence_shard_enabled: - output = sequence_model_parallel_all_gather(output, dim=1) + hidden_gen = sequence_model_parallel_all_gather(hidden_gen, dim=1) if seq_shard_pad > 0: - output = output[:, :seq_len_orig, :] + hidden_gen = hidden_gen[:, :seq_len_orig, :] - return self.unpatchify(output, T, H, W) + s_video = seq_len_orig - extra_frames + output, _ = self.proj_out(hidden_gen[:, :s_video, :]) + video_pred = self.unpatchify(output, T, H, W) + + extra_outputs: list[torch.Tensor] = [] + idx = s_video + if action_frames > 0: + action_hidden = hidden_gen[:, idx : idx + action_frames, :] + extra_outputs.append(self.action_proj_out(action_hidden, action_domain_ids)) + idx += action_frames + if sound_frames > 0: + sound_hidden = hidden_gen[:, idx:, :] + sound_output, _ = self.audio_proj_out(sound_hidden) + extra_outputs.append(sound_output.permute(0, 2, 1).contiguous()) + + return (video_pred, *extra_outputs) def preprocess_loaded_state_dict( self, iterator: Iterable[tuple[str, torch.Tensor]] diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/cosmos3_avae.py b/python/sglang/multimodal_gen/runtime/models/vaes/cosmos3_avae.py new file mode 100644 index 000000000..b797ffd6a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/vaes/cosmos3_avae.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Decoder-only audio tokenizer for the Cosmos3 sound modality.""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +from torch import nn +from torch.nn.utils import weight_norm + + +class Snake1d(nn.Module): + def __init__(self, hidden_dim: int, logscale: bool = True) -> None: + super().__init__() + self.alpha = nn.Parameter(torch.zeros(1, hidden_dim, 1)) + self.beta = nn.Parameter(torch.zeros(1, hidden_dim, 1)) + self.logscale = logscale + + def forward(self, x: torch.Tensor) -> torch.Tensor: + shape = x.shape + alpha = torch.exp(self.alpha) if self.logscale else self.alpha + beta = torch.exp(self.beta) if self.logscale else self.beta + x = x.reshape(shape[0], shape[1], -1) + x = x + (beta + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + return x.reshape(shape) + + +class OobleckResidualUnit(nn.Module): + def __init__(self, dim: int, dilation: int = 1) -> None: + super().__init__() + pad = ((7 - 1) * dilation) // 2 + self.snake1 = Snake1d(dim) + self.conv1 = weight_norm( + nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad) + ) + self.snake2 = Snake1d(dim) + self.conv2 = weight_norm(nn.Conv1d(dim, dim, kernel_size=1)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y = self.conv1(self.snake1(x)) + y = self.conv2(self.snake2(y)) + pad = (x.shape[-1] - y.shape[-1]) // 2 + if pad > 0: + x = x[..., pad:-pad] + return x + y + + +class OobleckDecoderBlock(nn.Module): + def __init__( + self, + input_dim: int, + output_dim: int, + stride: int, + output_padding: int, + ) -> None: + super().__init__() + self.snake1 = Snake1d(input_dim) + self.conv_t1 = weight_norm( + nn.ConvTranspose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + output_padding=output_padding, + ) + ) + self.res_unit1 = OobleckResidualUnit(output_dim, dilation=1) + self.res_unit2 = OobleckResidualUnit(output_dim, dilation=3) + self.res_unit3 = OobleckResidualUnit(output_dim, dilation=9) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.snake1(x) + x = self.conv_t1(x) + x = self.res_unit1(x) + x = self.res_unit2(x) + return self.res_unit3(x) + + +class OobleckDecoder(nn.Module): + def __init__( + self, + channels: int, + input_channels: int, + audio_channels: int, + upsampling_ratios: list[int], + channel_multiples: list[int], + ) -> None: + super().__init__() + strides = upsampling_ratios + mults = [1] + list(channel_multiples) + + self.conv1 = weight_norm( + nn.Conv1d(input_channels, channels * mults[-1], kernel_size=7, padding=3) + ) + + blocks = [] + for i, stride in enumerate(strides): + blocks.append( + OobleckDecoderBlock( + input_dim=channels * mults[len(strides) - i], + output_dim=channels * mults[len(strides) - i - 1], + stride=stride, + output_padding=stride % 2, + ) + ) + self.block = nn.ModuleList(blocks) + self.snake1 = Snake1d(channels) + self.conv2 = weight_norm( + nn.Conv1d(channels, audio_channels, kernel_size=7, padding=3, bias=False) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv1(x) + for layer in self.block: + x = layer(x) + x = self.snake1(x) + return self.conv2(x) + + +def _cfg(config: dict[str, Any], *keys: str, default: Any = None) -> Any: + for k in keys: + v = config.get(k) + if v is not None: + return v + return default + + +class Cosmos3AVAEAudioTokenizer(nn.Module): + """Cosmos3 audio tokenizer: latents → waveform via an Oobleck decoder stack.""" + + def __init__(self, config: dict[str, Any]) -> None: + super().__init__() + self.sample_rate = int( + _cfg(config, "sampling_rate", "sample_rate", default=48000) + ) + self.audio_channels = int( + _cfg( + config, + "dec_out_channels", + "audio_channels", + default=2 if bool(config.get("stereo", True)) else 1, + ) + ) + self.latent_channels = int( + _cfg(config, "vocoder_input_dim", "io_channels", "latent_ch", default=64) + ) + dec_strides = [ + int(s) for s in _cfg(config, "dec_strides", default=[2, 4, 5, 6, 8]) + ] + self.hop_size = int( + _cfg( + config, + "hop_size", + default=math.prod(dec_strides) if dec_strides else 1920, + ) + ) + stride_product = math.prod(dec_strides) + if stride_product != self.hop_size: + raise ValueError( + "Cosmos3 AVAE dec_strides product must equal hop_size: " + f"product={stride_product}, hop_size={self.hop_size}." + ) + + norm = str(_cfg(config, "normalization_type", default="none")) + if bool(_cfg(config, "normalize_latents", default=False)) and norm == "none": + norm = "tanh" + self.normalization_type = norm + self.tanh_input_scale = float(_cfg(config, "tanh_input_scale", default=1.5)) + self.tanh_output_scale = float(_cfg(config, "tanh_output_scale", default=3.5)) + self.tanh_clamp = float(_cfg(config, "tanh_clamp", default=0.995)) + + self.decoder = OobleckDecoder( + channels=int(_cfg(config, "dec_dim", default=320)), + input_channels=self.latent_channels, + audio_channels=self.audio_channels, + upsampling_ratios=list(reversed(dec_strides)), + channel_multiples=list( + _cfg(config, "dec_c_mults", default=[1, 2, 4, 8, 16]) + ), + ) + + @property + def temporal_compression_factor(self) -> int: + return self.hop_size + + def get_latent_num_samples(self, num_audio_samples: int) -> int: + return int(num_audio_samples) // self.hop_size + + def get_audio_num_samples(self, num_latent_samples: int) -> int: + return int(num_latent_samples) * self.hop_size + + def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor: + if self.normalization_type == "tanh": + in_dtype = latent.dtype + x = torch.clamp( + latent.float() / self.tanh_output_scale, + -self.tanh_clamp, + self.tanh_clamp, + ) + return (torch.atanh(x) * self.tanh_input_scale).to(in_dtype) + if self.normalization_type != "none": + raise ValueError( + f"Unsupported AVAE normalization_type={self.normalization_type!r}." + ) + return latent + + @torch.no_grad() + def decode(self, latent: torch.Tensor) -> torch.Tensor: + squeeze = latent.ndim == 2 + if squeeze: + latent = latent.unsqueeze(0) + decoder_dtype = next(self.decoder.parameters()).dtype + decoder_device = next(self.decoder.parameters()).device + z = self._denormalize_latent(latent.to(decoder_device)).to(decoder_dtype) + audio = self.decoder(z).clamp(-1.0, 1.0).to(latent.dtype) + return audio.squeeze(0) if squeeze else audio + + +EntryClass = Cosmos3AVAEAudioTokenizer diff --git a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py index 2b47d2276..6e65aa449 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py @@ -41,8 +41,18 @@ class Cosmos3Pipeline(ComposedPipelineBase): "vae", "transformer", "scheduler", + "sound_tokenizer", ] + def load_modules(self, server_args, loaded_modules=None): + # Visual-only Cosmos3 checkpoints ship no sound_tokenizer; require it + # only when the checkpoint actually provides one. + if "sound_tokenizer" not in self._load_config(): + self._required_config_modules = [ + m for m in self._required_config_modules if m != "sound_tokenizer" + ] + return super().load_modules(server_args, loaded_modules) + def create_pipeline_stages(self, server_args: ServerArgs) -> None: """Create Cosmos3 pipeline stages. @@ -58,6 +68,7 @@ class Cosmos3Pipeline(ComposedPipelineBase): vae = self.get_module("vae") transformer = self.get_module("transformer") scheduler = self.get_module("scheduler") + sound_tokenizer = self.get_module("sound_tokenizer") guardrails_disabled = ( os.environ.get("SGLANG_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" @@ -86,7 +97,11 @@ class Cosmos3Pipeline(ComposedPipelineBase): self.add_stage(Cosmos3LatentPreparationStage(vae, transformer)) self.add_stage(Cosmos3TimestepPreparationStage(scheduler)) self.add_stage(Cosmos3DenoisingStage(transformer, scheduler, server_args)) - self.add_stage(Cosmos3DecodingStage(vae, guardrails=guardrails_on)) + self.add_stage( + Cosmos3DecodingStage( + vae, guardrails=guardrails_on, sound_tokenizer=sound_tokenizer + ) + ) logger.info( "Cosmos3 pipeline stages created successfully (guardrails=%s)", diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index d4d1247e6..066577df3 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -85,6 +85,7 @@ class Req: vae_image: torch.Tensor | PIL.Image.Image | None = None pixel_values: torch.Tensor | PIL.Image.Image | None = None preprocessed_image: torch.Tensor | None = None + preprocessed_video: torch.Tensor | None = None output_file_ext: str | None = None # Primary encoder embeddings @@ -136,6 +137,9 @@ class Req: # Audio Parameters generate_audio: bool = True + # Action Latents (Cosmos3 action-conditioned generation) + action_latents: torch.Tensor | None = None + raw_latent_shape: torch.Tensor | None = None did_sp_shard_latents: bool = False sp_video_start_frame: int = 0 @@ -415,6 +419,10 @@ class OutputBatch: raw_frame_metadata: dict[str, Any] | None = None audio: torch.Tensor | None = None audio_sample_rate: int | None = None + action_pred: torch.Tensor | None = None + action_mode: str | None = None + action_domain_id: int | None = None + action_raw_action_dim: int | None = None trajectory_timesteps: torch.Tensor | None = None trajectory_latents: torch.Tensor | None = None rollout_trajectory_data: RolloutTrajectoryData | None = None diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py index 1d21f5e26..1492e41f3 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py @@ -4,10 +4,13 @@ prep, denoising, decode. Cosmos3 has no separate text encoder — text is tokenized with Qwen2's chat template and embedded inside the transformer's UND pathway. The same -``Cosmos3Pipeline`` serves T2V, I2V, and T2I; mode is dispatched per-request -from ``batch.data_type`` and the presence of ``batch.preprocessed_image``. +``Cosmos3Pipeline`` serves T2V, I2V, V2V, and T2I; mode is dispatched +per-request from ``batch.data_type`` and the presence of +``batch.preprocessed_image`` / ``batch.preprocessed_video``. """ +import copy +import json from typing import Any import numpy as np @@ -27,11 +30,23 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_sp_world_size, ) from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.models.vision_utils import load_video from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, StageParallelismType, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_action import ( + ACTION_MODE_FORWARD_DYNAMICS, + ACTION_MODE_INVERSE_DYNAMICS, + ACTION_MODES, + EMBODIMENT_TO_DOMAIN_ID, + build_action_prompt, + denormalize_action, + get_raw_action_dim, + load_action_stats, + normalize_action, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( StageValidators as V, ) @@ -57,12 +72,40 @@ COSMOS3_IMAGE_SYSTEM_PROMPT = ( "You are a helpful assistant who will generate images from a given prompt." ) +# Per-mode flow-shift defaults, applied only when the request and pipeline +# config leave flow_shift unset. +COSMOS3_T2I_FLOW_SHIFT = 3.0 +COSMOS3_I2V_FLOW_SHIFT = 10.0 +COSMOS3_T2V_FLOW_SHIFT = 10.0 +COSMOS3_V2V_FLOW_SHIFT = 10.0 +COSMOS3_ACTION_FLOW_SHIFT = 10.0 + + +def _resize_crop_pil( + image: PIL.Image.Image, target_w: int, target_h: int +) -> PIL.Image.Image: + """Aspect-preserving resize then center-crop to ``target_w x target_h``.""" + scale = max(target_w / image.width, target_h / image.height) + resize_w = int(np.ceil(scale * image.width)) + resize_h = int(np.ceil(scale * image.height)) + image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS) + left = (resize_w - target_w) // 2 + top = (resize_h - target_h) // 2 + return image.crop((left, top, left + target_w, top + target_h)) + + +def _pil_to_normalized_tensor(image: PIL.Image.Image) -> torch.Tensor: + """PIL RGB → ``[3, H, W]`` float32 tensor in ``[-1, 1]``.""" + arr = np.asarray(image, dtype=np.float32) / 127.5 - 1.0 + return torch.from_numpy(arr).permute(2, 0, 1).contiguous() + class Cosmos3ImagePreprocessStage(PipelineStage): - """Load, aspect-resize, and center-crop the I2V conditioning image. + """Load, aspect-resize, and center-crop the conditioning input. - No-op when the request has no image (T2V / T2I). The output is a - ``[1, 3, H, W]`` tensor in ``[-1, 1]`` written to ``batch.preprocessed_image``. + For I2V: writes ``[1, 3, H, W]`` to ``batch.preprocessed_image``. + For V2V: writes ``[1, 3, T_in, H, W]`` to ``batch.preprocessed_video``. + No-op for T2V / T2I. """ parallelism_type = StageParallelismType.REPLICATED @@ -74,26 +117,89 @@ class Cosmos3ImagePreprocessStage(PipelineStage): image_path = batch.image_path if isinstance(image_path, list): image_path = image_path[0] if image_path else None - if not isinstance(image_path, str) or not image_path: + video_path = batch.video_path + if isinstance(video_path, list): + video_path = video_path[0] if video_path else None + + if image_path and video_path: + raise ValueError( + "Cosmos3 accepts either --image-path (I2V) or --video-path " + "(V2V), not both" + ) + + target_h, target_w = batch.height, batch.width + + if isinstance(image_path, str) and image_path: + image = PIL.Image.open(image_path).convert("RGB") + image = _resize_crop_pil(image, target_w, target_h) + batch.preprocessed_image = _pil_to_normalized_tensor(image).unsqueeze(0) + self.log_info(f"Preprocessed conditioning image to {target_w}x{target_h}") return batch - image = PIL.Image.open(image_path).convert("RGB") - target_h, target_w = batch.height, batch.width - scale = max(target_w / image.width, target_h / image.height) - resize_w = int(np.ceil(scale * image.width)) - resize_h = int(np.ceil(scale * image.height)) - image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS) - left = (resize_w - target_w) // 2 - top = (resize_h - target_h) // 2 - image = image.crop((left, top, left + target_w, top + target_h)) + if isinstance(video_path, str) and video_path: + frames = load_video(video_path) + if not frames: + raise ValueError(f"No frames decoded from video: {video_path!r}") - arr = np.asarray(image, dtype=np.float32) / 127.5 - 1.0 - tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).contiguous() + keep = ( + getattr(batch.sampling_params, "condition_video_keep", "first") + or "first" + ) + if keep not in ("first", "last"): + raise ValueError( + f"condition_video_keep must be 'first' or 'last', got {keep!r}" + ) + cond_indexes = self._resolve_condition_indexes(batch) + # Encode the full output-length video so that the latent positions + # we lock match what the decoder will reconstruct at those frame + # indices. Encoding only the first ``max_idx*4+1`` frames produces + # an out-of-distribution latent for the locked slots and decodes + # to noise. + num_source_frames = max(cond_indexes) * 4 + 1 + num_target_frames = batch.num_frames + if keep == "last": + frames = frames[-num_source_frames:] + else: + frames = frames[:num_source_frames] + if len(frames) < num_source_frames: + frames = frames + [frames[-1]] * (num_source_frames - len(frames)) + if len(frames) < num_target_frames: + frames = frames + [frames[-1]] * (num_target_frames - len(frames)) + + processed = [ + _pil_to_normalized_tensor( + _resize_crop_pil(f.convert("RGB"), target_w, target_h) + ) + for f in frames + ] + video_tensor = torch.stack(processed, dim=1).unsqueeze(0).contiguous() + batch.preprocessed_video = video_tensor + self.log_info( + f"Preprocessed conditioning video to " + f"{video_tensor.shape[2]}x{target_h}x{target_w} " + f"(keep={keep}, source frames={num_source_frames}, padded to {num_target_frames})" + ) - batch.preprocessed_image = tensor - self.log_info(f"Preprocessed conditioning image to {target_w}x{target_h}") return batch + @staticmethod + def _resolve_condition_indexes(batch: Req) -> list[int]: + """Resolve condition_frame_indexes for V2V (default ``[0, 1]``). + + Inverse-dynamics action mode conditions on the whole input video, so + every latent frame is locked. + """ + if ( + getattr(batch.sampling_params, "action_mode", None) + == ACTION_MODE_INVERSE_DYNAMICS + ): + num_latent_frames = (batch.num_frames - 1) // 4 + 1 + return list(range(num_latent_frames)) + cond_indexes = getattr(batch.sampling_params, "condition_frame_indexes", None) + if not cond_indexes: + return [0, 1] + return sorted(set(int(i) for i in cond_indexes)) + class Cosmos3TokenizationStage(PipelineStage): """Tokenization stage for Cosmos3. @@ -207,6 +313,21 @@ class Cosmos3TokenizationStage(PipelineStage): COSMOS3_IMAGE_SYSTEM_PROMPT if is_image_gen else COSMOS3_VIDEO_SYSTEM_PROMPT ) + # Action mode uses a structured JSON caption with neither a system + # prompt nor the duration suffix. + if getattr(batch.sampling_params, "action_mode", None) is not None: + prompt = build_action_prompt( + prompt, + getattr(batch.sampling_params, "action_view_point", "ego_view"), + num_frames, + fps, + batch.height, + batch.width, + ) + use_system_prompt = False + use_duration_template = False + self.log_info(f"Action prompt: {prompt}") + # Apply duration template if enabled (no temporal concept for T2I). if use_duration_template and not is_image_gen and num_frames > 1: duration = num_frames / fps @@ -250,10 +371,13 @@ class Cosmos3TokenizationStage(PipelineStage): class Cosmos3LatentPreparationStage(PipelineStage): """Initialize the noisy latent for Cosmos3. - T2V / T2I produce pure Gaussian noise. I2V VAE-encodes the conditioning - image, replaces frame 0 of the latent with the encoded image, and stashes - a per-frame velocity mask plus the clean frame-0 latent for the denoiser - to re-inject after each scheduler step. + T2V / T2I produce pure Gaussian noise. I2V / V2V VAE-encode the + conditioning input, write the resulting latents at the conditioned + frame indexes, and stash a per-frame velocity mask plus the full + condition latent so the denoiser can re-blend after each scheduler step. + I2V is the special case of conditioning at frame ``[0]`` with the image + expanded across the temporal axis; V2V conditions at ``[0, 1]`` (or a + user-supplied list) with frames from the input video. """ parallelism_type = StageParallelismType.REPLICATED @@ -316,26 +440,55 @@ class Cosmos3LatentPreparationStage(PipelineStage): noise = torch.randn(shape, generator=generator, device=device, dtype=dtype) - is_i2v = ( - batch.preprocessed_image is not None and batch.data_type == DataType.VIDEO - ) + is_video_gen = batch.data_type == DataType.VIDEO + has_image_cond = batch.preprocessed_image is not None and is_video_gen + has_video_cond = batch.preprocessed_video is not None and is_video_gen - if is_i2v: + if has_image_cond or has_video_cond: vae_dtype = next(self.vae.parameters()).dtype - pixel_video = batch.preprocessed_image.unsqueeze(2).to( - device=device, dtype=vae_dtype - ) - with torch.no_grad(): - cond_latent = self._vae_encode(pixel_video).to(dtype) + if has_video_cond: + pixel_input = batch.preprocessed_video.to( + device=device, dtype=vae_dtype + ) + cond_indexes = Cosmos3ImagePreprocessStage._resolve_condition_indexes( + batch + ) + else: + pixel_input = batch.preprocessed_image.unsqueeze(2).to( + device=device, dtype=vae_dtype + ) + cond_indexes = [0] + + with torch.no_grad(): + cond_latent = self._vae_encode(pixel_input).to(dtype) + + max_idx = max(cond_indexes) + if max_idx >= num_latent_frames: + raise ValueError( + f"condition_frame_indexes={cond_indexes} exceeds the " + f"latent frame count {num_latent_frames} for " + f"num_frames={batch.num_frames}" + ) + + condition_latents = torch.zeros_like(noise) condition_mask = torch.zeros( 1, 1, num_latent_frames, 1, 1, device=device, dtype=dtype ) - condition_mask[:, :, 0, :, :] = 1.0 - latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise - batch.image_latent = cond_latent[:, :, 0:1, :, :].clone() + for idx in cond_indexes: + src = min(idx, cond_latent.shape[2] - 1) + condition_latents[:, :, idx, :, :] = cond_latent[:, :, src, :, :] + condition_mask[:, :, idx, :, :] = 1.0 + + latents = ( + condition_mask * condition_latents + (1.0 - condition_mask) * noise + ) + batch.extra["condition_latents"] = condition_latents batch.extra["velocity_mask"] = 1.0 - condition_mask - self.log_info("Prepared I2V latents with frame-0 conditioning") + mode = "V2V" if has_video_cond else "I2V" + self.log_info( + f"Prepared {mode} latents with conditioning at frames {cond_indexes}" + ) else: latents = noise @@ -347,8 +500,172 @@ class Cosmos3LatentPreparationStage(PipelineStage): batch.extra["vae_scale_factor_spatial"] = vae_scale_factor_spatial self.log_info(f"Prepared latents with shape {shape}") + + sound_duration = float(getattr(batch, "sound_duration", 0.0) or 0.0) + if sound_duration > 0.0: + if not getattr(self.transformer, "sound_gen", False): + raise ValueError( + "sound generation was requested (sound_duration > 0) but the " + "loaded Cosmos3 checkpoint has no sound modality (sound_gen is " + "False)." + ) + sound_latent_fps = self.transformer.sound_latent_fps + sound_latent_frames = max(1, round(sound_duration * sound_latent_fps)) + sound_shape = (1, self.transformer.sound_dim, sound_latent_frames) + batch.audio_latents = torch.randn( + sound_shape, generator=generator, device=device, dtype=dtype + ) + self.log_info(f"Prepared sound latents with shape {sound_shape}") + + action_mode = getattr(batch.sampling_params, "action_mode", None) + if action_mode is not None: + if getattr(self.transformer, "action_dim", None) is None: + raise ValueError( + "action_mode is set but the loaded Cosmos3 checkpoint has no " + "action modality (action_gen is False)." + ) + self._prepare_action_latents(batch, generator, device, dtype) return batch + @staticmethod + def _resolve_domain_id(batch: Req) -> int: + """Resolve action embodiment domain ID; required for action generation.""" + domain_id = getattr(batch.sampling_params, "domain_id", None) + if domain_id is not None: + domain_id = int(domain_id) + if domain_id < 0: + raise ValueError(f"domain_id must be non-negative, got {domain_id}") + return domain_id + domain_name = getattr(batch.sampling_params, "domain_name", None) + if domain_name: + key = str(domain_name).strip().lower() + if key not in EMBODIMENT_TO_DOMAIN_ID: + raise ValueError( + f"Unknown action domain name {domain_name!r}. " + f"Valid names: {sorted(EMBODIMENT_TO_DOMAIN_ID)}" + ) + return EMBODIMENT_TO_DOMAIN_ID[key] + raise ValueError( + "Cosmos3 action generation requires --domain-id or --domain-name." + ) + + def _prepare_action_latents( + self, + batch: Req, + generator, + device: torch.device, + dtype: torch.dtype, + ) -> None: + """Prepare action latents and conditioning, writing them onto ``batch``. + + Action tokens run at frame rate (no temporal compression), so the chunk + length is ``num_frames - 1`` with ``start_frame_offset=1`` so each action + aligns with the frame it drives. + + Three modes: + - ``forward_dynamics``: the user supplies the action; all tokens are + clean conditioning (velocity mask 0) and the model predicts video. + - ``policy`` / ``inverse_dynamics``: actions are denoised from noise + (velocity mask 1); ``raw_action_dim`` is required. + """ + sp = batch.sampling_params + mode = str(sp.action_mode).strip().lower() + if mode not in ACTION_MODES: + raise ValueError( + f"Unsupported action_mode={sp.action_mode!r}; " + f"expected one of {sorted(ACTION_MODES)}" + ) + action_dim = self.transformer.action_dim + num_frames = batch.num_frames + + action_chunk_size = num_frames - 1 if num_frames > 1 else 1 + action_offset = 1 if action_chunk_size == num_frames - 1 else 0 + + domain_id = self._resolve_domain_id(batch) + raw_action_dim = getattr(sp, "raw_action_dim", None) + if raw_action_dim is None: + embodiment = getattr(sp, "domain_name", None) + if embodiment: + raw_action_dim = get_raw_action_dim(embodiment) + + if mode == ACTION_MODE_FORWARD_DYNAMICS: + raw = getattr(sp, "action", None) + if raw is None: + raise ValueError( + "action_mode='forward_dynamics' requires an 'action' array " + "(list[list[float]] of shape [T, D])." + ) + if isinstance(raw, str): + raw = json.loads(raw) + action = torch.as_tensor(np.asarray(raw), dtype=torch.float32) + if action.ndim == 3 and action.shape[0] == 1: + action = action.squeeze(0) + if action.ndim != 2: + raise ValueError( + f"action must have shape [T, D], got {tuple(action.shape)}" + ) + if action.shape[0] < action_chunk_size: + pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1) + action = torch.cat([action, pad], dim=0) + elif action.shape[0] > action_chunk_size: + action = action[:action_chunk_size] + if raw_action_dim is None: + raw_action_dim = int(action.shape[-1]) + stats_path = getattr(sp, "action_stats_path", None) + if stats_path is not None: + method = getattr(sp, "action_normalization", "quantile") + action = normalize_action(action, method, load_action_stats(stats_path)) + if action.shape[-1] < action_dim: + pad = torch.zeros(action.shape[0], action_dim - action.shape[-1]) + action = torch.cat([action, pad], dim=-1) + clean_action = action.to(device=device, dtype=dtype).unsqueeze(0) + else: + if raw_action_dim is None: + raise ValueError(f"action_mode={mode!r} requires --raw-action-dim.") + clean_action = torch.zeros( + 1, action_chunk_size, action_dim, device=device, dtype=dtype + ) + + raw_action_dim = int(raw_action_dim) + if not 0 < raw_action_dim <= action_dim: + raise ValueError( + f"raw_action_dim must be in [1, {action_dim}], got {raw_action_dim}" + ) + + # condition_mask marks clean (given) action tokens. forward_dynamics + # conditions on the whole action sequence; the others denoise it fully. + condition_mask = torch.zeros( + 1, action_chunk_size, 1, device=device, dtype=dtype + ) + if mode == ACTION_MODE_FORWARD_DYNAMICS: + condition_mask[:] = 1.0 + + noise = torch.randn( + 1, + action_chunk_size, + action_dim, + generator=generator, + device=device, + dtype=dtype, + ) + noise[:, :, raw_action_dim:] = 0.0 + clean_action[:, :, raw_action_dim:] = 0.0 + action_latents = condition_mask * clean_action + (1.0 - condition_mask) * noise + + batch.action_latents = action_latents + batch.extra["action_domain_ids"] = torch.tensor( + [domain_id], dtype=torch.long, device=device + ) + batch.extra["action_velocity_mask"] = 1.0 - condition_mask + batch.extra["action_condition_latents"] = clean_action + batch.extra["raw_action_dim"] = raw_action_dim + batch.extra["action_start_frame_offset"] = action_offset + self.log_info( + f"Prepared action latents with shape {tuple(action_latents.shape)} " + f"(mode={mode}, domain_id={domain_id}, raw_action_dim={raw_action_dim}, " + f"start_frame_offset={action_offset})" + ) + class Cosmos3TimestepPreparationStage(PipelineStage): """ @@ -362,9 +679,21 @@ class Cosmos3TimestepPreparationStage(PipelineStage): def __init__(self, scheduler): super().__init__() self.scheduler = scheduler - self.default_flow_shift = getattr( - getattr(scheduler, "config", None), "flow_shift", None - ) + + def _default_flow_shift_for_mode(self, batch: Req) -> float | None: + """Resolve the per-mode default flow_shift for the request. + + Matches cosmos-framework's built-in per-mode sample defaults. + """ + if getattr(batch.sampling_params, "action_mode", None) is not None: + return COSMOS3_ACTION_FLOW_SHIFT + if batch.data_type == DataType.IMAGE: + return COSMOS3_T2I_FLOW_SHIFT + if batch.preprocessed_image is not None: + return COSMOS3_I2V_FLOW_SHIFT + if batch.preprocessed_video is not None: + return COSMOS3_V2V_FLOW_SHIFT + return COSMOS3_T2V_FLOW_SHIFT def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Prepare scheduler timesteps.""" @@ -374,7 +703,7 @@ class Cosmos3TimestepPreparationStage(PipelineStage): if flow_shift is None: flow_shift = server_args.pipeline_config.flow_shift if flow_shift is None: - flow_shift = self.default_flow_shift + flow_shift = self._default_flow_shift_for_mode(batch) if flow_shift is not None and hasattr(self.scheduler, "set_shift"): self.scheduler.set_shift(float(flow_shift)) @@ -502,7 +831,13 @@ class Cosmos3DenoisingStage(PipelineStage): noisy_frame_mask: torch.Tensor | None = None, max_text_seq_len: int | None = None, current_timestep: int | None = None, - ) -> torch.Tensor: + sound_latents: torch.Tensor | None = None, + action_latents: torch.Tensor | None = None, + action_domain_ids: torch.Tensor | None = None, + action_noisy_mask: torch.Tensor | None = None, + action_fps: float | None = None, + action_start_frame_offset: int = 1, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Run transformer forward pass. Args: @@ -514,7 +849,7 @@ class Cosmos3DenoisingStage(PipelineStage): fps: Video frame rate cache_key: Key for the UND K/V cache. Use "cond" for conditional and "uncond" for unconditional to enable cache reuse across steps. - noisy_frame_mask: Optional [B, 1, T, 1, 1] I2V conditioning mask. + noisy_frame_mask: Optional [B, 1, T, 1, 1] I2V / V2V conditioning mask. """ if current_timestep is None: current_timestep = int(timestep.flatten()[0].item()) @@ -529,6 +864,12 @@ class Cosmos3DenoisingStage(PipelineStage): cache_key=cache_key, noisy_frame_mask=noisy_frame_mask, max_text_seq_len=max_text_seq_len, + sound_latents=sound_latents, + action_latents=action_latents, + action_domain_ids=action_domain_ids, + action_noisy_mask=action_noisy_mask, + action_fps=action_fps, + action_start_frame_offset=action_start_frame_offset, ) def _manage_device_placement(self, server_args: ServerArgs): @@ -565,6 +906,14 @@ class Cosmos3DenoisingStage(PipelineStage): self._manage_device_placement(server_args) latents = batch.latents + sound_latents = batch.audio_latents + action_latents = getattr(batch, "action_latents", None) + action_domain_ids = batch.extra.get("action_domain_ids") + action_velocity_mask = batch.extra.get("action_velocity_mask") + action_condition_latents = batch.extra.get("action_condition_latents") + action_raw_dim = batch.extra.get("raw_action_dim") + action_start_frame_offset = batch.extra.get("action_start_frame_offset", 1) + action_fps = getattr(batch.sampling_params, "action_fps", None) timesteps = batch.timesteps guidance_scale = batch.guidance_scale @@ -575,12 +924,28 @@ class Cosmos3DenoisingStage(PipelineStage): video_shape = batch.extra["video_shape"] fps = batch.extra.get("fps", 24.0) velocity_mask = batch.extra.get("velocity_mask") - image_latent = batch.image_latent + condition_latents = batch.extra.get("condition_latents") guidance_interval = getattr(batch.sampling_params, "guidance_interval", None) do_cfg = guidance_scale > 1.0 enable_cfg_parallel = server_args.enable_cfg_parallel and do_cfg + if action_latents is not None and enable_cfg_parallel: + raise NotImplementedError( + "Cosmos3 action generation does not support CFG parallel yet" + ) + + # Use separate scheduler instances for action/sound: UniPC keeps a + # per-call output history sized to the last sample, so video (5D), + # action (3D), and sound (3D) steps must not share state. + sound_scheduler = None + if sound_latents is not None: + sound_scheduler = copy.deepcopy(self.scheduler) + sound_scheduler.set_timesteps(len(timesteps), device=timesteps.device) + action_scheduler = None + if action_latents is not None: + action_scheduler = copy.deepcopy(self.scheduler) + action_scheduler.set_timesteps(len(timesteps), device=timesteps.device) cfg_rank = get_classifier_free_guidance_rank() if enable_cfg_parallel else 0 cfg_world_size = ( get_classifier_free_guidance_world_size() if enable_cfg_parallel else 1 @@ -644,6 +1009,12 @@ class Cosmos3DenoisingStage(PipelineStage): cond_text_seq_len=batch.extra["cond_text_seq_len"], uncond_text_seq_len=batch.extra["uncond_text_seq_len"], current_timestep=i, + sound_latents=sound_latents, + action_latents=action_latents, + action_domain_ids=action_domain_ids, + action_noisy_mask=action_velocity_mask, + action_fps=action_fps, + action_start_frame_offset=action_start_frame_offset, ) elif effective_scale == 1.0: noise_pred = self._run_transformer( @@ -657,6 +1028,12 @@ class Cosmos3DenoisingStage(PipelineStage): noisy_frame_mask=velocity_mask, max_text_seq_len=batch.extra["cond_text_seq_len"], current_timestep=i, + sound_latents=sound_latents, + action_latents=action_latents, + action_domain_ids=action_domain_ids, + action_noisy_mask=action_velocity_mask, + action_fps=action_fps, + action_start_frame_offset=action_start_frame_offset, ) else: noise_pred = self._predict_noise_cfg_batched( @@ -675,6 +1052,12 @@ class Cosmos3DenoisingStage(PipelineStage): batch.extra["uncond_text_seq_len"], ), current_timestep=i, + sound_latents=sound_latents, + action_latents=action_latents, + action_domain_ids=action_domain_ids, + action_noisy_mask=action_velocity_mask, + action_fps=action_fps, + action_start_frame_offset=action_start_frame_offset, ) else: noise_pred = self._run_transformer( @@ -688,11 +1071,30 @@ class Cosmos3DenoisingStage(PipelineStage): noisy_frame_mask=velocity_mask, max_text_seq_len=batch.extra["cond_text_seq_len"], current_timestep=i, + sound_latents=sound_latents, + action_latents=action_latents, + action_domain_ids=action_domain_ids, + action_noisy_mask=action_velocity_mask, + action_fps=action_fps, + action_start_frame_offset=action_start_frame_offset, ) - # I2V: zero-velocity at conditioned frames so the scheduler keeps - # them clean; UniPC's predictor-corrector still rescales the - # sample, so we re-inject the clean image latent below. + # Unpack multi-modality outputs; ordering is (video[, action][, sound]). + action_noise_pred = None + sound_noise_pred = None + if isinstance(noise_pred, tuple): + out_idx = 1 + video_noise_pred = noise_pred[0] + if action_latents is not None: + action_noise_pred = noise_pred[out_idx] + out_idx += 1 + if sound_latents is not None: + sound_noise_pred = noise_pred[out_idx] + noise_pred = video_noise_pred + + # I2V / V2V: zero-velocity at conditioned frames so the scheduler + # keeps them clean; UniPC's predictor-corrector still rescales the + # sample, so we re-blend the clean condition latents below. if velocity_mask is not None: noise_pred = noise_pred * velocity_mask @@ -703,13 +1105,53 @@ class Cosmos3DenoisingStage(PipelineStage): return_dict=False, )[0] - if image_latent is not None: - latents[:, :, 0:1, :, :] = image_latent + if action_noise_pred is not None: + # Zero the velocity at conditioned (clean) action tokens and at + # padding dims so the scheduler only denoises the active slots, + # then re-blend the clean condition after the step. + if action_velocity_mask is not None: + action_noise_pred = action_noise_pred * action_velocity_mask + if ( + action_raw_dim is not None + and action_raw_dim < action_noise_pred.shape[-1] + ): + action_noise_pred[..., action_raw_dim:] = 0.0 + action_latents = action_scheduler.step( + action_noise_pred, + t, + action_latents, + return_dict=False, + )[0] + if ( + action_condition_latents is not None + and action_velocity_mask is not None + ): + action_latents = ( + action_velocity_mask * action_latents + + (1.0 - action_velocity_mask) * action_condition_latents + ) + + if sound_noise_pred is not None: + sound_latents = sound_scheduler.step( + sound_noise_pred, + t, + sound_latents, + return_dict=False, + )[0] + + if condition_latents is not None and velocity_mask is not None: + latents = ( + velocity_mask * latents + (1.0 - velocity_mask) * condition_latents + ) if batch.profile and not batch.is_warmup: self.step_profile() batch.latents = latents + if action_latents is not None: + batch.action_latents = action_latents + if sound_latents is not None: + batch.audio_latents = sound_latents self.log_info("Denoising complete") return batch @@ -727,7 +1169,13 @@ class Cosmos3DenoisingStage(PipelineStage): noisy_frame_mask: torch.Tensor | None = None, max_text_seq_len: int | None = None, current_timestep: int | None = None, - ) -> torch.Tensor: + sound_latents: torch.Tensor | None = None, + action_latents: torch.Tensor | None = None, + action_domain_ids: torch.Tensor | None = None, + action_noisy_mask: torch.Tensor | None = None, + action_fps: float | None = None, + action_start_frame_offset: int = 1, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Run CFG by stacking both branches into a batch_size=2 forward. Halves the kernel-launch count vs running cond and uncond serially. @@ -743,8 +1191,28 @@ class Cosmos3DenoisingStage(PipelineStage): if noisy_frame_mask is not None else None ) + sound_batched = ( + torch.cat([sound_latents, sound_latents], dim=0) + if sound_latents is not None + else None + ) + action_batched = ( + torch.cat([action_latents, action_latents], dim=0) + if action_latents is not None + else None + ) + action_domain_ids_batched = ( + torch.cat([action_domain_ids, action_domain_ids], dim=0) + if action_domain_ids is not None + else None + ) + action_noisy_mask_batched = ( + torch.cat([action_noisy_mask, action_noisy_mask], dim=0) + if action_noisy_mask is not None + else None + ) - noise_pred = self._run_transformer( + out = self._run_transformer( latents=latents_batched, timestep=timestep_batched, text_ids=text_ids_batched, @@ -755,13 +1223,21 @@ class Cosmos3DenoisingStage(PipelineStage): noisy_frame_mask=mask_batched, max_text_seq_len=max_text_seq_len, current_timestep=current_timestep, + sound_latents=sound_batched, + action_latents=action_batched, + action_domain_ids=action_domain_ids_batched, + action_noisy_mask=action_noisy_mask_batched, + action_fps=action_fps, + action_start_frame_offset=action_start_frame_offset, ) - noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2, dim=0) - # CFG: uncond + g·(cond − uncond). - return noise_pred_uncond + guidance_scale * ( - noise_pred_cond - noise_pred_uncond - ) + def _cfg_combine(pred: torch.Tensor) -> torch.Tensor: + uncond, cond = pred.chunk(2, dim=0) + return uncond + guidance_scale * (cond - uncond) + + if isinstance(out, tuple): + return tuple(_cfg_combine(p) for p in out) + return _cfg_combine(out) def _predict_noise_cfg_parallel( self, @@ -779,45 +1255,54 @@ class Cosmos3DenoisingStage(PipelineStage): cond_text_seq_len: int | None = None, uncond_text_seq_len: int | None = None, current_timestep: int | None = None, - ) -> torch.Tensor: + sound_latents: torch.Tensor | None = None, + action_latents: torch.Tensor | None = None, + action_domain_ids: torch.Tensor | None = None, + action_noisy_mask: torch.Tensor | None = None, + action_fps: float | None = None, + action_start_frame_offset: int = 1, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Run CFG with one branch per CFG rank, combined by all-reduce. Rank 0 runs the conditional branch and contributes ``g·cond`` to the sum; rank 1 runs the unconditional branch and contributes ``(1−g)·uncond``. The all-reduce sum is exactly the standard CFG result. Each rank keeps its own UND K/V cache (``"cond"`` / - ``"uncond"``). + ``"uncond"``). When sound/action modalities are present the forward + returns a per-modality tuple; each branch scales every modality by its + coefficient and the reduction combines them element-wise. """ if cfg_rank == 0: - noise_pred = self._run_transformer( - latents=latents, - timestep=timestep, - text_ids=cond_text_ids, - text_mask=cond_text_mask, - video_shape=video_shape, - fps=fps, - cache_key="cond", - noisy_frame_mask=noisy_frame_mask, - max_text_seq_len=cond_text_seq_len, - current_timestep=current_timestep, - ) - partial = guidance_scale * noise_pred + text_ids, text_mask, cache_key = cond_text_ids, cond_text_mask, "cond" + text_seq_len = cond_text_seq_len + coeff = guidance_scale else: - noise_pred = self._run_transformer( - latents=latents, - timestep=timestep, - text_ids=uncond_text_ids, - text_mask=uncond_text_mask, - video_shape=video_shape, - fps=fps, - cache_key="uncond", - noisy_frame_mask=noisy_frame_mask, - max_text_seq_len=uncond_text_seq_len, - current_timestep=current_timestep, - ) - partial = (1.0 - guidance_scale) * noise_pred + text_ids, text_mask, cache_key = uncond_text_ids, uncond_text_mask, "uncond" + text_seq_len = uncond_text_seq_len + coeff = 1.0 - guidance_scale - return cfg_model_parallel_all_reduce(partial) + out = self._run_transformer( + latents=latents, + timestep=timestep, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=fps, + cache_key=cache_key, + noisy_frame_mask=noisy_frame_mask, + max_text_seq_len=text_seq_len, + current_timestep=current_timestep, + sound_latents=sound_latents, + action_latents=action_latents, + action_domain_ids=action_domain_ids, + action_noisy_mask=action_noisy_mask, + action_fps=action_fps, + action_start_frame_offset=action_start_frame_offset, + ) + + if isinstance(out, tuple): + return tuple(cfg_model_parallel_all_reduce(coeff * p) for p in out) + return cfg_model_parallel_all_reduce(coeff * out) class Cosmos3DecodingStage(PipelineStage): @@ -830,12 +1315,13 @@ class Cosmos3DecodingStage(PipelineStage): parallelism_type = StageParallelismType.REPLICATED - def __init__(self, vae, guardrails: bool = False): + def __init__(self, vae, guardrails: bool = False, sound_tokenizer=None): super().__init__() self.vae = vae self._latents_mean = None self._latents_std = None self._guardrails = guardrails + self.sound_tokenizer = sound_tokenizer if guardrails: from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import ( _init_guardrails, @@ -935,7 +1421,50 @@ class Cosmos3DecodingStage(PipelineStage): elif not is_image_gen: self.log_info(f"Postprocessed video tensor shape: {output.shape}") + audio = None + audio_sample_rate = None + if self.sound_tokenizer is not None and batch.audio_latents is not None: + if server_args.vae_cpu_offload: + self.sound_tokenizer.to(device) + with torch.no_grad(): + decoded_audio = self.sound_tokenizer.decode( + batch.audio_latents.to(device) + ) + audio = decoded_audio.float().cpu() + audio_sample_rate = self.sound_tokenizer.sample_rate + if server_args.vae_cpu_offload and not getattr(batch, "is_warmup", False): + self.sound_tokenizer.to("cpu", non_blocking=True) + self.log_info( + f"Decoded audio tensor shape: {tuple(audio.shape)} @ {audio_sample_rate} Hz" + ) + + action_pred = None + if getattr(batch, "action_latents", None) is not None: + raw_action_dim = batch.extra.get("raw_action_dim") + action_pred = batch.action_latents.float().cpu() + if raw_action_dim is not None: + action_pred = action_pred[:, :, :raw_action_dim] + stats_path = getattr(batch.sampling_params, "action_stats_path", None) + if stats_path is not None: + method = getattr( + batch.sampling_params, "action_normalization", "quantile" + ) + action_pred = denormalize_action( + action_pred, method, load_action_stats(stats_path) + ) + self.log_info(f"Action predictions shape: {tuple(action_pred.shape)}") + return OutputBatch( output=output, + audio=audio, + audio_sample_rate=audio_sample_rate, + action_pred=action_pred, + action_mode=getattr(batch.sampling_params, "action_mode", None), + action_domain_id=getattr(batch.sampling_params, "domain_id", None), + action_raw_action_dim=( + batch.extra.get("raw_action_dim") + if getattr(batch, "extra", None) + else None + ), metrics=batch.metrics if hasattr(batch, "metrics") else None, ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_action.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_action.py new file mode 100644 index 000000000..a573673eb --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_action.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 action modality helpers: domain mapping, mode constants, the +structured JSON caption, and dataset-derived action (de)normalization. +""" + +import json +import math +from pathlib import Path + +import numpy as np +import torch + +ACTION_MODE_POLICY = "policy" +ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics" +ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics" +ACTION_MODES = { + ACTION_MODE_POLICY, + ACTION_MODE_FORWARD_DYNAMICS, + ACTION_MODE_INVERSE_DYNAMICS, +} + +EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = { + "no_action": 0, + "av": 1, + "camera_pose": 2, + "hand_pose": 3, + "pusht": 4, + "libero": 5, + "umi": 6, + "bridge_orig_lerobot": 7, + "droid_lerobot": 8, + "robomind-franka": 8, + "galbot": 9, + "robomind-franka-dual": 12, + "robomind-ur": 13, + "agibotworld": 15, + "agibot_gear_gripper": 15, + "agibot_gear_gripper_ext": 15, + "fractal": 20, +} + +# Embodiment -> real (unpadded) action channel count. Channels beyond this are +# zero-padding up to the model's action_dim. +EMBODIMENT_TO_RAW_ACTION_DIM: dict[str, int] = { + "av": 9, + "camera_pose": 9, + "pusht": 2, + "umi": 10, + "bridge_orig_lerobot": 10, + "droid_lerobot": 10, + "robomind-franka": 10, + "robomind-franka-dual": 20, + "robomind-ur": 10, + "agibotworld": 29, + "fractal": 10, +} + +# Canonical (width, height) targets per resolution tier and aspect ratio, used +# to render the aspect_ratio field of the action caption. +VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = { + "256": { + "1,1": (256, 256), + "4,3": (320, 256), + "3,4": (256, 320), + "16,9": (320, 192), + "9,16": (192, 320), + }, + "480": { + "1,1": (640, 640), + "4,3": (736, 544), + "3,4": (544, 736), + "16,9": (832, 480), + "9,16": (480, 832), + }, + "704": { + "1,1": (960, 960), + "4,3": (1088, 832), + "3,4": (832, 1088), + "16,9": (1280, 704), + "9,16": (704, 1280), + }, + "720": { + "1,1": (960, 960), + "4,3": (1104, 832), + "3,4": (832, 1104), + "16,9": (1280, 720), + "9,16": (720, 1280), + }, +} + +VIEWPOINT_TEMPLATES: dict[str, str] = { + "ego_view": "This video is captured from a first-person perspective looking at the scene.", + "third_person_view": "This video is captured from a third-person perspective looking towards the agent from the front.", + "wrist_view": "This video is captured from a wrist-mounted camera.", + "concat_view": "This video contains concatenated views from multiple camera perspectives.", +} + +_STAT_KEYS = {"mean", "std", "min", "max", "q01", "q99"} + + +def get_raw_action_dim(embodiment: str) -> int: + key = embodiment.lower().strip() + if key not in EMBODIMENT_TO_RAW_ACTION_DIM: + raise ValueError( + f"No raw action dim for Cosmos3 embodiment {embodiment!r}. Expected one " + f"of {sorted(EMBODIMENT_TO_RAW_ACTION_DIM)}." + ) + return EMBODIMENT_TO_RAW_ACTION_DIM[key] + + +def canonical_aspect_ratio(width: int, height: int) -> str: + """Canonical ``"W,H"`` aspect string for the action caption.""" + for sizes in VIDEO_RES_SIZE_INFO.values(): + for aspect, (cand_w, cand_h) in sizes.items(): + if width == cand_w and height == cand_h: + return aspect + divisor = math.gcd(width, height) + if divisor == 0: + raise ValueError( + f"width and height must be non-zero, got width={width}, height={height}." + ) + return f"{width // divisor},{height // divisor}" + + +def build_action_prompt( + description: str, + view_point: str, + num_frames: int, + fps: float, + height: int, + width: int, +) -> str: + """Render the structured JSON action caption the action checkpoints expect.""" + duration_seconds = num_frames / fps + minutes, secs = divmod(round(duration_seconds), 60) + if description and description[-1] not in ".!?": + description = description + "." + prompt = { + "cinematography": { + "framing": VIEWPOINT_TEMPLATES.get( + view_point, VIEWPOINT_TEMPLATES["ego_view"] + ) + }, + "actions": [{"time": f"0:00-{minutes}:{secs:02d}", "description": description}], + "duration": f"{int(duration_seconds)}s", + "fps": float(fps), + "resolution": {"H": int(height), "W": int(width)}, + "aspect_ratio": canonical_aspect_ratio(int(width), int(height)), + } + return json.dumps(prompt) + + +def load_action_stats( + stats_path: str, stats_key: str = "global" +) -> dict[str, torch.Tensor]: + """Load per-channel action normalization stats from a JSON file.""" + path = Path(stats_path) + if not path.exists(): + raise FileNotFoundError( + f"Action normalization stats not found at {stats_path}." + ) + raw = json.loads(path.read_text()) + if stats_key in raw: + raw = raw[stats_key] + return { + k: torch.as_tensor(np.array(v, dtype=np.float32)) + for k, v in raw.items() + if k in _STAT_KEYS + } + + +def normalize_action( + action: torch.Tensor, method: str, stats: dict[str, torch.Tensor] +) -> torch.Tensor: + if method == "quantile": + q01, q99 = stats["q01"].to(action), stats["q99"].to(action) + return (2.0 * (action - q01) / (q99 - q01).clamp(min=1e-8) - 1.0).clamp( + -1.0, 1.0 + ) + if method == "meanstd": + return (action - stats["mean"].to(action)) / stats["std"].to(action).clamp( + min=1e-8 + ) + if method == "minmax": + lo, hi = stats["min"].to(action), stats["max"].to(action) + return (2.0 * (action - lo) / (hi - lo).clamp(min=1e-8) - 1.0).clamp(-1.0, 1.0) + raise ValueError(f"Unknown action normalization method {method!r}.") + + +def denormalize_action( + action: torch.Tensor, method: str, stats: dict[str, torch.Tensor] +) -> torch.Tensor: + if method == "quantile": + q01, q99 = stats["q01"].to(action), stats["q99"].to(action) + return (action + 1.0) / 2.0 * (q99 - q01) + q01 + if method == "meanstd": + return action * stats["std"].to(action) + stats["mean"].to(action) + if method == "minmax": + lo, hi = stats["min"].to(action), stats["max"].to(action) + return (action + 1.0) / 2.0 * (hi - lo) + lo + raise ValueError(f"Unknown action normalization method {method!r}.") diff --git a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py index 428896ff3..10859b05c 100644 --- a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py @@ -162,6 +162,7 @@ def get_quant_config( packed_modules_mapping: Dict[str, List[str]] = {}, reverse_param_names_mapping: Dict[str, List[str]] = {}, remap_prefix: Dict[str, str] | None = None, + quant_ignore_remap: Optional[Dict[str, str]] = None, ) -> QuantizationConfig: quant_cfg = find_quant_modelslim_config(model_config, component_model_path) if quant_cfg is not None: @@ -191,7 +192,16 @@ def get_quant_config( hf_quant_config = getattr(model_config, "compression_config", None) if hf_quant_config is not None: hf_quant_config["packed_modules_mapping"] = packed_modules_mapping - return quant_cls.from_config(hf_quant_config) + is_modelopt_fp8 = ( + hf_quant_config.get("quant_method") == "modelopt" + and "FP8" in str(hf_quant_config.get("quant_algo", "")).upper() + ) + extra_kwargs = ( + {"ignore_remap": quant_ignore_remap} + if quant_ignore_remap and is_modelopt_fp8 + else {} + ) + return quant_cls.from_config(hf_quant_config, **extra_kwargs) model_name_or_path = model_config["model_path"] hf_folder = model_name_or_path diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 6eb3eb152..54cf3f589 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -483,7 +483,6 @@ COSMOS3_NANO_CI_sampling_params = DiffusionSamplingParams( "num_inference_steps": 35, "seed": 0, "max_sequence_length": 128, - "flow_shift": 10.0, "extra_args": { "guardrails": False, "use_resolution_template": False, diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index af9afaf7a..f59ec5093 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -34,7 +34,7 @@ if TYPE_CHECKING: logger = init_logger(__name__) -SGL_TEST_FILES_CI_DATA_REVISION = "77bd016251220fee8917a30ec92e89da03794a8a" +SGL_TEST_FILES_CI_DATA_REVISION = "9a64abec5a7517a9f2b04ac1b4eab4173adb2d38" if current_platform.is_npu(): SGL_TEST_FILES_CI_DATA_REVISION = "6b62f4b6825c76a25fd2ba28248df68f2b400e65" diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py index 3d120e1f8..42c9f7a82 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py @@ -2,10 +2,11 @@ """Unit tests for Cosmos3 config, weight mapping, and sampling params.""" import importlib.util +import types import unittest from unittest import mock -from fastapi import HTTPException +import torch from sglang.multimodal_gen.configs.models.dits.cosmos3video import ( _build_cosmos3_param_names_mapping, @@ -22,9 +23,29 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( VideoGenerationsRequest, ) from sglang.multimodal_gen.runtime.entrypoints.openai.video_api import ( - _reject_unsupported_cosmos3_modes, + _cosmos3_sampling_param_kwargs, + _resolve_sound_duration, + _resolve_video_path, +) +from sglang.multimodal_gen.runtime.loader.component_loaders import scheduler_loader +from sglang.multimodal_gen.runtime.loader.component_loaders.scheduler_loader import ( + SchedulerLoader, ) from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping +from sglang.multimodal_gen.runtime.models.dits.cosmos3video import ( + DomainAwareLinear, + compute_mrope_position_ids_action, + compute_mrope_position_ids_sound, + compute_mrope_position_ids_vision, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3 import ( + Cosmos3ImagePreprocessStage, + Cosmos3LatentPreparationStage, + Cosmos3TimestepPreparationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_action import ( + EMBODIMENT_TO_DOMAIN_ID, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import ( is_cosmos_guardrail_available, ) @@ -54,16 +75,16 @@ class TestCosmos3ParamNamesMapping(unittest.TestCase): key, idx, total = _apply(self.fn, "norm.weight") self.assertEqual(key, "") - def test_audio_proj_in_dropped(self): - key, *_ = _apply(self.fn, "audio_proj_in.weight") - self.assertEqual(key, "") - - def test_action_proj_in_dropped(self): - key, *_ = _apply(self.fn, "action_proj_in.weight") - self.assertEqual(key, "") - # --- top-level pass-through --- + def test_audio_proj_in_passthrough(self): + key, *_ = _apply(self.fn, "audio_proj_in.weight") + self.assertEqual(key, "audio_proj_in.weight") + + def test_action_proj_in_passthrough(self): + key, *_ = _apply(self.fn, "action_proj_in.fc.weight") + self.assertEqual(key, "action_proj_in.fc.weight") + def test_embed_tokens(self): key, *_ = _apply(self.fn, "embed_tokens.weight") self.assertEqual(key, "language_model.embed_tokens.weight") @@ -204,6 +225,87 @@ class TestCosmos3AdjustNumFrames(unittest.TestCase): self.assertEqual(self.cfg.adjust_num_frames(2), 1) +class TestCosmos3SchedulerConfig(unittest.TestCase): + """Verify Cosmos3 scheduler class and flow-shift defaults.""" + + def test_config_overrides_checkpoint_scheduler_class(self): + cfg = Cosmos3Config() + self.assertEqual(cfg.scheduler_class_override, "FlowUniPCMultistepScheduler") + self.assertIsNone(cfg.flow_shift) + + def test_scheduler_loader_uses_configured_class_override(self): + class FakeScheduler: + def __init__(self, **config): + self.config = config + + server_args = types.SimpleNamespace( + pipeline_config=types.SimpleNamespace( + scheduler_class_override="FlowUniPCMultistepScheduler", + flow_shift=None, + ) + ) + with ( + mock.patch.object( + scheduler_loader, + "get_diffusers_component_config", + return_value={"_class_name": "CheckpointScheduler", "foo": "bar"}, + ), + mock.patch.object( + scheduler_loader.ModelRegistry, + "resolve_model_cls", + return_value=(FakeScheduler, None), + ) as resolve, + ): + scheduler = SchedulerLoader().load_customized("unused", server_args) + + resolve.assert_called_once_with("FlowUniPCMultistepScheduler") + self.assertEqual(scheduler.config["foo"], "bar") + + @staticmethod + def _stage(): + stage = Cosmos3TimestepPreparationStage.__new__(Cosmos3TimestepPreparationStage) + stage.scheduler = types.SimpleNamespace( + config=types.SimpleNamespace(flow_shift=1.0) + ) + return stage + + @staticmethod + def _batch(**kwargs): + sp_kwargs = kwargs.pop("sp_kwargs", {}) + return types.SimpleNamespace( + sampling_params=Cosmos3SamplingParams(prompt="t", **sp_kwargs), + data_type=kwargs.pop("data_type", DataType.VIDEO), + preprocessed_image=kwargs.pop("preprocessed_image", None), + preprocessed_video=kwargs.pop("preprocessed_video", None), + ) + + def test_per_mode_flow_shift_defaults(self): + stage = self._stage() + self.assertEqual( + stage._default_flow_shift_for_mode(self._batch(data_type=DataType.IMAGE)), + 3.0, + ) + self.assertEqual( + stage._default_flow_shift_for_mode( + self._batch(preprocessed_image=torch.empty(1)) + ), + 10.0, + ) + self.assertEqual( + stage._default_flow_shift_for_mode( + self._batch(preprocessed_video=torch.empty(1)) + ), + 10.0, + ) + self.assertEqual(stage._default_flow_shift_for_mode(self._batch()), 10.0) + self.assertEqual( + stage._default_flow_shift_for_mode( + self._batch(sp_kwargs={"action_mode": "policy"}) + ), + 10.0, + ) + + class TestCosmos3SamplingParamsDataType(unittest.TestCase): """Verify num_frames==1 flips data_type to IMAGE before file name derivation.""" @@ -246,9 +348,9 @@ class TestCosmos3ModelResolution(unittest.TestCase): class TestCosmos3OpenAIProtocol(unittest.TestCase): - """Verify Cosmos3-only knobs stay out of the stable request schema.""" + """Verify Cosmos3 modality knobs are exposed by the video HTTP schema.""" - def test_cosmos3_private_fields_are_extra_fields(self): + def test_cosmos3_template_fields_remain_extra_fields(self): for request_cls in (ImageGenerationsRequest, VideoGenerationsRequest): with self.subTest(request_cls=request_cls.__name__): self.assertIn("max_sequence_length", request_cls.model_fields) @@ -258,22 +360,66 @@ class TestCosmos3OpenAIProtocol(unittest.TestCase): self.assertNotIn("use_system_prompt", request_cls.model_fields) self.assertNotIn("use_guardrails", request_cls.model_fields) - self.assertNotIn("generate_sound", VideoGenerationsRequest.model_fields) - self.assertNotIn("sound_duration", VideoGenerationsRequest.model_fields) + def test_cosmos3_modal_fields_pass_through_as_extras(self): + for field_name in ("video_path", "video_url"): + with self.subTest(field_name=field_name): + self.assertIn(field_name, VideoGenerationsRequest.model_fields) - def test_unsupported_cosmos3_modes_allow_falsy_extra_fields(self): + modal_values = { + "generate_sound": True, + "sound_duration": 3.0, + "condition_frame_indexes": [0, 2], + "condition_frame_indexes_vision": [0, 2], + "condition_video_keep": "last", + "action_mode": "policy", + "domain_id": 1, + "domain_name": "umi", + "raw_action_dim": 9, + "action_fps": 30.0, + "action": [0.0, 1.0], + "action_view_point": "ego_view", + "action_normalization": "mean_std", + } + req = VideoGenerationsRequest(prompt="test", **modal_values) + for field_name, value in modal_values.items(): + with self.subTest(field_name=field_name): + self.assertNotIn(field_name, VideoGenerationsRequest.model_fields) + self.assertEqual(getattr(req, field_name), value) + + def test_cosmos3_http_aliases_map_to_sampling_params(self): req = VideoGenerationsRequest( prompt="test", - generate_sound=False, - action_mode="", - condition_frame_indexes_vision=[], - condition_video_keep={}, + video_url="https://example.com/input.mp4", + generate_sound=True, + condition_frame_indexes_vision=[0, 2], + condition_video_keep="last", + action_mode="policy", + domain_name="umi", + raw_action_dim=9, + action_fps=30.0, + action_view_point="ego_view", ) - _reject_unsupported_cosmos3_modes(req, "nvidia/Cosmos3-Nano") - req = VideoGenerationsRequest(prompt="test", generate_sound=True) - with self.assertRaises(HTTPException): - _reject_unsupported_cosmos3_modes(req, "nvidia/Cosmos3-Nano") + self.assertEqual(_resolve_video_path(req), "https://example.com/input.mp4") + + kwargs = _cosmos3_sampling_param_kwargs(req, num_frames=48, fps=24) + self.assertEqual(kwargs["sound_duration"], 2.0) + self.assertEqual(kwargs["condition_frame_indexes"], [0, 2]) + self.assertEqual(kwargs["condition_video_keep"], "last") + self.assertEqual(kwargs["action_mode"], "policy") + self.assertEqual(kwargs["domain_name"], "umi") + self.assertEqual(kwargs["raw_action_dim"], 9) + self.assertEqual(kwargs["action_fps"], 30.0) + self.assertEqual(kwargs["action_view_point"], "ego_view") + + def test_generate_sound_false_disables_sound_duration(self): + req = VideoGenerationsRequest( + prompt="test", generate_sound=False, sound_duration=3.0 + ) + self.assertEqual( + _resolve_sound_duration(req, num_frames=48, fps=24), + 0.0, + ) class TestCosmos3Guardrails(unittest.TestCase): @@ -296,5 +442,306 @@ class TestCosmos3Guardrails(unittest.TestCase): self.assertFalse(is_cosmos_guardrail_available()) +class TestCosmos3MRoPE(unittest.TestCase): + """mRoPE position-ID computation for vision / sound / action token grids.""" + + DEVICE = torch.device("cpu") + + def test_vision_default_args_unchanged(self): + # With base_temporal_compression_factor=None and start_frame_offset=0 + # the token and base rates cancel, so t-index == frame index. + pos, _ = compute_mrope_position_ids_vision( + grid_t=21, + grid_h=1, + grid_w=1, + temporal_offset=0, + device=self.DEVICE, + fps=24.0, + base_fps=24.0, + temporal_compression_factor=4, + ) + self.assertEqual(tuple(pos.shape), (3, 21)) + self.assertAlmostEqual(float(pos[0, 0]), 0.0, places=5) + self.assertAlmostEqual(float(pos[0, 20]), 20.0, places=5) + + def test_sound_grid_shape_and_scaling(self): + pos, _ = compute_mrope_position_ids_sound( + grid_t=10, + temporal_offset=0, + sound_latent_fps=25.0, + device=self.DEVICE, + base_fps=24.0, + temporal_compression_factor_sound=1, + ) + # (T, 1, 1) grid -> spatial axes are all zero. + self.assertEqual(tuple(pos.shape), (3, 10)) + self.assertTrue(torch.all(pos[1] == 0)) + self.assertTrue(torch.all(pos[2] == 0)) + # t-index = i / sound_fps * base_fps = i / 25 * 24. + self.assertAlmostEqual(float(pos[0, 5]), 5 / 25 * 24, places=4) + + def test_action_uses_video_base_compression_and_offset(self): + # Action runs at frame rate (tcf=1) but is scaled by the video's + # base_temporal_compression_factor=4, and shifted by start_frame_offset. + pos, _ = compute_mrope_position_ids_action( + grid_t=16, + temporal_offset=0, + action_fps=10.0, + device=self.DEVICE, + base_fps=24.0, + base_temporal_compression_factor=4, + start_frame_offset=1, + ) + self.assertEqual(tuple(pos.shape), (3, 16)) + self.assertTrue(torch.all(pos[1] == 0)) + self.assertTrue(torch.all(pos[2] == 0)) + # t-index[i] = (i + start_frame_offset) / action_fps * (base_fps / base_tcf) + # = (i + 1) / 10 * (24 / 4) = (i + 1) * 0.6 + self.assertAlmostEqual(float(pos[0, 0]), 0.6, places=4) + self.assertAlmostEqual(float(pos[0, 15]), 16 * 0.6, places=4) + + def test_action_offset_zero(self): + pos, _ = compute_mrope_position_ids_action( + grid_t=8, + temporal_offset=0, + action_fps=10.0, + device=self.DEVICE, + base_fps=24.0, + base_temporal_compression_factor=4, + start_frame_offset=0, + ) + self.assertAlmostEqual(float(pos[0, 0]), 0.0, places=5) + + def test_action_aligns_with_video_positions(self): + # Action frames at frame rate should share the video's temporal frame: + # every 4th action token lands on the next video latent-frame position. + media_offset = 100 + vid, _ = compute_mrope_position_ids_vision( + grid_t=5, + grid_h=1, + grid_w=1, + temporal_offset=media_offset, + device=self.DEVICE, + fps=24.0, + base_fps=24.0, + temporal_compression_factor=4, + ) + act, _ = compute_mrope_position_ids_action( + grid_t=16, + temporal_offset=media_offset, + action_fps=24.0, + device=self.DEVICE, + base_fps=24.0, + base_temporal_compression_factor=4, + start_frame_offset=0, + ) + # video latent frame 1 sits at media_offset+1; action frame 4 (4 frames + # per latent at tcf=4) lands at the same temporal position. + self.assertAlmostEqual(float(vid[0, 1]), float(act[0, 4]), places=4) + + +class TestCosmos3DomainAwareLinear(unittest.TestCase): + """Per-domain action projection.""" + + def test_rank3_and_rank2_shapes(self): + layer = DomainAwareLinear(input_size=7, output_size=64, num_domains=32) + x3 = torch.randn(2, 16, 7) + out3 = layer(x3, torch.tensor([1, 5])) + self.assertEqual(tuple(out3.shape), (2, 16, 64)) + x2 = torch.randn(3, 7) + out2 = layer(x2, torch.tensor([0, 1, 2])) + self.assertEqual(tuple(out2.shape), (3, 64)) + + def test_distinct_domains_give_distinct_outputs(self): + torch.manual_seed(0) + layer = DomainAwareLinear(input_size=4, output_size=8, num_domains=4) + x = torch.randn(1, 3, 4) + out_a = layer(x, torch.tensor([0])) + out_b = layer(x, torch.tensor([2])) + self.assertFalse(torch.allclose(out_a, out_b)) + + def test_scalar_domain_id_promoted(self): + layer = DomainAwareLinear(input_size=4, output_size=8, num_domains=4) + out = layer(torch.randn(1, 2, 4), torch.tensor(3)) + self.assertEqual(tuple(out.shape), (1, 2, 8)) + + +class TestCosmos3ConditionIndexes(unittest.TestCase): + """Vision condition-frame resolution across V2V and action modes.""" + + @staticmethod + def _batch(num_frames=61, **sp_kwargs): + sp = Cosmos3SamplingParams(prompt="t", num_frames=num_frames, **sp_kwargs) + return types.SimpleNamespace(sampling_params=sp, num_frames=num_frames) + + def test_v2v_default(self): + idx = Cosmos3ImagePreprocessStage._resolve_condition_indexes(self._batch()) + self.assertEqual(idx, [0, 1]) + + def test_v2v_explicit_sorted_unique(self): + idx = Cosmos3ImagePreprocessStage._resolve_condition_indexes( + self._batch(condition_frame_indexes=[2, 0, 2]) + ) + self.assertEqual(idx, [0, 2]) + + def test_inverse_dynamics_conditions_all_latent_frames(self): + # 61 frames -> (61-1)//4 + 1 = 16 latent frames, all locked. + idx = Cosmos3ImagePreprocessStage._resolve_condition_indexes( + self._batch(num_frames=61, action_mode="inverse_dynamics") + ) + self.assertEqual(idx, list(range(16))) + + +class TestCosmos3DomainResolution(unittest.TestCase): + """Embodiment domain-id resolution for action generation.""" + + @staticmethod + def _batch(**sp_kwargs): + sp = Cosmos3SamplingParams(prompt="t", **sp_kwargs) + return types.SimpleNamespace(sampling_params=sp) + + def test_explicit_domain_id(self): + self.assertEqual( + Cosmos3LatentPreparationStage._resolve_domain_id(self._batch(domain_id=7)), + 7, + ) + + def test_domain_name_lookup(self): + self.assertEqual( + Cosmos3LatentPreparationStage._resolve_domain_id( + self._batch(domain_name="av") + ), + EMBODIMENT_TO_DOMAIN_ID["av"], + ) + self.assertEqual( + Cosmos3LatentPreparationStage._resolve_domain_id( + self._batch(domain_name="umi") + ), + EMBODIMENT_TO_DOMAIN_ID["umi"], + ) + + def test_missing_domain_raises(self): + with self.assertRaises(ValueError): + Cosmos3LatentPreparationStage._resolve_domain_id(self._batch()) + + def test_unknown_domain_name_raises(self): + with self.assertRaises(ValueError): + Cosmos3LatentPreparationStage._resolve_domain_id( + self._batch(domain_name="not_a_robot") + ) + + +class TestCosmos3ActionLatentPrep(unittest.TestCase): + """Action latent / mask preparation per action mode.""" + + @classmethod + def setUpClass(cls): + # Bypass PipelineStage.__init__ (needs global server args); only the + # transformer's action_dim and log_info are used by _prepare_action_latents. + cls.stage = Cosmos3LatentPreparationStage.__new__(Cosmos3LatentPreparationStage) + cls.stage.transformer = types.SimpleNamespace(action_dim=64) + cls.stage.log_info = lambda *a, **k: None + cls.device = torch.device("cpu") + cls.dtype = torch.float32 + + def _run(self, num_frames=17, **sp_kwargs): + sp = Cosmos3SamplingParams(prompt="t", num_frames=num_frames, **sp_kwargs) + batch = types.SimpleNamespace( + sampling_params=sp, num_frames=num_frames, extra={} + ) + gen = torch.Generator(device=self.device).manual_seed(0) + self.stage._prepare_action_latents(batch, gen, self.device, self.dtype) + return batch + + def test_forward_dynamics_clean_conditioning(self): + batch = self._run( + action_mode="forward_dynamics", + domain_name="agibotworld", + action=[[0.1] * 29 for _ in range(16)], + ) + # action_chunk_size = num_frames - 1 = 16; padded to action_dim 64. + self.assertEqual(tuple(batch.action_latents.shape), (1, 16, 64)) + # raw_action_dim is derived from the embodiment (agibotworld -> 29). + self.assertEqual(batch.extra["raw_action_dim"], 29) + self.assertEqual(batch.extra["action_start_frame_offset"], 1) + self.assertEqual( + int(batch.extra["action_domain_ids"][0]), + EMBODIMENT_TO_DOMAIN_ID["agibotworld"], + ) + # forward_dynamics: action is clean conditioning -> velocity mask all zero. + self.assertTrue(torch.all(batch.extra["action_velocity_mask"] == 0)) + + def test_policy_denoises_from_noise(self): + batch = self._run( + action_mode="policy", domain_name="droid_lerobot", raw_action_dim=10 + ) + self.assertEqual(tuple(batch.action_latents.shape), (1, 16, 64)) + self.assertEqual(batch.extra["raw_action_dim"], 10) + # policy: action fully denoised -> velocity mask all one. + self.assertTrue(torch.all(batch.extra["action_velocity_mask"] == 1)) + # padding dims beyond raw_action_dim start at zero. + self.assertTrue(torch.all(batch.action_latents[:, :, 10:] == 0)) + + def test_inverse_dynamics_denoises_from_noise(self): + batch = self._run( + num_frames=61, + action_mode="inverse_dynamics", + domain_name="av", + raw_action_dim=9, + ) + self.assertEqual(tuple(batch.action_latents.shape), (1, 60, 64)) + self.assertTrue(torch.all(batch.extra["action_velocity_mask"] == 1)) + + def test_forward_dynamics_requires_action(self): + with self.assertRaises(ValueError): + self._run(action_mode="forward_dynamics", domain_name="agibotworld") + + def test_policy_requires_raw_action_dim(self): + # policy has no input action to infer from, so it needs raw_action_dim + # from either the embodiment or an explicit value. With only a numeric + # domain_id (no embodiment name) and no raw_action_dim, it must raise. + with self.assertRaises(ValueError): + self._run(action_mode="policy", domain_id=0) + + def test_policy_raw_action_dim_from_embodiment(self): + # droid_lerobot -> 10, so policy no longer needs an explicit raw dim. + batch = self._run(action_mode="policy", domain_name="droid_lerobot") + self.assertEqual(batch.extra["raw_action_dim"], 10) + + def test_unknown_action_mode_raises(self): + with self.assertRaises(ValueError): + self._run(action_mode="teleport", domain_id=0) + + +class TestCosmos3ModalitySamplingParams(unittest.TestCase): + """Sound / V2V / action sampling-param fields and defaults.""" + + def test_sound_duration_default_and_set(self): + self.assertEqual(Cosmos3SamplingParams(prompt="t").sound_duration, 0.0) + self.assertEqual( + Cosmos3SamplingParams(prompt="t", sound_duration=3.0).sound_duration, 3.0 + ) + + def test_v2v_fields(self): + sp = Cosmos3SamplingParams( + prompt="t", video_path="in.mp4", condition_frame_indexes=[0, 1] + ) + self.assertEqual(sp.video_path, "in.mp4") + self.assertEqual(sp.condition_frame_indexes, [0, 1]) + self.assertEqual(sp.condition_video_keep, "first") + + def test_action_fields_default_none(self): + sp = Cosmos3SamplingParams(prompt="t") + for field in ( + "action_mode", + "domain_id", + "domain_name", + "raw_action_dim", + "action_fps", + "action", + ): + self.assertIsNone(getattr(sp, field)) + + if __name__ == "__main__": unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index b0255d7ad..14a18f9b9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -52,6 +52,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( ModelOptFp4Config, + ModelOptFp8Config, _prepare_nvfp4_weight_bytes, ) from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( @@ -456,6 +457,37 @@ class TestTransformerQuantHelpers(unittest.TestCase): ["single_transformer_blocks.*.proj_mlp*"], ) + def test_modelopt_fp8_hf_config_uses_general_modelopt_fp8(self): + config = get_quant_config( + { + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "FP8", + "ignore": ["vae2llm", "llm2vae"], + } + }, + "/unused/component/path", + quant_ignore_remap={"vae2llm": "proj_in", "llm2vae": "proj_out"}, + ) + + self.assertIsInstance(config, ModelOptFp8Config) + self.assertEqual(config.exclude_modules, ["proj_in", "proj_out"]) + + def test_modelopt_fp8_explicit_config_uses_general_modelopt_fp8(self): + config = get_quant_config( + { + "quantization_config": { + "quant_method": "modelopt_fp8", + "quant_algo": "FP8", + "ignore": ["proj_out"], + } + }, + "/unused/component/path", + ) + + self.assertIsInstance(config, ModelOptFp8Config) + self.assertEqual(config.exclude_modules, ["proj_out"]) + @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_rank", return_value=0) @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_size", return_value=1) @patch( diff --git a/python/sglang/multimodal_gen/test/unit/test_wan_pipeline_config.py b/python/sglang/multimodal_gen/test/unit/test_wan_pipeline_config.py new file mode 100644 index 000000000..39fc1d71c --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_wan_pipeline_config.py @@ -0,0 +1,19 @@ +from types import SimpleNamespace + +import torch + +from sglang.multimodal_gen.configs.pipeline_configs.wan import WanT2V480PConfig + + +def test_wan_prompt_embed_accessors_return_transformer_tensor(): + prompt_embeds = torch.empty(1, 2, 3) + negative_prompt_embeds = torch.empty(1, 2, 3) + batch = SimpleNamespace( + prompt_embeds=[prompt_embeds], + negative_prompt_embeds=[negative_prompt_embeds], + ) + + config = WanT2V480PConfig() + + assert config.get_pos_prompt_embeds(batch) is prompt_embeds + assert config.get_neg_prompt_embeds(batch) is negative_prompt_embeds