[feat] Add Cosmos Predict2.5 distilled profile to DreamVerse

apps/dreamverse/README.md

@@ -138,6 +138,27 @@ dreamverse-server --host 0.0.0.0 --port 8009
 The Dreamverse backend defaults to `0.0.0.0:8009` and starts one GPU worker on
 the first visible GPU by default.
 
+### Cosmos Predict2.5 distilled (experimental)
+
+Dreamverse can run the FastVideo-converted Cosmos Predict2.5 2B distilled
+package as an independent text-to-world segment generator. Point the runtime at
+the converted package and disable startup warmup for the first validation run:
+
+```bash
+export DREAMVERSE_MODEL_ID=cosmos25-distilled
+export DREAMVERSE_MODEL_PATH=/path/to/Cosmos-Predict2.5-2B-Distilled-TrigFlow-FastVideo
+export FASTVIDEO_ENABLE_STARTUP_WARMUP=0
+export ENABLE_TORCH_COMPILE=0
+dreamverse-server --host 0.0.0.0 --port 8009
+```
+
+This initial integration uses Torch SDPA with BF16 weights and the validated
+Cosmos profile: 704x1280, 77 frames at 16 FPS, and four distilled steps. Cosmos
+does not produce audio, so Dreamverse muxes a matching silent track to preserve
+the existing browser streaming contract. Image input, LoRA, and cross-segment
+continuation are not supported in this first profile; later prompts generate
+independent clips.
+
 ### Check Readiness
 
 In another shell, verify that the backend process is alive:

apps/dreamverse/dreamverse/av_streaming.py

