fix(diffusion): re-inject clean first-frame image during OmniDreams context-noise KV cache write

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",
 ]

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"

python/sglang/multimodal_gen/configs/models/vaes/__init__.py

@@ -6,12 +6,16 @@
 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",
     "HunyuanVAEConfig",
     "StableDiffusion3VAEConfig",
     "WanVAEConfig",
     "Hunyuan3DVAEConfig",
+    "OmniDreamsVAEConfig",
 ]

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
+    )

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",

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,
+        )

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",
 ]

python/sglang/multimodal_gen/configs/sample/omnidreams.py

@@ -0,0 +1,40 @@
+# 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 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

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -947,6 +947,18 @@ def add_argument(*name_or_flags, **kwargs):
                 '--image-path "img1.png" "img2.png"'
             ),
         )
+        add_argument(
+            "--hdmap-path",
+            type=str,
+            nargs="+",
+            help=(
+                "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(
             "--action",
             type=str,

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,
@@ -125,6 +128,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,
@@ -395,8 +399,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 = []
@@ -635,6 +669,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,

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):