From 9d673191a9bcbf5d88dbeb719089582c5a015adc Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Wed, 22 Jul 2026 20:26:08 +0000 Subject: [PATCH 1/2] Support split WebRTC timing and THWC frame chunks Add shared WebRTC runtime hooks for integrations whose input sampling cadence differs from their output video cadence. The session manager now uses optional input FPS, input frame count, and steady output frame count hooks while keeping the legacy runtime methods as defaults. Add periodic WebRTC performance logging, support THWC uint8 realtime frame chunks, and relax block KV cache rolling so cache sizes do not need to be divisible by chunk size. Cover the shared behavior with focused realtime, WebRTC manager, and KV cache tests. --- .../flashdreams/core/attention/kvcache.py | 40 ++-- .../flashdreams/serving/realtime/input.py | 8 +- .../flashdreams/serving/realtime/media.py | 15 +- .../flashdreams/serving/webrtc/manager.py | 152 ++++++++++++++- .../flashdreams/serving/webrtc/runtime.py | 10 + flashdreams/tests/test_kvcache.py | 32 ++-- flashdreams/tests/test_realtime_serving.py | 12 ++ flashdreams/tests/test_webrtc_manager.py | 178 ++++++++++++++++++ 8 files changed, 399 insertions(+), 48 deletions(-) diff --git a/flashdreams/flashdreams/core/attention/kvcache.py b/flashdreams/flashdreams/core/attention/kvcache.py index 6b6ae639..25d02429 100644 --- a/flashdreams/flashdreams/core/attention/kvcache.py +++ b/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,17 +158,19 @@ 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 @@ -181,6 +178,7 @@ def _roll_local_window_left(self) -> None: 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() + self._n_cached = total_size def _current_chunk_overlaps_sink(self) -> bool: assert self._curr_chunk_idx is not None, ( @@ -286,8 +284,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 +333,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 diff --git a/flashdreams/flashdreams/serving/realtime/input.py b/flashdreams/flashdreams/serving/realtime/input.py index 1b488bb6..6baf92dd 100644 --- a/flashdreams/flashdreams/serving/realtime/input.py +++ b/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 diff --git a/flashdreams/flashdreams/serving/realtime/media.py b/flashdreams/flashdreams/serving/realtime/media.py index 266e7ac5..3be277d0 100644 --- a/flashdreams/flashdreams/serving/realtime/media.py +++ b/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: + 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)}" ) diff --git a/flashdreams/flashdreams/serving/webrtc/manager.py b/flashdreams/flashdreams/serving/webrtc/manager.py index d0b7d896..6ad20743 100644 --- a/flashdreams/flashdreams/serving/webrtc/manager.py +++ b/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, diff --git a/flashdreams/flashdreams/serving/webrtc/runtime.py b/flashdreams/flashdreams/serving/webrtc/runtime.py index d9cf0d03..55c1b6a5 100644 --- a/flashdreams/flashdreams/serving/webrtc/runtime.py +++ b/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: ... diff --git a/flashdreams/tests/test_kvcache.py b/flashdreams/tests/test_kvcache.py index 6e0ee738..90042e33 100644 --- a/flashdreams/tests/test_kvcache.py +++ b/flashdreams/tests/test_kvcache.py @@ -91,20 +91,16 @@ def update(self, chunk_idx: int, k: torch.Tensor, v: torch.Tensor) -> None: self._cache_v[:, -length:] = v return - if self._cache_k.shape[1] == self.total_size: - tail_k = self._cache_k[:, self.sink_size + self.chunk_size :] - tail_v = self._cache_v[:, self.sink_size + self.chunk_size :] - window_k = torch.cat([tail_k, k], dim=1)[:, -self.window_size :] - window_v = torch.cat([tail_v, v], dim=1)[:, -self.window_size :] - self._cache_k = torch.cat( - [self._cache_k[:, : self.sink_size], window_k], dim=1 - ) - self._cache_v = torch.cat( - [self._cache_v[:, : self.sink_size], window_v], dim=1 - ) - else: - self._cache_k = torch.cat([self._cache_k, k], dim=1) - self._cache_v = torch.cat([self._cache_v, v], dim=1) + sink_k = self._cache_k[:, : self.sink_size] + sink_v = self._cache_v[:, : self.sink_size] + window_k = torch.cat([self._cache_k[:, self.sink_size :], k], dim=1)[ + :, -self.window_size : + ] + window_v = torch.cat([self._cache_v[:, self.sink_size :], v], dim=1)[ + :, -self.window_size : + ] + self._cache_k = torch.cat([sink_k, window_k], dim=1) + self._cache_v = torch.cat([sink_v, window_v], dim=1) self._prev_chunk_idx += 1 def cached_k(self) -> torch.Tensor: @@ -176,7 +172,9 @@ def dtype() -> torch.dtype: @pytest.mark.ci_cpu -@pytest.mark.parametrize("sink_size,window_size", [(0, 8), (0, 24), (3, 5), (3, 21)]) +@pytest.mark.parametrize( + "sink_size,window_size", [(0, 8), (0, 24), (3, 5), (3, 21), (1, 16)] +) def test_block_kvcache_matches_baseline( device: torch.device, dtype: torch.dtype, @@ -350,7 +348,9 @@ def test_kvcache_relative_rope_cp_freqs_match_cache_chunks() -> None: not torch.cuda.is_available(), reason="CUDA required for cudagraph test" ) @pytest.mark.ci_gpu -@pytest.mark.parametrize("sink_size,window_size", [(0, 8), (0, 24), (3, 5), (3, 21)]) +@pytest.mark.parametrize( + "sink_size,window_size", [(0, 8), (0, 24), (3, 5), (3, 21), (1, 16)] +) def test_block_kvcache_cudagraph_matches_baseline( dtype: torch.dtype, sink_size: int, diff --git a/flashdreams/tests/test_realtime_serving.py b/flashdreams/tests/test_realtime_serving.py index 5c89d6c3..1a0646f9 100644 --- a/flashdreams/tests/test_realtime_serving.py +++ b/flashdreams/tests/test_realtime_serving.py @@ -144,6 +144,18 @@ def test_realtime_media_supports_omnidreams_uint8_layout() -> None: assert frames[1][0, 0, 0] == 255 +def test_realtime_media_supports_thwc_uint8_chunks() -> None: + chunk = torch.zeros((2, 4, 5, 3), dtype=torch.uint8) + chunk[1, 0, 0] = torch.tensor([10, 20, 30], dtype=torch.uint8) + + frames = tensor_chunk_to_rgb_frames(chunk) + + assert len(frames) == 2 + assert frames[0].shape == (4, 5, 3) + assert frames[0].dtype == np.uint8 + assert frames[1][0, 0].tolist() == [10, 20, 30] + + @pytest.mark.parametrize( "chunk", [ diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index 69bb05ce..439293cf 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -12,6 +12,7 @@ import torch from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS +from flashdreams.serving.webrtc import manager as manager_module from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, ManagedWebRTCSession, @@ -75,6 +76,11 @@ def sample_chunk( return [(0.0, 0.0, frozenset({"w"}))], [0.0] +class _CountingVideoTrack(_FakeVideoTrack): + async def enqueue_chunk(self, chunk: Any) -> int: + return int(chunk.shape[0]) + + class _BaseTestManager(BaseWebRTCSessionManager): def _model_name(self) -> str: return "fake-model" @@ -90,6 +96,45 @@ def _make_manager( ) +def test_runtime_frame_timing_hooks_default_to_legacy_methods() -> None: + class _LegacyRuntime: + def peek_next_chunk_num_frames(self) -> int: + return 2 + + def peek_steady_chunk_num_frames(self) -> int: + return 3 + + runtime = _LegacyRuntime() + manager = _make_manager(_BaseTestManager, runtime) + + assert manager._runtime_input_fps(runtime) == pytest.approx(30.0) + assert manager._runtime_next_input_num_frames(runtime) == 2 + assert manager._runtime_steady_output_num_frames(runtime) == 3 + + +def test_runtime_frame_timing_hooks_can_split_input_and_output() -> None: + class _SplitRuntime: + def peek_input_fps(self) -> float: + return 6.0 + + def peek_next_input_num_frames(self) -> int: + return 4 + + def peek_steady_output_num_frames(self) -> int: + return 16 + + runtime = _SplitRuntime() + manager = _make_manager(_BaseTestManager, runtime) + resampler = manager._make_resampler_at_fps( + start_v=0.0, + fps=manager._runtime_input_fps(runtime), + ) + + assert resampler.dt == pytest.approx(1.0 / 6.0) + assert manager._runtime_next_input_num_frames(runtime) == 4 + assert manager._runtime_steady_output_num_frames(runtime) == 16 + + def _managed_session( runtime: Any, ) -> tuple[ManagedWebRTCSession, _FakeVideoTrack, _FakePeerConnection, _FakeChannel]: @@ -231,6 +276,139 @@ def _chunk_done_extra(self) -> dict[str, Any]: assert payload["resolution"] == {"width": 8, "height": 4} +@pytest.mark.asyncio +async def test_generation_worker_uses_split_input_and_output_frame_counts() -> None: + class _SplitResampler: + dt = 0.0 + next_chunk_start_v = 0.0 + + def __init__(self) -> None: + self.sampled_num_frames: list[int] = [] + + def sample_chunk( + self, num_frames: int + ) -> tuple[list[tuple[float, float, frozenset[str]]], list[float]]: + self.sampled_num_frames.append(num_frames) + return ( + [(0.0, 0.0, frozenset({"w"}))], + [float(index) for index in range(num_frames)], + ) + + class _SplitRuntime: + def __init__(self) -> None: + self.managed_session: ManagedWebRTCSession | None = None + self.frame_times: list[float] | None = None + + def peek_input_fps(self) -> float: + return 6.0 + + def peek_next_input_num_frames(self) -> int: + return 2 + + async def generate_chunk( + self, *, segments: Any, frame_times: list[float] + ) -> WebRTCStepResult: + del segments + self.frame_times = frame_times + if self.managed_session is not None: + self.managed_session.closed = True + return WebRTCStepResult( + chunk_index=0, + num_frames=5, + video_chunk=torch.zeros((5, 1, 1, 3, 2, 2), dtype=torch.uint8), + stats=None, + ) + + runtime = _SplitRuntime() + manager = _make_manager(_BaseTestManager, runtime) + managed, _video_track, _peer, channel = _managed_session(runtime) + resampler = _SplitResampler() + managed.video_track = _CountingVideoTrack() # ty:ignore[invalid-assignment] + managed.resampler = resampler # ty:ignore[invalid-assignment] + runtime.managed_session = managed + manager._active_session = managed + + task = asyncio.create_task(manager._generation_worker(managed_session=managed)) + managed.generation_task = task + await asyncio.wait_for(task, timeout=5.0) + + assert resampler.sampled_num_frames == [2] + assert runtime.frame_times == [0.0, 1.0] + chunk_done = [ + json.loads(message) + for message in channel.messages + if json.loads(message).get("type") == "chunk_done" + ] + assert len(chunk_done) == 1 + assert chunk_done[0]["num_frames"] == 5 + assert chunk_done[0]["enqueued_frames"] == 5 + assert chunk_done[0]["play_ms"] == pytest.approx(5 * 1000 / 30, abs=0.05) + + +@pytest.mark.asyncio +async def test_generation_worker_logs_periodic_perf_stats( + monkeypatch: pytest.MonkeyPatch, +) -> None: + perf_logs: list[tuple[str, tuple[Any, ...]]] = [] + + def _record_info(message: str, *args: Any, **_kwargs: Any) -> None: + if message.startswith("WebRTC perf"): + perf_logs.append((message, args)) + + class _StatsRuntime: + def __init__(self) -> None: + self.managed_session: ManagedWebRTCSession | None = None + self.chunk_index = 0 + + def peek_next_chunk_num_frames(self) -> int: + return 1 + + async def generate_chunk( + self, *, segments: Any, frame_times: Any + ) -> WebRTCStepResult: + del segments, frame_times + chunk_index = self.chunk_index + self.chunk_index += 1 + if chunk_index >= 2 and self.managed_session is not None: + self.managed_session.closed = True + return WebRTCStepResult( + chunk_index=chunk_index, + num_frames=4, + video_chunk=torch.zeros((4, 1, 1, 3, 2, 2), dtype=torch.uint8), + stats={ + "model_step_s": 0.02, + "denoise_s": 0.01, + "decode_s": 0.004, + "pixel_post_s": 0.003, + "gpu_to_cpu_copy_s": 0.002, + "compile_denoise_active": 1.0, + "compile_denoise_start_step": 3.0, + "cache_frames": 13.0, + "cache_tokens": 512.0, + }, + ) + + class _FrequentLogManager(_BaseTestManager): + _perf_log_interval_chunks = 2 + + monkeypatch.setattr(manager_module.logger, "info", _record_info) + runtime = _StatsRuntime() + manager = _make_manager(_FrequentLogManager, runtime) + managed, _video_track, _peer, _channel = _managed_session(runtime) + runtime.managed_session = managed + manager._active_session = managed + + task = asyncio.create_task(manager._generation_worker(managed_session=managed)) + managed.generation_task = task + await asyncio.wait_for(task, timeout=5.0) + + assert [args[0] for _message, args in perf_logs] == [0, 2] + assert "compile_active" in perf_logs[0][0] + assert "pixel_post_ms" in perf_logs[0][0] + assert "copy_ms" in perf_logs[0][0] + assert perf_logs[0][1][-2:] == (13, 512) + + @pytest.mark.asyncio async def test_create_answer_raises_busy_with_subclass_message() -> None: class _BusyManager(_BaseTestManager): From f9d377f50d45055599b7c4978e7650f53244c0ec Mon Sep 17 00:00:00 2001 From: Jesse Archer Date: Wed, 22 Jul 2026 21:32:52 +0000 Subject: [PATCH 2/2] Address GitHub PR review feedback Tighten THWC chunk detection so legacy TCHW tensors with width 3 are not misclassified, and add a regression test for that ambiguous shape. Document the BlockKVCache roll/update invariant and include Ruff's import ordering fix from the failed CPU lint job. --- flashdreams/flashdreams/core/attention/kvcache.py | 3 +++ flashdreams/flashdreams/serving/realtime/media.py | 2 +- flashdreams/tests/test_realtime_serving.py | 11 +++++++++++ flashdreams/tests/test_webrtc_manager.py | 2 +- 4 files changed, 16 insertions(+), 2 deletions(-) diff --git a/flashdreams/flashdreams/core/attention/kvcache.py b/flashdreams/flashdreams/core/attention/kvcache.py index 25d02429..5673a32a 100644 --- a/flashdreams/flashdreams/core/attention/kvcache.py +++ b/flashdreams/flashdreams/core/attention/kvcache.py @@ -178,6 +178,9 @@ def _roll_local_window_left(self, shift_size: int) -> None: 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: diff --git a/flashdreams/flashdreams/serving/realtime/media.py b/flashdreams/flashdreams/serving/realtime/media.py index 3be277d0..9b93eaca 100644 --- a/flashdreams/flashdreams/serving/realtime/media.py +++ b/flashdreams/flashdreams/serving/realtime/media.py @@ -141,7 +141,7 @@ 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: + if video_chunk.shape[-1] == 3 and video_chunk.shape[1] != 3: return rgb_array_to_uint8_frames( video_chunk, layout="thwc", diff --git a/flashdreams/tests/test_realtime_serving.py b/flashdreams/tests/test_realtime_serving.py index 1a0646f9..33b21476 100644 --- a/flashdreams/tests/test_realtime_serving.py +++ b/flashdreams/tests/test_realtime_serving.py @@ -156,6 +156,17 @@ def test_realtime_media_supports_thwc_uint8_chunks() -> None: assert frames[1][0, 0].tolist() == [10, 20, 30] +def test_realtime_media_prefers_tchw_when_width_is_three() -> None: + chunk = torch.zeros((2, 3, 4, 3), dtype=torch.uint8) + chunk[1, :, 0, 0] = torch.tensor([10, 20, 30], dtype=torch.uint8) + + frames = tensor_chunk_to_rgb_frames(chunk) + + assert len(frames) == 2 + assert frames[1].shape == (4, 3, 3) + assert frames[1][0, 0].tolist() == [10, 20, 30] + + @pytest.mark.parametrize( "chunk", [ diff --git a/flashdreams/tests/test_webrtc_manager.py b/flashdreams/tests/test_webrtc_manager.py index 439293cf..ebcafdbb 100644 --- a/flashdreams/tests/test_webrtc_manager.py +++ b/flashdreams/tests/test_webrtc_manager.py @@ -11,8 +11,8 @@ import pytest import torch -from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS from flashdreams.serving.webrtc import manager as manager_module +from flashdreams.serving.webrtc.controls import WSAD_SUPPORTED_KEYS from flashdreams.serving.webrtc.manager import ( BaseWebRTCSessionManager, ManagedWebRTCSession,