@@ -130,6 +130,7 @@ def stream_fmp4(
     frames: list[np.ndarray],
     audio: object,
     audio_sample_rate: int | None,
+    fps: int = TARGET_FPS,
     stream_id: str,
     timings: dict,
     head_trim_frames: int = 0,
@@ -145,6 +146,7 @@ def stream_fmp4(
         frames: RGB24 video frames as HxWx3 uint8 arrays.
         audio: 1D/2D tensor or ndarray, float values in [-1, 1].
         audio_sample_rate: sample rate of ``audio``.
+        fps: video frame rate used for overlap trimming and muxing.
         stream_id: caller-supplied identifier carried on every event.
         timings: dict mutated in place with ffmpeg/stream timing metrics.
         head_trim_frames: video frames to drop from the start
@@ -172,6 +174,8 @@ def stream_fmp4(
         return False, "audio_sample_rate is None"
     if FFMPEG_BIN is None:
         return False, "ffmpeg not found"
+    if fps <= 0:
+        return False, f"fps must be positive, got {fps}"
 
     if head_trim_audio_frames is None:
         head_trim_audio_frames = head_trim_frames
@@ -192,12 +196,12 @@ def stream_fmp4(
     out_frames = (frames[head_trim_frames:] if head_trim_frames > 0 else frames)
     sample_rate = int(audio_sample_rate)
     if head_trim_audio_frames > 0:
-        trim_start_samples = int(round((head_trim_audio_frames / float(TARGET_FPS)) * sample_rate))
+        trim_start_samples = int(round((head_trim_audio_frames / float(fps)) * sample_rate))
         if trim_start_samples >= audio_int16.shape[0]:
             return False, ("audio too short after overlap trim: "
                            f"trim_start_samples={trim_start_samples}"
                            f", audio_samples={audio_int16.shape[0]}")
-        keep_samples = int(round((len(out_frames) / float(TARGET_FPS)) * sample_rate))
+        keep_samples = int(round((len(out_frames) / float(fps)) * sample_rate))
         trim_end_samples = min(
             audio_int16.shape[0],
             trim_start_samples + keep_samples,
@@ -228,7 +232,7 @@ def stream_fmp4(
         "-s:v",
         f"{width}x{height}",
         "-r",
-        str(TARGET_FPS),
+        str(fps),
         "-i",
         "pipe:0",
         "-i",

apps/dreamverse/dreamverse/config.py

@@ -55,23 +55,58 @@ def _resolve_frontend_static_dir_candidates() -> tuple[str, ...]:
 MODEL_REGISTRY = {
     "fast-ltx2": {
         "name": "FastLTX2",
+        "family": "ltx2",
         "model_path": "FastVideo/LTX2-Distilled-Diffusers",
         "config_model_path": "FastVideo/LTX2-Distilled-Diffusers",
         "lora_repo": "FastVideo/LTX2-OmniNFT-LoRA",
+        "attention_backend": "FLASH_ATTN",
+        "supports_audio": True,
+        "supports_continuation": True,
+        "supports_lora": True,
+        "uses_refine": True,
+        "transformer_quant": "NVFP4",
+        "compile": True,
     },
     "fast-ltx23": {
         "name": "FastLTX23",
+        "family": "ltx2",
         "model_path": "FastVideo/LTX-2.3-Distilled-Diffusers",
         "config_model_path": "FastVideo/LTX-2.3-Distilled-Diffusers",
         "lora_repo": "FastVideo/LTX-2.3-OmniNFT-LoRA",
+        "attention_backend": "FLASH_ATTN",
+        "supports_audio": True,
+        "supports_continuation": True,
+        "supports_lora": True,
+        "uses_refine": True,
+        "transformer_quant": "NVFP4",
+        "compile": True,
+    },
+    "cosmos25-distilled": {
+        "name": "Cosmos Predict2.5 Distilled",
+        "family": "cosmos25_distilled",
+        "model_path": "FastVideo/Cosmos-Predict2.5-2B-Distilled-TrigFlow",
+        "config_model_path": "FastVideo/Cosmos-Predict2.5-2B-Distilled-TrigFlow",
+        "attention_backend": "TORCH_SDPA",
+        "supports_audio": False,
+        "supports_continuation": False,
+        "supports_lora": False,
+        "uses_refine": False,
+        "transformer_quant": None,
+        "compile": False,
+        "height": 704,
+        "width": 1280,
+        "num_frames": 77,
+        "fps": 16,
+        "num_inference_steps": 4,
+        "seed": 42,
     },
 }
 
-DEFAULT_MODEL_ID = "fast-ltx2"
-
-ACTIVE_MODEL_ID = (os.getenv("DREAMVERSE_MODEL_ID", "").strip() or DEFAULT_MODEL_ID)
+BUILTIN_DEFAULT_MODEL_ID = "fast-ltx2"
+ACTIVE_MODEL_ID = (os.getenv("DREAMVERSE_MODEL_ID", "").strip() or BUILTIN_DEFAULT_MODEL_ID)
 if ACTIVE_MODEL_ID not in MODEL_REGISTRY:
-    ACTIVE_MODEL_ID = DEFAULT_MODEL_ID
+    ACTIVE_MODEL_ID = BUILTIN_DEFAULT_MODEL_ID
+DEFAULT_MODEL_ID = ACTIVE_MODEL_ID
 
 # Active model configuration
 MODEL_CONFIG = MODEL_REGISTRY[ACTIVE_MODEL_ID]

apps/dreamverse/dreamverse/gpu_pool.py

@@ -162,8 +162,6 @@ def gpu_worker_process(
     generation to VideoGenerationWorker; AV muxing to av_streaming.
     """
     os.environ["CUDA_VISIBLE_DEVICES"] = cuda_device
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
-
     from dreamverse.video_generation import VideoGenerationWorker
 
     worker = VideoGenerationWorker(gpu_id)
@@ -212,6 +210,7 @@ def _publish(event: StreamEvent) -> None:
                         frames=step_result.frames,
                         audio=step_result.audio,
                         audio_sample_rate=step_result.audio_sample_rate,
+                        fps=step_result.fps,
                         stream_id=stream_id,
                         timings=step_result.timings,
                         head_trim_frames=head_trim_frames,

apps/dreamverse/dreamverse/tests/test_config.py

@@ -144,6 +144,42 @@ def test_config_uses_five_minute_session_timeout(monkeypatch):
     assert module.SESSION_TIMEOUT_SECONDS == 300
 
 
+def test_config_selects_cosmos25_distilled_profile(monkeypatch, tmp_path):
+    _set_required_prompt_keys(monkeypatch)
+    model_path = tmp_path / "cosmos25-distilled"
+    monkeypatch.setenv("DREAMVERSE_MODEL_ID", "cosmos25-distilled")
+    monkeypatch.setenv("DREAMVERSE_MODEL_PATH", str(model_path))
+
+    module = _load_config_module()
+
+    expected_config = {
+        **module.MODEL_REGISTRY["cosmos25-distilled"],
+        "model_path": str(model_path),
+        "config_model_path": str(model_path),
+    }
+    assert module.ACTIVE_MODEL_ID == "cosmos25-distilled"
+    assert module.DEFAULT_MODEL_ID == "cosmos25-distilled"
+    assert expected_config == module.MODEL_CONFIG
+    assert module.MODEL_CONFIG["attention_backend"] == "TORCH_SDPA"
+    assert module.MODEL_CONFIG["transformer_quant"] is None
+    assert module.MODEL_CONFIG["supports_audio"] is False
+    assert module.MODEL_CONFIG["supports_continuation"] is False
+    assert module.MODEL_CONFIG["num_inference_steps"] == 4
+    assert module.MODEL_CONFIG["fps"] == 16
+
+
+def test_config_unknown_model_falls_back_to_builtin_default(monkeypatch):
+    _set_required_prompt_keys(monkeypatch)
+    monkeypatch.setenv("DREAMVERSE_MODEL_ID", "not-a-model")
+    monkeypatch.delenv("DREAMVERSE_MODEL_PATH", raising=False)
+
+    module = _load_config_module()
+
+    assert module.ACTIVE_MODEL_ID == "fast-ltx2"
+    assert module.DEFAULT_MODEL_ID == "fast-ltx2"
+    assert module.MODEL_CONFIG["family"] == "ltx2"
+
+
 def test_config_rejects_invalid_prompt_provider(monkeypatch):
     monkeypatch.setenv("FASTVIDEO_PROMPT_PROVIDER", "unsupported")
     _set_required_prompt_keys(monkeypatch)

apps/dreamverse/dreamverse/tests/test_video_generation_cosmos25.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from dreamverse.video_generation import VideoGenerationWorker
+
+
+def _cosmos_config() -> dict:
+    return {
+        "name": "Cosmos Predict2.5 Distilled",
+        "family": "cosmos25_distilled",
+        "model_path": "/models/cosmos25-distilled",
+        "supports_audio": False,
+        "supports_continuation": False,
+        "supports_lora": False,
+        "height": 704,
+        "width": 1280,
+        "num_frames": 77,
+        "fps": 16,
+        "num_inference_steps": 4,
+        "seed": 42,
+    }
+
+
+class _RecordingGenerator:
+    def __init__(self) -> None:
+        self.calls: list[dict] = []
+
+    def generate_video(self, **kwargs):
+        self.calls.append(kwargs)
+        return {
+            "frames": [np.zeros((8, 8, 3), dtype=np.uint8) for _ in range(9)],
+            "generation_time": 0.25,
+        }
+
+
+def _worker(monkeypatch: pytest.MonkeyPatch) -> tuple[VideoGenerationWorker, _RecordingGenerator]:
+    monkeypatch.setattr(torch.cuda, "synchronize", lambda: None)
+    worker = VideoGenerationWorker(gpu_id=0)
+    generator = _RecordingGenerator()
+    worker.generator = generator
+    worker.current_model_config = _cosmos_config()
+    return worker, generator
+
+
+def test_cosmos25_step_uses_distilled_profile_and_synthesizes_silence(monkeypatch):
+    worker, generator = _worker(monkeypatch)
+
+    result = worker.generate_step(
+        "A robot crossing a desert",
+        segment_idx=2,
+        image_path=None,
+        reset_conditioning=False,
+    )
+
+    assert len(generator.calls) == 1
+    request = generator.calls[0]
+    assert request["height"] == 704
+    assert request["width"] == 1280
+    assert request["num_frames"] == 77
+    assert request["fps"] == 16
+    assert request["num_inference_steps"] == 4
+    assert request["seed"] == 42
+    assert request["return_frames"] is True
+    assert request["save_video"] is False
+    assert "ltx2_image_crf" not in request
+    assert "conditioning_images" not in request
+    assert "audio_latents" not in request
+
+    assert result.fps == 16
+    assert result.audio_sample_rate == 24000
+    assert isinstance(result.audio, torch.Tensor)
+    assert result.audio.dtype is torch.float32
+    assert result.audio.shape == (13500,)
+    assert torch.count_nonzero(result.audio) == 0
+    assert result.head_trim_frames == 0
+    assert result.head_trim_audio_frames == 0
+    assert worker.continuation.video_images is None
+    assert worker.continuation.audio_latents is None
+
+
+def test_cosmos25_step_rejects_initial_image(monkeypatch):
+    worker, generator = _worker(monkeypatch)
+
+    with pytest.raises(RuntimeError, match="text-to-world only"):
+        worker.generate_step(
+            "Animate this image",
+            segment_idx=1,
+            image_path="input.png",
+            reset_conditioning=True,
+        )
+
+    assert generator.calls == []
+
+
+def test_cosmos25_warmup_generates_one_independent_segment(monkeypatch):
+    worker, _generator = _worker(monkeypatch)
+    calls: list[tuple[int, bool]] = []
+
+    def fake_generate_step(prompt, segment_idx, image_path, reset_conditioning):
+        del prompt, image_path
+        calls.append((segment_idx, reset_conditioning))
+        return SimpleNamespace(timings={"e2e_latency_ms": 123.0})
+
+    monkeypatch.setattr(worker, "generate_step", fake_generate_step)
+
+    timings = worker.warmup("A warmup prompt")
+
+    assert calls == [(1, True)]
+    assert timings["warmup_segment1_ms"] == 123.0
+    assert timings["warmup_total_ms"] >= 0.0
+
+
+def test_ltx_profile_keeps_existing_generation_defaults(monkeypatch):
+    worker, generator = _worker(monkeypatch)
+    worker.current_model_config = {
+        "name": "FastLTX2",
+        "family": "ltx2",
+        "model_path": "/models/ltx2",
+        "supports_audio": True,
+        "supports_continuation": True,
+        "supports_lora": True,
+    }
+
+    result = worker.generate_step(
+        "A reference image comes alive",
+        segment_idx=1,
+        image_path="input.png",
+        reset_conditioning=True,
+    )
+
+    request = generator.calls[0]
+    assert request["height"] == 1088
+    assert request["width"] == 1920
+    assert request["num_frames"] == 121
+    assert request["fps"] == 24
+    assert request["num_inference_steps"] == 5
+    assert request["seed"] == 10
+    assert request["ltx2_image_crf"] == 0.0
+    assert request["image_path"] == "input.png"
+    assert result.audio is None
+    assert result.audio_sample_rate is None
+    assert result.fps == 24