diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py index a15b46b3a..048bd2349 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py @@ -2,50 +2,28 @@ from __future__ import annotations -import os -import tempfile from typing import TYPE_CHECKING, Any -from fastapi import WebSocket - from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( RealtimeEvent, RealtimeVideoGenerationsRequest, ) from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import ( + BaseRealtimeModelAdapter, RealtimeChunkInputs, - RealtimeModelAdapter, + build_realtime_sampling_params, + save_realtime_first_frame, ) -from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_output_adapter import ( - RawRGBRealtimeOutputAdapter, - RealtimeFrameSendStats, -) -from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( - build_sampling_params, - save_image_to_path, -) -from sglang.multimodal_gen.runtime.entrypoints.utils import ( - prepare_request, -) -from sglang.multimodal_gen.runtime.realtime.camera_controls import ( +from sglang.multimodal_gen.runtime.realtime.control_signals import ControlSignalQueue +from sglang.multimodal_gen.runtime.realtime.states import ( RealtimeCameraControlState, ) -from sglang.multimodal_gen.runtime.realtime.condition_events import ( - ConditionEvent, - ConditionEventQueue, - ControlSignal, -) -from sglang.multimodal_gen.runtime.server_args import get_global_server_args if TYPE_CHECKING: from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import ( GenerateSession, RealtimeChunkContext, ) - from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import ( - OutputBatch, - Req, - ) from sglang.multimodal_gen.runtime.server_args import ServerArgs @@ -60,53 +38,39 @@ class LingBotWorldRealtimeState(RealtimeCameraControlState): script_maxlen=512, max_transitions=512, ) - self.events = ConditionEventQueue(max_events={"prompt": 1}) + self.prompt_queue = ControlSignalQueue(max_events={"prompt": 1}) def clear(self) -> None: super().clear() - self.events.clear() + self.prompt_queue.clear() def receive_prompt(self, prompt: str, *, event_id: int | None = None) -> None: - self.events.push( - ConditionEvent( - kind="prompt", - payload=ControlSignal( - kind="prompt", - payload=prompt, - seq_id=event_id, - ), - ) - ) + self.prompt_queue.push("prompt", prompt, event_id=event_id) - def receive_camera_event_payload( + def receive_camera_control_event_payload( self, payload: Any, *, event_id: int | None, ) -> str: - return super().receive_camera_event_payload( + return super().receive_camera_control_event_payload( payload, event_id=event_id, validate_camera_actions=LingBotWorldRealtimeAdapter._validate_camera_actions, ) def sample_prompt(self) -> str: - prompt = self.events.pop_latest("prompt") + prompt = self.prompt_queue.pop_latest("prompt") if not isinstance(prompt, str): raise ValueError("prompt event payload must be a string") - self.latest_sampled_event_id = self.events.last_sampled_seq_id("prompt") + self.latest_sampled_event_id = self.prompt_queue.last_sampled_seq_id("prompt") return prompt def has_prompt(self) -> bool: - return self.events.has_events("prompt") + return self.prompt_queue.has_events("prompt") -class LingBotWorldRealtimeAdapter(RealtimeModelAdapter): - name = "lingbot_world" - - def __init__(self): - self.output_adapter = RawRGBRealtimeOutputAdapter() - +class LingBotWorldRealtimeAdapter(BaseRealtimeModelAdapter): def create_state(self) -> LingBotWorldRealtimeState: return LingBotWorldRealtimeState() @@ -125,26 +89,11 @@ class LingBotWorldRealtimeAdapter(RealtimeModelAdapter): camera_actions = condition_inputs.get("camera_actions") if camera_actions is not None: state = self._state(session) - state.receive_camera_script(self._validate_camera_actions(camera_actions)) + state.receive_camera_action_script( + self._validate_camera_actions(camera_actions) + ) - if request.first_frame is None: - return - - server_args = get_global_server_args() - if server_args.input_save_path is not None: - uploads_dir = server_args.input_save_path - os.makedirs(uploads_dir, exist_ok=True) - else: - if session.input_temp_dir is None: - session.input_temp_dir = tempfile.mkdtemp(prefix="sglang_input_") - uploads_dir = session.input_temp_dir - - target_path = os.path.join(uploads_dir, f"{session.id}_first_frame") - image_path = await save_image_to_path(request.first_frame, target_path) - request.first_frame = image_path - - async def wait_for_next_chunk(self, session: GenerateSession) -> None: - del session + await save_realtime_first_frame(session, request) @staticmethod def _validate_camera_actions(payload: Any) -> list[list[str]]: @@ -164,7 +113,7 @@ class LingBotWorldRealtimeAdapter(RealtimeModelAdapter): ) -> str: state = self._state(session) if event.kind == "camera_actions": - return state.receive_camera_event_payload( + return state.receive_camera_control_event_payload( event.payload, event_id=event.event_id, ) @@ -175,9 +124,10 @@ class LingBotWorldRealtimeAdapter(RealtimeModelAdapter): return f"kind=prompt, prompt_len={len(event.payload)}" raise ValueError(f"unsupported event kind: {event.kind}") - def _sample_chunk_inputs( + def sample_chunk_inputs( self, session: GenerateSession, + server_args: ServerArgs, chunk: RealtimeChunkContext, chunk_size: int, ) -> RealtimeChunkInputs: @@ -201,13 +151,13 @@ class LingBotWorldRealtimeAdapter(RealtimeModelAdapter): condition_inputs["camera_actions"] = camera_actions return RealtimeChunkInputs(prompt=prompt, condition_inputs=condition_inputs) - def _build_sampling_params( + def build_sampling_params( self, session: GenerateSession, + server_args: ServerArgs, chunk: RealtimeChunkContext, chunk_inputs: RealtimeChunkInputs, chunk_size: int, - server_args: ServerArgs, ): request = session.request if request is None: @@ -219,42 +169,16 @@ class LingBotWorldRealtimeAdapter(RealtimeModelAdapter): chunk_size=chunk_size, ) - return build_sampling_params( + return build_realtime_sampling_params( chunk.request_id, - prompt=chunk_inputs.prompt, - size=request.size, + request=request, + chunk_inputs=chunk_inputs, num_frames=num_frames, - fps=request.fps, - image_path=request.first_frame, - output_file_name=chunk.request_id, - save_output=False, - seed=request.seed, - generator_device=request.generator_device, num_inference_steps=( request.num_inference_steps or LINGBOT_REALTIME_DEFAULT_NUM_INFERENCE_STEPS ), - guidance_scale=request.guidance_scale, - guidance_scale_2=request.guidance_scale_2, - negative_prompt=request.negative_prompt, - enable_teacache=request.enable_teacache, - enable_frame_interpolation=request.enable_frame_interpolation, - frame_interpolation_exp=request.frame_interpolation_exp, - frame_interpolation_scale=request.frame_interpolation_scale, - frame_interpolation_model_path=request.frame_interpolation_model_path, - enable_upscaling=request.enable_upscaling, - upscaling_model_path=request.upscaling_model_path, - upscaling_scale=request.upscaling_scale, - diffusers_kwargs=request.diffusers_kwargs, - profile=request.profile, - num_profiled_timesteps=request.num_profiled_timesteps, - profile_all_stages=request.profile_all_stages, - perf_dump_path=request.perf_dump_path, - output_path=request.output_path, - output_compression=request.output_compression, - output_quality=request.output_quality, - condition_inputs=chunk_inputs.condition_inputs, - realtime_chunk_size=chunk_size, + chunk_size=chunk_size, ) @staticmethod @@ -276,59 +200,10 @@ class LingBotWorldRealtimeAdapter(RealtimeModelAdapter): required_num_frames = (required_latent_frames - 1) * temporal_ratio + 1 return max(int(request.num_frames or 0), required_num_frames) - def prepare_next_request( - self, - session: GenerateSession, - server_args: ServerArgs, - chunk: RealtimeChunkContext, - ) -> Req: - """build a new request for the next chunk""" - pipeline_config = server_args.pipeline_config - chunk_size = int(pipeline_config.dit_config.arch_config.num_frames_per_block) - chunk_inputs = self._sample_chunk_inputs(session, chunk, chunk_size) - sampling_params = self._build_sampling_params( - session, - chunk, - chunk_inputs, - chunk_size, - server_args, - ) - batch = prepare_request( - server_args=server_args, - sampling_params=sampling_params, - ) - batch.session = session.realtime_session - batch.realtime_session_id = session.id - batch.return_raw_frames = True - batch.block_idx = chunk.index - batch.realtime_event_id = self._state(session).latest_sampled_event_id - if session.request is not None: - batch.realtime_output_format = session.request.realtime_output_format - batch.realtime_preview_max_width = ( - session.request.realtime_preview_max_width - ) - batch.realtime_output_pacing = bool(session.request.realtime_output_pacing) - batch.realtime_causal_sink_size = session.request.realtime_causal_sink_size - batch.realtime_causal_kv_cache_num_frames = ( - session.request.realtime_causal_kv_cache_num_frames - ) - return batch + def get_realtime_event_id(self, session: GenerateSession) -> int | None: + return self._state(session).latest_sampled_event_id - async def send_output( - self, - ws: WebSocket, - session: GenerateSession, - result: OutputBatch, - batch: Req, - ) -> RealtimeFrameSendStats: - return await self.output_adapter.send(ws, session, result, batch) - - def on_chunk_complete(self, session: GenerateSession, result: OutputBatch) -> None: - del result - session.generate_chunk_completed() - - def dispose(self, session: GenerateSession) -> None: + def clear_state(self, session: GenerateSession) -> None: state = session.adapter_state if isinstance(state, LingBotWorldRealtimeState): state.clear() - self.output_adapter.reset() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/sana_wm_realtime_adapter.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/sana_wm_realtime_adapter.py index 3ba2967aa..dd8e40948 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/sana_wm_realtime_adapter.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/sana_wm_realtime_adapter.py @@ -2,30 +2,18 @@ from __future__ import annotations -import hashlib -import os -import tempfile from typing import TYPE_CHECKING, Any -from fastapi import WebSocket - from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( RealtimeEvent, RealtimeVideoGenerationsRequest, ) from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import ( + BaseRealtimeModelAdapter, RealtimeChunkInputs, - RealtimeModelAdapter, + build_realtime_sampling_params, + save_realtime_first_frame, ) -from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_output_adapter import ( - RawRGBRealtimeOutputAdapter, - RealtimeFrameSendStats, -) -from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( - build_sampling_params, - save_image_to_path, -) -from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.base import ( normalize_sana_wm_camera_actions, parse_sana_wm_action_string, @@ -34,10 +22,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.s from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.self_forcing import ( SanaWMSelfForcingSampler, ) -from sglang.multimodal_gen.runtime.realtime.camera_controls import ( +from sglang.multimodal_gen.runtime.realtime.states import ( RealtimeCameraControlState, ) -from sglang.multimodal_gen.runtime.server_args import get_global_server_args if TYPE_CHECKING: from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import ( @@ -46,7 +33,6 @@ if TYPE_CHECKING: ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import ( OutputBatch, - Req, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs @@ -77,25 +63,20 @@ class SanaWMRealtimeAdapterState(RealtimeCameraControlState): super().clear() self.base_condition_inputs.clear() - def receive_camera_event_payload( + def receive_camera_control_event_payload( self, payload: Any, *, event_id: int | None, ) -> str: - return super().receive_camera_event_payload( + return super().receive_camera_control_event_payload( payload, event_id=event_id, validate_camera_actions=SanaWMRealtimeAdapter._validate_camera_actions, ) -class SanaWMRealtimeAdapter(RealtimeModelAdapter): - name = "sana_wm_realtime" - - def __init__(self): - self.output_adapter = RawRGBRealtimeOutputAdapter() - +class SanaWMRealtimeAdapter(BaseRealtimeModelAdapter): def create_state(self) -> SanaWMRealtimeAdapterState: return SanaWMRealtimeAdapterState() @@ -122,9 +103,6 @@ class SanaWMRealtimeAdapter(RealtimeModelAdapter): session: GenerateSession, request: RealtimeVideoGenerationsRequest, ) -> None: - if request.first_frame is None: - raise ValueError("SANA-WM realtime requires first_frame") - request.size = request.size or SANA_WM_DEFAULT_SIZE if request.num_frames is not None: request.num_frames = int(request.num_frames) @@ -157,42 +135,21 @@ class SanaWMRealtimeAdapter(RealtimeModelAdapter): if camera_actions is not None and action is not None: raise ValueError("pass only one of camera_actions or action") if camera_actions is not None: - state.receive_camera_event_payload(camera_actions, event_id=None) + state.receive_camera_control_event_payload(camera_actions, event_id=None) if action is not None: if not isinstance(action, str) or not action: raise ValueError("action condition input must be a non-empty string") - state.receive_camera_script( + state.receive_camera_action_script( parse_sana_wm_action_string(action), event_id=None ) state.base_condition_inputs = condition_inputs - server_args = get_global_server_args() - if server_args.input_save_path is not None: - uploads_dir = server_args.input_save_path - os.makedirs(uploads_dir, exist_ok=True) - else: - if session.input_temp_dir is None: - session.input_temp_dir = tempfile.mkdtemp(prefix="sglang_input_") - uploads_dir = session.input_temp_dir - - if isinstance( - request.first_frame, str - ) and request.first_frame.lower().startswith(("http://", "https://")): - suffix = os.path.splitext(request.first_frame.split("?", 1)[0])[1] - digest = hashlib.sha256(request.first_frame.encode("utf-8")).hexdigest()[ - :16 - ] - target_path = os.path.join(uploads_dir, f"realtime_ref_{digest}{suffix}") - if os.path.exists(target_path): - request.first_frame = target_path - return - else: - target_path = os.path.join(uploads_dir, f"{session.id}_first_frame") - image_path = await save_image_to_path(request.first_frame, target_path) - request.first_frame = image_path - - async def wait_for_next_chunk(self, session: GenerateSession) -> None: - del session + await save_realtime_first_frame( + session, + request, + required_error="SANA-WM realtime requires first_frame", + cache_remote_urls=True, + ) def ingest_event( self, @@ -201,7 +158,7 @@ class SanaWMRealtimeAdapter(RealtimeModelAdapter): ) -> str: state = self._state(session) if event.kind == "camera_actions": - return state.receive_camera_event_payload( + return state.receive_camera_control_event_payload( event.payload, event_id=event.event_id, ) @@ -209,16 +166,23 @@ class SanaWMRealtimeAdapter(RealtimeModelAdapter): if not isinstance(event.payload, str) or not event.payload: raise ValueError("action event payload must be a non-empty string") camera_actions = parse_sana_wm_action_string(event.payload) - state.receive_camera_script(camera_actions, event_id=event.event_id) + state.receive_camera_action_script(camera_actions, event_id=event.event_id) return f"kind=action, frames={len(camera_actions)}" raise ValueError(f"unsupported event kind: {event.kind}") - def _sample_chunk_inputs( + def sample_chunk_inputs( self, session: GenerateSession, + server_args: ServerArgs, chunk: RealtimeChunkContext, - action_chunk_size: int, + chunk_size: int, ) -> RealtimeChunkInputs: + action_chunk_size = self._action_chunk_size( + session, + server_args, + chunk, + chunk_size, + ) state = self._state(session) request = session.request if request is None: @@ -233,9 +197,10 @@ class SanaWMRealtimeAdapter(RealtimeModelAdapter): condition_inputs=condition_inputs, ) - def _build_sampling_params( + def build_sampling_params( self, session: GenerateSession, + server_args: ServerArgs, chunk: RealtimeChunkContext, chunk_inputs: RealtimeChunkInputs, chunk_size: int, @@ -244,49 +209,22 @@ class SanaWMRealtimeAdapter(RealtimeModelAdapter): if request is None: raise ValueError("realtime request is not initialized") - return build_sampling_params( + return build_realtime_sampling_params( chunk.request_id, - prompt=chunk_inputs.prompt, - size=request.size, + request=request, + chunk_inputs=chunk_inputs, num_frames=request.num_frames, - fps=request.fps, - image_path=request.first_frame, - output_file_name=chunk.request_id, - save_output=False, - seed=request.seed, - generator_device=request.generator_device, num_inference_steps=request.num_inference_steps, - guidance_scale=request.guidance_scale, - guidance_scale_2=request.guidance_scale_2, - negative_prompt=request.negative_prompt, - enable_teacache=request.enable_teacache, - enable_frame_interpolation=request.enable_frame_interpolation, - frame_interpolation_exp=request.frame_interpolation_exp, - frame_interpolation_scale=request.frame_interpolation_scale, - frame_interpolation_model_path=request.frame_interpolation_model_path, - enable_upscaling=request.enable_upscaling, - upscaling_model_path=request.upscaling_model_path, - upscaling_scale=request.upscaling_scale, - diffusers_kwargs=request.diffusers_kwargs, - profile=request.profile, - num_profiled_timesteps=request.num_profiled_timesteps, - profile_all_stages=request.profile_all_stages, - perf_dump_path=request.perf_dump_path, - output_path=request.output_path, - output_compression=request.output_compression, - output_quality=request.output_quality, - condition_inputs=chunk_inputs.condition_inputs, - realtime_chunk_size=chunk_size, + chunk_size=chunk_size, ) - def prepare_next_request( + def _action_chunk_size( self, session: GenerateSession, server_args: ServerArgs, chunk: RealtimeChunkContext, - ) -> Req: - arch_config = server_args.pipeline_config.dit_config.arch_config - chunk_size = int(getattr(arch_config, "num_frames_per_block", 3)) + chunk_size: int, + ) -> int: temporal_compression = int( server_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio ) @@ -310,50 +248,17 @@ class SanaWMRealtimeAdapter(RealtimeModelAdapter): action_chunk_size = ( segments[idx + 1] - segments[idx] ) * temporal_compression - chunk_inputs = self._sample_chunk_inputs(session, chunk, action_chunk_size) - sampling_params = self._build_sampling_params( - session, - chunk, - chunk_inputs, - chunk_size, - ) - batch = prepare_request( - server_args=server_args, - sampling_params=sampling_params, - ) - batch.session = session.realtime_session - batch.realtime_session_id = session.id - batch.return_raw_frames = True - batch.block_idx = chunk.index - batch.realtime_event_id = self._state(session).latest_sampled_event_id - if session.request is not None: - # Forward the full transport config like the LingBot adapter does — - # the shared RawRGB output adapter / realtime_video_api consume - # preview width + pacing too; dropping them silently disabled both - # features for SANA-WM sessions. - batch.realtime_output_format = session.request.realtime_output_format - batch.realtime_preview_max_width = ( - session.request.realtime_preview_max_width - ) - batch.realtime_output_pacing = bool(session.request.realtime_output_pacing) - return batch + return action_chunk_size - async def send_output( - self, - ws: WebSocket, - session: GenerateSession, - result: OutputBatch, - batch: Req, - ) -> RealtimeFrameSendStats: - return await self.output_adapter.send(ws, session, result, batch) + def get_realtime_event_id(self, session: GenerateSession) -> int | None: + return self._state(session).latest_sampled_event_id def on_chunk_complete(self, session: GenerateSession, result: OutputBatch) -> None: if session.request is not None and self._raw_frame_count(result) == 0: session.request.max_chunks = session.generate_chunk_cnt + 1 session.generate_chunk_completed() - def dispose(self, session: GenerateSession) -> None: + def clear_state(self, session: GenerateSession) -> None: state = session.adapter_state if isinstance(state, SanaWMRealtimeAdapterState): state.clear() - self.output_adapter.reset() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/generate_session.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/generate_session.py index 6df62efa7..d59fea048 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/generate_session.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/generate_session.py @@ -15,7 +15,7 @@ from sglang.multimodal_gen.runtime.realtime.session import ( if TYPE_CHECKING: from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import ( - RealtimeModelAdapter, + BaseRealtimeModelAdapter, ) @@ -27,6 +27,8 @@ class RealtimeChunkContext: class GenerateSession: + """A realtime generation session""" + def __init__(self): self.id = uuid4().hex self.request: RealtimeVideoGenerationsRequest | None = None @@ -34,12 +36,12 @@ class GenerateSession: self.generate_chunk_cnt = 0 self.current_chunk: RealtimeChunkContext | None = None self.realtime_session = RealtimeSession() - self.adapter: RealtimeModelAdapter | None = None + self.adapter: BaseRealtimeModelAdapter | None = None self.adapter_state: Any = None self.output_pace_next_send_at: float | None = None self.output_pace_last_event_id: int | None = None - def set_adapter(self, adapter: RealtimeModelAdapter): + def set_adapter(self, adapter: BaseRealtimeModelAdapter): self.adapter = adapter self.adapter_state = adapter.create_state() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_adapter.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_adapter.py index 6a95a9116..6b6c71ddb 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_adapter.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_adapter.py @@ -2,8 +2,11 @@ from __future__ import annotations +import hashlib +import os +import tempfile from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any from fastapi import WebSocket @@ -12,8 +15,17 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( RealtimeVideoGenerationsRequest, ) from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_output_adapter import ( + RawRGBRealtimeOutputAdapter, RealtimeFrameSendStats, ) +from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( + build_sampling_params, + save_image_to_path, +) +from sglang.multimodal_gen.runtime.entrypoints.utils import ( + prepare_request, +) +from sglang.multimodal_gen.runtime.server_args import get_global_server_args if TYPE_CHECKING: from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.generate_session import ( @@ -29,35 +41,210 @@ if TYPE_CHECKING: @dataclass(slots=True) class RealtimeChunkInputs: + """Sampled from realtime control state, consumed by the Req""" + prompt: str condition_inputs: dict[str, Any] = field(default_factory=dict) -class RealtimeModelAdapter(Protocol): - name: str +async def save_realtime_first_frame( + session: GenerateSession, + request: RealtimeVideoGenerationsRequest, + *, + required_error: str | None = None, + cache_remote_urls: bool = False, +) -> None: + first_frame = request.first_frame + if first_frame is None: + if required_error is not None: + raise ValueError(required_error) + return - def create_state(self) -> Any: ... + server_args = get_global_server_args() + if server_args.input_save_path is not None: + uploads_dir = server_args.input_save_path + os.makedirs(uploads_dir, exist_ok=True) + else: + if session.input_temp_dir is None: + session.input_temp_dir = tempfile.mkdtemp(prefix="sglang_input_") + uploads_dir = session.input_temp_dir + + if ( + cache_remote_urls + and isinstance(first_frame, str) + and first_frame.lower().startswith(("http://", "https://")) + ): + suffix = os.path.splitext(first_frame.split("?", 1)[0])[1] + digest = hashlib.sha256(first_frame.encode("utf-8")).hexdigest()[:16] + target_path = os.path.join(uploads_dir, f"realtime_ref_{digest}{suffix}") + if os.path.exists(target_path): + request.first_frame = target_path + return + else: + target_path = os.path.join(uploads_dir, f"{session.id}_first_frame") + + request.first_frame = await save_image_to_path(first_frame, target_path) + + +def build_realtime_sampling_params( + request_id: str, + *, + request: RealtimeVideoGenerationsRequest, + chunk_inputs: RealtimeChunkInputs, + num_frames: int | None, + num_inference_steps: int | None, + chunk_size: int, +): + return build_sampling_params( + request_id, + prompt=chunk_inputs.prompt, + size=request.size, + num_frames=num_frames, + fps=request.fps, + image_path=request.first_frame, + output_file_name=request_id, + save_output=False, + seed=request.seed, + generator_device=request.generator_device, + num_inference_steps=num_inference_steps, + guidance_scale=request.guidance_scale, + guidance_scale_2=request.guidance_scale_2, + negative_prompt=request.negative_prompt, + enable_teacache=request.enable_teacache, + enable_frame_interpolation=request.enable_frame_interpolation, + frame_interpolation_exp=request.frame_interpolation_exp, + frame_interpolation_scale=request.frame_interpolation_scale, + frame_interpolation_model_path=request.frame_interpolation_model_path, + enable_upscaling=request.enable_upscaling, + upscaling_model_path=request.upscaling_model_path, + upscaling_scale=request.upscaling_scale, + diffusers_kwargs=request.diffusers_kwargs, + profile=request.profile, + num_profiled_timesteps=request.num_profiled_timesteps, + profile_all_stages=request.profile_all_stages, + perf_dump_path=request.perf_dump_path, + output_path=request.output_path, + output_compression=request.output_compression, + output_quality=request.output_quality, + condition_inputs=chunk_inputs.condition_inputs, + realtime_chunk_size=chunk_size, + ) + + +class BaseRealtimeModelAdapter: + def __init__(self): + self.output_adapter = RawRGBRealtimeOutputAdapter() async def on_init( self, session: GenerateSession, request: RealtimeVideoGenerationsRequest, - ) -> None: ... + ) -> None: + raise NotImplementedError - async def wait_for_next_chunk(self, session: GenerateSession) -> None: ... + def create_state(self) -> Any: + """create a state for managing runtime states""" + raise NotImplementedError def ingest_event( self, session: GenerateSession, event: RealtimeEvent, - ) -> str: ... + ) -> str: + """ + Ingest a realtime endpoint event and install it into the model's realtime control queues + """ + raise NotImplementedError + + async def wait_for_next_chunk(self, session: GenerateSession) -> None: + del session + + def get_chunk_size( + self, + session: GenerateSession, + server_args: ServerArgs, + chunk: RealtimeChunkContext, + ) -> int: + del session, chunk + arch_config = server_args.pipeline_config.dit_config.arch_config + return int(getattr(arch_config, "num_frames_per_block", 3)) + + def sample_chunk_inputs( + self, + session: GenerateSession, + server_args: ServerArgs, + chunk: RealtimeChunkContext, + chunk_size: int, + ) -> RealtimeChunkInputs: + raise NotImplementedError + + def build_sampling_params( + self, + session: GenerateSession, + server_args: ServerArgs, + chunk: RealtimeChunkContext, + chunk_inputs: RealtimeChunkInputs, + chunk_size: int, + ): + raise NotImplementedError + + def get_realtime_event_id(self, session: GenerateSession) -> int | None: + del session + return None def prepare_next_request( self, session: GenerateSession, server_args: ServerArgs, chunk: RealtimeChunkContext, - ) -> Req: ... + ) -> Req: + chunk_size = self.get_chunk_size(session, server_args, chunk) + chunk_inputs = self.sample_chunk_inputs( + session, + server_args, + chunk, + chunk_size, + ) + sampling_params = self.build_sampling_params( + session, + server_args, + chunk, + chunk_inputs, + chunk_size, + ) + batch = prepare_request( + server_args=server_args, + sampling_params=sampling_params, + ) + self.apply_realtime_request_fields( + batch, + session, + chunk, + event_id=self.get_realtime_event_id(session), + ) + return batch + + def apply_realtime_request_fields( + self, + batch: Req, + session: GenerateSession, + chunk: RealtimeChunkContext, + *, + event_id: int | None, + ) -> None: + batch.realtime_session_id = session.id + batch.return_raw_frames = True + batch.block_idx = chunk.index + batch.realtime_event_id = event_id + if session.request is None: + return + batch.realtime_output_format = session.request.realtime_output_format + batch.realtime_preview_max_width = session.request.realtime_preview_max_width + batch.realtime_output_pacing = bool(session.request.realtime_output_pacing) + batch.realtime_causal_sink_size = session.request.realtime_causal_sink_size + batch.realtime_causal_kv_cache_num_frames = ( + session.request.realtime_causal_kv_cache_num_frames + ) async def send_output( self, @@ -65,10 +252,17 @@ class RealtimeModelAdapter(Protocol): session: GenerateSession, result: OutputBatch, batch: Req, - ) -> RealtimeFrameSendStats: ... + ) -> RealtimeFrameSendStats: + """send the generate output (usually frames) back via websocket""" + return await self.output_adapter.send(ws, session, result, batch) - def on_chunk_complete( - self, session: GenerateSession, result: OutputBatch - ) -> None: ... + def on_chunk_complete(self, session: GenerateSession, result: OutputBatch) -> None: + del result + session.generate_chunk_completed() - def dispose(self, session: GenerateSession) -> None: ... + def clear_state(self, session: GenerateSession) -> None: + del session + + def dispose(self, session: GenerateSession) -> None: + self.clear_state(session) + self.output_adapter.reset() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_output_adapter.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_output_adapter.py index 349a3b879..ba567c8a2 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_output_adapter.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_output_adapter.py @@ -11,7 +11,7 @@ import msgspec.msgpack from fastapi import WebSocket from PIL import Image -from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.timing import ( +from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.timer import ( RealtimeStageTimer, ) from sglang.multimodal_gen.runtime.utils.realtime_video import ( diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_video_api.py index 0aade381a..305a40818 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_video_api.py @@ -22,7 +22,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_output_a from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.registry import ( get_realtime_model_adapter, ) -from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.timing import ( +from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.timer import ( RealtimeStageTimer, ) from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( @@ -181,7 +181,7 @@ async def _generate_loop(ws: WebSocket, session: GenerateSession): adapter.on_chunk_complete(session, result) if pending_send_task is not None: await pending_send_task - if batch.realtime_output_pacing: + if getattr(batch, "realtime_output_pacing", False): await _send_output_and_log( ws, session, @@ -306,7 +306,7 @@ async def _wait_for_realtime_output_slot( batch: "Req", result, ) -> float: - if not batch.realtime_output_pacing: + if not getattr(batch, "realtime_output_pacing", False): return 0.0 frame_count = _result_num_frames(result) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py index 851358542..4c682d0d7 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py @@ -5,20 +5,20 @@ from __future__ import annotations from typing import TYPE_CHECKING from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.realtime_adapter import ( - RealtimeModelAdapter, + BaseRealtimeModelAdapter, ) if TYPE_CHECKING: from sglang.multimodal_gen.runtime.server_args import ServerArgs -_REALTIME_ADAPTER_REGISTRY: dict[type, type[RealtimeModelAdapter]] = {} +_REALTIME_ADAPTER_REGISTRY: dict[type, type[BaseRealtimeModelAdapter]] = {} _BUILTIN_ADAPTERS_REGISTERED = False def register_realtime_model_adapter( pipeline_config_cls: type, - adapter_cls: type[RealtimeModelAdapter], + adapter_cls: type[BaseRealtimeModelAdapter], ) -> None: _REALTIME_ADAPTER_REGISTRY[pipeline_config_cls] = adapter_cls @@ -54,7 +54,7 @@ def _register_builtin_realtime_model_adapters() -> None: def get_realtime_model_adapter( server_args: ServerArgs, -) -> RealtimeModelAdapter: +) -> BaseRealtimeModelAdapter: _register_builtin_realtime_model_adapters() pipeline_config = server_args.pipeline_config diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/timing.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/timer.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/timing.py rename to python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/timer.py diff --git a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py index d3de28615..25a2a858f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py @@ -82,7 +82,9 @@ from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) -from sglang.multimodal_gen.runtime.realtime.causal_state import RealtimeCausalDiTState +from sglang.multimodal_gen.runtime.realtime.states import ( + get_realtime_causal_dit_state, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.utils import add_prefix @@ -1216,7 +1218,7 @@ class CausalLingBotWorldTransformer3DModel(CausalWanTransformer3DModel): return None session = getattr(forward_batch, "session", None) if session is not None: - state = session.get_or_create_state(RealtimeCausalDiTState) + state = get_realtime_causal_dit_state(session) return state.runtime_cache.setdefault(name, {}) extra = getattr(forward_batch, "extra", None) if extra is None: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py index eaae177dc..24abe136e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py @@ -29,7 +29,10 @@ from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) -from sglang.multimodal_gen.runtime.realtime.causal_state import RealtimeCausalDiTState +from sglang.multimodal_gen.runtime.realtime.states import ( + RealtimeCausalDiTState, + get_realtime_causal_dit_state, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -314,7 +317,7 @@ class CausalDMDDenoisingStage(DenoisingStage): batch: Req, ) -> tuple[RealtimeCausalDiTState, bool]: if batch.session is not None: - state = batch.session.get_or_create_state(RealtimeCausalDiTState) + state = get_realtime_causal_dit_state(batch.session) return state, True return RealtimeCausalDiTState(), False diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_chain.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_chain.py index daed8f247..3a0a06268 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_chain.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_chain.py @@ -9,7 +9,7 @@ Per-tick pipeline: SanaWMRealtimeLatentPrepStage -> batch.latents = this tick's pre-noised chunk(s); batch.extra["sana_wm_chunk_plan"] SanaWMCameraCondStage -> batch.extra camera_conditions/chunk_plucker - SanaWMStreamingDenoisingStage (session path; SanaWMStreamCacheState) + SanaWMStreamingDenoisingStage (session path; RealtimeCausalDiTState) SanaWMChunkedRefinerChainStage -> batch.latents = refined buffer SanaWMCausalDecodeChainStage -> OutputBatch (decodes past its frontier) @@ -30,10 +30,11 @@ from sglang.multimodal_gen.runtime.models.dits.sana_wm_components import ( compute_chunk_plucker, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req -from sglang.multimodal_gen.runtime.realtime.causal_state import ( - RealtimeCausalDecodeState, -) from sglang.multimodal_gen.runtime.realtime.session import BaseRealtimeState +from sglang.multimodal_gen.runtime.realtime.states import ( + RealtimeCausalDecodeState, + get_realtime_causal_dit_state, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.utils import PRECISION_TO_TYPE @@ -53,7 +54,7 @@ from .realtime_stage import ( SanaWMRealtimeStage, _motion_param, ) -from .streaming import SanaWMStreamCacheState, SanaWMStreamingDenoisingStage +from .streaming import SanaWMStreamingDenoisingStage from .streaming_refiner import RefinerChunkRunner @@ -219,7 +220,7 @@ class SanaWMRealtimeLatentPrepStage(SanaWMRealtimeStage): session = self.require_session(batch, context="SANA-WM realtime chain") inputs = session.get_or_create_state(SanaWMSessionInputsState) noise = session.get_or_create_state(SanaWMNoiseState) - cache = session.get_or_create_state(SanaWMStreamCacheState) + cache = get_realtime_causal_dit_state(session) first_latent = batch.image_latent if first_latent is None: raise ValueError("cond-frame latent missing (run the encode stage first)") @@ -336,7 +337,7 @@ class SanaWMCameraCondStage(SanaWMRealtimeStage): ) session = self.require_session(batch, context="SANA-WM realtime chain") inputs = session.get_or_create_state(SanaWMSessionInputsState) - cache = session.get_or_create_state(SanaWMStreamCacheState) + cache = get_realtime_causal_dit_state(session) plan = list(batch.extra.get("sana_wm_chunk_plan") or []) target_latent = cache.chunk_indices[-1] + sum(plan) if cache.chunk_idx == 0 and plan: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py index 085795dad..f94eef1f2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py @@ -43,7 +43,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.s SanaWMSelfForcingSampler, SanaWMSelfForcingSamplerConfig, ) -from sglang.multimodal_gen.runtime.realtime.causal_state import RealtimeCausalDiTState +from sglang.multimodal_gen.runtime.realtime.states import ( + get_realtime_causal_dit_state, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.utils import PRECISION_TO_TYPE @@ -159,31 +161,6 @@ def self_forcing_denoise_chunk( return updated_cache -class SanaWMStreamCacheState(RealtimeCausalDiTState): - """Per-session streaming DiT state, framework-pattern (cf. LingBot). - - The Wan-shaped ``kv_cache`` field stays None — SANA-WM's cache is the - heterogeneous per-block 10-slot list (GDN recurrent matrix states + softmax - concat windows + conv tails), carried in the fields below.""" - - def __init__(self): - super().__init__() - # kv[chunk][block][slot] — grown per chunk, stale chunks evicted. - self.stream_kv_cache: list = [] - self.chunk_indices: list[int] = [0] - # Growing stage-1 latent buffer (cond frame + denoised chunks). - self.latents: torch.Tensor | None = None - # Per-session flow-Euler scheduler (fresh shift=1.0; persists across ticks). - self.scheduler: FlowMatchEulerDiscreteScheduler | None = None - - def dispose(self) -> None: - super().dispose() - self.stream_kv_cache = [] - self.chunk_indices = [0] - self.latents = None - self.scheduler = None - - class SanaWMStreamingDenoisingStage(CausalDMDDenoisingStage): """Autoregressive self-forcing streaming denoise — SANA-WM's causal-DMD variant. @@ -263,7 +240,7 @@ class SanaWMStreamingDenoisingStage(CausalDMDDenoisingStage): kv_cache, chunk_idx, valid, num_cached_blocks, num_blocks ) - # Realtime per-chunk path (sessions): per-session state in SanaWMStreamCacheState. + # Realtime per-chunk path (sessions): per-session state in RealtimeCausalDiTState. @torch.no_grad() def _forward_realtime_chunk(self, batch: Req, server_args: ServerArgs) -> Req: if batch.latents is None or batch.latents.ndim != 5: @@ -276,7 +253,9 @@ class SanaWMStreamingDenoisingStage(CausalDMDDenoisingStage): target_dtype = PRECISION_TO_TYPE.get( getattr(pcfg, "dit_precision", "bf16"), torch.bfloat16 ) - state = batch.session.get_or_create_state(SanaWMStreamCacheState) + if batch.session is None: + raise ValueError("SANA-WM realtime denoising requires a realtime session") + state = get_realtime_causal_dit_state(batch.session) if batch.block_idx == 0 and state.latents is not None: state.dispose() # session restart on chunk 0 (mirrors the base stage) @@ -292,6 +271,12 @@ class SanaWMStreamingDenoisingStage(CausalDMDDenoisingStage): ) if state.scheduler is None: state.scheduler = FlowMatchEulerDiscreteScheduler(shift=1.0) + kv_cache = state.kv_cache + if kv_cache is None: + # SANA-WM stores its heterogeneous per-block 10-slot stream cache in + # the generic causal DiT kv_cache slot. + kv_cache = [] + state.kv_cache = kv_cache # Device-only move, NO dtype cast: Module.to(dtype=...) would cast the # DiT's complex RoPE buffers to real, discarding the imaginary part @@ -318,11 +303,11 @@ class SanaWMStreamingDenoisingStage(CausalDMDDenoisingStage): start_f = state.chunk_indices[-1] if chunk_idx > 0 else 0 end_f = (start_f + n) if chunk_idx > 0 else n state.chunk_indices.append(end_f) - state.stream_kv_cache.append( + kv_cache.append( [[None] * _NUM_STREAM_CACHE_SLOTS for _ in range(num_blocks)] ) chunk_kv, sink_num = self._accumulate_kv_cache( - state.stream_kv_cache, + kv_cache, chunk_idx, state.chunk_indices, sampler_cfg.num_cached_blocks, @@ -339,7 +324,7 @@ class SanaWMStreamingDenoisingStage(CausalDMDDenoisingStage): if sink_start > 0: valid = [0] + list(range(sink_start, chunk_idx)) self._evict_stale_kv_cache( - state.stream_kv_cache, + kv_cache, chunk_idx, valid, sampler_cfg.num_cached_blocks, @@ -404,7 +389,7 @@ class SanaWMStreamingDenoisingStage(CausalDMDDenoisingStage): attn_metadata=None, forward_batch=batch, ): - state.stream_kv_cache[chunk_idx] = self_forcing_denoise_chunk( + kv_cache[chunk_idx] = self_forcing_denoise_chunk( transformer=transformer, scheduler=state.scheduler, sigmas=sc.explicit_sigmas, diff --git a/python/sglang/multimodal_gen/runtime/realtime/__init__.py b/python/sglang/multimodal_gen/runtime/realtime/__init__.py index cfdcb609f..e8fcb49d4 100644 --- a/python/sglang/multimodal_gen/runtime/realtime/__init__.py +++ b/python/sglang/multimodal_gen/runtime/realtime/__init__.py @@ -1,35 +1,43 @@ # SPDX-License-Identifier: Apache-2.0 -"""session-scoped realtime state, control events, and runtime-only helpers""" +"""session-scoped realtime state, control signals, and runtime-only helpers""" -from sglang.multimodal_gen.runtime.realtime.camera_controls import ( - RealtimeCameraControlState, -) -from sglang.multimodal_gen.runtime.realtime.causal_state import RealtimeCausalDiTState -from sglang.multimodal_gen.runtime.realtime.condition_events import ( - ConditionEvent, - ConditionEventQueue, - ConditionSamplingParams, +from sglang.multimodal_gen.runtime.realtime.control_signals import ( + ControlScriptQueue, ControlSignal, - ControlStateSamplingQueue, + ControlSignalQueue, + ControlSignalSamplingParams, + ControlStateQueue, ControlStateTransition, + ParsedControlEventPayload, + parse_control_event_payload, ) from sglang.multimodal_gen.runtime.realtime.session import ( BaseRealtimeState, RealtimeSession, RealtimeSessionCache, ) +from sglang.multimodal_gen.runtime.realtime.states import ( + RealtimeCameraControlState, + RealtimeCausalDecodeState, + RealtimeCausalDiTState, + get_realtime_causal_dit_state, +) __all__ = [ "BaseRealtimeState", - "ConditionEvent", - "ConditionEventQueue", - "ConditionSamplingParams", + "ControlScriptQueue", "ControlSignal", - "ControlStateSamplingQueue", + "ControlSignalQueue", + "ControlSignalSamplingParams", + "ControlStateQueue", "ControlStateTransition", + "ParsedControlEventPayload", "RealtimeCameraControlState", + "RealtimeCausalDecodeState", "RealtimeCausalDiTState", "RealtimeSession", "RealtimeSessionCache", + "get_realtime_causal_dit_state", + "parse_control_event_payload", ] diff --git a/python/sglang/multimodal_gen/runtime/realtime/camera_controls.py b/python/sglang/multimodal_gen/runtime/realtime/camera_controls.py deleted file mode 100644 index 3f70990de..000000000 --- a/python/sglang/multimodal_gen/runtime/realtime/camera_controls.py +++ /dev/null @@ -1,201 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from collections import deque -from collections.abc import Callable -from typing import Any - -from sglang.multimodal_gen.runtime.realtime.condition_events import ( - ControlSignal, - ControlStateSamplingQueue, - ControlStateTransition, -) - -CameraActionNormalizer = Callable[[list[Any]], list[str]] -CameraActionValidator = Callable[[Any], list[list[str]]] - - -def _identity_actions(actions: list[Any]) -> list[str]: - return list(actions) - - -class RealtimeCameraControlState: - """Session-local camera-control buffer shared by realtime model adapters. - - Camera controls arrive in two shapes: - - 1. Script mode: ``list[list[str]]`` where each item is one output-frame's - held actions. The script is consumed once from a FIFO and padded with - neutral ``[]`` frames after it runs out. - 2. State mode: timestamped transitions such as "W is currently held". - ``ControlStateSamplingQueue`` expands that continuous state into the next - chunk and can pulse a short key press for a minimum number of frames. - - The two modes are intentionally exclusive. A new script clears state mode, - and new state transitions clear script mode, so adapters never merge two - camera timelines accidentally. ``sample_camera_actions`` returns ``None`` - only when no control should be sent; otherwise it returns exactly - ``chunk_size`` frames, with ``[]`` meaning neutral/no-op for that frame. - """ - - def __init__( - self, - *, - min_pulse_items: int = 1, - script_maxlen: int = 512, - max_transitions: int = 512, - normalize_state_actions: CameraActionNormalizer = _identity_actions, - ) -> None: - self.camera_state = ControlStateSamplingQueue( - default_item=[], - min_pulse_items=min_pulse_items, - max_transitions=max_transitions, - ) - self.camera_script_queue: deque[ControlSignal] = deque(maxlen=script_maxlen) - self.latest_sampled_event_id: int | None = None - self._normalize_state_actions = normalize_state_actions - - def clear(self) -> None: - """Reset all camera controls owned by this realtime session.""" - self.camera_state.clear() - self.camera_script_queue.clear() - self.latest_sampled_event_id = None - - def receive_camera_script( - self, - camera_actions: list[list[str]], - *, - event_id: int | None = None, - ) -> None: - """Replace active controls with a finite per-frame script.""" - self.camera_script_queue.clear() - self.camera_state.clear() - for actions in camera_actions: - self.camera_script_queue.append( - ControlSignal( - kind="camera_actions", - payload=list(actions), - seq_id=event_id, - ) - ) - - def receive_camera_state_transitions( - self, - transitions: list[ControlStateTransition], - ) -> None: - """Replace the script with continuous state transitions.""" - self.camera_script_queue.clear() - self.camera_state.push_many(transitions) - - def receive_camera_actions( - self, - camera_actions: list[list[str]], - *, - event_id: int | None = None, - ) -> None: - self.receive_camera_script(camera_actions, event_id=event_id) - - def receive_camera_state( - self, - actions: list[str], - *, - event_id: int | None = None, - timestamp_ms: int | None = None, - ) -> None: - self.receive_camera_state_transitions( - [ - self._camera_state_transition( - actions, - event_id=event_id, - timestamp_ms=timestamp_ms, - ) - ] - ) - - def receive_camera_event_payload( - self, - payload: Any, - *, - event_id: int | None, - validate_camera_actions: CameraActionValidator, - ) -> str: - """Parse an external camera event and install it as script or state.""" - if isinstance(payload, dict) and payload.get("mode") == "state": - transitions = self._camera_transitions_from_event_payload( - payload, - event_id=event_id, - ) - self.receive_camera_state_transitions(transitions) - return f"kind=camera_actions, mode=state, transitions={len(transitions)}" - - camera_actions = validate_camera_actions(payload) - self.receive_camera_script(camera_actions, event_id=event_id) - return f"kind=camera_actions, mode=script, frames={len(camera_actions)}" - - def sample_camera_actions(self, chunk_size: int) -> list[list[str]] | None: - """Return the next chunk-sized camera action window. - - Script mode has priority because it represents an explicit finite - timeline. State mode is sampled only when no script is pending. - """ - if self.camera_script_queue: - return self._sample_camera_script(chunk_size) - action_list = self.camera_state.sample_chunk(chunk_size) - if action_list is None: - return None - self.latest_sampled_event_id = self.camera_state.latest_sampled_seq_id() - return [list(actions) for actions in action_list] - - def _sample_camera_script(self, chunk_size: int) -> list[list[str]]: - chunk: list[list[str]] = [] - latest_event_id = self.latest_sampled_event_id - while self.camera_script_queue and len(chunk) < chunk_size: - signal = self.camera_script_queue.popleft() - chunk.append(list(signal.payload)) - latest_event_id = signal.seq_id - while len(chunk) < chunk_size: - chunk.append([]) - self.latest_sampled_event_id = latest_event_id - return chunk - - def _camera_state_transition( - self, - actions: list[Any], - *, - event_id: int | None, - timestamp_ms: int | None, - ) -> ControlStateTransition: - return ControlStateTransition( - payload=self._normalize_state_actions(actions), - seq_id=event_id, - timestamp_ms=timestamp_ms, - ) - - def _camera_transitions_from_event_payload( - self, - payload: dict[str, Any], - *, - event_id: int | None, - ) -> list[ControlStateTransition]: - transitions = payload.get("transitions") - if not isinstance(transitions, list): - raise ValueError("camera_actions state payload requires transitions") - result = [] - for transition in transitions: - if not isinstance(transition, dict): - raise ValueError("camera_actions transition must be a map") - actions = transition.get("actions") - if not isinstance(actions, list): - raise ValueError("camera_actions transition actions must be a list") - timestamp_ms = transition.get("client_ts_ms") - if timestamp_ms is not None: - timestamp_ms = int(timestamp_ms) - result.append( - self._camera_state_transition( - actions, - event_id=event_id, - timestamp_ms=timestamp_ms, - ) - ) - return result diff --git a/python/sglang/multimodal_gen/runtime/realtime/causal_state.py b/python/sglang/multimodal_gen/runtime/realtime/causal_state.py deleted file mode 100644 index b54a13f2f..000000000 --- a/python/sglang/multimodal_gen/runtime/realtime/causal_state.py +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from sglang.multimodal_gen.runtime.realtime.session import BaseRealtimeState - - -class RealtimeCausalDiTState(BaseRealtimeState): - """persist causal DiT cache and frame position across realtime chunks""" - - def __init__(self): - super().__init__() - self.kv_cache = None - self.crossattn_cache = None - self.runtime_cache: dict = {} - self.current_chunk_start_frame: int = 0 - self.chunk_idx: int = 0 - - def dispose(self) -> None: - self.kv_cache = None - self.crossattn_cache = None - self.runtime_cache.clear() - self.current_chunk_start_frame = 0 - self.chunk_idx = 0 - - -class RealtimeCausalDecodeState(BaseRealtimeState): - """persist causal VAE decode cache and output frontier across chunks""" - - def __init__(self): - super().__init__() - self.conv_cache: dict | None = None - self.next_dec_idx: int = 0 - - def dispose(self) -> None: - self.conv_cache = None - self.next_dec_idx = 0 diff --git a/python/sglang/multimodal_gen/runtime/realtime/condition_events.py b/python/sglang/multimodal_gen/runtime/realtime/condition_events.py deleted file mode 100644 index d18c589a6..000000000 --- a/python/sglang/multimodal_gen/runtime/realtime/condition_events.py +++ /dev/null @@ -1,289 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import copy -from collections import deque -from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from typing import Any - -_MISSING = object() - - -@dataclass(frozen=True) -class ControlSignal: - kind: str - payload: Any - timestamp_ms: int | None = None - seq_id: int | None = None - - -@dataclass(frozen=True) -class ControlStateTransition: - payload: Any - timestamp_ms: int | None = None - seq_id: int | None = None - - -@dataclass(frozen=True) -class ConditionEvent: - """transport envelope for one or more same-kind control signals""" - - kind: str - payload: Any - - def iter_signals(self, expand_payload: bool = True): - items = ( - self.payload - if self._should_expand_payload(self.payload, expand_payload) - else (self.payload,) - ) - for item in items: - if isinstance(item, ControlSignal): - if item.kind != self.kind: - raise ValueError( - "control signal kind " - f"{item.kind!r} does not match event kind {self.kind!r}" - ) - yield item - else: - yield ControlSignal(kind=self.kind, payload=item) - - @staticmethod - def _should_expand_payload(payload: Any, expand_payload: bool) -> bool: - return ( - expand_payload - and isinstance(payload, Sequence) - and not isinstance(payload, (str, bytes, bytearray)) - ) - - -@dataclass(frozen=True) -class ConditionSamplingParams: - chunk_size: int - default_item: Any = _MISSING - repeat_last: bool = True - repeat_last_across_empty_chunks: bool = False - expand_payload: bool = True - - -class ConditionEventQueue: - """per-session queue for prompt, camera, audio, and future events - - all events are stored here for per-chunk sampling - """ - - def __init__( - self, - max_events: int | Mapping[str, int] = 512, - ) -> None: - self._max_events = max_events - self._events: dict[str, deque[ConditionEvent]] = {} - self._pending_signals: dict[str, deque[ControlSignal]] = {} - self._last_payloads: dict[str, Any] = {} - self._last_sampled_seq_ids: dict[str, int | None] = {} - self._seen_kinds: set[str] = set() - - def push(self, event: ConditionEvent) -> None: - queue = self._queue_for(event.kind) - queue.append(event) - self._seen_kinds.add(event.kind) - - def replace(self, event: ConditionEvent) -> None: - self.clear_kind(event.kind) - self.push(event) - - def pop_latest(self, kind: str) -> Any | None: - queue = self._events.get(kind) - if not queue: - return None - latest_payload = None - has_signal = False - for signal in queue.pop().iter_signals(): - latest_payload = signal.payload - has_signal = True - self._last_sampled_seq_ids[kind] = signal.seq_id - queue.clear() - self._seen_kinds.add(kind) - if not has_signal: - return None - return latest_payload - - def has_events(self, kind: str) -> bool: - queue = self._events.get(kind) - pending = self._pending_signals.get(kind) - return bool(queue) or bool(pending) - - def sample_chunk( - self, - kind: str, - params: ConditionSamplingParams, - ) -> list[Any] | None: - """samples a list of actions for a chunk - - Args: - params: the sampling strategy - - """ - if params.chunk_size <= 0: - return None - - chunk: list[Any] = [] - pending = self._pending_signals.get(kind) - self._drain_signals(kind, pending, chunk, params.chunk_size) - - queue = self._events.get(kind) - while len(chunk) < params.chunk_size and queue: - event = queue.popleft() - signals = deque(event.iter_signals(params.expand_payload)) - self._drain_signals(kind, signals, chunk, params.chunk_size) - if signals: - self._pending_signals[kind] = signals - - if len(chunk) == 0 and kind not in self._seen_kinds: - if params.default_item is _MISSING: - return None - return [params.default_item for _ in range(params.chunk_size)] - - if len(chunk) == 0: - if params.repeat_last_across_empty_chunks and kind in self._last_payloads: - return [self._last_payloads[kind] for _ in range(params.chunk_size)] - if params.default_item is _MISSING: - return None - return [params.default_item for _ in range(params.chunk_size)] - - if not params.repeat_last: - return chunk - - pad_item = self._last_payloads.get(kind, params.default_item) - if pad_item is _MISSING: - return chunk - while len(chunk) < params.chunk_size: - chunk.append(pad_item) - return chunk - - def clear(self) -> None: - self._events.clear() - self._pending_signals.clear() - self._last_payloads.clear() - self._last_sampled_seq_ids.clear() - self._seen_kinds.clear() - - def clear_kind(self, kind: str) -> None: - self._events.pop(kind, None) - self._pending_signals.pop(kind, None) - self._last_payloads.pop(kind, None) - self._last_sampled_seq_ids.pop(kind, None) - self._seen_kinds.discard(kind) - - def last_sampled_seq_id(self, kind: str) -> int | None: - return self._last_sampled_seq_ids.get(kind) - - def _queue_for(self, kind: str) -> deque[ConditionEvent]: - queue = self._events.get(kind) - if queue is None: - if isinstance(self._max_events, Mapping): - maxlen = self._max_events.get(kind, 512) - else: - maxlen = self._max_events - queue = deque(maxlen=maxlen) - self._events[kind] = queue - return queue - - def _drain_signals( - self, - kind: str, - signals: deque[ControlSignal] | None, - chunk: list[Any], - chunk_size: int, - ) -> None: - while signals and len(chunk) < chunk_size: - signal = signals.popleft() - chunk.append(signal.payload) - self._last_payloads[kind] = signal.payload - self._last_sampled_seq_ids[kind] = signal.seq_id - - -class ControlStateSamplingQueue: - """state-based control sampler for realtime inputs""" - - def __init__( - self, - *, - default_item: Any, - min_pulse_items: int = 1, - max_transitions: int = 512, - ) -> None: - self.default_item = default_item - self.min_pulse_items = min_pulse_items - self._pending: deque[ControlStateTransition] = deque(maxlen=max_transitions) - self._current_item = default_item - self._current_seq_id: int | None = None - self._latest_sampled_seq_id: int | None = None - - def clear(self) -> None: - self._pending.clear() - self._current_item = self.default_item - self._current_seq_id = None - self._latest_sampled_seq_id = None - - def push(self, transition: ControlStateTransition) -> None: - self._pending.append(transition) - - def push_many(self, transitions: Sequence[ControlStateTransition]) -> None: - for transition in transitions: - self.push(transition) - - def sample_chunk(self, chunk_size: int) -> list[Any] | None: - if chunk_size <= 0: - return None - - transitions = self._drain_pending() - if not transitions: - self._latest_sampled_seq_id = self._current_seq_id - return [self._copy_item(self._current_item) for _ in range(chunk_size)] - - pulse = self._latest_non_default_transition(transitions) - final = transitions[-1] - self._current_item = final.payload - self._current_seq_id = final.seq_id - - if pulse is not None and pulse.payload != final.payload: - pulse_items = min(self.min_pulse_items, chunk_size) - chunk = [self._copy_item(pulse.payload) for _ in range(pulse_items)] - chunk.extend( - self._copy_item(final.payload) for _ in range(chunk_size - pulse_items) - ) - self._latest_sampled_seq_id = ( - final.seq_id if len(chunk) > pulse_items else pulse.seq_id - ) - return chunk - - self._latest_sampled_seq_id = final.seq_id - return [self._copy_item(final.payload) for _ in range(chunk_size)] - - def latest_sampled_seq_id(self) -> int | None: - return self._latest_sampled_seq_id - - def _drain_pending(self) -> list[ControlStateTransition]: - transitions = list(self._pending) - self._pending.clear() - return transitions - - def _latest_non_default_transition( - self, - transitions: Sequence[ControlStateTransition], - ) -> ControlStateTransition | None: - for transition in reversed(transitions): - if transition.payload != self.default_item: - return transition - return None - - @staticmethod - def _copy_item(item: Any) -> Any: - if isinstance(item, list): - return copy.deepcopy(item) - if isinstance(item, dict): - return copy.deepcopy(item) - return item diff --git a/python/sglang/multimodal_gen/runtime/realtime/control_signals.py b/python/sglang/multimodal_gen/runtime/realtime/control_signals.py new file mode 100644 index 000000000..da3f1e681 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/realtime/control_signals.py @@ -0,0 +1,464 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Realtime control signal primitives. + +This module owns the small, model-agnostic data structures that turn external +realtime control inputs into chunk-sized payloads consumed by model adapters. +It intentionally does not know about cameras, LingBot, or SANA-WM semantics; +callers provide validation/normalization functions and decide how sampled +payloads map to request ``condition_inputs``. + +There are two control modes: + +* Script mode is a finite per-frame timeline, for example + ``[["w"], ["w"], [], ...]``. It is already expanded by the caller, so + ``ControlScriptQueue`` consumes it in order and pads the tail with a neutral + default item when configured. Script mode is useful for tests, presets, and + deterministic replay. +* State mode is a level-triggered stream of transitions, for example "these + keys are currently held". ``ControlStateQueue`` keeps the latest state and + samples a stable chunk from it; a short non-default pulse can be preserved + even when press/release transitions arrive between two render chunks. State + mode is the natural shape for live keyboard/gamepad controls. + +``ControlSignalQueue`` is the lower-level FIFO for discrete signals. It is used +directly for one-shot controls such as prompt updates, and wrapped by +``ControlScriptQueue`` for finite timelines. ``ControlStateQueue`` is separate +because held controls need stateful sampling rather than FIFO draining. +""" + +from __future__ import annotations + +import copy +from collections import deque +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +_MISSING = object() + + +@dataclass(frozen=True) +class ControlSignal: + """a discrete, minimal input signal for a specific control kind. + + Generated from script-mode inputs such as receive_camera_action_script + Consumed by chunk samplers such as sample_camera_actions + + """ + + # camera_actions, prompt + kind: str + # two formats: + # 1. script mode: a sequence of flatten actions (e.g., [["w"], ["w"], ["a"], []]) + # 2. state mode: a sequence of state changes (e.g., "actions": ["w"], "client_ts_ms": 1000) + payload: Any + # the timestep of this signal + timestamp_ms: int | None = None + seq_id: int | None = None + + +@dataclass(frozen=True) +class ControlStateTransition: + payload: Any + timestamp_ms: int | None = None + seq_id: int | None = None + + +@dataclass(frozen=True) +class ControlSignalSamplingParams: + chunk_size: int + default_item: Any = _MISSING + repeat_last: bool = True + repeat_last_across_empty_chunks: bool = False + + +@dataclass(frozen=True) +class ParsedControlEventPayload: + mode: str + payload: Any + + +ControlStatePayloadNormalizer = Callable[[list[Any]], Any] +ControlScriptPayloadValidator = Callable[[Any], Any] + + +def parse_control_event_payload( + payload: Any, + *, + event_id: int | None, + kind: str, + normalize_state_payload: ControlStatePayloadNormalizer, + validate_script_payload: ControlScriptPayloadValidator, +) -> ParsedControlEventPayload: + """parse external control event from endpoint""" + if isinstance(payload, dict) and payload.get("mode") == "state": + return ParsedControlEventPayload( + mode="state", + payload=_control_state_transitions_from_event_payload( + payload, + event_id=event_id, + kind=kind, + normalize_state_payload=normalize_state_payload, + ), + ) + return ParsedControlEventPayload( + mode="script", + payload=validate_script_payload(payload), + ) + + +def _control_state_transitions_from_event_payload( + payload: dict[str, Any], + *, + event_id: int | None, + kind: str, + normalize_state_payload: ControlStatePayloadNormalizer, +) -> list[ControlStateTransition]: + transitions = payload.get("transitions") + if not isinstance(transitions, list): + raise ValueError(f"{kind} state payload requires transitions") + result = [] + for transition in transitions: + if not isinstance(transition, dict): + raise ValueError(f"{kind} transition must be a map") + actions = transition.get("actions") + if not isinstance(actions, list): + raise ValueError(f"{kind} transition actions must be a list") + timestamp_ms = transition.get("client_ts_ms") + if timestamp_ms is not None: + timestamp_ms = int(timestamp_ms) + result.append( + ControlStateTransition( + payload=normalize_state_payload(actions), + seq_id=event_id, + timestamp_ms=timestamp_ms, + ) + ) + return result + + +class ControlSignalQueue: + """FIFO storage for discrete realtime control signals + + Script-mode controls and one-shot signals are already expressed as discrete + payloads, so sampling only consumes queued signals and applies the requested + padding strategy. + """ + + def __init__( + self, + max_events: int | Mapping[str, int] = 512, + ) -> None: + self._max_events = max_events + # [control_kind, deque of control signals] + self._signals: dict[str, deque[ControlSignal]] = {} + self._last_payloads: dict[str, Any] = {} + self._last_sampled_seq_ids: dict[str, int | None] = {} + self._seen_kinds: set[str] = set() + + def push( + self, + kind: str, + payload: Any, + *, + event_id: int | None = None, + timestamp_ms: int | None = None, + expand_payload: bool = True, + ) -> None: + queue = self._queue_for(kind) + for signal in self._iter_signals( + kind, + payload, + event_id=event_id, + timestamp_ms=timestamp_ms, + expand_payload=expand_payload, + ): + queue.append(signal) + self._seen_kinds.add(kind) + + def replace( + self, + kind: str, + payload: Any, + *, + event_id: int | None = None, + timestamp_ms: int | None = None, + expand_payload: bool = True, + ) -> None: + self.clear_kind(kind) + self.push( + kind, + payload, + event_id=event_id, + timestamp_ms=timestamp_ms, + expand_payload=expand_payload, + ) + + def pop_latest(self, kind: str) -> Any | None: + queue = self._signals.get(kind) + if not queue: + return None + signal = queue.pop() + latest_payload = signal.payload + self._last_payloads[kind] = signal.payload + self._last_sampled_seq_ids[kind] = signal.seq_id + queue.clear() + self._seen_kinds.add(kind) + return latest_payload + + def has_events(self, kind: str) -> bool: + queue = self._signals.get(kind) + return bool(queue) + + def sample_chunk( + self, + kind: str, + params: ControlSignalSamplingParams, + ) -> list[Any] | None: + """sample queued signals for one realtime chunk""" + if params.chunk_size <= 0: + return None + + chunk: list[Any] = [] + queue = self._signals.get(kind) + self._drain_signals(kind, queue, chunk, params.chunk_size) + + if len(chunk) == 0 and kind not in self._seen_kinds: + if params.default_item is _MISSING: + return None + return [params.default_item for _ in range(params.chunk_size)] + + if len(chunk) == 0: + if params.repeat_last_across_empty_chunks and kind in self._last_payloads: + return [self._last_payloads[kind] for _ in range(params.chunk_size)] + if params.default_item is _MISSING: + return None + return [params.default_item for _ in range(params.chunk_size)] + + if not params.repeat_last: + return chunk + + pad_item = self._last_payloads.get(kind, params.default_item) + if pad_item is _MISSING: + return chunk + while len(chunk) < params.chunk_size: + chunk.append(pad_item) + return chunk + + def clear(self) -> None: + self._signals.clear() + self._last_payloads.clear() + self._last_sampled_seq_ids.clear() + self._seen_kinds.clear() + + def clear_kind(self, kind: str) -> None: + self._signals.pop(kind, None) + self._last_payloads.pop(kind, None) + self._last_sampled_seq_ids.pop(kind, None) + self._seen_kinds.discard(kind) + + def last_sampled_seq_id(self, kind: str) -> int | None: + return self._last_sampled_seq_ids.get(kind) + + def _queue_for(self, kind: str) -> deque[ControlSignal]: + queue = self._signals.get(kind) + if queue is None: + if isinstance(self._max_events, Mapping): + maxlen = self._max_events.get(kind, 512) + else: + maxlen = self._max_events + queue = deque(maxlen=maxlen) + self._signals[kind] = queue + return queue + + def _iter_signals( + self, + kind: str, + payload: Any, + *, + event_id: int | None, + timestamp_ms: int | None, + expand_payload: bool, + ): + items = ( + payload + if self._should_expand_payload(payload, expand_payload) + else (payload,) + ) + for item in items: + if isinstance(item, ControlSignal): + if item.kind != kind: + raise ValueError( + "control signal kind " + f"{item.kind!r} does not match queue kind {kind!r}" + ) + yield item + else: + yield ControlSignal( + kind=kind, + payload=item, + timestamp_ms=timestamp_ms, + seq_id=event_id, + ) + + @staticmethod + def _should_expand_payload(payload: Any, expand_payload: bool) -> bool: + return ( + expand_payload + and isinstance(payload, Sequence) + and not isinstance(payload, (str, bytes, bytearray)) + ) + + def _drain_signals( + self, + kind: str, + signals: deque[ControlSignal] | None, + chunk: list[Any], + chunk_size: int, + ) -> None: + while signals and len(chunk) < chunk_size: + signal = signals.popleft() + chunk.append(signal.payload) + self._last_payloads[kind] = signal.payload + self._last_sampled_seq_ids[kind] = signal.seq_id + + +class ControlScriptQueue: + """Script-mode queue for finite per-frame control timelines.""" + + def __init__( + self, + kind: str, + *, + max_events: int = 512, + default_item: Any = _MISSING, + ) -> None: + self.kind = kind + self.default_item = default_item + self._signals = ControlSignalQueue(max_events={kind: max_events}) + + def clear(self) -> None: + self._signals.clear_kind(self.kind) + + def push_script( + self, + script: Sequence[Any], + *, + event_id: int | None = None, + ) -> None: + self.clear() + self._signals.push(self.kind, script, event_id=event_id) + + def has_script(self) -> bool: + return self._signals.has_events(self.kind) + + def sample_script(self, chunk_size: int) -> list[Any]: + chunk = self._signals.sample_chunk( + self.kind, + ControlSignalSamplingParams( + chunk_size=chunk_size, + default_item=self.default_item, + repeat_last=False, + ), + ) + if chunk is None: + chunk = [] + while len(chunk) < chunk_size and self.default_item is not _MISSING: + chunk.append(self._copy_item(self.default_item)) + return chunk + + def last_sampled_seq_id(self) -> int | None: + return self._signals.last_sampled_seq_id(self.kind) + + @staticmethod + def _copy_item(item: Any) -> Any: + if isinstance(item, list): + return copy.deepcopy(item) + if isinstance(item, dict): + return copy.deepcopy(item) + return item + + +class ControlStateQueue: + """State-mode sampler for level-triggered realtime controls.""" + + def __init__( + self, + *, + default_item: Any, + min_pulse_items: int = 1, + max_transitions: int = 512, + ) -> None: + self.default_item = default_item + self.min_pulse_items = min_pulse_items + self._pending: deque[ControlStateTransition] = deque(maxlen=max_transitions) + self._current_item = default_item + self._current_seq_id: int | None = None + self._latest_sampled_seq_id: int | None = None + + def clear(self) -> None: + self._pending.clear() + self._current_item = self.default_item + self._current_seq_id = None + self._latest_sampled_seq_id = None + + def push(self, transition: ControlStateTransition) -> None: + self._pending.append(transition) + + def push_many(self, transitions: Sequence[ControlStateTransition]) -> None: + for transition in transitions: + self.push(transition) + + def sample_chunk(self, chunk_size: int) -> list[Any] | None: + if chunk_size <= 0: + return None + + transitions = self._drain_pending() + if not transitions: + self._latest_sampled_seq_id = self._current_seq_id + return [self._copy_item(self._current_item) for _ in range(chunk_size)] + + pulse = self._latest_non_default_transition(transitions) + final = transitions[-1] + self._current_item = final.payload + self._current_seq_id = final.seq_id + + if pulse is not None and pulse.payload != final.payload: + pulse_items = min(self.min_pulse_items, chunk_size) + chunk = [self._copy_item(pulse.payload) for _ in range(pulse_items)] + chunk.extend( + self._copy_item(final.payload) for _ in range(chunk_size - pulse_items) + ) + self._latest_sampled_seq_id = ( + final.seq_id if len(chunk) > pulse_items else pulse.seq_id + ) + return chunk + + self._latest_sampled_seq_id = final.seq_id + return [self._copy_item(final.payload) for _ in range(chunk_size)] + + def latest_sampled_seq_id(self) -> int | None: + return self._latest_sampled_seq_id + + def _drain_pending(self) -> list[ControlStateTransition]: + transitions = list(self._pending) + self._pending.clear() + return transitions + + def _latest_non_default_transition( + self, + transitions: Sequence[ControlStateTransition], + ) -> ControlStateTransition | None: + for transition in reversed(transitions): + if transition.payload != self.default_item: + return transition + return None + + @staticmethod + def _copy_item(item: Any) -> Any: + if isinstance(item, list): + return copy.deepcopy(item) + if isinstance(item, dict): + return copy.deepcopy(item) + return item diff --git a/python/sglang/multimodal_gen/runtime/realtime/session.py b/python/sglang/multimodal_gen/runtime/realtime/session.py index e1f2a9e19..8fa2818d6 100644 --- a/python/sglang/multimodal_gen/runtime/realtime/session.py +++ b/python/sglang/multimodal_gen/runtime/realtime/session.py @@ -33,6 +33,7 @@ class RealtimeSession: def get_or_create_state( self, state_cls: type[BaseRealtimeState] ) -> BaseRealtimeState: + """returns the BaseRealtimeState instance hold by the current RealtimeSession""" state = self._states.get(state_cls) if state is None: state = state_cls() diff --git a/python/sglang/multimodal_gen/runtime/realtime/states/__init__.py b/python/sglang/multimodal_gen/runtime/realtime/states/__init__.py new file mode 100644 index 000000000..89be329cb --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/realtime/states/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""reusable session-scoped state implementations for realtime pipelines""" + +from sglang.multimodal_gen.runtime.realtime.states.camera_control import ( + RealtimeCameraControlState, +) +from sglang.multimodal_gen.runtime.realtime.states.causal import ( + RealtimeCausalDecodeState, + RealtimeCausalDiTState, + get_realtime_causal_dit_state, +) + +__all__ = [ + "RealtimeCameraControlState", + "RealtimeCausalDecodeState", + "RealtimeCausalDiTState", + "get_realtime_causal_dit_state", +] diff --git a/python/sglang/multimodal_gen/runtime/realtime/states/camera_control.py b/python/sglang/multimodal_gen/runtime/realtime/states/camera_control.py new file mode 100644 index 000000000..506d6b999 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/realtime/states/camera_control.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +from sglang.multimodal_gen.runtime.realtime.control_signals import ( + ControlScriptQueue, + ControlStateQueue, + ControlStateTransition, + parse_control_event_payload, +) + +CameraActionNormalizer = Callable[[list[Any]], list[str]] +CameraActionValidator = Callable[[Any], list[list[str]]] + + +def _identity_actions(actions: list[Any]) -> list[str]: + return list(actions) + + +class RealtimeCameraControlState: + """Session-local camera-control buffer shared by realtime model adapters. + + Camera controls arrive in two shapes: + + 1. Script mode: ``list[list[str]]`` where each item is one output-frame's + held actions. The script is consumed once from a FIFO and padded with + neutral ``[]`` frames after it runs out. + 2. State mode: timestamped transitions such as "W is currently held". + ``ControlStateQueue`` expands that continuous state into the next + chunk and can pulse a short key press for a minimum number of frames. + + The two modes are intentionally exclusive. A new script clears state mode, + and new state transitions clear script mode, so adapters never merge two + camera timelines accidentally. ``sample_camera_actions`` returns ``None`` + only when no control should be sent; otherwise it returns exactly + ``chunk_size`` frames, with ``[]`` meaning neutral/no-op for that frame. + """ + + def __init__( + self, + *, + min_pulse_items: int = 1, + script_maxlen: int = 512, + max_transitions: int = 512, + normalize_state_actions: CameraActionNormalizer = _identity_actions, + ) -> None: + # stores state-mode control signals + self.camera_state_queue = ControlStateQueue( + default_item=[], + min_pulse_items=min_pulse_items, + max_transitions=max_transitions, + ) + # stores script-mode control signals + # script-mode signals take precedence over state-mode signals, see sample_camera_actions + self.camera_script_queue = ControlScriptQueue( + "camera_actions", + max_events=script_maxlen, + default_item=[], + ) + self.latest_sampled_event_id: int | None = None + self._normalize_state_actions = normalize_state_actions + + def clear(self) -> None: + """Reset all camera controls owned by this realtime session.""" + self.camera_state_queue.clear() + self.camera_script_queue.clear() + self.latest_sampled_event_id = None + + def receive_camera_action_script( + self, + camera_actions: list[list[str]], + *, + event_id: int | None = None, + ) -> None: + """Replace active controls with a finite per-frame script.""" + self.camera_state_queue.clear() + self.camera_script_queue.push_script( + [list(actions) for actions in camera_actions], + event_id=event_id, + ) + + def receive_camera_state_transitions( + self, + transitions: list[ControlStateTransition], + ) -> None: + """Replace the script with continuous state transitions.""" + self.camera_script_queue.clear() + self.camera_state_queue.push_many(transitions) + + def receive_camera_state( + self, + actions: list[str], + *, + event_id: int | None = None, + timestamp_ms: int | None = None, + ) -> None: + self.receive_camera_state_transitions( + [ + ControlStateTransition( + payload=self._normalize_state_actions(actions), + seq_id=event_id, + timestamp_ms=timestamp_ms, + ) + ] + ) + + def receive_camera_control_event_payload( + self, + payload: Any, + *, + event_id: int | None, + validate_camera_actions: CameraActionValidator, + ) -> str: + """Parse an external camera event (from endpoint) and install it as script or state.""" + parsed = parse_control_event_payload( + payload, + event_id=event_id, + kind="camera_actions", + normalize_state_payload=self._normalize_state_actions, + validate_script_payload=validate_camera_actions, + ) + if parsed.mode == "state": + transitions = parsed.payload + self.receive_camera_state_transitions(transitions) + return f"kind=camera_actions, mode=state, transitions={len(transitions)}" + + camera_actions = parsed.payload + self.receive_camera_action_script(camera_actions, event_id=event_id) + return f"kind=camera_actions, mode=script, frames={len(camera_actions)}" + + def sample_camera_actions(self, chunk_size: int) -> list[list[str]] | None: + """Core method, return the next chunk-sized camera action window. + + Script mode has priority because it represents an explicit finite + timeline. State mode is sampled only when no script is pending. + """ + # Script mode wins: it is an explicit finite timeline and should not be + # merged with held-key state from the live control path. + if self.camera_script_queue.has_script(): + return self._sample_camera_script(chunk_size) + # State mode is the WebUI path: held controls persist across chunks until + # a new transition changes the current state. + action_list = self.camera_state_queue.sample_chunk(chunk_size) + if action_list is None: + return None + self.latest_sampled_event_id = self.camera_state_queue.latest_sampled_seq_id() + return [list(actions) for actions in action_list] + + def _sample_camera_script(self, chunk_size: int) -> list[list[str]]: + chunk = self.camera_script_queue.sample_script(chunk_size) + chunk = [list(actions) for actions in chunk] + self.latest_sampled_event_id = self.camera_script_queue.last_sampled_seq_id() + return chunk diff --git a/python/sglang/multimodal_gen/runtime/realtime/states/causal.py b/python/sglang/multimodal_gen/runtime/realtime/states/causal.py new file mode 100644 index 000000000..3b1c2a083 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/realtime/states/causal.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any, TypeVar, cast + +from sglang.multimodal_gen.runtime.realtime.session import ( + BaseRealtimeState, + RealtimeSession, +) + + +class RealtimeCausalDiTState(BaseRealtimeState): + """persist causal DiT cache, chunk frontier, and output buffer""" + + def __init__(self): + super().__init__() + self.kv_cache: Any = None + self.crossattn_cache: Any = None + self.runtime_cache: dict = {} + self.current_chunk_start_frame: int = 0 + self.chunk_idx: int = 0 + self.chunk_indices: list[int] = [0] + self.latents: Any = None + self.scheduler: Any = None + + def dispose(self) -> None: + self.kv_cache = None + self.crossattn_cache = None + self.runtime_cache.clear() + self.current_chunk_start_frame = 0 + self.chunk_idx = 0 + self.chunk_indices = [0] + self.latents = None + self.scheduler = None + + +RealtimeCausalDiTStateT = TypeVar( + "RealtimeCausalDiTStateT", bound=RealtimeCausalDiTState +) + + +def get_realtime_causal_dit_state( + session: RealtimeSession, + state_cls: type[RealtimeCausalDiTStateT] = RealtimeCausalDiTState, +) -> RealtimeCausalDiTStateT: + return cast(RealtimeCausalDiTStateT, session.get_or_create_state(state_cls)) + + +class RealtimeCausalDecodeState(BaseRealtimeState): + """persist causal VAE decode cache and output frontier across chunks""" + + def __init__(self): + super().__init__() + self.conv_cache: dict | None = None + self.next_dec_idx: int = 0 + + def dispose(self) -> None: + self.conv_cache = None + self.next_dec_idx = 0 diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_condition_events.py b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_condition_events.py deleted file mode 100644 index 2c4091020..000000000 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_condition_events.py +++ /dev/null @@ -1,250 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from sglang.multimodal_gen.runtime.realtime.condition_events import ( - ConditionEvent, - ConditionEventQueue, - ConditionSamplingParams, - ControlSignal, - ControlStateSamplingQueue, - ControlStateTransition, -) - - -def test_condition_event_queue_samples_chunk_and_repeats_last_item(): - queue = ConditionEventQueue() - queue.push(ConditionEvent(kind="camera_actions", payload=[["w"], ["d"]])) - - chunk = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=4, default_item=[]), - ) - - assert chunk == [["w"], ["d"], ["d"], ["d"]] - - -def test_condition_event_contains_multiple_same_kind_control_signals(): - event = ConditionEvent( - kind="camera_actions", - payload=[ - ControlSignal(kind="camera_actions", payload=["w"]), - ControlSignal(kind="camera_actions", payload=["d"]), - ], - ) - signals = list(event.iter_signals()) - - assert [signal.kind for signal in signals] == [ - "camera_actions", - "camera_actions", - ] - assert [signal.payload for signal in signals] == [["w"], ["d"]] - - -def test_condition_event_queue_samples_control_signal_payloads(): - queue = ConditionEventQueue() - queue.push( - ConditionEvent( - kind="camera_actions", - payload=[ - ControlSignal(kind="camera_actions", payload=["w"]), - ControlSignal(kind="camera_actions", payload=["a"]), - ControlSignal(kind="camera_actions", payload=["s"]), - ], - ) - ) - - first = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=2, default_item=[]), - ) - second = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=2, default_item=[]), - ) - - assert first == [["w"], ["a"]] - assert second == [["s"], ["s"]] - - -def test_condition_event_queue_preserves_event_remainder_across_chunks(): - queue = ConditionEventQueue() - queue.push(ConditionEvent(kind="camera_actions", payload=[["w"], ["a"], ["s"]])) - - first = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=2, default_item=[]), - ) - second = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=2, default_item=[]), - ) - - assert first == [["w"], ["a"]] - assert second == [["s"], ["s"]] - - -def test_condition_event_queue_does_not_persist_last_signal_across_empty_chunks(): - queue = ConditionEventQueue() - queue.push(ConditionEvent(kind="camera_actions", payload=[["w"]])) - - first = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=2, default_item=[]), - ) - second = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=2, default_item=[]), - ) - - assert first == [["w"], ["w"]] - assert second == [[], []] - - -def test_condition_event_queue_can_repeat_last_signal_across_empty_chunks(): - queue = ConditionEventQueue() - queue.push(ConditionEvent(kind="camera_actions", payload=[["w"]])) - - first = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams( - chunk_size=2, - default_item=[], - repeat_last_across_empty_chunks=True, - ), - ) - second = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams( - chunk_size=2, - default_item=[], - repeat_last_across_empty_chunks=True, - ), - ) - queue.push(ConditionEvent(kind="camera_actions", payload=[[]])) - third = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams( - chunk_size=2, - default_item=[], - repeat_last_across_empty_chunks=True, - ), - ) - - assert first == [["w"], ["w"]] - assert second == [["w"], ["w"]] - assert third == [[], []] - - -def test_condition_event_queue_tracks_sampled_signal_seq_id(): - queue = ConditionEventQueue() - queue.push( - ConditionEvent( - kind="camera_actions", - payload=[ - ControlSignal(kind="camera_actions", payload=["w"], seq_id=7), - ControlSignal(kind="camera_actions", payload=[], seq_id=8), - ], - ) - ) - - first = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams( - chunk_size=1, - default_item=[], - repeat_last_across_empty_chunks=True, - ), - ) - first_seq_id = queue.last_sampled_seq_id("camera_actions") - second = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams( - chunk_size=1, - default_item=[], - repeat_last_across_empty_chunks=True, - ), - ) - second_seq_id = queue.last_sampled_seq_id("camera_actions") - - assert first == [["w"]] - assert first_seq_id == 7 - assert second == [[]] - assert second_seq_id == 8 - - -def test_condition_event_queue_replace_clears_pending_signals(): - queue = ConditionEventQueue() - queue.push( - ConditionEvent(kind="camera_actions", payload=[["w"], ["w"], ["w"], ["w"]]) - ) - - first = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=2, default_item=[]), - ) - queue.replace(ConditionEvent(kind="camera_actions", payload=[["d"]])) - second = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=3, default_item=[]), - ) - third = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=3, default_item=[]), - ) - - assert first == [["w"], ["w"]] - assert second == [["d"], ["d"], ["d"]] - assert third == [[], [], []] - - -def test_condition_event_queue_returns_none_without_default_item(): - queue = ConditionEventQueue() - - chunk = queue.sample_chunk("audio", ConditionSamplingParams(chunk_size=2)) - - assert chunk is None - - -def test_condition_event_queue_empty_event_switches_to_default_item(): - queue = ConditionEventQueue() - queue.push(ConditionEvent(kind="camera_actions", payload=[])) - - chunk = queue.sample_chunk( - "camera_actions", - ConditionSamplingParams(chunk_size=3, default_item=[]), - ) - - assert chunk == [[], [], []] - - -def test_control_state_sampling_queue_preserves_short_pulse(): - queue = ControlStateSamplingQueue(default_item=[], min_pulse_items=1) - queue.push(ControlStateTransition(payload=["w"], seq_id=7)) - queue.push(ControlStateTransition(payload=[], seq_id=8)) - - chunk = queue.sample_chunk(3) - - assert chunk == [["w"], [], []] - assert queue.latest_sampled_seq_id() == 8 - assert queue.sample_chunk(3) == [[], [], []] - - -def test_control_state_sampling_queue_holds_current_state_without_backlog(): - queue = ControlStateSamplingQueue(default_item=[], min_pulse_items=1) - queue.push(ControlStateTransition(payload=["w"], seq_id=7)) - - assert queue.sample_chunk(3) == [["w"], ["w"], ["w"]] - assert queue.latest_sampled_seq_id() == 7 - assert queue.sample_chunk(3) == [["w"], ["w"], ["w"]] - assert queue.latest_sampled_seq_id() == 7 - - -def test_control_state_sampling_queue_compacts_many_transitions(): - queue = ControlStateSamplingQueue(default_item=[], min_pulse_items=1) - queue.push(ControlStateTransition(payload=["w"], seq_id=7)) - queue.push(ControlStateTransition(payload=["w", "d"], seq_id=8)) - queue.push(ControlStateTransition(payload=["d"], seq_id=9)) - - chunk = queue.sample_chunk(3) - - assert chunk == [["d"], ["d"], ["d"]] - assert queue.latest_sampled_seq_id() == 9 diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_control_signals.py b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_control_signals.py new file mode 100644 index 000000000..e95044911 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_control_signals.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 + +from sglang.multimodal_gen.runtime.realtime.control_signals import ( + ControlScriptQueue, + ControlSignal, + ControlSignalQueue, + ControlSignalSamplingParams, + ControlStateQueue, + ControlStateTransition, + parse_control_event_payload, +) + + +def test_control_signal_queue_samples_chunk_and_repeats_last_item(): + queue = ControlSignalQueue() + queue.push("camera_actions", [["w"], ["d"]]) + + chunk = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=4, default_item=[]), + ) + + assert chunk == [["w"], ["d"], ["d"], ["d"]] + + +def test_control_signal_queue_accepts_multiple_same_kind_control_signals(): + queue = ControlSignalQueue() + queue.push( + "camera_actions", + [ + ControlSignal(kind="camera_actions", payload=["w"]), + ControlSignal(kind="camera_actions", payload=["d"]), + ], + ) + + chunk = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=2, default_item=[]), + ) + + assert chunk == [["w"], ["d"]] + + +def test_control_signal_queue_samples_control_signal_payloads(): + queue = ControlSignalQueue() + queue.push( + "camera_actions", + [ + ControlSignal(kind="camera_actions", payload=["w"]), + ControlSignal(kind="camera_actions", payload=["a"]), + ControlSignal(kind="camera_actions", payload=["s"]), + ], + ) + + first = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=2, default_item=[]), + ) + second = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=2, default_item=[]), + ) + + assert first == [["w"], ["a"]] + assert second == [["s"], ["s"]] + + +def test_control_signal_queue_preserves_event_remainder_across_chunks(): + queue = ControlSignalQueue() + queue.push("camera_actions", [["w"], ["a"], ["s"]]) + + first = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=2, default_item=[]), + ) + second = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=2, default_item=[]), + ) + + assert first == [["w"], ["a"]] + assert second == [["s"], ["s"]] + + +def test_control_signal_queue_does_not_persist_last_signal_across_empty_chunks(): + queue = ControlSignalQueue() + queue.push("camera_actions", [["w"]]) + + first = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=2, default_item=[]), + ) + second = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=2, default_item=[]), + ) + + assert first == [["w"], ["w"]] + assert second == [[], []] + + +def test_control_signal_queue_can_repeat_last_signal_across_empty_chunks(): + queue = ControlSignalQueue() + queue.push("camera_actions", [["w"]]) + + first = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams( + chunk_size=2, + default_item=[], + repeat_last_across_empty_chunks=True, + ), + ) + second = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams( + chunk_size=2, + default_item=[], + repeat_last_across_empty_chunks=True, + ), + ) + queue.push("camera_actions", [[]]) + third = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams( + chunk_size=2, + default_item=[], + repeat_last_across_empty_chunks=True, + ), + ) + + assert first == [["w"], ["w"]] + assert second == [["w"], ["w"]] + assert third == [[], []] + + +def test_control_signal_queue_tracks_sampled_signal_seq_id(): + queue = ControlSignalQueue() + queue.push( + "camera_actions", + [ + ControlSignal(kind="camera_actions", payload=["w"], seq_id=7), + ControlSignal(kind="camera_actions", payload=[], seq_id=8), + ], + ) + + first = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams( + chunk_size=1, + default_item=[], + repeat_last_across_empty_chunks=True, + ), + ) + first_seq_id = queue.last_sampled_seq_id("camera_actions") + second = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams( + chunk_size=1, + default_item=[], + repeat_last_across_empty_chunks=True, + ), + ) + second_seq_id = queue.last_sampled_seq_id("camera_actions") + + assert first == [["w"]] + assert first_seq_id == 7 + assert second == [[]] + assert second_seq_id == 8 + + +def test_control_signal_queue_replace_clears_pending_signals(): + queue = ControlSignalQueue() + queue.push("camera_actions", [["w"], ["w"], ["w"], ["w"]]) + + first = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=2, default_item=[]), + ) + queue.replace("camera_actions", [["d"]]) + second = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=3, default_item=[]), + ) + third = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=3, default_item=[]), + ) + + assert first == [["w"], ["w"]] + assert second == [["d"], ["d"], ["d"]] + assert third == [[], [], []] + + +def test_control_signal_queue_returns_none_without_default_item(): + queue = ControlSignalQueue() + + chunk = queue.sample_chunk("audio", ControlSignalSamplingParams(chunk_size=2)) + + assert chunk is None + + +def test_control_signal_queue_empty_event_switches_to_default_item(): + queue = ControlSignalQueue() + queue.push("camera_actions", []) + + chunk = queue.sample_chunk( + "camera_actions", + ControlSignalSamplingParams(chunk_size=3, default_item=[]), + ) + + assert chunk == [[], [], []] + + +def test_control_script_queue_samples_script_and_pads_default_item(): + queue = ControlScriptQueue("camera_actions", default_item=[]) + queue.push_script([["w"], ["d"]], event_id=7) + + chunk = queue.sample_script(3) + + assert chunk == [["w"], ["d"], []] + assert queue.last_sampled_seq_id() == 7 + assert not queue.has_script() + + +def test_parse_control_event_payload_normalizes_state_transitions(): + parsed = parse_control_event_payload( + { + "mode": "state", + "transitions": [ + {"actions": ["W"], "client_ts_ms": 100}, + {"actions": [], "client_ts_ms": 120}, + ], + }, + event_id=11, + kind="camera_actions", + normalize_state_payload=lambda actions: [ + str(action).lower() for action in actions + ], + validate_script_payload=lambda payload: payload, + ) + + assert parsed.mode == "state" + assert parsed.payload == [ + ControlStateTransition(payload=["w"], timestamp_ms=100, seq_id=11), + ControlStateTransition(payload=[], timestamp_ms=120, seq_id=11), + ] + + +def test_parse_control_event_payload_validates_script_payload(): + parsed = parse_control_event_payload( + [["w"], ["d"]], + event_id=7, + kind="camera_actions", + normalize_state_payload=lambda actions: actions, + validate_script_payload=lambda payload: [list(actions) for actions in payload], + ) + + assert parsed.mode == "script" + assert parsed.payload == [["w"], ["d"]] + + +def test_control_state_queue_preserves_short_pulse(): + queue = ControlStateQueue(default_item=[], min_pulse_items=1) + queue.push(ControlStateTransition(payload=["w"], seq_id=7)) + queue.push(ControlStateTransition(payload=[], seq_id=8)) + + chunk = queue.sample_chunk(3) + + assert chunk == [["w"], [], []] + assert queue.latest_sampled_seq_id() == 8 + assert queue.sample_chunk(3) == [[], [], []] + + +def test_control_state_queue_holds_current_state_without_backlog(): + queue = ControlStateQueue(default_item=[], min_pulse_items=1) + queue.push(ControlStateTransition(payload=["w"], seq_id=7)) + + assert queue.sample_chunk(3) == [["w"], ["w"], ["w"]] + assert queue.latest_sampled_seq_id() == 7 + assert queue.sample_chunk(3) == [["w"], ["w"], ["w"]] + assert queue.latest_sampled_seq_id() == 7 + + +def test_control_state_queue_compacts_many_transitions(): + queue = ControlStateQueue(default_item=[], min_pulse_items=1) + queue.push(ControlStateTransition(payload=["w"], seq_id=7)) + queue.push(ControlStateTransition(payload=["w", "d"], seq_id=8)) + queue.push(ControlStateTransition(payload=["d"], seq_id=9)) + + chunk = queue.sample_chunk(3) + + assert chunk == [["d"], ["d"], ["d"]] + assert queue.latest_sampled_seq_id() == 9 diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py index 6cb8cc8e2..29216d062 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py @@ -17,6 +17,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( RealtimeVideoGenerationsRequest, ) from sglang.multimodal_gen.runtime.entrypoints.openai.realtime import ( + realtime_adapter, realtime_video_api, ) from sglang.multimodal_gen.runtime.entrypoints.openai.realtime.adapters import ( @@ -45,16 +46,16 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.input_validati RealtimeInputValidationStage, RealtimeInputValidationState, ) -from sglang.multimodal_gen.runtime.realtime.causal_state import ( - RealtimeCausalDecodeState, -) -from sglang.multimodal_gen.runtime.realtime.condition_events import ( +from sglang.multimodal_gen.runtime.realtime.control_signals import ( ControlStateTransition, ) from sglang.multimodal_gen.runtime.realtime.session import ( BaseRealtimeState, RealtimeSessionCache, ) +from sglang.multimodal_gen.runtime.realtime.states import ( + RealtimeCausalDecodeState, +) from sglang.multimodal_gen.runtime.utils.realtime_video import ( RAW_RGB_CONTENT_TYPE, ) @@ -153,11 +154,11 @@ def test_realtime_session_cache_rejects_missing_nonzero_chunk(): raise AssertionError("expected missing realtime session to fail") -def test_lingbot_realtime_state_uses_generic_condition_queue(): +def test_lingbot_realtime_state_uses_control_script_and_prompt_queues(): state = lingbot_realtime.LingBotWorldRealtimeState() assert state.sample_camera_actions(3) == [[], [], []] - state.receive_camera_actions([["w"], ["a"], ["s"], ["d"]]) + state.receive_camera_action_script([["w"], ["a"], ["s"], ["d"]]) assert state.sample_camera_actions(3) == [["w"], ["a"], ["s"]] assert state.sample_camera_actions(3) == [["d"], [], []] assert state.sample_camera_actions(3) == [[], [], []] @@ -168,6 +169,17 @@ def test_lingbot_realtime_state_uses_generic_condition_queue(): assert not state.has_prompt() +def test_lingbot_realtime_camera_script_replaces_state_queue(): + state = lingbot_realtime.LingBotWorldRealtimeState() + + state.receive_camera_state(["w"], event_id=7) + state.receive_camera_action_script([["d"]], event_id=8) + + assert state.sample_camera_actions(3) == [["d"], [], []] + assert state.latest_sampled_event_id == 8 + assert state.sample_camera_actions(3) == [[], [], []] + + def test_lingbot_realtime_camera_events_preserve_short_presses(): state = lingbot_realtime.LingBotWorldRealtimeState() @@ -210,7 +222,7 @@ def test_lingbot_realtime_camera_state_compacts_multiple_pending_updates(): def test_sana_wm_realtime_camera_state_uses_sana_normalizer(): state = sana_wm_realtime.SanaWMRealtimeAdapterState() - result = state.receive_camera_event_payload( + result = state.receive_camera_control_event_payload( { "mode": "state", "transitions": [ @@ -230,10 +242,10 @@ def test_sana_wm_realtime_adapter_preserves_requested_size(): async def fake_save_image_to_path(image, target_path): return target_path - old_save_image_to_path = sana_wm_realtime.save_image_to_path - old_get_global_server_args = sana_wm_realtime.get_global_server_args - sana_wm_realtime.save_image_to_path = fake_save_image_to_path - sana_wm_realtime.get_global_server_args = lambda: SimpleNamespace( + old_save_image_to_path = realtime_adapter.save_image_to_path + old_get_global_server_args = realtime_adapter.get_global_server_args + realtime_adapter.save_image_to_path = fake_save_image_to_path + realtime_adapter.get_global_server_args = lambda: SimpleNamespace( input_save_path=None ) try: @@ -251,8 +263,8 @@ def test_sana_wm_realtime_adapter_preserves_requested_size(): assert request.size == "832x480" finally: - sana_wm_realtime.save_image_to_path = old_save_image_to_path - sana_wm_realtime.get_global_server_args = old_get_global_server_args + realtime_adapter.save_image_to_path = old_save_image_to_path + realtime_adapter.get_global_server_args = old_get_global_server_args def test_lingbot_realtime_adapter_ingests_generic_events(): @@ -596,15 +608,16 @@ def test_lingbot_realtime_adapter_prepares_chunk_request(monkeypatch): prompt=sampling_params.prompt, condition_inputs=dict(sampling_params.condition_inputs), realtime_chunk_size=sampling_params.realtime_chunk_size, + session=None, ) monkeypatch.setattr( - lingbot_realtime, + realtime_adapter, "build_sampling_params", fake_build_sampling_params, ) monkeypatch.setattr( - lingbot_realtime, + realtime_adapter, "prepare_request", fake_prepare_backend_request, ) @@ -623,7 +636,7 @@ def test_lingbot_realtime_adapter_prepares_chunk_request(monkeypatch): assert batch.request_id == chunk.request_id assert batch.condition_inputs == {"camera_actions": [["w"], ["w"], ["w"]]} assert batch.realtime_chunk_size == 3 - assert batch.session is session.realtime_session + assert batch.session is None assert batch.realtime_session_id == session.id assert batch.block_idx == 0 assert batch.return_raw_frames is True @@ -749,7 +762,7 @@ def test_lingbot_realtime_adapter_sends_stale_output_for_client_cutover(): session = GenerateSession() session.set_adapter(adapter) state = adapter._state(session) - state.receive_camera_actions([["d"]], event_id=7) + state.receive_camera_action_script([["d"]], event_id=7) calls = [] async def fake_send(ws, session_arg, result_arg, batch_arg): diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py index cbed7d8bc..66a3ec82a 100644 --- a/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py @@ -25,10 +25,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.s from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_stage import ( SanaWMRealtimeStage, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import ( - SanaWMStreamCacheState, -) from sglang.multimodal_gen.runtime.realtime.session import RealtimeSession +from sglang.multimodal_gen.runtime.realtime.states import ( + get_realtime_causal_dit_state, +) from sglang.multimodal_gen.runtime.server_args import set_global_server_args MC = 8 @@ -162,7 +162,7 @@ def test_latent_prep_plan_and_noise_discipline(_global_args): inputs.latent_t = 5 # fixed horizon: segments [0, 3, 5] for nfpb=2 inputs.num_frame_per_block = 2 inputs.sink_size = 1 - cache = session.get_or_create_state(SanaWMStreamCacheState) + cache = get_realtime_causal_dit_state(session) # Tick 0: front-loaded chunk 0 (cond + remainder) + full-horizon buffer. batch = stage.forward(_batch(session, 0, fl), _server_args()) @@ -202,7 +202,7 @@ def test_latent_prep_open_ended_uniform_chunk0(_global_args): inputs.latent_t = None # open-ended inputs.num_frame_per_block = 3 inputs.sink_size = 1 - session.get_or_create_state(SanaWMStreamCacheState) + get_realtime_causal_dit_state(session) batch = stage.forward(_batch(session, 0, fl), _server_args()) noise = session.get_or_create_state(SanaWMNoiseState) diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_realtime_path.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_realtime_path.py index 4f42793b5..2e85089e1 100644 --- a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_realtime_path.py +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_realtime_path.py @@ -23,10 +23,12 @@ from sglang.multimodal_gen.configs.models.dits.sana_wm import ( from sglang.multimodal_gen.runtime import server_args as _sa_mod from sglang.multimodal_gen.runtime.models.dits.sana_wm import SanaWMTransformer3DModel from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import ( - SanaWMStreamCacheState, SanaWMStreamingDenoisingStage, ) from sglang.multimodal_gen.runtime.realtime.session import RealtimeSession +from sglang.multimodal_gen.runtime.realtime.states import ( + get_realtime_causal_dit_state, +) from sglang.multimodal_gen.runtime.server_args import set_global_server_args MC = 8 @@ -124,10 +126,10 @@ def test_realtime_path_multi_tick_carries_state(_global_args): # Tick 0 (chunk 0): conditioning frame + 2 new frames. chunk0 = torch.cat([fl, _noise(1, MC, 2, 2, 2, seed=1)], dim=2) out = stage.forward(_tick(session, 0, chunk0, [3]), server_args) - state = session.get_or_create_state(SanaWMStreamCacheState) + state = get_realtime_causal_dit_state(session) assert out.latents.shape[2] == 3 assert state.chunk_idx == 1 and state.chunk_indices == [0, 3] - assert state.stream_kv_cache[0][0][0] is not None # GDN state stored + assert state.kv_cache[0][0][0] is not None # GDN state stored # The condition frame is held fixed. assert torch.allclose(state.latents[:, :, 0].cpu(), fl[:, :, 0]) @@ -159,13 +161,13 @@ def test_realtime_path_evicts_stale_kv(_global_args): stage.forward( _tick(session, i, _noise(1, MC, 2, 2, 2, seed=10 + i), [2]), server_args ) - state = session.get_or_create_state(SanaWMStreamCacheState) + state = get_realtime_causal_dit_state(session) assert state.chunk_indices == [0, 3, 5, 7, 9, 11] def _has_any(entry): return any(slot is not None for block in entry for slot in block) - kept = [i for i, e in enumerate(state.stream_kv_cache) if _has_any(e)] + kept = [i for i, e in enumerate(state.kv_cache) if _has_any(e)] # Sink chunk + the last num_cached_blocks chunks (accumulate's read window). assert kept == [0, 3, 4] @@ -183,7 +185,7 @@ def test_realtime_path_is_deterministic(_global_args): stage.forward( _tick(session, 1, _noise(1, MC, 2, 2, 2, seed=6), [2]), server_args ) - return session.get_or_create_state(SanaWMStreamCacheState).latents.cpu() + return get_realtime_causal_dit_state(session).latents.cpu() torch.manual_seed(1234) a = _run() @@ -199,7 +201,7 @@ def test_realtime_path_resets_on_block_zero(_global_args): chunk0 = torch.cat([fl, _noise(1, MC, 2, 2, 2, seed=1)], 2) stage.forward(_tick(session, 0, chunk0, [3]), server_args) stage.forward(_tick(session, 1, _noise(1, MC, 2, 2, 2, seed=2), [2]), server_args) - state = session.get_or_create_state(SanaWMStreamCacheState) + state = get_realtime_causal_dit_state(session) assert state.chunk_idx == 2 # block_idx == 0 restarts the session in place.