clean up + formatting + try to reduce latency

Author: ArielG-NV

apps/interactive_drive/README.md

@@ -55,8 +55,7 @@ uv sync --package flashdreams-omnidreams --extra interactive-drive
 uv run flashdreams-run-v2 interactive-drive-omnidreams --mode webrtc --port 8089
 ```
 
-Use `interactive-drive-omnidreams-perf` instead for the native-accelerated,
-performance-tuned configuration.
+Use `interactive-drive-omnidreams-perf`/`interactive-drive-omnidreams-fast-perf` instead for the native-accelerated, performance-tuned configurations.
 
 Forward so that you may connect via `<ip>:8089` by adding `--host 0.0.0.0`
 

apps/interactive_drive/interactive_drive/core.py

@@ -24,7 +24,7 @@
 from flashdreams.infra.postprocess import VideoPostprocessChainConfig
 from flashdreams.plugins.registry import discover_postprocess_presets
 from flashdreams.runtime.keyboard import normalize_key
-from flashdreams.runtime_v2.session_desc import SessionDesc
+from flashdreams.runtime_v2.session_desc import BackpressureMode, SessionDesc
 from flashdreams.runtime_v2.step_result import StepResult
 from flashdreams.runtime_v2.user_input_event import (
     GamepadUserInputEvent,
@@ -453,7 +453,11 @@ def __init__(
         self._config: InteractiveDriveConfig | None = None
         self._desc = SessionDesc(
             output_layout=VideoTensorLayout.tchw,
-            frames_per_second_for_ui=30,
+            backpressure_mode=BackpressureMode.BLOCK,
+            # If too low we start building up a 
+            # backlog of frames which raises percived
+            # latency.
+            frames_per_second_for_ui=75,
             frames_per_second_for_step=defaults.fps,
             video_width=defaults.width,
             video_height=defaults.height,

apps/interactive_drive/interactive_drive/tests/test_application.py

@@ -29,7 +29,7 @@
 from interactive_drive.input.keyboard import command_from_snapshot
 from interactive_drive.types import ControlSnapshot
 
-from flashdreams.runtime_v2.session_desc import PresentationMode
+from flashdreams.runtime_v2.session_desc import BackpressureMode, PresentationMode
 from flashdreams.runtime_v2.user_input_event import (
     GamepadUserInputEvent,
     KeyboardInputState,
@@ -299,6 +299,16 @@ def test_interactive_drive_no_ui_skips_ui_and_bev(tmp_path: Path) -> None:
     assert session._registered_ui_loop is None
 
 
+def test_interactive_drive_drops_stale_chunks_for_low_latency(tmp_path: Path) -> None:
+    scene = tmp_path / "local.usdz"
+    scene.touch()
+    app = InteractiveDriveApplication()
+
+    app.init(["--scene", str(scene)])
+
+    assert app.session_desc().backpressure_mode is BackpressureMode.DROP_OLDEST
+
+
 def test_interactive_drive_resolves_default_scene_when_omitted(
     tmp_path: Path,
     monkeypatch: pytest.MonkeyPatch,

integrations_v2/omnidreams/config.py

@@ -8,7 +8,6 @@
 from typing import cast
 
 import torch
-
 from omnidreams.impl.pipeline import OmnidreamsPipelineConfig
 from omnidreams.impl.transformer import CosmosTransformerConfig
 from omnidreams.impl.transformer.network import CosmosDiTNetworkConfig

integrations_v2/omnidreams/tests/test_omnidreams_native_vae.py

@@ -189,9 +189,13 @@ def ensure(path: Path) -> Path:
 
     monkeypatch.setattr(exporter, "ensure_lightvae_fp8_state", ensure)
     config = OMNIDREAMS_FAST_PERF_PIPELINE_CONFIG
+    image_encoder = config.image_encoder
+    encoder = config.encoder
+    assert isinstance(image_encoder, OmnidreamsWanVAEEncoderConfig)
+    assert isinstance(encoder, OmnidreamsWanVAEEncoderConfig)
 
-    assert _native_vae_fp8_state_path(config.image_encoder) == str(expected)
-    assert _native_vae_fp8_state_path(config.encoder) == str(expected)
+    assert _native_vae_fp8_state_path(image_encoder) == str(expected)
+    assert _native_vae_fp8_state_path(encoder) == str(expected)
     assert calls == [
         Path(DEFAULT_LIGHTVAE_FP8_STATE_PATH),
         Path(DEFAULT_LIGHTVAE_FP8_STATE_PATH),

integrations_v2/omnidreams/tests/test_recipe_configs.py

@@ -29,6 +29,8 @@
     OMNIDREAMS_PERF_PIPELINE_CONFIG,
     OMNIDREAMS_PIPELINE_CONFIG,
 )
+from omnidreams.impl.transformer import CosmosTransformerConfig
+from omnidreams.impl.vae_native import OmnidreamsWanVAEEncoderConfig
 
 from flashdreams.api_v2.application import IApplication
 
@@ -46,16 +48,22 @@ def test_pipeline_configs_are_keyed_by_name() -> None:
 
 def test_fast_perf_uses_native_vae_when_available() -> None:
     config = OMNIDREAMS_FAST_PERF_PIPELINE_CONFIG
-
-    assert config.image_encoder.native_vae_acceleration == "required"
-    assert config.image_encoder.native_vae_backend == "fp8"
-    assert config.image_encoder.native_vae_fp8_state_path is None
-    assert config.image_encoder.native_vae_fp8_auto_export is True
-    assert config.encoder.native_vae_acceleration == "required"
-    assert config.encoder.native_vae_backend == "fp8"
-    assert config.encoder.native_vae_fp8_state_path is None
-    assert config.encoder.native_vae_fp8_auto_export is True
-    assert config.diffusion_model.transformer.native_dit_acceleration == "required"
+    image_encoder = config.image_encoder
+    encoder = config.encoder
+    transformer = config.diffusion_model.transformer
+    assert isinstance(image_encoder, OmnidreamsWanVAEEncoderConfig)
+    assert isinstance(encoder, OmnidreamsWanVAEEncoderConfig)
+    assert isinstance(transformer, CosmosTransformerConfig)
+
+    assert image_encoder.native_vae_acceleration == "required"
+    assert image_encoder.native_vae_backend == "fp8"
+    assert image_encoder.native_vae_fp8_state_path is None
+    assert image_encoder.native_vae_fp8_auto_export is True
+    assert encoder.native_vae_acceleration == "required"
+    assert encoder.native_vae_backend == "fp8"
+    assert encoder.native_vae_fp8_state_path is None
+    assert encoder.native_vae_fp8_auto_export is True
+    assert transformer.native_dit_acceleration == "required"
 
 
 def test_application_defaults_are_owned_by_each_adapter() -> None: