minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] refactor: refactor realtime control state and adapters (#27698)

Browse Source
This commit is contained in:
Mick
2026-06-10 12:27:16 +08:00
committed by GitHub
parent 047e5df3b1
commit 56f06278c6
26 changed files with 1381 additions and 1174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/lingbot_world_realtime_adapter.py
+30 -155
View File
@@ -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()
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/adapters/sana_wm_realtime_adapter.py
+39 -134
View File
@@ -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()
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/generate_session.py
+5 -3
View File
@@ -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()
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_adapter.py
+207 -13
View File
@@ -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()
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_output_adapter.py
+1 -1
View File
@@ -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 (
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/realtime_video_api.py
+3 -3
View File
@@ -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)
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/registry.py
+4 -4
View File
@@ -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
python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/timing.py → python/sglang/multimodal_gen/runtime/entrypoints/openai/realtime/timer.py
View File
python/sglang/multimodal_gen/runtime/models/dits/lingbot_world.py
+4 -2
View File
@@ -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:
python/sglang/multimodal_gen/runtime/pipelines_core/stages/causal_denoising.py
+5 -2
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_chain.py
+8 -7
View File
@@ -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:
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/streaming.py
+17 -32
View File
@@ -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,
python/sglang/multimodal_gen/runtime/realtime/__init__.py
+22 -14
View File
@@ -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",
]
python/sglang/multimodal_gen/runtime/realtime/camera_controls.py
-201
View File
@@ -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
python/sglang/multimodal_gen/runtime/realtime/causal_state.py
-35
View File
@@ -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
python/sglang/multimodal_gen/runtime/realtime/condition_events.py
-289
View File
@@ -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
python/sglang/multimodal_gen/runtime/realtime/control_signals.py
+464
View File
@@ -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
python/sglang/multimodal_gen/runtime/realtime/session.py
+1
View File
@@ -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()
python/sglang/multimodal_gen/runtime/realtime/states/__init__.py
+19
View File
@@ -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",
]
python/sglang/multimodal_gen/runtime/realtime/states/camera_control.py
+156
View File
@@ -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
python/sglang/multimodal_gen/runtime/realtime/states/causal.py
+58
View File
@@ -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
python/sglang/multimodal_gen/test/unit/realtime/test_realtime_condition_events.py
-250
View File
@@ -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
python/sglang/multimodal_gen/test/unit/realtime/test_realtime_control_signals.py
+294
View File
@@ -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
python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py
+30 -17
View File
@@ -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):
python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py
+5 -5
View File
@@ -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)
python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_realtime_path.py
+9 -7
View File
@@ -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.