Add Omnidream interactive upsampler client

integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py

@@ -29,6 +29,10 @@
     FlashdreamsWorldModelSession,
 )
 from omnidreams.interactive_drive.world_model.manifest import WorldModelManifest
+from omnidreams.interactive_drive.world_model.uplift import (
+    UpliftStreamClient,
+    UpliftStreamConfig,
+)
 from PIL import Image
 
 _FIRST_STEADY_STATE_WARMUP_MESSAGE = "Optimizing world model..."
@@ -55,6 +59,7 @@ def __init__(
         self._scene: SceneBundle | None = None
         self._next_chunk_count = 0
         self._debug_first_chunk_condition_frames: tuple[np.ndarray, ...] | None = None
+        self._uplift: UpliftStreamClient | None = self._build_uplift_client()
 
     @property
     def can_prewarm(self) -> bool:
@@ -107,6 +112,7 @@ def load_scene(self, scene: SceneBundle) -> None:
         self._session.prepare_for_scene(
             initial_rgb=scene.initial_rgb, prompt=scene.prompt
         )
+        self._reset_uplift_stream()
         prepare_end = time.perf_counter()
         logger.info(
             "[world-model] load_scene "
@@ -154,10 +160,12 @@ def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
             scene.initial_rgb, condition_frames, scene.prompt
         )
         model_end = time.perf_counter()
+        uplift_status = self._submit_uplift(model_frames)
         merged_frames = self._merge_frames(
             display_frames,
             model_frames,
             annotate_first_transition=True,
+            status_message=uplift_status,
         )
         merge_end = time.perf_counter()
         logger.info(
@@ -193,7 +201,12 @@ def render_next_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
         condition_frames = [frame.rgb_host_uint8 for frame in raster_chunk.frames]
         model_frames = self._session.continue_generation(condition_frames)
         model_end = time.perf_counter()
-        merged_frames = self._merge_frames(raster_chunk.frames, model_frames)
+        uplift_status = self._submit_uplift(model_frames)
+        merged_frames = self._merge_frames(
+            raster_chunk.frames,
+            model_frames,
+            status_message=uplift_status,
+        )
         merge_end = time.perf_counter()
         self._next_chunk_count += 1
         total_ms = (merge_end - chunk_start) * 1000.0
@@ -228,15 +241,58 @@ def render_next_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
     def reset(self) -> None:
         self._session.reset()
         self._next_chunk_count = 0
+        self._reset_uplift_stream()
 
     def reset_scene_conditioning(self) -> None:
         self._session.reset(clear_precomputed_embeddings=True)
         self._next_chunk_count = 0
+        self._reset_uplift_stream()
 
     def close(self) -> None:
+        if self._uplift is not None:
+            self._uplift.close()
+            self._uplift = None
         self._session.close()
         self._rasterizer.cleanup()
 
+    def _build_uplift_client(self) -> UpliftStreamClient | None:
+        if not self._manifest.upsampling_enabled:
+            return None
+        config = UpliftStreamConfig(
+            server=self._manifest.upsampling_server,
+            scale=self._manifest.upsampling_scale,
+            sparse_ratio=self._manifest.upsampling_sparse_ratio,
+            input_format=self._manifest.upsampling_input_format,  # type: ignore[arg-type]
+            input_jpeg_quality=self._manifest.upsampling_input_jpeg_quality,
+            return_frames=self._manifest.upsampling_return_frames,
+            max_queue_chunks=self._manifest.upsampling_max_queue_chunks,
+            max_message_mb=self._manifest.upsampling_max_message_mb,
+        )
+        logger.info(
+            "[world-model] uplift enabled server={} scale={}x queue_depth={} "
+            "input_format={} return_frames={}",
+            config.server,
+            config.scale,
+            config.max_queue_chunks,
+            config.input_format,
+            config.return_frames,
+        )
+        client = UpliftStreamClient(config)
+        client.start()
+        return client
+
+    def _reset_uplift_stream(self) -> None:
+        if not self._manifest.upsampling_enabled:
+            return
+        if self._uplift is not None:
+            self._uplift.close()
+        self._uplift = self._build_uplift_client()
+
+    def _submit_uplift(self, model_frames: list[object]) -> str | None:
+        if self._uplift is None:
+            return None
+        return self._uplift.submit(model_frames)
+
     def _require_scene(self) -> SceneBundle:
         if self._scene is None:
             raise RuntimeError(
@@ -271,6 +327,7 @@ def _merge_frames(
         model_frames: Sequence[object],
         *,
         annotate_first_transition: bool = False,
+        status_message: str | None = None,
     ) -> tuple[PresentedFrame, ...]:
         if len(raster_frames) != len(model_frames):
             raise ValueError(
@@ -283,6 +340,14 @@ def _merge_frames(
         for index, (raster_frame, model_rgb) in enumerate(
             zip(raster_frames, model_frames, strict=True)
         ):
+            frame_status = _combine_status_messages(
+                (
+                    _FIRST_STEADY_STATE_WARMUP_MESSAGE
+                    if annotate_first_transition and index == last_index
+                    else None
+                ),
+                status_message,
+            )
             merged.append(
                 PresentedFrame(
                     timestamp_us=raster_frame.timestamp_us,
@@ -292,11 +357,7 @@ def _merge_frames(
                     depth_native=raster_frame.depth_native,
                     model_rgb_host_uint8=model_rgb,
                     bev_host_uint8=raster_frame.bev_host_uint8,
-                    status_message=(
-                        _FIRST_STEADY_STATE_WARMUP_MESSAGE
-                        if annotate_first_transition and index == last_index
-                        else None
-                    ),
+                    status_message=frame_status,
                 )
             )
         return tuple(merged)
@@ -314,3 +375,10 @@ def _log_prompt_handoff(stage: str, scene: SceneBundle) -> None:
         f"length={len(prompt)} "
         f"text={prompt_text!r}",
     )
+
+
+def _combine_status_messages(*messages: str | None) -> str | None:
+    parts = [message for message in messages if message]
+    if not parts:
+        return None
+    return " | ".join(parts)

integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model.yaml

@@ -26,7 +26,14 @@ local_attn_size: 6
 skip_finalize_kv_cache: false
 sink_size: 0
 denoising_steps: [1000, 500]
-upsampling_enabled: false
-upsampling_scale: 4
+upsampling_enabled: true
+upsampling_server: 127.0.0.1:8090
+upsampling_scale: 2
+upsampling_sparse_ratio: 1.5
+upsampling_input_format: jpeg
+upsampling_input_jpeg_quality: 90
+upsampling_return_frames: false
+upsampling_max_queue_chunks: 16
+upsampling_max_message_mb: 512
 device: cuda:0
 seed_for_every_rollout:

integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_perf.yaml

@@ -28,7 +28,14 @@ skip_finalize_kv_cache: true
 sink_size: 0
 denoising_steps: [1000, 100]
 upsampling_enabled: false
+upsampling_server: 127.0.0.1:8090
 upsampling_scale: 4
+upsampling_sparse_ratio: 0.0
+upsampling_input_format: jpeg
+upsampling_input_jpeg_quality: 90
+upsampling_return_frames: false
+upsampling_max_queue_chunks: 4
+upsampling_max_message_mb: 512
 device: cuda:0
 seed_for_every_rollout:
 

integrations/omnidreams/omnidreams/interactive_drive/world_model/flashdreams_adapter.py

@@ -31,10 +31,6 @@ def _select_config_name(manifest: WorldModelManifest) -> str:
     Returns a key from ``omnidreams.config.OMNIDREAMS_CONFIGS``
     (i.e. the same slug ``flashdreams-run`` accepts as its first positional arg).
     """
-    if manifest.upsampling_enabled:
-        raise NotImplementedError(
-            "flashdreams interactive-drive path does not support upsampling."
-        )
     if manifest.sink_size != 0:
         raise NotImplementedError(
             "flashdreams interactive-drive path currently supports sink_size=0 only."

integrations/omnidreams/omnidreams/interactive_drive/world_model/manifest.py

@@ -157,7 +157,14 @@ class WorldModelManifest:
     sink_size: int = 0
     denoising_steps: list[int] = field(default_factory=lambda: [1000, 500])
     upsampling_enabled: bool = False
+    upsampling_server: str = "127.0.0.1:8090"
     upsampling_scale: int = 4
+    upsampling_sparse_ratio: float = 0.0
+    upsampling_input_format: str = "jpeg"
+    upsampling_input_jpeg_quality: int = 90
+    upsampling_return_frames: bool = False
+    upsampling_max_queue_chunks: int = 4
+    upsampling_max_message_mb: int = 512
     device: str = "cuda:0"
     seed_for_every_rollout: int | None = None
     native_dit_acceleration: str = "disabled"
@@ -207,7 +214,14 @@ def load_world_model_manifest(path: str | Path) -> WorldModelManifest:
         sink_size=int(data.get("sink_size", 0)),
         denoising_steps=[int(x) for x in data.get("denoising_steps", [1000, 500])],
         upsampling_enabled=bool(data.get("upsampling_enabled", False)),
+        upsampling_server=str(data.get("upsampling_server", "127.0.0.1:8090")),
         upsampling_scale=int(data.get("upsampling_scale", 4)),
+        upsampling_sparse_ratio=float(data.get("upsampling_sparse_ratio", 0.0)),
+        upsampling_input_format=str(data.get("upsampling_input_format", "jpeg")),
+        upsampling_input_jpeg_quality=int(data.get("upsampling_input_jpeg_quality", 90)),
+        upsampling_return_frames=bool(data.get("upsampling_return_frames", False)),
+        upsampling_max_queue_chunks=int(data.get("upsampling_max_queue_chunks", 4)),
+        upsampling_max_message_mb=int(data.get("upsampling_max_message_mb", 512)),
         device=str(data.get("device", "cuda:0")),
         seed_for_every_rollout=(
             int(data["seed_for_every_rollout"])