Support split WebRTC timing and THWC frame chunks

flashdreams/flashdreams/core/attention/kvcache.py

@@ -36,14 +36,13 @@ class BlockKVCache:
     non-overlapping: each update adds one chunk of ``chunk_size`` tokens at the
     next logical position in the full sequence.
 
-    Note: Currently only supports ``total_size`` (``sink_size + window_size``) divisible by ``chunk_size``.
-
     Phases:
         - Filling: cache not yet full; tokens are written contiguously;
           ``cached_k()`` / ``cached_v()`` return only the valid prefix.
-        - Steady-state: cache full; each new chunk triggers a left-roll of the
-          local window and overwrites the rightmost positions;
-          ``cached_k()`` / ``cached_v()`` return the full buffer.
+        - Steady-state: if adding a chunk would exceed the fixed cache size, the
+          local window rolls left by the overflow amount and the new chunk
+          overwrites the rightmost positions; ``cached_k()`` / ``cached_v()``
+          return the full buffer.
 
     The argument ``chunk_idx`` (0, 1, 2, ...) is the index of the new chunk in the full
     sequence (not an index into the cache). If ``chunk_idx`` is greater than
@@ -150,10 +149,6 @@ def __post_init__(self) -> None:
             f"k_shape[seq_dim] ({self.k_shape[self.seq_dim]}) must equal 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 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)
 
