diff --git a/docker/Dockerfile b/docker/Dockerfile index 0451e21b3..aa8c71896 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -261,6 +261,24 @@ RUN --mount=type=cache,target=/root/.cache/pip \ # stage does not COPY; force a pip copy into /usr/local so it survives the stage split. RUN python3 -m pip install --ignore-installed --no-deps distro +# Cosmos3 guardrails default ON: sglang enables them whenever `cosmos_guardrail` +# is importable (Cosmos3Pipeline.create_pipeline_stages); opt out at runtime with +# SGLANG_DISABLE_COSMOS3_GUARDRAILS=1. Installed with --no-deps because the +# package metadata conflicts with our pins in ways that don't matter at runtime +# (imageio>=2.37 vs our ==2.36) or would be harmful (opencv-python — also pulled +# by retinaface-py — would shadow our opencv-python-headless cv2). The deps not +# already in the image are installed explicitly; the find_spec check mirrors +# sglang's availability probe plus cosmos_guardrail's own imports. +RUN --mount=type=cache,target=/root/.cache/pip \ + python3 -m pip install --no-deps \ + cosmos-guardrail==0.3.1 \ + retinaface-py==0.0.2 \ + && python3 -m pip install \ + "better-profanity==0.7.0" \ + "nltk==3.9.1" \ + "peft==0.18.1" \ + && python3 -c "import importlib.util as u; missing = [m for m in ('cosmos_guardrail', 'retinaface', 'better_profanity', 'nltk', 'cv2', 'peft', 'sentencepiece') if u.find_spec(m) is None]; assert not missing, f'guardrail deps missing: {missing}'" + ######################################################## # PARALLEL STAGE 2: HPC-Ops Builder (needs torch_deps) ######################################################## diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py index 6ad73fa18..838e71d89 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py @@ -113,7 +113,7 @@ class Cosmos3Config(PipelineConfig): vae_precision: str = "bf16" # Pipeline-level (not sampling) knobs. - max_sequence_length: int = 512 + max_sequence_length: int = 4096 use_duration_template: bool = True use_system_prompt: bool = False diff --git a/python/sglang/multimodal_gen/configs/sample/cosmos3.py b/python/sglang/multimodal_gen/configs/sample/cosmos3.py index d0c1be460..7ae7e14a9 100644 --- a/python/sglang/multimodal_gen/configs/sample/cosmos3.py +++ b/python/sglang/multimodal_gen/configs/sample/cosmos3.py @@ -10,7 +10,7 @@ so the file extension and decode path agree. """ from dataclasses import dataclass, field -from typing import Any +from typing import Any, ClassVar from sglang.multimodal_gen.configs.sample.sampling_params import ( DataType, @@ -74,6 +74,61 @@ class Cosmos3SamplingParams(SamplingParams): condition_frame_indexes: list[int] | None = None condition_video_keep: str = "first" + # Transfer (control-video) conditioning. ``control_path`` points to one or + # more pre-computed control videos (e.g. edge / blur / depth / seg / wsm + # maps). When set, each control clip is VAE-encoded and packed as clean + # vision tokens that prefix the target clip in the GEN sequence; multiple + # paths drive multi-hint transfer (e.g. edge + depth). Control clips reuse + # ``proj_in``, so every Cosmos3 checkpoint supports transfer. + control_path: str | list[str] | None = None + + # Optional hint type(s) parallel to ``control_path`` (one of + # ``edge`` / ``blur`` / ``depth`` / ``seg`` / ``wsm``). Used only to apply + # tuned per-hint defaults (``guidance`` / ``control_guidance`` / ``shift``) + # when exactly one control input is given and the user left those unset. + control_hint: str | list[str] | None = None + + # Control-CFG scale for transfer. ``1.0`` (default) disables the extra + # control-dropped forward; values > 1.0 amplify the control map's influence + # by blending the with-control and without-control predictions on the + # generated span: ``cond_nc + control_guidance * (cond_full - cond_nc)``. + control_guidance: float = 1.0 + + # Optional timestep window ``(lo, hi)`` restricting where control-CFG is + # applied (analogous to ``guidance_interval`` for text CFG). ``None`` applies + # it at every step. + control_guidance_interval: tuple[float, float] | None = None + + # Long-video transfer controls. Chunks overlap by + # ``num_conditional_frames`` pixel frames; overlap frames from the previous + # decoded chunk are kept clean in the next chunk. + num_video_frames_per_chunk: int = 93 + num_conditional_frames: int = 1 + num_first_chunk_conditional_frames: int = 0 + max_frames: int = 5000 + show_control_condition: bool = False + show_input: bool = False + share_vision_temporal_positions: bool = True + + # Tuned per-hint defaults applied when exactly one control input is given + # and the corresponding field was not set explicitly (mirrors the + # cosmos-framework ``_TRANSFER_DEFAULTS`` table). ``shift`` maps to + # ``flow_shift``. Multi-hint transfer keeps the request's own values. + _TRANSFER_DEFAULTS: ClassVar[dict[str, dict[str, float | int]]] = { + "edge": {"guidance": 3.0, "control_guidance": 1.5, "shift": 10.0}, + "blur": {"guidance": 3.0, "control_guidance": 1.5, "shift": 10.0}, + "depth": {"guidance": 3.0, "control_guidance": 1.5, "shift": 10.0}, + "seg": {"guidance": 3.0, "control_guidance": 2.0, "shift": 10.0}, + "wsm": { + "guidance": 1.0, + "control_guidance": 3.0, + "shift": 10.0, + "num_frames": 101, + "fps": 10, + "num_video_frames_per_chunk": 101, + }, + } + supported_resolutions: list[tuple[int, int]] | None = field( default_factory=lambda: [ (1280, 720), @@ -103,6 +158,108 @@ class Cosmos3SamplingParams(SamplingParams): action_stats_path: str | None = None action_normalization: str = "quantile" + @classmethod + def video_request_extra_fields(cls) -> frozenset[str]: + return frozenset( + { + "generate_sound", + "sound_duration", + "condition_frame_indexes", + "condition_frame_indexes_vision", + "condition_video_keep", + "control_path", + "control_hint", + "control_guidance", + "control_guidance_interval", + "num_video_frames_per_chunk", + "num_conditional_frames", + "num_first_chunk_conditional_frames", + "max_frames", + "show_control_condition", + "show_input", + "share_vision_temporal_positions", + "action_mode", + "domain_id", + "domain_name", + "raw_action_dim", + "action_fps", + "action", + "action_view_point", + "action_normalization", + } + ) + + def _resolve_control_paths(self) -> list[str]: + cp = self.control_path + if cp is None: + return [] + if isinstance(cp, str): + return [cp] if cp else [] + return [p for p in cp if isinstance(p, str) and p] + + def _resolve_control_hints(self) -> list[str]: + hint = self.control_hint + if hint is None: + return [] + hints = [hint] if isinstance(hint, str) else list(hint) + hints = [h for h in hints if h] + for h in hints: + if h not in self._TRANSFER_DEFAULTS: + raise ValueError( + f"Unknown control_hint {h!r}; expected one of " + f"{sorted(self._TRANSFER_DEFAULTS)}" + ) + return hints + + def _apply_transfer_hint_defaults(self) -> None: + """Fill tuned per-hint defaults for a single, typed control input. + + Mirrors cosmos-framework: defaults apply only when there is exactly one + control input with a known hint type, and only to fields the user did + not pass explicitly (tracked via ``_explicit_fields``). Multi-hint + transfer keeps the request's own ``guidance`` / ``control_guidance`` / + ``flow_shift``. + """ + if len(self._resolve_control_paths()) != 1: + return + hints = self._resolve_control_hints() + if len(hints) != 1: + return + defaults = self._TRANSFER_DEFAULTS.get(hints[0]) + if defaults is None: + return + explicit = getattr(self, "_explicit_fields", None) or set() + if "control_guidance" not in explicit: + self.control_guidance = defaults["control_guidance"] + if "guidance_scale" not in explicit: + self.guidance_scale = defaults["guidance"] + if "flow_shift" not in explicit and self.flow_shift is None: + self.flow_shift = defaults["shift"] + for field_name in ( + "num_frames", + "fps", + "num_video_frames_per_chunk", + ): + if field_name in defaults and field_name not in explicit: + setattr(self, field_name, defaults[field_name]) + + @classmethod + def lower_video_request_kwargs( + cls, request: Any, kwargs: dict[str, Any] + ) -> dict[str, Any]: + """Apply defaults that the generic video endpoint pre-resolves.""" + hint = kwargs.get("control_hint") + paths = kwargs.get("control_path") + hints = [hint] if isinstance(hint, str) else list(hint or []) + control_paths = [paths] if isinstance(paths, str) else list(paths or []) + if len(control_paths) == 1 and hints == ["wsm"]: + defaults = cls._TRANSFER_DEFAULTS["wsm"] + if getattr(request, "num_frames", None) is None: + kwargs["num_frames"] = defaults["num_frames"] + if getattr(request, "fps", None) is None: + kwargs["fps"] = defaults["fps"] + return kwargs + 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. @@ -130,6 +287,37 @@ class Cosmos3SamplingParams(SamplingParams): ) action_output = self.action_mode != "forward_dynamics" + # Apply transfer per-hint defaults before the base resolves remaining + # fields (e.g. flow_shift per mode), so an unset flow_shift can pick up + # the hint's tuned shift. + self._apply_transfer_hint_defaults() + control_paths = self._resolve_control_paths() + if control_paths: + if pipeline_config.distilled_sigmas is not None: + raise ValueError( + "Cosmos3 distilled checkpoints do not support transfer inference" + ) + if pipeline_config.is_edge: + raise ValueError( + "Cosmos3 Edge checkpoints do not support transfer inference" + ) + if self.num_frames == 1: + raise ValueError( + "Cosmos3 transfer inference is supported only for video outputs" + ) + if self.image_path is not None: + raise ValueError( + "Cosmos3 transfer accepts control videos and an optional source " + "video, not an image input" + ) + if self.action_mode is not None: + raise ValueError( + "Cosmos3 transfer cannot be combined with action generation" + ) + if float(self.sound_duration or 0.0) > 0.0: + raise ValueError( + "Cosmos3 transfer cannot be combined with sound generation" + ) super()._adjust(server_args) # Policy and inverse dynamics produce actions. Forward dynamics consumes @@ -142,6 +330,39 @@ class Cosmos3SamplingParams(SamplingParams): self.output_file_name = None self.output_compression = 0 + def _validate(self) -> None: + super()._validate() + paths = self._resolve_control_paths() + hints = self._resolve_control_hints() + if hints and len(hints) != len(paths): + raise ValueError( + "control_hint must contain exactly one entry per control_path " + f"(got {len(hints)} hint(s) for {len(paths)} path(s))" + ) + if self.control_guidance_interval is not None: + if len(self.control_guidance_interval) != 2: + raise ValueError( + "control_guidance_interval must contain exactly two values" + ) + lo, hi = self.control_guidance_interval + if float(lo) > float(hi): + raise ValueError( + "control_guidance_interval must be ordered as (low, high)" + ) + if self.num_video_frames_per_chunk <= 0: + raise ValueError("num_video_frames_per_chunk must be positive") + if self.num_conditional_frames < 0: + raise ValueError("num_conditional_frames must be non-negative") + if self.num_conditional_frames >= self.num_video_frames_per_chunk: + raise ValueError( + "num_conditional_frames must be smaller than " + "num_video_frames_per_chunk" + ) + if self.num_first_chunk_conditional_frames < 0: + raise ValueError("num_first_chunk_conditional_frames must be non-negative") + if self.max_frames <= 0: + raise ValueError("max_frames must be positive") + def _guidance_is_explicit(self) -> bool: explicit = getattr(self, "_explicit_fields", None) return explicit is not None and "guidance_scale" in explicit 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 319632a5a..b0f2b6718 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -117,19 +117,6 @@ _MULTIPART_EXTRA_FORM_FIELDS = ( "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", "quality", ) @@ -265,6 +252,49 @@ def _coerce_optional_int_list(value: Any) -> list[int] | None: return [int(value)] +def _coerce_optional_float_list(value: Any) -> list[float] | 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 [float(item) for item in value] + return [float(value)] + + +def _coerce_optional_bool(value: Any) -> bool | None: + value = _parse_form_extra_value(value) + if value is None or (isinstance(value, str) and not value.strip()): + return None + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + raise ValueError(f"Invalid boolean value: {value!r}") + return bool(value) + + +def _coerce_optional_str_list(value: Any) -> str | list[str] | None: + """Coerce a control_path/control_hint value to str or list[str]. + + Accepts a JSON list (``["edge.mp4", "depth.mp4"]``), a plain string, or a + native list. Empty values resolve to ``None`` so unset fields don't override + sampling-param defaults. + """ + value = _parse_form_extra_value(value) + if value is None: + return None + if isinstance(value, (list, tuple)): + items = [str(item) for item in value if str(item).strip()] + return items or None + if isinstance(value, str): + return value if value.strip() else None + return str(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: @@ -324,6 +354,41 @@ def _cosmos3_sampling_param_kwargs( if condition_frame_indexes is not None: kwargs["condition_frame_indexes"] = condition_frame_indexes + # Transfer (control-video) conditioning. + control_path = _coerce_optional_str_list(_request_value(req, "control_path")) + if control_path is not None: + kwargs["control_path"] = control_path + control_hint = _coerce_optional_str_list(_request_value(req, "control_hint")) + if control_hint is not None: + kwargs["control_hint"] = control_hint + control_guidance = _request_value(req, "control_guidance") + if control_guidance is not None: + kwargs["control_guidance"] = float(control_guidance) + control_guidance_interval = _coerce_optional_float_list( + _request_value(req, "control_guidance_interval") + ) + if control_guidance_interval is not None: + kwargs["control_guidance_interval"] = tuple(control_guidance_interval) + + for name in ( + "num_video_frames_per_chunk", + "num_conditional_frames", + "num_first_chunk_conditional_frames", + "max_frames", + ): + value = _parse_form_extra_value(_request_value(req, name)) + if value is not None and value != "": + kwargs[name] = int(value) + + for name in ( + "show_control_condition", + "show_input", + "share_vision_temporal_positions", + ): + value = _coerce_optional_bool(_request_value(req, name)) + if value is not None: + kwargs[name] = value + for name in ( "condition_video_keep", "action_mode", diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index e1161acea..013dfae83 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -206,6 +206,13 @@ def _maybe_dequantize_fp8( return full_tensor +def _move_to_device_preserving_meta(model: nn.Module, device: torch.device) -> None: + # Buffers absent from the checkpoint (e.g. cosmos3's RoPE inv_freq) are + # still on the meta device here and .to() cannot copy out of meta; leave + # them for the model's post_load_weights() to rebuild on the real device. + model._apply(lambda t: t if t.is_meta else t.to(device)) + + def register_fsdp_entrypoints(model: torch.nn.Module) -> None: """Let FSDP2 unshard around forward passes that bypass ``__call__``. @@ -435,7 +442,7 @@ def maybe_load_fsdp_model( # 3. postprocessing if weight_postprocess_device is not None: # move to device to perform postprocessing - model.to(weight_postprocess_device) + _move_to_device_preserving_meta(model, weight_postprocess_device) for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index 66dd66c15..6ef0bd291 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py @@ -7,7 +7,7 @@ cross-attends from noisy visual tokens to that cache at every denoising step. """ import math -from collections.abc import Iterable, Iterator +from collections.abc import Iterable, Iterator, Sequence from typing import Any import torch @@ -1333,12 +1333,30 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): action_frames: int = 0, action_fps: float | None = None, action_start_frame_offset: int = 1, + control_frames: int | Sequence[int] = 0, + share_vision_temporal_positions: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: - """Compute mRoPE position IDs for UND text and GEN visual + action + sound tokens.""" + """Compute mRoPE position IDs for UND text and GEN tokens. + + The GEN sequence is ordered ``[control, video, action, sound]``. + Control and target videos either share matching temporal coordinates or + occupy consecutive temporal ranges, according to + ``share_vision_temporal_positions``. + """ B = text_mask.shape[0] S_text = text_mask.shape[1] text_lengths = text_mask.sum(dim=1).long() effective_fps = fps if fps is not None and T > 1 else None + control_frame_counts = ( + [control_frames] + if isinstance(control_frames, int) and control_frames > 0 + else ( + [int(count) for count in control_frames if int(count) > 0] + if not isinstance(control_frames, int) + else [] + ) + ) + vision_frame_counts = [*control_frame_counts, T] text_pos_list = [] vis_pos_list = [] @@ -1348,16 +1366,28 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): 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=media_offset, - device=device, - fps=effective_fps, - base_fps=self.base_fps, - temporal_compression_factor=self.temporal_compression_factor, - ) + vision_offset = media_offset + vision_pos_blocks = [] + for frame_count in vision_frame_counts: + temporal_offset = ( + media_offset if share_vision_temporal_positions else vision_offset + ) + vision_pos, vision_offset = compute_mrope_position_ids_vision( + frame_count, + Hp, + Wp, + temporal_offset=temporal_offset, + device=device, + fps=effective_fps, + base_fps=self.base_fps, + temporal_compression_factor=self.temporal_compression_factor, + ) + vision_pos_blocks.append(vision_pos) + + pos_dtype = vision_pos_blocks[0].dtype + for pos in vision_pos_blocks[1:]: + pos_dtype = torch.promote_types(pos_dtype, pos.dtype) + v_pos = torch.cat([pos.to(pos_dtype) for pos in vision_pos_blocks], dim=1) if action_frames > 0: a_pos, _ = compute_mrope_position_ids_action( action_frames, @@ -1394,9 +1424,8 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): text_pos_list.append(t_pos) 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_gen] - + text_pos_ids = torch.stack(text_pos_list, dim=1).to(device) + vis_pos_ids = torch.stack(vis_pos_list, dim=1).to(device) return text_pos_ids, vis_pos_ids def reset_cache(self, cache_key: str | None = None): @@ -1443,6 +1472,8 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): action_noisy_mask: torch.Tensor | None = None, action_fps: float | None = None, action_start_frame_offset: int = 1, + control_latents: torch.Tensor | list[torch.Tensor] | None = None, + transfer_share_vision_temporal_positions: bool = True, **kwargs, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Forward pass for denoising. @@ -1473,6 +1504,11 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): 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). + control_latents: Optional [B, C, T_ctrl, H, W] control-video latents + (transfer / control-net conditioning). They are patchified and + projected with the shared ``proj_in``, prepended to the GEN + sequence as clean (noise-free) tokens that share the video's + temporal positions, and excluded from the output projection. Returns: [B, C, T, H, W] velocity prediction, or a tuple @@ -1505,8 +1541,44 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): device=action_latents.device, ) + # Transfer / control-video conditioning: one or more control clips + # (e.g. edge + depth) are packed as clean vision tokens that prefix the + # target clip in the GEN sequence. Each clip reuses ``proj_in`` and the + # shared transformer; blocks are concatenated in input order. + control_clips: list[torch.Tensor] = [] + if control_latents is not None: + control_clips = ( + list(control_latents) + if isinstance(control_latents, (list, tuple)) + else [control_latents] + ) + control_frame_counts: list[int] = [] + hidden_control_blocks: list[torch.Tensor] = [] + control_token_len = 0 + for clip in control_clips: + _, _, c_frames, Hc, Wc = clip.shape + if (Hc, Wc) != (H, W): + raise ValueError( + "control_latents spatial dims " + f"{(Hc, Wc)} must match hidden_states {(H, W)}" + ) + block, _ = self.proj_in( + self.patchify(clip.to(hidden_states.dtype), c_frames, Hc, Wc) + ) + hidden_control_blocks.append(block) + control_frame_counts.append(c_frames) + control_token_len += block.shape[1] + has_control = len(hidden_control_blocks) > 0 + hidden_control = ( + torch.cat(hidden_control_blocks, dim=1) if has_control else None + ) + extra_frames = action_frames + sound_frames sequence_shard_enabled = self.sp_size > 1 + # When a control clip is present we always assemble the combined GEN + # stream (control prefix + video [+ action] [+ sound]) instead of the + # video-only fast path. + use_assembly_path = extra_frames > 0 or has_control # Add timestep embedding (computed in float32 for numerical stability, then cast back) time_embed = self.time_embedder(timestep.float()) @@ -1531,7 +1603,7 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): .to(hidden_gen.dtype) ) - if extra_frames == 0: + if not use_assembly_path: # Video-only: shard the visual tokens, then add the timestep # embedding on the local shard. if sequence_shard_enabled: @@ -1567,15 +1639,21 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): hidden_gen = hidden_gen + time_embed.unsqueeze(1) else: # 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. + # ([control,] 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) + # Control tokens are clean conditioning: prepend them WITHOUT a + # timestep embedding so the GEN tokens can attend to the raw control + # map. They are stripped before the output projection. + if has_control: + hidden_gen = torch.cat([hidden_control, hidden_gen], dim=1) + if action_latents is not None: hidden_action = self.action_proj_in( action_latents.to(hidden_gen.dtype), action_domain_ids @@ -1648,6 +1726,10 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): action_frames=action_frames, action_fps=action_fps if action_fps is not None else fps, action_start_frame_offset=action_start_frame_offset, + control_frames=control_frame_counts, + share_vision_temporal_positions=( + transfer_share_vision_temporal_positions + ), ) # UND K/V cache is kept FULL on all ranks (not sharded). Text # sequence is short, so memory impact is minimal, and the GEN @@ -1722,7 +1804,7 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): hidden_gen = hidden_gen + residual hidden_gen = self.norm_moe_gen(hidden_gen) - if extra_frames == 0: + if not use_assembly_path: # 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 @@ -1741,12 +1823,15 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): if seq_shard_pad > 0: hidden_gen = hidden_gen[:, :seq_len_orig, :] - s_video = seq_len_orig - extra_frames - output, _ = self.proj_out(hidden_gen[:, :s_video, :]) + # Sequence layout: [control prefix | video | action | sound]. Control + # tokens are conditioning only and produce no output. + s_video = seq_len_orig - extra_frames - control_token_len + video_start = control_token_len + output, _ = self.proj_out(hidden_gen[:, video_start : video_start + s_video, :]) video_pred = self.unpatchify(output, T, H, W) extra_outputs: list[torch.Tensor] = [] - idx = s_video + idx = video_start + 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)) @@ -1756,6 +1841,10 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin): sound_output, _ = self.audio_proj_out(sound_hidden) extra_outputs.append(sound_output.permute(0, 2, 1).contiguous()) + # Control-only conditioning (no action/sound): keep the bare-tensor + # return type identical to the video-only path. + if not extra_outputs: + return video_pred return (video_pred, *extra_outputs) def preprocess_loaded_state_dict( diff --git a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py index 3509f5383..040dd86b9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py @@ -100,7 +100,11 @@ class Cosmos3Pipeline(LoRAPipeline, ComposedPipelineBase): self.add_stage(Cosmos3TextGuardrailStage()) self.add_stage(Cosmos3LatentPreparationStage(vae, transformer)) self.add_stage(Cosmos3TimestepPreparationStage(scheduler)) - self.add_stage(Cosmos3DenoisingStage(transformer, scheduler, server_args)) + self.add_stage( + Cosmos3DenoisingStage( + transformer, scheduler, server_args=server_args, vae=vae + ) + ) self.add_stage( Cosmos3DecodingStage( vae, guardrails=guardrails_on, sound_tokenizer=sound_tokenizer 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 ee496fe86..220916090 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 @@ -11,12 +11,14 @@ per-request from ``batch.data_type`` and the presence of import copy import json +import math from typing import Any import numpy as np import PIL.Image import torch import torch.nn as nn +import torch.nn.functional as F from sglang.multimodal_gen.configs.sample.sampling_params import DataType from sglang.multimodal_gen.runtime.distributed import get_local_torch_device @@ -139,6 +141,59 @@ def _pil_to_normalized_tensor(image: PIL.Image.Image) -> torch.Tensor: return torch.from_numpy(arr).permute(2, 0, 1).contiguous() +def _pil_to_uint8_tensor(image: PIL.Image.Image) -> torch.Tensor: + arr = np.asarray(image, dtype=np.uint8).copy() + return torch.from_numpy(arr).permute(2, 0, 1).contiguous() + + +def _resize_center_crop_uint8_cthw( + frames: torch.Tensor, height: int, width: int +) -> torch.Tensor: + """Resize and center-crop ``uint8 [3, T, H, W]`` transfer frames.""" + if frames.ndim != 4 or frames.shape[0] != 3: + raise ValueError( + "Transfer frames must have shape [3, T, H, W], got " + f"{tuple(frames.shape)}" + ) + orig_h, orig_w = int(frames.shape[2]), int(frames.shape[3]) + scale = max(width / orig_w, height / orig_h) + resize_h = int(np.ceil(scale * orig_h)) + resize_w = int(np.ceil(scale * orig_w)) + frames_tchw = frames.permute(1, 0, 2, 3).to(dtype=torch.float32) + resized = F.interpolate( + frames_tchw, + size=(resize_h, resize_w), + mode="bilinear", + align_corners=False, + ) + top = (resize_h - height) // 2 + left = (resize_w - width) // 2 + cropped = resized[:, :, top : top + height, left : left + width] + return ( + cropped.round().clamp(0, 255).to(torch.uint8).permute(1, 0, 2, 3).contiguous() + ) + + +def _pad_transfer_frames(video: torch.Tensor, target_frames: int) -> torch.Tensor: + """Pad ``[1, 3, T, H, W]`` with reflected temporal content.""" + if video.ndim != 5 or video.shape[0] != 1 or video.shape[1] != 3: + raise ValueError( + "Transfer video must have shape [1, 3, T, H, W], got " + f"{tuple(video.shape)}" + ) + if target_frames <= 0: + raise ValueError("Transfer target frame count must be positive") + video = video[:, :, :target_frames] + if video.shape[2] == 0: + raise ValueError("Transfer video cannot be empty") + while video.shape[2] < target_frames: + remaining = target_frames - video.shape[2] + reflected = video.flip(dims=[2]) + pad_len = min(max(video.shape[2] - 1, 1), remaining) + video = torch.cat([video, reflected[:, :, :pad_len]], dim=2) + return video.contiguous() + + class Cosmos3ImagePreprocessStage(PipelineStage): """Load, aspect-resize, and center-crop the conditioning input. @@ -146,6 +201,8 @@ class Cosmos3ImagePreprocessStage(PipelineStage): policy requests write ``[B, 3, H, W]``; regular visual generation remains single-image conditioned. For V2V: writes ``[1, 3, T_in, H, W]`` to ``batch.preprocessed_video``. + For transfer: writes ``[1, 3, T, H, W]`` control pixels to + ``batch.extra["preprocessed_control"]`` (independent of I2V / V2V). No-op for T2V / T2I. """ @@ -154,6 +211,136 @@ class Cosmos3ImagePreprocessStage(PipelineStage): def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: return VerificationResult() + def _load_control_video( + self, + control_path: str, + target_w: int, + target_h: int, + max_frames: int, + ) -> torch.Tensor: + """Load transfer media as CPU ``uint8 [1, 3, T, H, W]``.""" + frames = load_video(control_path) + if not frames: + raise ValueError(f"No frames decoded from transfer video: {control_path!r}") + frames = frames[:max_frames] + frames_cthw = torch.stack( + [_pil_to_uint8_tensor(frame.convert("RGB")) for frame in frames], + dim=1, + ) + return ( + _resize_center_crop_uint8_cthw(frames_cthw, height=target_h, width=target_w) + .unsqueeze(0) + .contiguous() + ) + + @staticmethod + def _get_transfer_num_chunks( + total_frames: int, frames_per_chunk: int, conditional_frames: int + ) -> tuple[int, int]: + if total_frames <= frames_per_chunk: + return 1, frames_per_chunk + stride = frames_per_chunk - conditional_frames + if stride <= 0: + raise ValueError( + "num_conditional_frames must be smaller than " + "num_video_frames_per_chunk" + ) + remaining = total_frames - frames_per_chunk + return 1 + math.ceil(remaining / stride), stride + + def _prepare_transfer_plan( + self, + batch: Req, + control_paths: list[str], + source_video_path: str | None, + ) -> None: + max_frames = int(batch.sampling_params.max_frames) + controls = [ + self._load_control_video( + path, batch.width, batch.height, max_frames=max_frames + ) + for path in control_paths + ] + total_frames = min(batch.num_frames, controls[0].shape[2], max_frames) + if total_frames <= 0: + raise ValueError("Cosmos3 transfer requires at least one control frame") + + requested_chunk_frames = int(batch.sampling_params.num_video_frames_per_chunk) + chunk_frames = ( + 1 + if total_frames == 1 + else (math.ceil((requested_chunk_frames - 1) / 4) * 4 + 1) + ) + num_chunks, stride = self._get_transfer_num_chunks( + total_frames, + chunk_frames, + int(batch.sampling_params.num_conditional_frames), + ) + padded_frames = max(total_frames, chunk_frames) + controls = [ + _pad_transfer_frames(control, padded_frames) for control in controls + ] + + source_video = None + if source_video_path: + source_video = self._load_control_video( + source_video_path, + batch.width, + batch.height, + max_frames=max_frames, + ) + source_video = _pad_transfer_frames(source_video, padded_frames) + if ( + batch.sampling_params.num_first_chunk_conditional_frames > 0 + and source_video is None + ): + raise ValueError( + "num_first_chunk_conditional_frames > 0 requires video_path" + ) + + batch.num_frames = total_frames + batch.extra["preprocessed_control"] = controls + batch.extra["preprocessed_transfer_video"] = source_video + batch.extra["transfer_plan"] = { + "total_frames": total_frames, + "chunk_frames": chunk_frames, + "num_chunks": num_chunks, + "stride": stride, + } + self.log_info( + f"Prepared transfer plan with {len(controls)} control(s), " + f"{total_frames} output frames, {num_chunks} chunk(s) of " + f"{chunk_frames} frames" + ) + + @staticmethod + def _normalize_control_paths(control_path: Any) -> list[str]: + """Normalize ``control_path`` (str / list / None) to a list of paths. + + Multiple paths drive multi-hint transfer (e.g. ``[edge.mp4, depth.mp4]``): + each is VAE-encoded into its own control-latent block and all blocks + prefix the target clip in the GEN sequence. + """ + if control_path is None: + return [] + if isinstance(control_path, str): + if not control_path.strip(): + raise ValueError("control_path is an empty string") + return [control_path] + if isinstance(control_path, (list, tuple)): + paths: list[str] = [] + for i, p in enumerate(control_path): + if not isinstance(p, str) or not p.strip(): + raise ValueError( + f"control_path[{i}] must be a non-empty string, got {p!r}" + ) + paths.append(p) + return paths + raise ValueError( + "control_path must be a string or list of strings, got " + f"{type(control_path).__name__}" + ) + def forward(self, batch: Req, server_args: ServerArgs) -> Req: image_path = batch.image_path video_path = batch.video_path @@ -167,6 +354,17 @@ class Cosmos3ImagePreprocessStage(PipelineStage): if isinstance(video_path, list): video_path = video_path[0] if video_path else None + control_paths = self._normalize_control_paths( + getattr(batch.sampling_params, "control_path", None) + ) + if control_paths: + if image_path is not None: + raise ValueError( + "Cosmos3 transfer accepts an optional source video, not an image" + ) + self._prepare_transfer_plan(batch, control_paths, video_path) + return batch + if image_path and video_path: raise ValueError( "Cosmos3 accepts either --image-path (I2V) or --video-path " @@ -368,7 +566,7 @@ class Cosmos3TokenizationStage(PipelineStage): negative_prompt = batch.negative_prompt or COSMOS3_DEFAULT_NEGATIVE_PROMPT # Get parameters - max_sequence_length = getattr(batch, "max_sequence_length", None) or 512 + max_sequence_length = getattr(batch, "max_sequence_length", None) or 4096 use_duration_template = getattr(batch, "use_duration_template", None) if use_duration_template is None: use_duration_template = getattr( @@ -380,7 +578,9 @@ class Cosmos3TokenizationStage(PipelineStage): server_args.pipeline_config, "use_system_prompt", False ) fps = batch.fps or 24.0 - num_frames = batch.num_frames + num_frames = batch.extra.get("transfer_plan", {}).get( + "chunk_frames", batch.num_frames + ) is_image_gen = batch.data_type == DataType.IMAGE system_prompt = ( COSMOS3_IMAGE_SYSTEM_PROMPT if is_image_gen else COSMOS3_VIDEO_SYSTEM_PROMPT @@ -506,8 +706,12 @@ class Cosmos3LatentPreparationStage(PipelineStage): vae_scale_factor_temporal = getattr(self.vae.config, "scale_factor_temporal", 4) vae_scale_factor_spatial = getattr(self.vae.config, "scale_factor_spatial", 16) + transfer_plan = batch.extra.get("transfer_plan") + pixel_num_frames = ( + transfer_plan["chunk_frames"] if transfer_plan else batch.num_frames + ) num_channels_latents = self.transformer.latent_channel - num_latent_frames = (batch.num_frames - 1) // vae_scale_factor_temporal + 1 + num_latent_frames = (pixel_num_frames - 1) // vae_scale_factor_temporal + 1 height_latent = batch.height // vae_scale_factor_spatial width_latent = batch.width // vae_scale_factor_spatial @@ -524,6 +728,19 @@ class Cosmos3LatentPreparationStage(PipelineStage): width_latent, ) + if transfer_plan is not None: + batch.latents = torch.zeros(shape, device=device, dtype=dtype) + batch.raw_latent_shape = shape + batch.extra["video_shape"] = ( + num_latent_frames, + height_latent, + width_latent, + ) + batch.extra["vae_scale_factor_temporal"] = vae_scale_factor_temporal + batch.extra["vae_scale_factor_spatial"] = vae_scale_factor_spatial + self.log_info(f"Prepared transfer latent shape {shape}") + return batch + generator = batch.generator if generator is None and batch.seed is not None: generator = torch.Generator(device=device).manual_seed(batch.seed) @@ -594,6 +811,36 @@ class Cosmos3LatentPreparationStage(PipelineStage): self.log_info(f"Prepared latents with shape {shape}") + # Transfer (control-video) conditioning: VAE-encode each control clip + # into clean latents the transformer prepends to the GEN sequence. Stored + # as a list (one block per hint) so multi-hint transfer (edge + depth …) + # threads through the denoiser uniformly with the single-hint case. + preprocessed_control = batch.extra.get("preprocessed_control") + if preprocessed_control is not None: + control_blocks = ( + preprocessed_control + if isinstance(preprocessed_control, list) + else [preprocessed_control] + ) + vae_dtype = next(self.vae.parameters()).dtype + control_latents_list: list[torch.Tensor] = [] + for control_pixels_t in control_blocks: + control_pixels = control_pixels_t.to(device=device, dtype=vae_dtype) + with torch.no_grad(): + control_latent = self._vae_encode(control_pixels).to(dtype) + if control_latent.shape[-2:] != latents.shape[-2:]: + raise ValueError( + "control latent spatial dims " + f"{tuple(control_latent.shape[-2:])} must match the target " + f"latents {tuple(latents.shape[-2:])}" + ) + control_latents_list.append(control_latent) + batch.extra["control_latents"] = control_latents_list + self.log_info( + f"Prepared {len(control_latents_list)} control latent block(s) " + f"with shape {tuple(control_latents_list[0].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): @@ -868,11 +1115,18 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): parallelism_type = StageParallelismType.REPLICATED - def __init__(self, transformer, scheduler, server_args: ServerArgs | None = None): + def __init__( + self, + transformer, + scheduler, + server_args: ServerArgs | None = None, + vae=None, + ): super().__init__() self.transformer = transformer self.scheduler = scheduler self.server_args = server_args + self.vae = vae self._logged_parallel_config = False self._logged_cfg_split = False @@ -971,6 +1225,7 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): action_noisy_mask: torch.Tensor | None = None, action_fps: float | None = None, action_start_frame_offset: int = 1, + control_latents: list[torch.Tensor] | None = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Run transformer forward pass. @@ -1004,6 +1259,10 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): action_noisy_mask=action_noisy_mask, action_fps=action_fps, action_start_frame_offset=action_start_frame_offset, + control_latents=control_latents, + transfer_share_vision_temporal_positions=getattr( + self, "_share_vision_temporal_positions", True + ), ) @staticmethod @@ -1020,8 +1279,207 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): lo, hi = interval return lo <= t_scalar <= hi - def forward(self, batch: Req, server_args: ServerArgs) -> Req: - """Run the denoising loop with CFG and optional I2V conditioning.""" + def _normalize_transfer_video( + self, video: torch.Tensor, device: torch.device, dtype: torch.dtype + ) -> torch.Tensor: + return video.to(device=device, dtype=dtype).div(127.5).sub(1.0) + + def _encode_transfer_video( + self, video: torch.Tensor, output_dtype: torch.dtype + ) -> torch.Tensor: + vae_dtype = next(self.vae.parameters()).dtype + with torch.no_grad(): + latent = self.vae.encode(video.to(dtype=vae_dtype)).mode() + mean = ( + torch.as_tensor(self.vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(latent.device, latent.dtype) + ) + std = ( + torch.as_tensor(self.vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(latent.device, latent.dtype) + ) + return ((latent - mean) / std).to(output_dtype) + + def _decode_transfer_latents(self, latents: torch.Tensor) -> torch.Tensor: + vae_dtype = next(self.vae.parameters()).dtype + latents = latents.to(vae_dtype) + mean = ( + torch.as_tensor(self.vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(latents.device, vae_dtype) + ) + std = ( + torch.as_tensor(self.vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(latents.device, vae_dtype) + ) + with torch.no_grad(): + decoded = self.vae.decode(latents * std + mean) + if hasattr(decoded, "sample"): + decoded = decoded.sample + elif isinstance(decoded, tuple): + decoded = decoded[0] + return decoded + + def _prepare_transfer_chunk( + self, + batch: Req, + chunk_id: int, + previous_output: torch.Tensor | None, + generator: torch.Generator | None, + ) -> int: + plan = batch.extra["transfer_plan"] + chunk_frames = int(plan["chunk_frames"]) + start_frame = chunk_id * int(plan["stride"]) + end_frame = min(start_frame + chunk_frames, int(plan["total_frames"])) + device = batch.latents.device + latent_dtype = batch.latents.dtype + vae_dtype = next(self.vae.parameters()).dtype + + control_norms = [ + self._normalize_transfer_video( + _pad_transfer_frames( + control[:, :, start_frame:end_frame], chunk_frames + ), + device, + vae_dtype, + ) + for control in batch.extra["preprocessed_control"] + ] + target_norm = torch.zeros_like(control_norms[0]) + current_conditional_frames = 0 + + source_video = batch.extra.get("preprocessed_transfer_video") + if ( + chunk_id == 0 + and batch.sampling_params.num_first_chunk_conditional_frames > 0 + ): + current_conditional_frames = min( + int(batch.sampling_params.num_first_chunk_conditional_frames), + source_video.shape[2], + chunk_frames, + ) + source_norm = self._normalize_transfer_video( + source_video[:, :, :current_conditional_frames], device, vae_dtype + ) + target_norm[:, :, :current_conditional_frames] = source_norm + elif chunk_id > 0 and previous_output is not None: + current_conditional_frames = min( + int(batch.sampling_params.num_conditional_frames), + previous_output.shape[2], + chunk_frames, + ) + if current_conditional_frames > 0: + target_norm[:, :, :current_conditional_frames] = previous_output[ + :, :, -current_conditional_frames: + ].to(target_norm) + + if 0 < current_conditional_frames < chunk_frames: + target_norm[:, :, current_conditional_frames:] = target_norm[ + :, :, current_conditional_frames - 1 : current_conditional_frames + ].expand(-1, -1, chunk_frames - current_conditional_frames, -1, -1) + + control_latents = [ + self._encode_transfer_video(control, latent_dtype) + for control in control_norms + ] + condition_latents = self._encode_transfer_video(target_norm, latent_dtype) + noise = torch.randn( + condition_latents.shape, + generator=generator, + device=device, + dtype=latent_dtype, + ) + condition_mask = torch.zeros( + 1, + 1, + condition_latents.shape[2], + 1, + 1, + device=device, + dtype=latent_dtype, + ) + if current_conditional_frames > 0: + temporal_scale = int(batch.extra["vae_scale_factor_temporal"]) + latent_conditional_frames = ( + current_conditional_frames - 1 + ) // temporal_scale + 1 + condition_mask[:, :, :latent_conditional_frames] = 1.0 + + velocity_mask = 1.0 - condition_mask + batch.latents = condition_mask * condition_latents + velocity_mask * noise + batch.raw_latent_shape = tuple(batch.latents.shape) + batch.extra["video_shape"] = tuple(batch.latents.shape[2:]) + batch.extra["condition_latents"] = condition_mask * condition_latents + batch.extra["velocity_mask"] = velocity_mask + batch.extra["control_latents"] = control_latents + return current_conditional_frames + + def _forward_transfer(self, batch: Req, server_args: ServerArgs) -> Req: + if self.vae is None: + raise RuntimeError("Cosmos3 Transfer denoising requires the pipeline VAE") + + device = batch.latents.device + generator = batch.generator + if generator is None and batch.seed is not None: + generator = torch.Generator(device=device).manual_seed(batch.seed) + + plan = batch.extra["transfer_plan"] + output_chunks = [] + previous_output = None + for chunk_id in range(int(plan["num_chunks"])): + with self.use_declared_component( + component_name="vae", + module=self.vae, + phase="transfer_encode", + ): + current_conditional_frames = self._prepare_transfer_chunk( + batch, chunk_id, previous_output, generator + ) + self.scheduler.set_timesteps( + batch.num_inference_steps, device=batch.latents.device + ) + batch.timesteps = self.scheduler.timesteps + with self.use_declared_component( + component_name="transformer", + module=self.transformer, + phase="denoise", + ): + self._denoise_once(batch, server_args, generator=generator) + with self.use_declared_component( + component_name="vae", + module=self.vae, + phase="transfer_decode", + ): + previous_output = self._decode_transfer_latents(batch.latents).clamp( + -1, 1 + ) + if chunk_id == 0: + output_chunks.append(previous_output) + else: + output_chunks.append(previous_output[:, :, current_conditional_frames:]) + + batch.extra["transfer_decoded_output"] = torch.cat(output_chunks, dim=2)[ + :, :, : int(plan["total_frames"]) + ] + return batch + + def _denoise_once( + self, + batch: Req, + server_args: ServerArgs, + generator: torch.Generator | None = None, + ) -> Req: + """Run one denoising loop with CFG and optional conditioning.""" + self._share_vision_temporal_positions = bool( + getattr( + batch.sampling_params, + "share_vision_temporal_positions", + True, + ) + ) latents = batch.latents sound_latents = batch.audio_latents action_latents = getattr(batch, "action_latents", None) @@ -1037,7 +1495,8 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): # 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: + generator = batch.generator if generator is None and batch.seed is not None: generator = torch.Generator(device=latents.device).manual_seed(batch.seed) @@ -1049,7 +1508,14 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): fps = batch.extra.get("fps", 24.0) velocity_mask = batch.extra.get("velocity_mask") condition_latents = batch.extra.get("condition_latents") + control_latents = batch.extra.get("control_latents") guidance_interval = getattr(batch.sampling_params, "guidance_interval", None) + control_guidance = getattr(batch.sampling_params, "control_guidance", 1.0) + if control_guidance is None: + control_guidance = 1.0 + control_guidance_interval = getattr( + batch.sampling_params, "control_guidance_interval", None + ) # Rollout requests carry a per-request scheduler bound by the timestep stage. scheduler = batch.scheduler if batch.scheduler is not None else self.scheduler @@ -1083,8 +1549,13 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): ) do_cfg = guidance_scale > 1.0 + # Control-CFG runs even when text guidance is off (its own extra + # control-dropped forward), so it can drive CFG parallel on its own. + any_control_cfg = control_latents is not None and control_guidance != 1.0 - enable_cfg_parallel = server_args.enable_cfg_parallel and do_cfg + enable_cfg_parallel = server_args.enable_cfg_parallel and ( + do_cfg or any_control_cfg + ) if action_latents is not None and enable_cfg_parallel: raise NotImplementedError( "Cosmos3 action generation does not support CFG parallel yet" @@ -1147,10 +1618,74 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): effective_scale = ( guidance_scale if self._cfg_active_at(t, guidance_interval) else 1.0 ) + # Transfer control-CFG: active only when a control video is present, + # ``control_guidance != 1.0``, and the step is inside the (optional) + # control window. It needs a second control-dropped forward, so it + # owns the prediction for the step and composes text CFG internally. + control_cfg_active = ( + control_latents is not None + and control_guidance != 1.0 + and self._cfg_active_at(t, control_guidance_interval) + ) - if do_cfg: - if enable_cfg_parallel: - noise_pred = self._predict_noise_cfg_parallel( + if control_cfg_active: + # Control-CFG owns the step: 2 branches (text guidance off) or 3 + # (text guidance on), distributed across CFG ranks and reduced by + # ``_predict_noise_cfg`` (sequential per rank, no batching). + branches = self._control_cfg_branches( + cond_text_ids, + cond_text_mask, + uncond_text_ids, + uncond_text_mask, + cond_text_seq_len=batch.extra["cond_text_seq_len"], + uncond_text_seq_len=batch.extra["uncond_text_seq_len"], + control_latents=control_latents, + text_guidance_scale=effective_scale, + control_guidance_scale=control_guidance, + ) + noise_pred = self._predict_noise_cfg( + branches, + latents=latents, + timestep=timestep, + video_shape=video_shape, + fps=fps, + cfg_rank=cfg_rank, + cfg_world_size=cfg_world_size, + noisy_frame_mask=velocity_mask, + 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 do_cfg and effective_scale != 1.0: + cond_text_seq_len = batch.extra["cond_text_seq_len"] + uncond_text_seq_len = batch.extra["uncond_text_seq_len"] + text_seq_lens_differ = cond_text_seq_len != uncond_text_seq_len + if ( + text_seq_lens_differ + and 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" + ) + single_cfg_rank_control_free = ( + cfg_world_size == 1 and control_latents is None + ) + can_batch_text_cfg = ( + single_cfg_rank_control_free and not text_seq_lens_differ + ) + if can_batch_text_cfg: + # Single-CFG-rank, control-free text CFG: one batched forward + # (lower launch overhead, no control tokens to duplicate). + noise_pred = self._predict_noise_cfg_batched( latents=latents, timestep=timestep, cond_text_ids=cond_text_ids, @@ -1160,10 +1695,8 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): video_shape=video_shape, fps=fps, guidance_scale=effective_scale, - cfg_rank=cfg_rank, noisy_frame_mask=velocity_mask, - cond_text_seq_len=batch.extra["cond_text_seq_len"], - uncond_text_seq_len=batch.extra["uncond_text_seq_len"], + max_text_seq_len=cond_text_seq_len, current_timestep=i, sound_latents=sound_latents, action_latents=action_latents, @@ -1172,17 +1705,24 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): action_fps=action_fps, action_start_frame_offset=action_start_frame_offset, ) - elif effective_scale == 1.0: - noise_pred = self._run_transformer( + elif single_cfg_rank_control_free: + # Keep each branch at its native text length, but preserve the + # canonical CFG operation order used by the batched path. The + # algebraically equivalent coefficient sum rounds differently + # in BF16 and changes deterministic generation results. + noise_pred = self._predict_noise_text_cfg_serial( latents=latents, timestep=timestep, - text_ids=cond_text_ids, - text_mask=cond_text_mask, + 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, - cache_key="cond", + guidance_scale=effective_scale, noisy_frame_mask=velocity_mask, - max_text_seq_len=batch.extra["cond_text_seq_len"], + cond_text_seq_len=cond_text_seq_len, + uncond_text_seq_len=uncond_text_seq_len, current_timestep=i, sound_latents=sound_latents, action_latents=action_latents, @@ -1192,19 +1732,28 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): action_start_frame_offset=action_start_frame_offset, ) else: + # CFG parallel or control passthrough: distribute unbatched + # branches across ranks. Separate forwards preserve each + # branch's native text length. + branches = self._text_cfg_branches( + cond_text_ids, + cond_text_mask, + uncond_text_ids, + uncond_text_mask, + guidance_scale=effective_scale, + cond_text_seq_len=cond_text_seq_len, + uncond_text_seq_len=uncond_text_seq_len, + control_latents=control_latents, + ) noise_pred = self._predict_noise_cfg( + branches, 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=effective_scale, + cfg_rank=cfg_rank, + cfg_world_size=cfg_world_size, noisy_frame_mask=velocity_mask, - 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, @@ -1214,6 +1763,8 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): action_start_frame_offset=action_start_frame_offset, ) else: + # No CFG this step (guidance off or outside the CFG window): a + # single conditional forward, run identically on every rank. noise_pred = self._run_transformer( latents=latents, timestep=timestep, @@ -1231,6 +1782,7 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): action_noisy_mask=action_velocity_mask, action_fps=action_fps, action_start_frame_offset=action_start_frame_offset, + control_latents=control_latents, ) # Unpack multi-modality outputs; ordering is (video[, action][, sound]). @@ -1337,8 +1889,100 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): self.log_info("Denoising complete") return batch + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + if batch.extra.get("transfer_plan") is not None: + return self._forward_transfer(batch, server_args) + with self.use_declared_component( + component_name="transformer", + module=self.transformer, + phase="denoise", + ): + return self._denoise_once(batch, server_args) + def _predict_noise_cfg( self, + branches: list[dict], + *, + latents: torch.Tensor, + timestep: torch.Tensor, + video_shape: tuple[int, int, int], + fps: float, + cfg_rank: int, + cfg_world_size: int, + noisy_frame_mask: torch.Tensor | None = None, + current_timestep: 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, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """Combine CFG branches as the weighted sum ``sum_b coeff_b * f(branch_b)``. + + Both text CFG and transfer control-CFG are linear in their per-branch + forwards, so each is expressed as a list of branches (text ids/mask, + control latents, UND cache key, coeff) and reduced here. ``branches`` is + identical on every CFG rank. + + Branches are distributed round-robin across the ``cfg_world_size`` CFG + ranks. A rank runs its branches sequentially and a final sum all-reduce + combines ranks. Forwards are never batched, which preserves each text + branch's native length and bounds activation memory with control inputs. + """ + acc = None + for i, branch in enumerate(branches): + if i % cfg_world_size != cfg_rank: + continue + out = self._run_transformer( + latents=latents, + timestep=timestep, + text_ids=branch["text_ids"], + text_mask=branch["text_mask"], + video_shape=video_shape, + fps=fps, + cache_key=branch["cache_key"], + noisy_frame_mask=noisy_frame_mask, + max_text_seq_len=branch["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, + control_latents=branch["control_latents"], + ) + coeff = branch["coeff"] + if isinstance(out, tuple): + scaled = tuple(coeff * prediction for prediction in out) + acc = ( + scaled + if acc is None + else tuple( + total + contribution + for total, contribution in zip(acc, scaled, strict=True) + ) + ) + else: + scaled = coeff * out + acc = scaled if acc is None else acc + scaled + + if acc is None: + # More ranks than branches: contribute zeros to the all-reduce. + acc = self._zero_like_output(latents, action_latents, sound_latents) + + if cfg_world_size > 1: + if isinstance(acc, tuple): + return tuple( + cfg_model_parallel_all_reduce(prediction) for prediction in acc + ) + return cfg_model_parallel_all_reduce(acc) + return acc + + def _predict_noise_text_cfg_serial( + self, + *, latents: torch.Tensor, timestep: torch.Tensor, cond_text_ids: torch.Tensor, @@ -1348,78 +1992,66 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): video_shape: tuple[int, int, int], fps: float, guidance_scale: float, - cond_text_seq_len: int, - uncond_text_seq_len: int, - **kwargs: Any, + noisy_frame_mask: torch.Tensor | None = None, + cond_text_seq_len: int | None = None, + uncond_text_seq_len: int | None = None, + current_timestep: 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, ) -> torch.Tensor | tuple[torch.Tensor, ...]: - """Run CFG, batching the two branches only when that is lossless. + """Run control-free text CFG as separate native-length forwards. - 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. + Keep the canonical uncond + g * (cond - uncond) operation order. + Although a coefficient-weighted sum is algebraically equivalent, it is + not numerically equivalent in BF16 and would make results depend on + whether the two prompts happen to tokenize to the same 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" - ) - + common_kwargs = { + "latents": latents, + "timestep": timestep, + "video_shape": video_shape, + "fps": fps, + "noisy_frame_mask": noisy_frame_mask, + "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, + "control_latents": None, + } 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, + **common_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, + **common_kwargs, ) - def _cfg_combine( + def _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) + return tuple(_combine(c, u) for c, u in zip(cond, uncond, strict=True)) + return _combine(cond, uncond) def _predict_noise_cfg_batched( self, + *, latents: torch.Tensor, timestep: torch.Tensor, cond_text_ids: torch.Tensor, @@ -1439,145 +2071,227 @@ class Cosmos3DenoisingStage(PipelineStage, RolloutDenoisingMixin): 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. + """Run CFG as one ``batch_size=2`` forward stacking both branches (``[uncond, cond]``). - Halves the kernel-launch count vs running cond and uncond serially. - Order is ``[uncond, cond]`` so the chunk-and-combine math below - matches the standard CFG formula. + Kept only for the non-parallel, control-free text path: one batched + forward has lower kernel-launch overhead than two serial ones, and + doubling the GEN tokens is cheap here. The caller does not route control + through this (batching the larger control-in forwards risks OOM on big + models — that path uses ``_predict_noise_cfg`` instead) and CFG parallel + splits branches across ranks rather than batching. """ - latents_batched = torch.cat([latents, latents], dim=0) - text_ids_batched = torch.cat([uncond_text_ids, cond_text_ids], dim=0) - text_mask_batched = torch.cat([uncond_text_mask, cond_text_mask], dim=0) - timestep_batched = torch.cat([timestep, timestep], dim=0) - mask_batched = ( + latents_b = torch.cat([latents, latents], dim=0) + text_ids_b = torch.cat([uncond_text_ids, cond_text_ids], dim=0) + text_mask_b = torch.cat([uncond_text_mask, cond_text_mask], dim=0) + timestep_b = torch.cat([timestep, timestep], dim=0) + mask_b = ( torch.cat([noisy_frame_mask, noisy_frame_mask], dim=0) if noisy_frame_mask is not None else None ) - sound_batched = ( + sound_b = ( torch.cat([sound_latents, sound_latents], dim=0) if sound_latents is not None else None ) - action_batched = ( + action_b = ( torch.cat([action_latents, action_latents], dim=0) if action_latents is not None else None ) - action_domain_ids_batched = ( + action_domain_b = ( torch.cat([action_domain_ids, action_domain_ids], dim=0) if action_domain_ids is not None else None ) - action_noisy_mask_batched = ( + action_mask_b = ( torch.cat([action_noisy_mask, action_noisy_mask], dim=0) if action_noisy_mask is not None else None ) out = self._run_transformer( - latents=latents_batched, - timestep=timestep_batched, - text_ids=text_ids_batched, - text_mask=text_mask_batched, + latents=latents_b, + timestep=timestep_b, + text_ids=text_ids_b, + text_mask=text_mask_b, video_shape=video_shape, fps=fps, cache_key="cfg_batched", - noisy_frame_mask=mask_batched, + noisy_frame_mask=mask_b, 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, + sound_latents=sound_b, + action_latents=action_b, + action_domain_ids=action_domain_b, + action_noisy_mask=action_mask_b, action_fps=action_fps, action_start_frame_offset=action_start_frame_offset, + control_latents=None, ) - def _cfg_combine(pred: torch.Tensor) -> torch.Tensor: - uncond, cond = pred.chunk(2, dim=0) + def _combine(o: torch.Tensor) -> torch.Tensor: + uncond, cond = o.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) + return tuple(_combine(p) for p in out) + return _combine(out) - def _predict_noise_cfg_parallel( - self, + @staticmethod + def _zero_like_output( latents: torch.Tensor, - timestep: torch.Tensor, + action_latents: torch.Tensor | None, + sound_latents: torch.Tensor | None, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """Zero prediction matching the forward's (video[, action][, sound]) layout.""" + zeros = [torch.zeros_like(latents)] + if action_latents is not None: + zeros.append(torch.zeros_like(action_latents)) + if sound_latents is not None: + zeros.append(torch.zeros_like(sound_latents)) + return zeros[0] if len(zeros) == 1 else tuple(zeros) + + @staticmethod + def _text_cfg_branches( 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, - cfg_rank: int, - noisy_frame_mask: torch.Tensor | None = None, - cond_text_seq_len: int | None = None, - uncond_text_seq_len: int | None = None, - current_timestep: 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, - ) -> torch.Tensor | tuple[torch.Tensor, ...]: - """Run CFG with one branch per CFG rank, combined by all-reduce. + cond_text_seq_len: int | None, + uncond_text_seq_len: int | None, + control_latents: list[torch.Tensor] | None, + ) -> list[dict]: + """Standard text CFG as two branches: ``g*cond + (1-g)*uncond``. - 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"``). 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. + Control latents (if any) pass through both branches unchanged — text CFG + does not drop the control map; that is what control-CFG adds. """ - if cfg_rank == 0: - text_ids, text_mask, cache_key = cond_text_ids, cond_text_mask, "cond" - text_seq_len = cond_text_seq_len - coeff = guidance_scale - else: - 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 [ + { + "cache_key": "cond", + "text_ids": cond_text_ids, + "text_mask": cond_text_mask, + "text_seq_len": cond_text_seq_len, + "control_latents": control_latents, + "coeff": guidance_scale, + }, + { + "cache_key": "uncond", + "text_ids": uncond_text_ids, + "text_mask": uncond_text_mask, + "text_seq_len": uncond_text_seq_len, + "control_latents": control_latents, + "coeff": 1.0 - guidance_scale, + }, + ] - 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, - ) + @staticmethod + def _control_cfg_branches( + cond_text_ids: torch.Tensor, + cond_text_mask: torch.Tensor, + uncond_text_ids: torch.Tensor, + uncond_text_mask: torch.Tensor, + *, + cond_text_seq_len: int | None, + uncond_text_seq_len: int | None, + control_latents: list[torch.Tensor] | None, + text_guidance_scale: float, + control_guidance_scale: float, + ) -> list[dict]: + """Transfer control-CFG (optionally composed with text CFG) as branches. - 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) + Two conditional forwards share the cond-text branch but differ in whether + the control map is packed in: + + - ``cond_full`` — control clips in (the standard transfer forward) + - ``cond_nc`` — control clips dropped (``control_latents=None``) + + mixed on the generated span as ``cond = cond_nc + cg*(cond_full - + cond_nc)``. When text CFG is also active a third (uncond, control-in) + forward composes the text blend ``pred = uncond + g*(cond - uncond)`` on + top. Expanding both gives the coefficient-weighted sum reduced by + ``_predict_noise_cfg``:: + + pred = g*cg*cond_full + g*(1-cg)*cond_nc + (1-g)*uncond + + (with ``g = 1`` collapsing to ``cg*cond_full + (1-cg)*cond_nc``). + + Branch order places the two control-in forwards first and ``cond_nc`` + second so the round-robin split in ``_predict_noise_cfg`` lands the + control-in pair on rank 0 and ``cond_nc`` on rank 1 under 2-rank CFG. + ``cond_full`` and ``cond_nc`` reuse distinct UND cache keys (``"cond"`` / + ``"cond_nc"``) because their GEN rope layout differs. + """ + g = text_guidance_scale + cg = control_guidance_scale + cond_full = { + "cache_key": "cond", + "text_ids": cond_text_ids, + "text_mask": cond_text_mask, + "text_seq_len": cond_text_seq_len, + "control_latents": control_latents, + } + cond_nc = { + "cache_key": "cond_nc", + "text_ids": cond_text_ids, + "text_mask": cond_text_mask, + "text_seq_len": cond_text_seq_len, + "control_latents": None, + } + if g == 1.0: + cond_full["coeff"] = cg + cond_nc["coeff"] = 1.0 - cg + return [cond_full, cond_nc] + cond_full["coeff"] = g * cg + cond_nc["coeff"] = g * (1.0 - cg) + uncond = { + "cache_key": "uncond", + "text_ids": uncond_text_ids, + "text_mask": uncond_text_mask, + "text_seq_len": uncond_text_seq_len, + "control_latents": control_latents, + "coeff": 1.0 - g, + } + return [cond_full, cond_nc, uncond] def component_uses( self, server_args: ServerArgs, stage_name: str | None = None ) -> list[ComponentUse]: - return [ + stage_name = self._component_stage_name(stage_name) + uses = [ ComponentUse( - self._component_stage_name(stage_name), + stage_name, "transformer", + phase="denoise", preferred_ready_after_request=True, memory_intensive=True, + start_at_stage_entry=False, ) ] + if self.vae is not None: + uses = [ + ComponentUse( + stage_name, + "vae", + phase="transfer_encode", + allow_prefetch=False, + keep_ready_after_warmup=True, + start_at_stage_entry=False, + ), + *uses, + ComponentUse( + stage_name, + "vae", + phase="transfer_decode", + allow_prefetch=False, + keep_ready_after_warmup=True, + start_at_stage_entry=False, + ), + ] + return uses class Cosmos3DecodingStage(PipelineStage): @@ -1656,6 +2370,32 @@ class Cosmos3DecodingStage(PipelineStage): return video.squeeze(2).permute(0, 2, 3, 1).cpu().numpy() return video.permute(0, 2, 3, 4, 1).cpu().numpy() + @staticmethod + def _compose_transfer_display(output: torch.Tensor, batch: Req) -> torch.Tensor: + plan = batch.extra.get("transfer_plan") + if plan is None: + return output + total_frames = int(plan["total_frames"]) + if batch.sampling_params.show_control_condition: + controls = batch.extra["preprocessed_control"] + controls = controls if isinstance(controls, list) else [controls] + normalized_controls = [ + control[:, :, :total_frames] + .to(device=output.device, dtype=output.dtype) + .div(255.0) + for control in controls + ] + output = torch.cat([*normalized_controls, output], dim=-1) + source = batch.extra.get("preprocessed_transfer_video") + if batch.sampling_params.show_input and source is not None: + normalized_source = ( + source[:, :, :total_frames] + .to(device=output.device, dtype=output.dtype) + .div(255.0) + ) + output = torch.cat([normalized_source, output], dim=-1) + return output + def forward(self, batch: Req, server_args: ServerArgs): """Decode latents to video, or to a single image for T2I.""" from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import ( @@ -1723,12 +2463,15 @@ class Cosmos3DecodingStage(PipelineStage): ) device = batch.latents.device - with self.use_declared_component(component_name="vae", module=self.vae): - with torch.no_grad(): - decoded = self._decode_latents(batch.latents) + decoded = batch.extra.get("transfer_decoded_output") + if decoded is None: + with self.use_declared_component(component_name="vae", module=self.vae): + with torch.no_grad(): + decoded = self._decode_latents(batch.latents) self.log_debug("Decoded tensor shape: %s", decoded.shape) output = self._postprocess_tensor(decoded) + output = self._compose_transfer_display(output, batch) if self._guardrails and batch.use_guardrails is not False: from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_guardrails import ( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py index a475aec5f..343e89db9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py @@ -10,12 +10,17 @@ Enabled by default when available; opt out with from __future__ import annotations +import hashlib import importlib.util +import os +import shutil from functools import lru_cache +from pathlib import Path import numpy as np import torch +from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, @@ -35,6 +40,68 @@ def is_cosmos_guardrail_available() -> bool: return importlib.util.find_spec("cosmos_guardrail") is not None +def _mirror_symlinked_nltk_data() -> None: + """Make the guardrail's nltk_data readable under NLTK's hardened opener. + + ``CosmosSafetyChecker`` registers its HF-hub snapshot's + ``blocklist/nltk_data`` directory on ``nltk.data.path``. Hub snapshot files + are symlinks into the blob store, which NLTK builds that ship the + ``pathsec`` hardened opener refuse to follow (O_NOFOLLOW, CWE-59 TOCTOU + guard) — every text-safety check then fails with "refusing to follow a + symlink at open time". Mirror each symlink-containing search entry to a + plain-file copy and register the mirror ahead of the original, so NLTK's + data lookup resolves to real files first. No-op for NLTK builds without + the hardened opener and for search entries that are already plain files. + """ + try: + import nltk.data + except ImportError: + return + + mirror_root = Path(envs.SGLANG_DIFFUSION_CACHE_ROOT) / "nltk_data_deref" + for entry in list(nltk.data.path): + try: + root = Path(entry) + if not root.is_dir(): + continue + if not any(p.is_symlink() for p in root.rglob("*")): + continue + mirror = mirror_root / hashlib.sha256(str(root).encode()).hexdigest()[:16] + if str(mirror) in nltk.data.path: + continue + if not mirror.is_dir(): + # Every GPU worker process runs this at pipeline construction on + # a shared filesystem, so stage under a per-process name and + # publish with an atomic rename; whichever process publishes + # first wins and the others adopt its mirror. + staging = mirror.with_name(f"{mirror.name}.{os.getpid()}.tmp") + shutil.rmtree(staging, ignore_errors=True) + # symlinks=False dereferences: the copy holds real file contents. + shutil.copytree(root, staging, symlinks=False) + try: + staging.rename(mirror) + except OSError: + shutil.rmtree(staging, ignore_errors=True) + if not mirror.is_dir(): + raise + nltk.data.path.insert(nltk.data.path.index(entry), str(mirror)) + logger.info( + "Mirrored symlinked nltk_data %s -> %s (hardened-NLTK compatibility)", + root, + mirror, + ) + except OSError as exc: + # Best-effort: an unwritable cache root must not break guardrail + # init. Hardened-NLTK builds may still fail at check time; plain + # NLTK builds work fine without the mirror. + logger.warning( + "Could not mirror symlinked nltk_data %s under %s: %s", + entry, + mirror_root, + exc, + ) + + def _init_guardrails(offload_to_cpu: bool = False) -> None: global _checker if _checker is not None: @@ -50,6 +117,7 @@ def _init_guardrails(offload_to_cpu: bool = False) -> None: "Initializing Cosmos3 guardrails (offload_to_cpu=%s) ...", offload_to_cpu ) _checker = CosmosSafetyChecker() + _mirror_symlinked_nltk_data() idle_device = "cpu" if offload_to_cpu else current_platform.device_type for runner in (_checker.text_guardrail, _checker.video_guardrail): if runner is None or not hasattr(runner, "models"): diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py index 67348c15f..df2d9090a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py @@ -18,7 +18,10 @@ 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.configs.sample.sampling_params import ( + DataType, + SamplingParams, +) from sglang.multimodal_gen.registry import ( _PIPELINE_REGISTRY, _discover_and_register_pipelines, @@ -36,6 +39,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( ) from sglang.multimodal_gen.runtime.entrypoints.openai.video_api import ( _cosmos3_sampling_param_kwargs, + _multipart_video_extras, _resolve_sound_duration, _resolve_video_path, ) @@ -45,6 +49,7 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.scheduler_loader imp ) from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.models.dits.cosmos3video import ( + Cosmos3OmniTransformer, DomainAwareLinear, _can_enable_t1_fused_qk_norm_rope, compute_mrope_position_ids_action, @@ -59,6 +64,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.c Cosmos3TimestepPreparationStage, Cosmos3TokenizationStage, _inject_caption_metadata, + _pad_transfer_frames, + _resize_center_crop_uint8_cthw, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.cosmos3_action import ( EMBODIMENT_TO_DOMAIN_ID, @@ -877,7 +884,7 @@ class TestCosmos3ModelResolution(unittest.TestCase): class TestCosmos3OpenAIProtocol(unittest.TestCase): - """Verify Cosmos3 modality knobs are exposed by the video HTTP schema.""" + """Verify Cosmos3 modality knobs stay model-specific video extras.""" def test_cosmos3_template_fields_remain_extra_fields(self): for request_cls in (ImageGenerationsRequest, VideoGenerationsRequest): @@ -889,31 +896,45 @@ class TestCosmos3OpenAIProtocol(unittest.TestCase): self.assertNotIn("use_system_prompt", request_cls.model_fields) self.assertNotIn("use_guardrails", request_cls.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) - - 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(): + def test_cosmos3_modal_fields_are_model_specific_video_extras(self): + for field_name in ( + "generate_sound", + "sound_duration", + "condition_frame_indexes", + "condition_frame_indexes_vision", + "condition_video_keep", + "control_path", + "control_hint", + "control_guidance", + "control_guidance_interval", + "num_video_frames_per_chunk", + "num_conditional_frames", + "num_first_chunk_conditional_frames", + "max_frames", + "show_control_condition", + "show_input", + "share_vision_temporal_positions", + "action_mode", + "domain_id", + "domain_name", + "raw_action_dim", + "action_fps", + "action", + "action_view_point", + "action_normalization", + ): with self.subTest(field_name=field_name): self.assertNotIn(field_name, VideoGenerationsRequest.model_fields) - self.assertEqual(getattr(req, field_name), value) + self.assertIn( + field_name, Cosmos3SamplingParams.video_request_extra_fields() + ) + + self.assertIn("video_path", VideoGenerationsRequest.model_fields) + self.assertIn("video_url", VideoGenerationsRequest.model_fields) + self.assertNotIn("action_stats_path", VideoGenerationsRequest.model_fields) + self.assertNotIn( + "action_stats_path", Cosmos3SamplingParams.video_request_extra_fields() + ) def test_cosmos3_http_aliases_map_to_sampling_params(self): req = VideoGenerationsRequest( @@ -927,6 +948,17 @@ class TestCosmos3OpenAIProtocol(unittest.TestCase): raw_action_dim=9, action_fps=30.0, action_view_point="ego_view", + control_path=["edge.mp4", "depth.mp4"], + control_hint=["edge", "depth"], + control_guidance=1.5, + control_guidance_interval=[0.0, 500.0], + num_video_frames_per_chunk=97, + num_conditional_frames=5, + num_first_chunk_conditional_frames=2, + max_frames=1200, + show_control_condition="true", + show_input="false", + share_vision_temporal_positions="false", ) self.assertEqual(_resolve_video_path(req), "https://example.com/input.mp4") @@ -940,6 +972,48 @@ class TestCosmos3OpenAIProtocol(unittest.TestCase): self.assertEqual(kwargs["raw_action_dim"], 9) self.assertEqual(kwargs["action_fps"], 30.0) self.assertEqual(kwargs["action_view_point"], "ego_view") + self.assertEqual(kwargs["control_path"], ["edge.mp4", "depth.mp4"]) + self.assertEqual(kwargs["control_hint"], ["edge", "depth"]) + self.assertEqual(kwargs["control_guidance"], 1.5) + self.assertEqual(kwargs["control_guidance_interval"], (0.0, 500.0)) + self.assertEqual(kwargs["num_video_frames_per_chunk"], 97) + self.assertEqual(kwargs["num_conditional_frames"], 5) + self.assertEqual(kwargs["num_first_chunk_conditional_frames"], 2) + self.assertEqual(kwargs["max_frames"], 1200) + self.assertTrue(kwargs["show_control_condition"]) + self.assertFalse(kwargs["show_input"]) + self.assertFalse(kwargs["share_vision_temporal_positions"]) + + def test_cosmos3_multipart_extras_are_model_specific(self): + raw_form = { + "generate_sound": "true", + "control_path": '["edge.mp4", "depth.mp4"]', + "control_hint": '["edge", "depth"]', + "action_mode": "policy", + "action_stats_path": "/tmp/action_stats.json", + } + + generic = _multipart_video_extras( + raw_form, + extra_body=None, + extra_params=None, + sampling_params_cls=SamplingParams, + ) + self.assertNotIn("generate_sound", generic) + self.assertNotIn("control_path", generic) + self.assertNotIn("action_mode", generic) + + cosmos = _multipart_video_extras( + raw_form, + extra_body=None, + extra_params=None, + sampling_params_cls=Cosmos3SamplingParams, + ) + self.assertIs(cosmos["generate_sound"], True) + self.assertEqual(cosmos["control_path"], ["edge.mp4", "depth.mp4"]) + self.assertEqual(cosmos["control_hint"], ["edge", "depth"]) + self.assertEqual(cosmos["action_mode"], "policy") + self.assertNotIn("action_stats_path", cosmos) def test_generate_sound_false_disables_sound_duration(self): req = VideoGenerationsRequest( @@ -1069,6 +1143,403 @@ class TestCosmos3MRoPE(unittest.TestCase): self.assertAlmostEqual(float(vid[0, 1]), float(act[0, 4]), places=4) +class TestCosmos3Transfer(unittest.TestCase): + """Transfer (control-video) conditioning: rope packing, control-CFG, defaults.""" + + DEVICE = torch.device("cpu") + + @staticmethod + def _transformer_self() -> Cosmos3OmniTransformer: + model = Cosmos3OmniTransformer.__new__(Cosmos3OmniTransformer) + model.temporal_margin = 15000 + model.base_fps = 24.0 + model.temporal_compression_factor = 4 + model.sound_latent_fps = 25.0 + model.temporal_compression_factor_sound = 1 + return model + + def test_single_control_shares_video_positions(self): + model = self._transformer_self() + text_mask = torch.ones(1, 4) + T, Hp, Wp = 3, 2, 2 + tpc = T * Hp * Wp + + _, base = model._compute_rope_position_ids( + text_mask, T, Hp, Wp, fps=None, device=self.DEVICE, control_frames=0 + ) + _, with_ctrl = model._compute_rope_position_ids( + text_mask, T, Hp, Wp, fps=None, device=self.DEVICE, control_frames=T + ) + self.assertEqual(tuple(with_ctrl.shape), (3, 1, 2 * tpc)) + # control prefix == video block, and video block is unchanged. + self.assertTrue(torch.equal(with_ctrl[:, :, :tpc], with_ctrl[:, :, tpc:])) + self.assertTrue(torch.equal(with_ctrl[:, :, tpc:], base)) + + def test_multi_control_prepended_in_order(self): + model = self._transformer_self() + text_mask = torch.ones(1, 4) + T, Hp, Wp = 3, 2, 2 + tpc = T * Hp * Wp + + _, base = model._compute_rope_position_ids( + text_mask, T, Hp, Wp, fps=None, device=self.DEVICE, control_frames=0 + ) + _, multi = model._compute_rope_position_ids( + text_mask, T, Hp, Wp, fps=None, device=self.DEVICE, control_frames=[T, T] + ) + self.assertEqual(tuple(multi.shape), (3, 1, 3 * tpc)) + c0 = multi[:, :, :tpc] + c1 = multi[:, :, tpc : 2 * tpc] + vid = multi[:, :, 2 * tpc :] + self.assertTrue(torch.equal(c0, vid)) + self.assertTrue(torch.equal(c1, vid)) + self.assertTrue(torch.equal(vid, base)) + + def test_control_positions_can_be_sequential(self): + model = self._transformer_self() + text_mask = torch.ones(1, 4) + T, Hp, Wp = 3, 1, 1 + + _, positions = model._compute_rope_position_ids( + text_mask, + T, + Hp, + Wp, + fps=None, + device=self.DEVICE, + control_frames=[T, T], + share_vision_temporal_positions=False, + ) + c0 = positions[0, 0, :T] + c1 = positions[0, 0, T : 2 * T] + video = positions[0, 0, 2 * T :] + self.assertLess(c0.max().item(), c1.min().item()) + self.assertLess(c1.max().item(), video.min().item()) + + def test_transfer_chunk_count_and_reflection_padding(self): + self.assertEqual( + Cosmos3ImagePreprocessStage._get_transfer_num_chunks(93, 93, 1), + (1, 93), + ) + self.assertEqual( + Cosmos3ImagePreprocessStage._get_transfer_num_chunks(186, 93, 1), + (3, 92), + ) + frames = torch.tensor([0, 1, 2], dtype=torch.uint8).view(1, 1, 3, 1, 1) + frames = frames.expand(1, 3, 3, 1, 1) + padded = _pad_transfer_frames(frames, 5) + self.assertEqual(padded[0, 0, :, 0, 0].tolist(), [0, 1, 2, 2, 1]) + + def test_transfer_resize_matches_vllm_omni(self): + frames = torch.arange(3 * 2 * 2 * 3, dtype=torch.uint8).reshape(3, 2, 2, 3) + + actual = _resize_center_crop_uint8_cthw(frames, height=3, width=3) + + resized = torch.nn.functional.interpolate( + frames.permute(1, 0, 2, 3).float(), + size=(3, 5), + mode="bilinear", + align_corners=False, + ) + expected = ( + resized[:, :, :, 1:4] + .round() + .clamp(0, 255) + .to(torch.uint8) + .permute(1, 0, 2, 3) + .contiguous() + ) + self.assertTrue(torch.equal(actual, expected)) + self.assertEqual(tuple(actual.shape), (3, 2, 3, 3)) + self.assertEqual(actual.dtype, torch.uint8) + + def test_transfer_chunks_stitch_without_duplicate_overlap(self): + class Scheduler: + def __init__(self): + self.timesteps = torch.tensor([]) + self.calls = 0 + + def set_timesteps(self, steps, device): + self.calls += 1 + self.timesteps = torch.arange(steps, device=device) + + stage = Cosmos3DenoisingStage.__new__(Cosmos3DenoisingStage) + stage.vae = torch.nn.Linear(1, 1, bias=False) + stage.transformer = torch.nn.Linear(1, 1, bias=False) + stage.scheduler = Scheduler() + stage._prepare_transfer_chunk = mock.Mock(side_effect=[0, 1]) + stage._denoise_once = mock.Mock(side_effect=lambda batch, *_a, **_k: batch) + stage._decode_transfer_latents = mock.Mock( + side_effect=[ + torch.arange(5).view(1, 1, 5, 1, 1).float() / 10, + torch.arange(5, 10).view(1, 1, 5, 1, 1).float() / 10, + ] + ) + generator = torch.Generator(device="cpu").manual_seed(7) + batch = types.SimpleNamespace( + latents=torch.zeros(1), + generator=generator, + seed=7, + num_inference_steps=2, + is_warmup=False, + extra={ + "transfer_plan": { + "num_chunks": 2, + "total_frames": 9, + } + }, + ) + server_args = types.SimpleNamespace(vae_cpu_offload=False) + stage.server_args = server_args + + stage._forward_transfer(batch, server_args) + + stitched = batch.extra["transfer_decoded_output"].flatten() + self.assertTrue( + torch.allclose( + stitched, + torch.tensor([0.0, 0.1, 0.2, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9]), + ) + ) + self.assertEqual(stage.scheduler.calls, 2) + for call in stage._prepare_transfer_chunk.call_args_list: + self.assertIs(call.args[3], generator) + for call in stage._denoise_once.call_args_list: + self.assertIs(call.kwargs["generator"], generator) + + def test_transfer_display_composes_input_controls_and_output(self): + output = torch.full((1, 3, 2, 1, 2), 0.5) + control = torch.full((1, 3, 2, 1, 2), 255, dtype=torch.uint8) + source = torch.zeros((1, 3, 2, 1, 2), dtype=torch.uint8) + batch = types.SimpleNamespace( + sampling_params=types.SimpleNamespace( + show_control_condition=True, + show_input=True, + ), + extra={ + "transfer_plan": {"total_frames": 2}, + "preprocessed_control": [control], + "preprocessed_transfer_video": source, + }, + ) + + composed = Cosmos3DecodingStage._compose_transfer_display(output, batch) + + self.assertEqual(tuple(composed.shape), (1, 3, 2, 1, 6)) + self.assertTrue(torch.equal(composed[..., :2], torch.zeros_like(output))) + self.assertTrue(torch.equal(composed[..., 2:4], torch.ones_like(output))) + self.assertTrue(torch.equal(composed[..., 4:], output)) + + def test_control_cfg_blend_math(self): + stage = Cosmos3DenoisingStage.__new__(Cosmos3DenoisingStage) + + # Per-branch forward values (each run un-batched, bs=1): cond_full + # (control in, cond text) -> 20; cond_nc (control dropped) -> 10; uncond + # (control in, uncond text) -> 2. The unified executor reduces the + # coefficient-weighted branch sum built by _control_cfg_branches. + def fake_run(**kw): + bs = kw["latents"].shape[0] + if kw["control_latents"] is None: + return torch.full((bs,), 10.0) # cond_nc, control dropped + if kw["cache_key"] == "uncond": + return torch.full((bs,), 2.0) # uncond, control in + return torch.full((bs,), 20.0) # cond_full, control in + + stage._run_transformer = fake_run + cond_text_ids = torch.zeros(1) + cond_text_mask = torch.ones(1) + uncond_text_ids = torch.zeros(1) + uncond_text_mask = torch.ones(1) + control_latents = [torch.zeros(1)] + + def _run(text_g, control_g): + branches = stage._control_cfg_branches( + cond_text_ids, + cond_text_mask, + uncond_text_ids, + uncond_text_mask, + cond_text_seq_len=None, + uncond_text_seq_len=None, + control_latents=control_latents, + text_guidance_scale=text_g, + control_guidance_scale=control_g, + ) + return stage._predict_noise_cfg( + branches, + latents=torch.zeros(1), + timestep=torch.zeros(1), + video_shape=(1, 1, 1), + fps=24.0, + cfg_rank=0, + cfg_world_size=1, + ) + + # control-CFG only (g=1): cg*cond_full + (1-cg)*cond_nc + # = 2*20 + (-1)*10 = 30 + out = _run(text_g=1.0, control_g=2.0) + self.assertTrue(torch.allclose(out, torch.full((1,), 30.0))) + + # control-CFG + text CFG (g=3): g*cg*cond_full + g*(1-cg)*cond_nc + # + (1-g)*uncond = 6*20 + (-3)*10 + (-2)*2 = 86 + out2 = _run(text_g=3.0, control_g=2.0) + self.assertTrue(torch.allclose(out2, torch.full((1,), 86.0))) + + def test_text_cfg_batched_single_gpu(self): + """Single-GPU text CFG batches both branches into one bs=2 forward.""" + stage = Cosmos3DenoisingStage.__new__(Cosmos3DenoisingStage) + + def fake_run(**kw): + self.assertIsNone(kw["control_latents"]) + self.assertEqual(kw["latents"].shape[0], 2) # batched [uncond, cond] + return torch.tensor([2.0, 20.0]) + + stage._run_transformer = fake_run + out = stage._predict_noise_cfg_batched( + latents=torch.zeros(1), + timestep=torch.zeros(1), + cond_text_ids=torch.zeros(1), + cond_text_mask=torch.ones(1), + uncond_text_ids=torch.zeros(1), + uncond_text_mask=torch.ones(1), + video_shape=(1, 1, 1), + fps=24.0, + guidance_scale=3.0, + ) + # uncond + g*(cond - uncond) = 2 + 3*(20-2) = 56 + self.assertTrue(torch.allclose(out, torch.full((1,), 56.0))) + + def test_control_cfg_parallel_distribution(self): + """Round-robin branch distribution + all-reduce equals the serial blend.""" + stage = Cosmos3DenoisingStage.__new__(Cosmos3DenoisingStage) + seen = {} + + def make_run(rank): + seen[rank] = [] + + def fake_run(**kw): + seen[rank].append(kw["cache_key"]) + bs = kw["latents"].shape[0] + if kw["control_latents"] is None: + return torch.full((bs,), 10.0) # cond_nc + if kw["cache_key"] == "uncond": + return torch.full((bs,), 2.0) # uncond + return torch.full((bs,), 20.0) # cond_full + + return fake_run + + def run(rank, world): + stage._run_transformer = make_run(rank) + branches = stage._control_cfg_branches( + torch.zeros(1), + torch.ones(1), + torch.zeros(1), + torch.ones(1), + cond_text_seq_len=None, + uncond_text_seq_len=None, + control_latents=[torch.zeros(1)], + text_guidance_scale=3.0, + control_guidance_scale=2.0, + ) + return stage._predict_noise_cfg( + branches, + latents=torch.zeros(1), + timestep=torch.zeros(1), + video_shape=(1, 1, 1), + fps=24.0, + cfg_rank=rank, + cfg_world_size=world, + ) + + target = ( + "sglang.multimodal_gen.runtime.pipelines_core.stages." + "model_specific_stages.cosmos3.cfg_model_parallel_all_reduce" + ) + # Stand in for the all-reduce with identity so each rank returns its own + # partial; summing them mimics the cross-rank reduction. + with mock.patch(target, side_effect=lambda x: x): + # 2 ranks: rank 0 runs the two control-in forwards (cond_full + + # uncond), rank 1 runs only the control-dropped cond_nc. + r0 = run(0, 2) + r1 = run(1, 2) + self.assertEqual(seen[0], ["cond", "uncond"]) + self.assertEqual(seen[1], ["cond_nc"]) + self.assertTrue(torch.allclose(r0 + r1, torch.full((1,), 86.0))) + + seen.clear() + # 4 ranks for 3 branches: rank 3 is idle and contributes zeros. + parts = [run(r, 4) for r in range(4)] + self.assertEqual(seen[0], ["cond"]) + self.assertEqual(seen[1], ["cond_nc"]) + self.assertEqual(seen[2], ["uncond"]) + self.assertEqual(seen[3], []) + self.assertTrue(torch.allclose(parts[3], torch.zeros(1))) + self.assertTrue(torch.allclose(sum(parts), torch.full((1,), 86.0))) + + def test_single_hint_defaults_applied(self): + sp = Cosmos3SamplingParams( + prompt="t", control_path="edge.mp4", control_hint="edge" + ) + sp._explicit_fields = {"prompt", "control_path", "control_hint"} + sp._apply_transfer_hint_defaults() + self.assertEqual(sp.control_guidance, 1.5) + self.assertEqual(sp.guidance_scale, 3.0) + self.assertEqual(sp.flow_shift, 10.0) + + def test_wsm_uses_transfer_spec_defaults(self): + sp = Cosmos3SamplingParams( + prompt="t", control_path="wsm.mp4", control_hint="wsm" + ) + sp._explicit_fields = {"prompt", "control_path", "control_hint"} + sp._apply_transfer_hint_defaults() + self.assertEqual(sp.guidance_scale, 1.0) + self.assertEqual(sp.control_guidance, 3.0) + self.assertEqual(sp.flow_shift, 10.0) + self.assertEqual(sp.num_frames, 101) + self.assertEqual(sp.fps, 10) + self.assertEqual(sp.num_video_frames_per_chunk, 101) + + lowered = Cosmos3SamplingParams.lower_video_request_kwargs( + types.SimpleNamespace(num_frames=None, fps=None), + { + "control_path": "wsm.mp4", + "control_hint": "wsm", + "num_frames": 121, + "fps": 24, + }, + ) + self.assertEqual(lowered["num_frames"], 101) + self.assertEqual(lowered["fps"], 10) + + def test_explicit_value_overrides_hint_default(self): + sp = Cosmos3SamplingParams( + prompt="t", + control_path="seg.mp4", + control_hint="seg", + control_guidance=4.2, + ) + sp._explicit_fields = {"control_path", "control_hint", "control_guidance"} + sp._apply_transfer_hint_defaults() + self.assertEqual(sp.control_guidance, 4.2) + + def test_multi_hint_defaults_not_applied(self): + sp = Cosmos3SamplingParams( + prompt="t", + control_path=["edge.mp4", "depth.mp4"], + control_hint=["edge", "depth"], + ) + base_guidance = sp.guidance_scale + sp._explicit_fields = {"control_path", "control_hint"} + sp._apply_transfer_hint_defaults() + self.assertEqual(sp.control_guidance, 1.0) + self.assertEqual(sp.guidance_scale, base_guidance) + + def test_unknown_hint_rejected(self): + with self.assertRaises(ValueError): + Cosmos3SamplingParams( + prompt="t", control_path="x.mp4", control_hint="bogus" + ) + + class TestCosmos3DomainAwareLinear(unittest.TestCase): """Per-domain action projection.""" diff --git a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py index 11ff7fa0c..b2022cf11 100644 --- a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py +++ b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py @@ -48,6 +48,23 @@ class _CustomEntrypointModel(_UniformDtypeModel): pass +class _MetaBufferModel(_UniformDtypeModel): + """Mirrors cosmos3: a non-checkpoint buffer stays on meta until + post_load_weights() rebuilds it.""" + + def __init__(self) -> None: + super().__init__() + self.proj = ReplicatedLinear(4, 4, bias=False) + self.register_buffer("inv_freq", torch.empty(2), persistent=False) + + def post_load_weights(self) -> None: + if self.inv_freq.is_meta: + device = next(self.parameters()).device + self.register_buffer( + "inv_freq", torch.ones(2, device=device), persistent=False + ) + + class TestFSDPMixedPrecisionPolicy(unittest.TestCase): def test_quant_config_detection_uses_the_runtime_instance(self): self.assertTrue(fsdp_load._is_bitsandbytes_quant_config(BitsAndBytesConfig())) @@ -218,6 +235,35 @@ class TestOrdinaryWeightLoading(unittest.TestCase): torch.testing.assert_close(model.proj.weight, checkpoint_weight) +class TestDevicePostprocessMove(unittest.TestCase): + def test_postprocess_move_preserves_meta_buffers_for_post_load_weights(self): + # The pre-postprocess device move must not copy buffers that are + # still on meta awaiting post_load_weights() (cosmos3's RoPE inv_freq). + load_plan = WeightLoadPlan( + checkpoint_load_device=torch.device("cpu"), + weight_postprocess_device=torch.device("cpu"), + ) + checkpoint_weight = torch.arange(16, dtype=torch.float32).reshape(4, 4) + + with patch.object(fsdp_load.current_platform, "is_mps", return_value=False): + model = fsdp_load.maybe_load_fsdp_model( + model_cls=_MetaBufferModel, + init_params={}, + weight_dir_list=[], + device=torch.device("cpu"), + hsdp_replicate_dim=1, + hsdp_shard_dim=1, + param_dtype=torch.float32, + reduce_dtype=torch.float32, + weight_load_plan=load_plan, + weights_iterator=iter((("proj.weight", checkpoint_weight),)), + ) + + self.assertFalse(model.inv_freq.is_meta) + torch.testing.assert_close(model.inv_freq, torch.ones(2)) + torch.testing.assert_close(model.proj.weight, checkpoint_weight) + + class TestRankLocalSafetensorsRead(unittest.TestCase): def _source( self, diff --git a/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py new file mode 100644 index 000000000..3f42243fa --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_modelopt_fp8_layerwise_offload_load.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Serialized ModelOpt FP8 checkpoints must postprocess on device even under +layerwise offload: requantize_with_max_scale() runs scaled_fp8_quant(), a +CUDA-only kernel, so a CPU-resident postprocess must never come back.""" + +import unittest + +import torch +from torch import nn + +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized, +) +from sglang.multimodal_gen.runtime.layers.linear import MergedColumnParallelLinear +from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( + ModelOptFp8Config, +) +from sglang.multimodal_gen.runtime.loader import fsdp_load +from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan +from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import ( + ensure_distributed_env_defaults, +) + +_IN_FEATURES = 32 +_SHARD_OUT = 16 + + +class _FusedFp8Model(nn.Module): + """Minimal stand-in for a serialized ModelOpt FP8 DiT: one fused linear + with per-shard scales, plus a cosmos3-style non-checkpoint meta buffer.""" + + param_names_mapping = {} + _fsdp_forward_methods: tuple[str, ...] = () + + def __init__(self, quant_config: ModelOptFp8Config) -> None: + super().__init__() + self.qkv = MergedColumnParallelLinear( + input_size=_IN_FEATURES, + output_sizes=[_SHARD_OUT, _SHARD_OUT], + bias=False, + quant_config=quant_config, + prefix="qkv", + ) + self.register_buffer("inv_freq", torch.empty(4), persistent=False) + + def post_load_weights(self) -> None: + if self.inv_freq.is_meta: + device = next(self.parameters()).device + self.register_buffer( + "inv_freq", + torch.ones(4, dtype=torch.float32, device=device), + persistent=False, + ) + + +def _make_serialized_fp8_checkpoint() -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """Quantize a reference weight shard-by-shard, as ModelOpt exports do.""" + torch.manual_seed(0) + weight_ref = torch.randn(2 * _SHARD_OUT, _IN_FEATURES, dtype=torch.float32) * 0.05 + + fp8_max = torch.finfo(torch.float8_e4m3fn).max + shard_scales = [] + fp8_shards = [] + for shard in weight_ref.split(_SHARD_OUT, dim=0): + scale = shard.abs().max() / fp8_max + shard_scales.append(scale) + fp8_shards.append((shard / scale).to(torch.float8_e4m3fn)) + + state_dict = { + "qkv.weight": torch.cat(fp8_shards, dim=0), + "qkv.weight_scale": torch.stack(shard_scales), + "qkv.input_scale": torch.tensor([0.5, 0.5], dtype=torch.float32), + } + return state_dict, weight_ref + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA for scaled_fp8_quant") +class TestModelOptFp8LayerwiseOffloadLoad(unittest.TestCase): + def test_serialized_checkpoint_loads_with_component_starting_on_cpu(self): + if not model_parallel_is_initialized(): + ensure_distributed_env_defaults() + maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1) + + state_dict, weight_ref = _make_serialized_fp8_checkpoint() + expected_max_scale = state_dict["qkv.weight_scale"].max() + + # The plan a layerwise-offload component gets when + # _needs_device_weight_postprocess() returns True: load and postprocess + # on GPU, then defer the CPU placement. + load_plan = WeightLoadPlan.for_component( + checkpoint_load_device=torch.device("cuda"), + needs_device_weight_postprocess=True, + component_starts_on_cpu=True, + ) + self.assertTrue(load_plan.defer_cpu_placement) + + model = fsdp_load.maybe_load_fsdp_model( + model_cls=_FusedFp8Model, + init_params={ + "quant_config": ModelOptFp8Config(is_checkpoint_fp8_serialized=True) + }, + weight_dir_list=[], + device=torch.device("cuda"), + hsdp_replicate_dim=1, + hsdp_shard_dim=1, + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + component_starts_on_cpu=True, + weight_load_plan=load_plan, + weights_iterator=iter(state_dict.items()), + ) + + # Postprocess ran: the weight was requantized to the shared max scale + # and rebound transposed. + weight = model.qkv.weight + self.assertEqual(weight.dtype, torch.float8_e4m3fn) + self.assertEqual(tuple(weight.shape), (_IN_FEATURES, 2 * _SHARD_OUT)) + weight_scale = model.qkv.weight_scale + torch.testing.assert_close( + weight_scale.flatten(), + expected_max_scale.expand(weight_scale.numel()), + check_device=False, + ) + torch.testing.assert_close( + model.qkv.input_scale.flatten().max(), torch.tensor(0.5), check_device=False + ) + + # The round trip through both quantizations stays close to the source. + # Loose on purpose: this guards against garbage (wrong scale, wrong + # shard order), not fp8 precision. + dequant = weight.t().float().cpu() * expected_max_scale + torch.testing.assert_close( + dequant, weight_ref, rtol=0.5, atol=float(expected_max_scale) * 8 + ) + + # Layerwise offload contract: the component lands on CPU afterwards, + # with the non-checkpoint buffer rebuilt. + self.assertEqual(weight.device.type, "cpu") + self.assertFalse(model.inv_freq.is_meta) + self.assertEqual(model.inv_freq.device.type, "cpu") + + +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 4a847ec7a..9212a64eb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -895,12 +895,19 @@ class TestTransformerQuantHelpers(unittest.TestCase): warning.assert_called_once() - def test_modelopt_fp8_serialized_checkpoint_needs_device_postprocess(self): + def test_modelopt_fp8_always_needs_device_weight_postprocess(self): + # Even a serialized checkpoint requantizes fused shards through + # scaled_fp8_quant(), which cannot process CPU tensors. self.assertTrue( _needs_device_weight_postprocess( ModelOptFp8Config(is_checkpoint_fp8_serialized=True) ) ) + self.assertTrue( + _needs_device_weight_postprocess( + ModelOptFp8Config(is_checkpoint_fp8_serialized=False) + ) + ) def test_online_fp8_needs_device_weight_postprocess(self): self.assertTrue(_needs_device_weight_postprocess(Fp8Config()))