From 62562cb0f18751d1a399a1436a311ec767c6f5ce Mon Sep 17 00:00:00 2001 From: Cerdore Date: Tue, 9 Jun 2026 00:00:08 +0800 Subject: [PATCH 1/5] feat(diffusion): add OmniDreams 2B DiT autoregressive video world model Add OmniDreams pipeline to multimodal_gen, supporting: - Flat checkpoint loading (570-key .pt, non-diffusers layout) with post-load parameter fusion matching FlashDreams conventions - Autoregressive rollout with BlockKVCache (sink + rolling window) and per-chunk KV-cache lifecycle - 3D NeoX RoPE (44:42:42) with shift_t for chunk positioning - 2-step flow-match scheduler (sigmas {1.0, 0.8036, 0.0}) - Text conditioning via full_concat 100352 embedding (Cosmos-Reason1-7B) with explicit attention mask - Per-chunk VAE latent concatenation with single-pass Wan 2.1 decode for temporal continuity - HDMap/trajectory conditioning pipeline (per-chunk VAE-encode) - Tensor parallelism via ColumnParallelLinear/RowParallelLinear - Precomputed cross-attention K/V caching - HTTP API fields for hdmap_path and num_views - 43/43 CPU unit tests covering component construction, RoPE, BlockKVCache, scheduler, DiT forward, denoising stage, registry, and regression guards - GPU CI coverage with consistency thresholds and perf baselines Relates to #27214 --- .../configs/models/dits/__init__.py | 2 + .../configs/models/dits/omnidreams.py | 62 ++ .../configs/models/vaes/__init__.py | 6 +- .../configs/models/vaes/wanvae.py | 27 + .../configs/pipeline_configs/__init__.py | 4 + .../configs/pipeline_configs/omnidreams.py | 64 ++ .../multimodal_gen/configs/sample/__init__.py | 2 + .../configs/sample/omnidreams.py | 38 + .../configs/sample/sampling_params.py | 11 + python/sglang/multimodal_gen/registry.py | 47 +- .../runtime/entrypoints/openai/protocol.py | 3 + .../runtime/entrypoints/openai/video_api.py | 34 + .../runtime/models/dits/omnidreams.py | 927 ++++++++++++++++++ .../runtime/models/dits/omnidreams_kvcache.py | 296 ++++++ .../runtime/models/dits/omnidreams_rope.py | 142 +++ .../models/encoders/omnidreams_text.py | 74 ++ .../scheduling_omnidreams_flow_match.py | 148 +++ .../runtime/pipelines/omnidreams_pipeline.py | 358 +++++++ .../model_specific_stages/omnidreams.py | 702 +++++++++++++ .../test/server/consistency_threshold.json | 6 + .../multimodal_gen/test/server/gpu_cases.py | 15 + .../test/server/perf_baselines.json | 8 + .../test/server/testcase_configs.py | 13 + .../sglang/multimodal_gen/test/test_utils.py | 3 + .../test/unit/data/omnidreams_dit_keys.txt | 570 +++++++++++ .../test/unit/test_omnidreams_components.py | 654 ++++++++++++ .../test/unit/test_omnidreams_regression.py | 261 +++++ .../test/unit/test_omnidreams_scaffold.py | 95 ++ python/sglang/utils.py | 2 + 29 files changed, 4571 insertions(+), 3 deletions(-) create mode 100644 python/sglang/multimodal_gen/configs/models/dits/omnidreams.py create mode 100644 python/sglang/multimodal_gen/configs/pipeline_configs/omnidreams.py create mode 100644 python/sglang/multimodal_gen/configs/sample/omnidreams.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/omnidreams.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/omnidreams_kvcache.py create mode 100644 python/sglang/multimodal_gen/runtime/models/dits/omnidreams_rope.py create mode 100644 python/sglang/multimodal_gen/runtime/models/encoders/omnidreams_text.py create mode 100644 python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_omnidreams_flow_match.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines/omnidreams_pipeline.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py create mode 100644 python/sglang/multimodal_gen/test/unit/data/omnidreams_dit_keys.txt create mode 100644 python/sglang/multimodal_gen/test/unit/test_omnidreams_components.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_omnidreams_regression.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_omnidreams_scaffold.py diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py index 39ab3fd6df18..dff3ed6cc9b0 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/dits/__init__.py @@ -10,6 +10,7 @@ ) from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig +from sglang.multimodal_gen.configs.models.dits.omnidreams import OmniDreamsDiTConfig from sglang.multimodal_gen.configs.models.dits.stablediffusion3 import ( StableDiffusion3TransformerConfig, ) @@ -25,5 +26,6 @@ "Hunyuan3DDiTConfig", "MOVAAudioConfig", "MOVAVideoConfig", + "OmniDreamsDiTConfig", "StableDiffusion3TransformerConfig", ] diff --git a/python/sglang/multimodal_gen/configs/models/dits/omnidreams.py b/python/sglang/multimodal_gen/configs/models/dits/omnidreams.py new file mode 100644 index 000000000000..9046253411a6 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/omnidreams.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +"""DiT config for NVIDIA OmniDreams (Cosmos-Predict2.5-2B based autoregressive +video world model; production runtime = FlashDreams). + +Architecture facts mirror FlashDreams ``CosmosDiTNetworkConfig`` for the +``2b_res720p_30fps_i2v_hdmap_distilled`` checkpoint (HDMap single-view variant): +``additional_concat_ch=16`` enables HDMap conditioning, cross-view attention is +off. The flat checkpoint key names match the submodule tree one-to-one, so +``param_names_mapping`` is the identity (empty dict). +""" + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +@dataclass +class OmniDreamsDiTArchConfig(DiTArchConfig): + # --- Cosmos DiT architecture (FlashDreams CosmosDiTNetworkConfig) --- + in_channels: int = 16 + out_channels: int = 16 + patch_spatial: int = 2 + patch_temporal: int = 1 + model_channels: int = 2048 + num_blocks: int = 28 + num_heads: int = 16 + mlp_ratio: float = 4.0 + concat_padding_mask: bool = True + use_adaln_lora: bool = True + adaln_lora_dim: int = 256 + use_crossattn_projection: bool = True + crossattn_proj_in_channels: int = 100352 + crossattn_emb_channels: int = 1024 + timestep_scale: float = 0.001 + # HDMap variant: 16 extra latent channels routed through additional_patch_embedding. + # Overrides the FlashDreams CosmosDiTNetworkConfig default of 0 (HDMap disabled). + additional_concat_ch: int = 16 + # Cross-view attention is disabled for the single-view checkpoint. + enable_cross_view_attn: bool = False + view_condition_dim: int = 16 + n_cameras_emb: int = 7 + + # Checkpoint keys equal submodule names -> identity mappings. + param_names_mapping: dict = field(default_factory=dict) + reverse_param_names_mapping: dict = field(default_factory=dict) + + def __post_init__(self) -> None: + super().__post_init__() + # BaseDiT-required instance attrs (also surfaced via ModelConfig.__getattr__). + self.hidden_size = self.model_channels + self.num_attention_heads = self.num_heads + self.num_channels_latents = self.out_channels + + @property + def head_dim(self) -> int: + return self.model_channels // self.num_heads + + +@dataclass +class OmniDreamsDiTConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=OmniDreamsDiTArchConfig) + prefix: str = "OmniDreams" diff --git a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py index 3438b1b8937f..95553bbaf195 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/__init__.py @@ -6,7 +6,10 @@ from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import ( StableDiffusion3VAEConfig, ) -from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig +from sglang.multimodal_gen.configs.models.vaes.wanvae import ( + OmniDreamsVAEConfig, + WanVAEConfig, +) __all__ = [ "DacVAEConfig", @@ -14,4 +17,5 @@ "StableDiffusion3VAEConfig", "WanVAEConfig", "Hunyuan3DVAEConfig", + "OmniDreamsVAEConfig", ] diff --git a/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py b/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py index caf59c407663..f322b07d92d3 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/wanvae.py @@ -94,3 +94,30 @@ def get_vae_scale_factor(self): # Wan VAE does not expose block_out_channels like SD-style VAEs. # Its spatial downsample factor is explicitly defined by scale_factor_spatial. return self.arch_config.scale_factor_spatial + + +# ---- OmniDreams (Cosmos-Predict2.5-based) VAE config ---- + + +@dataclass +class OmniDreamsVAEArchConfig(WanVAEArchConfig): + """VAE arch config for OmniDreams (Cosmos-Predict2.5 latent space). + + Inherits the Wan 2.1 VAE architecture (same encoder/decoder/z_dim=16), but + the Cosmos-Predict2.5 latent distribution may differ from Wan's training + distribution. Override ``latents_mean`` / ``latents_std`` if GPU validation + shows a mismatch; otherwise the Wan defaults are a safe fallback (the + encode/decode scaling is self-consistent). + + TODO(gpu): numerically validate latent mean/std against a FlashDreams dump; + if they diverge from Wan 2.1, replace these tuples with the + OmniDreams-specific values. Current values = Wan 2.1 (same-behavior + fallback). + """ + + +@dataclass +class OmniDreamsVAEConfig(WanVAEConfig): + arch_config: OmniDreamsVAEArchConfig = field( + default_factory=OmniDreamsVAEArchConfig + ) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index 60de1090e5f7..f3bcdb2474d5 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -36,6 +36,9 @@ ) from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig +from sglang.multimodal_gen.configs.pipeline_configs.omnidreams import ( + OmniDreamsPipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import ( StableDiffusion3PipelineConfig, @@ -67,6 +70,7 @@ "SanaPipelineConfig", "SlidingTileAttnConfig", "MOVAPipelineConfig", + "OmniDreamsPipelineConfig", "StableDiffusion3PipelineConfig", "WanT2V480PConfig", "WanI2V480PConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/omnidreams.py b/python/sglang/multimodal_gen/configs/pipeline_configs/omnidreams.py new file mode 100644 index 000000000000..2804b6c58949 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/omnidreams.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pipeline config for NVIDIA OmniDreams. + +Phase 0 wires the static structure (DiT config, VAE reuse, task type) and the +2-step flow-match sigma schedule. The denoising/decoding callbacks used at GPU +time are added in later phases. +""" + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.omnidreams import OmniDreamsDiTConfig +from sglang.multimodal_gen.configs.models.vaes.wanvae import OmniDreamsVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ModelTaskType, + PipelineConfig, +) + + +def warp_flow_match_sigmas( + denoising_timesteps: tuple[int, ...] = (1000, 450), + flow_shift: float = 5.0, + sigma_min: float = 0.0, +) -> list[float]: + """OmniDreams 2-step flow-match sigma schedule. + + Each raw timestep ``t`` maps to ``s = t / 1000`` then is warped by + ``shift*s / (1 + (shift-1)*s)``; ``sigma_min`` is appended as the final + target. With the distilled defaults this yields ``[1.0, 0.8036, 0.0]``. + """ + sigmas = [ + flow_shift * (t / 1000.0) / (1.0 + (flow_shift - 1.0) * (t / 1000.0)) + for t in denoising_timesteps + ] + sigmas.append(sigma_min) + return sigmas + + +@dataclass +class OmniDreamsPipelineConfig(PipelineConfig): + task_type: ModelTaskType = ModelTaskType.I2V + # CFG disabled for the distilled checkpoint. + should_use_guidance: bool = False + # Native bf16 DiT; VAE in fp32 for numerical stability. + dit_precision: str = "bf16" + vae_precision: str = "fp32" + # Flow-match warp shift (also drives warp_flow_match_sigmas). + flow_shift: float | None = 5.0 + + dit_config: OmniDreamsDiTConfig = field(default_factory=OmniDreamsDiTConfig) + # A.5: OmniDreams uses a Cosmos-Predict2.5-based latent space; the + # latents_mean/std defaults match Wan 2.1 (safe fallback). Override in + # OmniDreamsVAEArchConfig once GPU validation confirms the correct values. + vae_config: OmniDreamsVAEConfig = field(default_factory=OmniDreamsVAEConfig) + + # 2-step distilled flow-match schedule. + denoising_timesteps: tuple[int, ...] = (1000, 450) + sigma_min: float = 0.0 + + def denoising_sigmas(self) -> list[float]: + return warp_flow_match_sigmas( + self.denoising_timesteps, + self.flow_shift if self.flow_shift is not None else 5.0, + self.sigma_min, + ) diff --git a/python/sglang/multimodal_gen/configs/sample/__init__.py b/python/sglang/multimodal_gen/configs/sample/__init__.py index 047622e3308f..684426dfde9f 100644 --- a/python/sglang/multimodal_gen/configs/sample/__init__.py +++ b/python/sglang/multimodal_gen/configs/sample/__init__.py @@ -4,10 +4,12 @@ DiffusersGenericSamplingParams, ) from sglang.multimodal_gen.configs.sample.ideogram import Ideogram4SamplingParams +from sglang.multimodal_gen.configs.sample.omnidreams import OmniDreamsSamplingParams from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams __all__ = [ "SamplingParams", "DiffusersGenericSamplingParams", "Ideogram4SamplingParams", + "OmniDreamsSamplingParams", ] diff --git a/python/sglang/multimodal_gen/configs/sample/omnidreams.py b/python/sglang/multimodal_gen/configs/sample/omnidreams.py new file mode 100644 index 000000000000..87dcae64cd22 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/omnidreams.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Sampling params for NVIDIA OmniDreams (autoregressive video world model).""" + +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.sample.sampling_params import ( + DataType, + SamplingParams, +) + + +@dataclass +class OmniDreamsSamplingParams(SamplingParams): + data_type: DataType = DataType.VIDEO + + # 720p single-view i2v defaults (latent grid is /8 spatial, /4 temporal). + height: int = 704 + width: int = 1280 + # 2-step distilled flow-match schedule; CFG disabled. + num_inference_steps: int = 2 + guidance_scale: float = 1.0 + + supported_resolutions: list[tuple[int, int]] | None = field( + default_factory=lambda: [(1280, 704)] + ) + + # --- Autoregressive rollout knobs (see FlashDreams streaming inference) --- + # Number of latent frames produced per chunk. + len_t: int = 2 + # Rolling KV-cache window (in latent frames) and permanent sink size. + window_size_t: int = 6 + sink_size_t: int = 0 + # Raw timestep injected as context noise on cached/clean frames. + context_noise: int = 128 + + # HD-map / trajectory conditioning input (Phase 4). A single path (broadcast + # to every AR chunk) or a per-chunk list of paths. ``None`` disables HDMap. + hdmap_path: str | list[str] | None = None diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index f559fa5c59ae..65f31e281a97 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -897,6 +897,17 @@ def add_argument(*name_or_flags, **kwargs): '--image-path "img1.png" "img2.png"' ), ) + add_argument( + "--hdmap-path", + type=str, + nargs="+", + help=( + "Path(s) to HD-map / trajectory conditioning input(s) for " + "OmniDreams autoregressive video generation. Pass a single path " + "(broadcast to every chunk) or one per autoregressive chunk, e.g.: " + '--hdmap-path "chunk0.png" "chunk1.png"' + ), + ) add_argument( "--moba-config-path", type=str, diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 0871eaf5a83c..d5e5a140c3ff 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -69,6 +69,9 @@ MOVA360PConfig, MOVA720PConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.omnidreams import ( + OmniDreamsPipelineConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImageEditPipelineConfig, QwenImageEditPlus_2511_PipelineConfig, @@ -124,6 +127,7 @@ MOVA_360P_SamplingParams, MOVA_720P_SamplingParams, ) +from sglang.multimodal_gen.configs.sample.omnidreams import OmniDreamsSamplingParams from sglang.multimodal_gen.configs.sample.qwenimage import ( QwenImage2512SamplingParams, QwenImageEditPlusSamplingParams, @@ -393,8 +397,38 @@ def _get_config_info( model_id = _MODEL_HF_PATH_TO_NAME[registered_model_hf_id] return _CONFIG_REGISTRY.get(model_id) - # 3. Use detectors - config = maybe_download_model_index(model_path) + # 3. Use detectors. + # 3a. Path-based detection for NON-diffusers local checkpoints only (those + # without a model_index.json, e.g. the flat OmniDreams .pt). This is + # gated on the absence of model_index.json so that a substring detector + # (e.g. "sana") cannot hijack a legitimate diffusers model whose path + # happens to contain that substring — diffusers models still go through + # 3b where the pipeline _class_name disambiguates. + is_local_non_diffusers = os.path.isdir(model_path) and not os.path.isfile( + os.path.join(model_path, "model_index.json") + ) + if is_local_non_diffusers: + path_matched = [ + model_id + for model_id, detector in _MODEL_NAME_DETECTORS + if detector(model_path.lower()) + ] + if path_matched: + if len(path_matched) > 1: + logger.warning( + "More than one model name matched by path, using the first" + ) + return _CONFIG_REGISTRY.get(path_matched[0]) + + # 3b. Fall back to diffusers model_index.json + pipeline-class-name detection. + try: + config = maybe_download_model_index(model_path) + except ValueError: + logger.debug( + "diffusers model_index.json resolution failed for '%s'; no match.", + model_path, + ) + return None pipeline_name = config.get("_class_name", "").lower() matched_model_names = [] @@ -633,6 +667,15 @@ def get_model_info( # Registration of model configs def _register_configs(): + # OmniDreams (NVIDIA autoregressive video world model, flat .pt DiT) + register_configs( + sampling_param_cls=OmniDreamsSamplingParams, + pipeline_config_cls=OmniDreamsPipelineConfig, + hf_model_paths=["nvidia/omni-dreams-models"], + model_detectors=[ + lambda path: "omnidreams" in path.lower() or "omni-dreams" in path.lower() + ], + ) # LTX-2 register_configs( sampling_param_cls=LTX2SamplingParams, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py index f3f651149e0c..70eaca4d5334 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py @@ -127,6 +127,9 @@ class VideoGenerationsRequest(BaseModel): diffusers_kwargs: Optional[Dict[str, Any]] = None # kwargs for diffusers backend # Performance profiling perf_dump_path: Optional[str] = None + # OmniDreams / HDMap conditioning (Phase 4) + hdmap_path: Optional[Union[str, List[str]]] = None + num_views: Optional[int] = Field(default=None, ge=1, le=64) class VideoListResponse(BaseModel): 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 4ee93723b136..366313c2e790 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -64,9 +64,40 @@ def _parse_form_extra_value(value: Any) -> Any: return value +def _validate_http_hdmap_path(hdmap_path) -> None: + """Reject raw local filesystem paths for HDMap over the HTTP API. + + ``hdmap_path`` is fed to ``load_image``, which opens local files directly — + a raw server path from an untrusted HTTP body is an arbitrary-file-read + vector. Over HTTP, only remote URLs (http/https) or data URLs are allowed. + CLI callers build sampling params directly (bypassing this) and may still + pass trusted local paths. + """ + if hdmap_path is None: + return + entries = hdmap_path if isinstance(hdmap_path, list) else [hdmap_path] + for entry in entries: + if not isinstance(entry, str): + continue + low = entry.strip().lower() + if not ( + low.startswith("http://") + or low.startswith("https://") + or low.startswith("data:") + ): + raise HTTPException( + status_code=400, + detail=( + "hdmap_path must be an http(s) or data URL over the HTTP API; " + "local filesystem paths are not permitted." + ), + ) + + def _build_video_sampling_params(request_id: str, request: VideoGenerationsRequest): """Resolve video-specific defaults (fps, seconds → num_frames) then delegate to the shared build_sampling_params.""" + _validate_http_hdmap_path(request.hdmap_path) seconds = request.seconds if request.seconds is not None else DEFAULT_VIDEO_SECONDS fps = request.fps if request.fps is not None else DEFAULT_FPS num_frames = request.num_frames if request.num_frames is not None else fps * seconds @@ -110,6 +141,9 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque output_quality=request.output_quality, perf_dump_path=request.perf_dump_path, diffusers_kwargs=request.diffusers_kwargs, + # OmniDreams HDMap / trajectory conditioning (Phase 4). + hdmap_path=request.hdmap_path, + num_views=request.num_views, ) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/omnidreams.py b/python/sglang/multimodal_gen/runtime/models/dits/omnidreams.py new file mode 100644 index 000000000000..10ad3657849d --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/omnidreams.py @@ -0,0 +1,927 @@ +# SPDX-License-Identifier: Apache-2.0 +"""OmniDreams DiT (Cosmos DiT) for SGLang. + +Submodule names are checkpoint-exact: they mirror FlashDreams +``omnidreams.transformer.impl.modules`` / ``network`` so the flat ``.pt`` loads +with an identity parameter mapping. + +The module is constructed with PRE-FUSION shapes to match the raw checkpoint: +- ``x_embedder`` keeps the always-zero inference padding-mask channel + (in_features = 72 for the HDMap variant), and +- ``final_layer.linear`` keeps the Cosmos ``(kt kh kw c)`` patch-shuffle order. + +Both fusions run once in :meth:`post_load_weights` after ``load_state_dict``, +matching FlashDreams ``update_parameters_after_loading_checkpoint``. + +The denoising forward pass (RoPE, SDPA, KV-cache) is implemented in a later +phase; Phase-0 scaffolding only needs checkpoint-exact construction + fusion. +""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor + +from sglang.multimodal_gen.configs.models.dits.base import DiTConfig +from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT +from sglang.multimodal_gen.runtime.models.dits.omnidreams_kvcache import BlockKVCache + +# RoPE primitives live in omnidreams_rope; re-exported here for callers/tests. +from sglang.multimodal_gen.runtime.models.dits.omnidreams_rope import ( # noqa: F401 + ROPE_IS_NEOX_STYLE, + RotaryPositionEmbedding3D, + apply_rope_freqs, + rope_dims, +) + + +# ---- TP / distributed helpers (safe to call before distributed init) ------- # +def _use_tp() -> bool: + """True if tensor parallelism is initialised and > 1 rank.""" + try: + from sglang.multimodal_gen.runtime.distributed import get_tp_world_size + + return get_tp_world_size() > 1 + except (ImportError, AssertionError, RuntimeError): + return False + + +def _tp_size() -> int: + try: + from sglang.multimodal_gen.runtime.distributed import get_tp_world_size + + return max(1, get_tp_world_size()) + except (ImportError, AssertionError, RuntimeError): + return 1 + + +def _tp_col_linear( + in_f: int, out_f: int, bias: bool = True, gather_output: bool = False +): + """Column-parallel linear when TP is active, else a plain ``nn.Linear``.""" + if _use_tp(): + from sglang.multimodal_gen.runtime.layers.linear import ColumnParallelLinear + + return ColumnParallelLinear(in_f, out_f, bias=bias, gather_output=gather_output) + return nn.Linear(in_f, out_f, bias=bias) + + +def _tp_row_linear(in_f: int, out_f: int, bias: bool = True): + """Row-parallel linear when TP is active, else a plain ``nn.Linear``.""" + if _use_tp(): + from sglang.multimodal_gen.runtime.layers.linear import RowParallelLinear + + return RowParallelLinear( + in_f, out_f, bias=bias, reduce_results=True, input_is_parallel=True + ) + return nn.Linear(in_f, out_f, bias=bias) + + +def _divide(a: int, b: int) -> int: + """Integer division that asserts divisibility.""" + q, r = divmod(a, b) + if r != 0: + raise ValueError(f"{a} is not divisible by {b}") + return q + + +def _local_heads(num_heads: int) -> int: + """Number of attention heads visible to this TP rank.""" + return _divide(num_heads, _tp_size()) + + +def _sp_size() -> int: + try: + from sglang.multimodal_gen.runtime.distributed import get_sp_world_size + + return get_sp_world_size() + except (ImportError, AssertionError, RuntimeError): + return 1 + + +# --------------------------------------------------------------------------- # +# Building blocks (checkpoint-exact module names) # +# --------------------------------------------------------------------------- # +class GPT2FeedForward(nn.Module): + """GPT-2 style FFN with GELU (submodules: ``layer1``/``layer2``).""" + + def __init__(self, d_model: int, d_ff: int) -> None: + super().__init__() + self.activation = nn.GELU() + self.layer1 = nn.Linear(d_model, d_ff, bias=False) + self.layer2 = nn.Linear(d_ff, d_model, bias=False) + + def forward(self, x: Tensor) -> Tensor: + return self.layer2(self.activation(self.layer1(x))) + + +class Timesteps(nn.Module): + """Sinusoidal timestep embedding (non-persistent buffer -> no ckpt key).""" + + SINUSOIDAL_FREQ_BASE = 10000 + emb: Tensor + + def __init__(self, num_channels: int) -> None: + super().__init__() + self.num_channels = num_channels + self.reset_emb() + + def reset_emb(self, device: torch.device | str | None = None) -> None: + """(Re)create the non-persistent ``emb`` frequency table. + + Carries no checkpoint key, so a meta-device construction (the production + load path) leaves it empty; ``post_load_weights`` calls this to + rematerialize it on the loaded param device. Kept in float32 regardless + of model dtype, matching the original sinusoidal precision. + """ + half_dim = self.num_channels // 2 + exponent = -math.log(self.SINUSOIDAL_FREQ_BASE) * torch.arange( + half_dim, dtype=torch.float32, device=device + ) + exponent = exponent / half_dim + self.register_buffer("emb", torch.exp(exponent), persistent=False) + + def forward(self, timesteps: Tensor) -> Tensor: + emb = timesteps.unsqueeze(-1) * self.emb + return torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1) + + +class TimestepEmbedding(nn.Module): + """Timestep MLP with optional AdaLN-LoRA (submodules ``linear_1``/``linear_2``). + + When ``use_adaln_lora`` is True, ``linear_1`` has no bias and forward returns + ``(raw_sinusoidal_input, lora_out)``. + """ + + def __init__( + self, in_features: int, out_features: int, use_adaln_lora: bool = True + ) -> None: + super().__init__() + self.use_adaln_lora = use_adaln_lora + self.linear_1 = nn.Linear(in_features, out_features, bias=not use_adaln_lora) + self.activation = nn.SiLU() + out_dim = 3 * out_features if use_adaln_lora else out_features + self.linear_2 = nn.Linear(out_features, out_dim, bias=False) + + def forward(self, x: Tensor) -> tuple[Tensor, Tensor | None]: + # Timesteps emits float32 sinusoids for precision; cast to the MLP + # param dtype before the linears (and before the raw embedding is + # returned for RMSNorm / AdaLN), matching the model's running dtype. + x = x.to(self.linear_1.weight.dtype) + out = self.linear_2(self.activation(self.linear_1(x))) + if self.use_adaln_lora: + return x, out + return out, None + + +class PatchEmbed(nn.Module): + """Patch embed: ``proj = Sequential(Identity, Linear(in_features, out))``. + + The leading ``Identity`` is a placeholder kept for checkpoint key + compatibility (the learnable linear lives at ``proj.1``). + """ + + def __init__( + self, + spatial_patch_size: int, + temporal_patch_size: int, + in_channels: int, + out_channels: int, + ) -> None: + super().__init__() + self.spatial_patch_size = spatial_patch_size + self.temporal_patch_size = temporal_patch_size + self.in_channels = in_channels + self.proj = nn.Sequential( + nn.Identity(), + nn.Linear(self._compute_in_features(), out_channels, bias=False), + ) + + def _compute_in_features(self) -> int: + return self.in_channels * self.temporal_patch_size * self.spatial_patch_size**2 + + def forward(self, x: Tensor) -> Tensor: + return self.proj(x) + + +class FinalLayer(nn.Module): + """Final AdaLN layer (submodules ``layer_norm``/``linear``/``adaln_modulation``).""" + + NUM_ADALN_CHUNKS = 2 + + def __init__( + self, + hidden_size: int, + spatial_patch_size: int, + temporal_patch_size: int, + out_channels: int, + use_adaln_lora: bool = False, + adaln_lora_dim: int = 256, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.use_adaln_lora = use_adaln_lora + self.layer_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + patch_dim = spatial_patch_size**2 * temporal_patch_size * out_channels + self.linear = nn.Linear(hidden_size, patch_dim, bias=False) + modulation_out_dim = self.NUM_ADALN_CHUNKS * hidden_size + if use_adaln_lora: + self.adaln_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, adaln_lora_dim, bias=False), + nn.Linear(adaln_lora_dim, modulation_out_dim, bias=False), + ) + else: + self.adaln_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, modulation_out_dim, bias=False), + ) + + def forward( + self, x: Tensor, emb: Tensor, adaln_lora: Tensor | None = None + ) -> Tensor: + """Final AdaLN + projection. ``x``: [B, L, D], ``emb``: [B, D].""" + B, L, D = x.shape + emb_ = emb.reshape(B, 1, D) + if self.use_adaln_lora: + assert adaln_lora is not None + al = adaln_lora.reshape(B, 1, 3 * D) + modulation = self.adaln_modulation(emb_) + al[..., : 2 * self.hidden_size] + shift, scale = modulation.chunk(2, dim=-1) + else: + shift, scale = self.adaln_modulation(emb_).chunk(2, dim=-1) + x = self.layer_norm(x) * (1.0 + scale) + shift + return self.linear(x) + + +class OmniDreamsAttention(nn.Module): + """Multi-head attention block (self or cross), with optional TP support. + + Submodules: ``q_proj``/``k_proj``/``v_proj``/``output_proj`` (all bias-free) + and per-head ``q_norm``/``k_norm`` RMSNorms. The attention op itself carries + no parameters, so it is intentionally not a registered submodule (keeps the + checkpoint key set identical to FlashDreams). + + When TP is active, ``q_proj``/``k_proj``/``v_proj`` use column-parallel + projection (each rank sees ``local_num_heads`` heads) and ``output_proj`` + uses row-parallel projection (all-reduces across ranks). The per-head + RMSNorms already operate on individual heads so they are naturally TP-local. + """ + + def __init__( + self, + query_dim: int, + context_dim: int | None, + n_heads: int, + head_dim: int, + ) -> None: + super().__init__() + context_dim = query_dim if context_dim is None else context_dim + self.n_heads = n_heads + self.head_dim = head_dim + self.local_num_heads = _local_heads(n_heads) + self._is_tp = _use_tp() + + inner_dim_full = head_dim * n_heads + self.q_proj = _tp_col_linear( + query_dim, inner_dim_full, bias=False, gather_output=False + ) + self.k_proj = _tp_col_linear( + context_dim, inner_dim_full, bias=False, gather_output=False + ) + self.v_proj = _tp_col_linear( + context_dim, inner_dim_full, bias=False, gather_output=False + ) + self.output_proj = _tp_row_linear(inner_dim_full, query_dim, bias=False) + self.q_norm = nn.RMSNorm(head_dim, eps=1e-6) + self.k_norm = nn.RMSNorm(head_dim, eps=1e-6) + + def _project_qkv(self, x: Tensor, ctx: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Project Q/K/V, handling the TP ColumnParallelLinear return convention.""" + if self._is_tp: + # TP linears (ColumnParallelLinear) return ``(out, bias)`` tuples. + q_raw, _ = self.q_proj(x) + k_raw, _ = self.k_proj(ctx) + v_raw, _ = self.v_proj(ctx) + else: + q_raw = self.q_proj(x) + k_raw = self.k_proj(ctx) + v_raw = self.v_proj(ctx) + return q_raw, k_raw, v_raw + + def forward( + self, + x: Tensor, + rope_freqs: Tensor | None = None, + context: Tensor | None = None, + kv_cache=None, + cross_kv: tuple[Tensor, Tensor] | None = None, + ) -> Tensor: + """Attention with an optional KV-cache window (autoregressive self-attn) + and optional precomputed cross-attention K/V (Phase 6 caching). + + Q/K are per-head RMSNorm'd, then (self-attn only) rotated by RoPE before + a full bidirectional SDPA (scale 1/sqrt(head_dim)). Cross-attention passes + ``context`` as K/V source and no RoPE. + + When ``kv_cache`` (a :class:`BlockKVCache`) is given, this is an AR + self-attention step: the current chunk's post-RoPE K and V are written + into the cache (``seq_dim=1``) and Q attends over the full cached window + ``[sink | window]``. The cache's ``before_update``/``after_update`` (window + roll + bookkeeping) are driven by the denoising stage, not here -- here we + only ``update`` (write) and read ``cached_k``/``cached_v``. Cross-chunk + causality comes solely from the window (no causal mask); each chunk's K is + rotated by its absolute position via ``shift_t`` before being cached, so + cached K and current Q keep the correct relative rotation. + + When ``cross_kv`` (a ``(K,V)`` tuple) is given, the K/V projections and + K-norm are skipped for the cross-attention path and the precomputed + tensors are used directly. The cached K already has ``k_norm`` applied. + This avoids redundant projection of the same text context in every AR + forward (28 blocks × num_chunks × 3 calls/chunk). ``context`` is still + required for the Q projection from ``x``. + """ + ctx = x if context is None else context + B, L, _ = x.shape + n, d = self.local_num_heads, self.head_dim + + if cross_kv is not None: + # Precomputed cross-attn K/V — skip k_proj/v_proj/k_norm. + # Q still needs projection since it depends on the current x. + q_raw = self.q_proj(x)[0] if self._is_tp else self.q_proj(x) + q = self.q_norm(q_raw.reshape(B, L, n, d)) + k, v = cross_kv + else: + q_raw, k_raw, v_raw = self._project_qkv(x, ctx) + Lk = ctx.shape[1] + q = self.q_norm(q_raw.reshape(B, L, n, d)) + k = self.k_norm(k_raw.reshape(B, Lk, n, d)) + v = v_raw.reshape(B, Lk, n, d) + + if rope_freqs is not None: + q = apply_rope_freqs(q, rope_freqs) + k = apply_rope_freqs(k, rope_freqs) + if kv_cache is not None: + # Write this chunk's (post-RoPE) K/V, then attend over the window. + kv_cache.update(k, v) + k = kv_cache.cached_k() + v = kv_cache.cached_v() + # SDPA expects [B, n, S, d]; default scale = 1/sqrt(d), no mask. + out = F.scaled_dot_product_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + ) + out = out.transpose(1, 2).reshape(B, L, n * d) + if self._is_tp: + out, _ = self.output_proj(out) + return out + return self.output_proj(out) + + +class OmniDreamsBlock(nn.Module): + """Cosmos transformer block: self-attn -> (cross-view-attn) -> cross-attn -> MLP. + + Cross-view attention is gated behind ``enable_cross_view_attn`` (Phase 5, + default off for the single-view checkpoint). When enabled each camera + view's tokens attend over all views at the same temporal position via a + dense bidirectional attention (no RoPE, no causal mask), and per-view + AdaLN modulation terms are added to the timestep-conditioned shift/scale/ + gate biases. + """ + + def __init__( + self, + x_dim: int, + context_dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + use_adaln_lora: bool = False, + adaln_lora_dim: int = 256, + enable_cross_view_attn: bool = False, + ) -> None: + super().__init__() + self.x_dim = x_dim + self.use_adaln_lora = use_adaln_lora + self.enable_cross_view_attn = enable_cross_view_attn + head_dim = x_dim // num_heads + + self.layer_norm_self_attn = nn.LayerNorm( + x_dim, elementwise_affine=False, eps=1e-6 + ) + self.self_attn = OmniDreamsAttention(x_dim, None, num_heads, head_dim) + + # Cross-view attention (Phase 5) — learnable LayerNorm (unlike AdaLN). + if enable_cross_view_attn: + self.layer_norm_cross_view_attn = nn.LayerNorm( + x_dim, elementwise_affine=True, eps=1e-6 + ) + self.cross_view_attn = OmniDreamsAttention( + x_dim, x_dim, num_heads, head_dim + ) + + self.layer_norm_cross_attn = nn.LayerNorm( + x_dim, elementwise_affine=False, eps=1e-6 + ) + self.cross_attn = OmniDreamsAttention(x_dim, context_dim, num_heads, head_dim) + + self.layer_norm_mlp = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6) + self.mlp = GPT2FeedForward(x_dim, int(x_dim * mlp_ratio)) + + def _make_adaln_mod() -> nn.Sequential: + if use_adaln_lora: + return nn.Sequential( + nn.SiLU(), + nn.Linear(x_dim, adaln_lora_dim, bias=False), + nn.Linear(adaln_lora_dim, 3 * x_dim, bias=False), + ) + return nn.Sequential(nn.SiLU(), nn.Linear(x_dim, 3 * x_dim, bias=False)) + + self.adaln_modulation_self_attn = _make_adaln_mod() + self.adaln_modulation_cross_attn = _make_adaln_mod() + self.adaln_modulation_mlp = _make_adaln_mod() + + @staticmethod + def _expand_view_mod(view_tensor: Tensor, B: int, V: int, D: int) -> Tensor: + """Expand per-view modulation ``[B, V, D]`` into ``[B, V, 1, D]``.""" + return view_tensor.reshape(B, V, 1, D) + + def forward( + self, + x: Tensor, + emb: Tensor, + adaln_lora: Tensor | None, + rope_freqs: Tensor | None, + context: Tensor, + self_attn_kv_cache=None, + cross_attn_kv: tuple[Tensor, Tensor] | None = None, + view_embedding_proj: Tensor | None = None, + ) -> Tensor: + """One transformer block on a single chunk. + + Args: + x: ``[B, L, D]`` tokens (flat across views & frames). + emb: ``[B, D]`` timestep embedding. + adaln_lora: ``[B, 3D]`` AdaLN-LoRA term. + rope_freqs: ``[L, 1, 1, head_dim]`` RoPE freqs for self-attention. + context: ``[B, Lctx, D]`` cross-attention key/value source. + self_attn_kv_cache: optional :class:`BlockKVCache`. + cross_attn_kv: optional precomputed ``(K,V)`` tuple for cross-attn. + view_embedding_proj: optional ``[B, V, 9D]`` view modulation tensor + (Phase 5 cross-view attention). ``None`` for single-view. + """ + B, L, D = x.shape + emb_ = emb.reshape(B, 1, D) + + if self.use_adaln_lora: + assert adaln_lora is not None + al = adaln_lora.reshape(B, 1, 3 * D) + shift_s, scale_s, gate_s = ( + self.adaln_modulation_self_attn(emb_) + al + ).chunk(3, dim=-1) + shift_c, scale_c, gate_c = ( + self.adaln_modulation_cross_attn(emb_) + al + ).chunk(3, dim=-1) + shift_m, scale_m, gate_m = (self.adaln_modulation_mlp(emb_) + al).chunk( + 3, dim=-1 + ) + else: + shift_s, scale_s, gate_s = self.adaln_modulation_self_attn(emb_).chunk( + 3, dim=-1 + ) + shift_c, scale_c, gate_c = self.adaln_modulation_cross_attn(emb_).chunk( + 3, dim=-1 + ) + shift_m, scale_m, gate_m = self.adaln_modulation_mlp(emb_).chunk(3, dim=-1) + + # Cross-view modulation (Phase 5): additive per-view AdaLN bias. + if self.enable_cross_view_attn and view_embedding_proj is not None: + V = view_embedding_proj.shape[1] + ( + view_shift_s, + view_scale_s, + view_gate_s, + view_shift_c, + view_scale_c, + view_gate_c, + view_shift_m, + view_scale_m, + view_gate_m, + ) = view_embedding_proj.chunk(9, dim=-1) + shift_s = shift_s + self._expand_view_mod(view_shift_s, B, V, D) + scale_s = scale_s + self._expand_view_mod(view_scale_s, B, V, D) + gate_s = gate_s + self._expand_view_mod(view_gate_s, B, V, D) + shift_c = shift_c + self._expand_view_mod(view_shift_c, B, V, D) + scale_c = scale_c + self._expand_view_mod(view_scale_c, B, V, D) + gate_c = gate_c + self._expand_view_mod(view_gate_c, B, V, D) + shift_m = shift_m + self._expand_view_mod(view_shift_m, B, V, D) + scale_m = scale_m + self._expand_view_mod(view_scale_m, B, V, D) + gate_m = gate_m + self._expand_view_mod(view_gate_m, B, V, D) + + normed = self.layer_norm_self_attn(x) * (1 + scale_s) + shift_s + x = x + gate_s * self.self_attn( + normed, rope_freqs=rope_freqs, kv_cache=self_attn_kv_cache + ) + + # Cross-view attention (Phase 5): each view attends over all views at + # the same temporal position (dense bidirectional, no RoPE, no gate). + if self.enable_cross_view_attn and view_embedding_proj is not None: + x_cv = self._cross_view_attn_forward(x, L, B, D) + x = x + x_cv + + normed = self.layer_norm_cross_attn(x) * (1 + scale_c) + shift_c + x = x + gate_c * self.cross_attn( + normed, context=context, cross_kv=cross_attn_kv + ) + + normed = self.layer_norm_mlp(x) * (1 + scale_m) + shift_m + x = x + gate_m * self.mlp(normed) + return x + + def _cross_view_attn_forward(self, x: Tensor, L: int, B: int, D: int) -> Tensor: + """Cross-view attention (Phase 5) — not yet implemented. + + The intended behavior reshapes the flat ``[B, L, D]`` tokens into + ``[B, V, T*HW, D]`` and, for each temporal position ``t``, lets every + view's queries attend over the concatenated K/V of all views at that + same ``t`` (no RoPE, no causal mask). The view count ``V`` must be + threaded in from the caller, which is not wired up yet, so we fail + loudly instead of silently running global attention over all tokens. + """ + raise NotImplementedError( + "Cross-view attention (enable_cross_view_attn=True) is not yet " + "supported: temporal-position-restricted attention is unimplemented. " + "Run with the default enable_cross_view_attn=False." + ) + + +# --------------------------------------------------------------------------- # +# Top-level DiT # +# --------------------------------------------------------------------------- # +class OmniDreamsDiT(BaseDiT): + """OmniDreams Cosmos DiT (2.06B, DiT-only, autoregressive video world model). + + Supports TP (tensor parallelism via ``ColumnParallelLinear``/``RowParallelLinear`` + head sharding), SP (sequence parallelism — guarded: SP init is detected and + rejected with a clear error since the autoregressive chunk loop is not yet + SP-aware), and optional cross-view attention (Phase 5, gated by config + ``enable_cross_view_attn``). + """ + + _fsdp_shard_conditions = [lambda n, m: isinstance(m, OmniDreamsBlock)] + _compile_conditions = [lambda n, m: isinstance(m, OmniDreamsBlock)] + param_names_mapping: dict = {} + reverse_param_names_mapping: dict = {} + + def __init__( + self, config: DiTConfig, hf_config: dict | None = None, **kwargs + ) -> None: + super().__init__(config, hf_config or {}, **kwargs) + arch = config.arch_config + self.arch = arch + + # +1 for the per-frame condition mask, +1 for the (training) padding mask. + in_channels = arch.in_channels + 1 + if arch.concat_padding_mask: + in_channels += 1 + + self.x_embedder = PatchEmbed( + arch.patch_spatial, arch.patch_temporal, in_channels, arch.model_channels + ) + if arch.additional_concat_ch > 0: + self.additional_patch_embedding = PatchEmbed( + arch.patch_spatial, + arch.patch_temporal, + arch.additional_concat_ch, + arch.model_channels, + ) + + self.t_embedder = nn.Sequential( + Timesteps(arch.model_channels), + TimestepEmbedding( + arch.model_channels, + arch.model_channels, + use_adaln_lora=arch.use_adaln_lora, + ), + ) + self.t_embedding_norm = nn.RMSNorm(arch.model_channels, eps=1e-6) + + # Phase 5: cross-view attention (default off for single-view checkpoint). + _cv_enabled = getattr(arch, "enable_cross_view_attn", False) + self.blocks = nn.ModuleList( + [ + OmniDreamsBlock( + x_dim=arch.model_channels, + context_dim=arch.crossattn_emb_channels, + num_heads=arch.num_heads, + mlp_ratio=arch.mlp_ratio, + use_adaln_lora=arch.use_adaln_lora, + adaln_lora_dim=arch.adaln_lora_dim, + enable_cross_view_attn=_cv_enabled, + ) + for _ in range(arch.num_blocks) + ] + ) + + self.final_layer = FinalLayer( + hidden_size=arch.model_channels, + spatial_patch_size=arch.patch_spatial, + temporal_patch_size=arch.patch_temporal, + out_channels=arch.out_channels, + use_adaln_lora=arch.use_adaln_lora, + adaln_lora_dim=arch.adaln_lora_dim, + ) + + if arch.use_crossattn_projection: + self.crossattn_proj = nn.Sequential( + nn.Linear( + arch.crossattn_proj_in_channels, + arch.crossattn_emb_channels, + bias=True, + ), + nn.GELU(), + ) + + # Cross-view attention (Phase 5): embedder + projection network. + # Stays None for the single-view checkpoint (no params, no forward cost). + if _cv_enabled: + n_cameras = getattr(arch, "n_cameras_emb", 7) + self.adaln_view_embedder = nn.Embedding(n_cameras, arch.model_channels) + self.adaln_view_proj = nn.Linear( + arch.model_channels, arch.model_channels * 9 + ) + else: + self.adaln_view_embedder = None + self.adaln_view_proj = None + + self._is_shuffle_op_fused = False + self._is_padding_mask_fused = False + + # Phase 6 SP guard: store sp_size for forward-time detection. + self._sp_size = _sp_size() + + # BaseDiT-required instance attributes. + self.hidden_size = arch.model_channels + self.num_attention_heads = arch.num_heads + self.num_channels_latents = arch.out_channels + self.__post_init__() + + # ----- load-time weight fusions (mirror FlashDreams) ------------------- # + def _fuse_padding_mask_into_patch_embed(self) -> None: + """Drop the always-zero inference padding-mask channels (72 -> 68).""" + if not self.arch.concat_padding_mask or self._is_padding_mask_fused: + return + self.x_embedder.in_channels -= 1 + in_channels_to_keep = self.x_embedder._compute_in_features() + proj_linear = self.x_embedder.proj[1] + proj_linear.weight.data = proj_linear.weight.data[ + :, :in_channels_to_keep + ].contiguous() + if proj_linear.bias is not None: + proj_linear.bias.data = proj_linear.bias.data[ + :in_channels_to_keep + ].contiguous() + self._is_padding_mask_fused = True + + def _fuse_shuffle_op_into_last_layer(self) -> None: + """Fold the Cosmos ``(kt kh kw c) -> (c kt kh kw)`` shuffle into the last linear.""" + if self._is_shuffle_op_fused: + return + self.final_layer.linear.weight.data = rearrange( + self.final_layer.linear.weight, + "(kt kh kw c) in_dim -> (c kt kh kw) in_dim", + kt=self.arch.patch_temporal, + kh=self.arch.patch_spatial, + kw=self.arch.patch_spatial, + c=self.arch.out_channels, + ).contiguous() + if self.final_layer.linear.bias is not None: + self.final_layer.linear.bias.data = rearrange( + self.final_layer.linear.bias, + "(kt kh kw c) -> (c kt kh kw)", + kt=self.arch.patch_temporal, + kh=self.arch.patch_spatial, + kw=self.arch.patch_spatial, + c=self.arch.out_channels, + ).contiguous() + self._is_shuffle_op_fused = True + + def post_load_weights(self) -> None: + self._fuse_padding_mask_into_patch_embed() + self._fuse_shuffle_op_into_last_layer() + self._materialize_nonpersistent_buffers() + + def _materialize_nonpersistent_buffers(self) -> None: + """Rematerialize non-persistent buffers left empty by meta construction. + + Buffers registered with ``persistent=False`` (e.g. the sinusoidal + ``Timesteps.emb`` table) have no checkpoint key, so the state-dict + loader never fills them and they stay on the meta device after a + meta-init load. Recompute them on the loaded parameter device. + """ + device = self.x_embedder.proj[1].weight.device + for module in self.modules(): + if isinstance(module, Timesteps): + module.reset_emb(device=device) + + def init_kv_caches( + self, + batch_size: int, + chunk_tokens: int, + window_tokens: int, + sink_tokens: int = 0, + device: torch.device | str = "cpu", + dtype: torch.dtype = torch.bfloat16, + ) -> list[BlockKVCache]: + """Build one :class:`BlockKVCache` per transformer block for AR rollout. + + Token counts are in *patchified tokens*, i.e. ``frames * Hp * Wp``: + ``chunk_tokens`` = ``len_t`` latent frames, ``window_tokens`` = + ``window_size_t`` frames, ``sink_tokens`` = ``sink_size_t`` frames (all + already multiplied by the per-frame spatial token count). K/V are stored + with ``seq_dim=1`` and shape ``[B, sink+window, local_heads, head_dim]`` + (TP-sharded heads when tensor parallelism is active). + """ + n = _local_heads(self.arch.num_heads) + d = self.arch.model_channels // self.arch.num_heads + total = sink_tokens + window_tokens + shape = (batch_size, total, n, d) + return [ + BlockKVCache( + k_shape=shape, + v_shape=shape, + seq_dim=1, + chunk_size=chunk_tokens, + window_size=window_tokens, + sink_size=sink_tokens, + device=device, + dtype=dtype, + ) + for _ in range(self.arch.num_blocks) + ] + + @torch.no_grad() + def precompute_cross_attn_kv( + self, context: torch.Tensor + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Precompute cross-attention K/V for every block once per prompt. + + Returns a list indexed by block, each ``(K,V)`` with per-head + RMSNorm applied to K. When TP is active, each rank sees only its + local heads (``local_num_heads``). This is used by the AR denoising + stage to avoid redundant ``k_proj(ctx)``/``v_proj(ctx)`` in every + forward call (28 blocks × num_chunks × 3 calls/chunk of wasted matmul). + + Phase 6. Invoke once before the chunk loop; pass to + :meth:`forward` as ``cross_attn_kv=result``. + """ + result: list[tuple[torch.Tensor, torch.Tensor]] = [] + for block in self.blocks: + attn = block.cross_attn + n, d = attn.local_num_heads, attn.head_dim + k_raw, _ = ( + attn.k_proj(context) if _use_tp() else (attn.k_proj(context), None) + ) + v_raw, _ = ( + attn.v_proj(context) if _use_tp() else (attn.v_proj(context), None) + ) + if isinstance(k_raw, tuple): + k_raw = k_raw[0] + if isinstance(v_raw, tuple): + v_raw = v_raw[0] + k = attn.k_norm(k_raw.reshape(context.shape[0], context.shape[1], n, d)) + v = v_raw.reshape(context.shape[0], context.shape[1], n, d) + result.append((k, v)) + return result + + def patchify(self, video: Tensor) -> Tensor: + """[B, C, T, H, W] -> [B, T*Hp*Wp, C*kt*kh*kw] (channel-major packing).""" + return rearrange( + video, + "b c (t kt) (h kh) (w kw) -> b (t h w) (c kt kh kw)", + kt=self.arch.patch_temporal, + kh=self.arch.patch_spatial, + kw=self.arch.patch_spatial, + ) + + def unpatchify( + self, tokens: Tensor, grid_t: int, grid_h: int, grid_w: int + ) -> Tensor: + """[B, L, out*kt*kh*kw] -> [B, out, T, H, W]. + + Uses the simple ``(c kt kh kw)`` unpack because the Cosmos channel + shuffle is already folded into ``final_layer.linear`` by + :meth:`post_load_weights`. + """ + return rearrange( + tokens, + "b (t h w) (c kt kh kw) -> b c (t kt) (h kh) (w kw)", + t=grid_t, + h=grid_h, + w=grid_w, + kt=self.arch.patch_temporal, + kh=self.arch.patch_spatial, + kw=self.arch.patch_spatial, + c=self.arch.out_channels, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + guidance=None, + *, + condition_video_input_mask: torch.Tensor, + rope_freqs: torch.Tensor, + hdmap_condition: torch.Tensor | None = None, + kv_caches: list | None = None, + cross_attn_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None, + view_indices: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + """Denoising forward (single-chunk, or one autoregressive chunk). + + Args: + hidden_states: patchified latent tokens ``[B, L, in_channels*kt*kh*kw]``. + encoder_hidden_states: text context ``[B, Lctx, crossattn_proj_in_channels]`` + (projected to ``crossattn_emb_channels`` by ``crossattn_proj``). + timestep: scalar (warped) timestep; scaled by ``timestep_scale`` here. + condition_video_input_mask: patchified per-frame condition mask + ``[B, L, kt*kh*kw]`` (concatenated onto the latent channels). + rope_freqs: ``[L, 1, 1, head_dim]`` 3D-RoPE freqs for self-attention. + For AR rollout pass ``shift_t(ar_idx)`` so the chunk is rotated by + its absolute temporal position. + hdmap_condition: patchified HDMap ``[B, L, additional_concat_ch*kt*kh*kw]``. + kv_caches: optional per-block list of :class:`BlockKVCache` (length + ``num_blocks``) enabling the autoregressive self-attention window. + ``None`` runs the plain single-chunk path. + cross_attn_kv: optional per-block list of precomputed ``(K,V)`` tuples + (Phase 6 caching). When given, bypasses ``k_proj``/``v_proj``/ + ``k_norm`` in the cross-attention path. + view_indices: optional ``[B, V]`` long tensor of camera view indices + (Phase 5 cross-view attention). ``None`` for single-view mode. + + Returns: + Patchified flow prediction ``[B, L, out_channels*kt*kh*kw]``. + Call :meth:`unpatchify` to recover ``[B, out, T, H, W]``. + """ + assert self._is_padding_mask_fused and self._is_shuffle_op_fused, ( + "call post_load_weights() before forward (fuses padding-mask + " + "last-layer shuffle)" + ) + # Phase 6: the outer AR loop (dynamic chunk count + BlockKVCache ops) + # is not torch.compile-safe. The hot-path OmniDreamsBlocks ARE compiled + # via _compile_conditions, which is sufficient. + assert not torch.compiler.is_compiling(), ( + "OmniDreamsDiT.forward() is not torch.compile-safe (dynamic chunk " + "loop + KV cache operations break fullgraph). Individual blocks are " + "compiled via _compile_conditions." + ) + # Phase 6: SP (ulysses/ring sequence parallelism) is not yet supported + # for the autoregressive chunk loop. Guard with a clear error. + if self._sp_size > 1: + raise RuntimeError( + "Sequence parallelism (SP) is not yet supported for OmniDreams. " + "Run with --ulysses-degree 1 --ring-degree 1." + ) + + timestep = timestep * self.arch.timestep_scale + + x = torch.cat([hidden_states, condition_video_input_mask], dim=-1) + x = self.x_embedder(x) + if self.arch.additional_concat_ch > 0: + assert hdmap_condition is not None, "HDMap variant requires hdmap_condition" + x = x + self.additional_patch_embedding(hdmap_condition) + + t_emb, adaln_lora = self.t_embedder(timestep) + t_emb = self.t_embedding_norm(t_emb) + batch = x.shape[0] + t_emb = t_emb.reshape(1, -1).expand(batch, -1) + if adaln_lora is not None: + adaln_lora = adaln_lora.reshape(1, -1).expand(batch, -1) + + # Phase 5: compute cross-view modulation once per forward. + view_embedding_proj: Tensor | None = None + if view_indices is not None and self.adaln_view_proj is not None: + view_emb = self.adaln_view_embedder(view_indices) # [B, V, D] + view_embedding_proj = self.adaln_view_proj(view_emb) # [B, V, 9D] + + context = self.crossattn_proj(encoder_hidden_states) + for i, block in enumerate(self.blocks): + x = block( + x, + t_emb, + adaln_lora, + rope_freqs, + context, + self_attn_kv_cache=None if kv_caches is None else kv_caches[i], + cross_attn_kv=None if cross_attn_kv is None else cross_attn_kv[i], + view_embedding_proj=view_embedding_proj, + ) + return self.final_layer(x, t_emb, adaln_lora) + + +EntryClass = OmniDreamsDiT diff --git a/python/sglang/multimodal_gen/runtime/models/dits/omnidreams_kvcache.py b/python/sglang/multimodal_gen/runtime/models/dits/omnidreams_kvcache.py new file mode 100644 index 000000000000..bb5a8b8c813a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/omnidreams_kvcache.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Block KV cache for OmniDreams autoregressive rollout. + +Faithful port of the FlashDreams production ``BlockKVCache`` +(``flashdreams/core/attention/kvcache.py``): causal attention with a fixed-size +local window plus optional sink tokens. + +Layout along ``seq_dim``: ``[sink tokens | local window tokens]``. Sink tokens +are never evicted; the local window rolls left by ``chunk_size`` once full. +Cross-chunk causality in OmniDreams comes ONLY from this window (the attention +op itself is full bidirectional SDPA with no causal mask). + +Differences from the FlashDreams original: +- ``torch.sym_min`` / ``torch.sym_max`` are replaced with the plain builtins. + The bounds (``_n_cached``, ``total_size``, ``chunk_size``) are Python ints, so + the two are eager-equivalent; the ``sym_*`` variants only matter for symbolic + CUDA-graph tracing, and plain ``min``/``max`` keep this importable + testable + on CPU without a CUDA build. + +Per-step usage:: + + cache.before_update(chunk_idx) # roll window if steady-state + cache.update(k, v) # write this chunk's K/V + k_all, v_all = cache.cached_k(), cache.cached_v() + cache.after_update(chunk_idx) # bookkeeping +""" + +from dataclasses import dataclass, field + +import torch +from torch import Tensor +from typing_extensions import Self + + +@dataclass +class BlockKVCache: + """KV cache for causal attention with a fixed-size local window + sink. + + Keys/values may have arbitrary shape ``[..., total_size, ...]``; the rolling + (sequence) dimension is ``seq_dim`` (may be negative). ``total_size`` equals + ``sink_size + window_size`` and must be divisible by ``chunk_size``. Chunks + are non-overlapping: each update appends one ``chunk_size``-token chunk at the + next logical position. + + ``chunk_idx`` (0, 1, 2, ...) is the chunk's index in the full sequence, not a + cache offset. A ``chunk_idx`` of ``prev + 1`` appends (or, in steady-state, + writes after a left-roll); a ``chunk_idx`` equal to ``prev`` overwrites the + same positions (used to refresh K/V after a renoise/finalize pass). + """ + + k_shape: tuple[int, ...] + v_shape: tuple[int, ...] + seq_dim: int + chunk_size: int + window_size: int + sink_size: int = 0 + device: torch.device | str = torch.device("cpu") + dtype: torch.dtype = torch.float32 + + _prev_chunk_idx: int = -1 + _curr_chunk_idx: int | None = None + _n_cached: int = 0 + + _k: Tensor = field(init=False) + _v: Tensor = field(init=False) + + # ----- properties ------------------------------------------------------- # + @property + def size(self) -> int: + """Number of valid cached tokens visible to attention.""" + if self._curr_chunk_idx is None: + return self._n_cached + return self._visible_end() + + @property + def write_end(self) -> int: + """Right edge of the current chunk in the physical cache layout.""" + assert ( + self._curr_chunk_idx is not None + ), "Must call before_update() before write_end" + return self.size + + # ----- construction ----------------------------------------------------- # + @classmethod + def from_tensor(cls, k: Tensor, v: Tensor, seq_dim: int) -> Self: + """Build a single-chunk cache pre-filled with the given K/V tensors.""" + cache = cls( + k_shape=tuple(k.shape), + v_shape=tuple(v.shape), + seq_dim=seq_dim, + chunk_size=k.shape[seq_dim], + window_size=k.shape[seq_dim], + device=k.device, + dtype=k.dtype, + ) + cache.before_update(0) + cache.update(k, v) + cache.after_update(0) + cache._curr_chunk_idx = 0 + return cache + + def __post_init__(self) -> None: + assert ( + self.k_shape[:-1] == self.v_shape[:-1] + ), "k and v must have the same shape except for the last dimension" + + tensor_dim = len(self.k_shape) + assert ( + -tensor_dim <= self.seq_dim < tensor_dim + ), f"seq_dim must be in [-{tensor_dim}, {tensor_dim}), got {self.seq_dim}" + # Normalize seq_dim to a non-negative index. + self.seq_dim = self.seq_dim if self.seq_dim >= 0 else self.seq_dim + tensor_dim + + assert self.sink_size >= 0, "sink_size must be non-negative" + + expected_length = self.sink_size + self.window_size + assert self.k_shape[self.seq_dim] == expected_length, ( + f"k_shape[seq_dim] ({self.k_shape[self.seq_dim]}) must equal " + f"sink_size + window_size ({expected_length})" + ) + assert (self.window_size + self.sink_size) % self.chunk_size == 0, ( + f"window_size + sink_size ({self.window_size + self.sink_size}) must be " + f"divisible by chunk_size ({self.chunk_size})" + ) + + self._k = torch.empty(self.k_shape, device=self.device, dtype=self.dtype) + self._v = torch.empty(self.v_shape, device=self.device, dtype=self.dtype) + + # ----- internal helpers ------------------------------------------------- # + def _seq_slice(self, start: int | None, end: int | None) -> tuple[slice | int, ...]: + """Index tuple selecting ``[start:end]`` on ``seq_dim``, all else full.""" + idx: list[slice | int] = [slice(None)] * len(self.k_shape) + idx[self.seq_dim] = slice(start, end) + return tuple(idx) + + def _roll_local_window_left(self) -> None: + """Shift the local window left by ``chunk_size`` (steady-state only).""" + total_size = self._k.shape[self.seq_dim] + assert ( + total_size == self._n_cached + ), f"Expected full cache: {total_size=} != {self._n_cached=}" + tokens_to_keep = self.window_size - self.chunk_size + if tokens_to_keep > 0: + src_start = self.sink_size + self.chunk_size + src_end = total_size + dst_start = self.sink_size + dst_end = self.sink_size + tokens_to_keep + dst_slice = self._seq_slice(dst_start, dst_end) + src_slice = self._seq_slice(src_start, src_end) + self._k[dst_slice] = self._k[src_slice].clone() + self._v[dst_slice] = self._v[src_slice].clone() + + def _current_chunk_overlaps_sink(self) -> bool: + assert ( + self._curr_chunk_idx is not None + ), "Must call before_update() before checking sink overlap" + return ( + self.sink_size > 0 + and self._curr_chunk_idx * self.chunk_size < self.sink_size + ) + + def _current_write_bounds(self) -> tuple[int, int]: + """Physical cache range written by the current update.""" + assert ( + self._curr_chunk_idx is not None + ), "Must call before_update() before computing write bounds" + total_size = self._k.shape[self.seq_dim] + assert ( + self.chunk_size <= total_size + ), f"chunk_size ({self.chunk_size}) must be <= cache size ({total_size})" + if self._curr_chunk_idx == self._prev_chunk_idx + 1: + write_start = min(self._n_cached, total_size - self.chunk_size) + write_end = write_start + self.chunk_size + elif self._curr_chunk_idx == self._prev_chunk_idx: + write_end = min(self._n_cached, total_size) + write_start = max(write_end - self.chunk_size, 0) + else: + raise ValueError( + f"{self._curr_chunk_idx=} should be either " + f"{self._prev_chunk_idx + 1} or {self._prev_chunk_idx}." + ) + return write_start, write_end + + def _write_current_chunk(self, k: Tensor, v: Tensor) -> None: + """Write the current chunk through a filling/steady-compatible path.""" + write_start, write_end = self._current_write_bounds() + read_start = 0 + read_end = write_end - write_start + + if ( + self.sink_size > 0 + and not self._current_chunk_overlaps_sink() + and write_start < self.sink_size + ): + write_start = self.sink_size + keep_size = write_end - write_start + read_end = self.chunk_size + read_start = read_end - keep_size + + sl_read = self._seq_slice(read_start, read_end) + sl_write = self._seq_slice(write_start, write_end) + self._k[sl_write] = k[sl_read] + self._v[sl_write] = v[sl_read] + + def _visible_end(self) -> int: + """Right edge of cached tokens visible to attention this update.""" + assert ( + self._curr_chunk_idx is not None + ), "Must call before_update() before computing visible cache size" + total_size = self._k.shape[self.seq_dim] + if self._curr_chunk_idx == self._prev_chunk_idx + 1: + return min(self._n_cached + self.chunk_size, total_size) + if self._curr_chunk_idx == self._prev_chunk_idx: + return min(self._n_cached, total_size) + raise ValueError( + f"{self._curr_chunk_idx=} should be either " + f"{self._prev_chunk_idx + 1} or {self._prev_chunk_idx}." + ) + + # ----- public lifecycle ------------------------------------------------- # + def is_steady_state(self) -> bool: + """True if the cache is full (steady-state phase).""" + assert ( + self._curr_chunk_idx is not None + ), "Must call before_update() before is_steady_state()" + total_size = self._k.shape[self.seq_dim] + is_full = total_size == self._n_cached + is_overlapping_with_sink = ( + self.sink_size > 0 + and self._curr_chunk_idx * self.chunk_size < self.sink_size + ) + return is_full and not is_overlapping_with_sink + + def before_update(self, chunk_idx: int) -> None: + """Prepare the cache before writing (roll window if steady-state).""" + assert ( + self._curr_chunk_idx is None + ), "Must call after_update() before before_update()" + self._curr_chunk_idx = chunk_idx + if chunk_idx == self._prev_chunk_idx: + return + assert chunk_idx == self._prev_chunk_idx + 1, ( + "Expected the new chunk_idx to be +1 from the previous chunk_idx, " + f"got {chunk_idx} != {self._prev_chunk_idx} + 1" + ) + if self.is_steady_state(): + self._roll_local_window_left() + + def update(self, k: Tensor, v: Tensor) -> None: + """Write the new chunk's K/V into the cache.""" + assert ( + self._curr_chunk_idx is not None + ), "Must call before_update() before update()" + chunk_size_k = k.shape[self.seq_dim] + chunk_size_v = v.shape[self.seq_dim] + assert chunk_size_k == self.chunk_size, ( + f"Expected input k chunk_size {self.chunk_size} at seq_dim " + f"{self.seq_dim}, got {chunk_size_k}" + ) + assert chunk_size_v == self.chunk_size, ( + f"Expected input v chunk_size {self.chunk_size} at seq_dim " + f"{self.seq_dim}, got {chunk_size_v}" + ) + self._write_current_chunk(k, v) + + def after_update(self, chunk_idx: int) -> None: + """Finalize bookkeeping after writing the chunk.""" + assert ( + chunk_idx == self._curr_chunk_idx + ), f"Expected chunk_idx to be {self._curr_chunk_idx}, got {chunk_idx}" + if self._curr_chunk_idx == self._prev_chunk_idx + 1: + if not self.is_steady_state(): + self._n_cached += self.chunk_size + self._prev_chunk_idx += 1 + elif self._curr_chunk_idx == self._prev_chunk_idx: + pass + else: + raise ValueError( + f"{self._curr_chunk_idx=} should be either " + f"{self._prev_chunk_idx + 1} or {self._prev_chunk_idx}." + ) + self._curr_chunk_idx = None + + def cached_k(self) -> Tensor: + """Cached keys (valid prefix while filling, full buffer in steady-state).""" + return self._k[self._seq_slice(0, self.size)] + + def cached_v(self) -> Tensor: + """Cached values (valid prefix while filling, full buffer in steady-state).""" + return self._v[self._seq_slice(0, self.size)] + + def reset(self) -> None: + """Reset the cache to its initial empty state.""" + self._prev_chunk_idx = -1 + self._n_cached = 0 + self._curr_chunk_idx = None diff --git a/python/sglang/multimodal_gen/runtime/models/dits/omnidreams_rope.py b/python/sglang/multimodal_gen/runtime/models/dits/omnidreams_rope.py new file mode 100644 index 000000000000..95855a338414 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/omnidreams_rope.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +"""3D rotary position embedding for OmniDreams (pure-torch NeoX port). + +Faithful port of FlashDreams ``flashdreams.core.attention.rope`` + +``rope_kernel`` (the fused Triton kernel) into device-agnostic PyTorch so the +single-chunk and autoregressive forward paths run without the Triton dependency. + +Key facts (verified against FlashDreams source for the distilled single-view +checkpoint): +- head_dim 128 splits T:H:W = 44:42:42 (``dim_h = dim_w = head_dim//6*2``). +- NeoX / non-interleaved rotation: pairs are ``(d, d + D/2)``. +- NTK extrapolation: theta = 10000 * ratio**(dim/(dim-2)); H/W ratio = 3.0, T = 1.0. +- ``shift_t(ar_idx)`` advances time positions by ``ar_idx * len_t`` (used for the + KV-cache window in autoregressive rollout). Keys are rotated *before* being + written to the cache (standard RoPE, not the cache-relative variant). +""" + +from __future__ import annotations + +import torch +from einops import repeat +from torch import Tensor + +# OmniDreams uses non-interleaved (NeoX) rotation: the pair is (d, d + D/2). +ROPE_IS_NEOX_STYLE = True + + +def rope_dims(head_dim: int) -> tuple[int, int, int]: + """Return the (T, H, W) split of a head dim for 3D RoPE. + + For head_dim=128 this yields (44, 42, 42): each spatial axis takes + ``head_dim // 6 * 2`` and time takes the remainder. + """ + dim_h = dim_w = head_dim // 6 * 2 + dim_t = head_dim - dim_h - dim_w + return dim_t, dim_h, dim_w + + +def _compute_freqs( + dim: int, + extrapolation_ratio: float = 1.0, + device: torch.device | str = "cpu", +) -> Tensor: + """Base RoPE frequencies for one axis with NTK extrapolation (shape [dim//2]).""" + dim_range = ( + torch.arange(0, dim, 2, dtype=torch.float32, device=device)[: (dim // 2)] / dim + ) + ntk_factor = extrapolation_ratio ** (dim / (dim - 2)) + theta = 10000.0 * ntk_factor + return 1.0 / (theta**dim_range) + + +class RotaryPositionEmbedding3D: + """Standard 3D RoPE with unbounded autoregressive time positions. + + ``shift_t`` returns a full-width ``[L, 1, 1, head_dim]`` frequency tensor + (memory layout (T, H, W)) suitable for :func:`apply_rope_freqs`. + """ + + def __init__( + self, + head_dim: int, + len_h: int, + len_w: int, + len_t: int, + h_extrapolation_ratio: float = 1.0, + w_extrapolation_ratio: float = 1.0, + t_extrapolation_ratio: float = 1.0, + device: torch.device | str = "cpu", + ) -> None: + self.head_dim = head_dim + self.len_h = len_h + self.len_w = len_w + self.len_t = len_t + self.device = device + + dim_w = dim_h = head_dim // 6 * 2 + dim_t = head_dim - (dim_h + dim_w) + self.raw_freqs_h = _compute_freqs(dim_h, h_extrapolation_ratio, device) + self.raw_freqs_w = _compute_freqs(dim_w, w_extrapolation_ratio, device) + self.raw_freqs_t = _compute_freqs(dim_t, t_extrapolation_ratio, device) + + self.freqs_t, self.freqs_h, self.freqs_w = self._freq_components_for_len(len_t) + + def _freq_components_for_len(self, len_t: int) -> tuple[Tensor, Tensor, Tensor]: + seq_t = torch.arange(len_t, dtype=torch.float32, device=self.device) + seq_h = torch.arange(self.len_h, dtype=torch.float32, device=self.device) + seq_w = torch.arange(self.len_w, dtype=torch.float32, device=self.device) + freqs_t = repeat( + torch.outer(seq_t, self.raw_freqs_t), + "t d -> (t h w) 1 1 d", + h=self.len_h, + w=self.len_w, + ) + freqs_h = repeat( + torch.outer(seq_h, self.raw_freqs_h), + "h d -> (t h w) 1 1 d", + t=len_t, + w=self.len_w, + ) + freqs_w = repeat( + torch.outer(seq_w, self.raw_freqs_w), + "w d -> (t h w) 1 1 d", + t=len_t, + h=self.len_h, + ) + return freqs_t, freqs_h, freqs_w + + def _cat_freqs(self, freqs_t: Tensor, freqs_h: Tensor, freqs_w: Tensor) -> Tensor: + # Non-interleaved (NeoX): [t, h, w] repeated twice along the last dim. + return torch.cat([freqs_t, freqs_h, freqs_w] * 2, dim=-1) + + def shift_t(self, autoregressive_index: int = 0) -> Tensor: + """Frequencies for AR chunk ``autoregressive_index`` (offset = idx * len_t).""" + offset = autoregressive_index * self.len_t + freqs_t = self.freqs_t + offset * self.raw_freqs_t + return self._cat_freqs(freqs_t, self.freqs_h, self.freqs_w) + + +def apply_rope_freqs(x: Tensor, freqs: Tensor) -> Tensor: + """Apply NeoX 3D RoPE to ``x``. + + Mirrors the FlashDreams fused kernel (non-interleaved branch): + out[a] = x[a] * cos(f) - x[b] * sin(f) + out[b] = x[b] * cos(f) + x[a] * sin(f) + with ``(a, b) = (d, d + D/2)`` and ``f`` the first-half angles of ``freqs``. + + Args: + x: ``[B, S, H, D]`` query or key. + freqs: ``[S, 1, 1, D]`` full-width frequencies from ``shift_t`` (the + first and second halves are identical by construction). + Returns: + Rotated tensor of shape ``[B, S, H, D]`` (cos/sin computed in fp32). + """ + seq_len = freqs.shape[0] + half = x.shape[-1] // 2 + f = freqs[..., :half].reshape(seq_len, half).view(1, seq_len, 1, half) + cos = f.cos().to(x.dtype) + sin = f.sin().to(x.dtype) + a = x[..., :half] + b = x[..., half:] + return torch.cat([a * cos - b * sin, b * cos + a * sin], dim=-1) diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/omnidreams_text.py b/python/sglang/multimodal_gen/runtime/models/encoders/omnidreams_text.py new file mode 100644 index 000000000000..885043304643 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/encoders/omnidreams_text.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +"""OmniDreams text-embedding transform (Cosmos-Reason1-7B ``full_concat``). + +OmniDreams conditions cross-attention on a 100352-dim text embedding built from +the Cosmos-Reason1-7B (Qwen2.5-VL, 28 transformer layers x 3584 hidden) hidden +states. The transform is a faithful port of FlashDreams +``flashdreams/infra/encoder/text/cosmos_reason1.py``: + +1. Run the LM with ``output_hidden_states=True`` -> tuple of ``num_layers + 1`` + tensors (index 0 is the embedding layer). +2. Drop index 0 (embedding layer); per-token mean/std-normalize each of the + remaining 28 layers independently. +3. Concatenate along the feature dim -> ``28 * 3584 = 100352``. + +This module holds the **pure tensor transform** (CPU-testable, no HF model). The +GPU-only encoder wrapper that actually runs Cosmos-Reason1-7B lives in the +pipeline's text-encoding stage and calls :func:`full_concat_embeddings` on the +LM's ``output_hidden_states``. +""" + +from collections.abc import Sequence + +import torch +from torch import Tensor + +# Cosmos-Reason1-7B: 28 transformer layers x 3584 hidden = 100352 concat dim. +COSMOS_REASON1_NUM_LAYERS = 28 +COSMOS_REASON1_HIDDEN = 3584 +FULL_CONCAT_DIM = COSMOS_REASON1_NUM_LAYERS * COSMOS_REASON1_HIDDEN # 100352 + +_NORM_EPS = 1e-8 + + +def mean_normalize(tensor: Tensor) -> Tensor: + """Per-token mean/std normalization over the last (feature) dim. + + Matches FlashDreams ``_mean_normalize``: ``std`` uses torch's default + (unbiased / Bessel-corrected) estimator. + """ + return (tensor - tensor.mean(dim=-1, keepdim=True)) / ( + tensor.std(dim=-1, keepdim=True) + _NORM_EPS + ) + + +def full_concat_embeddings( + hidden_states: Sequence[Tensor], + *, + skip_embedding_layer: bool = True, +) -> Tensor: + """Build the OmniDreams ``full_concat`` text embedding. + + Args: + hidden_states: LM ``output_hidden_states`` -- a sequence of + ``num_layers + 1`` tensors each ``[B, L, H]`` (index 0 is the + embedding layer). For Cosmos-Reason1-7B this is 29 tensors of + ``[B, L, 3584]``. + skip_embedding_layer: drop ``hidden_states[0]`` (the embedding layer) + before normalizing, as FlashDreams does. + + Returns: + ``[B, L, num_transformer_layers * H]`` -- ``[B, L, 100352]`` for + Cosmos-Reason1-7B. + """ + start = 1 if skip_embedding_layer else 0 + layers = [ + mean_normalize(hidden_states[i]) for i in range(start, len(hidden_states)) + ] + if not layers: + raise ValueError( + "full_concat_embeddings got no transformer layers to concatenate; " + f"received {len(hidden_states)} hidden_states with " + f"skip_embedding_layer={skip_embedding_layer}" + ) + return torch.cat(layers, dim=-1) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_omnidreams_flow_match.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_omnidreams_flow_match.py new file mode 100644 index 000000000000..c291b75597a8 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_omnidreams_flow_match.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +"""OmniDreams self-forcing flow-match scheduler (port of FlashDreams ``fm.py``). + +Why a dedicated scheduler instead of reusing +``scheduling_self_forcing_flow_match.SelfForcingFlowMatchScheduler``: + +1. Schedule construction differs. OmniDreams warps the FULL ``num_train_timesteps`` + table and selects entries by ``denoising_timesteps`` (e.g. [1000, 450] -> + indices [0, 550] -> sigmas [1.0, 0.8036]). Building a 2-point linspace + directly instead yields [1.0, 0.8333] -- a silent ~3.7% sigma error. +2. Step semantics differ. OmniDreams uses self-forcing renoise + (``clean = noisy - sigma*flow``; ``noisy = (1-sigma)*clean + sigma*noise`` + with FRESH noise each step), not the deterministic Euler + ``prev = sample + v*(sigma_next - sigma)``. + +Distilled single-view defaults: num_inference_steps=2, +denoising_timesteps=[1000, 450], shift=5.0, sigma_min=0.0 -> sigmas +[1.0, 0.8036], warped network timesteps [1000.0, 803.57]. +""" + +from __future__ import annotations + +from collections.abc import Callable + +import torch +from torch import Tensor + + +def warp_sigmas(sigmas: Tensor, shift: float) -> Tensor: + """DiffSynth schedule warp: ``shift * s / (1 + (shift - 1) * s)``.""" + return shift * sigmas / (1.0 + (shift - 1.0) * sigmas) + + +class OmniDreamsFlowMatchScheduler: + """Self-forcing flow-match scheduler for OmniDreams. + + Attributes: + denoising_sigmas: per-step sigmas, shape ``[num_inference_steps]``. + denoising_step_list: per-step (warped) timesteps fed to the network. + """ + + def __init__( + self, + num_inference_steps: int = 2, + denoising_timesteps: tuple[int, ...] = (1000, 450), + shift: float = 5.0, + sigma_max: float = 1.0, + sigma_min: float = 0.0, + num_train_timesteps: int = 1000, + extra_one_step: bool = True, + warp_denoising_step: bool = True, + device: torch.device | str = "cpu", + ) -> None: + assert num_inference_steps == len(denoising_timesteps), ( + f"num_inference_steps ({num_inference_steps}) must equal " + f"len(denoising_timesteps) ({len(denoising_timesteps)})" + ) + N = num_train_timesteps + self.num_train_timesteps = N + self.num_inference_steps = num_inference_steps + + if extra_one_step: + base = torch.linspace(sigma_max, sigma_min, N + 1, dtype=torch.float32)[:-1] + else: + base = torch.linspace(sigma_max, sigma_min, N, dtype=torch.float32) + full_sigmas = warp_sigmas(base, shift) + full_timesteps = full_sigmas * N + + idxs = [N - t for t in denoising_timesteps] + if warp_denoising_step: + step_list = [full_timesteps[i].item() if i < N else 0.0 for i in idxs] + sigma_list = [full_sigmas[i if i < N else N - 1].item() for i in idxs] + else: + step_list = [float(t) for t in denoising_timesteps] + snapped = [ + int(torch.argmin((full_timesteps - t).abs()).item()) for t in step_list + ] + sigma_list = [full_sigmas[i].item() for i in snapped] + + self.denoising_step_list = torch.tensor( + step_list, dtype=torch.float32, device=device + ) + self.denoising_sigmas = torch.tensor( + sigma_list, dtype=torch.float32, device=device + ) + self._full_sigmas = full_sigmas.to(device) + self._full_timesteps = full_timesteps.to(device) + + def to(self, device: torch.device | str) -> "OmniDreamsFlowMatchScheduler": + self.denoising_step_list = self.denoising_step_list.to(device) + self.denoising_sigmas = self.denoising_sigmas.to(device) + self._full_sigmas = self._full_sigmas.to(device) + self._full_timesteps = self._full_timesteps.to(device) + return self + + def sample( + self, + initial_noise: Tensor, + predict_flow: Callable[[Tensor, Tensor], Tensor], + rng: torch.Generator | None = None, + ) -> Tensor: + """Self-forcing denoising loop; returns the clean ``x0`` estimate. + + ``predict_flow(noisy, timestep)`` returns the network's flow prediction. + Iteration 0 trusts ``initial_noise`` as the sigma=1 sample; later + iterations re-noise the previous clean estimate with FRESH noise. + """ + input_dtype = initial_noise.dtype + noisy = initial_noise + clean: Tensor | None = None + for i in range(self.denoising_step_list.shape[0]): + sigma = self.denoising_sigmas[i] + timestep = self.denoising_step_list[i].to(dtype=input_dtype) + if i > 0: + assert clean is not None + noise = torch.empty_like(noisy).normal_(generator=rng) + noisy = ((1.0 - sigma) * clean + sigma * noise).to(input_dtype) + flow = predict_flow(noisy, timestep) + clean = noisy - sigma * flow + assert clean is not None, "denoising_step_list is empty" + return clean.to(input_dtype) + + def add_noise( + self, + clean_input: Tensor, + timestep: Tensor, + rng: torch.Generator | None = None, + ) -> Tensor: + """Forward corruption at an arbitrary timestep (snapped to the table). + + Used for context-noise on cached/clean frames (raw timestep 128 -> + sigma ~= 0.128). + """ + assert timestep.shape == (), f"expected scalar timestep, got {timestep.shape}" + idx = torch.argmin( + (self._full_timesteps - timestep.to(self._full_timesteps.dtype)).abs() + ).reshape(1) + sigma = self._full_sigmas.index_select(0, idx).reshape(()) + noise = torch.empty_like(clean_input).normal_(generator=rng) + return ((1.0 - sigma) * clean_input + sigma * noise).to(clean_input.dtype) + + def sigma_for_timestep(self, timestep: float) -> float: + """Return the warped-table sigma nearest to a raw ``timestep`` (for tests/debug).""" + idx = int(torch.argmin((self._full_timesteps - float(timestep)).abs()).item()) + return float(self._full_sigmas[idx].item()) + + +EntryClass = OmniDreamsFlowMatchScheduler diff --git a/python/sglang/multimodal_gen/runtime/pipelines/omnidreams_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/omnidreams_pipeline.py new file mode 100644 index 000000000000..d917799d3a88 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/omnidreams_pipeline.py @@ -0,0 +1,358 @@ +# SPDX-License-Identifier: Apache-2.0 +"""NVIDIA OmniDreams pipeline (autoregressive video world model). + +The checkpoint (``single_view/2b_res720p_30fps_i2v_hdmap_distilled.pt``) is a +flat, DiT-only ``.pt`` (570 keys, bf16) -- not a diffusers layout. This pipeline +therefore overrides ``_load_config`` (fabricates a model_index-like dict) and +``load_modules`` (loads the flat DiT directly, mirroring the Hunyuan3D +precedent). The VAE (Wan 2.1) and text encoder (Cosmos-Reason1-7B / Qwen2.5-VL) +are loaded alongside it. + +Stage layout (Hybrid monolithic, autoregressive): +``BeforeDenoising -> OmniDreamsDenoising (AR rollout) -> standard DecodingStage``. +The denoising stage concatenates the AR chunks into ``batch.latents``; the +standard single-pass VAE decode lets the Wan VAE's causal temporal feature cache +flow across chunk boundaries, giving correct continuity and the FlashDreams +frame counts (chunk0 -> 1+(len_t-1)*4, each later chunk -> len_t*4). +""" + +from __future__ import annotations + +import glob +import os +from itertools import chain +from typing import Any + +import torch +import torch.nn as nn + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.loader.fsdp_load import ( + load_model_from_full_model_state_dict, + set_default_torch_dtype, +) +from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping +from sglang.multimodal_gen.runtime.models.dits.omnidreams import OmniDreamsDiT +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_omnidreams_flow_match import ( # noqa: E501 + OmniDreamsFlowMatchScheduler, +) +from sglang.multimodal_gen.runtime.models.vaes.wanvae import AutoencoderKLWan +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.omnidreams import ( # noqa: E501 + OmniDreamsBeforeDenoisingStage, + OmniDreamsDenoisingStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + +# Default in-repo location of the distilled single-view checkpoint. +_DEFAULT_CKPT_RELPATH = "single_view/2b_res720p_30fps_i2v_hdmap_distilled.pt" + +# Cosmos-Reason1-7B text encoder (Qwen2.5-VL), pinned for numerical parity with +# FlashDreams (flashdreams/infra/encoder/text/cosmos_reason1.py). +_TEXT_ENCODER_ID = "nvidia/Cosmos-Reason1-7B" +_TEXT_ENCODER_REVISION = "3210bec0495fdc7a8d3dbb8d58da5711eab4b423" + +# Wan 2.1 VAE weights subdirectory candidates (relative to the model path). +_VAE_RELDIRS = ("vae", "wan_vae", "Wan2.1_VAE") + + +class OmniDreamsPipeline(ComposedPipelineBase): + pipeline_name = "OmniDreamsPipeline" + is_video_pipeline = True + _required_config_modules = [ + "transformer", + "vae", + "text_encoder", + "tokenizer", + "scheduler", + ] + + def _load_config(self) -> dict[str, Any]: + return { + "_class_name": self.pipeline_name, + "_diffusers_version": "0.0.0", + "transformer": ["sglang", "OmniDreamsDiT"], + "vae": ["diffusers", "AutoencoderKLWan"], + "text_encoder": ["transformers", "AutoModel"], + "tokenizer": ["transformers", "AutoProcessor"], + "scheduler": ["sglang", "OmniDreamsFlowMatchScheduler"], + } + + # ----- path resolution -------------------------------------------------- # + @staticmethod + def _resolve_ckpt_path(model_path: str) -> str: + """Locate the flat DiT ``.pt`` from a directory, repo, or direct file.""" + if os.path.isfile(model_path): + return model_path + candidate = os.path.join(model_path, _DEFAULT_CKPT_RELPATH) + if os.path.isfile(candidate): + return candidate + matches = sorted( + glob.glob(os.path.join(model_path, "**", "*.pt"), recursive=True) + ) + if matches: + return matches[0] + raise FileNotFoundError( + f"OmniDreams checkpoint (.pt) not found under {model_path}" + ) + + @staticmethod + def _resolve_vae_path(model_path: str) -> str: + """Locate the Wan 2.1 VAE weights (a flat ``.pth`` or a weights dir). + + Looks for a ``vae``-like subdirectory first, then any ``*VAE*.pth`` / + ``*vae*.safetensors`` under the model path. The caller (GPU bring-up) + may also point ``model_path`` directly at a Wan VAE directory. + """ + base = model_path if os.path.isdir(model_path) else os.path.dirname(model_path) + for sub in _VAE_RELDIRS: + cand = os.path.join(base, sub) + if os.path.isdir(cand) or os.path.isfile(cand): + return cand + for pattern in ("**/*vae*.safetensors", "**/*VAE*.safetensors"): + matches = sorted(glob.glob(os.path.join(base, pattern), recursive=True)) + if matches: + return matches[0] + raise FileNotFoundError( + f"Diffusers-format Wan 2.1 VAE not found under {base}. Place a " + f"diffusers Wan VAE (*.safetensors + config.json) under a 'vae/' " + f"subdirectory or pass its path explicitly." + ) + + # ----- component loaders ------------------------------------------------ # + @classmethod + def _load_flat_dit( + cls, + dit_config: Any, + ckpt_path: str, + device: torch.device, + dtype: torch.dtype, + ) -> nn.Module: + """Instantiate OmniDreamsDiT on meta, load the flat .pt, then fuse weights. + + The flat checkpoint keys equal the submodule names, so the param mapping + is the identity. The custom loader bypasses the generic loader's + post-load hook, so ``post_load_weights`` (72->68 padding-mask fuse + + last-layer shuffle fuse) is invoked explicitly here. + """ + with set_default_torch_dtype(dtype), torch.device("meta"): + model = OmniDreamsDiT(config=dit_config, hf_config={}) + + weights = torch.load(ckpt_path, map_location="cpu", weights_only=True) + weight_iterator = ((k, v) for k, v in weights.items()) + mapping_fn = get_param_names_mapping(model.param_names_mapping) + + load_model_from_full_model_state_dict( + model, + weight_iterator, + device, + dtype, + strict=True, + param_names_mapping=mapping_fn, + ) + + model.post_load_weights() + + for name, p in chain(model.named_parameters(), model.named_buffers()): + if p.is_meta: + raise RuntimeError(f"Unexpected param/buffer {name} on meta device.") + if isinstance(p, nn.Parameter): + p.requires_grad = False + + return model.eval() + + @classmethod + def _load_wan_vae( + cls, + vae_config: Any, + vae_path: str, + device: torch.device, + dtype: torch.dtype, + ) -> nn.Module: + """Build the SGLang Wan 2.1 VAE and load diffusers-format weights. + + SGLang's ``AutoencoderKLWan`` uses the diffusers WanVAE key naming, so + the VAE must be supplied in **diffusers format**: a ``vae/`` directory + with ``*.safetensors`` (+ ``config.json``) as exported by diffusers. + + The original lightx2v flat ``Wan2.1_VAE.pth`` uses a different + (original-Wan) key naming and is intentionally **not** remapped here — + converting between the two schemes is exactly what diffusers' + ``convert_wan_to_diffusers`` already does. Point this at a diffusers Wan + VAE (any ``Wan2.1-*-Diffusers/vae``) instead of the flat ``.pth``. + """ + with set_default_torch_dtype(dtype): + vae = AutoencoderKLWan(vae_config) + + state = cls._read_vae_state_dict(vae_path) + + try: + vae.load_state_dict(state, strict=True) + except RuntimeError as exc: + raise RuntimeError( + f"Failed to load Wan VAE weights from '{vae_path}' into the " + "diffusers-format AutoencoderKLWan. Supply the VAE in diffusers " + "format (a 'vae/' directory with *.safetensors + config.json). " + "The flat lightx2v 'Wan2.1_VAE.pth' uses original-Wan key names; " + "convert it with diffusers' convert_wan_to_diffusers first.\n" + f"Underlying error: {exc}" + ) from exc + return vae.to(device=device, dtype=dtype).eval() + + @staticmethod + def _read_vae_state_dict(vae_path: str) -> dict[str, torch.Tensor]: + """Read a VAE state dict from a diffusers ``*.safetensors`` dir/file.""" + if os.path.isdir(vae_path): + files = sorted(glob.glob(os.path.join(vae_path, "*.safetensors"))) + if not files: + raise FileNotFoundError( + f"No *.safetensors found under '{vae_path}'. Supply a " + "diffusers-format Wan VAE directory." + ) + from safetensors.torch import load_file as safetensors_load_file + + state: dict[str, torch.Tensor] = {} + for f in files: + state.update(safetensors_load_file(f)) + elif vae_path.endswith(".safetensors"): + from safetensors.torch import load_file as safetensors_load_file + + state = safetensors_load_file(vae_path) + else: + state = torch.load(vae_path, map_location="cpu", weights_only=True) + + # Some checkpoints nest the state dict under "model"/"state_dict". + if "state_dict" in state and isinstance(state["state_dict"], dict): + state = state["state_dict"] + return state + + @staticmethod + def _resolve_text_encoder_src(model_path: str) -> tuple[str, str | None]: + """Resolve the Cosmos-Reason1-7B source. + + Prefers a local ``/text_encoder`` directory (offline / + mirrored deployments), returning ``(local_dir, None)``; otherwise falls + back to the pinned HF id + revision. + """ + if os.path.isdir(model_path): + local = os.path.join(model_path, "text_encoder") + if os.path.isfile(os.path.join(local, "config.json")): + return local, None + return _TEXT_ENCODER_ID, _TEXT_ENCODER_REVISION + + @classmethod + def _load_text_encoder( + cls, model_path: str, device: torch.device + ) -> tuple[Any, Any]: + """Load Cosmos-Reason1-7B (Qwen2.5-VL) + processor. + + Uses a local ``/text_encoder`` dir when present, else the + pinned HF id + revision (see :meth:`_resolve_text_encoder_src`). + """ + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + src, revision = cls._resolve_text_encoder_src(model_path) + logger.info("OmniDreams: loading text encoder from %s", src) + processor = AutoProcessor.from_pretrained(src, revision=revision) + text_encoder = ( + Qwen2_5_VLForConditionalGeneration.from_pretrained( + src, + revision=revision, + torch_dtype=torch.bfloat16, # canonical kwarg; avoids fp32 load + ) + .eval() + .requires_grad_(False) + .to(device) + ) + return text_encoder, processor + + def load_modules( + self, + server_args: ServerArgs, + loaded_modules: dict[str, torch.nn.Module] | None = None, + ) -> dict[str, Any]: + pipeline_config = server_args.pipeline_config + device = get_local_torch_device() + dit_dtype = PRECISION_TO_TYPE[pipeline_config.dit_precision] + vae_dtype = PRECISION_TO_TYPE[pipeline_config.vae_precision] + + model_path = server_args.model_path + ckpt_path = self._resolve_ckpt_path(model_path) + logger.info("OmniDreams: loading flat DiT from %s", ckpt_path) + transformer = self._load_flat_dit( + pipeline_config.dit_config, ckpt_path, device, dit_dtype + ) + + vae_path = self._resolve_vae_path(model_path) + logger.info("OmniDreams: loading Wan 2.1 VAE from %s", vae_path) + vae = self._load_wan_vae( + pipeline_config.vae_config, vae_path, device, vae_dtype + ) + + text_encoder, tokenizer = self._load_text_encoder(model_path, device) + + scheduler = OmniDreamsFlowMatchScheduler( + num_inference_steps=len(pipeline_config.denoising_timesteps), + denoising_timesteps=tuple(pipeline_config.denoising_timesteps), + shift=( + pipeline_config.flow_shift + if pipeline_config.flow_shift is not None + else 5.0 + ), + sigma_min=pipeline_config.sigma_min, + device=device, + ) + + # Phase 6: populate memory budgets (GiB, approximate) for the + # ComponentResidencyManager offload scheduler. Exact values are + # TODO(gpu): measure on the target GPU with torch.cuda.memory_stats(). + self.memory_usages = { + "transformer": 4.0, # ~2B params in bf16 ≈ 4 GiB + "text_encoder": 14.0, # Cosmos-Reason1-7B ≈ 14 GiB + "vae": 1.0, # Wan 2.1 VAE ≈ 1 GiB + } + + return { + "transformer": transformer, + "vae": vae, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "scheduler": scheduler, + } + + def create_pipeline_stages(self, server_args: ServerArgs): + config = server_args.pipeline_config + self.add_stage( + stage_name="omnidreams_before_denoising", + stage=OmniDreamsBeforeDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + text_encoder=self.get_module("text_encoder"), + tokenizer=self.get_module("tokenizer"), + vae=self.get_module("vae"), + config=config, + ), + ) + self.add_stage( + stage_name="omnidreams_denoising", + stage=OmniDreamsDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + vae=self.get_module("vae"), + ), + ) + # Standard single-pass VAE decode. The denoising stage concatenates the + # AR chunks into batch.latents; the Wan VAE's causal temporal feature + # cache flows frame-to-frame within one decode() call, giving correct + # cross-chunk continuity and the FlashDreams frame counts (chunk0 -> + # 1+(len_t-1)*4, each later chunk -> len_t*4). + self.add_standard_decoding_stage(stage_name="omnidreams_decoding") + + +EntryClass = OmniDreamsPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py new file mode 100644 index 000000000000..0a83412d8011 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py @@ -0,0 +1,702 @@ +# SPDX-License-Identifier: Apache-2.0 +"""OmniDreams pre-processing + autoregressive denoising stages. + +Two model-specific stages implement the OmniDreams (FlashDreams port) inference +contract (decoding uses the standard single-pass ``DecodingStage``): + +* :class:`OmniDreamsBeforeDenoisingStage` -- text encoding (Cosmos-Reason1-7B + ``full_concat`` -> 100352), i2v reference-frame VAE latent + condition mask, + 2-step self-forcing schedule, and all ``DenoisingStage.verify_input`` fields. + AR geometry/conditioning are stashed in ``batch.extra["omnidreams"]``. + +* :class:`OmniDreamsDenoisingStage` -- a full ``forward()`` override (the base + ``DenoisingStage`` loop is single-pass and stateless) that drives the + autoregressive rollout: per-chunk ``BlockKVCache`` lifecycle, ``shift_t`` + RoPE, the 2-step self-forcing denoise, and the context-noise re-forward that + writes the *clean* chunk's K/V into the cache (mirroring FlashDreams + ``DiffusionModel.finalize`` and SGLang ``CausalDMDDenoisingStage``). + +Tensors flow in PATCHIFIED token space ``[B, L, D]`` during the denoise loop +(matching the DiT forward), and each chunk's clean latent is unpatchified to +``[B, C, T, H, W]``. The chunks are concatenated into ``batch.latents`` for the +standard single-pass decode, whose Wan VAE causal feature cache flows across +chunk boundaries (correct continuity + FlashDreams frame counts). + +A remaining HD-map VAE-encode numerics note is flagged inline with ``TODO(gpu)``. +""" + +from __future__ import annotations + +import PIL.Image +import torch + +from sglang.multimodal_gen.runtime.distributed import ( + get_local_torch_device, +) +from sglang.multimodal_gen.runtime.models.dits.omnidreams_rope import ( + RotaryPositionEmbedding3D, +) +from sglang.multimodal_gen.runtime.models.encoders.omnidreams_text import ( + full_concat_embeddings, +) +from sglang.multimodal_gen.runtime.models.vision_utils import ( + load_image, + normalize, + numpy_to_pt, + pil_to_numpy, + resize, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( + StageValidators as V, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import ( + VerificationResult, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.utils import PRECISION_TO_TYPE + +logger = init_logger(__name__) + +# Cosmos-Reason1 prompt template (FlashDreams cosmos_reason1.py). +_SYSTEM_PROMPT = ( + "You are a helpful assistant who will provide prompts to an image generator." +) +_TEXT_MAX_LENGTH = 512 +# Upper bound on autoregressive chunks per request. Bounds the rollout loop +# length (and thus GPU memory/compute) against an unbounded ``num_frames`` from +# the HTTP API. ~256 chunks * len_t(2) * 4 = ~2048 pixel frames. +_MAX_AR_CHUNKS = 256 + + +# --------------------------------------------------------------------------- # +# Pre-processing stage # +# --------------------------------------------------------------------------- # +class OmniDreamsBeforeDenoisingStage(PipelineStage): + """Consolidated pre-processing for the OmniDreams pipeline. + + Populates every field ``DenoisingStage.verify_input`` checks plus the AR + rollout state consumed by :class:`OmniDreamsDenoisingStage`. + """ + + def __init__( + self, + transformer, + scheduler=None, + text_encoder=None, + tokenizer=None, + vae=None, + config=None, + ) -> None: + super().__init__() + self.transformer = transformer + self.scheduler = scheduler + self.text_encoder = text_encoder + self.tokenizer = tokenizer + self.vae = vae + self.config = config + + # ----- helpers ---------------------------------------------------------- # + @torch.no_grad() + def _encode_text(self, prompt: str, device: torch.device) -> torch.Tensor: + """Cosmos-Reason1-7B text -> ``[1, L, 100352]`` full_concat embedding.""" + messages = [ + {"role": "system", "content": _SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ] + input_ids = self.tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=False, + return_tensors="pt", + ) + # Newer transformers return a BatchEncoding (dict-like) rather than a + # bare tensor; normalize to the input_ids tensor. + if not isinstance(input_ids, torch.Tensor): + input_ids = input_ids["input_ids"] + # Pad/truncate to a fixed context length (FlashDreams uses 512). Build + # an attention mask so the encoder does not attend to padding tokens + # (otherwise short prompts get corrupted embeddings). + pad_id = getattr(self.tokenizer, "pad_token_id", None) or 0 + valid_len = input_ids.shape[1] + if valid_len < _TEXT_MAX_LENGTH: + pad = input_ids.new_full( + (input_ids.shape[0], _TEXT_MAX_LENGTH - valid_len), pad_id + ) + input_ids = torch.cat([input_ids, pad], dim=1) + attention_mask = torch.cat( + [ + input_ids.new_ones((input_ids.shape[0], valid_len)), + input_ids.new_zeros( + (input_ids.shape[0], _TEXT_MAX_LENGTH - valid_len) + ), + ], + dim=1, + ) + else: + input_ids = input_ids[:, :_TEXT_MAX_LENGTH] + attention_mask = input_ids.new_ones(input_ids.shape) + input_ids = input_ids.to(device) + attention_mask = attention_mask.to(device) + out = self.text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=True, + ) + return full_concat_embeddings(out.hidden_states) + + @staticmethod + def _resolve_input_image(batch: Req): + """First non-None image-like input on the request, in priority order. + + Accepts a preprocessed tensor, raw pixels, the standard + ``condition_image``/``vae_image`` fields, or an ``image_path`` to load. + Returns ``None`` when the request carries no reference image (text-only / + unconditioned). Uses explicit None-checks -- ``a or b`` chaining raises on + tensor operands. + """ + for attr in ( + "preprocessed_image", + "pixel_values", + "condition_image", + "vae_image", + ): + val = getattr(batch, attr, None) + if val is not None: + return val + image_path = getattr(batch, "image_path", None) + if image_path is None: + return None + if isinstance(image_path, (list, tuple)): + return image_path[0] if len(image_path) > 0 else None + return image_path + + @staticmethod + def _preprocess_pixels( + image, height: int, width: int, device: torch.device, dtype: torch.dtype + ) -> torch.Tensor | None: + """Pixels (PIL / path / tensor) -> ``[1, 3, 1, H, W]`` in ``[-1, 1]``. + + (The batch dim is synthesized when the input lacks it; the Wan VAE encodes + a single image at a time.) + + Mirrors ``ImageVAEEncodingStage.preprocess`` (``image_encoding.py``): + PIL/path -> ``resize(h, w)`` -> ``pil_to_numpy`` -> ``numpy_to_pt`` -> + ``normalize``. Already-tensor inputs pass through (assumed pre-sized) and + skip the ``[0,1]->[-1,1]`` normalize when already signed (``min() < 0``), + matching the standard stage's ``do_normalize`` heuristic. A temporal axis + is added so the Wan VAE sees a single-frame clip. + """ + if isinstance(image, (str, bytes)): + image = load_image(image) + if isinstance(image, PIL.Image.Image): + image = resize(image, height, width) + x = numpy_to_pt(pil_to_numpy(image)) # [1,3,H,W] in [0,1] + x = normalize(x) # -> [-1,1] + elif torch.is_tensor(image): + x = image + if x.dim() == 3: # [3,H,W] -> [1,3,H,W] + x = x.unsqueeze(0) + if x.min() >= 0: # assume [0,1] pixels -> [-1,1] + x = normalize(x) + else: + return None + x = x.to(device=device, dtype=dtype) + if x.dim() == 4: # [B,3,H,W] -> [B,3,1,H,W] + x = x.unsqueeze(2) + return x + + @torch.no_grad() + def _encode_reference_image( + self, + batch: Req, + device: torch.device, + vae_dtype: torch.dtype, + height: int, + width: int, + ) -> torch.Tensor | None: + """VAE-encode the i2v reference frame -> normalized latent ``[B,16,1,h,w]``. + + Returns ``None`` (text-only / unconditioned smoke) when no reference + image is present on the request. + """ + image = self._resolve_input_image(batch) + if image is None: + logger.warning( + "OmniDreams: no reference image on request; running unconditioned. " + "Provide condition_image/image_path (or pixel_values) for i2v." + ) + return None + + x = self._preprocess_pixels(image, height, width, device, vae_dtype) + if x is None: + return None + return self._vae_encode_normalized(x) + + @torch.no_grad() + def _vae_encode_normalized(self, x: torch.Tensor) -> torch.Tensor: + """``[B,3,T,H,W]`` pixels -> latent normalized into the DiT space. + + Encode, take the distribution mode (deterministic), then ``(z-mean)/std``. + Shared by the i2v reference frame and the HD-map conditioning encode. + """ + latent_dist = self.vae.encode(x) + latent = ( + latent_dist.mode() + if hasattr(latent_dist, "mode") + else latent_dist.sample() if hasattr(latent_dist, "sample") else latent_dist + ) + mean = torch.tensor( + self.vae.latents_mean, device=latent.device, dtype=latent.dtype + ).view(1, -1, 1, 1, 1) + std = torch.tensor( + self.vae.latents_std, device=latent.device, dtype=latent.dtype + ).view(1, -1, 1, 1, 1) + return (latent - mean) / std + + @torch.no_grad() + def _encode_hdmap( + self, + batch: Req, + device: torch.device, + vae_dtype: torch.dtype, + dit_dtype: torch.dtype, + num_chunks: int, + height: int, + width: int, + ) -> list[torch.Tensor] | None: + """Per-chunk HD-map conditioning -> ``list[num_chunks]`` of patchified tokens. + + HD-map is a *per-chunk* driving condition (each latent chunk has its own + HD-map frames), so this returns a list indexed by chunk -- not one tensor + shared across chunks. Returns ``None`` when the request carries no HD-map + input, in which case the AR stage falls back to zeros (HDMap disabled). + + Accepts ``batch.hdmap_path`` / ``batch.hdmap_pixels`` as either a single + input (broadcast to every chunk) or a per-chunk list. Each entry runs the + same pixel-preprocess + VAE-encode + patchify path as the reference image. + + TODO(gpu): the HD-map pixel -> 16ch-latent VAE numerics are validated on + GPU; this encode path only runs when real HD-map input is supplied. + """ + hdmap = getattr(batch, "hdmap_path", None) + if hdmap is None: + hdmap = getattr(batch, "hdmap_pixels", None) + if hdmap is None: + return None + + per_chunk = list(hdmap) if isinstance(hdmap, (list, tuple)) else None + tokens: list[torch.Tensor] = [] + for ci in range(num_chunks): + if per_chunk is not None: + src = per_chunk[ci] if ci < len(per_chunk) else per_chunk[-1] + else: + src = hdmap + x = self._preprocess_pixels(src, height, width, device, vae_dtype) + if x is None: + logger.warning( + "OmniDreams: HD-map chunk %d preprocessed to None; disabling " + "HDMap (all chunks fall back to zeros). Check hdmap input.", + ci, + ) + return None + latent = self._vae_encode_normalized(x).to(dit_dtype) + # [B,16,t,h,w] -> [B, L, additional_concat_ch*pdim] via the DiT patchify. + tokens.append(self.transformer.patchify(latent)) + return tokens + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + config = server_args.pipeline_config + device = get_local_torch_device() + dit_dtype = PRECISION_TO_TYPE[config.dit_precision] + vae_dtype = PRECISION_TO_TYPE[config.vae_precision] + arch = config.dit_config.arch_config + + # --- geometry (patchified token grid) --- + height = int(batch.height) + width = int(batch.width) + sp = config.vae_config.arch_config.scale_factor_spatial # 8 + latent_h = height // sp + latent_w = width // sp + hp = latent_h // arch.patch_spatial + wp = latent_w // arch.patch_spatial + len_t = int(getattr(batch, "len_t", getattr(config, "len_t", 2))) + tokens_per_frame = hp * wp + chunk_tokens = len_t * tokens_per_frame + + # --- text conditioning (100352) --- + prompt = batch.prompt if isinstance(batch.prompt, str) else str(batch.prompt) + text_embeds = self._encode_text(prompt, device).to(dit_dtype) + batch.prompt_embeds = [text_embeds] + batch.negative_prompt_embeds = None + batch.image_embeds = [] + batch.do_classifier_free_guidance = False + + # --- i2v reference latent -> patchified frame-0 token block --- + image_latent = self._encode_reference_image( + batch, device, vae_dtype, height, width + ) + if image_latent is not None: + image_latent = image_latent.to(dit_dtype) + # [B,16,1,h,w] -> [B, hp*wp, 16*pdim] via the DiT patchify. + image_token = self.transformer.patchify(image_latent) + else: + image_token = None + batch.image_latent = None # AR stage assembles batch.latents itself. + + # --- schedule (2-step self-forcing) --- + scheduler = self.scheduler + batch.scheduler = scheduler + batch.timesteps = scheduler.denoising_step_list.to(device) + batch.sigmas = scheduler.denoising_sigmas.tolist() + batch.num_inference_steps = int(scheduler.denoising_step_list.shape[0]) + batch.guidance_scale = 1.0 + batch.eta = 0.0 + if batch.generator is None: + seed = int(getattr(batch, "seed", None) or 0) + batch.generator = torch.Generator(device=device).manual_seed(seed) + + # --- AR rollout state for the denoising stage --- + num_chunks = self._compute_num_chunks(batch, len_t) + # Per-chunk HD-map tokens (None -> AR stage uses zeros / HDMap disabled). + hdmap_tokens = self._encode_hdmap( + batch, device, vae_dtype, dit_dtype, num_chunks, height, width + ) + batch.extra["omnidreams"] = { + "hp": hp, + "wp": wp, + "len_t": len_t, + "tokens_per_frame": tokens_per_frame, + "chunk_tokens": chunk_tokens, + "latent_h": latent_h, + "latent_w": latent_w, + "num_chunks": num_chunks, + "window_size_t": int(getattr(batch, "window_size_t", 6)), + "sink_size_t": int(getattr(batch, "sink_size_t", 0)), + "context_noise": float(getattr(batch, "context_noise", 128)), + "image_token": image_token, # [B, hp*wp, in*pdim] or None + # Per-chunk HD-map tokens: None, or list[num_chunks] of + # [B, chunk_tokens, additional_concat_ch*pdim]. + "hdmap_tokens": hdmap_tokens, + } + # raw_latent_shape lets SDPA-path attn metadata stay a no-op. + batch.raw_latent_shape = ( + text_embeds.shape[0], + chunk_tokens, + arch.out_channels, + latent_h, + latent_w, + ) + return batch + + @staticmethod + def _compute_num_chunks(batch: Req, len_t: int) -> int: + """Latent chunks needed for the requested pixel-frame count. + + FlashDreams decode: chunk 0 -> ``1 + (len_t-1)*4`` frames, chunk>=1 -> + ``len_t*4`` frames (temporal compression 4). + """ + num_frames = int(getattr(batch, "num_frames", None) or 0) + tc = 4 + if num_frames <= 0: + n = max(1, int(getattr(batch, "num_chunks", 1))) + else: + first = 1 + (len_t - 1) * tc + if num_frames <= first: + n = 1 + else: + n = 1 + -(-(num_frames - first) // (len_t * tc)) # ceil division + if n > _MAX_AR_CHUNKS: + logger.warning( + "OmniDreams: requested %d AR chunks exceeds the cap %d; clamping " + "(num_frames=%d). Raise _MAX_AR_CHUNKS if longer rollouts are needed.", + n, + _MAX_AR_CHUNKS, + num_frames, + ) + n = _MAX_AR_CHUNKS + return n + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("prompt", batch.prompt, V.string_or_list_strings) + return result + + +# --------------------------------------------------------------------------- # +# Autoregressive denoising stage # +# --------------------------------------------------------------------------- # +class OmniDreamsDenoisingStage(DenoisingStage): + """Autoregressive rollout (full ``forward()`` override). + + The base ``DenoisingStage`` runs a single diffusers-style ``scheduler.step`` + loop with no cross-chunk state. OmniDreams instead rolls over latent chunks, + each chunk: roll the per-block KV window, run a 2-step self-forcing denoise, + then re-forward the *clean* chunk at the context-noise timestep to write the + authoritative (clean) K/V into the cache. + """ + + def __init__(self, transformer, scheduler, vae=None) -> None: + super().__init__(transformer, scheduler, vae=vae) + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + # Phase 6 guard: TP is supported via column/row parallel layers in the + # DiT, but SP (ulysses/ring) is not yet supported for the AR chunk loop. + # Guarded with try/except since SP is not initialized in CPU tests. + try: + from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_sp_world_size, + ) + + assert get_sp_world_size() <= 1, ( + "Sequence parallelism (SP) is not yet supported for OmniDreams. " + "Run with --ulysses-degree 1 --ring-degree 1." + ) + except AssertionError as e: + if "not initialized" in str(e): + pass # SP not booted — single-GPU (or CPU test), safe. + else: + raise + + # The downstream single-pass decode of the concatenated AR latents + # relies on the Wan VAE's causal temporal feature cache flowing across + # frames within one decode() call. A VAE without feature caching would + # change the frame-count semantics, so fail loudly rather than silently. + if self.vae is not None and not getattr(self.vae, "use_feature_cache", True): + raise RuntimeError( + "OmniDreams AR rollout requires a Wan VAE with " + "use_feature_cache=True for correct streaming frame counts." + ) + + config = server_args.pipeline_config + device = get_local_torch_device() + dit_dtype = PRECISION_TO_TYPE[config.dit_precision] + arch = config.dit_config.arch_config + st = batch.extra["omnidreams"] + + hp, wp, len_t = st["hp"], st["wp"], st["len_t"] + tokens_per_frame = st["tokens_per_frame"] + chunk_tokens = st["chunk_tokens"] + num_chunks = st["num_chunks"] + context_noise = st["context_noise"] + head_dim = arch.model_channels // arch.num_heads + in_d = arch.in_channels * arch.patch_temporal * arch.patch_spatial**2 + hdmap_d = ( + arch.additional_concat_ch * arch.patch_temporal * arch.patch_spatial**2 + ) + mask_d = arch.patch_temporal * arch.patch_spatial**2 + + scheduler = batch.scheduler if batch.scheduler is not None else self.scheduler + scheduler = scheduler.to(device) + text = batch.prompt_embeds[0].to(device=device, dtype=dit_dtype) + B = text.shape[0] + gen = batch.generator + # ``.normal_(generator=)`` requires the generator to live on the noise + # tensor's device; re-seed onto ``device`` if the caller passed a CPU one. + if gen is not None and gen.device != device: + gen = torch.Generator(device=device).manual_seed(gen.initial_seed()) + + # 3D RoPE for one chunk (NeoX 44:42:42; h/w extrapolate 3.0, t 1.0). + rope = RotaryPositionEmbedding3D( + head_dim=head_dim, + len_h=hp, + len_w=wp, + len_t=len_t, + h_extrapolation_ratio=3.0, + w_extrapolation_ratio=3.0, + t_extrapolation_ratio=1.0, + device=device, + ) + + # One BlockKVCache per transformer block (token counts = frames * hp*wp). + caches = self.transformer.init_kv_caches( + batch_size=B, + chunk_tokens=chunk_tokens, + window_tokens=st["window_size_t"] * tokens_per_frame, + sink_tokens=st["sink_size_t"] * tokens_per_frame, + device=device, + dtype=dit_dtype, + ) + + # Frame-0 conditioning (i2v): channel mask (into the DiT) + inject mask + # (pins the clean reference latent) + the reference token block. + frame0 = tokens_per_frame + image_token = st["image_token"] + if image_token is not None: + image_full = torch.zeros( + B, chunk_tokens, in_d, device=device, dtype=dit_dtype + ) + image_full[:, :frame0, :] = image_token.to(device=device, dtype=dit_dtype) + inject_mask = torch.zeros( + B, chunk_tokens, 1, device=device, dtype=dit_dtype + ) + inject_mask[:, :frame0, :] = 1.0 + cond_mask_c0 = torch.zeros( + B, chunk_tokens, mask_d, device=device, dtype=dit_dtype + ) + cond_mask_c0[:, :frame0, :] = 1.0 + else: + image_full = inject_mask = None + cond_mask_c0 = torch.zeros( + B, chunk_tokens, mask_d, device=device, dtype=dit_dtype + ) + cond_mask_zero = torch.zeros( + B, chunk_tokens, mask_d, device=device, dtype=dit_dtype + ) + hdmap_zero = torch.zeros( + B, chunk_tokens, hdmap_d, device=device, dtype=dit_dtype + ) + + # Phase 6: precompute cross-attn K/V once per prompt (text context is + # static; avoids redundant k_proj/v_proj in every forward call). + # Feed the projected context (crossattn_proj applied) so the cached K/V + # match the per-block input dimensions. + cross_attn_kv = self.transformer.precompute_cross_attn_kv( + self.transformer.crossattn_proj(text) + ) + + # Phase 5: compute view_indices for cross-view attention (optional). + # Default: single-view (V=1). Multi-view is gated by + # arch.enable_cross_view_attn and num_views on the request. + view_count = int(getattr(batch, "num_views", 1) or 1) + # Bound num_views to the camera-embedding table to prevent an + # out-of-range index (and reject nonsensical/abusive values). + n_cameras = int(getattr(arch, "n_cameras_emb", 1)) + if view_count < 1 or view_count > n_cameras: + raise ValueError( + f"num_views={view_count} out of range [1, {n_cameras}] " + "(n_cameras_emb)." + ) + view_indices: torch.Tensor | None = None + if view_count > 1 and self.transformer.adaln_view_embedder is not None: + view_indices = ( + torch.arange(view_count, device=device, dtype=torch.long) + .unsqueeze(0) + .expand(B, -1) + ) # [B, V] + + # Loop-invariant context-noise timestep tensor (same scalar every chunk). + ctx_noise_t = torch.tensor(context_noise, device=device, dtype=dit_dtype) + + latent_chunks: list[torch.Tensor] = [] + for chunk_idx in range(num_chunks): + rope_freqs = rope.shift_t(chunk_idx) + is_first = chunk_idx == 0 + cond_mask = cond_mask_c0 if is_first else cond_mask_zero + # HD-map is per-chunk: index this chunk's tokens (None -> zeros, i.e. + # HDMap disabled). Explicit None-check -- a tensor in ``or`` raises. + if st["hdmap_tokens"] is None: + hdmap_chunk = hdmap_zero + else: + hdmap_chunk = st["hdmap_tokens"][chunk_idx].to( + device=device, dtype=dit_dtype + ) + pin = is_first and image_full is not None + + def predict_flow(noisy: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + if pin: + noisy = noisy * (1.0 - inject_mask) + image_full * inject_mask + return self.transformer( + hidden_states=noisy, + encoder_hidden_states=text, + timestep=t, + condition_video_input_mask=cond_mask, + rope_freqs=rope_freqs, + hdmap_condition=hdmap_chunk, + kv_caches=caches, + cross_attn_kv=cross_attn_kv, + view_indices=view_indices, + ) + + for c in caches: + c.before_update(chunk_idx) + + noise = torch.empty( + B, chunk_tokens, in_d, device=device, dtype=dit_dtype + ).normal_(generator=gen) + clean = scheduler.sample(noise, predict_flow=predict_flow, rng=gen) + if pin: + clean = clean * (1.0 - inject_mask) + image_full * inject_mask + + # Authoritative cache write: re-forward the CLEAN chunk at the + # context-noise timestep so the cache holds in-distribution K/V. + ctx_latent = scheduler.add_noise( + clean, + ctx_noise_t, + rng=gen, + ) + self.transformer( + hidden_states=ctx_latent, + encoder_hidden_states=text, + timestep=ctx_noise_t, + condition_video_input_mask=cond_mask, + rope_freqs=rope_freqs, + hdmap_condition=hdmap_chunk, + kv_caches=caches, + cross_attn_kv=cross_attn_kv, + view_indices=view_indices, + ) + + for c in caches: + c.after_update(chunk_idx) + + # [B, L, out*pdim] -> [B, out, len_t, h, w]. + latent_chunks.append(self.transformer.unpatchify(clean, len_t, hp, wp)) + + # Concatenate the AR chunks into the full latent sequence. The standard + # DecodingStage decodes this in a single pass; the Wan VAE's causal + # temporal feature cache flows across chunk boundaries, yielding correct + # continuity and FlashDreams frame counts. + batch.latents = torch.cat(latent_chunks, dim=2) + + # Phase 6: SP post-process — latents may need gathering when SP is + # eventually supported. Currently a no-op (SP is guarded at entry). + batch.latents = self._postprocess_sp_latents(batch, server_args) + + return batch + + def _postprocess_sp_latents( + self, batch: Req, server_args: ServerArgs + ) -> torch.Tensor: + """Gather sharded latents when SP is active (future, currently no-op). + + When SP is enabled, each rank outputs partial-sequence latents that + must be all-gathered along the time dimension. This is a placeholder + that returns the latents as-is when SP is not active. + """ + try: + from sglang.multimodal_gen.runtime.distributed import ( + get_sp_world_size, + sequence_model_parallel_all_gather, + ) + + if get_sp_world_size() > 1 and getattr( + batch, "did_sp_shard_latents", False + ): + return sequence_model_parallel_all_gather(batch.latents, dim=2) + except (ImportError, AssertionError): + pass + return batch.latents + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)]) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check("generator", batch.generator, V.generator_or_list_generators) + result.add_check( + "omnidreams_extra", + batch.extra.get("omnidreams"), + lambda x: isinstance(x, dict), + ) + return result + + def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + return result diff --git a/python/sglang/multimodal_gen/test/server/consistency_threshold.json b/python/sglang/multimodal_gen/test/server/consistency_threshold.json index 15cba9f06ef9..53b00ee6f9e3 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_threshold.json +++ b/python/sglang/multimodal_gen/test/server/consistency_threshold.json @@ -121,6 +121,12 @@ "psnr_threshold": 24.0, "mean_abs_diff_threshold": 10.0 }, + "omnidreams_2b_i2v": { + "clip_threshold": 0.90, + "ssim_threshold": 0.85, + "psnr_threshold": 20.0, + "mean_abs_diff_threshold": 15.0 + }, "ltx_2_two_stage_t2v": { "clip_threshold": 0.90, "ssim_threshold": 0.89, diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index e1ac033786d4..a276436a38e4 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -26,6 +26,7 @@ MULTI_FRAME_I2I_sampling_params, MULTI_IMAGE_TI2I_sampling_params, MULTI_IMAGE_TI2I_UPLOAD_sampling_params, + OMNIDREAMS_I2V_sampling_params, T2I_sampling_params, T2V_sampling_params, _make_modelopt_ci_case, @@ -39,6 +40,7 @@ DEFAULT_FLUX_2_KLEIN_BASE_4B_MODEL_NAME_FOR_TEST, DEFAULT_JOYAI_IMAGE_EDIT_MODEL_NAME_FOR_TEST, DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST, + DEFAULT_OMNIDREAMS_2B_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_EDIT_2509_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_EDIT_2511_MODEL_NAME_FOR_TEST, DEFAULT_QWEN_IMAGE_EDIT_MODEL_NAME_FOR_TEST, @@ -797,6 +799,19 @@ ), run_component_accuracy_check=False, ), + # === OmniDreams (NVIDIA autoregressive video world model) === + DiffusionTestCase( + "omnidreams_2b_i2v", + DiffusionServerArgs( + model_path=DEFAULT_OMNIDREAMS_2B_MODEL_NAME_FOR_TEST, + modality="video", + ), + OMNIDREAMS_I2V_sampling_params, + run_perf_check=True, + run_consistency_check=True, + run_component_accuracy_check=False, + run_models_api_check=False, + ), ] if not current_platform.is_hip(): diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 025ec7b5ef4c..c2a033d4ea09 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -2913,6 +2913,14 @@ "expected_median_denoise_ms": 0.0, "estimated_full_test_time_s": 592.3 }, + "omnidreams_2b_i2v": { + "stages_ms": {}, + "denoise_step_ms": {}, + "expected_e2e_ms": 0.0, + "expected_avg_denoise_ms": 0.0, + "expected_median_denoise_ms": 0.0, + "estimated_full_test_time_s": 135.0 + }, "wan22_modelopt_nvfp4_t2v": { "stages_ms": {}, "denoise_step_ms": {}, diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 28f9756b12da..e8ee024f2361 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -482,6 +482,19 @@ def from_req_perf_record( extras={"preset": "V4_QUALITY_48", "seed": 0}, ) +OMNIDREAMS_I2V_sampling_params = DiffusionSamplingParams( + prompt="A car drives down a sunny city street.", + image_path=( + "https://is1-ssl.mzstatic.com/image/thumb/Music114/v4/5f/fa/56/" + "5ffa56c2-ea1f-7a17-6bad-192ff9b6476d/825646124206.jpg/600x600bb.jpg" + ), + output_size="1280x704", + direct_url_test=True, + num_frames=13, + fps=30, + extras={"seed": 42}, +) + MODELOPT_T2I_CI_sampling_params = DiffusionSamplingParams( prompt="Doraemon is eating dorayaki", output_size="768x768", diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index cb53038b60d4..a47821194633 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -151,6 +151,9 @@ def roberta_processing_compat(*processor_args, **processor_kwargs): DEFAULT_WAN_2_2_T2V_A14B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.2-T2V-A14B-Diffusers" DEFAULT_WAN_2_2_I2V_A14B_MODEL_NAME_FOR_TEST = "Wan-AI/Wan2.2-I2V-A14B-Diffusers" +# OmniDreams (NVIDIA autoregressive video world model) +DEFAULT_OMNIDREAMS_2B_MODEL_NAME_FOR_TEST = "nvidia/omni-dreams-models" + # MOVA video generation models DEFAULT_MOVA_360P_MODEL_NAME_FOR_TEST = "OpenMOSS-Team/MOVA-360p" diff --git a/python/sglang/multimodal_gen/test/unit/data/omnidreams_dit_keys.txt b/python/sglang/multimodal_gen/test/unit/data/omnidreams_dit_keys.txt new file mode 100644 index 000000000000..f3bab9495765 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/data/omnidreams_dit_keys.txt @@ -0,0 +1,570 @@ +additional_patch_embedding.proj.1.weight +blocks.0.adaln_modulation_cross_attn.1.weight +blocks.0.adaln_modulation_cross_attn.2.weight +blocks.0.adaln_modulation_mlp.1.weight +blocks.0.adaln_modulation_mlp.2.weight +blocks.0.adaln_modulation_self_attn.1.weight +blocks.0.adaln_modulation_self_attn.2.weight +blocks.0.cross_attn.k_norm.weight +blocks.0.cross_attn.k_proj.weight +blocks.0.cross_attn.output_proj.weight +blocks.0.cross_attn.q_norm.weight +blocks.0.cross_attn.q_proj.weight +blocks.0.cross_attn.v_proj.weight +blocks.0.mlp.layer1.weight +blocks.0.mlp.layer2.weight +blocks.0.self_attn.k_norm.weight +blocks.0.self_attn.k_proj.weight +blocks.0.self_attn.output_proj.weight +blocks.0.self_attn.q_norm.weight +blocks.0.self_attn.q_proj.weight +blocks.0.self_attn.v_proj.weight +blocks.1.adaln_modulation_cross_attn.1.weight +blocks.1.adaln_modulation_cross_attn.2.weight +blocks.1.adaln_modulation_mlp.1.weight +blocks.1.adaln_modulation_mlp.2.weight +blocks.1.adaln_modulation_self_attn.1.weight +blocks.1.adaln_modulation_self_attn.2.weight +blocks.1.cross_attn.k_norm.weight +blocks.1.cross_attn.k_proj.weight +blocks.1.cross_attn.output_proj.weight +blocks.1.cross_attn.q_norm.weight +blocks.1.cross_attn.q_proj.weight +blocks.1.cross_attn.v_proj.weight +blocks.1.mlp.layer1.weight +blocks.1.mlp.layer2.weight +blocks.1.self_attn.k_norm.weight +blocks.1.self_attn.k_proj.weight +blocks.1.self_attn.output_proj.weight +blocks.1.self_attn.q_norm.weight +blocks.1.self_attn.q_proj.weight +blocks.1.self_attn.v_proj.weight +blocks.10.adaln_modulation_cross_attn.1.weight +blocks.10.adaln_modulation_cross_attn.2.weight +blocks.10.adaln_modulation_mlp.1.weight +blocks.10.adaln_modulation_mlp.2.weight +blocks.10.adaln_modulation_self_attn.1.weight +blocks.10.adaln_modulation_self_attn.2.weight +blocks.10.cross_attn.k_norm.weight +blocks.10.cross_attn.k_proj.weight +blocks.10.cross_attn.output_proj.weight +blocks.10.cross_attn.q_norm.weight +blocks.10.cross_attn.q_proj.weight +blocks.10.cross_attn.v_proj.weight +blocks.10.mlp.layer1.weight +blocks.10.mlp.layer2.weight +blocks.10.self_attn.k_norm.weight +blocks.10.self_attn.k_proj.weight +blocks.10.self_attn.output_proj.weight +blocks.10.self_attn.q_norm.weight +blocks.10.self_attn.q_proj.weight +blocks.10.self_attn.v_proj.weight +blocks.11.adaln_modulation_cross_attn.1.weight +blocks.11.adaln_modulation_cross_attn.2.weight +blocks.11.adaln_modulation_mlp.1.weight +blocks.11.adaln_modulation_mlp.2.weight +blocks.11.adaln_modulation_self_attn.1.weight +blocks.11.adaln_modulation_self_attn.2.weight +blocks.11.cross_attn.k_norm.weight +blocks.11.cross_attn.k_proj.weight +blocks.11.cross_attn.output_proj.weight +blocks.11.cross_attn.q_norm.weight +blocks.11.cross_attn.q_proj.weight +blocks.11.cross_attn.v_proj.weight +blocks.11.mlp.layer1.weight +blocks.11.mlp.layer2.weight +blocks.11.self_attn.k_norm.weight +blocks.11.self_attn.k_proj.weight +blocks.11.self_attn.output_proj.weight +blocks.11.self_attn.q_norm.weight +blocks.11.self_attn.q_proj.weight +blocks.11.self_attn.v_proj.weight +blocks.12.adaln_modulation_cross_attn.1.weight +blocks.12.adaln_modulation_cross_attn.2.weight +blocks.12.adaln_modulation_mlp.1.weight +blocks.12.adaln_modulation_mlp.2.weight +blocks.12.adaln_modulation_self_attn.1.weight +blocks.12.adaln_modulation_self_attn.2.weight +blocks.12.cross_attn.k_norm.weight +blocks.12.cross_attn.k_proj.weight +blocks.12.cross_attn.output_proj.weight +blocks.12.cross_attn.q_norm.weight +blocks.12.cross_attn.q_proj.weight +blocks.12.cross_attn.v_proj.weight +blocks.12.mlp.layer1.weight +blocks.12.mlp.layer2.weight +blocks.12.self_attn.k_norm.weight +blocks.12.self_attn.k_proj.weight +blocks.12.self_attn.output_proj.weight +blocks.12.self_attn.q_norm.weight +blocks.12.self_attn.q_proj.weight +blocks.12.self_attn.v_proj.weight +blocks.13.adaln_modulation_cross_attn.1.weight +blocks.13.adaln_modulation_cross_attn.2.weight +blocks.13.adaln_modulation_mlp.1.weight +blocks.13.adaln_modulation_mlp.2.weight +blocks.13.adaln_modulation_self_attn.1.weight +blocks.13.adaln_modulation_self_attn.2.weight +blocks.13.cross_attn.k_norm.weight +blocks.13.cross_attn.k_proj.weight +blocks.13.cross_attn.output_proj.weight +blocks.13.cross_attn.q_norm.weight +blocks.13.cross_attn.q_proj.weight +blocks.13.cross_attn.v_proj.weight +blocks.13.mlp.layer1.weight +blocks.13.mlp.layer2.weight +blocks.13.self_attn.k_norm.weight +blocks.13.self_attn.k_proj.weight +blocks.13.self_attn.output_proj.weight +blocks.13.self_attn.q_norm.weight +blocks.13.self_attn.q_proj.weight +blocks.13.self_attn.v_proj.weight +blocks.14.adaln_modulation_cross_attn.1.weight +blocks.14.adaln_modulation_cross_attn.2.weight +blocks.14.adaln_modulation_mlp.1.weight +blocks.14.adaln_modulation_mlp.2.weight +blocks.14.adaln_modulation_self_attn.1.weight +blocks.14.adaln_modulation_self_attn.2.weight +blocks.14.cross_attn.k_norm.weight +blocks.14.cross_attn.k_proj.weight +blocks.14.cross_attn.output_proj.weight +blocks.14.cross_attn.q_norm.weight +blocks.14.cross_attn.q_proj.weight +blocks.14.cross_attn.v_proj.weight +blocks.14.mlp.layer1.weight +blocks.14.mlp.layer2.weight +blocks.14.self_attn.k_norm.weight +blocks.14.self_attn.k_proj.weight +blocks.14.self_attn.output_proj.weight +blocks.14.self_attn.q_norm.weight +blocks.14.self_attn.q_proj.weight +blocks.14.self_attn.v_proj.weight +blocks.15.adaln_modulation_cross_attn.1.weight +blocks.15.adaln_modulation_cross_attn.2.weight +blocks.15.adaln_modulation_mlp.1.weight +blocks.15.adaln_modulation_mlp.2.weight +blocks.15.adaln_modulation_self_attn.1.weight +blocks.15.adaln_modulation_self_attn.2.weight +blocks.15.cross_attn.k_norm.weight +blocks.15.cross_attn.k_proj.weight +blocks.15.cross_attn.output_proj.weight +blocks.15.cross_attn.q_norm.weight +blocks.15.cross_attn.q_proj.weight +blocks.15.cross_attn.v_proj.weight +blocks.15.mlp.layer1.weight +blocks.15.mlp.layer2.weight +blocks.15.self_attn.k_norm.weight +blocks.15.self_attn.k_proj.weight +blocks.15.self_attn.output_proj.weight +blocks.15.self_attn.q_norm.weight +blocks.15.self_attn.q_proj.weight +blocks.15.self_attn.v_proj.weight +blocks.16.adaln_modulation_cross_attn.1.weight +blocks.16.adaln_modulation_cross_attn.2.weight +blocks.16.adaln_modulation_mlp.1.weight +blocks.16.adaln_modulation_mlp.2.weight +blocks.16.adaln_modulation_self_attn.1.weight +blocks.16.adaln_modulation_self_attn.2.weight +blocks.16.cross_attn.k_norm.weight +blocks.16.cross_attn.k_proj.weight +blocks.16.cross_attn.output_proj.weight +blocks.16.cross_attn.q_norm.weight +blocks.16.cross_attn.q_proj.weight +blocks.16.cross_attn.v_proj.weight +blocks.16.mlp.layer1.weight +blocks.16.mlp.layer2.weight +blocks.16.self_attn.k_norm.weight +blocks.16.self_attn.k_proj.weight +blocks.16.self_attn.output_proj.weight +blocks.16.self_attn.q_norm.weight +blocks.16.self_attn.q_proj.weight +blocks.16.self_attn.v_proj.weight +blocks.17.adaln_modulation_cross_attn.1.weight +blocks.17.adaln_modulation_cross_attn.2.weight +blocks.17.adaln_modulation_mlp.1.weight +blocks.17.adaln_modulation_mlp.2.weight +blocks.17.adaln_modulation_self_attn.1.weight +blocks.17.adaln_modulation_self_attn.2.weight +blocks.17.cross_attn.k_norm.weight +blocks.17.cross_attn.k_proj.weight +blocks.17.cross_attn.output_proj.weight +blocks.17.cross_attn.q_norm.weight +blocks.17.cross_attn.q_proj.weight +blocks.17.cross_attn.v_proj.weight +blocks.17.mlp.layer1.weight +blocks.17.mlp.layer2.weight +blocks.17.self_attn.k_norm.weight +blocks.17.self_attn.k_proj.weight +blocks.17.self_attn.output_proj.weight +blocks.17.self_attn.q_norm.weight +blocks.17.self_attn.q_proj.weight +blocks.17.self_attn.v_proj.weight +blocks.18.adaln_modulation_cross_attn.1.weight +blocks.18.adaln_modulation_cross_attn.2.weight +blocks.18.adaln_modulation_mlp.1.weight +blocks.18.adaln_modulation_mlp.2.weight +blocks.18.adaln_modulation_self_attn.1.weight +blocks.18.adaln_modulation_self_attn.2.weight +blocks.18.cross_attn.k_norm.weight +blocks.18.cross_attn.k_proj.weight +blocks.18.cross_attn.output_proj.weight +blocks.18.cross_attn.q_norm.weight +blocks.18.cross_attn.q_proj.weight +blocks.18.cross_attn.v_proj.weight +blocks.18.mlp.layer1.weight +blocks.18.mlp.layer2.weight +blocks.18.self_attn.k_norm.weight +blocks.18.self_attn.k_proj.weight +blocks.18.self_attn.output_proj.weight +blocks.18.self_attn.q_norm.weight +blocks.18.self_attn.q_proj.weight +blocks.18.self_attn.v_proj.weight +blocks.19.adaln_modulation_cross_attn.1.weight +blocks.19.adaln_modulation_cross_attn.2.weight +blocks.19.adaln_modulation_mlp.1.weight +blocks.19.adaln_modulation_mlp.2.weight +blocks.19.adaln_modulation_self_attn.1.weight +blocks.19.adaln_modulation_self_attn.2.weight +blocks.19.cross_attn.k_norm.weight +blocks.19.cross_attn.k_proj.weight +blocks.19.cross_attn.output_proj.weight +blocks.19.cross_attn.q_norm.weight +blocks.19.cross_attn.q_proj.weight +blocks.19.cross_attn.v_proj.weight +blocks.19.mlp.layer1.weight +blocks.19.mlp.layer2.weight +blocks.19.self_attn.k_norm.weight +blocks.19.self_attn.k_proj.weight +blocks.19.self_attn.output_proj.weight +blocks.19.self_attn.q_norm.weight +blocks.19.self_attn.q_proj.weight +blocks.19.self_attn.v_proj.weight +blocks.2.adaln_modulation_cross_attn.1.weight +blocks.2.adaln_modulation_cross_attn.2.weight +blocks.2.adaln_modulation_mlp.1.weight +blocks.2.adaln_modulation_mlp.2.weight +blocks.2.adaln_modulation_self_attn.1.weight +blocks.2.adaln_modulation_self_attn.2.weight +blocks.2.cross_attn.k_norm.weight +blocks.2.cross_attn.k_proj.weight +blocks.2.cross_attn.output_proj.weight +blocks.2.cross_attn.q_norm.weight +blocks.2.cross_attn.q_proj.weight +blocks.2.cross_attn.v_proj.weight +blocks.2.mlp.layer1.weight +blocks.2.mlp.layer2.weight +blocks.2.self_attn.k_norm.weight +blocks.2.self_attn.k_proj.weight +blocks.2.self_attn.output_proj.weight +blocks.2.self_attn.q_norm.weight +blocks.2.self_attn.q_proj.weight +blocks.2.self_attn.v_proj.weight +blocks.20.adaln_modulation_cross_attn.1.weight +blocks.20.adaln_modulation_cross_attn.2.weight +blocks.20.adaln_modulation_mlp.1.weight +blocks.20.adaln_modulation_mlp.2.weight +blocks.20.adaln_modulation_self_attn.1.weight +blocks.20.adaln_modulation_self_attn.2.weight +blocks.20.cross_attn.k_norm.weight +blocks.20.cross_attn.k_proj.weight +blocks.20.cross_attn.output_proj.weight +blocks.20.cross_attn.q_norm.weight +blocks.20.cross_attn.q_proj.weight +blocks.20.cross_attn.v_proj.weight +blocks.20.mlp.layer1.weight +blocks.20.mlp.layer2.weight +blocks.20.self_attn.k_norm.weight +blocks.20.self_attn.k_proj.weight +blocks.20.self_attn.output_proj.weight +blocks.20.self_attn.q_norm.weight +blocks.20.self_attn.q_proj.weight +blocks.20.self_attn.v_proj.weight +blocks.21.adaln_modulation_cross_attn.1.weight +blocks.21.adaln_modulation_cross_attn.2.weight +blocks.21.adaln_modulation_mlp.1.weight +blocks.21.adaln_modulation_mlp.2.weight +blocks.21.adaln_modulation_self_attn.1.weight +blocks.21.adaln_modulation_self_attn.2.weight +blocks.21.cross_attn.k_norm.weight +blocks.21.cross_attn.k_proj.weight +blocks.21.cross_attn.output_proj.weight +blocks.21.cross_attn.q_norm.weight +blocks.21.cross_attn.q_proj.weight +blocks.21.cross_attn.v_proj.weight +blocks.21.mlp.layer1.weight +blocks.21.mlp.layer2.weight +blocks.21.self_attn.k_norm.weight +blocks.21.self_attn.k_proj.weight +blocks.21.self_attn.output_proj.weight +blocks.21.self_attn.q_norm.weight +blocks.21.self_attn.q_proj.weight +blocks.21.self_attn.v_proj.weight +blocks.22.adaln_modulation_cross_attn.1.weight +blocks.22.adaln_modulation_cross_attn.2.weight +blocks.22.adaln_modulation_mlp.1.weight +blocks.22.adaln_modulation_mlp.2.weight +blocks.22.adaln_modulation_self_attn.1.weight +blocks.22.adaln_modulation_self_attn.2.weight +blocks.22.cross_attn.k_norm.weight +blocks.22.cross_attn.k_proj.weight +blocks.22.cross_attn.output_proj.weight +blocks.22.cross_attn.q_norm.weight +blocks.22.cross_attn.q_proj.weight +blocks.22.cross_attn.v_proj.weight +blocks.22.mlp.layer1.weight +blocks.22.mlp.layer2.weight +blocks.22.self_attn.k_norm.weight +blocks.22.self_attn.k_proj.weight +blocks.22.self_attn.output_proj.weight +blocks.22.self_attn.q_norm.weight +blocks.22.self_attn.q_proj.weight +blocks.22.self_attn.v_proj.weight +blocks.23.adaln_modulation_cross_attn.1.weight +blocks.23.adaln_modulation_cross_attn.2.weight +blocks.23.adaln_modulation_mlp.1.weight +blocks.23.adaln_modulation_mlp.2.weight +blocks.23.adaln_modulation_self_attn.1.weight +blocks.23.adaln_modulation_self_attn.2.weight +blocks.23.cross_attn.k_norm.weight +blocks.23.cross_attn.k_proj.weight +blocks.23.cross_attn.output_proj.weight +blocks.23.cross_attn.q_norm.weight +blocks.23.cross_attn.q_proj.weight +blocks.23.cross_attn.v_proj.weight +blocks.23.mlp.layer1.weight +blocks.23.mlp.layer2.weight +blocks.23.self_attn.k_norm.weight +blocks.23.self_attn.k_proj.weight +blocks.23.self_attn.output_proj.weight +blocks.23.self_attn.q_norm.weight +blocks.23.self_attn.q_proj.weight +blocks.23.self_attn.v_proj.weight +blocks.24.adaln_modulation_cross_attn.1.weight +blocks.24.adaln_modulation_cross_attn.2.weight +blocks.24.adaln_modulation_mlp.1.weight +blocks.24.adaln_modulation_mlp.2.weight +blocks.24.adaln_modulation_self_attn.1.weight +blocks.24.adaln_modulation_self_attn.2.weight +blocks.24.cross_attn.k_norm.weight +blocks.24.cross_attn.k_proj.weight +blocks.24.cross_attn.output_proj.weight +blocks.24.cross_attn.q_norm.weight +blocks.24.cross_attn.q_proj.weight +blocks.24.cross_attn.v_proj.weight +blocks.24.mlp.layer1.weight +blocks.24.mlp.layer2.weight +blocks.24.self_attn.k_norm.weight +blocks.24.self_attn.k_proj.weight +blocks.24.self_attn.output_proj.weight +blocks.24.self_attn.q_norm.weight +blocks.24.self_attn.q_proj.weight +blocks.24.self_attn.v_proj.weight +blocks.25.adaln_modulation_cross_attn.1.weight +blocks.25.adaln_modulation_cross_attn.2.weight +blocks.25.adaln_modulation_mlp.1.weight +blocks.25.adaln_modulation_mlp.2.weight +blocks.25.adaln_modulation_self_attn.1.weight +blocks.25.adaln_modulation_self_attn.2.weight +blocks.25.cross_attn.k_norm.weight +blocks.25.cross_attn.k_proj.weight +blocks.25.cross_attn.output_proj.weight +blocks.25.cross_attn.q_norm.weight +blocks.25.cross_attn.q_proj.weight +blocks.25.cross_attn.v_proj.weight +blocks.25.mlp.layer1.weight +blocks.25.mlp.layer2.weight +blocks.25.self_attn.k_norm.weight +blocks.25.self_attn.k_proj.weight +blocks.25.self_attn.output_proj.weight +blocks.25.self_attn.q_norm.weight +blocks.25.self_attn.q_proj.weight +blocks.25.self_attn.v_proj.weight +blocks.26.adaln_modulation_cross_attn.1.weight +blocks.26.adaln_modulation_cross_attn.2.weight +blocks.26.adaln_modulation_mlp.1.weight +blocks.26.adaln_modulation_mlp.2.weight +blocks.26.adaln_modulation_self_attn.1.weight +blocks.26.adaln_modulation_self_attn.2.weight +blocks.26.cross_attn.k_norm.weight +blocks.26.cross_attn.k_proj.weight +blocks.26.cross_attn.output_proj.weight +blocks.26.cross_attn.q_norm.weight +blocks.26.cross_attn.q_proj.weight +blocks.26.cross_attn.v_proj.weight +blocks.26.mlp.layer1.weight +blocks.26.mlp.layer2.weight +blocks.26.self_attn.k_norm.weight +blocks.26.self_attn.k_proj.weight +blocks.26.self_attn.output_proj.weight +blocks.26.self_attn.q_norm.weight +blocks.26.self_attn.q_proj.weight +blocks.26.self_attn.v_proj.weight +blocks.27.adaln_modulation_cross_attn.1.weight +blocks.27.adaln_modulation_cross_attn.2.weight +blocks.27.adaln_modulation_mlp.1.weight +blocks.27.adaln_modulation_mlp.2.weight +blocks.27.adaln_modulation_self_attn.1.weight +blocks.27.adaln_modulation_self_attn.2.weight +blocks.27.cross_attn.k_norm.weight +blocks.27.cross_attn.k_proj.weight +blocks.27.cross_attn.output_proj.weight +blocks.27.cross_attn.q_norm.weight +blocks.27.cross_attn.q_proj.weight +blocks.27.cross_attn.v_proj.weight +blocks.27.mlp.layer1.weight +blocks.27.mlp.layer2.weight +blocks.27.self_attn.k_norm.weight +blocks.27.self_attn.k_proj.weight +blocks.27.self_attn.output_proj.weight +blocks.27.self_attn.q_norm.weight +blocks.27.self_attn.q_proj.weight +blocks.27.self_attn.v_proj.weight +blocks.3.adaln_modulation_cross_attn.1.weight +blocks.3.adaln_modulation_cross_attn.2.weight +blocks.3.adaln_modulation_mlp.1.weight +blocks.3.adaln_modulation_mlp.2.weight +blocks.3.adaln_modulation_self_attn.1.weight +blocks.3.adaln_modulation_self_attn.2.weight +blocks.3.cross_attn.k_norm.weight +blocks.3.cross_attn.k_proj.weight +blocks.3.cross_attn.output_proj.weight +blocks.3.cross_attn.q_norm.weight +blocks.3.cross_attn.q_proj.weight +blocks.3.cross_attn.v_proj.weight +blocks.3.mlp.layer1.weight +blocks.3.mlp.layer2.weight +blocks.3.self_attn.k_norm.weight +blocks.3.self_attn.k_proj.weight +blocks.3.self_attn.output_proj.weight +blocks.3.self_attn.q_norm.weight +blocks.3.self_attn.q_proj.weight +blocks.3.self_attn.v_proj.weight +blocks.4.adaln_modulation_cross_attn.1.weight +blocks.4.adaln_modulation_cross_attn.2.weight +blocks.4.adaln_modulation_mlp.1.weight +blocks.4.adaln_modulation_mlp.2.weight +blocks.4.adaln_modulation_self_attn.1.weight +blocks.4.adaln_modulation_self_attn.2.weight +blocks.4.cross_attn.k_norm.weight +blocks.4.cross_attn.k_proj.weight +blocks.4.cross_attn.output_proj.weight +blocks.4.cross_attn.q_norm.weight +blocks.4.cross_attn.q_proj.weight +blocks.4.cross_attn.v_proj.weight +blocks.4.mlp.layer1.weight +blocks.4.mlp.layer2.weight +blocks.4.self_attn.k_norm.weight +blocks.4.self_attn.k_proj.weight +blocks.4.self_attn.output_proj.weight +blocks.4.self_attn.q_norm.weight +blocks.4.self_attn.q_proj.weight +blocks.4.self_attn.v_proj.weight +blocks.5.adaln_modulation_cross_attn.1.weight +blocks.5.adaln_modulation_cross_attn.2.weight +blocks.5.adaln_modulation_mlp.1.weight +blocks.5.adaln_modulation_mlp.2.weight +blocks.5.adaln_modulation_self_attn.1.weight +blocks.5.adaln_modulation_self_attn.2.weight +blocks.5.cross_attn.k_norm.weight +blocks.5.cross_attn.k_proj.weight +blocks.5.cross_attn.output_proj.weight +blocks.5.cross_attn.q_norm.weight +blocks.5.cross_attn.q_proj.weight +blocks.5.cross_attn.v_proj.weight +blocks.5.mlp.layer1.weight +blocks.5.mlp.layer2.weight +blocks.5.self_attn.k_norm.weight +blocks.5.self_attn.k_proj.weight +blocks.5.self_attn.output_proj.weight +blocks.5.self_attn.q_norm.weight +blocks.5.self_attn.q_proj.weight +blocks.5.self_attn.v_proj.weight +blocks.6.adaln_modulation_cross_attn.1.weight +blocks.6.adaln_modulation_cross_attn.2.weight +blocks.6.adaln_modulation_mlp.1.weight +blocks.6.adaln_modulation_mlp.2.weight +blocks.6.adaln_modulation_self_attn.1.weight +blocks.6.adaln_modulation_self_attn.2.weight +blocks.6.cross_attn.k_norm.weight +blocks.6.cross_attn.k_proj.weight +blocks.6.cross_attn.output_proj.weight +blocks.6.cross_attn.q_norm.weight +blocks.6.cross_attn.q_proj.weight +blocks.6.cross_attn.v_proj.weight +blocks.6.mlp.layer1.weight +blocks.6.mlp.layer2.weight +blocks.6.self_attn.k_norm.weight +blocks.6.self_attn.k_proj.weight +blocks.6.self_attn.output_proj.weight +blocks.6.self_attn.q_norm.weight +blocks.6.self_attn.q_proj.weight +blocks.6.self_attn.v_proj.weight +blocks.7.adaln_modulation_cross_attn.1.weight +blocks.7.adaln_modulation_cross_attn.2.weight +blocks.7.adaln_modulation_mlp.1.weight +blocks.7.adaln_modulation_mlp.2.weight +blocks.7.adaln_modulation_self_attn.1.weight +blocks.7.adaln_modulation_self_attn.2.weight +blocks.7.cross_attn.k_norm.weight +blocks.7.cross_attn.k_proj.weight +blocks.7.cross_attn.output_proj.weight +blocks.7.cross_attn.q_norm.weight +blocks.7.cross_attn.q_proj.weight +blocks.7.cross_attn.v_proj.weight +blocks.7.mlp.layer1.weight +blocks.7.mlp.layer2.weight +blocks.7.self_attn.k_norm.weight +blocks.7.self_attn.k_proj.weight +blocks.7.self_attn.output_proj.weight +blocks.7.self_attn.q_norm.weight +blocks.7.self_attn.q_proj.weight +blocks.7.self_attn.v_proj.weight +blocks.8.adaln_modulation_cross_attn.1.weight +blocks.8.adaln_modulation_cross_attn.2.weight +blocks.8.adaln_modulation_mlp.1.weight +blocks.8.adaln_modulation_mlp.2.weight +blocks.8.adaln_modulation_self_attn.1.weight +blocks.8.adaln_modulation_self_attn.2.weight +blocks.8.cross_attn.k_norm.weight +blocks.8.cross_attn.k_proj.weight +blocks.8.cross_attn.output_proj.weight +blocks.8.cross_attn.q_norm.weight +blocks.8.cross_attn.q_proj.weight +blocks.8.cross_attn.v_proj.weight +blocks.8.mlp.layer1.weight +blocks.8.mlp.layer2.weight +blocks.8.self_attn.k_norm.weight +blocks.8.self_attn.k_proj.weight +blocks.8.self_attn.output_proj.weight +blocks.8.self_attn.q_norm.weight +blocks.8.self_attn.q_proj.weight +blocks.8.self_attn.v_proj.weight +blocks.9.adaln_modulation_cross_attn.1.weight +blocks.9.adaln_modulation_cross_attn.2.weight +blocks.9.adaln_modulation_mlp.1.weight +blocks.9.adaln_modulation_mlp.2.weight +blocks.9.adaln_modulation_self_attn.1.weight +blocks.9.adaln_modulation_self_attn.2.weight +blocks.9.cross_attn.k_norm.weight +blocks.9.cross_attn.k_proj.weight +blocks.9.cross_attn.output_proj.weight +blocks.9.cross_attn.q_norm.weight +blocks.9.cross_attn.q_proj.weight +blocks.9.cross_attn.v_proj.weight +blocks.9.mlp.layer1.weight +blocks.9.mlp.layer2.weight +blocks.9.self_attn.k_norm.weight +blocks.9.self_attn.k_proj.weight +blocks.9.self_attn.output_proj.weight +blocks.9.self_attn.q_norm.weight +blocks.9.self_attn.q_proj.weight +blocks.9.self_attn.v_proj.weight +crossattn_proj.0.bias +crossattn_proj.0.weight +final_layer.adaln_modulation.1.weight +final_layer.adaln_modulation.2.weight +final_layer.linear.weight +t_embedder.1.linear_1.weight +t_embedder.1.linear_2.weight +t_embedding_norm.weight +x_embedder.proj.1.weight diff --git a/python/sglang/multimodal_gen/test/unit/test_omnidreams_components.py b/python/sglang/multimodal_gen/test/unit/test_omnidreams_components.py new file mode 100644 index 000000000000..f9d2da2ebdec --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_omnidreams_components.py @@ -0,0 +1,654 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU component tests for the OmniDreams port (no checkpoint, no GPU). + +Covers the pure-torch building blocks that the GPU phases depend on: +- 3D NeoX RoPE (``omnidreams_rope``): layout, rotation correctness, ``shift_t``. +- ``BlockKVCache``: fill -> roll -> steady-state, sink retention, overwrite. +- ``OmniDreamsFlowMatchScheduler``: 2-step sigmas, self-forcing ``sample``, ``add_noise``. +- Cosmos-Reason1 ``full_concat_embeddings``: drop embedding layer, per-layer norm, 100352. +- A tiny-config ``OmniDreamsDiT`` end-to-end forward (single-chunk + KV-cache path). +- The ``OmniDreamsDenoisingStage`` autoregressive rollout orchestration. + +The structural/fixture checks live in ``test_omnidreams_scaffold.py``. +""" + +import types +from collections import Counter + +import numpy as np +import PIL.Image +import torch + +from sglang.multimodal_gen.configs.models.dits.omnidreams import ( + OmniDreamsDiTArchConfig, + OmniDreamsDiTConfig, +) +from sglang.multimodal_gen.runtime.models.dits.omnidreams import OmniDreamsDiT +from sglang.multimodal_gen.runtime.models.dits.omnidreams_kvcache import BlockKVCache +from sglang.multimodal_gen.runtime.models.dits.omnidreams_rope import ( + ROPE_IS_NEOX_STYLE, + RotaryPositionEmbedding3D, + apply_rope_freqs, + rope_dims, +) +from sglang.multimodal_gen.runtime.models.encoders.omnidreams_text import ( + FULL_CONCAT_DIM, + full_concat_embeddings, + mean_normalize, +) +from sglang.multimodal_gen.runtime.models.schedulers.scheduling_omnidreams_flow_match import ( # noqa: E501 + OmniDreamsFlowMatchScheduler, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.omnidreams import ( # noqa: E501 + OmniDreamsBeforeDenoisingStage, +) + + +# --------------------------------------------------------------------- RoPE -- # +def test_rope_dims_44_42_42_neox(): + assert rope_dims(128) == (44, 42, 42) + assert sum(rope_dims(128)) == 128 + assert ROPE_IS_NEOX_STYLE is True + + +def test_apply_rope_matches_neox_reference_and_preserves_norm(): + torch.manual_seed(0) + emb = RotaryPositionEmbedding3D( + head_dim=128, + len_h=4, + len_w=5, + len_t=2, + h_extrapolation_ratio=3.0, + w_extrapolation_ratio=3.0, + t_extrapolation_ratio=1.0, + ) + freqs = emb.shift_t(0) + L = freqs.shape[0] + assert L == 2 * 4 * 5 + assert freqs.shape[-1] == 128 + # NeoX builds freqs as [first | first] -> the two halves are identical. + assert torch.allclose(freqs[..., :64], freqs[..., 64:]) + + x = torch.randn(1, L, 16, 128) + out = apply_rope_freqs(x, freqs) + half = 64 + f = freqs[..., :half].reshape(L, half).view(1, L, 1, half) + cos, sin = f.cos(), f.sin() + a, b = x[..., :half], x[..., half:] + ref = torch.cat([a * cos - b * sin, b * cos + a * sin], dim=-1) + assert torch.allclose(out, ref, atol=1e-6) + # rotations are norm-preserving + assert torch.allclose(out.norm(dim=-1), x.norm(dim=-1), atol=1e-4) + + +def test_shift_t_advances_only_time_frequencies(): + emb = RotaryPositionEmbedding3D(head_dim=128, len_h=4, len_w=5, len_t=2) + dim_t_half = rope_dims(128)[0] // 2 # 22 + f0, f1 = emb.shift_t(0), emb.shift_t(1) + # time band (first dim_t_half angles) changes with ar_idx ... + assert not torch.allclose(f0[..., :dim_t_half], f1[..., :dim_t_half]) + # ... spatial bands (h then w) do not. + assert torch.allclose(f0[..., dim_t_half:64], f1[..., dim_t_half:64]) + + +# ----------------------------------------------------------------- KV cache -- # +def _chunk(val, B=1, n=2, d=3, size=2): + t = torch.full((B, size, n, d), float(val)) + return t, t.clone() + + +def test_kv_cache_fill_roll_steady_state(): + c = BlockKVCache( + k_shape=(1, 4, 2, 3), + v_shape=(1, 4, 2, 3), + seq_dim=1, + chunk_size=2, + window_size=4, + sink_size=0, + ) + # chunk 0: [10,10], visible 2 + c.before_update(0) + k, v = _chunk(10) + c.update(k, v) + ck = c.cached_k().clone() + c.after_update(0) + assert ck.shape[1] == 2 and bool((ck == 10).all()) + # chunk 1: [10,10,11,11], visible 4 (full) + c.before_update(1) + k, v = _chunk(11) + c.update(k, v) + ck = c.cached_k().clone() + c.after_update(1) + assert ck.shape[1] == 4 + assert bool((ck[:, :2] == 10).all()) and bool((ck[:, 2:] == 11).all()) + # chunk 2: steady-state roll left -> drop 10, keep 11, add 12 + c.before_update(2) + k, v = _chunk(12) + c.update(k, v) + ck = c.cached_k().clone() + c.after_update(2) + assert bool((ck[:, :2] == 11).all()) and bool((ck[:, 2:] == 12).all()) + + +def test_kv_cache_overwrite_same_chunk_idx(): + c = BlockKVCache( + k_shape=(1, 4, 2, 3), + v_shape=(1, 4, 2, 3), + seq_dim=1, + chunk_size=2, + window_size=4, + sink_size=0, + ) + for idx, val in enumerate([10, 11]): + c.before_update(idx) + k, v = _chunk(val) + c.update(k, v) + c.after_update(idx) + # rewriting the same chunk_idx refreshes the rightmost slots in place + c.before_update(1) + k, v = _chunk(99) + c.update(k, v) + ck = c.cached_k().clone() + c.after_update(1) + assert bool((ck[:, 2:] == 99).all()) and bool((ck[:, :2] == 10).all()) + + +def test_kv_cache_sink_is_never_evicted(): + c = BlockKVCache( + k_shape=(1, 6, 2, 3), + v_shape=(1, 6, 2, 3), + seq_dim=1, + chunk_size=2, + window_size=4, + sink_size=2, + ) + ck = None + for idx, val in enumerate([1, 2, 3, 4, 5]): + c.before_update(idx) + k, v = _chunk(val) + c.update(k, v) + ck = c.cached_k().clone() + c.after_update(idx) + # sink keeps the first chunk; rolling window keeps the last two chunks + assert bool((ck[:, :2] == 1).all()) + assert bool((ck[:, 2:4] == 4).all()) and bool((ck[:, 4:6] == 5).all()) + + +def test_kv_cache_from_tensor_roundtrip(): + kk, vv = torch.randn(1, 3, 2, 3), torch.randn(1, 3, 2, 3) + fc = BlockKVCache.from_tensor(kk, vv, seq_dim=1) + assert torch.allclose(fc.cached_k(), kk) and torch.allclose(fc.cached_v(), vv) + + +# ---------------------------------------------------------------- scheduler -- # +def test_scheduler_two_step_sigmas_and_timesteps(): + s = OmniDreamsFlowMatchScheduler() + sig = s.denoising_sigmas.tolist() + ts = s.denoising_step_list.tolist() + assert len(sig) == 2 + assert abs(sig[0] - 1.0) < 1e-6 + assert abs(sig[1] - 0.8036) < 1e-3 + assert abs(ts[0] - 1000.0) < 1e-3 + assert abs(ts[1] - 803.57) < 1.0 + # context-noise raw timestep 128 snaps to sigma ~= 0.128 + assert abs(s.sigma_for_timestep(128) - 0.128) < 0.02 + + +def test_scheduler_sample_ideal_flow_recovers_target(): + s = OmniDreamsFlowMatchScheduler() + init = torch.randn(1, 4) + target = torch.randn(1, 4) + step2sigma = { + round(t, 4): sg + for t, sg in zip(s.denoising_step_list.tolist(), s.denoising_sigmas.tolist()) + } + + def ideal_flow(noisy, t): + return (noisy - target) / step2sigma[round(float(t), 4)] + + out = s.sample(init, predict_flow=ideal_flow, rng=torch.Generator().manual_seed(7)) + assert torch.allclose(out, target, atol=1e-4) + + +def test_scheduler_add_noise_scales_by_sigma(): + s = OmniDreamsFlowMatchScheduler() + clean = torch.zeros(1, 4096) + noisy = s.add_noise( + clean, torch.tensor(128.0), rng=torch.Generator().manual_seed(3) + ) + # clean == 0 -> noisy == sigma * noise -> std approximates sigma + assert abs(noisy.std().item() - s.sigma_for_timestep(128)) < 0.05 + + +# ------------------------------------------------------------------- text ---- # +def test_full_concat_dim_and_per_layer_normalization(): + torch.manual_seed(0) + hs = [torch.randn(2, 7, 3584) for _ in range(29)] # 28 layers + embedding layer + emb = full_concat_embeddings(hs) + assert emb.shape == (2, 7, FULL_CONCAT_DIM) == (2, 7, 100352) + block0 = emb[..., :3584] + assert block0.mean(dim=-1).abs().max().item() < 1e-4 + assert (block0.std(dim=-1) - 1.0).abs().max().item() < 1e-2 + # dropping the embedding layer == normalizing layers 1..28 and concatenating + manual = torch.cat([mean_normalize(hs[i]) for i in range(1, 29)], dim=-1) + assert torch.allclose(emb, manual, atol=1e-6) + + +# ------------------------------------------------------- tiny DiT forward ---- # +def _tiny_dit() -> OmniDreamsDiT: + """A small CPU-constructible OmniDreamsDiT for end-to-end forward testing. + + head_dim = 24/2 = 12 keeps the RoPE 6-way split valid (dim_t/h/w = 4/4/4). + """ + arch = OmniDreamsDiTArchConfig( + in_channels=4, + out_channels=4, + model_channels=24, + num_blocks=2, + num_heads=2, + mlp_ratio=2.0, + adaln_lora_dim=8, + crossattn_proj_in_channels=32, + crossattn_emb_channels=16, + additional_concat_ch=4, + ) + model = OmniDreamsDiT(config=OmniDreamsDiTConfig(arch_config=arch), hf_config={}) + model.post_load_weights() # fuse padding-mask (24->20) + last-layer shuffle + return model.eval() + + +def _tiny_inputs(model, grid=(2, 2, 2), B=1, lctx=5): + arch = model.arch + gt, gh, gw = grid + L = gt * gh * gw + pdim = arch.patch_temporal * arch.patch_spatial**2 # kt*kh*kw + hidden = torch.randn(B, L, arch.in_channels * pdim) + cond_mask = torch.zeros(B, L, pdim) + hdmap = torch.randn(B, L, arch.additional_concat_ch * pdim) + ctx = torch.randn(B, lctx, arch.crossattn_proj_in_channels) + head_dim = arch.model_channels // arch.num_heads + rope = RotaryPositionEmbedding3D(head_dim=head_dim, len_h=gh, len_w=gw, len_t=gt) + return hidden, cond_mask, hdmap, ctx, rope, (gt, gh, gw, L) + + +@torch.no_grad() +def test_tiny_dit_single_chunk_forward_and_unpatchify(): + torch.manual_seed(0) + model = _tiny_dit() + hidden, cond_mask, hdmap, ctx, rope, (gt, gh, gw, L) = _tiny_inputs(model) + out = model( + hidden_states=hidden, + encoder_hidden_states=ctx, + timestep=torch.tensor([500.0]), + condition_video_input_mask=cond_mask, + rope_freqs=rope.shift_t(0), + hdmap_condition=hdmap, + ) + pdim = model.arch.patch_temporal * model.arch.patch_spatial**2 + assert out.shape == (1, L, model.arch.out_channels * pdim) + assert torch.isfinite(out).all() + video = model.unpatchify(out, gt, gh, gw) + assert video.shape == ( + 1, + model.arch.out_channels, + gt * model.arch.patch_temporal, + gh * model.arch.patch_spatial, + gw * model.arch.patch_spatial, + ) + + +@torch.no_grad() +def test_tiny_dit_autoregressive_kv_cache_path(): + torch.manual_seed(0) + model = _tiny_dit() + hidden, cond_mask, hdmap, ctx, rope, (gt, gh, gw, L) = _tiny_inputs(model) + # window holds two chunks so chunk 1's Q (L) attends K/V of length 2L. + caches = model.init_kv_caches( + batch_size=1, + chunk_tokens=L, + window_tokens=2 * L, + sink_tokens=0, + dtype=torch.float32, # match the float32 CPU test model (bf16 in production) + ) + + def run_chunk(idx): + for c in caches: + c.before_update(idx) + out = model( + hidden_states=hidden, + encoder_hidden_states=ctx, + timestep=torch.tensor([500.0]), + condition_video_input_mask=cond_mask, + rope_freqs=rope.shift_t(idx), + hdmap_condition=hdmap, + kv_caches=caches, + ) + for c in caches: + c.after_update(idx) + return out + + out0 = run_chunk(0) + out1 = run_chunk(1) + pdim = model.arch.patch_temporal * model.arch.patch_spatial**2 + expected = (1, L, model.arch.out_channels * pdim) + assert out0.shape == expected and out1.shape == expected + assert torch.isfinite(out0).all() and torch.isfinite(out1).all() + # chunk 1 attends a larger cached window than chunk 0 -> outputs differ. + assert not torch.allclose(out0, out1) + + +# ------------------------------------------------ AR denoising rollout ------ # +def _ar_stage_and_args(arch, dit, scheduler, monkeypatch): + """Build an OmniDreamsDenoisingStage bypassing the heavy base __init__. + + Forces CPU device (production runs uniformly on cuda) and fakes the minimal + server_args.pipeline_config the AR forward reads. + """ + import sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.omnidreams as od_stage + from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.omnidreams import ( # noqa: E501 + OmniDreamsDenoisingStage, + ) + + monkeypatch.setattr(od_stage, "get_local_torch_device", lambda: torch.device("cpu")) + stage = OmniDreamsDenoisingStage.__new__(OmniDreamsDenoisingStage) + stage.transformer = dit + stage.scheduler = scheduler + stage.vae = None + server_args = types.SimpleNamespace( + pipeline_config=types.SimpleNamespace( + dit_precision="fp32", + dit_config=types.SimpleNamespace(arch_config=arch), + ) + ) + return stage, server_args + + +def _ar_batch( + arch, image_token, num_chunks, text, gen, hp=2, wp=2, len_t=2, window_size_t=2 +): + tokens_per_frame = hp * wp + return types.SimpleNamespace( + scheduler=None, + prompt_embeds=[text], + generator=gen, + latents=None, + extra={ + "omnidreams": { + "hp": hp, + "wp": wp, + "len_t": len_t, + "tokens_per_frame": tokens_per_frame, + "chunk_tokens": len_t * tokens_per_frame, + "latent_h": hp * arch.patch_spatial, + "latent_w": wp * arch.patch_spatial, + "num_chunks": num_chunks, + "window_size_t": window_size_t, + "sink_size_t": 0, + "context_noise": 128.0, + "image_token": image_token, + "hdmap_tokens": None, + } + }, + ) + + +@torch.no_grad() +def test_ar_denoising_unconditioned_rollout(monkeypatch): + torch.manual_seed(0) + arch = OmniDreamsDiTArchConfig( + in_channels=4, + out_channels=4, + model_channels=24, + num_blocks=2, + num_heads=2, + mlp_ratio=2.0, + adaln_lora_dim=8, + crossattn_proj_in_channels=32, + crossattn_emb_channels=16, + additional_concat_ch=4, + ) + dit = OmniDreamsDiT(config=OmniDreamsDiTConfig(arch_config=arch), hf_config={}) + dit.post_load_weights() + dit.eval() + sched = OmniDreamsFlowMatchScheduler() + stage, server_args = _ar_stage_and_args(arch, dit, sched, monkeypatch) + + text = torch.randn(1, 5, arch.crossattn_proj_in_channels) + gen = torch.Generator().manual_seed(1) + batch = _ar_batch(arch, image_token=None, num_chunks=3, text=text, gen=gen) + out = stage.forward(batch, server_args) + assert tuple(out.latents.shape) == (1, 4, 3 * 2, 2 * 2, 2 * 2) + assert torch.isfinite(out.latents).all() + + +@torch.no_grad() +def test_ar_denoising_i2v_pins_frame0(monkeypatch): + torch.manual_seed(0) + arch = OmniDreamsDiTArchConfig( + in_channels=4, + out_channels=4, + model_channels=24, + num_blocks=2, + num_heads=2, + mlp_ratio=2.0, + adaln_lora_dim=8, + crossattn_proj_in_channels=32, + crossattn_emb_channels=16, + additional_concat_ch=4, + ) + dit = OmniDreamsDiT(config=OmniDreamsDiTConfig(arch_config=arch), hf_config={}) + dit.post_load_weights() + dit.eval() + sched = OmniDreamsFlowMatchScheduler() + stage, server_args = _ar_stage_and_args(arch, dit, sched, monkeypatch) + + in_d = arch.in_channels * arch.patch_temporal * arch.patch_spatial**2 # 16 + tokens_per_frame = 4 + image_token = torch.randn(1, tokens_per_frame, in_d) + text = torch.randn(1, 5, arch.crossattn_proj_in_channels) + gen = torch.Generator().manual_seed(1) + batch = _ar_batch(arch, image_token=image_token, num_chunks=2, text=text, gen=gen) + out = stage.forward(batch, server_args) + assert tuple(out.latents.shape) == (1, 4, 2 * 2, 2 * 2, 2 * 2) + assert torch.isfinite(out.latents).all() + # chunk-0 frame-0 must equal the (unpatchified) pinned reference latent. + ref_f0 = dit.unpatchify( + torch.cat([image_token, torch.zeros(1, tokens_per_frame, in_d)], dim=1), 2, 2, 2 + )[:, :, 0] + assert torch.allclose(out.latents[:, :, 0], ref_f0, atol=1e-4) + + +@torch.no_grad() +def test_ar_denoising_window_roll_many_chunks(monkeypatch): + torch.manual_seed(0) + arch = OmniDreamsDiTArchConfig( + in_channels=4, + out_channels=4, + model_channels=24, + num_blocks=2, + num_heads=2, + mlp_ratio=2.0, + adaln_lora_dim=8, + crossattn_proj_in_channels=32, + crossattn_emb_channels=16, + additional_concat_ch=4, + ) + dit = OmniDreamsDiT(config=OmniDreamsDiTConfig(arch_config=arch), hf_config={}) + dit.post_load_weights() + dit.eval() + sched = OmniDreamsFlowMatchScheduler() + stage, server_args = _ar_stage_and_args(arch, dit, sched, monkeypatch) + + text = torch.randn(1, 5, arch.crossattn_proj_in_channels) + gen = torch.Generator().manual_seed(2) + # window of 4 latent frames (2 chunks) exercises the steady-state left-roll. + batch = _ar_batch( + arch, image_token=None, num_chunks=4, text=text, gen=gen, window_size_t=4 + ) + out = stage.forward(batch, server_args) + assert tuple(out.latents.shape) == (1, 4, 4 * 2, 2 * 2, 2 * 2) + assert torch.isfinite(out.latents).all() + + +# ----------------------------------------------- A.2 reference preprocess ---- # +_PRE = OmniDreamsBeforeDenoisingStage._preprocess_pixels +_CPU = torch.device("cpu") + + +def test_reference_preprocess_pil_resizes_and_normalizes(): + # off-target-size PIL [0,255] -> resized [1,3,1,H,W] in [-1,1]. + pil = PIL.Image.fromarray((np.random.rand(40, 60, 3) * 255).astype("uint8")) + out = _PRE(pil, height=32, width=48, device=_CPU, dtype=torch.float32) + assert tuple(out.shape) == (1, 3, 1, 32, 48) + assert out.min() >= -1.0 - 1e-4 and out.max() <= 1.0 + 1e-4 + assert out.min() < 0.0 # mapped into the signed VAE input range + + +def test_reference_preprocess_from_path(tmp_path): + pil = PIL.Image.fromarray((np.random.rand(20, 30, 3) * 255).astype("uint8")) + p = tmp_path / "ref.png" + pil.save(p) + out = _PRE(str(p), height=16, width=24, device=_CPU, dtype=torch.float32) + assert tuple(out.shape) == (1, 3, 1, 16, 24) + + +def test_reference_preprocess_signed_tensor_passthrough(): + # already-[-1,1] tensor: no re-normalize, just gains the temporal axis. + t = torch.rand(1, 3, 32, 48) * 2 - 1 + out = _PRE(t, height=32, width=48, device=_CPU, dtype=torch.float32) + assert tuple(out.shape) == (1, 3, 1, 32, 48) + assert torch.allclose(out[:, :, 0], t) + + +def test_reference_preprocess_unsigned_tensor_is_normalized(): + # [0,1] 3D tensor -> [1,3,1,H,W] normalized into [-1,1]. + u = torch.rand(3, 16, 16) + out = _PRE(u, height=16, width=16, device=_CPU, dtype=torch.float32) + assert tuple(out.shape) == (1, 3, 1, 16, 16) + assert out.min() < 0.0 + assert torch.allclose(out[0, :, 0], u * 2 - 1) + + +# ----------------------------------------------- A.1 per-chunk HDMap -------- # +class _RecordingDiT: + """Wraps a DiT, recording the ``hdmap_condition`` of each forward call. + + Delegates every other attribute (``init_kv_caches``/``patchify``/...) to the + wrapped module so the AR stage runs unchanged. + """ + + def __init__(self, dit): + self._dit = dit + self.hdmap_calls: list[torch.Tensor] = [] + + def __call__(self, *args, **kwargs): + self.hdmap_calls.append(kwargs.get("hdmap_condition")) + return self._dit(*args, **kwargs) + + def __getattr__(self, name): + return getattr(self._dit, name) + + +def _ar_setup(monkeypatch, num_chunks, window_size_t=4): + arch = OmniDreamsDiTArchConfig( + in_channels=4, + out_channels=4, + model_channels=24, + num_blocks=2, + num_heads=2, + mlp_ratio=2.0, + adaln_lora_dim=8, + crossattn_proj_in_channels=32, + crossattn_emb_channels=16, + additional_concat_ch=4, + ) + dit = OmniDreamsDiT(config=OmniDreamsDiTConfig(arch_config=arch), hf_config={}) + dit.post_load_weights() + dit.eval() + rec = _RecordingDiT(dit) + sched = OmniDreamsFlowMatchScheduler() + stage, server_args = _ar_stage_and_args(arch, rec, sched, monkeypatch) + text = torch.randn(1, 5, arch.crossattn_proj_in_channels) + gen = torch.Generator().manual_seed(3) + batch = _ar_batch( + arch, + image_token=None, + num_chunks=num_chunks, + text=text, + gen=gen, + window_size_t=window_size_t, + ) + return arch, stage, server_args, rec, batch + + +@torch.no_grad() +def test_ar_hdmap_per_chunk_indexing(monkeypatch): + torch.manual_seed(0) + num_chunks = 3 + arch, stage, server_args, rec, batch = _ar_setup(monkeypatch, num_chunks) + + pdim = arch.patch_temporal * arch.patch_spatial**2 + hdmap_d = arch.additional_concat_ch * pdim + chunk_tokens = batch.extra["omnidreams"]["chunk_tokens"] + # Each chunk's HD-map is a constant tensor tagged with its chunk index. + hdmap_tokens = [ + torch.full((1, chunk_tokens, hdmap_d), float(ci)) for ci in range(num_chunks) + ] + batch.extra["omnidreams"]["hdmap_tokens"] = hdmap_tokens + + stage.forward(batch, server_args) + + # Each chunk is tagged with its index; verify routing without hard-coding the + # per-chunk call count (predict_flow x steps + 1 cache-write). + assert rec.hdmap_calls and all(h is not None for h in rec.hdmap_calls) + tags = [float(h.flatten()[0]) for h in rec.hdmap_calls] + counts = Counter(tags) + # distinct tags == {0..n-1}: each chunk received its OWN tensor (not shared); + assert sorted(counts) == [float(ci) for ci in range(num_chunks)] + # uniform #calls per chunk + non-decreasing order -> each chunk's calls are + # contiguous and never cross-contaminated with another chunk's tensor. + assert len(set(counts.values())) == 1 + assert tags == sorted(tags) + + +@torch.no_grad() +def test_ar_hdmap_none_falls_back_to_zeros(monkeypatch): + torch.manual_seed(0) + arch, stage, server_args, rec, batch = _ar_setup(monkeypatch, num_chunks=2) + assert batch.extra["omnidreams"]["hdmap_tokens"] is None + stage.forward(batch, server_args) + # Disabled HDMap -> every forward gets an all-zero condition. + assert rec.hdmap_calls + assert all(bool((h == 0).all()) for h in rec.hdmap_calls) + + +def test_encode_hdmap_broadcast_clamp_and_none(monkeypatch): + """`_encode_hdmap` assembles the per-chunk list (broadcast / clamp / None). + + Stubs preprocess+encode+patchify so the test isolates the per-chunk control + flow (the VAE numerics are a GPU concern). Each src tag flows through as a + 1-element tensor so chunk identity is traceable. + """ + stage = OmniDreamsBeforeDenoisingStage.__new__(OmniDreamsBeforeDenoisingStage) + stage.transformer = types.SimpleNamespace(patchify=lambda latent: latent) + monkeypatch.setattr( + stage, + "_preprocess_pixels", + lambda src, h, w, d, dt: torch.tensor([float(src)]), + ) + monkeypatch.setattr(stage, "_vae_encode_normalized", lambda x: x) + dev = torch.device("cpu") + + # Single (non-list) input broadcasts to every chunk. + b1 = types.SimpleNamespace(hdmap_path=7, hdmap_pixels=None) + toks = stage._encode_hdmap(b1, dev, torch.float32, torch.float32, 3, 16, 16) + assert [float(t) for t in toks] == [7.0, 7.0, 7.0] + + # Per-chunk list shorter than num_chunks clamps to the last entry. + b2 = types.SimpleNamespace(hdmap_path=[1, 2], hdmap_pixels=None) + toks2 = stage._encode_hdmap(b2, dev, torch.float32, torch.float32, 4, 16, 16) + assert [float(t) for t in toks2] == [1.0, 2.0, 2.0, 2.0] + + # No HD-map input -> None (AR stage falls back to zeros). + b3 = types.SimpleNamespace(hdmap_path=None, hdmap_pixels=None) + assert stage._encode_hdmap(b3, dev, torch.float32, torch.float32, 2, 16, 16) is None diff --git a/python/sglang/multimodal_gen/test/unit/test_omnidreams_regression.py b/python/sglang/multimodal_gen/test/unit/test_omnidreams_regression.py new file mode 100644 index 000000000000..045680d0b305 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_omnidreams_regression.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU regression tests pinning fixes made during OmniDreams GPU bring-up. + +Each test guards a specific defect found while validating OmniDreams end-to-end +on GPU (no checkpoint, no GPU required): + +1. Meta-init load materializes the non-persistent sinusoidal ``emb`` buffer + (otherwise the production meta->load path leaves it on the meta device and + ``_load_dit_model`` raises / ``.to()`` fails). +2. ``TimestepEmbedding`` casts the float32 sinusoid to the MLP param dtype + (otherwise a bf16 model hits a float!=bf16 matmul error). +3. The registry resolves a non-diffusers local checkpoint via a path detector, + and the path-detector short-circuit is gated to dirs WITHOUT model_index.json + (so it cannot hijack a normal diffusers model). +4. ``_compute_num_chunks`` maps ``num_frames`` -> chunk count and caps the AR + rollout length (``_MAX_AR_CHUNKS``) against unbounded requests. +5. ``apply_chat_template`` BatchEncoding output is normalized to the input_ids + tensor (newer transformers return a dict-like, not a bare tensor). +6. ``_read_vae_state_dict`` reads diffusers safetensors; ``_load_wan_vae`` raises + a helpful error for a non-diffusers (original-Wan) state dict. +7. ``_resolve_text_encoder_src`` prefers a local ``text_encoder`` dir, else the + pinned HF id + revision. +""" + +import json +import types +from itertools import chain + +import pytest +import torch + +from sglang.multimodal_gen.configs.models.dits.omnidreams import ( + OmniDreamsDiTArchConfig, + OmniDreamsDiTConfig, +) +from sglang.multimodal_gen.runtime.models.dits.omnidreams import ( + OmniDreamsDiT, + TimestepEmbedding, + Timesteps, +) +from sglang.multimodal_gen.runtime.models.encoders.omnidreams_text import ( + COSMOS_REASON1_HIDDEN, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.omnidreams import ( # noqa: E501 + _MAX_AR_CHUNKS, + _TEXT_MAX_LENGTH, + OmniDreamsBeforeDenoisingStage, +) + + +def _tiny_arch() -> OmniDreamsDiTArchConfig: + return OmniDreamsDiTArchConfig( + in_channels=4, + out_channels=4, + model_channels=24, + num_blocks=2, + num_heads=2, + mlp_ratio=2.0, + adaln_lora_dim=8, + crossattn_proj_in_channels=32, + crossattn_emb_channels=16, + additional_concat_ch=4, + ) + + +# ---- 1. meta-init buffer materialization ----------------------------------- # +def test_meta_init_materializes_nonpersistent_buffers(): + with torch.device("meta"): + model = OmniDreamsDiT( + config=OmniDreamsDiTConfig(arch_config=_tiny_arch()), hf_config={} + ) + # Simulate the production load path: materialize params/buffers on a real + # device (mirrors the FSDP loader), then run the post-load hook. + model.to_empty(device="cpu") + model.post_load_weights() + + on_meta = [ + n + for n, p in chain(model.named_parameters(), model.named_buffers()) + if p.is_meta + ] + assert not on_meta, f"params/buffers left on meta: {on_meta}" + + ts = model.t_embedder[0] + assert isinstance(ts, Timesteps) + assert ts.emb.device.type == "cpu" + assert torch.isfinite(ts.emb).all() + assert ts.emb.shape == (ts.num_channels // 2,) + + +# ---- 2. TimestepEmbedding dtype cast --------------------------------------- # +def test_timestep_embedding_casts_sinusoid_to_param_dtype(): + te = TimestepEmbedding(16, 16, use_adaln_lora=True).to(torch.bfloat16) + sinusoid_fp32 = torch.randn(16, dtype=torch.float32) + raw, lora = te(sinusoid_fp32) # must not raise float != bf16 + assert raw.dtype == torch.bfloat16 # raw embedding cast for RMSNorm/AdaLN + assert lora.dtype == torch.bfloat16 + assert torch.isfinite(raw).all() and torch.isfinite(lora).all() + + +# ---- 3. registry resolution + gated short-circuit -------------------------- # +def test_registry_resolves_nondiffusers_local_omnidreams(tmp_path): + import sglang.multimodal_gen.registry as reg + + local = tmp_path / "omni-dreams" + local.mkdir() # non-diffusers: no model_index.json + + reg._get_config_info.cache_clear() + info = reg._get_config_info(str(local)) + assert info is not None + assert info.pipeline_config_cls.__name__ == "OmniDreamsPipelineConfig" + + +def test_registry_path_detector_gated_to_non_diffusers(tmp_path): + import sglang.multimodal_gen.registry as reg + + # A diffusers-style dir (has model_index.json) with a neutral name and an + # unknown _class_name must NOT be short-circuited by the path detectors; + # it falls through to model_index resolution and returns None (no match), + # proving step 3a did not fire for a model_index.json dir. + d = tmp_path / "plain-model" + d.mkdir() + (d / "model_index.json").write_text( + json.dumps({"_class_name": "ZzzUnknownPipeline"}) + ) + + reg._get_config_info.cache_clear() + info = reg._get_config_info(str(d)) + assert info is None + + +# ---- 4. num_frames -> chunk mapping + AR cap ------------------------------- # +@pytest.mark.parametrize( + "num_frames,expected", + [(5, 1), (6, 2), (13, 2), (14, 3), (21, 3)], # len_t=2: first=5, step=8 +) +def test_compute_num_chunks_boundaries(num_frames, expected): + batch = types.SimpleNamespace(num_frames=num_frames) + assert ( + OmniDreamsBeforeDenoisingStage._compute_num_chunks(batch, len_t=2) == expected + ) + + +def test_compute_num_chunks_caps_ar_loop(): + batch = types.SimpleNamespace(num_frames=10_000_000) + assert ( + OmniDreamsBeforeDenoisingStage._compute_num_chunks(batch, len_t=2) + == _MAX_AR_CHUNKS + ) + + +# ---- 5. tokenizer BatchEncoding normalization ------------------------------ # +def test_encode_text_normalizes_batchencoding(): + n_layers = 3 # tiny stand-in for the 28 transformer layers + hidden = COSMOS_REASON1_HIDDEN + + class _DictTokenizer: + pad_token_id = 0 + + def apply_chat_template(self, messages, **kwargs): + # Newer transformers return a BatchEncoding (dict-like), not a tensor. + return {"input_ids": torch.zeros(1, 10, dtype=torch.long)} + + def _text_encoder( + input_ids=None, + attention_mask=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + # 1 embedding layer + n_layers transformer layers, each [B, L, H]. + L = input_ids.shape[1] + hs = [torch.randn(1, L, hidden) for _ in range(n_layers + 1)] + return types.SimpleNamespace(hidden_states=hs) + + stage = OmniDreamsBeforeDenoisingStage.__new__(OmniDreamsBeforeDenoisingStage) + stage.tokenizer = _DictTokenizer() + stage.text_encoder = _text_encoder + + out = stage._encode_text("a prompt", torch.device("cpu")) + assert out.shape == (1, _TEXT_MAX_LENGTH, n_layers * hidden) + assert torch.isfinite(out).all() + + +# ---- 6. VAE state-dict reader + helpful error ------------------------------ # +def test_read_vae_state_dict_safetensors_file_and_dir(tmp_path): + from safetensors.torch import save_file + + from sglang.multimodal_gen.runtime.pipelines.omnidreams_pipeline import ( + OmniDreamsPipeline, + ) + + tensors = {"w": torch.zeros(2, 3)} + f = tmp_path / "vae.safetensors" + save_file(tensors, str(f)) + + sd_file = OmniDreamsPipeline._read_vae_state_dict(str(f)) + assert set(sd_file.keys()) == {"w"} + + d = tmp_path / "vae" + d.mkdir() + save_file(tensors, str(d / "diffusion_pytorch_model.safetensors")) + sd_dir = OmniDreamsPipeline._read_vae_state_dict(str(d)) + assert set(sd_dir.keys()) == {"w"} + + empty = tmp_path / "empty" + empty.mkdir() + with pytest.raises(FileNotFoundError): + OmniDreamsPipeline._read_vae_state_dict(str(empty)) + + +def test_load_wan_vae_raises_helpful_error_on_key_mismatch(tmp_path, monkeypatch): + from safetensors.torch import save_file + + import sglang.multimodal_gen.runtime.pipelines.omnidreams_pipeline as pipe_mod + from sglang.multimodal_gen.runtime.pipelines.omnidreams_pipeline import ( + OmniDreamsPipeline, + ) + + class _FakeVAE: + def __init__(self, config): + pass + + def load_state_dict(self, state, strict=True): + raise RuntimeError("Missing key(s) in state_dict: ...") + + def to(self, *a, **k): + return self + + def eval(self): + return self + + monkeypatch.setattr(pipe_mod, "AutoencoderKLWan", _FakeVAE) + + f = tmp_path / "vae.safetensors" + save_file({"original_wan_key": torch.zeros(1)}, str(f)) + + with pytest.raises(RuntimeError, match="diffusers format"): + OmniDreamsPipeline._load_wan_vae( + object(), str(f), torch.device("cpu"), torch.float32 + ) + + +# ---- 7. text-encoder source resolution ------------------------------------- # +def test_resolve_text_encoder_src_prefers_local_then_hf(tmp_path): + from sglang.multimodal_gen.runtime.pipelines.omnidreams_pipeline import ( + _TEXT_ENCODER_ID, + _TEXT_ENCODER_REVISION, + OmniDreamsPipeline, + ) + + # No local text_encoder dir -> falls back to the pinned HF id + revision. + src, rev = OmniDreamsPipeline._resolve_text_encoder_src(str(tmp_path)) + assert src == _TEXT_ENCODER_ID and rev == _TEXT_ENCODER_REVISION + + # Local text_encoder/config.json present -> use the local dir, no revision. + te = tmp_path / "text_encoder" + te.mkdir() + (te / "config.json").write_text("{}") + src2, rev2 = OmniDreamsPipeline._resolve_text_encoder_src(str(tmp_path)) + assert src2 == str(te) and rev2 is None diff --git a/python/sglang/multimodal_gen/test/unit/test_omnidreams_scaffold.py b/python/sglang/multimodal_gen/test/unit/test_omnidreams_scaffold.py new file mode 100644 index 000000000000..4837522f88cb --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_omnidreams_scaffold.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Phase-0 scaffold checks for the OmniDreams DiT (CPU-only, no checkpoint). + +Constructs the DiT on the meta device (no memory for ~2B params) and validates +the checkpoint-exact structure against an independently-derived authoritative +key fixture, plus the pre/post-fusion shapes, the 2-step flow-match sigmas, and +the 3D-RoPE layout. +""" + +import os + +import torch + +from sglang.multimodal_gen.configs.models.dits.omnidreams import OmniDreamsDiTConfig +from sglang.multimodal_gen.configs.pipeline_configs.omnidreams import ( + OmniDreamsPipelineConfig, + warp_flow_match_sigmas, +) +from sglang.multimodal_gen.runtime.models.dits.omnidreams import ( + ROPE_IS_NEOX_STYLE, + OmniDreamsDiT, + rope_dims, +) + +_KEY_FIXTURE = os.path.join( + os.path.dirname(__file__), "data", "omnidreams_dit_keys.txt" +) + + +def _load_fixture_keys() -> set[str]: + with open(_KEY_FIXTURE) as f: + return {line.strip() for line in f if line.strip()} + + +def _build_meta_model() -> OmniDreamsDiT: + with torch.device("meta"): + return OmniDreamsDiT(config=OmniDreamsDiTConfig(), hf_config={}) + + +def test_state_dict_matches_authoritative_key_fixture(): + model = _build_meta_model() + keys = set(model.state_dict().keys()) + expected = _load_fixture_keys() + assert len(expected) == 570 + assert ( + keys == expected + ), f"missing={sorted(expected - keys)} extra={sorted(keys - expected)}" + + +def test_unique_bias_is_crossattn_proj(): + model = _build_meta_model() + biases = [k for k in model.state_dict() if k.endswith(".bias")] + assert biases == ["crossattn_proj.0.bias"] + + +def test_pre_fusion_shapes(): + model = _build_meta_model() + sd = model.state_dict() + # x_embedder keeps the padding-mask channel pre-fusion: (16 + 1 + 1) * 2 * 2 = 72. + assert tuple(sd["x_embedder.proj.1.weight"].shape) == (2048, 72) + # HDMap embed: 16 * 2 * 2 = 64 in-features. + assert tuple(sd["additional_patch_embedding.proj.1.weight"].shape) == (2048, 64) + # Final layer pre-shuffle: patch_dim = 2*2*1*16 = 64. + assert tuple(sd["final_layer.linear.weight"].shape) == (64, 2048) + assert tuple(sd["crossattn_proj.0.weight"].shape) == (1024, 100352) + + +def test_post_load_weights_fuses_in_place(): + model = _build_meta_model() + pre_keys = set(model.state_dict().keys()) + model.post_load_weights() + sd = model.state_dict() + # Padding-mask channels dropped: 72 -> 68. + assert tuple(sd["x_embedder.proj.1.weight"].shape) == (2048, 68) + # Shuffle fuse is a reorder; shape is preserved. + assert tuple(sd["final_layer.linear.weight"].shape) == (64, 2048) + # Fusion must not add or remove parameters. + assert set(sd.keys()) == pre_keys + assert model._is_padding_mask_fused and model._is_shuffle_op_fused + + +def test_two_step_flow_match_sigmas(): + sigmas = warp_flow_match_sigmas() + assert len(sigmas) == 3 + assert abs(sigmas[0] - 1.0) < 1e-9 + assert abs(sigmas[1] - 0.8036) < 1e-3 + assert sigmas[2] == 0.0 + # The pipeline config exposes the same schedule. + assert OmniDreamsPipelineConfig().denoising_sigmas() == sigmas + + +def test_rope_layout_neox_44_42_42(): + assert rope_dims(128) == (44, 42, 42) + assert sum(rope_dims(128)) == 128 + assert ROPE_IS_NEOX_STYLE is True diff --git a/python/sglang/utils.py b/python/sglang/utils.py index c7dc1ae7822a..55303793b7d1 100644 --- a/python/sglang/utils.py +++ b/python/sglang/utils.py @@ -34,6 +34,8 @@ KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS: dict[str, str] = { "hunyuan3d": "Hunyuan3D2Pipeline", "flux.2-dev-nvfp4": "Flux2NvfpPipeline", + "omnidreams": "OmniDreamsPipeline", + "omni-dreams": "OmniDreamsPipeline", "comfy-org/ideogram-4": "Ideogram4Nvfp4Pipeline", "comfy-org--ideogram-4": "Ideogram4Nvfp4Pipeline", } From 1143fd2caa4b85b50b41e58f5b0ee59ee70d01c1 Mon Sep 17 00:00:00 2001 From: Cerdore Date: Wed, 10 Jun 2026 02:01:50 +0800 Subject: [PATCH 2/5] fix(diffusion): prevent OOM when OmniDreams loads with CPU offload --- .../runtime/pipelines/omnidreams_pipeline.py | 26 +++++- .../model_specific_stages/omnidreams.py | 79 +++++++++++++++++-- 2 files changed, 95 insertions(+), 10 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/omnidreams_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/omnidreams_pipeline.py index d917799d3a88..3f8ec65cb70e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/omnidreams_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/omnidreams_pipeline.py @@ -282,20 +282,40 @@ def load_modules( dit_dtype = PRECISION_TO_TYPE[pipeline_config.dit_precision] vae_dtype = PRECISION_TO_TYPE[pipeline_config.vae_precision] + # Honor the CPU-offload flags at load time: each flagged component is + # staged on CPU so the three heavy weights (the 2B DiT, the 7B text + # encoder, and the Wan VAE) never need to co-reside on the GPU while + # loading. The ComponentResidencyManager then brings each one to the GPU + # only around its use-site (declared in the stages' ``component_uses``) + # and releases it afterwards. Without this the custom loaders push all + # three straight to the GPU, so a small-VRAM card OOMs while loading the + # text encoder even though --text-encoder-cpu-offload was requested. + cpu_device = torch.device("cpu") + + def _load_device(offload_flag: object) -> torch.device: + offload = bool(offload_flag) and not server_args.use_fsdp_inference + return cpu_device if offload else device + + dit_device = _load_device(server_args.dit_cpu_offload) + vae_device = _load_device(server_args.vae_cpu_offload) + text_encoder_device = _load_device(server_args.text_encoder_cpu_offload) + model_path = server_args.model_path ckpt_path = self._resolve_ckpt_path(model_path) logger.info("OmniDreams: loading flat DiT from %s", ckpt_path) transformer = self._load_flat_dit( - pipeline_config.dit_config, ckpt_path, device, dit_dtype + pipeline_config.dit_config, ckpt_path, dit_device, dit_dtype ) vae_path = self._resolve_vae_path(model_path) logger.info("OmniDreams: loading Wan 2.1 VAE from %s", vae_path) vae = self._load_wan_vae( - pipeline_config.vae_config, vae_path, device, vae_dtype + pipeline_config.vae_config, vae_path, vae_device, vae_dtype ) - text_encoder, tokenizer = self._load_text_encoder(model_path, device) + text_encoder, tokenizer = self._load_text_encoder( + model_path, text_encoder_device + ) scheduler = OmniDreamsFlowMatchScheduler( num_inference_steps=len(pipeline_config.denoising_timesteps), diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py index 0a83412d8011..18e427eedab4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py @@ -33,6 +33,9 @@ from sglang.multimodal_gen.runtime.distributed import ( get_local_torch_device, ) +from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( + ComponentUse, +) from sglang.multimodal_gen.runtime.models.dits.omnidreams_rope import ( RotaryPositionEmbedding3D, ) @@ -99,6 +102,26 @@ def __init__( self.vae = vae self.config = config + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + """Declare the text encoder + VAE so the residency manager can stage + them on the GPU only around their use-sites (and offload them again + afterwards) when ``--text-encoder-cpu-offload`` / ``--vae-cpu-offload`` + are set. The DiT is only used here via the weightless ``patchify`` + rearrange, so it is not declared (its weights stay wherever loaded). + """ + stage_name = self._component_stage_name(stage_name) + vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] + return [ + ComponentUse(stage_name=stage_name, component_name="text_encoder"), + ComponentUse( + stage_name=stage_name, + component_name="vae", + target_dtype=vae_dtype, + ), + ] + # ----- helpers ---------------------------------------------------------- # @torch.no_grad() def _encode_text(self, prompt: str, device: torch.device) -> torch.Tensor: @@ -331,16 +354,20 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: # --- text conditioning (100352) --- prompt = batch.prompt if isinstance(batch.prompt, str) else str(batch.prompt) - text_embeds = self._encode_text(prompt, device).to(dit_dtype) + with self.use_declared_component( + component_name="text_encoder", module=self.text_encoder + ): + text_embeds = self._encode_text(prompt, device).to(dit_dtype) batch.prompt_embeds = [text_embeds] batch.negative_prompt_embeds = None batch.image_embeds = [] batch.do_classifier_free_guidance = False # --- i2v reference latent -> patchified frame-0 token block --- - image_latent = self._encode_reference_image( - batch, device, vae_dtype, height, width - ) + with self.use_declared_component(component_name="vae", module=self.vae): + image_latent = self._encode_reference_image( + batch, device, vae_dtype, height, width + ) if image_latent is not None: image_latent = image_latent.to(dit_dtype) # [B,16,1,h,w] -> [B, hp*wp, 16*pdim] via the DiT patchify. @@ -364,9 +391,10 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: # --- AR rollout state for the denoising stage --- num_chunks = self._compute_num_chunks(batch, len_t) # Per-chunk HD-map tokens (None -> AR stage uses zeros / HDMap disabled). - hdmap_tokens = self._encode_hdmap( - batch, device, vae_dtype, dit_dtype, num_chunks, height, width - ) + with self.use_declared_component(component_name="vae", module=self.vae): + hdmap_tokens = self._encode_hdmap( + batch, device, vae_dtype, dit_dtype, num_chunks, height, width + ) batch.extra["omnidreams"] = { "hp": hp, "wp": wp, @@ -444,6 +472,27 @@ class OmniDreamsDenoisingStage(DenoisingStage): def __init__(self, transformer, scheduler, vae=None) -> None: super().__init__(transformer, scheduler, vae=vae) + def component_uses( + self, server_args: ServerArgs, stage_name: str | None = None + ) -> list[ComponentUse]: + """Declare only the DiT for residency scheduling. The base + ``DenoisingStage`` would also declare the VAE (this stage receives one + to assert ``use_feature_cache``), but the AR rollout never runs the VAE + here, so declaring it would needlessly hold it on the GPU through the + denoise loop. The VAE encode/decode use-sites live in the before- and + decoding-stages instead. + """ + stage_name = self._component_stage_name(stage_name) + return [ + ComponentUse( + stage_name=stage_name, + component_name="transformer", + phase="transformer", + preferred_ready_after_request=True, + memory_intensive=True, + ) + ] + @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> Req: # Phase 6 guard: TP is supported via column/row parallel layers in the @@ -474,6 +523,19 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: "use_feature_cache=True for correct streaming frame counts." ) + # Bring the DiT onto the GPU for the whole AR rollout when it is being + # CPU-offloaded (no-op when it is already resident). The manager keeps a + # single DiT resident afterwards; on the error path the request-level + # finish_request still releases it, so an explicit try/finally is not + # needed here. + residency_manager = self._component_residency_manager + transformer_use = None + if residency_manager is not None: + transformer_use = self._declared_component_use( + component_name="transformer" + ) + residency_manager.begin_use(transformer_use, self.transformer) + config = server_args.pipeline_config device = get_local_torch_device() dit_dtype = PRECISION_TO_TYPE[config.dit_precision] @@ -659,6 +721,9 @@ def predict_flow(noisy: torch.Tensor, t: torch.Tensor) -> torch.Tensor: # eventually supported. Currently a no-op (SP is guarded at entry). batch.latents = self._postprocess_sp_latents(batch, server_args) + if residency_manager is not None and transformer_use is not None: + residency_manager.end_use(transformer_use, self.transformer) + return batch def _postprocess_sp_latents( From 28972969caa671d1e76551c4b5ba7a3f3825cd0a Mon Sep 17 00:00:00 2001 From: Cerdore Date: Wed, 10 Jun 2026 02:07:59 +0800 Subject: [PATCH 3/5] chore: apply black formatting --- .../pipelines_core/stages/model_specific_stages/omnidreams.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py index 18e427eedab4..c172200ab080 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py @@ -531,9 +531,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: residency_manager = self._component_residency_manager transformer_use = None if residency_manager is not None: - transformer_use = self._declared_component_use( - component_name="transformer" - ) + transformer_use = self._declared_component_use(component_name="transformer") residency_manager.begin_use(transformer_use, self.transformer) config = server_args.pipeline_config From 227575a77ca5c7d125032577b527937fde747334 Mon Sep 17 00:00:00 2001 From: Cerdore Date: Sat, 13 Jun 2026 14:05:40 +0800 Subject: [PATCH 4/5] refactor(diffusion): complete OmniDreams AR stage implementation with HD-map conditioning Consolidate pre-processing (text encode, i2v VAE, HD-map per-chunk encoding) into OmniDreamsBeforeDenoisingStage. Implement full autoregressive rollout in OmniDreamsDenoisingStage: per-block BlockKVCache lifecycle, 3D NeoX RoPE shift_t, 2-step self-forcing denoise (sample + context-noise re-forward), frame-0 i2v pinning, and per-chunk HD-map indexing. Add comprehensive unit tests covering RoPE correctness, KV-cache fill/roll/steady-state, scheduler 2-step sigmas, Cosmos-Reason1 100352-dim text embedding, tiny DiT end-to-end forward, AR unconditioned/i2v/multi-chunk rollout, HD-map per-frame slicing, and reference image preprocessing. Constraint: 2-step distilled model trades temporal smoothness for real-time speed; blur is expected behavior from Self-Forcing + DMD distillation, not a bug. Confidence: high Scope-risk: moderate (AR loop is new; GPU validation still needed) Not-tested: end-to-end generation on full OmniDreams 2B checkpoint; HD-map VAE numerics on GPU (flagged with TODO) Co-Authored-By: Claude --- .../configs/sample/omnidreams.py | 6 +- .../configs/sample/sampling_params.py | 9 +- .../model_specific_stages/omnidreams.py | 160 ++++++++++++++--- .../test/unit/test_omnidreams_components.py | 170 ++++++++++++++++-- 4 files changed, 293 insertions(+), 52 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/sample/omnidreams.py b/python/sglang/multimodal_gen/configs/sample/omnidreams.py index 87dcae64cd22..a5a79946c181 100644 --- a/python/sglang/multimodal_gen/configs/sample/omnidreams.py +++ b/python/sglang/multimodal_gen/configs/sample/omnidreams.py @@ -33,6 +33,8 @@ class OmniDreamsSamplingParams(SamplingParams): # Raw timestep injected as context noise on cached/clean frames. context_noise: int = 128 - # HD-map / trajectory conditioning input (Phase 4). A single path (broadcast - # to every AR chunk) or a per-chunk list of paths. ``None`` disables HDMap. + # HD-map conditioning input -- OmniDreams' central per-frame control signal. + # Accepts a video path (``.mp4``/``.gif``/...; decoded to per-frame rasters), + # a per-frame list of image paths, or -- degenerate fallback -- a single image + # broadcast across every frame (no temporal motion). ``None`` disables HDMap. hdmap_path: str | list[str] | None = None diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 7f1ba1930a4e..78e855dc6015 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -952,10 +952,11 @@ def add_argument(*name_or_flags, **kwargs): type=str, nargs="+", help=( - "Path(s) to HD-map / trajectory conditioning input(s) for " - "OmniDreams autoregressive video generation. Pass a single path " - "(broadcast to every chunk) or one per autoregressive chunk, e.g.: " - '--hdmap-path "chunk0.png" "chunk1.png"' + "HD-map conditioning input for OmniDreams autoregressive video " + "generation (the central per-frame control signal). Pass a video " + "path (decoded to per-frame rasters), a per-frame list of image " + "paths, or a single image (broadcast fallback, no motion), e.g.: " + '--hdmap-path scene_hdmap.mp4 OR --hdmap-path f0.png f1.png ...' ), ) add_argument( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py index c172200ab080..b1bb48ac3643 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py @@ -44,6 +44,7 @@ ) from sglang.multimodal_gen.runtime.models.vision_utils import ( load_image, + load_video, normalize, numpy_to_pt, pil_to_numpy, @@ -73,6 +74,9 @@ # length (and thus GPU memory/compute) against an unbounded ``num_frames`` from # the HTTP API. ~256 chunks * len_t(2) * 4 = ~2048 pixel frames. _MAX_AR_CHUNKS = 256 +# HD-map inputs ending in one of these are decoded as a per-frame raster video; +# any other single string is treated as one image (degenerate broadcast). +_HDMAP_VIDEO_EXTS = (".mp4", ".gif", ".webm", ".mov", ".mkv", ".avi") # --------------------------------------------------------------------------- # @@ -233,6 +237,30 @@ def _preprocess_pixels( x = x.unsqueeze(2) return x + def _preprocess_hdmap_clip( + self, + frames: list, + height: int, + width: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor | None: + """Per-frame HD-map rasters -> a single clip ``[B, 3, T, H, W]`` in ``[-1, 1]``. + + Each frame runs the same single-frame ``_preprocess_pixels`` path + (PIL/path/tensor -> resize -> normalize) and the results are stacked on + the temporal axis. Returns ``None`` if any frame fails to preprocess. + """ + per_frame: list[torch.Tensor] = [] + for f in frames: + x = self._preprocess_pixels(f, height, width, device, dtype) + if x is None: + return None + per_frame.append(x) # [B,3,1,H,W] + if not per_frame: + return None + return torch.cat(per_frame, dim=2) # [B,3,T,H,W] + @torch.no_grad() def _encode_reference_image( self, @@ -281,6 +309,34 @@ def _vae_encode_normalized(self, x: torch.Tensor) -> torch.Tensor: ).view(1, -1, 1, 1, 1) return (latent - mean) / std + @staticmethod + def _resolve_hdmap_frames(hdmap) -> list | None: + """Resolve an HD-map input to a per-frame raster list, or ``None`` for the + degenerate single-image case (caller broadcasts one raster). + + * video path (``_HDMAP_VIDEO_EXTS``) -> decoded to per-frame PIL images; + * ``list`` / ``tuple`` -> per-frame sequence as-is, **except** a single + video-path element is expanded via ``load_video`` (CLI ``--hdmap-path`` + always passes a list, even for a single arg); + * single image (path / PIL / tensor) -> ``None``. + """ + if isinstance(hdmap, (list, tuple)): + items = list(hdmap) + if len(items) == 0: + return None + # CLI passes --hdmap-path as a list even for a single arg; expand a + # lone video-path element via load_video so it becomes per-frame rasters. + if ( + len(items) == 1 + and isinstance(items[0], str) + and items[0].lower().endswith(_HDMAP_VIDEO_EXTS) + ): + return load_video(items[0]) + return items + if isinstance(hdmap, str) and hdmap.lower().endswith(_HDMAP_VIDEO_EXTS): + return load_video(hdmap) + return None + @torch.no_grad() def _encode_hdmap( self, @@ -289,22 +345,28 @@ def _encode_hdmap( vae_dtype: torch.dtype, dit_dtype: torch.dtype, num_chunks: int, + len_t: int, height: int, width: int, ) -> list[torch.Tensor] | None: - """Per-chunk HD-map conditioning -> ``list[num_chunks]`` of patchified tokens. - - HD-map is a *per-chunk* driving condition (each latent chunk has its own - HD-map frames), so this returns a list indexed by chunk -- not one tensor - shared across chunks. Returns ``None`` when the request carries no HD-map - input, in which case the AR stage falls back to zeros (HDMap disabled). - - Accepts ``batch.hdmap_path`` / ``batch.hdmap_pixels`` as either a single - input (broadcast to every chunk) or a per-chunk list. Each entry runs the - same pixel-preprocess + VAE-encode + patchify path as the reference image. - - TODO(gpu): the HD-map pixel -> 16ch-latent VAE numerics are validated on - GPU; this encode path only runs when real HD-map input is supplied. + """Per-frame HD-map conditioning -> ``list[num_chunks]`` of patchified tokens. + + HD-map is OmniDreams' central per-frame control signal (lane lines + actor + boxes rendered at the ego pose); the generated viewpoint changes *because* + the raster shifts frame-to-frame. The full per-frame raster sequence is + VAE-encoded as one causal clip -- matching the output latent temporal + layout -- then sliced into ``num_chunks`` groups of ``len_t`` latent + frames. Returns ``None`` when the request carries no HD-map input, in + which case the AR stage falls back to zeros (HDMap disabled). + + Accepts ``batch.hdmap_path`` / ``batch.hdmap_pixels`` as a video path + (decoded per-frame), a per-frame list of rasters, or -- as a degenerate + back-compat / smoke fallback -- a single image broadcast across every + latent frame (no temporal motion). + + TODO(gpu): the HD-map pixel -> 16ch-latent VAE numerics and the causal + multi-frame temporal compression are validated on GPU; this encode path + only runs when real HD-map input is supplied. """ hdmap = getattr(batch, "hdmap_path", None) if hdmap is None: @@ -312,24 +374,66 @@ def _encode_hdmap( if hdmap is None: return None - per_chunk = list(hdmap) if isinstance(hdmap, (list, tuple)) else None - tokens: list[torch.Tensor] = [] - for ci in range(num_chunks): - if per_chunk is not None: - src = per_chunk[ci] if ci < len(per_chunk) else per_chunk[-1] - else: - src = hdmap - x = self._preprocess_pixels(src, height, width, device, vae_dtype) + # L latent frames total -> 1 + (L-1)*4 pixel frames (causal VAE, tc=4), + # matching the output chunk math (chunk0=1+(len_t-1)*4, later=len_t*4). + num_latent = num_chunks * len_t + total_pixel = 1 + (num_latent - 1) * 4 + + frames = self._resolve_hdmap_frames(hdmap) + if frames is None: + # Degenerate single-image fallback: one raster broadcast across all + # latent frames (no temporal motion -- back-compat / smoke only). + x = self._preprocess_pixels(hdmap, height, width, device, vae_dtype) if x is None: logger.warning( - "OmniDreams: HD-map chunk %d preprocessed to None; disabling " - "HDMap (all chunks fall back to zeros). Check hdmap input.", - ci, + "OmniDreams: HD-map preprocessed to None; disabling HDMap " + "(all chunks fall back to zeros). Check hdmap input." + ) + return None + latent = self._vae_encode_normalized(x).to(dit_dtype) # [B,16,1,h,w] + if num_latent > 1 and latent.ndim == 5 and latent.shape[2] == 1: + latent = latent.repeat(1, 1, num_latent, 1, 1) + else: + # Per-frame path: clamp/truncate to total_pixel, then encode the whole + # clip causally -> num_latent distinct latent frames. + if len(frames) < total_pixel: + logger.warning( + "OmniDreams: HD-map has %d frames but %d are needed for " + "%d chunks x len_t=%d; clamping (repeating last frame).", + len(frames), + total_pixel, + num_chunks, + len_t, + ) + frames = list(frames) + [frames[-1]] * (total_pixel - len(frames)) + else: + frames = list(frames[:total_pixel]) + clip = self._preprocess_hdmap_clip( + frames, height, width, device, vae_dtype + ) + if clip is None: + logger.warning( + "OmniDreams: HD-map clip preprocessed to None; disabling " + "HDMap (all chunks fall back to zeros). Check hdmap input." ) return None - latent = self._vae_encode_normalized(x).to(dit_dtype) - # [B,16,t,h,w] -> [B, L, additional_concat_ch*pdim] via the DiT patchify. - tokens.append(self.transformer.patchify(latent)) + latent = self._vae_encode_normalized(clip).to(dit_dtype) # [B,16,L,h,w] + + if latent.shape[2] != num_latent: + logger.warning( + "OmniDreams: HD-map encoded to %d latent frames, expected %d " + "(num_chunks=%d, len_t=%d). Check VAE temporal compression.", + latent.shape[2], + num_latent, + num_chunks, + len_t, + ) + # Slice into per-chunk groups of len_t latent frames, patchify each: + # [B,16,len_t,h,w] -> [B, chunk_tokens, additional_concat_ch*pdim]. + tokens: list[torch.Tensor] = [] + for ci in range(num_chunks): + chunk_latent = latent[:, :, ci * len_t : (ci + 1) * len_t] + tokens.append(self.transformer.patchify(chunk_latent)) return tokens @torch.no_grad() @@ -393,7 +497,7 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: # Per-chunk HD-map tokens (None -> AR stage uses zeros / HDMap disabled). with self.use_declared_component(component_name="vae", module=self.vae): hdmap_tokens = self._encode_hdmap( - batch, device, vae_dtype, dit_dtype, num_chunks, height, width + batch, device, vae_dtype, dit_dtype, num_chunks, len_t, height, width ) batch.extra["omnidreams"] = { "hp": hp, diff --git a/python/sglang/multimodal_gen/test/unit/test_omnidreams_components.py b/python/sglang/multimodal_gen/test/unit/test_omnidreams_components.py index f9d2da2ebdec..73920e3b61f8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_omnidreams_components.py +++ b/python/sglang/multimodal_gen/test/unit/test_omnidreams_components.py @@ -354,6 +354,7 @@ def _ar_stage_and_args(arch, dit, scheduler, monkeypatch): stage.transformer = dit stage.scheduler = scheduler stage.vae = None + stage._component_residency_manager = None server_args = types.SimpleNamespace( pipeline_config=types.SimpleNamespace( dit_precision="fp32", @@ -622,33 +623,166 @@ def test_ar_hdmap_none_falls_back_to_zeros(monkeypatch): assert all(bool((h == 0).all()) for h in rec.hdmap_calls) -def test_encode_hdmap_broadcast_clamp_and_none(monkeypatch): - """`_encode_hdmap` assembles the per-chunk list (broadcast / clamp / None). +def _hdmap_stage(monkeypatch): + """Stage with stubbed preprocess + a fake causal VAE (tc=4) + identity patchify. - Stubs preprocess+encode+patchify so the test isolates the per-chunk control - flow (the VAE numerics are a GPU concern). Each src tag flows through as a - 1-element tensor so chunk identity is traceable. + The fake VAE maps a ``T``-frame clip to ``1 + (T-1)//4`` latent frames (mirroring + the Wan VAE temporal compression), filling latent frame ``j`` with value ``j`` so + chunk slicing is verifiable. patchify is identity so the returned tokens are the + sliced latents themselves. (Real VAE numerics are a GPU concern.) """ stage = OmniDreamsBeforeDenoisingStage.__new__(OmniDreamsBeforeDenoisingStage) stage.transformer = types.SimpleNamespace(patchify=lambda latent: latent) monkeypatch.setattr( stage, "_preprocess_pixels", - lambda src, h, w, d, dt: torch.tensor([float(src)]), + lambda src, h, w, d, dt: torch.zeros(1, 3, 1, 2, 2), ) - monkeypatch.setattr(stage, "_vae_encode_normalized", lambda x: x) + + def fake_vae(clip): + t = clip.shape[2] if clip.dim() == 5 else 1 + n_latent = 1 + (t - 1) // 4 + return torch.cat( + [torch.full((1, 16, 1, 2, 2), float(j)) for j in range(n_latent)], dim=2 + ) + + monkeypatch.setattr(stage, "_vae_encode_normalized", fake_vae) + return stage + + +def test_encode_hdmap_per_frame_clip_slicing(monkeypatch): + """Per-frame HD-map (option 2): the full raster sequence is encoded once as a + causal clip and sliced into ``num_chunks`` groups of ``len_t`` *distinct* + latent frames (this is what makes the generated viewpoint move).""" + stage = _hdmap_stage(monkeypatch) + dev = torch.device("cpu") + num_chunks, len_t = 3, 2 + num_latent = num_chunks * len_t # 6 + total_pixel = 1 + (num_latent - 1) * 4 # 21 + + b = types.SimpleNamespace(hdmap_path=list(range(total_pixel)), hdmap_pixels=None) + toks = stage._encode_hdmap( + b, dev, torch.float32, torch.float32, num_chunks, len_t, 16, 16 + ) + assert len(toks) == num_chunks + for ci, t in enumerate(toks): + assert t.shape[2] == len_t # len_t latent frames per chunk + got = [float(t[0, 0, k, 0, 0]) for k in range(len_t)] + assert got == [float(ci * len_t + k) for k in range(len_t)] + # Frames must actually differ across the rollout (regression: static angle). + assert float(toks[0][0, 0, 0, 0, 0]) != float(toks[-1][0, 0, len_t - 1, 0, 0]) + + +def test_encode_hdmap_clamps_short_sequence(monkeypatch): + """Fewer frames than needed are clamped (last repeated); still yields the full + per-chunk token list with ``len_t`` frames each, without crashing.""" + stage = _hdmap_stage(monkeypatch) dev = torch.device("cpu") + num_chunks, len_t = 2, 2 + b = types.SimpleNamespace(hdmap_path=[0, 1, 2], hdmap_pixels=None) # short + toks = stage._encode_hdmap( + b, dev, torch.float32, torch.float32, num_chunks, len_t, 16, 16 + ) + assert len(toks) == num_chunks + assert all(t.shape[2] == len_t for t in toks) + + +def test_encode_hdmap_video_path_decoded_per_frame(monkeypatch): + """A video-path HD-map is decoded via ``load_video`` into a per-frame clip.""" + stage = _hdmap_stage(monkeypatch) + num_chunks, len_t = 2, 2 + total_pixel = 1 + (num_chunks * len_t - 1) * 4 + fake_frames = [PIL.Image.new("RGB", (2, 2)) for _ in range(total_pixel)] + monkeypatch.setattr( + "sglang.multimodal_gen.runtime.pipelines_core.stages." + "model_specific_stages.omnidreams.load_video", + lambda path: fake_frames, + ) + b = types.SimpleNamespace(hdmap_path="scene_hdmap.mp4", hdmap_pixels=None) + toks = stage._encode_hdmap( + b, torch.device("cpu"), torch.float32, torch.float32, num_chunks, len_t, 16, 16 + ) + assert len(toks) == num_chunks + assert all(t.shape[2] == len_t for t in toks) + + +def test_encode_hdmap_cli_list_wraps_video_path(monkeypatch): + """CLI ``--hdmap-path`` always passes a list, even for a single arg; a lone + video-path string inside that list must be expanded via ``load_video``.""" + stage = _hdmap_stage(monkeypatch) + num_chunks, len_t = 2, 2 + total_pixel = 1 + (num_chunks * len_t - 1) * 4 + fake_frames = [PIL.Image.new("RGB", (2, 2)) for _ in range(total_pixel)] + monkeypatch.setattr( + "sglang.multimodal_gen.runtime.pipelines_core.stages." + "model_specific_stages.omnidreams.load_video", + lambda path: fake_frames, + ) + # Simulate CLI --hdmap-path "scene.mp4" -> list with one string element. + b = types.SimpleNamespace(hdmap_path=["scene_hdmap.mp4"], hdmap_pixels=None) + toks = stage._encode_hdmap( + b, torch.device("cpu"), torch.float32, torch.float32, num_chunks, len_t, 16, 16 + ) + assert len(toks) == num_chunks + assert all(t.shape[2] == len_t for t in toks) + # Confirm it took the per-frame path (all tokens come from distinct VAE frames). + # The first and last chunk should differ because our fake VAE encodes each frame + # with a distinct value. + assert not torch.equal(toks[0], toks[-1]), ( + "CLI list-wrapped video path should produce distinct per-chunk tokens; " + "fallback broadcast would make them all identical" + ) - # Single (non-list) input broadcasts to every chunk. - b1 = types.SimpleNamespace(hdmap_path=7, hdmap_pixels=None) - toks = stage._encode_hdmap(b1, dev, torch.float32, torch.float32, 3, 16, 16) - assert [float(t) for t in toks] == [7.0, 7.0, 7.0] - # Per-chunk list shorter than num_chunks clamps to the last entry. - b2 = types.SimpleNamespace(hdmap_path=[1, 2], hdmap_pixels=None) - toks2 = stage._encode_hdmap(b2, dev, torch.float32, torch.float32, 4, 16, 16) - assert [float(t) for t in toks2] == [1.0, 2.0, 2.0, 2.0] +def test_encode_hdmap_single_image_broadcast_fallback(monkeypatch): + """A single image (non-video string) degenerates to broadcasting one raster + across every latent frame (back-compat / smoke; no temporal motion).""" + stage = _hdmap_stage(monkeypatch) + num_chunks, len_t = 3, 2 + b = types.SimpleNamespace(hdmap_path="frame.png", hdmap_pixels=None) + toks = stage._encode_hdmap( + b, torch.device("cpu"), torch.float32, torch.float32, num_chunks, len_t, 16, 16 + ) + assert len(toks) == num_chunks + for t in toks: + assert t.shape[2] == len_t + assert bool((t == 0).all()) # all broadcast from the single latent frame + + +def test_encode_hdmap_none_returns_none(monkeypatch): + """No HD-map input -> None (AR stage falls back to zeros).""" + stage = _hdmap_stage(monkeypatch) + b = types.SimpleNamespace(hdmap_path=None, hdmap_pixels=None) + assert ( + stage._encode_hdmap( + b, torch.device("cpu"), torch.float32, torch.float32, 2, 2, 16, 16 + ) + is None + ) + - # No HD-map input -> None (AR stage falls back to zeros). - b3 = types.SimpleNamespace(hdmap_path=None, hdmap_pixels=None) - assert stage._encode_hdmap(b3, dev, torch.float32, torch.float32, 2, 16, 16) is None +def test_encode_hdmap_broadcasts_single_frame_across_len_t(monkeypatch): + """A single HD-map image (1 latent frame) must be tiled to ``len_t`` frames so + the patchified token count matches ``chunk_tokens`` (regression: a real + hdmap input previously produced tokens_per_frame and crashed the DiT add). + """ + stage = OmniDreamsBeforeDenoisingStage.__new__(OmniDreamsBeforeDenoisingStage) + # Real patchify so token length reflects the latent temporal extent. + stage.transformer = OmniDreamsDiT.__new__(OmniDreamsDiT) + stage.transformer.arch = OmniDreamsDiTArchConfig( + in_channels=16, out_channels=16, patch_spatial=2, patch_temporal=1 + ) + # VAE-encode stub returns a single-frame latent [B=1, C=16, t=1, h=2, w=2] + # -> tokens_per_frame = (h/ps)*(w/ps) = 1. + monkeypatch.setattr( + stage, "_preprocess_pixels", lambda src, h, w, d, dt: torch.zeros(1, 3, 1, 2, 2) + ) + monkeypatch.setattr( + stage, "_vae_encode_normalized", lambda x: torch.zeros(1, 16, 1, 2, 2) + ) + dev = torch.device("cpu") + b = types.SimpleNamespace(hdmap_path=1, hdmap_pixels=None) + tokens_per_frame = (2 // 2) * (2 // 2) # = 1 + for len_t in (1, 2, 3): + toks = stage._encode_hdmap(b, dev, torch.float32, torch.float32, 1, len_t, 2, 2) + assert toks[0].shape[1] == len_t * tokens_per_frame From 721ccb49ea2055a2008ed41d82cd59b004aa2e2e Mon Sep 17 00:00:00 2001 From: Cerdore Date: Sat, 13 Jun 2026 16:40:15 +0800 Subject: [PATCH 5/5] fix(diffusion): re-inject clean first-frame image during OmniDreams context-noise KV cache write At AR step 0, the context-noise forward (finalize_kv_cache) must store CLEAN K/V for frame 0, matching FlashDreams _maybe_inject_image() inside _predict_branch which fires during finalize_kv_cache -> predict_flow. Without this re-injection, frame-0 KV cache entries are noise-corrupted at sigma ~0.13, causing a distribution shift that accumulates into progressively blurry video across AR chunks. --- .../pipelines_core/stages/model_specific_stages/omnidreams.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py index b1bb48ac3643..73ca0974062c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py @@ -795,6 +795,8 @@ def predict_flow(noisy: torch.Tensor, t: torch.Tensor) -> torch.Tensor: ctx_noise_t, rng=gen, ) + if pin: + ctx_latent = ctx_latent * (1.0 - inject_mask) + image_full * inject_mask self.transformer( hidden_states=ctx_latent, encoder_hidden_states=text,