diff --git a/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py index 6a8365465..6356cbd05 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py @@ -6,7 +6,7 @@ from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig -def _build_cosmos3_param_names_mapping() -> dict: +def _build_cosmos3_param_names_mapping(gated_mlp: bool = True) -> dict: """Map diffusers-format Cosmos3 weights to the sglang model namespace. Source keys (diffusers transformer ckpt) → target keys (sglang model): @@ -15,23 +15,26 @@ def _build_cosmos3_param_names_mapping() -> dict: layers.X.input_layernorm_moe_gen.weight -> gen_layers.X.input_layernorm.weight layers.X.self_attn.{to_q,to_k,to_v}.weight -> language_model.layers.X.self_attn.to_qkv.weight (concat dim 0) layers.X.self_attn.{add_q,add_k,add_v}_proj.weight -> gen_layers.X.cross_attention.to_qkv.weight (concat dim 0) - layers.X.mlp.{gate,up}_proj.weight -> language_model.layers.X.mlp.gate_up_proj.weight (concat dim 0) - layers.X.mlp_moe_gen.{gate,up}_proj.weight -> gen_layers.X.mlp.gate_up_proj.weight (concat dim 0) 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) + ``gated_mlp`` selects the MLP weight layout: SwiGLU checkpoints merge + ``mlp.{gate,up}_proj`` into a single ``gate_up_proj``; dense (squared-ReLU) + checkpoints ship only ``mlp.{up,down}_proj``, which flow through the + catch-all unchanged. + 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; both are skipped. - Audio and action keys (``audio_proj_*``, ``action_proj_*``, modality - embeds) pass through unchanged. + `norm.weight` and `lm_head.weight` are inherited from the text-pretrained + backbone 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. + mapping = { + # Inherited from text pretraining; unused at diffusion inference. r"^lm_head\.weight$": "", r"^norm\.weight$": "", # Top-level norms / embeddings. @@ -62,52 +65,61 @@ def _build_cosmos3_param_names_mapping() -> dict: r"^layers\.(\d+)\.self_attn\.norm_added_k\.(.*)$": r"gen_layers.\1.cross_attention.norm_k.\2", r"^layers\.(\d+)\.input_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.input_layernorm.\2", r"^layers\.(\d+)\.post_attention_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.post_attention_layernorm.\2", + } + + if gated_mlp: # GEN MLP gate/up merge into MergedColumnParallelLinear gate_up_proj. # Must precede the mlp_moe_gen catch-all below. - r"^layers\.(\d+)\.mlp_moe_gen\.gate_proj\.(.*)$": ( + mapping[r"^layers\.(\d+)\.mlp_moe_gen\.gate_proj\.(.*)$"] = ( r"gen_layers.\1.mlp.gate_up_proj.\2", 0, 2, - ), - r"^layers\.(\d+)\.mlp_moe_gen\.up_proj\.(.*)$": ( + ) + mapping[r"^layers\.(\d+)\.mlp_moe_gen\.up_proj\.(.*)$"] = ( r"gen_layers.\1.mlp.gate_up_proj.\2", 1, 2, - ), - r"^layers\.(\d+)\.mlp_moe_gen\.(.*)$": r"gen_layers.\1.mlp.\2", - # UND pathway: Q/K/V merge into to_qkv; remaining attention keys - # (to_out, norm_q, norm_k) and layernorms pass through the catch-all. - r"^layers\.(\d+)\.self_attn\.to_q\.(.*)$": ( - r"language_model.layers.\1.self_attn.to_qkv.\2", - 0, - 3, - ), - r"^layers\.(\d+)\.self_attn\.to_k\.(.*)$": ( - r"language_model.layers.\1.self_attn.to_qkv.\2", - 1, - 3, - ), - r"^layers\.(\d+)\.self_attn\.to_v\.(.*)$": ( - r"language_model.layers.\1.self_attn.to_qkv.\2", - 2, - 3, - ), + ) + + # GEN MLP catch-all: dense up/down_proj pass through unchanged. + mapping[r"^layers\.(\d+)\.mlp_moe_gen\.(.*)$"] = r"gen_layers.\1.mlp.\2" + + # UND pathway: Q/K/V merge into to_qkv; remaining attention keys + # (to_out, norm_q, norm_k) and layernorms pass through the catch-all. + mapping[r"^layers\.(\d+)\.self_attn\.to_q\.(.*)$"] = ( + r"language_model.layers.\1.self_attn.to_qkv.\2", + 0, + 3, + ) + mapping[r"^layers\.(\d+)\.self_attn\.to_k\.(.*)$"] = ( + r"language_model.layers.\1.self_attn.to_qkv.\2", + 1, + 3, + ) + mapping[r"^layers\.(\d+)\.self_attn\.to_v\.(.*)$"] = ( + r"language_model.layers.\1.self_attn.to_qkv.\2", + 2, + 3, + ) + + if gated_mlp: # UND MLP gate/up merge into MergedColumnParallelLinear gate_up_proj. # Must precede the layers catch-all below. - r"^layers\.(\d+)\.mlp\.gate_proj\.(.*)$": ( + mapping[r"^layers\.(\d+)\.mlp\.gate_proj\.(.*)$"] = ( r"language_model.layers.\1.mlp.gate_up_proj.\2", 0, 2, - ), - r"^layers\.(\d+)\.mlp\.up_proj\.(.*)$": ( + ) + mapping[r"^layers\.(\d+)\.mlp\.up_proj\.(.*)$"] = ( r"language_model.layers.\1.mlp.gate_up_proj.\2", 1, 2, - ), - # UND pathway: layernorms + remaining attention/mlp keys pass through - # under language_model.layers namespace. - r"^layers\.(\d+)\.(.*)$": r"language_model.layers.\1.\2", - } + ) + + # UND pathway: layernorms + remaining attention/mlp keys pass through + # under language_model.layers namespace. + mapping[r"^layers\.(\d+)\.(.*)$"] = r"language_model.layers.\1.\2" + return mapping @dataclass @@ -130,6 +142,12 @@ class Cosmos3VideoArchConfig(DiTArchConfig): num_key_value_heads: int = 8 # GQA head_dim: int = 128 intermediate_size: int = 12288 + # "relu2" selects a dense squared-ReLU MLP; anything else is SwiGLU. + hidden_act: str = "silu" + # Per-head QK-norm on the UND self-attention (text) pathway. + qk_norm_for_text: bool = True + # Normalize the UND keys before they feed the GEN cross-attention. + use_und_k_norm_for_gen: bool = False # Latent space configuration latent_patch_size: int = 2 @@ -180,6 +198,10 @@ class Cosmos3VideoArchConfig(DiTArchConfig): default_factory=lambda: {"vae2llm": "proj_in", "llm2vae": "proj_out"} ) + @property + def gated_mlp(self) -> bool: + return self.hidden_act != "relu2" + def __post_init__(self): super().__post_init__() # Diffusers configs nest the mrope sizes under `rope_scaling`; lift it. @@ -189,6 +211,9 @@ class Cosmos3VideoArchConfig(DiTArchConfig): self.num_channels_latents = self.out_channels # Patch latent dimension: (patch_size^2) * latent_channel self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel + # The MLP weight layout depends on the activation, which the checkpoint + # may override; rebuild the mapping after arch values are applied. + self.param_names_mapping = _build_cosmos3_param_names_mapping(self.gated_mlp) @dataclass diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py index 0a4034989..8065e2503 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py @@ -8,6 +8,8 @@ in the stages from ``num_frames`` and ``image_path``; T2I overrides ``data_type`` to ``IMAGE`` in :meth:`SamplingParams._adjust`. """ +import functools +import os from dataclasses import dataclass, field from sglang.multimodal_gen.configs.models import DiTConfig, VAEConfig @@ -18,6 +20,63 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import ( PipelineConfig, ) +COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense" + + +@functools.lru_cache(maxsize=None) +def is_edge_checkpoint(model_path: str) -> bool: + """Whether the checkpoint is the Edge (dense) variant. + + Read from the transformer config rather than the loaded arch so the answer + is available before the weights are on device (e.g. when resolving sampling + defaults in the client process). + """ + from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + get_diffusers_component_config, + ) + + config = get_diffusers_component_config( + component_path=os.path.join(model_path, "transformer") + ) + return ( + config.get("backbone_type") == COSMOS3_EDGE_BACKBONE_TYPE + or config.get("hidden_act") == "relu2" + ) + + +@functools.lru_cache(maxsize=None) +def _distilled_sampler_config(model_path: str) -> dict | None: + """The fixed-step sampler config for a distilled checkpoint, else ``None``. + + Distillation is a scheduler-only change: the checkpoint ships a + ``FlowMatchEulerDiscreteScheduler`` with an explicit fixed-step sigma + schedule instead of the multi-step FlowUniPC the other variants use. + """ + from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( + get_diffusers_component_config, + ) + + config = get_diffusers_component_config( + component_path=os.path.join(model_path, "scheduler") + ) + if config.get("_class_name") != "FlowMatchEulerDiscreteScheduler": + return None + sampler = config.get("fixed_step_sampler_config") + if not sampler or not sampler.get("t_list"): + return None + return sampler + + +def is_distilled_checkpoint(model_path: str) -> bool: + """Whether the checkpoint is a few-step distilled variant.""" + return _distilled_sampler_config(model_path) is not None + + +def get_distilled_sigmas(model_path: str) -> list[float] | None: + """The explicit fixed-step sigma schedule for a distilled checkpoint.""" + sampler = _distilled_sampler_config(model_path) + return list(sampler["t_list"]) if sampler is not None else None + @dataclass class Cosmos3Config(PipelineConfig): @@ -60,6 +119,11 @@ class Cosmos3Config(PipelineConfig): # names a server-side file. ``None`` disables normalization. action_stats_path: str | None = None + # Pre-computed once in update_config_from_dict from the resolved model_path. + # None until that point (e.g. in unit-test mocks that never call update_config_from_dict). + is_edge: bool | None = None + distilled_sigmas: list[float] | None = None + def __post_init__(self): self.vae_config.arch_config.z_dim = 48 # Encoder is needed for I2V; T2V/T2I never invoke it. @@ -70,6 +134,17 @@ class Cosmos3Config(PipelineConfig): self.vae_config.use_parallel_encode = False self.vae_config.use_parallel_decode = True + def update_config_from_dict(self, args, prefix: str = "") -> None: + super().update_config_from_dict(args, prefix) + # model_path is only populated here, after construction. Compute + # checkpoint variant flags once so per-request code reads them from the + # config rather than re-downloading the scheduler subfolder each time. + if self.model_path: + self.distilled_sigmas = get_distilled_sigmas(self.model_path) + self.is_edge = is_edge_checkpoint(self.model_path) + if self.distilled_sigmas is not None: + self.scheduler_class_override = None + def adjust_num_frames(self, num_frames: int) -> int: """Round ``num_frames`` so ``(n - 1) % 4 == 0`` for the VAE. diff --git a/python/sglang/multimodal_gen/configs/sample/cosmos3.py b/python/sglang/multimodal_gen/configs/sample/cosmos3.py index e55f46d32..d0c1be460 100644 --- a/python/sglang/multimodal_gen/configs/sample/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/sample/cosmos3.py @@ -17,17 +17,48 @@ from sglang.multimodal_gen.configs.sample.sampling_params import ( SamplingParams, ) +COSMOS3_DEFAULT_GUIDANCE_SCALE = 4.0 +COSMOS3_EDGE_T2I_GUIDANCE_SCALE = 7.0 +COSMOS3_EDGE_T2V_GUIDANCE_SCALE = 5.0 +COSMOS3_EDGE_T2V_WIDTH = 832 +COSMOS3_EDGE_T2V_HEIGHT = 480 +COSMOS3_EDGE_T2I_SIZE = 640 + +# Image generation applies guidance only over a high-noise window; guiding the +# low-noise steps degrades sample quality (Kynkaeaenniemi et al. 2024). Video +# modes guide at every step. +COSMOS3_T2I_GUIDANCE_INTERVAL = (400.0, 1000.0) + +# Edge is trained at 256p/480p only; larger frames push the spatial mRoPE grid +# past its trained range and shatter the output. +COSMOS3_EDGE_SUPPORTED_RESOLUTIONS = [ + (832, 480), + (480, 832), + (640, 480), + (480, 640), + (480, 480), + (640, 640), + (448, 256), + (256, 448), + (256, 256), +] + @dataclass class Cosmos3SamplingParams(SamplingParams): - """Cosmos3 sampling parameters (T2V defaults; also used for I2V / V2V / T2I).""" + """Cosmos3 sampling parameters (T2V defaults; also used for I2V / V2V / T2I). - height: int = 720 - width: int = 1280 + ``height``/``width`` default to ``None`` so the variant (Edge vs. base) can + pick the right resolution at request time in + :meth:`_resolve_variant_defaults`. + """ + + height: int | None = None + width: int | None = None num_frames: int = 81 fps: int = 24 - guidance_scale: float = 4.0 + guidance_scale: float = COSMOS3_DEFAULT_GUIDANCE_SCALE num_inference_steps: int = 35 negative_prompt: str = "" @@ -50,6 +81,7 @@ class Cosmos3SamplingParams(SamplingParams): (832, 480), (480, 832), (1024, 1024), + (640, 640), ] ) @@ -72,6 +104,18 @@ class Cosmos3SamplingParams(SamplingParams): action_normalization: str = "quantile" def _adjust(self, server_args) -> None: + # adjust distil and edge args — read from the pre-computed config fields + # so no checkpoint download happens at request time. + pipeline_config = server_args.pipeline_config + distilled_sigmas = pipeline_config.distilled_sigmas + if distilled_sigmas is not None: + self.num_inference_steps = len(distilled_sigmas) + self._resolve_variant_defaults( + bool(pipeline_config.is_edge), + is_distilled=distilled_sigmas is not None, + ) + + # adjust action args action_output = False if self.action_mode is not None: self.action_mode = str(self.action_mode).strip().lower() @@ -98,6 +142,41 @@ class Cosmos3SamplingParams(SamplingParams): self.output_file_name = None self.output_compression = 0 + def _guidance_is_explicit(self) -> bool: + explicit = getattr(self, "_explicit_fields", None) + return explicit is not None and "guidance_scale" in explicit + + def _resolve_variant_defaults( + self, is_edge: bool, is_distilled: bool = False + ) -> None: + """Fill unset resolution/guidance with the variant's defaults. + + Base resolution defaulting (``supported_resolutions[0]``) covers the + non-Edge path; only Edge and guidance need explicit handling here. + """ + is_t2i = self.num_frames == 1 + if is_distilled: + # Guidance is distilled into the model; run a single forward. + self.guidance_scale = 1.0 + elif is_edge and not self._guidance_is_explicit(): + self.guidance_scale = ( + COSMOS3_EDGE_T2I_GUIDANCE_SCALE + if is_t2i + else COSMOS3_EDGE_T2V_GUIDANCE_SCALE + ) + if is_t2i and not is_distilled and self.guidance_interval is None: + self.guidance_interval = COSMOS3_T2I_GUIDANCE_INTERVAL + if is_edge: + self.supported_resolutions = COSMOS3_EDGE_SUPPORTED_RESOLUTIONS + if self.height is None and self.width is None: + if is_t2i: + self.width = self.height = COSMOS3_EDGE_T2I_SIZE + else: + self.width, self.height = ( + COSMOS3_EDGE_T2V_WIDTH, + COSMOS3_EDGE_T2V_HEIGHT, + ) + def _set_output_file_name(self) -> None: # Action outputs never need a visual filename. This also avoids hashing # in-memory observation images while base visual adjustment is running. diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 2929f3286..df4931c4d 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -1088,8 +1088,8 @@ def _register_configs(): # Cosmos3 — single checkpoint serves T2V, I2V, and T2I. Mode is dispatched # per-request inside the pipeline from ``num_frames`` and ``image_path``. - # Both Nano (16B) and Super (64B) share the same pipeline; arch dimensions - # come from ``transformer/config.json`` via ``update_model_arch``. + # All variants share the same pipeline; arch dimensions (size, activation, + # QK-norm) come from ``transformer/config.json`` via ``update_model_arch``. register_configs( sampling_param_cls=Cosmos3SamplingParams, pipeline_config_cls=Cosmos3Config, @@ -1099,8 +1099,11 @@ def _register_configs(): "nvidia/Cosmos3-Super", "nvidia/Cosmos3-Super-Text2Image", "nvidia/Cosmos3-Super-Image2Video", + "nvidia/Cosmos3-Edge", ], - model_detectors=[lambda hf_id: "cosmos3omnidiffuserspipeline" in hf_id.lower()], + # Match both the new ``Cosmos3OmniPipeline`` and the legacy + # ``Cosmos3OmniDiffusersPipeline`` ``_class_name`` (diffusers rename). + model_detectors=[lambda hf_id: "cosmos3omni" in hf_id.lower()], ) # SANA diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/action/api.py b/python/sglang/multimodal_gen/runtime/entrypoints/action/api.py index dc2a7ea5f..db35e720a 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/action/api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/action/api.py @@ -2,8 +2,14 @@ from __future__ import annotations +import json +import os +from contextlib import nullcontext +from typing import Any + from fastapi import APIRouter, HTTPException, Request, Response, WebSocket +from sglang.multimodal_gen.configs.sample.sampling_params import generate_request_id from sglang.multimodal_gen.runtime.entrypoints.action.protocol import ( action_generation_response, action_metadata, @@ -15,11 +21,41 @@ from sglang.multimodal_gen.runtime.entrypoints.action.protocol import ( from sglang.multimodal_gen.runtime.entrypoints.action.ws_utils import ( run_action_msgpack_ws, ) +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( + flatten_extra_params, + save_image_to_path, + temp_dir_if_disabled, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.srt.utils.json_response import orjson_response router = APIRouter(prefix="/v1/actions", tags=["actions"]) +_ACTION_VIDEO_EXTENSIONS = { + ".avi", + ".gif", + ".m4v", + ".mkv", + ".mov", + ".mp4", + ".mpeg", + ".mpg", + ".webm", +} + +_ACTION_FORM_NON_PARAMETER_FIELDS = { + "extra_body", + "extra_params", + "id", + "image_reference", + "input_reference", + "prompt", + "reference_url", + "request_id", + "task", + "video_reference", +} + def _wants_msgpack(request: Request) -> bool: content_type = request.headers.get("content-type", "").lower() @@ -40,26 +76,154 @@ def _prefer_numpy_output(payload: dict) -> None: runtime.setdefault("output_format", "numpy") +def _parse_form_value(value: Any) -> Any: + if not isinstance(value, str): + return value + if not value.strip(): + return None + try: + return json.loads(value) + except Exception: + return value + + +def _parse_extra_params(value: Any, field_name: str) -> dict[str, Any]: + if value in (None, ""): + return {} + try: + parsed = json.loads(value) if isinstance(value, str) else value + except (json.JSONDecodeError, TypeError) as exc: + raise ValueError(f"{field_name} is not valid JSON") from exc + if not isinstance(parsed, dict): + raise ValueError(f"{field_name} must be a JSON object") + return flatten_extra_params(dict(parsed)) + + +def _is_form_upload(value: Any) -> bool: + return callable(getattr(value, "read", None)) and hasattr(value, "filename") + + +def _is_probably_video_upload(value: Any) -> bool: + content_type = (getattr(value, "content_type", "") or "").lower() + if content_type.startswith("video/"): + return True + filename = getattr(value, "filename", None) + if not filename: + return False + filename = str(filename).split("?", 1)[0].split("#", 1)[0] + return os.path.splitext(filename)[1].lower() in _ACTION_VIDEO_EXTENSIONS + + +async def _save_action_upload( + upload: Any, + request_id: str, + uploads_dir: str, + *, + fallback_name: str, +) -> str: + filename = getattr(upload, "filename", None) or fallback_name + target_path = os.path.join(uploads_dir, f"{request_id}_{filename}") + return await save_image_to_path(upload, target_path) + + +async def _multipart_action_payload(request: Request, uploads_dir: str) -> dict: + raw_form = await request.form() + request_id = ( + raw_form.get("request_id") or raw_form.get("id") or generate_request_id() + ) + + parameters: dict[str, Any] = {} + for key, value in raw_form.multi_items(): + if key in _ACTION_FORM_NON_PARAMETER_FIELDS or _is_form_upload(value): + continue + parameters[key] = _parse_form_value(value) + + parameters.update(_parse_extra_params(raw_form.get("extra_body"), "extra_body")) + parameters.update(_parse_extra_params(raw_form.get("extra_params"), "extra_params")) + flatten_extra_params(parameters) + + action_mode = str(parameters.get("action_mode", "policy")).strip().lower() + observation: dict[str, Any] = {} + + upload_fields = ( + ("video_reference", "video"), + ("image_reference", "image"), + ("input_reference", None), + ) + for field_name, forced_observation_key in upload_fields: + upload = raw_form.get(field_name) + if not _is_form_upload(upload): + continue + saved_path = await _save_action_upload( + upload, + str(request_id), + uploads_dir, + fallback_name=field_name, + ) + observation_key = forced_observation_key + if observation_key is None: + observation_key = ( + "video" + if action_mode == "inverse_dynamics" + or _is_probably_video_upload(upload) + else "input_reference" + ) + observation[observation_key] = saved_path + break + + for form_key, observation_key in ( + ("video_path", "video"), + ("video_url", "video"), + ("reference_url", "video" if action_mode == "inverse_dynamics" else "image"), + ( + "input_reference", + "video" if action_mode == "inverse_dynamics" else "input_reference", + ), + ): + value = raw_form.get(form_key) + if isinstance(value, str) and value and observation_key not in observation: + observation[observation_key] = value + + return { + "request_id": request_id, + "input": { + "task": raw_form.get("prompt") or raw_form.get("task") or "", + "observation": observation, + }, + "parameters": parameters, + } + + @router.post("/generations") async def create_action_generation(request: Request): server_args: ServerArgs = request.app.state.server_args + content_type = request.headers.get("content-type", "").lower() + is_multipart = "multipart/form-data" in content_type + input_context = ( + temp_dir_if_disabled(getattr(server_args, "input_save_path", None)) + if is_multipart + else nullcontext(None) + ) try: - if "msgpack" in request.headers.get("content-type", "").lower(): - payload = unpack_msgpack(await request.body()) - else: - payload = await request.json() - wants_msgpack = _wants_msgpack(request) - if wants_msgpack: - _prefer_numpy_output(payload) - output = await infer_action(payload, server_args) - if _response_format(payload) == "raw": - response = action_raw_response(output, preserve_numpy=wants_msgpack) - else: - response = action_generation_response( - output, - server_args, - preserve_numpy=wants_msgpack, - ) + with input_context as uploads_dir: + if is_multipart: + payload = await _multipart_action_payload(request, uploads_dir) + elif "msgpack" in content_type: + payload = unpack_msgpack(await request.body()) + else: + payload = await request.json() + wants_msgpack = _wants_msgpack(request) + if wants_msgpack: + _prefer_numpy_output(payload) + output = await infer_action(payload, server_args) + if _response_format(payload) == "raw": + response = action_raw_response(output, preserve_numpy=wants_msgpack) + else: + response = action_generation_response( + output, + server_args, + preserve_numpy=wants_msgpack, + ) except ValueError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc if wants_msgpack: 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 b28244c11..ffe213bb5 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -683,6 +683,15 @@ async def create_video( selected = value if value is not None else extra_from_form.get(name) return _parse_form_extra_value(selected) + def form_text_value(name: str, value: Any) -> Any: + """Resolve a text field without JSON-decoding it. + + Some models take a serialized JSON object as the prompt text, which + ``_parse_form_extra_value`` would turn back into a dict and fail + request validation. + """ + return value if value is not None else extra_from_form.get(name) + request_field_names = set(VideoGenerationsRequest.model_fields) extra_request_fields = { key: value @@ -708,7 +717,7 @@ async def create_video( num_frames=num_frames_val, seed=form_value("seed", seed), generator_device=form_value("generator_device", generator_device), - negative_prompt=form_value("negative_prompt", negative_prompt), + negative_prompt=form_text_value("negative_prompt", negative_prompt), num_inference_steps=form_value("num_inference_steps", num_inference_steps), guidance_scale=form_value("guidance_scale", guidance_scale), guidance_scale_2=form_value("guidance_scale_2", guidance_scale_2), 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 e0de44a5b..44b5b4139 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 @@ -1,3 +1,5 @@ +import inspect + from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, ) @@ -11,6 +13,26 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +def _supported_init_kwargs(scheduler_cls, config: dict) -> dict: + """Drop scheduler_config keys the constructor can't accept. + + Some checkpoints (e.g. distilled fixed-step schedulers) carry keys in + ``scheduler_config.json`` that are consumed downstream, not by the + scheduler ``__init__``. Constructors taking ``**kwargs`` accept anything + and are left untouched; otherwise unknown non-private keys would raise + ``TypeError``, so they are dropped (private ``_``-keys are handled by + ``register_to_config``). + """ + params = inspect.signature(scheduler_cls.__init__).parameters + if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()): + return config + dropped = [k for k in config if k not in params and not k.startswith("_")] + if dropped: + logger.debug("Ignoring unsupported scheduler config keys: %s", dropped) + return {k: v for k, v in config.items() if k not in dropped} + return config + + class SchedulerLoader(ComponentLoader): """Loader for scheduler.""" @@ -41,7 +63,7 @@ class SchedulerLoader(ComponentLoader): scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name) - scheduler = scheduler_cls(**config) + scheduler = scheduler_cls(**_supported_init_kwargs(scheduler_cls, config)) if server_args.pipeline_config.flow_shift is not None: scheduler.set_shift(server_args.pipeline_config.flow_shift) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index 45cfd0d1a..ee0916727 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py @@ -30,6 +30,7 @@ from sglang.multimodal_gen.runtime.layers.layernorm import ( apply_qk_norm_rope, ) from sglang.multimodal_gen.runtime.layers.linear import ( + ColumnParallelLinear, MergedColumnParallelLinear, ReplicatedLinear, RowParallelLinear, @@ -432,6 +433,57 @@ class Cosmos3GatedMLP(nn.Module): return out +class Cosmos3DenseMLP(nn.Module): + """Dense MLP with a squared-ReLU activation: ``down(relu(up(x)) ** 2)``.""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + prefix: str = "", + quant_config: QuantizationConfig | None = None, + ): + super().__init__() + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + gather_output=False, + quant_config=quant_config, + prefix=add_prefix("up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + input_is_parallel=True, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + up, _ = self.up_proj(x) + up = F.relu(up) + out, _ = self.down_proj(up * up) + return out + + +def _build_mlp( + hidden_act: str, + hidden_size: int, + intermediate_size: int, + prefix: str, + quant_config: QuantizationConfig | None, +) -> nn.Module: + mlp_cls = Cosmos3DenseMLP if hidden_act == "relu2" else Cosmos3GatedMLP + return mlp_cls( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + prefix=prefix, + quant_config=quant_config, + ) + + # ----------------------------------------------------------------------------- # Cosmos3 UND Causal Attention # ----------------------------------------------------------------------------- @@ -446,6 +498,9 @@ class Cosmos3CausalAttention(nn.Module): num_attention_heads: int, num_key_value_heads: int, head_dim: int, + qk_norm: bool = True, + use_k_norm_und_for_gen: bool = False, + rms_norm_eps: float = 1e-6, prefix: str = "", quant_config: QuantizationConfig | None = None, ): @@ -454,6 +509,7 @@ class Cosmos3CausalAttention(nn.Module): self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.head_dim = head_dim + self.qk_norm = qk_norm self.tp_size = get_tp_world_size() if num_attention_heads % self.tp_size != 0: raise ValueError( @@ -487,9 +543,16 @@ class Cosmos3CausalAttention(nn.Module): prefix=add_prefix("to_out", prefix), ) - # Per-head QK norm. - self.norm_q = RMSNorm(head_dim, eps=1e-6) - self.norm_k = RMSNorm(head_dim, eps=1e-6) + # Per-head QK norm (optional; some backbones omit it on text). + if qk_norm: + self.norm_q = RMSNorm(head_dim, eps=1e-6) + self.norm_k = RMSNorm(head_dim, eps=1e-6) + + # Dense-backbone variants normalize the keys handed to the GEN + # cross-attention separately from the reasoner self-attention keys. + self.k_norm_und_for_gen = ( + RMSNorm(head_dim, eps=rms_norm_eps) if use_k_norm_und_for_gen else None + ) def forward( self, @@ -526,14 +589,27 @@ class Cosmos3CausalAttention(nn.Module): :, ] - q = F.rms_norm( - q, (self.head_dim,), self.norm_q.weight, self.norm_q.variance_epsilon - ) - k = F.rms_norm( - k, (self.head_dim,), self.norm_k.weight, self.norm_k.variance_epsilon - ) + k_und = k + if self.qk_norm: + q = F.rms_norm( + q, (self.head_dim,), self.norm_q.weight, self.norm_q.variance_epsilon + ) + k = F.rms_norm( + k, (self.head_dim,), self.norm_k.weight, self.norm_k.variance_epsilon + ) q, k = _apply_qwen3_rope_from_cache(q, k, cos_sin_cache) + if self.k_norm_und_for_gen is not None: + k_gen = F.rms_norm( + k_und, + (self.head_dim,), + self.k_norm_und_for_gen.weight, + self.k_norm_und_for_gen.variance_epsilon, + ) + _, k_gen = _apply_qwen3_rope_from_cache(q, k_gen, cos_sin_cache) + else: + k_gen = k + out = F.scaled_dot_product_attention( q.transpose(1, 2), k.transpose(1, 2), @@ -544,7 +620,7 @@ class Cosmos3CausalAttention(nn.Module): out = out.transpose(1, 2).reshape(batch_size, seq_len, -1) out, _ = self.to_out(out) - return out, k, v + return out, k_gen, v # ----------------------------------------------------------------------------- @@ -697,6 +773,9 @@ class Cosmos3UndDecoderLayer(nn.Module): head_dim: int, intermediate_size: int, rms_norm_eps: float, + hidden_act: str, + qk_norm: bool, + use_k_norm_und_for_gen: bool, layer_idx: int, prefix: str = "", quant_config: QuantizationConfig | None = None, @@ -709,12 +788,16 @@ class Cosmos3UndDecoderLayer(nn.Module): num_attention_heads=num_attention_heads, num_key_value_heads=num_key_value_heads, head_dim=head_dim, + qk_norm=qk_norm, + use_k_norm_und_for_gen=use_k_norm_und_for_gen, + rms_norm_eps=rms_norm_eps, prefix=add_prefix("self_attn", prefix), quant_config=quant_config, ) self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.mlp = Cosmos3GatedMLP( + self.mlp = _build_mlp( + hidden_act=hidden_act, hidden_size=hidden_size, intermediate_size=intermediate_size, prefix=add_prefix("mlp", prefix), @@ -763,6 +846,7 @@ class Cosmos3GenDecoderLayer(nn.Module): head_dim: int, intermediate_size: int, rms_norm_eps: float, + hidden_act: str, layer_idx: int, prefix: str = "", quant_config: QuantizationConfig | None = None, @@ -782,7 +866,8 @@ class Cosmos3GenDecoderLayer(nn.Module): ) self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.mlp = Cosmos3GatedMLP( + self.mlp = _build_mlp( + hidden_act=hidden_act, hidden_size=hidden_size, intermediate_size=intermediate_size, prefix=add_prefix("mlp", prefix), @@ -844,6 +929,9 @@ class Cosmos3LanguageModel(nn.Module): rms_norm_eps: float, rope_theta: float, mrope_section: tuple[int, int, int], + hidden_act: str, + qk_norm: bool, + use_k_norm_und_for_gen: bool, quant_config: QuantizationConfig | None = None, ): super().__init__() @@ -869,6 +957,9 @@ class Cosmos3LanguageModel(nn.Module): head_dim=head_dim, intermediate_size=intermediate_size, rms_norm_eps=rms_norm_eps, + hidden_act=hidden_act, + qk_norm=qk_norm, + use_k_norm_und_for_gen=use_k_norm_und_for_gen, layer_idx=i, prefix=f"layers.{i}", quant_config=quant_config, @@ -955,6 +1046,14 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): 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 + self.hidden_act = arch.hidden_act + self.rope_theta = arch.rope_theta + + # The checkpoint may override the activation (and thus the MLP weight + # layout), so bind the arch-derived mappings on the instance. + self.param_names_mapping = arch.param_names_mapping + self.reverse_param_names_mapping = arch.reverse_param_names_mapping + self.lora_param_names_mapping = arch.lora_param_names_mapping # Ulysses sequence parallelism. When CFG-parallel is also enabled # the SP group only spans ranks that share a CFG context (cond or @@ -979,6 +1078,9 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): rms_norm_eps=arch.rms_norm_eps, rope_theta=arch.rope_theta, mrope_section=arch.mrope_section, + hidden_act=arch.hidden_act, + qk_norm=arch.qk_norm_for_text, + use_k_norm_und_for_gen=arch.use_und_k_norm_for_gen, quant_config=quant_config, ) @@ -1051,6 +1153,7 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): head_dim=arch.head_dim, intermediate_size=arch.intermediate_size, rms_norm_eps=arch.rms_norm_eps, + hidden_act=arch.hidden_act, layer_idx=i, prefix=f"gen_layers.{i}", quant_config=quant_config, @@ -1563,6 +1666,11 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): yield linear_target + ".input_scale", in_t.max() for name, tensor in iterator: + # ModelOpt calibration buffers are not inference parameters; the + # runtime FP8 path reads weight_scale/input_scale instead. Drop them + # so they don't reach the fused-linear concat below. + if "_quantizer." in name: + continue target_name, merge_index, num_to_merge = mapping_fn(name) if num_to_merge is None: yield target_name, tensor @@ -1606,9 +1714,8 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): rotary_emb = self.language_model.rotary_emb if rotary_emb.inv_freq.is_meta: dim = rotary_emb.head_dim - rope_theta = 5000000.0 # From config inv_freq = 1.0 / ( - rope_theta + self.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim) ) rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py index 1f25398ec..e38d4d792 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -533,7 +533,12 @@ class FlowMatchEulerDiscreteScheduler( else: if self.config.stochastic_sampling: x0 = sample - current_sigma * model_output - noise = torch.randn_like(sample) + noise = torch.randn( + sample.shape, + generator=generator, + device=sample.device, + dtype=sample.dtype, + ) prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise else: prev_sample = sample + dt * model_output diff --git a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py index 6e65aa449..490a48877 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py @@ -33,7 +33,10 @@ class Cosmos3Pipeline(ComposedPipelineBase): stages from ``batch.data_type`` and ``batch.preprocessed_image``. """ - pipeline_name = "Cosmos3OmniDiffusersPipeline" + # Canonical ``_class_name`` in newer checkpoints' ``model_index.json``. + # Older checkpoints declare ``Cosmos3OmniDiffusersPipeline`` (see the + # back-compat alias below). + pipeline_name = "Cosmos3OmniPipeline" is_video_pipeline = True _required_config_modules = [ @@ -109,4 +112,15 @@ class Cosmos3Pipeline(ComposedPipelineBase): ) -EntryClass = Cosmos3Pipeline +class Cosmos3OmniDiffusersPipeline(Cosmos3Pipeline): + """Back-compat alias for checkpoints whose ``model_index.json`` still + declares the legacy ``_class_name`` ``Cosmos3OmniDiffusersPipeline``. + + Registering both names lets old (Nano/Super) and new (Edge) checkpoints + resolve to the same native pipeline. + """ + + pipeline_name = "Cosmos3OmniDiffusersPipeline" + + +EntryClass = [Cosmos3Pipeline, Cosmos3OmniDiffusersPipeline] 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 ac24aaff9..54ca324a7 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 @@ -79,6 +79,35 @@ COSMOS3_I2V_FLOW_SHIFT = 10.0 COSMOS3_T2V_FLOW_SHIFT = 10.0 COSMOS3_V2V_FLOW_SHIFT = 10.0 COSMOS3_ACTION_FLOW_SHIFT = 10.0 +# Edge uses a single low flow-shift for every video mode (t2v/i2v/v2v). +COSMOS3_EDGE_VIDEO_FLOW_SHIFT = 3.0 + + +def _inject_caption_metadata( + prompt: str, num_frames: int, fps: float, height: int, width: int +) -> str | None: + """Add the generation metadata that Cosmos3's structured captions carry. + + Training captions always ship ``resolution``, plus ``duration``/``fps`` for + video, so a JSON caption without them is out of distribution. Returns + ``None`` when the prompt is not a JSON object; those prompts carry the same + metadata as trailing prose via the duration template instead. + """ + try: + caption = json.loads(prompt) + except (TypeError, ValueError): + return None + if not isinstance(caption, dict): + return None + + caption["resolution"] = {"H": int(height), "W": int(width)} + if num_frames > 1: + caption["duration"] = f"{int(num_frames / fps) if fps > 0 else 0}s" + caption["fps"] = float(fps) + else: + caption.pop("duration", None) + caption.pop("fps", None) + return json.dumps(caption) def _resize_crop_pil( @@ -327,6 +356,15 @@ class Cosmos3TokenizationStage(PipelineStage): use_system_prompt = False use_duration_template = False self.log_info(f"Action prompt: {prompt}") + else: + structured_caption = _inject_caption_metadata( + prompt, num_frames, fps, batch.height, batch.width + ) + if structured_caption is not None: + # The metadata is already in the caption; appending the prose + # template too would state the duration twice. + prompt = structured_caption + use_duration_template = False # Apply duration template if enabled (no temporal concept for T2I). if use_duration_template and not is_image_gen and num_frames > 1: @@ -680,15 +718,14 @@ class Cosmos3TimestepPreparationStage(PipelineStage): super().__init__() self.scheduler = scheduler - 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. - """ + def _default_flow_shift_for_mode(self, batch: Req, is_edge: bool) -> float | None: + """Resolve the per-mode default flow_shift for the request.""" 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 is_edge: + return COSMOS3_EDGE_VIDEO_FLOW_SHIFT if batch.preprocessed_image is not None: return COSMOS3_I2V_FLOW_SHIFT if batch.preprocessed_video is not None: @@ -698,12 +735,32 @@ class Cosmos3TimestepPreparationStage(PipelineStage): def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Prepare scheduler timesteps.""" device = get_local_torch_device() + + pipeline_config = server_args.pipeline_config + distilled_sigmas = pipeline_config.distilled_sigmas + if distilled_sigmas is not None: + # Distilled checkpoints carry an explicit fixed-step sigma schedule + # with the shift already baked in; drive the scheduler from it + # directly (step count == len(sigmas), num_inference_steps ignored). + # Reset shift so set_timesteps does not re-shift the baked-in sigmas. + if hasattr(self.scheduler, "set_shift"): + self.scheduler.set_shift(1.0) + self.scheduler.set_timesteps(sigmas=distilled_sigmas, device=device) + batch.timesteps = self.scheduler.timesteps + self.log_info( + f"Prepared {len(batch.timesteps)} distilled timesteps " + f"(sigmas={distilled_sigmas})" + ) + return batch + num_inference_steps = batch.num_inference_steps flow_shift = getattr(batch, "flow_shift", None) if flow_shift is None: - flow_shift = server_args.pipeline_config.flow_shift + flow_shift = pipeline_config.flow_shift if flow_shift is None: - flow_shift = self._default_flow_shift_for_mode(batch) + flow_shift = self._default_flow_shift_for_mode( + batch, bool(pipeline_config.is_edge) + ) if flow_shift is not None and hasattr(self.scheduler, "set_shift"): self.scheduler.set_shift(float(flow_shift)) @@ -742,6 +799,7 @@ class Cosmos3DenoisingStage(PipelineStage): self.scheduler = scheduler self.server_args = server_args self._logged_parallel_config = False + self._logged_cfg_split = False # Apply torch.compile if enabled if server_args is not None: @@ -917,6 +975,13 @@ class Cosmos3DenoisingStage(PipelineStage): timesteps = batch.timesteps guidance_scale = batch.guidance_scale + # Seed the scheduler's stochastic (SDE) noise from the request seed so it + # is identical on every sequence-parallel rank; otherwise each rank draws + # its own noise and the sharded latents diverge at the shard boundary. + generator = batch.generator + if generator is None and batch.seed is not None: + generator = torch.Generator(device=latents.device).manual_seed(batch.seed) + cond_text_ids = batch.extra["cond_text_ids"] cond_text_mask = batch.extra["cond_text_mask"] uncond_text_ids = batch.extra["uncond_text_ids"] @@ -1036,7 +1101,7 @@ class Cosmos3DenoisingStage(PipelineStage): action_start_frame_offset=action_start_frame_offset, ) else: - noise_pred = self._predict_noise_cfg_batched( + noise_pred = self._predict_noise_cfg( latents=latents, timestep=timestep, cond_text_ids=cond_text_ids, @@ -1047,10 +1112,8 @@ class Cosmos3DenoisingStage(PipelineStage): fps=fps, guidance_scale=effective_scale, noisy_frame_mask=velocity_mask, - max_text_seq_len=max( - batch.extra["cond_text_seq_len"], - batch.extra["uncond_text_seq_len"], - ), + 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, @@ -1102,6 +1165,7 @@ class Cosmos3DenoisingStage(PipelineStage): noise_pred, t, latents, + generator=generator, return_dict=False, )[0] @@ -1155,6 +1219,87 @@ class Cosmos3DenoisingStage(PipelineStage): self.log_info("Denoising complete") return batch + def _predict_noise_cfg( + self, + latents: torch.Tensor, + timestep: torch.Tensor, + cond_text_ids: torch.Tensor, + cond_text_mask: torch.Tensor, + uncond_text_ids: torch.Tensor, + uncond_text_mask: torch.Tensor, + video_shape: tuple[int, int, int], + fps: float, + guidance_scale: float, + cond_text_seq_len: int, + uncond_text_seq_len: int, + **kwargs: Any, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """Run CFG, batching the two branches only when that is lossless. + + A batched forward needs one shared text length, so the shorter prompt + gets right-padded. Those pad positions carry all-zero UND K/V, and the + GEN cross-attention runs unmasked over the full text K/V — a zero key + scores logit 0 and still takes softmax mass while contributing nothing, + which weakens the padded branch's conditioning. Fall back to one + forward per branch, each trimmed to its own length, whenever the + prompts differ in length. + """ + if cond_text_seq_len == uncond_text_seq_len: + return self._predict_noise_cfg_batched( + latents=latents, + timestep=timestep, + cond_text_ids=cond_text_ids, + cond_text_mask=cond_text_mask, + uncond_text_ids=uncond_text_ids, + uncond_text_mask=uncond_text_mask, + video_shape=video_shape, + fps=fps, + guidance_scale=guidance_scale, + max_text_seq_len=cond_text_seq_len, + **kwargs, + ) + + if not self._logged_cfg_split and not self._current_batch_is_warmup: + self._logged_cfg_split = True + self.log_info( + "Prompt and negative prompt tokenize to different lengths " + f"({cond_text_seq_len} vs {uncond_text_seq_len}); running the " + "CFG branches in separate forwards to keep padding out of the " + "cross-attention" + ) + + cond = 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", + max_text_seq_len=cond_text_seq_len, + **kwargs, + ) + uncond = 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", + max_text_seq_len=uncond_text_seq_len, + **kwargs, + ) + + def _cfg_combine( + cond_pred: torch.Tensor, uncond_pred: torch.Tensor + ) -> torch.Tensor: + return uncond_pred + guidance_scale * (cond_pred - uncond_pred) + + if isinstance(cond, tuple): + return tuple(_cfg_combine(c, u) for c, u in zip(cond, uncond, strict=True)) + return _cfg_combine(cond, uncond) + def _predict_noise_cfg_batched( self, latents: torch.Tensor, diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index 9f642ae0c..c216e930a 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -440,13 +440,31 @@ def load_dict(file_path): ) from e +def _split_hf_subfolder(path: str) -> tuple[str, str | None]: + """Split 'namespace/repo/subfolder' into (repo_id, subfolder), or return (path, None).""" + if os.path.isabs(path): + return path, None + parts = path.split("/") + if len(parts) > 2: + return "/".join(parts[:2]), "/".join(parts[2:]) + return path, None + + def prepare_diffusers_component_path_for_loading(component_path: str) -> str: """Download component repos if needed and patch legacy flat ModelOpt configs.""" - local_component_path = ( - maybe_download_model(component_path) - if not os.path.exists(component_path) - else component_path - ) + if os.path.exists(component_path): + local_component_path = component_path + else: + repo_id, subfolder = _split_hf_subfolder(component_path) + if subfolder is not None: + # component_path is 'namespace/repo/subfolder' — download only that subfolder + local_repo = maybe_download_model( + repo_id, + allow_patterns=[f"{subfolder}/**", f"{subfolder}/*"], + ) + local_component_path = os.path.join(local_repo, subfolder) + else: + local_component_path = maybe_download_model(component_path) config_path = os.path.join(local_component_path, "config.json") if not os.path.exists(config_path): return local_component_path diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py index 29cf01ca5..f303a3e53 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py @@ -2,6 +2,7 @@ """Unit tests for Cosmos3 config, weight mapping, and sampling params.""" import importlib.util +import json import types import unittest from unittest import mock @@ -12,9 +13,14 @@ from sglang.multimodal_gen.configs.models.dits.cosmos3video import ( _build_cosmos3_param_names_mapping, ) from sglang.multimodal_gen.configs.pipeline_configs.cosmos3 import Cosmos3Config -from sglang.multimodal_gen.configs.sample.cosmos3 import Cosmos3SamplingParams +from sglang.multimodal_gen.configs.sample.cosmos3 import ( + COSMOS3_EDGE_SUPPORTED_RESOLUTIONS, + Cosmos3SamplingParams, +) from sglang.multimodal_gen.configs.sample.sampling_params import DataType from sglang.multimodal_gen.registry import ( + _PIPELINE_REGISTRY, + _discover_and_register_pipelines, _get_config_info, get_non_diffusers_pipeline_name, ) @@ -48,6 +54,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.c Cosmos3ImagePreprocessStage, Cosmos3LatentPreparationStage, Cosmos3TimestepPreparationStage, + Cosmos3TokenizationStage, + _inject_caption_metadata, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_action import ( EMBODIMENT_TO_DOMAIN_ID, @@ -224,6 +232,42 @@ class TestCosmos3ParamNamesMapping(unittest.TestCase): self.assertNotIn("language_model", key) +class TestCosmos3DenseParamNamesMapping(unittest.TestCase): + """Dense (squared-ReLU) checkpoints ship no gate_proj; up/down_proj must + pass through unmerged.""" + + @classmethod + def setUpClass(cls): + cls.fn = staticmethod( + get_param_names_mapping(_build_cosmos3_param_names_mapping(gated_mlp=False)) + ) + + def test_und_mlp_up_proj_unmerged(self): + key, idx, _ = _apply(self.fn, "layers.1.mlp.up_proj.weight") + self.assertEqual(key, "language_model.layers.1.mlp.up_proj.weight") + self.assertIsNone(idx) + + def test_und_mlp_down_proj_unmerged(self): + key, idx, _ = _apply(self.fn, "layers.1.mlp.down_proj.weight") + self.assertEqual(key, "language_model.layers.1.mlp.down_proj.weight") + self.assertIsNone(idx) + + def test_gen_mlp_up_proj_unmerged(self): + key, idx, _ = _apply(self.fn, "layers.2.mlp_moe_gen.up_proj.weight") + self.assertEqual(key, "gen_layers.2.mlp.up_proj.weight") + self.assertIsNone(idx) + + def test_gen_mlp_down_proj_unmerged(self): + key, idx, _ = _apply(self.fn, "layers.2.mlp_moe_gen.down_proj.weight") + self.assertEqual(key, "gen_layers.2.mlp.down_proj.weight") + self.assertIsNone(idx) + + def test_qkv_merge_still_applies(self): + key, idx, total = _apply(self.fn, "layers.0.self_attn.to_q.weight") + self.assertEqual(key, "language_model.layers.0.self_attn.to_qkv.weight") + self.assertEqual((idx, total), (0, 3)) + + class TestCosmos3AdjustNumFrames(unittest.TestCase): """Verify VAE-aligned frame rounding in Cosmos3Config.""" @@ -306,29 +350,100 @@ class TestCosmos3SchedulerConfig(unittest.TestCase): 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)), + stage._default_flow_shift_for_mode( + self._batch(data_type=DataType.IMAGE), is_edge=False + ), 3.0, ) self.assertEqual( stage._default_flow_shift_for_mode( - self._batch(preprocessed_image=torch.empty(1)) + self._batch(preprocessed_image=torch.empty(1)), is_edge=False ), 10.0, ) self.assertEqual( stage._default_flow_shift_for_mode( - self._batch(preprocessed_video=torch.empty(1)) + self._batch(preprocessed_video=torch.empty(1)), is_edge=False ), 10.0, ) - self.assertEqual(stage._default_flow_shift_for_mode(self._batch()), 10.0) + self.assertEqual( + stage._default_flow_shift_for_mode(self._batch(), is_edge=False), 10.0 + ) self.assertEqual( stage._default_flow_shift_for_mode( - self._batch(sp_kwargs={"action_mode": "policy"}) + self._batch(sp_kwargs={"action_mode": "policy"}), is_edge=False ), 10.0, ) + def test_edge_flow_shift_default(self): + stage = self._stage() + # Edge uses 3.0 for T2I and every video mode (T2V/I2V/V2V); action stays high. + self.assertEqual( + stage._default_flow_shift_for_mode(self._batch(), is_edge=True), 3.0 + ) + self.assertEqual( + stage._default_flow_shift_for_mode( + self._batch(data_type=DataType.IMAGE), is_edge=True + ), + 3.0, + ) + self.assertEqual( + stage._default_flow_shift_for_mode( + self._batch(preprocessed_image=torch.empty(1)), is_edge=True + ), + 3.0, + ) + self.assertEqual( + stage._default_flow_shift_for_mode( + self._batch(preprocessed_video=torch.empty(1)), is_edge=True + ), + 3.0, + ) + self.assertEqual( + stage._default_flow_shift_for_mode( + self._batch(sp_kwargs={"action_mode": "policy"}), is_edge=True + ), + 10.0, + ) + + +class TestCosmos3EdgeSamplingDefaults(unittest.TestCase): + """Edge variant fills its own resolution/guidance defaults; base is untouched.""" + + def test_edge_t2v_defaults(self): + sp = Cosmos3SamplingParams(prompt="t", num_frames=81) + sp._resolve_variant_defaults(is_edge=True) + self.assertEqual((sp.width, sp.height), (832, 480)) + self.assertEqual(sp.guidance_scale, 5.0) + + def test_edge_t2i_defaults(self): + sp = Cosmos3SamplingParams(prompt="t", num_frames=1) + sp._resolve_variant_defaults(is_edge=True) + self.assertEqual((sp.width, sp.height), (640, 640)) + self.assertEqual(sp.guidance_scale, 7.0) + + def test_edge_restricts_supported_resolutions(self): + sp = Cosmos3SamplingParams(prompt="t", num_frames=1) + sp._resolve_variant_defaults(is_edge=True) + self.assertEqual(sp.supported_resolutions, COSMOS3_EDGE_SUPPORTED_RESOLUTIONS) + # Base high-res sizes are excluded so they trip the "unsupported" warning. + self.assertNotIn((1280, 720), sp.supported_resolutions) + self.assertNotIn((1024, 1024), sp.supported_resolutions) + + def test_non_edge_defers_resolution_to_base(self): + sp = Cosmos3SamplingParams(prompt="t", num_frames=81) + sp._resolve_variant_defaults(is_edge=False) + self.assertIsNone(sp.width) + self.assertIsNone(sp.height) + self.assertEqual(sp.guidance_scale, 4.0) + + def test_explicit_resolution_preserved_for_edge(self): + sp = Cosmos3SamplingParams(prompt="t", num_frames=81, width=1024, height=576) + sp._resolve_variant_defaults(is_edge=True) + self.assertEqual((sp.width, sp.height), (1024, 576)) + class TestCosmos3SamplingParamsDataType(unittest.TestCase): """Verify num_frames==1 flips data_type to IMAGE before file name derivation.""" @@ -525,6 +640,7 @@ class TestCosmos3ModelResolution(unittest.TestCase): "nvidia/Cosmos3-Super", "nvidia/Cosmos3-Super-Text2Image", "nvidia/Cosmos3-Super-Image2Video", + "nvidia/Cosmos3-Edge", ): with self.subTest(model_path=model_path): self.assertIsNone(get_non_diffusers_pipeline_name(model_path)) @@ -533,6 +649,31 @@ class TestCosmos3ModelResolution(unittest.TestCase): self.assertIs(config_info.sampling_param_cls, Cosmos3SamplingParams) self.assertIs(config_info.pipeline_config_cls, Cosmos3Config) + def test_class_name_detection_matches_legacy_and_new(self): + """Unregistered checkpoints resolve via ``_class_name``: both the legacy + ``Cosmos3OmniDiffusersPipeline`` and the current ``Cosmos3OmniPipeline`` + map to the native Cosmos3 config.""" + for idx, class_name in enumerate( + ("Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline") + ): + model_path = f"acme/mystery-ckpt-{idx}" + with self.subTest(class_name=class_name): + with mock.patch( + "sglang.multimodal_gen.registry.maybe_download_model_index", + return_value={"_class_name": class_name}, + ): + config_info = _get_config_info(model_path) + self.assertIsNotNone(config_info) + self.assertIs(config_info.pipeline_config_cls, Cosmos3Config) + + def test_legacy_and_new_pipeline_names_both_registered(self): + """Both ``_class_name`` spellings resolve to a native pipeline class so + old (Nano/Super) and new (Edge) checkpoints load.""" + _discover_and_register_pipelines() + for pipeline_name in ("Cosmos3OmniPipeline", "Cosmos3OmniDiffusersPipeline"): + with self.subTest(pipeline_name=pipeline_name): + self.assertIn(pipeline_name, _PIPELINE_REGISTRY) + class TestCosmos3OpenAIProtocol(unittest.TestCase): """Verify Cosmos3 modality knobs are exposed by the video HTTP schema.""" @@ -930,5 +1071,130 @@ class TestCosmos3ModalitySamplingParams(unittest.TestCase): self.assertIsNone(getattr(sp, field)) +class TestCosmos3CaptionMetadata(unittest.TestCase): + """Structured captions get generation metadata; prose prompts opt out.""" + + def test_video_caption_gets_resolution_duration_and_fps(self): + prompt = json.dumps({"temporal_caption": "a cone melts"}) + caption = json.loads( + _inject_caption_metadata( + prompt, num_frames=189, fps=24.0, height=480, width=832 + ) + ) + self.assertEqual(caption["resolution"], {"H": 480, "W": 832}) + # Whole seconds, matching the caption format the model was trained on. + self.assertEqual(caption["duration"], "7s") + self.assertEqual(caption["fps"], 24.0) + self.assertEqual(caption["temporal_caption"], "a cone melts") + + def test_image_caption_drops_temporal_fields(self): + prompt = json.dumps({"subjects": ["hands"], "duration": "5s", "fps": 24.0}) + caption = json.loads( + _inject_caption_metadata( + prompt, num_frames=1, fps=24.0, height=768, width=768 + ) + ) + self.assertEqual(caption["resolution"], {"H": 768, "W": 768}) + self.assertNotIn("duration", caption) + self.assertNotIn("fps", caption) + self.assertEqual(caption["subjects"], ["hands"]) + + def test_request_resolution_overrides_caption(self): + prompt = json.dumps({"resolution": {"H": 111, "W": 222}}) + caption = json.loads( + _inject_caption_metadata( + prompt, num_frames=1, fps=24.0, height=640, width=640 + ) + ) + self.assertEqual(caption["resolution"], {"H": 640, "W": 640}) + + def test_unrelated_caption_fields_are_preserved(self): + prompt = json.dumps({"aspect_ratio": "1,1", "subjects": ["a vase"]}) + caption = json.loads( + _inject_caption_metadata( + prompt, num_frames=1, fps=24.0, height=640, width=640 + ) + ) + self.assertEqual(caption["aspect_ratio"], "1,1") + self.assertEqual(caption["subjects"], ["a vase"]) + + def test_zero_fps_does_not_divide_by_zero(self): + prompt = json.dumps({"temporal_caption": "a cone melts"}) + caption = json.loads( + _inject_caption_metadata( + prompt, num_frames=81, fps=0.0, height=480, width=832 + ) + ) + self.assertEqual(caption["duration"], "0s") + + def test_non_structured_prompts_are_declined(self): + for prompt in ( + "A curious raccoon in a field of sunflowers.", + json.dumps(["not", "an", "object"]), + json.dumps("a bare string"), + "", + ): + with self.subTest(prompt=prompt): + self.assertIsNone( + _inject_caption_metadata( + prompt, num_frames=81, fps=24.0, height=480, width=832 + ) + ) + + +class TestCosmos3DurationTemplateSuppression(unittest.TestCase): + """A structured caption must not also get the prose duration suffix.""" + + def _run_prompt_stage(self, prompt: str, use_duration_template: bool) -> str: + """Drive Cosmos3TokenizationStage.forward far enough to see the prompt.""" + seen = {} + + def fake_tokenize(prompt_text, *args, **kwargs): + seen.setdefault("prompt", prompt_text) + return ( + torch.zeros(1, 4, dtype=torch.long), + torch.ones(1, 4, dtype=torch.long), + 4, + ) + + stage = Cosmos3TokenizationStage.__new__(Cosmos3TokenizationStage) + batch = types.SimpleNamespace( + prompt=prompt, + negative_prompt="bad", + max_sequence_length=512, + use_duration_template=use_duration_template, + use_system_prompt=False, + fps=24.0, + num_frames=189, + height=480, + width=832, + data_type=DataType.VIDEO, + sampling_params=types.SimpleNamespace(action_mode=None), + extra={}, + ) + server_args = types.SimpleNamespace( + pipeline_config=types.SimpleNamespace( + use_duration_template=True, use_system_prompt=False + ) + ) + with ( + mock.patch.object(stage, "_tokenize_prompt", fake_tokenize), + mock.patch.object(stage, "log_info"), + ): + stage.forward(batch, server_args) + return seen["prompt"] + + def test_structured_caption_stays_valid_json(self): + prompt = json.dumps({"temporal_caption": "a cone melts"}) + final = self._run_prompt_stage(prompt, use_duration_template=True) + caption = json.loads(final) # would raise if the suffix were appended + self.assertEqual(caption["duration"], "7s") + self.assertNotIn("seconds long", final) + + def test_prose_prompt_still_gets_the_duration_suffix(self): + final = self._run_prompt_stage("A curious raccoon.", use_duration_template=True) + self.assertIn("seconds long", final) + + if __name__ == "__main__": unittest.main()