@@ -163,24 +158,30 @@ def _seq_slice(self, start: int | None, end: int | None) -> tuple[slice | int, .
         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 tokens (steady-state only)."""
+    def _roll_local_window_left(self, shift_size: int) -> None:
+        """Shift valid local-window tokens left by ``shift_size`` tokens."""
         total_size = self._k.shape[self.seq_dim]
-        assert total_size == self._n_cached, (
-            f"Expected full cache: {total_size=} != {self._n_cached=}"
+        assert 0 < shift_size <= self.chunk_size, (
+            f"shift_size ({shift_size}) must be in (0, {self.chunk_size}]"
         )
-        tokens_to_keep = self.window_size - self.chunk_size
+        valid_end = min(self._n_cached, total_size)
+        valid_local_size = max(0, valid_end - self.sink_size)
+        tokens_to_keep = max(0, valid_local_size - shift_size)
 
         if tokens_to_keep > 0:
-            src_start = self.sink_size + self.chunk_size
-            src_end = total_size
+            src_start = self.sink_size + shift_size
+            src_end = src_start + tokens_to_keep
             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()
+        # before_update() only rolls when the next contiguous chunk would overflow
+        # this fixed buffer; update() must immediately write that chunk into the
+        # newly freed right edge.
+        self._n_cached = total_size
 
     def _current_chunk_overlaps_sink(self) -> bool:
         assert self._curr_chunk_idx is not None, (
@@ -286,8 +287,11 @@ def before_update(self, chunk_idx: int) -> None:
             "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()
+        total_size = self._k.shape[self.seq_dim]
+        if not self._current_chunk_overlaps_sink():
+            overflow = self._n_cached + self.chunk_size - total_size
+            if overflow > 0:
+                self._roll_local_window_left(overflow)
 
     def update(self, k: Tensor, v: Tensor) -> None:
         """
@@ -332,7 +336,8 @@ def after_update(self, chunk_idx: int) -> None:
             if self.is_steady_state():
                 pass
             else:
-                self._n_cached += self.chunk_size
+                total_size = self._k.shape[self.seq_dim]
+                self._n_cached = min(self._n_cached + self.chunk_size, total_size)
             self._prev_chunk_idx += 1
         elif self._curr_chunk_idx == self._prev_chunk_idx:
             pass

flashdreams/flashdreams/serving/realtime/input.py

@@ -135,21 +135,21 @@ class KeyboardResampler:
     def __init__(
         self,
         *,
-        fps: int,
+        fps: float,
         start_v: float = 0.0,
         supported_keys: frozenset[str] = DEFAULT_SUPPORTED_KEYS,
     ) -> None:
         if fps <= 0:
             raise ValueError("fps must be > 0")
-        self._fps = fps
-        self._dt = 1.0 / fps
+        self._fps = float(fps)
+        self._dt = 1.0 / self._fps
         self._supported_keys = supported_keys
         self.next_chunk_start_v = start_v
         self._event_log: deque[tuple[float, dict[str, str]]] = deque()
         self._carried_state = KeyboardState(supported_keys=supported_keys)
 
     @property
-    def fps(self) -> int:
+    def fps(self) -> float:
         return self._fps
 
     @property

flashdreams/flashdreams/serving/realtime/media.py

@@ -141,6 +141,18 @@ def tensor_chunk_to_rgb_frames(
         "minus_one_one" if video_chunk.is_floating_point() else "uint8"
     )
     if video_chunk.ndim == 4:
+        if video_chunk.shape[-1] == 3 and video_chunk.shape[1] != 3:
+            return rgb_array_to_uint8_frames(
+                video_chunk,
+                layout="thwc",
+                value_range=value_range,
+                sync_device=sync_device,
+            )
+        if video_chunk.shape[1] != 3:
+            raise ValueError(
+                "Expected video chunk [T, C, H, W] or [T, H, W, C] with 3 RGB "
+                f"channels, got {tuple(video_chunk.shape)}"
+            )
         return rgb_array_to_uint8_frames(
             video_chunk,
             layout="tchw",
@@ -155,7 +167,8 @@ def tensor_chunk_to_rgb_frames(
             sync_device=sync_device,
         )
     raise ValueError(
-        "Expected video chunk [T, C, H, W] or [1, 1, T, 3, H, W], "
+        "Expected video chunk [T, C, H, W], [T, H, W, C], "
+        "or [1, 1, T, 3, H, W], "
         f"got {tuple(video_chunk.shape)}"
     )
 

flashdreams/flashdreams/serving/webrtc/manager.py

@@ -53,11 +53,27 @@
 
 # How often the liveness watchdog wakes to re-check the elapsed-since-last-message.
 _CLIENT_LIVENESS_CHECK_INTERVAL_S = 1.0
+_DEFAULT_PERF_LOG_INTERVAL_CHUNKS = 5
 
 _RuntimeT = TypeVar("_RuntimeT", bound=WebRTCSessionRuntime)
 _RuntimeConfigT = TypeVar("_RuntimeConfigT", bound=WebRTCRuntimeConfig)
 
 
+def _stat_float(stats: dict[str, float], name: str, default: float = 0.0) -> float:
+    value = stats.get(name)
+    if value is None:
+        return default
+    return float(value)
+
+
+def _stat_ms(stats: dict[str, float], name: str, default_ms: float = 0.0) -> float:
+    return _stat_float(stats, name, default_ms / 1e3) * 1e3
+
+
+def _stat_int(stats: dict[str, float], name: str) -> int:
+    return int(round(_stat_float(stats, name)))
+
+
 class WebRTCControlSignal(IntEnum):
     """Rank-orchestration signals shared by the single-session runtimes."""
 
@@ -123,6 +139,7 @@ class BaseWebRTCSessionManager(Generic[_RuntimeT, _RuntimeConfigT]):
     _runtime_error_types: tuple[type[Exception], ...] = (RuntimeError,)
     _close_session_on_generation_error: bool = False
     _resampler_supported_keys: AbstractSet[str] | None = None
+    _perf_log_interval_chunks: int = _DEFAULT_PERF_LOG_INTERVAL_CHUNKS
 
     def __init__(
         self,
@@ -160,14 +177,72 @@ async def _reset_runtime_for_session(self, session_input: Any) -> None:
         await self._runtime.reset_for_new_session()
 
     def _make_resampler(self, *, start_v: float) -> KeyboardResampler:
+        return self._make_resampler_at_fps(start_v=start_v, fps=self.fps)
+
+    def _make_resampler_at_fps(
+        self, *, start_v: float, fps: float
+    ) -> KeyboardResampler:
         if self._resampler_supported_keys is None:
-            return KeyboardResampler(fps=self.fps, start_v=start_v)
+            return KeyboardResampler(fps=fps, start_v=start_v)
         return KeyboardResampler(
-            fps=self.fps,
+            fps=fps,
             start_v=start_v,
             supported_keys=frozenset(self._resampler_supported_keys),
         )
 
+    @staticmethod
+    def _positive_int_runtime_value(value: Any, *, label: str) -> int:
+        try:
+            parsed = int(value)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"{label} must be an integer.") from exc
+        if parsed <= 0:
+            raise ValueError(f"{label} must be > 0.")
+        return parsed
+
+    @staticmethod
+    def _positive_float_runtime_value(value: Any, *, label: str) -> float:
+        try:
+            parsed = float(value)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(f"{label} must be numeric.") from exc
+        if parsed <= 0.0:
+            raise ValueError(f"{label} must be > 0.")
+        return parsed
+
+    def _runtime_input_fps(self, runtime: Any) -> float:
+        method = getattr(runtime, "peek_input_fps", None)
+        if callable(method):
+            return self._positive_float_runtime_value(
+                method(),
+                label="peek_input_fps",
+            )
+        return float(self.fps)
+
+    def _runtime_next_input_num_frames(self, runtime: Any) -> int:
+        method = getattr(runtime, "peek_next_input_num_frames", None)
+        if callable(method):
+            return self._positive_int_runtime_value(
+                method(),
+                label="peek_next_input_num_frames",
+            )
+        return self._positive_int_runtime_value(
+            runtime.peek_next_chunk_num_frames(),
+            label="peek_next_chunk_num_frames",
+        )
+
+    def _runtime_steady_output_num_frames(self, runtime: Any) -> int:
+        method = getattr(runtime, "peek_steady_output_num_frames", None)
+        if callable(method):
+            return self._positive_int_runtime_value(
+                method(),
+                label="peek_steady_output_num_frames",
+            )
+        return self._positive_int_runtime_value(
+            runtime.peek_steady_chunk_num_frames(),
+            label="peek_steady_chunk_num_frames",
+        )
+
     def _register_extra_peer_handlers(self, peer_connection: Any) -> None:
         """Register optional extra peer-connection event handlers."""
 
@@ -289,13 +364,16 @@ async def _create_answer_with_runtime_ready_locked(
         # is throttled to the consumer's drain rate. AR step 0 emits fewer
         # frames than steady state; sizing to it would force a per-chunk
         # stall, so we size to the steady-state count.
-        num_frames = self._runtime.peek_steady_chunk_num_frames()
+        num_frames = self._runtime_steady_output_num_frames(self._runtime)
         video_track = BufferedVideoTrack(fps=self.fps, maxsize=num_frames)
         peer_connection.addTrack(video_track)
         # Start the resampler's virtual clock at 0; the real anchor is set
         # in the ``on_datachannel`` handler so chunk 0's window starts when
         # input can actually arrive.
-        resampler = self._make_resampler(start_v=0.0)
+        resampler = self._make_resampler_at_fps(
+            start_v=0.0,
+            fps=self._runtime_input_fps(self._runtime),
+        )
         loop = asyncio.get_running_loop()
         managed_session = ManagedWebRTCSession(
             runtime=self._runtime,
@@ -561,15 +639,19 @@ async def _generation_worker(
             "First action received; starting generation at start_v={:.3f}",
             resampler.next_chunk_start_v,
         )
+        perf_log_interval = max(0, int(self._perf_log_interval_chunks))
+        perf_window_start = loop.time()
+        perf_window_chunks = 0
+        perf_window_frames = 0
         try:
             while not managed_session.closed:
                 try:
-                    num_frames = runtime.peek_next_chunk_num_frames()
+                    input_num_frames = self._runtime_next_input_num_frames(runtime)
                 except self._runtime_error_types:
                     logger.exception("Runtime not ready; stopping generation worker.")
                     return
                 # Trigger when wallclock reaches the chunk's window end.
-                chunk_duration = num_frames * resampler.dt
+                chunk_duration = input_num_frames * resampler.dt
                 trigger_wall = resampler.next_chunk_start_v + chunk_duration
                 delay = trigger_wall - loop.time()
                 if delay > 0:
@@ -588,7 +670,7 @@ async def _generation_worker(
                     resampler.next_chunk_start_v = now - chunk_duration
 
                 t_before_gen = loop.time()
-                segments, frame_times = resampler.sample_chunk(num_frames)
+                segments, frame_times = resampler.sample_chunk(input_num_frames)
                 chunk_end_v = resampler.next_chunk_start_v
                 consumed_action_arrivals: list[float] = []
                 while (
@@ -624,11 +706,65 @@ async def _generation_worker(
                     if consumed_action_arrivals
                     else None
                 )
+                perf_window_chunks += 1
+                perf_window_frames += result.num_frames
+                if result.chunk_index == 0 or (
+                    perf_log_interval > 0
+                    and result.chunk_index % perf_log_interval == 0
+                ):
+                    interval_s = max(t_after_enqueue - perf_window_start, 1.0e-6)
+                    interval_fps = perf_window_frames / interval_s
+                    gen_fps = result.num_frames / max(
+                        t_after_gen - t_before_gen, 1.0e-6
+                    )
+                    stats = result.stats or {}
+                    logger.info(
+                        "WebRTC perf chunk={} interval_chunks={} frames={} "
+                        "gen_fps={:.1f} interval_fps={:.1f} playback_fps={} "
+                        "gen_ms={:.0f} enqueue_ms={:.0f} model_ms={:.0f} "
+                        "denoise_ms={:.0f} decode_ms={:.0f} pixel_post_ms={:.0f} "
+                        "copy_ms={:.0f} cache_ms={:.0f} "
+                        "cache_wait_ms={:.0f} cache_submit_ms={:.0f} "
+                        "queue_depth={} lag_ms={:.0f} control_latency_ms={} "
+                        "compile_active={} compile_start_step={} cuda_graph={} "
+                        "cache_frames={} cache_tokens={}",
+                        result.chunk_index,
+                        perf_window_chunks,
+                        perf_window_frames,
+                        gen_fps,
+                        interval_fps,
+                        video_track.fps,
+                        gen_ms,
+                        enqueue_ms,
+                        _stat_ms(stats, "model_step_s", gen_ms),
+                        _stat_ms(stats, "denoise_s"),
+                        _stat_ms(stats, "decode_s"),
+                        _stat_ms(stats, "pixel_post_s"),
+                        _stat_ms(stats, "gpu_to_cpu_copy_s"),
+                        _stat_ms(stats, "cache_seed_prune_s"),
+                        _stat_ms(stats, "cache_update_wait_s"),
+                        _stat_ms(stats, "cache_update_submit_s"),
+                        video_track.qsize(),
+                        lag_ms,
+                        "-"
+                        if control_latency_ms is None
+                        else f"{control_latency_ms:.0f}",
+                        _stat_int(stats, "compile_denoise_active"),
+                        _stat_int(stats, "compile_denoise_start_step"),
+                        _stat_int(stats, "cuda_graph_captured"),
+                        _stat_int(stats, "cache_frames"),
+                        _stat_int(stats, "cache_tokens"),
+                    )
+                    perf_window_start = t_after_enqueue
+                    perf_window_chunks = 0
+                    perf_window_frames = 0
                 logger.debug(
-                    "Chunk done chunk={} num_frames={} segments={} enqueued={} "
+                    "Chunk done chunk={} input_frames={} output_frames={} "
+                    "segments={} enqueued={} "
                     "gen_ms={:.1f} enqueue_ms={:.1f} play_ms={:.1f} queue_depth={} "
                     "lag_ms={:.1f}",
                     result.chunk_index,
+                    input_num_frames,
                     result.num_frames,
                     len(segments),
                     enqueued,

flashdreams/flashdreams/serving/webrtc/runtime.py

@@ -19,6 +19,7 @@ class WebRTCStepResult:
     """One generated chunk handed back by a WebRTC model runtime."""
 
     chunk_index: int
+    # Number of output video frames in ``video_chunk``.
     num_frames: int
     video_chunk: torch.Tensor
     stats: dict[str, float] | None
@@ -39,6 +40,15 @@ class WebRTCGenerationRuntime(Protocol):
     Integrations keep their model-specific state, checkpoints, conditioning,
     and cache logic inside their concrete runtime. The shared manager only
     needs this lifecycle and chunk-generation surface.
+
+    By default, ``peek_next_chunk_num_frames`` and
+    ``peek_steady_chunk_num_frames`` are used for both input sampling and
+    output queue sizing. Runtimes whose model input clock differs from their
+    output video clock may also implement these optional methods:
+
+    - ``peek_input_fps() -> float`` for the control/input sampling clock.
+    - ``peek_next_input_num_frames() -> int`` for the length of ``frame_times``.
+    - ``peek_steady_output_num_frames() -> int`` for video queue sizing.
     """
 
     async def initialize(self) -> None: ...