diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 571d7f4f1..db3f75ee0 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -7,9 +7,11 @@ from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig from fastvideo.configs.models.dits.stepvideo import StepVideoConfig from fastvideo.configs.models.dits.wanvideo import WanVideoConfig from fastvideo.configs.models.dits.hyworld import HYWorldConfig +from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig __all__ = [ "HunyuanVideoConfig", "HunyuanVideo15Config", "WanVideoConfig", "StepVideoConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", - "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig" + "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", + "WanGameVideoConfig" ] diff --git a/fastvideo/configs/models/dits/wangamevideo.py b/fastvideo/configs/models/dits/wangamevideo.py new file mode 100644 index 000000000..5d98a6c4d --- /dev/null +++ b/fastvideo/configs/models/dits/wangamevideo.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def is_blocks(n: str, m) -> bool: + return "blocks" in n and str.isdigit(n.split(".")[-1]) + + +@dataclass +class WanGameVideoArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks]) + + param_names_mapping: dict = field( + default_factory=lambda: { + r"^patch_embedding\.(.*)$": + r"patch_embedding.proj.\1", + r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": + r"condition_embedder.text_embedder.fc_in.\1", + r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": + r"condition_embedder.text_embedder.fc_out.\1", + r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": + r"condition_embedder.time_embedder.mlp.fc_in.\1", + r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": + r"condition_embedder.time_embedder.mlp.fc_out.\1", + r"^condition_embedder\.time_proj\.(.*)$": + r"condition_embedder.time_modulation.linear.\1", + r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$": + r"condition_embedder.image_embedder.ff.fc_in.\1", + r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$": + r"condition_embedder.image_embedder.ff.fc_out.\1", + r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": + r"blocks.\1.to_q.\2", + r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": + r"blocks.\1.to_k.\2", + r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": + r"blocks.\1.to_v.\2", + r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": + r"blocks.\1.to_out.\2", + r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": + r"blocks.\1.norm_q.\2", + r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": + r"blocks.\1.norm_k.\2", + r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": + r"blocks.\1.attn2.to_out.\2", + r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": + r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": + r"blocks.\1.ffn.fc_out.\2", + r"^blocks\.(\d+)\.norm2\.(.*)$": + r"blocks.\1.self_attn_residual_norm.norm.\2", + }) + + # Reverse mapping for saving checkpoints: custom -> hf + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + + # Some LoRA adapters use the original official layer names instead of hf layer names, + # so apply this before the param_names_mapping + lora_param_names_mapping: dict = field( + default_factory=lambda: { + r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2", + r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2", + r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2", + r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": + r"blocks.\1.attn1.to_out.0.\2", + r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", + r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", + r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", + r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": + r"blocks.\1.attn2.to_out.0.\2", + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + }) + + patch_size: tuple[int, int, int] = (1, 2, 2) + text_len = 512 + num_attention_heads: int = 40 + attention_head_dim: int = 128 + in_channels: int = 16 + out_channels: int = 16 + text_dim: int = 4096 + freq_dim: int = 256 + ffn_dim: int = 13824 + num_layers: int = 40 + cross_attn_norm: bool = True + qk_norm: str = "rms_norm_across_heads" + eps: float = 1e-6 + image_dim: int | None = None + added_kv_proj_dim: int | None = None + rope_max_seq_len: int = 1024 + pos_embed_seq_len: int | None = None + exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"]) + + # Wan MoE + boundary_ratio: float | None = None + + # Causal Wan + local_attn_size: int = 6 # Window size for temporal local attention (-1 indicates global attention) + sink_size: int = 0 # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache + num_frames_per_block: int = 3 + sliding_window_num_frames: int = 21 + + def __post_init__(self): + super().__post_init__() + self.out_channels = self.out_channels or self.in_channels + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = self.out_channels + + +@dataclass +class WanGameVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=WanGameVideoArchConfig) + + prefix: str = "WanGame" + + +@dataclass +class WanLingBotVideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=WanGameVideoArchConfig) + + prefix: str = "WanLingBot" diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index 3228f400a..494c63803 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -23,7 +23,7 @@ from fastvideo.configs.pipelines.wan import ( Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig, SelfForcingWanT2V480PConfig, WANV2VConfig, SelfForcingWan2_2_T2V480PConfig, - MatrixGameI2V480PConfig) + MatrixGameI2V480PConfig, WanGameI2V480PConfig) # isort: on from fastvideo.logger import init_logger from fastvideo.utils import (maybe_download_model_index, @@ -42,6 +42,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "FastVideo/HY-WorldPlay-Bidirectional-Diffusers": HYWorldConfig, "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig, "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig, + "weizhou03/Wan2.1-Game-Fun-1.3B-InP-Diffusers": WanGameI2V480PConfig, "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers": WANV2VConfig, "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig, "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig, @@ -93,6 +94,8 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = { lambda id: "hyworld" in id.lower(), "matrixgame": lambda id: "matrix-game" in id.lower() or "matrixgame" in id.lower(), + "wangameactionimagetovideo": + lambda id: "wangameactionimagetovideo" in id.lower(), "wanpipeline": lambda id: "wanpipeline" in id.lower(), "wanimagetovideo": @@ -124,6 +127,7 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = { "hunyuan": HunyuanConfig, # Base Hunyuan config as fallback for any Hunyuan variant "matrixgame": MatrixGameI2V480PConfig, + "wangameactionimagetovideo": WanGameI2V480PConfig, "hunyuan15": Hunyuan15T2V480PConfig, # Base Hunyuan15 config as fallback for any Hunyuan15 variant "hyworld": diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index 434839dcc..6241372c0 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -7,6 +7,7 @@ import torch from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig from fastvideo.configs.models.dits import WanVideoConfig from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig +from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPVisionConfig, T5Config, WAN2_1ControlCLIPVisionConfig) @@ -217,3 +218,23 @@ class MatrixGameI2V480PConfig(WanI2V480PConfig): context_noise: int = 0 num_frames_per_block: int = 3 # sliding_window_num_frames: int = 15 + + +# ============================================= +# ============= WanGame ====================== +# ============================================= +@dataclass +class WanGameI2V480PConfig(WanI2V480PConfig): + """Configuration for WanGame image-to-video pipeline.""" + dit_config: DiTConfig = field(default_factory=WanGameVideoConfig) + + image_encoder_config: EncoderConfig = field( + default_factory=WAN2_1ControlCLIPVisionConfig) + + is_causal: bool = True + flow_shift: float | None = 5.0 + dmd_denoising_steps: list[int] | None = field( + default_factory=lambda: [1000, 666, 333]) + warp_denoising_step: bool = True + context_noise: int = 0 + num_frames_per_block: int = 3 diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index 22014b06b..3983deaf6 100644 --- a/fastvideo/configs/sample/registry.py +++ b/fastvideo/configs/sample/registry.py @@ -92,6 +92,9 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = { "KyleShao/Cosmos-Predict2.5-2B-Diffusers": Cosmos_Predict2_5_2B_Diffusers_SamplingParam, + # WanGame models + "weizhou03/Wan2.1-Game-Fun-1.3B-InP-Diffusers": MatrixGame2_SamplingParam, + # MatrixGame2.0 models "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGame2_SamplingParam, "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGame2_SamplingParam, @@ -130,6 +133,8 @@ SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = { lambda id: "wandmdpipeline" in id.lower(), "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(), + "wangame": + lambda id: "wangame" in id.lower(), "matrixgame": lambda id: "matrixgame" in id.lower() or "matrix-game" in id.lower(), "turbodiffusion": @@ -157,6 +162,7 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = { "wandmdpipeline": FastWanT2V480P_SamplingParam, "wancausaldmdpipeline": SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam, "stepvideo": StepVideoT2VSamplingParam, + "wangame": MatrixGame2_SamplingParam, "matrixgame": MatrixGame2_SamplingParam, "turbodiffusion": TurboDiffusionT2V_1_3B_SamplingParam, # Default to T2V for fallback diff --git a/fastvideo/entrypoints/streaming_generator.py b/fastvideo/entrypoints/streaming_generator.py index fd52dc9dc..7441a603c 100644 --- a/fastvideo/entrypoints/streaming_generator.py +++ b/fastvideo/entrypoints/streaming_generator.py @@ -171,7 +171,7 @@ class StreamingVideoGenerator(VideoGenerator): def step( self, keyboard_cond: torch.Tensor, - mouse_cond: torch.Tensor) -> tuple[list[np.ndarray], Future | None]: + mouse_cond: torch.Tensor) -> tuple[list[np.ndarray], Future | None, dict | None]: if self.batch is None: raise RuntimeError("Call reset() before step()") @@ -195,7 +195,10 @@ class StreamingVideoGenerator(VideoGenerator): # Returns Future for block file, or None if no block_dir block_future = self.writer.add_frames(frames) - return frames, block_future + # Extract stage timings if available + stage_timings = getattr(output_batch, 'stage_timings', None) + + return frames, block_future, stage_timings async def step_async( self, keyboard_cond: torch.Tensor, diff --git a/fastvideo/models/dits/matrixgame/action_module.py b/fastvideo/models/dits/matrixgame/action_module.py index 6c06babae..5412c597b 100644 --- a/fastvideo/models/dits/matrixgame/action_module.py +++ b/fastvideo/models/dits/matrixgame/action_module.py @@ -496,14 +496,14 @@ class ActionModule(nn.Module): else: local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - kv_cache_keyboard["global_end_index"].item() local_start_index = local_end_index - num_new_tokens - assert k.shape[0] == S # BS == 1 or the cache should not be saved/ load method should be modified - kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k[:1] - kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v[:1] + assert k.shape[0] % S == 0 # BS >= 1 + kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k[::S] + kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v[::S] attn = self.keyboard_attn_layer( q, - kv_cache_keyboard["k"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat(S, 1, 1, 1), - kv_cache_keyboard["v"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat(S, 1, 1, 1), + kv_cache_keyboard["k"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat_interleave(S, dim=0), + kv_cache_keyboard["v"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat_interleave(S, dim=0), ) kv_cache_keyboard["global_end_index"].fill_(current_end) diff --git a/fastvideo/models/dits/matrixgame/causal_model.py b/fastvideo/models/dits/matrixgame/causal_model.py index 13806d2ac..d18334a9a 100644 --- a/fastvideo/models/dits/matrixgame/causal_model.py +++ b/fastvideo/models/dits/matrixgame/causal_model.py @@ -14,7 +14,7 @@ from torch.nn.attention.flex_attention import BlockMask from fastvideo.forward_context import get_forward_context flex_attention = torch.compile( - flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") + flex_attention, dynamic=True, mode="max-autotune-no-cudagraphs") from fastvideo.attention import LocalAttention from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig @@ -465,6 +465,7 @@ _DEFAULT_MATRIXGAME_CONFIG = MatrixGameWanVideoConfig() class CausalMatrixGameWanModel(BaseDiT): supports_action_input = True + _concatenates_image_latent = True # Model handles image_latent internally via forward_context _fsdp_shard_conditions = _DEFAULT_MATRIXGAME_CONFIG._fsdp_shard_conditions _compile_conditions = _DEFAULT_MATRIXGAME_CONFIG._compile_conditions diff --git a/fastvideo/models/dits/wangame/__init__.py b/fastvideo/models/dits/wangame/__init__.py new file mode 100644 index 000000000..ebe540626 --- /dev/null +++ b/fastvideo/models/dits/wangame/__init__.py @@ -0,0 +1,12 @@ +from .model import WanGameActionTransformer3DModel +from .causal_model import (CausalWanGameTransformer3DModel, + CausalWanTransformer3DModel) +from .hyworld_action_module import WanGameActionTimeImageEmbedding, WanGameActionSelfAttention + +__all__ = [ + "WanGameActionTransformer3DModel", + "CausalWanTransformer3DModel", + "CausalWanGameTransformer3DModel", + "WanGameActionTimeImageEmbedding", + "WanGameActionSelfAttention", +] diff --git a/fastvideo/models/dits/wangame/causal_model.py b/fastvideo/models/dits/wangame/causal_model.py new file mode 100644 index 000000000..1034ec6c2 --- /dev/null +++ b/fastvideo/models/dits/wangame/causal_model.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +from typing import Any + +import torch + +from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig +from fastvideo.models.dits.matrixgame.causal_model import CausalMatrixGameWanModel + +_DEFAULT_WANGAME_CAUSAL_CONFIG = WanGameVideoConfig() + + +class CausalWanTransformer3DModel(CausalMatrixGameWanModel): + supports_action_input = False + + _fsdp_shard_conditions = _DEFAULT_WANGAME_CAUSAL_CONFIG._fsdp_shard_conditions + _compile_conditions = _DEFAULT_WANGAME_CAUSAL_CONFIG._compile_conditions + _supported_attention_backends = _DEFAULT_WANGAME_CAUSAL_CONFIG._supported_attention_backends + param_names_mapping = _DEFAULT_WANGAME_CAUSAL_CONFIG.param_names_mapping + reverse_param_names_mapping = _DEFAULT_WANGAME_CAUSAL_CONFIG.reverse_param_names_mapping + lora_param_names_mapping = _DEFAULT_WANGAME_CAUSAL_CONFIG.lora_param_names_mapping + + def __init__(self, + config: WanGameVideoConfig, + hf_config: dict[str, Any], + **kwargs) -> None: + super().__init__(config=config, hf_config=hf_config, **kwargs) + + def _normalize_action_inputs( + self, + mouse_cond: torch.Tensor | None, + keyboard_cond: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + # drop + if len(getattr(self, "action_config", {})) == 0: + return None, None + return mouse_cond, keyboard_cond + + def _forward_inference( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + mouse_cond: torch.Tensor | None = None, + keyboard_cond: torch.Tensor | None = None, + kv_cache: dict | None = None, + kv_cache_mouse: dict | None = None, + kv_cache_keyboard: dict | None = None, + crossattn_cache: dict | None = None, + current_start: int = 0, + cache_start: int = 0, + start_frame: int = 0, + **kwargs, + ) -> torch.Tensor: + mouse_cond, keyboard_cond = self._normalize_action_inputs( + mouse_cond, keyboard_cond) + return super()._forward_inference( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + encoder_hidden_states_image=encoder_hidden_states_image, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond, + kv_cache=kv_cache, + kv_cache_mouse=kv_cache_mouse, + kv_cache_keyboard=kv_cache_keyboard, + crossattn_cache=crossattn_cache, + current_start=current_start, + cache_start=cache_start, + start_frame=start_frame, + **kwargs, + ) + + def _forward_train( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + mouse_cond: torch.Tensor | None = None, + keyboard_cond: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + mouse_cond, keyboard_cond = self._normalize_action_inputs( + mouse_cond, keyboard_cond) + return super()._forward_train( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + encoder_hidden_states_image=encoder_hidden_states_image, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond, + **kwargs, + ) + + +class CausalWanGameTransformer3DModel(CausalWanTransformer3DModel): + pass diff --git a/fastvideo/models/dits/wangame/hyworld_action_module.py b/fastvideo/models/dits/wangame/hyworld_action_module.py new file mode 100644 index 000000000..5fedff3f6 --- /dev/null +++ b/fastvideo/models/dits/wangame/hyworld_action_module.py @@ -0,0 +1,289 @@ +import math + +import torch +import torch.nn as nn + +from fastvideo.layers.visual_embedding import TimestepEmbedder, ModulateProjection, timestep_embedding +from fastvideo.platforms import AttentionBackendEnum +from fastvideo.attention import DistributedAttention +from fastvideo.forward_context import set_forward_context +from fastvideo.models.dits.wanvideo import WanImageEmbedding + +from fastvideo.models.dits.hyworld.camera_rope import prope_qkv +from fastvideo.layers.rotary_embedding import _apply_rotary_emb +from fastvideo.layers.mlp import MLP + +class WanGameActionTimeImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + image_embed_dim: int | None = None, + ): + super().__init__() + + self.time_freq_dim = time_freq_dim + self.time_embedder = TimestepEmbedder( + dim, frequency_embedding_size=time_freq_dim, act_layer="silu") + self.time_modulation = ModulateProjection(dim, + factor=6, + act_layer="silu") + + self.image_embedder = None + if image_embed_dim is not None: + self.image_embedder = WanImageEmbedding(image_embed_dim, dim) + + self.action_embedder = MLP( + time_freq_dim, + dim, + dim, + bias=True, + act_type="silu" + ) + # Initialize fc_in with kaiming_uniform (same as nn.Linear default) + nn.init.kaiming_uniform_(self.action_embedder.fc_in.weight, a=math.sqrt(5)) + # Initialize fc_out with zeros for residual-like behavior + nn.init.zeros_(self.action_embedder.fc_out.weight) + if self.action_embedder.fc_out.bias is not None: + nn.init.zeros_(self.action_embedder.fc_out.bias) + + def forward( + self, + timestep: torch.Tensor, + action: torch.Tensor, + encoder_hidden_states: torch.Tensor, # Kept for interface compatibility + encoder_hidden_states_image: torch.Tensor | None = None, + timestep_seq_len: int | None = None, + ): + """ + Args: + timestep: [B] diffusion timesteps (one per batch sample) + action: [B, T] action labels (one per frame per batch sample) + + Returns: + temb: [B*T, dim] combined timestep + action embedding + timestep_proj: [B*T, 6*dim] modulation projection + """ + # timestep: [B] -> temb: [B, dim] + temb = self.time_embedder(timestep, timestep_seq_len) + + # Handle action embedding for batch > 1 + # action shape: [B, T] where B=batch_size, T=num_frames + batch_size = action.shape[0] + num_frames = action.shape[1] + + # Compute action embeddings: [B, T] -> [B*T] -> [B*T, dim] + action_flat = action.flatten() # [B*T] + action_emb = timestep_embedding(action_flat, self.time_freq_dim) + action_embedder_dtype = next(iter(self.action_embedder.parameters())).dtype + if ( + action_emb.dtype != action_embedder_dtype + and action_embedder_dtype != torch.int8 + ): + action_emb = action_emb.to(action_embedder_dtype) + action_emb = self.action_embedder(action_emb).type_as(temb) # [B*T, dim] + + # Expand temb to match action_emb: [B, dim] -> [B, T, dim] -> [B*T, dim] + # Each batch's temb is repeated for all its frames + temb_expanded = temb.unsqueeze(1).expand(-1, num_frames, -1) # [B, T, dim] + temb_expanded = temb_expanded.reshape(batch_size * num_frames, -1) # [B*T, dim] + + # Add action embedding to expanded temb + temb = temb_expanded + action_emb # [B*T, dim] + + timestep_proj = self.time_modulation(temb) # [B*T, 6*dim] + + # MatrixGame does not use text embeddings, so we ignore encoder_hidden_states + + if encoder_hidden_states_image is not None: + assert self.image_embedder is not None + encoder_hidden_states_image = self.image_embedder( + encoder_hidden_states_image) + + encoder_hidden_states = torch.zeros((batch_size, 0, temb.shape[-1]), + device=temb.device, + dtype=temb.dtype) + + return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image + +class WanGameActionSelfAttention(nn.Module): + """ + Self-attention module with support for: + - Standard RoPE-based attention + - Camera PRoPE-based attention (when viewmats and Ks are provided) + - KV caching for autoregressive generation + """ + + def __init__(self, + dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm=True, + eps=1e-6) -> None: + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.qk_norm = qk_norm + self.eps = eps + self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560 + + # Scaled dot product attention (using DistributedAttention for SP support) + self.attn = DistributedAttention( + num_heads=num_heads, + head_size=self.head_dim, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA)) + + def forward(self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + kv_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + viewmats: torch.Tensor | None = None, + Ks: torch.Tensor | None = None, + is_cache: bool = False, + attention_mask: torch.Tensor | None = None): + """ + Forward pass with camera PRoPE attention combining standard RoPE and projective positional encoding. + + Args: + q, k, v: Query, key, value tensors [B, L, num_heads, head_dim] + freqs_cis: RoPE frequency cos/sin tensors + kv_cache: KV cache dict (may have None values for training) + current_start: Current position for KV cache + cache_start: Cache start position + viewmats: Camera view matrices for PRoPE [B, cameras, 4, 4] + Ks: Camera intrinsics for PRoPE [B, cameras, 3, 3] + is_cache: Whether to store to KV cache (for inference) + attention_mask: Attention mask [B, L] (1 = attend, 0 = mask) + """ + if cache_start is None: + cache_start = current_start + + # Apply RoPE manually + cos, sin = freqs_cis + query_rope = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v) + key_rope = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) + value_rope = v + + # # DEBUG: Check camera matrices + # if self.training and torch.distributed.get_rank() == 0: + # vm_info = f"viewmats={viewmats.shape if viewmats is not None else None}" + # ks_info = f"Ks={Ks.shape if Ks is not None else None}" + # vm_nonzero = (viewmats != 0).sum().item() if viewmats is not None else 0 + # ks_nonzero = (Ks != 0).sum().item() if Ks is not None else 0 + # print(f"[DEBUG] PRoPE input: {vm_info} nonzero={vm_nonzero}, {ks_info} nonzero={ks_nonzero}", flush=True) + + # Get PRoPE transformed q, k, v + query_prope, key_prope, value_prope, apply_fn_o = prope_qkv( + q.transpose(1, 2), # [B, num_heads, L, head_dim] + k.transpose(1, 2), + v.transpose(1, 2), + viewmats=viewmats, + Ks=Ks, + patches_x=40, # hardcoded for now + patches_y=22, + ) + # PRoPE returns [B, num_heads, L, head_dim], convert to [B, L, num_heads, head_dim] + query_prope = query_prope.transpose(1, 2) + key_prope = key_prope.transpose(1, 2) + value_prope = value_prope.transpose(1, 2) + + # # DEBUG: Check prope_qkv output + # if self.training and torch.distributed.get_rank() == 0: + # q_nz = (query_prope != 0).sum().item() + # k_nz = (key_prope != 0).sum().item() + # v_nz = (value_prope != 0).sum().item() + # print(f"[DEBUG] prope_qkv output: q_nonzero={q_nz}, k_nonzero={k_nz}, v_nonzero={v_nz}", flush=True) + + # KV cache handling + if kv_cache is not None: + cache_key = kv_cache.get("k", None) + cache_value = kv_cache.get("v", None) + + if cache_value is not None and not is_cache: + cache_key_rope, cache_key_prope = cache_key.chunk(2, dim=-1) + cache_value_rope, cache_value_prope = cache_value.chunk(2, dim=-1) + + key_rope = torch.cat([cache_key_rope, key_rope], dim=1) + value_rope = torch.cat([cache_value_rope, value_rope], dim=1) + key_prope = torch.cat([cache_key_prope, key_prope], dim=1) + value_prope = torch.cat([cache_value_prope, value_prope], dim=1) + + if is_cache: + # Store to cache (update input dict directly) + kv_cache["k"] = torch.cat([key_rope, key_prope], dim=-1) + kv_cache["v"] = torch.cat([value_rope, value_prope], dim=-1) + + # Concatenate rope and prope paths (matching original) + query_all = torch.cat([query_rope, query_prope], dim=0) + key_all = torch.cat([key_rope, key_prope], dim=0) + value_all = torch.cat([value_rope, value_prope], dim=0) + + # Check if Q and KV have different sequence lengths (KV cache mode) + # In this case, use LocalAttention (supports different Q/KV lengths) + if query_all.shape[1] != key_all.shape[1]: + # KV cache mode: Q has new tokens only, KV has cached + new tokens + # Use LocalAttention which supports different Q/KV lengths + # LocalAttention will use the appropriate backend (SageAttn, FlashAttn, etc.) + if not hasattr(self, '_kv_cache_attn'): + from fastvideo.attention import LocalAttention + self._kv_cache_attn = LocalAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=False, + supported_attention_backends=(AttentionBackendEnum.SAGE_ATTN, + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA) + ) + hidden_states_all = self._kv_cache_attn(query_all, key_all, value_all) + else: + # Same sequence length: use DistributedAttention (supports SP) + # Create default attention mask if not provided + # NOTE: query_all has shape [2*B, L, ...] (rope+prope concatenated), so mask needs 2*B + if attention_mask is None: + batch_size, seq_len = q.shape[0], q.shape[1] + attention_mask = torch.ones(batch_size * 2, seq_len, device=q.device, dtype=q.dtype) + + from fastvideo.attention.backends.sdpa import SDPAMetadataBuilder + attn_metadata_builder = SDPAMetadataBuilder + if q.dtype != torch.float32: + try: + from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder + attn_metadata_builder = FlashAttnMetadataBuilder + except ImportError: + pass # fall back to SDPA + attn_metadata = attn_metadata_builder().build( + current_timestep=0, + attn_mask=attention_mask, + ) + with set_forward_context(current_timestep=0, attn_metadata=attn_metadata): + hidden_states_all, _ = self.attn(query_all, key_all, value_all, attention_mask=attention_mask) + + hidden_states_rope, hidden_states_prope = hidden_states_all.chunk(2, dim=0) + + # # DEBUG: Check attention output and apply_fn_o + # if self.training and torch.distributed.get_rank() == 0: + # attn_all_nz = (hidden_states_all != 0).sum().item() + # rope_nz = (hidden_states_rope != 0).sum().item() + # prope_before = (hidden_states_prope != 0).sum().item() + # print(f"[DEBUG] attn output: all_nonzero={attn_all_nz}, rope_nonzero={rope_nz}, prope_before_apply={prope_before}", flush=True) + + hidden_states_prope = apply_fn_o(hidden_states_prope.transpose(1, 2)).transpose(1, 2) + + # # DEBUG: Check after apply_fn_o + # if self.training and torch.distributed.get_rank() == 0: + # prope_after = (hidden_states_prope != 0).sum().item() + # print(f"[DEBUG] prope_after_apply_fn_o={prope_after}", flush=True) + + return hidden_states_rope, hidden_states_prope \ No newline at end of file diff --git a/fastvideo/models/dits/wangame/model.py b/fastvideo/models/dits/wangame/model.py new file mode 100644 index 000000000..58839375c --- /dev/null +++ b/fastvideo/models/dits/wangame/model.py @@ -0,0 +1,427 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +import torch +import torch.nn as nn + +from fastvideo.configs.models.dits.wangamevideo import WanGameVideoConfig +from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, + RMSNorm, ScaleResidual, + ScaleResidualLayerNormScaleShift) +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.mlp import MLP +from fastvideo.layers.rotary_embedding import (_apply_rotary_emb, + get_rotary_pos_embed) +from fastvideo.layers.visual_embedding import PatchEmbed +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.wanvideo import WanI2VCrossAttention +from fastvideo.platforms import AttentionBackendEnum, current_platform + +# Import ActionModule +from fastvideo.models.dits.wangame.hyworld_action_module import WanGameActionTimeImageEmbedding, WanGameActionSelfAttention + +logger = init_logger(__name__) + + +class WanGameCrossAttention(WanI2VCrossAttention): + def forward(self, x, context, context_lens=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + """ + context_img = context + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) + k_img = self.norm_added_k(self.add_k_proj(context_img)[0]).view( + b, -1, n, d) + v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d) + img_x = self.attn(q, k_img, v_img) + + # output + x = img_x.flatten(2) + x, _ = self.to_out(x) + return x + +class WanGameActionTransformerBlock(nn.Module): + """ + Transformer block for WAN Action model with support for: + - Self-attention with RoPE and camera PRoPE + - Cross-attention with text/image context + - Feed-forward network with AdaLN modulation + """ + + def __init__(self, + dim: int, + ffn_dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = ""): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True) + self.to_k = ReplicatedLinear(dim, dim, bias=True) + self.to_v = ReplicatedLinear(dim, dim, bias=True) + self.to_out = ReplicatedLinear(dim, dim, bias=True) + + self.attn1 = WanGameActionSelfAttention( + dim, + num_heads, + local_attn_size=local_attn_size, + sink_size=sink_size, + qk_norm=qk_norm, + eps=eps) + + self.hidden_dim = dim + self.num_attention_heads = num_heads + self.local_attn_size = local_attn_size + dim_head = dim // num_heads + + if qk_norm == "rms_norm": + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + elif qk_norm == "rms_norm_across_heads": + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + else: + raise ValueError(f"QK Norm type {qk_norm} not supported") + + assert cross_attn_norm is True + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + compute_dtype=torch.float32) + + # 2. Cross-attention (I2V only for now) + self.attn2 = WanGameCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps) + # norm3 for FFN input + self.norm3 = LayerNormScaleShift(dim, norm_type="layer", eps=eps, + elementwise_affine=False) + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.mlp_residual = ScaleResidual() + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + # PRoPE output projection (initialized via add_discrete_action_parameters on the model) + self.to_out_prope = ReplicatedLinear(dim, dim, bias=True) + nn.init.zeros_(self.to_out_prope.weight) + if self.to_out_prope.bias is not None: + nn.init.zeros_(self.to_out_prope.bias) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + kv_cache: dict | None = None, + crossattn_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + viewmats: torch.Tensor | None = None, + Ks: torch.Tensor | None = None, + is_cache: bool = False, + ) -> torch.Tensor: + if hidden_states.dim() == 4: + hidden_states = hidden_states.squeeze(1) + + num_frames = temb.shape[1] + frame_seqlen = hidden_states.shape[1] // num_frames + bs, seq_length, _ = hidden_states.shape + orig_dtype = hidden_states.dtype + + # Cast temb to float32 for scale/shift computation + e = self.scale_shift_table + temb.float() + assert e.shape == (bs, num_frames, 6, self.hidden_dim) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=2) + + # 1. Self-attention + norm_hidden_states = (self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1 + scale_msa) + shift_msa).to(orig_dtype).flatten(1, 2) + + query, _ = self.to_q(norm_hidden_states) + key, _ = self.to_k(norm_hidden_states) + value, _ = self.to_v(norm_hidden_states) + + if self.norm_q is not None: + query = self.norm_q.forward_native(query) + if self.norm_k is not None: + key = self.norm_k.forward_native(key) + + query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + + # Self-attention with camera PRoPE + attn_output_rope, attn_output_prope = self.attn1( + query, key, value, freqs_cis, + kv_cache, current_start, cache_start, viewmats, Ks, + is_cache=is_cache + ) + # Combine rope and prope outputs + attn_output_rope = attn_output_rope.flatten(2) + attn_output_rope, _ = self.to_out(attn_output_rope) + attn_output_prope = attn_output_prope.flatten(2) + + # # DEBUG: Check if prope input is zero + # if self.training and torch.distributed.get_rank() == 0: + # prope_nonzero = (attn_output_prope != 0).sum().item() + # prope_total = attn_output_prope.numel() + # if prope_nonzero == 0: + # print(f"[DEBUG] to_out_prope INPUT is ALL ZEROS! shape={attn_output_prope.shape}", flush=True) + + attn_output_prope, _ = self.to_out_prope(attn_output_prope) + attn_output = attn_output_rope.squeeze(1) + attn_output_prope.squeeze(1) + + # Self-attention residual + norm in float32 + null_shift = null_scale = torch.zeros(1, device=hidden_states.device, dtype=torch.float32) + norm_hidden_states, hidden_states = self.self_attn_residual_norm( + hidden_states.float(), attn_output.float(), gate_msa, null_shift, null_scale) + hidden_states = hidden_states.type_as(attn_output) + norm_hidden_states = norm_hidden_states.type_as(attn_output) + + # 2. Cross-attention + attn_output = self.attn2(norm_hidden_states.to(orig_dtype), + context=encoder_hidden_states, + context_lens=None) + # Cross-attention residual in bfloat16 + hidden_states = hidden_states + attn_output + + # norm3 for FFN input in float32 + norm_hidden_states = self.norm3( + hidden_states.float(), c_shift_msa, c_scale_msa + ).type_as(hidden_states) + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states.to(orig_dtype)) + hidden_states = self.mlp_residual(hidden_states.float(), ff_output.float(), c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) # Cast back to original dtype + + return hidden_states + +class WanGameActionTransformer3DModel(BaseDiT): + """ + WAN Action Transformer 3D Model for video generation with action conditioning. + + Extends the base WAN video model with: + - Action embedding support for controllable generation + - camera PRoPE attention for 3D-aware generation + - KV caching for autoregressive inference + """ + _uses_discrete_action = True # Uses single action tensor [B, T] instead of keyboard/mouse KV caches + _kv_cache_head_dim_multiplier = 2 # PRoPE stores cat([rope, prope], dim=-1) in KV cache + + _fsdp_shard_conditions = WanGameVideoConfig()._fsdp_shard_conditions + _compile_conditions = WanGameVideoConfig()._compile_conditions + _supported_attention_backends = WanGameVideoConfig()._supported_attention_backends + param_names_mapping = WanGameVideoConfig().param_names_mapping + reverse_param_names_mapping = WanGameVideoConfig().reverse_param_names_mapping + lora_param_names_mapping = WanGameVideoConfig().lora_param_names_mapping + + def __init__(self, config: WanGameVideoConfig, hf_config: dict[str, Any]) -> None: + super().__init__(config=config, hf_config=hf_config) + + inner_dim = config.num_attention_heads * config.attention_head_dim + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.attention_head_dim = config.attention_head_dim + self.in_channels = config.in_channels + self.out_channels = config.out_channels + self.num_channels_latents = config.num_channels_latents + self.patch_size = config.patch_size + self.local_attn_size = config.local_attn_size + self.inner_dim = inner_dim + + # 1. Patch & position embedding + self.patch_embedding = PatchEmbed(in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False) + + # 2. Condition embeddings (with action support) + self.condition_embedder = WanGameActionTimeImageEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + image_embed_dim=config.image_dim, + ) + + # 3. Transformer blocks + self.blocks = nn.ModuleList([ + WanGameActionTransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.local_attn_size, + config.sink_size, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + supported_attention_backends=self._supported_attention_backends, + prefix=f"{config.prefix}.blocks.{i}") + for i in range(config.num_layers) + ]) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift(inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32) + self.proj_out = nn.Linear( + inner_dim, config.out_channels * math.prod(config.patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + self.gradient_checkpointing = False + + # Causal-specific + self.num_frame_per_block = config.arch_config.num_frames_per_block + assert self.num_frame_per_block <= 3 + + self.__post_init__() + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor], + guidance=None, + action: torch.Tensor | None = None, + viewmats: torch.Tensor | None = None, + Ks: torch.Tensor | None = None, + kv_cache: list[dict] | None = None, + crossattn_cache: list[dict] | None = None, + current_start: int = 0, + cache_start: int = 0, + start_frame: int = 0, + is_cache: bool = False, + **kwargs + ) -> torch.Tensor: + """ + Forward pass for both training and inference with KV caching. + + Args: + hidden_states: Video latents [B, C, T, H, W] + encoder_hidden_states: Text embeddings [B, L, D] + timestep: Timestep tensor + encoder_hidden_states_image: Optional image embeddings + action: Action tensor [B, T] for per-frame conditioning + viewmats: Camera view matrices for PRoPE [B, T, 4, 4] + Ks: Camera intrinsics for PRoPE [B, T, 3, 3] + kv_cache: KV cache for autoregressive inference (list of dicts per layer) + crossattn_cache: Cross-attention cache for inference + current_start: Current position for KV cache + cache_start: Cache start position + start_frame: RoPE offset for new frames in autoregressive mode + is_cache: If True, populate KV cache and return early (cache-only mode) + """ + orig_dtype = hidden_states.dtype + # if not isinstance(encoder_hidden_states, torch.Tensor): + # encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + # else: + # encoder_hidden_states_image = None + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000, + start_frame=start_frame + ) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + if timestep.dim() == 2: + # condition_embedder expects [B] timestep, not [B*T]. + # All frames share the same timestep, so take first column. + timestep = timestep[:, 0] + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, action, encoder_hidden_states, encoder_hidden_states_image=encoder_hidden_states_image) + + # condition_embedder returns: + # - temb: [B*T, dim] where T = post_patch_num_frames + # - timestep_proj: [B*T, 6*dim] + # Reshape to [B, T, 6, dim] for transformer blocks + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)) # [B*T, 6, dim] + timestep_proj = timestep_proj.view(batch_size, post_patch_num_frames, 6, self.hidden_size) # [B, T, 6, dim] + + encoder_hidden_states = encoder_hidden_states_image + + # Transformer blocks + for block_idx, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, + kv_cache[block_idx] if kv_cache else None, + crossattn_cache[block_idx] if crossattn_cache else None, + current_start, cache_start, + viewmats, Ks, is_cache) + else: + hidden_states = block( + hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, + kv_cache[block_idx] if kv_cache else None, + crossattn_cache[block_idx] if crossattn_cache else None, + current_start, cache_start, + viewmats, Ks, is_cache) + + # If cache-only mode, return early + if is_cache: + return kv_cache + + # Output norm, projection & unpatchify + # temb is [B*T, dim], reshape to [B, T, 1, dim] + temb = temb.view(batch_size, post_patch_num_frames, -1).unsqueeze(2) # [B, T, 1, dim] + + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, + post_patch_height, + post_patch_width, p_t, p_h, p_w, + -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return output \ No newline at end of file diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 954b96049..01d1546f1 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -44,6 +44,8 @@ _IMAGE_TO_VIDEO_DIT_MODELS = { "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), "MatrixGameWanModel": ("dits", "matrixgame", "MatrixGameWanModel"), "CausalMatrixGameWanModel": ("dits", "matrixgame", "CausalMatrixGameWanModel"), + "WanGameActionTransformer3DModel": ("dits", "wangame", "WanGameActionTransformer3DModel"), + "CausalWanGameTransformer3DModel": ("dits", "wangame", "CausalWanGameTransformer3DModel"), } _TEXT_ENCODER_MODELS = { diff --git a/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py index 837c47988..20dff65d8 100644 --- a/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py +++ b/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py @@ -96,28 +96,47 @@ class MatrixGameCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): self._vae_cache = None def streaming_step(self, keyboard_action, mouse_action) -> ForwardBatch: + import time denoiser = self._stage_name_mapping["denoising_stage"] ctx = denoiser._streaming_ctx assert ctx is not None, "streaming_ctx must be set" start_idx = ctx.start_index + + # Time DiT forward pass + torch.cuda.synchronize() + t0 = time.perf_counter() batch = denoiser.streaming_step(keyboard_action, mouse_action) + torch.cuda.synchronize() + dit_time_ms = (time.perf_counter() - t0) * 1000 + end_idx = ctx.start_index # Decode only the new generated block + vae_time_ms = 0.0 if end_idx > start_idx: current_latents = batch.latents[:, :, start_idx:end_idx, :, :] args = ctx.fastvideo_args decoder = self._stage_name_mapping["decoding_stage"] + + # Time VAE decode + torch.cuda.synchronize() + t0 = time.perf_counter() decoded_frames, self._vae_cache = decoder.streaming_decode( current_latents, args, cache=self._vae_cache, is_first_chunk=(start_idx == 0)) + torch.cuda.synchronize() + vae_time_ms = (time.perf_counter() - t0) * 1000 + batch.output = decoded_frames else: batch.output = None + # Store timings in batch for caller to access + batch.stage_timings = {"dit_ms": dit_time_ms, "vae_ms": vae_time_ms} + return batch def streaming_clear(self) -> None: diff --git a/fastvideo/pipelines/basic/wan/wangame_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wangame_causal_dmd_pipeline.py new file mode 100644 index 000000000..3904a97fa --- /dev/null +++ b/fastvideo/pipelines/basic/wan/wangame_causal_dmd_pipeline.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +"""WanGame causal DMD pipeline implementation.""" + +from fastvideo.fastvideo_args import FastVideoArgs +import torch +from fastvideo.logger import init_logger +from fastvideo.pipelines import ComposedPipelineBase, ForwardBatch, LoRAPipeline + +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + MatrixGameCausalDenoisingStage) +from fastvideo.pipelines.stages.image_encoding import ( + MatrixGameImageEncodingStage, MatrixGameImageVAEEncodingStage) + +logger = init_logger(__name__) + + +class WanGameCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): + _required_config_modules = [ + "vae", "transformer", "scheduler", "image_encoder", "image_processor" + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + if (self.get_module("text_encoder", None) is not None + and self.get_module("tokenizer", None) is not None): + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + if (self.get_module("image_encoder", None) is not None + and self.get_module("image_processor", None) is not None): + self.add_stage( + stage_name="image_encoding_stage", + stage=MatrixGameImageEncodingStage( + image_encoder=self.get_module("image_encoder"), + image_processor=self.get_module("image_processor"), + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer", None))) + + self.add_stage( + stage_name="image_latent_preparation_stage", + stage=MatrixGameImageVAEEncodingStage(vae=self.get_module("vae"))) + + self.add_stage(stage_name="denoising_stage", + stage=MatrixGameCausalDenoisingStage( + transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2", None), + scheduler=self.get_module("scheduler"), + pipeline=self, + vae=self.get_module("vae"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + logger.info("WanGameCausalDMDPipeline initialized") + + @torch.no_grad() + def streaming_reset(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs): + if not self.post_init_called: + self.post_init() + + stages_to_run = [ + "input_validation_stage", "prompt_encoding_stage", + "image_encoding_stage", "conditioning_stage", + "latent_preparation_stage", "image_latent_preparation_stage" + ] + + for stage_name in stages_to_run: + if stage_name in self._stage_name_mapping: + batch = self._stage_name_mapping[stage_name].forward( + batch, fastvideo_args) + + denoiser = self._stage_name_mapping["denoising_stage"] + denoiser.streaming_reset(batch, fastvideo_args) + self._vae_cache = None + + def streaming_step(self, keyboard_action, mouse_action) -> ForwardBatch: + denoiser = self._stage_name_mapping["denoising_stage"] + ctx = denoiser._streaming_ctx + assert ctx is not None, "streaming_ctx must be set" + + start_idx = ctx.start_index + batch = denoiser.streaming_step(keyboard_action, mouse_action) + end_idx = ctx.start_index + + if end_idx > start_idx: + current_latents = batch.latents[:, :, start_idx:end_idx, :, :] + args = ctx.fastvideo_args + decoder = self._stage_name_mapping["decoding_stage"] + decoded_frames, self._vae_cache = decoder.streaming_decode( + current_latents, + args, + cache=self._vae_cache, + is_first_chunk=(start_idx == 0)) + batch.output = decoded_frames + else: + batch.output = None + + return batch + + def streaming_clear(self) -> None: + denoiser = self._stage_name_mapping.get("denoising_stage") + if denoiser is not None and hasattr(denoiser, "streaming_clear"): + denoiser.streaming_clear() + self._vae_cache = None + +EntryClass = [WanGameCausalDMDPipeline] diff --git a/fastvideo/pipelines/basic/wan/wangame_i2v_pipeline.py b/fastvideo/pipelines/basic/wan/wangame_i2v_pipeline.py new file mode 100644 index 000000000..cd7e53d4d --- /dev/null +++ b/fastvideo/pipelines/basic/wan/wangame_i2v_pipeline.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Wan video diffusion pipeline implementation. + +This module contains an implementation of the Wan video diffusion pipeline +using the modular pipeline architecture. +""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.lora_pipeline import LoRAPipeline + +# isort: off +from fastvideo.pipelines.stages import ( + ImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage, + ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage, + TimestepPreparationStage) +# isort: on +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler) + +logger = init_logger(__name__) + + +class WanGameActionImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase): + + _required_config_modules = [ + "vae", "transformer", "scheduler", \ + "image_encoder", "image_processor" + ] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + self.modules["scheduler"] = FlowUniPCMultistepScheduler( + shift=fastvideo_args.pipeline_config.flow_shift) + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage( + stage_name="image_encoding_stage", + stage=ImageEncodingStage( + image_encoder=self.get_module("image_encoder"), + image_processor=self.get_module("image_processor"), + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"))) + + self.add_stage(stage_name="image_latent_preparation_stage", + stage=ImageVAEEncodingStage(vae=self.get_module("vae"))) + + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + +class WanLingBotImageToVideoPipeline(WanGameActionImageToVideoPipeline): + pass + + +EntryClass = [WanGameActionImageToVideoPipeline, WanLingBotImageToVideoPipeline] diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index dc00f8e42..64cddefaf 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -33,6 +33,8 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { "Cosmos2_5Pipeline": "cosmos", "MatrixGamePipeline": "matrixgame", "MatrixGameCausalDMDPipeline": "matrixgame", + "WanGameActionImageToVideoPipeline": "wan", + "WanGameCausalDMDPipeline": "wan", "LongCatPipeline": "longcat", "LongCatImageToVideoPipeline": "longcat", "LongCatVideoContinuationPipeline": "longcat", diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 1e0834fc2..4ececf37b 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -845,7 +845,18 @@ class MatrixGameImageVAEEncodingStage(ImageVAEEncodingStage): num_frames = batch.num_frames if isinstance( batch.num_frames, int) else batch.num_frames[0] + def _gpu_mem(label=""): + a = torch.cuda.memory_allocated() / 1024**3 + r = torch.cuda.memory_reserved() / 1024**3 + logger.info("VAE stage [%s]: alloc=%.2fGiB, reserved=%.2fGiB", label, a, r) + + _gpu_mem("before vae.to(GPU)") self.vae = self.vae.to(get_local_torch_device()) + _gpu_mem("after vae.to(GPU)") + logger.info("VAE use_feature_cache=%s, use_tiling=%s, use_temporal_tiling=%s", + getattr(self.vae, 'use_feature_cache', 'N/A'), + getattr(self.vae, 'use_tiling', 'N/A'), + getattr(self.vae, 'use_temporal_tiling', 'N/A')) # Process single image for I2V (latent dimensions computed but not used directly) @@ -906,15 +917,23 @@ class MatrixGameImageVAEEncodingStage(ImageVAEEncodingStage): vae_autocast_enabled = ( vae_dtype != torch.float32) and not fastvideo_args.disable_autocast - # Encode Image - with torch.autocast(device_type="cuda", - dtype=vae_dtype, - enabled=vae_autocast_enabled): + _gpu_mem("before encode") + logger.info("video_condition shape=%s dtype=%s device=%s", + video_condition.shape, video_condition.dtype, + video_condition.device) + + # Encode Image (no_grad avoids autograd graph that would retain all + # intermediate activations across the feature-cache loop iterations) + with torch.no_grad(), torch.autocast( + device_type="cuda", dtype=vae_dtype, + enabled=vae_autocast_enabled): if fastvideo_args.pipeline_config.vae_tiling: self.vae.enable_tiling() if not vae_autocast_enabled: video_condition = video_condition.to(vae_dtype) + _gpu_mem("right before vae.encode()") encoder_output = self.vae.encode(video_condition) + _gpu_mem("after vae.encode()") # MatrixGame uses deterministic VAE encode for the first-frame conditioning. # Sampling would inject random noise into the cond_concat tensor and destroy the action guidance. diff --git a/fastvideo/pipelines/stages/matrixgame_denoising.py b/fastvideo/pipelines/stages/matrixgame_denoising.py index 3b442b068..329720a15 100644 --- a/fastvideo/pipelines/stages/matrixgame_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame_denoising.py @@ -313,32 +313,43 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): attention_head_dim = getattr( self.transformer, 'attention_head_dim', self.transformer.hidden_size // num_attention_heads) + # WanGame PRoPE stores cat([rope, prope], dim=-1) in cache + head_dim_mult = getattr(self.transformer, + '_kv_cache_head_dim_multiplier', 1) + attention_head_dim = attention_head_dim * head_dim_mult if self.local_attn_size != -1: kv_cache_size = self.local_attn_size * self.frame_seq_length else: kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames + # WanGame uses None-initialized cache (concat-based); MatrixGame uses + # zero-tensor cache with index tracking. + use_none_init = head_dim_mult > 1 + for _ in range(self.num_transformer_blocks): - kv_cache.append({ - "k": - torch.zeros([ - batch_size, kv_cache_size, num_attention_heads, - attention_head_dim - ], - dtype=dtype, - device=device), - "v": - torch.zeros([ - batch_size, kv_cache_size, num_attention_heads, - attention_head_dim - ], - dtype=dtype, - device=device), - "global_end_index": - torch.tensor([0], dtype=torch.long, device=device), - "local_end_index": - torch.tensor([0], dtype=torch.long, device=device), - }) + if use_none_init: + kv_cache.append({"k": None, "v": None}) + else: + kv_cache.append({ + "k": + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device), + }) return kv_cache @@ -429,6 +440,227 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): }) return crossattn_cache + def _denoise_one_step( + self, + current_latents: torch.Tensor, + noise_latents_btchw: torch.Tensor, + batch: ForwardBatch, + start_index: int, + current_num_frames: int, + timestep: torch.Tensor, + step_idx: int, + next_timestep: torch.Tensor | None, + ctx: BlockProcessingContext, + action_kwargs: dict[str, Any], + noise_generator: Callable[[tuple, torch.dtype, int], torch.Tensor] + | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run one denoising step. Returns (updated_latents, updated_noise_latents_btchw).""" + prompt_embeds = batch.prompt_embeds + t_cur = timestep + + if ctx.boundary_timestep is not None and t_cur < ctx.boundary_timestep: + current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer + else: + current_model = self.transformer + + noise_latents = noise_latents_btchw.clone() + latent_model_input = current_latents.to(ctx.target_dtype) + + # For models that don't handle image_latent internally (e.g., WanGame), + # concatenate image conditioning along the channel dimension. + # MatrixGame handles this inside its forward() via forward_context. + if (batch.image_latent is not None + and not getattr(current_model, '_concatenates_image_latent', + False)): + image_cond = batch.image_latent.to(ctx.target_dtype) + end = start_index + current_num_frames + if image_cond.shape[2] >= end: + image_cond = image_cond[:, :, start_index:end] + elif image_cond.shape[2] > start_index: + image_cond = image_cond[:, :, start_index:] + pad_frames = current_num_frames - image_cond.shape[2] + if pad_frames > 0: + pad = torch.zeros(image_cond.shape[0], + image_cond.shape[1], + pad_frames, + image_cond.shape[3], + image_cond.shape[4], + device=image_cond.device, + dtype=image_cond.dtype) + image_cond = torch.cat([image_cond, pad], dim=2) + else: + image_cond = torch.zeros(image_cond.shape[0], + image_cond.shape[1], + current_num_frames, + image_cond.shape[3], + image_cond.shape[4], + device=image_cond.device, + dtype=image_cond.dtype) + latent_model_input = torch.cat([latent_model_input, image_cond], + dim=1) + + independent_first_frame = getattr(self.transformer, + 'independent_first_frame', False) + if batch.image_latent is not None and independent_first_frame and start_index == 0: + latent_model_input = torch.cat([ + latent_model_input, + batch.image_latent.to(ctx.target_dtype) + ], + dim=2) + + # t_expand needs to be [batch * frames] to match flattened pred_noise/noise_latents + t_expand = t_cur.repeat(latent_model_input.shape[0] * + current_num_frames) + + # Build attention metadata if VSA is available + if vsa_available and self.attn_backend == VideoSparseAttentionBackend: + self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls( + ) + if self.attn_metadata_builder_cls is not None: + self.attn_metadata_builder = self.attn_metadata_builder_cls( + ) + h, w = current_latents.shape[-2:] + attn_metadata = self.attn_metadata_builder.build( + current_timestep=step_idx, + raw_latent_shape=(current_num_frames, h, w), + patch_size=ctx.fastvideo_args.pipeline_config. + dit_config.patch_size, + STA_param=batch.STA_param, + VSA_sparsity=ctx.fastvideo_args.VSA_sparsity, + device=get_local_torch_device(), + ) + assert attn_metadata is not None, "attn_metadata cannot be None" + else: + attn_metadata = None + else: + attn_metadata = None + + with torch.autocast(device_type="cuda", + dtype=ctx.target_dtype, + enabled=ctx.autocast_enabled), \ + set_forward_context(current_timestep=step_idx, + attn_metadata=attn_metadata, + forward_batch=batch): + # Expand timestep to per-frame format [batch, num_frames] for causal model + t_expanded_noise = t_cur * torch.ones( + (latent_model_input.shape[0], current_num_frames), + device=latent_model_input.device, + dtype=torch.long) + + model_kwargs = { + "kv_cache": ctx.get_kv_cache(t_cur), + "crossattn_cache": ctx.crossattn_cache, + "current_start": start_index * self.frame_seq_length, + "start_frame": start_index, + } + + if self.use_action_module and current_model == self.transformer: + model_kwargs.update({ + "kv_cache_mouse": + ctx.kv_cache_mouse, + "kv_cache_keyboard": + ctx.kv_cache_keyboard, + }) + model_kwargs.update(action_kwargs) + + # For models with discrete action embedding (WanGame-style): + # convert keyboard one-hot to scalar action = trans_label * 9 + # and provide identity camera matrices for PRoPE + if getattr(current_model, '_uses_discrete_action', False): + dev = latent_model_input.device + b = latent_model_input.shape[0] + vae_ratio = 4 + start_pixel = (0 if start_index == 0 else 1 + vae_ratio * + (start_index - 1)) + if (batch.keyboard_cond is not None + and start_pixel < batch.keyboard_cond.shape[1]): + kbd = batch.keyboard_cond[:, start_pixel, :].to(dev) + has_key = kbd.sum(dim=-1) > 0 # [B] + trans_label = (kbd.argmax(dim=-1) + 1) * has_key.long() + else: + trans_label = torch.zeros(b, device=dev, dtype=torch.long) + action_value = trans_label * 9 # composite: trans * 9 + rot + model_kwargs["action"] = action_value.unsqueeze(1).expand( + -1, current_num_frames).to(dev) + # Identity viewmats and default normalized intrinsics for PRoPE + model_kwargs["viewmats"] = torch.eye( + 4, device=dev, dtype=ctx.target_dtype + ).unsqueeze(0).unsqueeze(0).expand( + b, current_num_frames, -1, -1).contiguous() + default_K = torch.tensor( + [[0.5052, 0.0, 0.5], + [0.0, 0.8979, 0.5], + [0.0, 0.0, 1.0]], + device=dev, dtype=ctx.target_dtype) + model_kwargs["Ks"] = default_K.unsqueeze(0).unsqueeze( + 0).expand(b, current_num_frames, -1, -1).contiguous() + + pred_noise_btchw = current_model( + latent_model_input, + prompt_embeds, + t_expanded_noise, + **ctx.image_kwargs, + **ctx.pos_cond_kwargs, + **model_kwargs, + ).permute(0, 2, 1, 3, 4) + + if ctx.boundary_timestep is not None and t_cur >= ctx.boundary_timestep: + pred_video_btchw = pred_noise_to_x_bound( + pred_noise=pred_noise_btchw.flatten(0, 1), + noise_input_latent=noise_latents.flatten(0, 1), + timestep=t_expand, + boundary_timestep=torch.ones_like(t_expand) * + ctx.boundary_timestep, + scheduler=self.scheduler).unflatten( + 0, pred_noise_btchw.shape[:2]) + else: + pred_video_btchw = pred_noise_to_pred_video( + pred_noise=pred_noise_btchw.flatten(0, 1), + noise_input_latent=noise_latents.flatten(0, 1), + timestep=t_expand, + scheduler=self.scheduler).unflatten( + 0, pred_noise_btchw.shape[:2]) + + if next_timestep is not None: + next_t = next_timestep * torch.ones( + [1], dtype=torch.long, device=pred_video_btchw.device) + + # Use custom noise generator if provided (for streaming), else generate + if noise_generator is not None: + noise = noise_generator(pred_video_btchw.shape, + pred_video_btchw.dtype, step_idx) + else: + noise = torch.randn( + pred_video_btchw.shape, + dtype=pred_video_btchw.dtype, + generator=(batch.generator[0] if isinstance( + batch.generator, list) else batch.generator)).to( + pred_video_btchw.device) + + noise_btchw = noise + if ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and step_idx < len( + ctx.high_noise_timesteps) - 1: + noise_latents_btchw = self.scheduler.add_noise_high( + pred_video_btchw.flatten(0, 1), + noise_btchw.flatten(0, 1), next_t, + torch.ones_like(next_t) * + ctx.boundary_timestep).unflatten( + 0, pred_video_btchw.shape[:2]) + elif ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and step_idx == len( + ctx.high_noise_timesteps) - 1: + noise_latents_btchw = pred_video_btchw + else: + noise_latents_btchw = self.scheduler.add_noise( + pred_video_btchw.flatten(0, + 1), noise_btchw.flatten(0, 1), + next_t).unflatten(0, pred_video_btchw.shape[:2]) + current_latents = noise_latents_btchw.permute(0, 2, 1, 3, 4) + else: + current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4) + + return current_latents, noise_latents_btchw + def _process_single_block( self, current_latents: torch.Tensor, @@ -442,144 +674,23 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): | None = None, progress_bar: Any | None = None, ) -> torch.Tensor: - prompt_embeds = batch.prompt_embeds noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) for i, t_cur in enumerate(timesteps): - if ctx.boundary_timestep is not None and t_cur < ctx.boundary_timestep: - current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer - else: - current_model = self.transformer - - noise_latents = noise_latents_btchw.clone() - latent_model_input = current_latents.to(ctx.target_dtype) - - independent_first_frame = getattr(self.transformer, - 'independent_first_frame', False) - if batch.image_latent is not None and independent_first_frame and start_index == 0: - latent_model_input = torch.cat([ - latent_model_input, - batch.image_latent.to(ctx.target_dtype) - ], - dim=2) - - # t_expand needs to be [batch * frames] to match flattened pred_noise/noise_latents - t_expand = t_cur.repeat(latent_model_input.shape[0] * - current_num_frames) - - # Build attention metadata if VSA is available - if vsa_available and self.attn_backend == VideoSparseAttentionBackend: - self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls( - ) - if self.attn_metadata_builder_cls is not None: - self.attn_metadata_builder = self.attn_metadata_builder_cls( - ) - h, w = current_latents.shape[-2:] - attn_metadata = self.attn_metadata_builder.build( - current_timestep=i, - raw_latent_shape=(current_num_frames, h, w), - patch_size=ctx.fastvideo_args.pipeline_config. - dit_config.patch_size, - STA_param=batch.STA_param, - VSA_sparsity=ctx.fastvideo_args.VSA_sparsity, - device=get_local_torch_device(), - ) - assert attn_metadata is not None, "attn_metadata cannot be None" - else: - attn_metadata = None - else: - attn_metadata = None - - with torch.autocast(device_type="cuda", - dtype=ctx.target_dtype, - enabled=ctx.autocast_enabled), \ - set_forward_context(current_timestep=i, - attn_metadata=attn_metadata, - forward_batch=batch): - # Expand timestep to per-frame format [batch, num_frames] for causal model - t_expanded_noise = t_cur * torch.ones( - (latent_model_input.shape[0], current_num_frames), - device=latent_model_input.device, - dtype=torch.long) - - model_kwargs = { - "kv_cache": ctx.get_kv_cache(t_cur), - "crossattn_cache": ctx.crossattn_cache, - "current_start": start_index * self.frame_seq_length, - "start_frame": start_index, - } - - if self.use_action_module and current_model == self.transformer: - model_kwargs.update({ - "kv_cache_mouse": - ctx.kv_cache_mouse, - "kv_cache_keyboard": - ctx.kv_cache_keyboard, - }) - model_kwargs.update(action_kwargs) - - pred_noise_btchw = current_model( - latent_model_input, - prompt_embeds, - t_expanded_noise, - **ctx.image_kwargs, - **ctx.pos_cond_kwargs, - **model_kwargs, - ).permute(0, 2, 1, 3, 4) - - if ctx.boundary_timestep is not None and t_cur >= ctx.boundary_timestep: - pred_video_btchw = pred_noise_to_x_bound( - pred_noise=pred_noise_btchw.flatten(0, 1), - noise_input_latent=noise_latents.flatten(0, 1), - timestep=t_expand, - boundary_timestep=torch.ones_like(t_expand) * - ctx.boundary_timestep, - scheduler=self.scheduler).unflatten( - 0, pred_noise_btchw.shape[:2]) - else: - pred_video_btchw = pred_noise_to_pred_video( - pred_noise=pred_noise_btchw.flatten(0, 1), - noise_input_latent=noise_latents.flatten(0, 1), - timestep=t_expand, - scheduler=self.scheduler).unflatten( - 0, pred_noise_btchw.shape[:2]) - - if i < len(timesteps) - 1: - next_timestep = timesteps[i + 1] * torch.ones( - [1], dtype=torch.long, device=pred_video_btchw.device) - - # Use custom noise generator if provided (for streaming), else generate - if noise_generator is not None: - noise = noise_generator(pred_video_btchw.shape, - pred_video_btchw.dtype, i) - else: - noise = torch.randn( - pred_video_btchw.shape, - dtype=pred_video_btchw.dtype, - generator=(batch.generator[0] if isinstance( - batch.generator, list) else batch.generator)).to( - pred_video_btchw.device) - - noise_btchw = noise - if ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and i < len( - ctx.high_noise_timesteps) - 1: - noise_latents_btchw = self.scheduler.add_noise_high( - pred_video_btchw.flatten(0, 1), - noise_btchw.flatten(0, 1), next_timestep, - torch.ones_like(next_timestep) * - ctx.boundary_timestep).unflatten( - 0, pred_video_btchw.shape[:2]) - elif ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and i == len( - ctx.high_noise_timesteps) - 1: - noise_latents_btchw = pred_video_btchw - else: - noise_latents_btchw = self.scheduler.add_noise( - pred_video_btchw.flatten(0, - 1), noise_btchw.flatten(0, 1), - next_timestep).unflatten(0, pred_video_btchw.shape[:2]) - current_latents = noise_latents_btchw.permute(0, 2, 1, 3, 4) - else: - current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4) + next_timestep = timesteps[i + 1] if i < len(timesteps) - 1 else None + current_latents, noise_latents_btchw = self._denoise_one_step( + current_latents=current_latents, + noise_latents_btchw=noise_latents_btchw, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + timestep=t_cur, + step_idx=i, + next_timestep=next_timestep, + ctx=ctx, + action_kwargs=action_kwargs, + noise_generator=noise_generator, + ) if progress_bar is not None: progress_bar.update() @@ -605,6 +716,36 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): dtype=torch.long) * int(context_noise) context_bcthw = current_latents.to(ctx.target_dtype) + # Channel concatenation for models that don't handle it internally + if (batch.image_latent is not None + and not getattr(self.transformer, '_concatenates_image_latent', + False)): + image_cond = batch.image_latent.to(ctx.target_dtype) + end = start_index + current_num_frames + if image_cond.shape[2] >= end: + image_cond = image_cond[:, :, start_index:end] + elif image_cond.shape[2] > start_index: + image_cond = image_cond[:, :, start_index:] + pad_frames = current_num_frames - image_cond.shape[2] + if pad_frames > 0: + pad = torch.zeros(image_cond.shape[0], + image_cond.shape[1], + pad_frames, + image_cond.shape[3], + image_cond.shape[4], + device=image_cond.device, + dtype=image_cond.dtype) + image_cond = torch.cat([image_cond, pad], dim=2) + else: + image_cond = torch.zeros(image_cond.shape[0], + image_cond.shape[1], + current_num_frames, + image_cond.shape[3], + image_cond.shape[4], + device=image_cond.device, + dtype=image_cond.dtype) + context_bcthw = torch.cat([context_bcthw, image_cond], dim=1) + with torch.autocast(device_type="cuda", dtype=ctx.target_dtype, enabled=ctx.autocast_enabled), \ @@ -628,6 +769,37 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): }) context_model_kwargs.update(action_kwargs) + # Discrete action and camera matrices for WanGame-style models + if getattr(self.transformer, '_uses_discrete_action', False): + dev = latents_device + b = context_bcthw.shape[0] + vae_ratio = 4 + start_pixel = (0 if start_index == 0 else 1 + vae_ratio * + (start_index - 1)) + if (batch.keyboard_cond is not None + and start_pixel < batch.keyboard_cond.shape[1]): + kbd = batch.keyboard_cond[:, start_pixel, :].to(dev) + has_key = kbd.sum(dim=-1) > 0 + trans_label = (kbd.argmax(dim=-1) + 1) * has_key.long() + else: + trans_label = torch.zeros(b, device=dev, dtype=torch.long) + action_value = trans_label * 9 + context_model_kwargs["action"] = ( + action_value.unsqueeze(1).expand( + -1, current_num_frames).to(dev)) + context_model_kwargs["viewmats"] = torch.eye( + 4, device=dev, dtype=ctx.target_dtype + ).unsqueeze(0).unsqueeze(0).expand( + b, current_num_frames, -1, -1).contiguous() + default_K = torch.tensor( + [[0.5052, 0.0, 0.5], + [0.0, 0.8979, 0.5], + [0.0, 0.0, 1.0]], + device=dev, dtype=ctx.target_dtype) + context_model_kwargs["Ks"] = default_K.unsqueeze( + 0).unsqueeze(0).expand( + b, current_num_frames, -1, -1).contiguous() + if ctx.boundary_timestep is not None and self.transformer_2 is not None: self.transformer_2( context_bcthw, @@ -641,6 +813,10 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): **ctx.pos_cond_kwargs, ) + # WanGame uses is_cache=True to store KV without prepending old cache + if getattr(self.transformer, '_uses_discrete_action', False): + context_model_kwargs["is_cache"] = True + self.transformer( context_bcthw, prompt_embeds, diff --git a/fastvideo/pipelines/stages/multi_user_engine.py b/fastvideo/pipelines/stages/multi_user_engine.py new file mode 100644 index 000000000..4c094ee92 --- /dev/null +++ b/fastvideo/pipelines/stages/multi_user_engine.py @@ -0,0 +1,695 @@ +"""Multi-user engine implementing ORCA-style batching for streaming generation. + +Groups users at the same (block_idx, denoising_step) and batches their forward +passes together, sharing model weights while maintaining per-user KV caches. +""" +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from typing import Any, Iterator + +import torch + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.matrixgame_denoising import ( + BlockProcessingContext, + MatrixGameCausalDenoisingStage, +) + +logger = init_logger(__name__) + + +@dataclass +class CompletedResult: + """Result for a user whose block generation finished.""" + user_id: str + output_batch: ForwardBatch + + +@dataclass +class UserSession: + """Per-user state for multi-user streaming.""" + user_id: str + ctx: BlockProcessingContext + vae_cache: list[torch.Tensor | None] | None + batch: ForwardBatch + fastvideo_args: FastVideoArgs + + # Denoising state (set when a step is in progress) + denoising_step: int = -1 # -1 means no pending work + current_latents: torch.Tensor | None = None + noise_latents_btchw: torch.Tensor | None = None + action_kwargs: dict[str, Any] = field(default_factory=dict) + block_start_index: int = 0 + block_num_frames: int = 0 + dit_elapsed_ms: float = 0.0 # accumulated DiT time across denoising steps + + +class MultiUserEngine: + """ORCA-style scheduler that batches multiple users on a single GPU.""" + + def __init__(self, pipeline, disable_batching: bool = False): + self.pipeline = pipeline + self.disable_batching = disable_batching + if not pipeline.post_init_called: + pipeline.post_init() + self.denoiser: MatrixGameCausalDenoisingStage = ( + pipeline._stage_name_mapping["denoising_stage"] + ) + self.decoder = pipeline._stage_name_mapping["decoding_stage"] + self.users: dict[str, UserSession] = {} + + def add_user( + self, + user_id: str, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> None: + """Initialize a new user session (runs preprocessing + streaming_reset).""" + def _gpu_mem(): + a = torch.cuda.memory_allocated() / 1024**3 + r = torch.cuda.memory_reserved() / 1024**3 + return f"alloc={a:.2f}GiB, reserved={r:.2f}GiB" + + logger.info("add_user(%s) start: %s", user_id, _gpu_mem()) + logger.info("add_user(%s) memory summary:\n%s", + user_id, torch.cuda.memory_summary()) + + if user_id in self.users: + logger.warning("User %s already exists, removing first", user_id) + self.remove_user(user_id) + + # Free any existing streaming state FIRST to make room on GPU. + # After warmup, the denoiser holds KV caches + noise pool. + import gc + logger.info("add_user(%s) denoiser._streaming_initialized=%s, ctx=%s", + user_id, self.denoiser._streaming_initialized, + self.denoiser._streaming_ctx is not None) + if self.denoiser._streaming_initialized: + ctx = self.denoiser._streaming_ctx + if ctx is not None: + # Explicitly delete large tensors to break any reference cycles + ctx.kv_cache1 = None + ctx.kv_cache2 = None + ctx.kv_cache_mouse = None + ctx.kv_cache_keyboard = None + ctx.crossattn_cache = None + ctx.noise_pool = None + if ctx.batch is not None: + ctx.batch.latents = None + ctx.batch.image_latent = None + ctx.batch.prompt_embeds = None + ctx.batch = None + self.denoiser._streaming_ctx = None + self.denoiser._streaming_initialized = False + logger.info("add_user(%s) after ctx cleanup: %s", user_id, _gpu_mem()) + gc.collect() + torch.cuda.empty_cache() + logger.info("add_user(%s) after gc+empty_cache: %s", user_id, _gpu_mem()) + logger.info("add_user(%s) post-cleanup summary:\n%s", + user_id, torch.cuda.memory_summary()) + + # Run preprocessing stages + if not self.pipeline.post_init_called: + self.pipeline.post_init() + + # Ensure VAE temporal tiling is enabled to avoid OOM when encoding + # the full video_condition [1,3,num_frames,H,W] through 3D conv layers. + # Enable on both the pipeline's VAE and the stage's VAE reference. + vae = self.pipeline.get_module("vae", None) + if vae is not None: + vae.enable_tiling(use_temporal_tiling=True) + logger.info("Pipeline VAE tiling: use_tiling=%s, temporal=%s, min_frames=%s, id=%s", + vae.use_tiling, vae.use_temporal_tiling, + getattr(vae, 'tile_sample_min_num_frames', 'N/A'), id(vae)) + img_stage = self.pipeline._stage_name_mapping.get("image_latent_preparation_stage") + if img_stage is not None and hasattr(img_stage, 'vae'): + img_stage.vae.enable_tiling(use_temporal_tiling=True) + logger.info("Stage VAE tiling: use_tiling=%s, temporal=%s, min_frames=%s, id=%s", + img_stage.vae.use_tiling, img_stage.vae.use_temporal_tiling, + getattr(img_stage.vae, 'tile_sample_min_num_frames', 'N/A'), + id(img_stage.vae)) + + stages_to_run = [ + "input_validation_stage", "prompt_encoding_stage", + "image_encoding_stage", "conditioning_stage", + "latent_preparation_stage", "image_latent_preparation_stage", + ] + for stage_name in stages_to_run: + if stage_name in self.pipeline._stage_name_mapping: + batch = self.pipeline._stage_name_mapping[stage_name].forward( + batch, fastvideo_args) + logger.info("add_user(%s) after %s: %s", + user_id, stage_name, _gpu_mem()) + + # Initialize denoiser state for this user (creates KV caches etc.) + # We call streaming_reset to set up the context, then steal it + logger.info("add_user(%s) before streaming_reset: %s", user_id, _gpu_mem()) + self.denoiser.streaming_reset(batch, fastvideo_args) + ctx = self.denoiser._streaming_ctx + assert ctx is not None + # Detach from denoiser so it doesn't interfere with other users + self.denoiser._streaming_ctx = None + self.denoiser._streaming_initialized = False + + self.users[user_id] = UserSession( + user_id=user_id, + ctx=ctx, + vae_cache=None, + batch=batch, + fastvideo_args=fastvideo_args, + ) + logger.info("Added user %s (total: %d)", user_id, len(self.users)) + + def remove_user(self, user_id: str) -> None: + """Remove a user session and free GPU memory.""" + session = self.users.pop(user_id, None) + if session is not None: + # Let GC handle tensor cleanup + session.ctx = None # type: ignore + session.vae_cache = None + session.current_latents = None + session.noise_latents_btchw = None + logger.info("Removed user %s (remaining: %d)", + user_id, len(self.users)) + + def submit_step( + self, + user_id: str, + keyboard_action: torch.Tensor | None, + mouse_action: torch.Tensor | None, + ) -> None: + """Queue a block generation request for a user.""" + session = self.users.get(user_id) + if session is None: + raise KeyError(f"User {user_id} not found") + + ctx = session.ctx + if ctx.block_idx >= len(ctx.block_sizes): + logger.warning("User %s has no more blocks to generate", user_id) + return + + batch = session.batch + latents = batch.latents + assert latents is not None + + current_num_frames = ctx.block_sizes[ctx.block_idx] + start_index = ctx.start_index + + current_latents = latents[:, :, start_index:start_index + + current_num_frames, :, :] + + # Update batch with new actions + if keyboard_action is not None or mouse_action is not None: + vae_ratio = 4 + start_frame = 0 if start_index == 0 else 1 + vae_ratio * ( + start_index - 1) + + if keyboard_action is not None: + n = keyboard_action.shape[1] + batch.keyboard_cond[:, start_frame:start_frame + + n] = keyboard_action.to( + batch.keyboard_cond.device) + if mouse_action is not None: + n = mouse_action.shape[1] + batch.mouse_cond[:, start_frame:start_frame + + n] = mouse_action.to(batch.mouse_cond.device) + + action_kwargs = self.denoiser._prepare_action_kwargs( + batch, start_index, current_num_frames) + + # Set up denoising state + session.denoising_step = 0 + session.dit_elapsed_ms = 0.0 + session.current_latents = current_latents + session.noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) + session.action_kwargs = action_kwargs + session.block_start_index = start_index + session.block_num_frames = current_num_frames + + def has_pending_work(self) -> bool: + """Check if any user has pending denoising work.""" + return any(s.denoising_step >= 0 for s in self.users.values()) + + def run_iteration(self) -> Iterator[CompletedResult]: + """Run one denoising step for all ready users (ORCA iteration-level scheduling). + + Groups users by (block_idx, denoising_step) and batches compatible groups. + Yields completed results immediately after each user's VAE decode finishes, + so responses can be sent without waiting for all users in the batch. + """ + # Group users with pending work by (block_idx, denoising_step) + groups: dict[tuple[int, int], list[UserSession]] = {} + for session in self.users.values(): + if session.denoising_step < 0: + continue + key = (session.ctx.block_idx, session.denoising_step) + groups.setdefault(key, []).append(session) + + for (block_idx, step_idx), sessions in groups.items(): + timesteps = sessions[0].ctx.timesteps + t_cur = timesteps[step_idx] + next_timestep = timesteps[step_idx + 1] if step_idx < len(timesteps) - 1 else None + + torch.cuda.synchronize() + t0 = time.perf_counter() + + user_ids = [s.user_id[:8] for s in sessions] + if len(sessions) == 1 or self.disable_batching: + # Process each user individually + for s in sessions: + logger.info("ORCA step: block=%d step=%d single user=%s", + block_idx, step_idx, s.user_id[:8]) + self._run_single_user_step(s, t_cur, step_idx, next_timestep) + else: + # Multiple users at same state - batch them + logger.info("ORCA step: block=%d step=%d BATCHED %d users=%s", + block_idx, step_idx, len(sessions), user_ids) + self._run_batched_step(sessions, t_cur, step_idx, next_timestep) + + torch.cuda.synchronize() + step_ms = (time.perf_counter() - t0) * 1000 + + # Advance step counters and finalize completed users + for s in sessions: + s.dit_elapsed_ms += step_ms + s.denoising_step += 1 + if s.denoising_step >= len(timesteps): + # Denoising complete — VAE decode and yield immediately + result = self._finalize_block(s) + if result is not None: + yield result + + def _run_single_user_step( + self, + session: UserSession, + timestep: torch.Tensor, + step_idx: int, + next_timestep: torch.Tensor | None, + ) -> None: + """Run one denoising step for a single user (no batching).""" + noise_gen = self._make_noise_generator(session) + + current_latents, noise_latents_btchw = self.denoiser._denoise_one_step( + current_latents=session.current_latents, + noise_latents_btchw=session.noise_latents_btchw, + batch=session.batch, + start_index=session.block_start_index, + current_num_frames=session.block_num_frames, + timestep=timestep, + step_idx=step_idx, + next_timestep=next_timestep, + ctx=session.ctx, + action_kwargs=session.action_kwargs, + noise_generator=noise_gen, + ) + session.current_latents = current_latents + session.noise_latents_btchw = noise_latents_btchw + + def _run_batched_step( + self, + sessions: list[UserSession], + timestep: torch.Tensor, + step_idx: int, + next_timestep: torch.Tensor | None, + ) -> None: + """Run one denoising step for multiple users batched together.""" + n = len(sessions) + + # Concatenate along batch dimension + batched_latents = torch.cat( + [s.current_latents for s in sessions], dim=0) + batched_noise = torch.cat( + [s.noise_latents_btchw for s in sessions], dim=0) + + # Build a merged ForwardBatch with concatenated per-user tensors + merged_batch = self._merge_batches(sessions) + + # Build merged context with concatenated KV caches + merged_ctx = self._merge_contexts(sessions, merged_batch) + + # Merge action kwargs + merged_action_kwargs = self._merge_action_kwargs(sessions) + + # Create noise generator that concatenates per-user noise + device = batched_latents.device + + def batched_noise_gen(shape, dtype, si): + noises = [] + for s in sessions: + gen = self._make_noise_generator(s) + if gen is not None: + per_user_shape = (1,) + shape[1:] + noises.append(gen(per_user_shape, dtype, si)) + else: + noises.append(torch.randn( + (1,) + shape[1:], dtype=dtype).to(device)) + return torch.cat(noises, dim=0) + + current_latents, noise_latents_btchw = self.denoiser._denoise_one_step( + current_latents=batched_latents, + noise_latents_btchw=batched_noise, + batch=merged_batch, + start_index=sessions[0].block_start_index, + current_num_frames=sessions[0].block_num_frames, + timestep=timestep, + step_idx=step_idx, + next_timestep=next_timestep, + ctx=merged_ctx, + action_kwargs=merged_action_kwargs, + noise_generator=batched_noise_gen, + ) + + # Unbatch results back to per-user tensors + latents_list = current_latents.split(1, dim=0) + noise_list = noise_latents_btchw.split(1, dim=0) + + # Restore per-user KV caches from the batched caches + self._split_contexts(merged_ctx, sessions) + + for i, s in enumerate(sessions): + s.current_latents = latents_list[i] + s.noise_latents_btchw = noise_list[i] + + def _merge_batches(self, sessions: list[UserSession]) -> ForwardBatch: + """Create a merged ForwardBatch with concatenated user tensors.""" + ref = sessions[0].batch + # prompt_embeds and image_embeds are shared (same game) + # latents, keyboard_cond, mouse_cond differ per user + merged = ForwardBatch.__new__(ForwardBatch) + merged.__dict__.update(ref.__dict__) + + # image_latent: cat along batch dim if present + if ref.image_latent is not None: + merged.image_latent = torch.cat( + [s.batch.image_latent for s in sessions], dim=0) + + return merged + + def _merge_contexts( + self, sessions: list[UserSession], + merged_batch: ForwardBatch | None = None, + ) -> BlockProcessingContext: + """Create a merged BlockProcessingContext with concatenated KV caches.""" + ref = sessions[0].ctx + n = len(sessions) + + # Concatenate KV caches along batch dimension + merged_kv1 = self._cat_kv_caches([s.ctx.kv_cache1 for s in sessions]) + merged_kv2 = None + if ref.kv_cache2 is not None: + merged_kv2 = self._cat_kv_caches( + [s.ctx.kv_cache2 for s in sessions]) + + merged_crossattn = self._cat_crossattn_caches( + [s.ctx.crossattn_cache for s in sessions]) + + merged_kv_mouse = None + merged_kv_keyboard = None + if ref.kv_cache_mouse is not None: + merged_kv_mouse = self._cat_action_kv_caches( + [s.ctx.kv_cache_mouse for s in sessions], + is_mouse=True) + if ref.kv_cache_keyboard is not None: + merged_kv_keyboard = self._cat_kv_caches( + [s.ctx.kv_cache_keyboard for s in sessions]) + + # image_kwargs: cat image embeds + merged_image_kwargs = dict(ref.image_kwargs) + if "encoder_hidden_states_image" in ref.image_kwargs: + embeds = ref.image_kwargs["encoder_hidden_states_image"] + if isinstance(embeds, list) and len(embeds) > 0 and torch.is_tensor(embeds[0]): + # Cat each tensor in the list across users + merged_embeds = [] + for idx in range(len(embeds)): + merged_embeds.append(torch.cat( + [s.ctx.image_kwargs["encoder_hidden_states_image"][idx] + for s in sessions], dim=0)) + merged_image_kwargs["encoder_hidden_states_image"] = merged_embeds + + if merged_batch is None: + merged_batch = self._merge_batches(sessions) + + return BlockProcessingContext( + batch=merged_batch, + block_idx=ref.block_idx, + start_index=ref.start_index, + kv_cache1=merged_kv1, + kv_cache2=merged_kv2, + kv_cache_mouse=merged_kv_mouse, + kv_cache_keyboard=merged_kv_keyboard, + crossattn_cache=merged_crossattn, + timesteps=ref.timesteps, + block_sizes=ref.block_sizes, + noise_pool=None, # noise handled per-user + fastvideo_args=ref.fastvideo_args, + target_dtype=ref.target_dtype, + autocast_enabled=ref.autocast_enabled, + boundary_timestep=ref.boundary_timestep, + high_noise_timesteps=ref.high_noise_timesteps, + context_noise=ref.context_noise, + image_kwargs=merged_image_kwargs, + pos_cond_kwargs=ref.pos_cond_kwargs, + ) + + def _cat_kv_caches( + self, caches_list: list[list[dict[str, Any]]] + ) -> list[dict[str, Any]]: + """Concatenate KV caches along batch dimension.""" + num_layers = len(caches_list[0]) + merged = [] + for layer_idx in range(num_layers): + merged.append({ + "k": torch.cat( + [c[layer_idx]["k"] for c in caches_list], dim=0), + "v": torch.cat( + [c[layer_idx]["v"] for c in caches_list], dim=0), + "global_end_index": caches_list[0][layer_idx]["global_end_index"], + "local_end_index": caches_list[0][layer_idx]["local_end_index"], + }) + return merged + + def _cat_action_kv_caches( + self, + caches_list: list[list[dict[str, Any]]], + is_mouse: bool = False, + ) -> list[dict[str, Any]]: + """Concatenate action KV caches. + + Mouse caches have shape [B*frame_seq_length, ...] so we need + to concatenate correctly. + """ + num_layers = len(caches_list[0]) + merged = [] + for layer_idx in range(num_layers): + if is_mouse: + # Mouse: [B*F, cache_size, heads, dim] -> [N*F, cache_size, heads, dim] + # Each user has [1*F, ...], so simple cat along dim 0 works + merged.append({ + "k": torch.cat( + [c[layer_idx]["k"] for c in caches_list], dim=0), + "v": torch.cat( + [c[layer_idx]["v"] for c in caches_list], dim=0), + "global_end_index": caches_list[0][layer_idx]["global_end_index"], + "local_end_index": caches_list[0][layer_idx]["local_end_index"], + }) + else: + merged.append({ + "k": torch.cat( + [c[layer_idx]["k"] for c in caches_list], dim=0), + "v": torch.cat( + [c[layer_idx]["v"] for c in caches_list], dim=0), + "global_end_index": caches_list[0][layer_idx]["global_end_index"], + "local_end_index": caches_list[0][layer_idx]["local_end_index"], + }) + return merged + + def _cat_crossattn_caches( + self, caches_list: list[list[dict[str, Any]]] + ) -> list[dict[str, Any]]: + """Concatenate cross-attention caches along batch dimension.""" + num_layers = len(caches_list[0]) + merged = [] + for layer_idx in range(num_layers): + merged.append({ + "k": torch.cat( + [c[layer_idx]["k"] for c in caches_list], dim=0), + "v": torch.cat( + [c[layer_idx]["v"] for c in caches_list], dim=0), + "is_init": caches_list[0][layer_idx]["is_init"], + }) + return merged + + def _split_contexts( + self, + merged_ctx: BlockProcessingContext, + sessions: list[UserSession], + ) -> None: + """Split merged KV caches back into per-user caches after forward pass.""" + n = len(sessions) + + self._split_kv_caches(merged_ctx.kv_cache1, + [s.ctx.kv_cache1 for s in sessions]) + if merged_ctx.kv_cache2 is not None: + self._split_kv_caches(merged_ctx.kv_cache2, + [s.ctx.kv_cache2 for s in sessions]) + + self._split_crossattn_caches(merged_ctx.crossattn_cache, + [s.ctx.crossattn_cache for s in sessions]) + + if merged_ctx.kv_cache_mouse is not None: + frame_seq_len = self.denoiser.frame_seq_length + self._split_action_kv_caches( + merged_ctx.kv_cache_mouse, + [s.ctx.kv_cache_mouse for s in sessions], + is_mouse=True, frame_seq_length=frame_seq_len) + if merged_ctx.kv_cache_keyboard is not None: + self._split_kv_caches(merged_ctx.kv_cache_keyboard, + [s.ctx.kv_cache_keyboard for s in sessions]) + + def _split_kv_caches( + self, + merged: list[dict[str, Any]], + per_user: list[list[dict[str, Any]]], + ) -> None: + """Copy data from merged KV cache back to per-user caches.""" + n = len(per_user) + for layer_idx in range(len(merged)): + k_split = merged[layer_idx]["k"].split(1, dim=0) + v_split = merged[layer_idx]["v"].split(1, dim=0) + for i in range(n): + per_user[i][layer_idx]["k"].copy_(k_split[i]) + per_user[i][layer_idx]["v"].copy_(v_split[i]) + per_user[i][layer_idx]["global_end_index"].copy_( + merged[layer_idx]["global_end_index"]) + per_user[i][layer_idx]["local_end_index"].copy_( + merged[layer_idx]["local_end_index"]) + + def _split_action_kv_caches( + self, + merged: list[dict[str, Any]], + per_user: list[list[dict[str, Any]]], + is_mouse: bool = False, + frame_seq_length: int = 1, + ) -> None: + """Split action KV caches back to per-user.""" + n = len(per_user) + for layer_idx in range(len(merged)): + if is_mouse: + # [N*F, ...] -> N chunks of [F, ...] + k_split = merged[layer_idx]["k"].split(frame_seq_length, dim=0) + v_split = merged[layer_idx]["v"].split(frame_seq_length, dim=0) + else: + k_split = merged[layer_idx]["k"].split(1, dim=0) + v_split = merged[layer_idx]["v"].split(1, dim=0) + for i in range(n): + per_user[i][layer_idx]["k"].copy_(k_split[i]) + per_user[i][layer_idx]["v"].copy_(v_split[i]) + per_user[i][layer_idx]["global_end_index"].copy_( + merged[layer_idx]["global_end_index"]) + per_user[i][layer_idx]["local_end_index"].copy_( + merged[layer_idx]["local_end_index"]) + + def _split_crossattn_caches( + self, + merged: list[dict[str, Any]], + per_user: list[list[dict[str, Any]]], + ) -> None: + """Split cross-attention caches back to per-user.""" + n = len(per_user) + for layer_idx in range(len(merged)): + k_split = merged[layer_idx]["k"].split(1, dim=0) + v_split = merged[layer_idx]["v"].split(1, dim=0) + for i in range(n): + per_user[i][layer_idx]["k"].copy_(k_split[i]) + per_user[i][layer_idx]["v"].copy_(v_split[i]) + per_user[i][layer_idx]["is_init"] = merged[layer_idx]["is_init"] + + def _merge_action_kwargs( + self, sessions: list[UserSession] + ) -> dict[str, Any]: + """Merge action kwargs by concatenating tensors along batch dim.""" + ref = sessions[0].action_kwargs + if not ref: + return {} + + merged: dict[str, Any] = {} + for key in ref: + vals = [s.action_kwargs[key] for s in sessions] + if torch.is_tensor(vals[0]): + merged[key] = torch.cat(vals, dim=0) + else: + # Scalars (like num_frame_per_block) - use ref value + merged[key] = vals[0] + return merged + + def _make_noise_generator(self, session: UserSession): + """Create a noise generator using the user's pre-allocated noise pool.""" + ctx = session.ctx + latents = session.batch.latents + + def noise_gen(shape, dtype, step_idx): + if ctx.noise_pool is not None and step_idx < len(ctx.noise_pool): + return ctx.noise_pool[step_idx][:, :shape[1], :, :, :].to( + latents.device) + else: + return torch.randn(shape, dtype=dtype).to(latents.device) + + return noise_gen + + def _finalize_block(self, session: UserSession) -> CompletedResult | None: + """Finalize a completed block: update context cache + VAE decode.""" + ctx = session.ctx + batch = session.batch + latents = batch.latents + start_index = session.block_start_index + current_num_frames = session.block_num_frames + + # Write denoised latents back + latents[:, :, start_index:start_index + + current_num_frames, :, :] = session.current_latents + + # Update KV caches with clean context + self.denoiser._update_context_cache( + current_latents=session.current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + ctx=ctx, + action_kwargs=session.action_kwargs, + context_noise=ctx.context_noise, + ) + + # Advance streaming state + old_start = ctx.start_index + ctx.start_index += current_num_frames + ctx.block_idx += 1 + + # VAE decode + torch.cuda.synchronize() + t0 = time.perf_counter() + + current_latents_for_vae = latents[:, :, start_index:start_index + + current_num_frames, :, :] + decoded_frames, session.vae_cache = self.decoder.streaming_decode( + current_latents_for_vae, + session.fastvideo_args, + cache=session.vae_cache, + is_first_chunk=(start_index == 0), + ) + torch.cuda.synchronize() + vae_ms = (time.perf_counter() - t0) * 1000 + + batch.output = decoded_frames + batch.stage_timings = {"dit_ms": session.dit_elapsed_ms, "vae_ms": vae_ms} + + # Reset denoising state + session.denoising_step = -1 + session.current_latents = None + session.noise_latents_btchw = None + + return CompletedResult(user_id=session.user_id, output_batch=batch) + diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index 0460a97e5..3545fae2c 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -75,6 +75,10 @@ class Worker: self.pipeline = build_pipeline(self.fastvideo_args) + def get_pipeline(self): + """Return the loaded pipeline instance.""" + return self.pipeline + def execute_forward(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: output_batch = self.pipeline.forward(forward_batch, self.fastvideo_args) diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py index cc09c5898..8af334cf2 100644 --- a/fastvideo/worker/multiproc_executor.py +++ b/fastvideo/worker/multiproc_executor.py @@ -39,13 +39,17 @@ logger = init_logger(__name__) class StreamingTaskType(str, Enum): """ Enumeration for different streaming task types. - + Inherits from str to allow string comparison for backward compatibility. """ RESET = "reset" STEP = "step" CLEAR = "clear" EXIT = "exit" + # Multi-user task types + USER_JOIN = "user_join" + USER_STEP = "user_step" + USER_LEAVE = "user_leave" @dataclass @@ -58,6 +62,8 @@ class StreamingTask: # For RESET tasks: batch: ForwardBatch | None = None fastvideo_args: FastVideoArgs | None = None + # For multi-user tasks: + user_id: str | None = None @dataclass @@ -66,6 +72,8 @@ class StreamingResult: task_type: StreamingTaskType output_batch: ForwardBatch | None = None error: Exception | None = None + # For multi-user results: + user_id: str | None = None class MultiprocExecutor(Executor): @@ -183,6 +191,13 @@ class MultiprocExecutor(Executor): self.collective_rpc("start_streaming_queue_loop") self._streaming_enabled = True + def enable_multi_user_streaming(self) -> None: + if self._streaming_enabled: + return + + self.collective_rpc("start_multi_user_streaming_loop") + self._streaming_enabled = True + def disable_streaming(self) -> None: if not self._streaming_enabled: return @@ -225,6 +240,42 @@ class MultiprocExecutor(Executor): self._streaming_input_queue.put( StreamingTask(task_type=StreamingTaskType.CLEAR)) + def submit_user_join(self, user_id: str, forward_batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> None: + if not self._streaming_enabled: + self.enable_streaming() + + self._streaming_input_queue.put( + StreamingTask( + task_type=StreamingTaskType.USER_JOIN, + user_id=user_id, + batch=forward_batch, + fastvideo_args=fastvideo_args, + )) + + def submit_user_step(self, user_id: str, + keyboard_action: torch.Tensor | None, + mouse_action: torch.Tensor | None) -> None: + if not self._streaming_enabled: + raise RuntimeError( + "Streaming mode not enabled. Call enable_streaming() first.") + + self._streaming_input_queue.put( + StreamingTask( + task_type=StreamingTaskType.USER_STEP, + user_id=user_id, + keyboard_action=keyboard_action, + mouse_action=mouse_action, + )) + + def submit_user_leave(self, user_id: str) -> None: + if self._streaming_enabled and self._streaming_input_queue is not None: + self._streaming_input_queue.put( + StreamingTask( + task_type=StreamingTaskType.USER_LEAVE, + user_id=user_id, + )) + def get_result(self, timeout: float | None = None) -> StreamingResult | None: if not self._streaming_enabled or self._streaming_output_queue is None: @@ -642,6 +693,11 @@ class WorkerMultiprocProc: {"status": "streaming_queue_loop_started"}) self.streaming_queue_loop() continue + if method == "start_multi_user_streaming_loop": + self.pipe.send( + {"status": "multi_user_streaming_loop_started"}) + self.multi_user_streaming_loop() + continue if method == 'execute_forward': forward_batch = kwargs['forward_batch'] fastvideo_args = kwargs['fastvideo_args'] @@ -715,11 +771,126 @@ class WorkerMultiprocProc: self.worker.execute_streaming_clear() self.streaming_output_queue.put( StreamingResult(task_type=StreamingTaskType.CLEAR)) + elif task.task_type in (StreamingTaskType.USER_JOIN, + StreamingTaskType.USER_STEP, + StreamingTaskType.USER_LEAVE): + # Switch to multi-user mode + self.multi_user_streaming_loop(first_task=task) + break except Exception as e: logger.error("Worker %d queue loop error: %s", self.rank, e) self.streaming_output_queue.put( StreamingResult(task_type=StreamingTaskType.STEP, error=e)) + def multi_user_streaming_loop( + self, first_task: StreamingTask | None = None + ) -> None: + """ORCA-style multi-user streaming loop. + + Processes USER_JOIN/USER_STEP/USER_LEAVE tasks and batches + compatible users together for each denoising step. + """ + from fastvideo.pipelines.stages.multi_user_engine import MultiUserEngine + + if self.streaming_input_queue is None or self.streaming_output_queue is None: + logger.error("Worker %d: streaming queues not initialized", + self.rank) + return + + pipeline = self.worker.get_pipeline() + try: + from ui.world_model.server.config import DISABLE_BATCHING + except ImportError: + DISABLE_BATCHING = False + engine = MultiUserEngine(pipeline, disable_batching=DISABLE_BATCHING) + logger.info("Worker %d: multi-user streaming loop started", self.rank) + + def handle_task(task: StreamingTask) -> bool: + """Handle a single task. Returns False if should exit.""" + if task.task_type == StreamingTaskType.EXIT: + return False + elif task.task_type == StreamingTaskType.USER_JOIN: + try: + engine.add_user( + task.user_id, task.batch, task.fastvideo_args) + self.streaming_output_queue.put( + StreamingResult( + task_type=StreamingTaskType.USER_JOIN, + user_id=task.user_id)) + except Exception as e: + logger.error("Worker %d user_join error for %s: %s", + self.rank, task.user_id, e) + self.streaming_output_queue.put( + StreamingResult( + task_type=StreamingTaskType.USER_JOIN, + user_id=task.user_id, error=e)) + elif task.task_type == StreamingTaskType.USER_STEP: + try: + engine.submit_step( + task.user_id, task.keyboard_action, task.mouse_action) + except Exception as e: + logger.error("Worker %d user_step submit error for %s: %s", + self.rank, task.user_id, e) + self.streaming_output_queue.put( + StreamingResult( + task_type=StreamingTaskType.USER_STEP, + user_id=task.user_id, error=e)) + elif task.task_type == StreamingTaskType.USER_LEAVE: + engine.remove_user(task.user_id) + self.streaming_output_queue.put( + StreamingResult( + task_type=StreamingTaskType.USER_LEAVE, + user_id=task.user_id)) + return True + + # Handle the first task that triggered the switch + if first_task is not None: + if not handle_task(first_task): + return + + while True: + # 1. Drain input queue (non-blocking) for new requests + while True: + try: + task = self.streaming_input_queue.get_nowait() + if not handle_task(task): + return + except queue.Empty: + break + + # 2. Run one ORCA iteration + if engine.has_pending_work(): + try: + completed = engine.run_iteration() + for result in completed: + self.streaming_output_queue.put( + StreamingResult( + task_type=StreamingTaskType.USER_STEP, + user_id=result.user_id, + output_batch=result.output_batch)) + except Exception as e: + logger.error("Worker %d ORCA iteration error: %s", + self.rank, e, exc_info=True) + # Reset all pending users to prevent infinite error loop + for uid, sess in engine.users.items(): + if sess.denoising_step >= 0: + sess.denoising_step = -1 + sess.current_latents = None + sess.noise_latents_btchw = None + self.streaming_output_queue.put( + StreamingResult( + task_type=StreamingTaskType.USER_STEP, + user_id=uid, error=e)) + + # 3. If no work, block on queue to avoid busy-waiting + if not engine.has_pending_work(): + try: + task = self.streaming_input_queue.get(timeout=0.1) + if not handle_task(task): + return + except queue.Empty: + continue + @staticmethod def setup_proc_title_and_log_prefix() -> None: dp_size = get_dp_group().world_size diff --git a/ui/world_model/client/src/App.svelte b/ui/world_model/client/src/App.svelte index e77c6ec7e..9c4a73a8e 100644 --- a/ui/world_model/client/src/App.svelte +++ b/ui/world_model/client/src/App.svelte @@ -28,6 +28,9 @@ right: false }; + let availableModels = []; + let selectedModelId = ''; + let frameBuffer = []; let isPlayingBuffer = false; let animationFrameId = null; @@ -150,6 +153,8 @@ ws.onopen = () => { connected = true; connecting = false; + // Send selected model to server + ws.send(JSON.stringify({ type: 'select_model', model_id: selectedModelId })); }; let pendingFrames = new Map(); // Track frames in decode order @@ -398,9 +403,19 @@ if (event.key === 'ArrowRight') pressedArrows.right = false; } - onMount(() => { + onMount(async () => { window.addEventListener('keydown', handleKeyDown); window.addEventListener('keyup', handleKeyUp); + + // Fetch available models + try { + const res = await fetch('/models'); + const data = await res.json(); + availableModels = data.models; + selectedModelId = data.default_model_id; + } catch (e) { + console.error('Failed to fetch models:', e); + } }); onDestroy(() => { @@ -463,7 +478,17 @@
-

Matrix-Game 2.0

+
+ {#if availableModels.length > 0} + + {:else} +

Loading...

+ {/if} +
{#if connected && gpuAssigned && timeLeft !== null} @@ -697,6 +722,31 @@ margin-bottom: 0.5rem; } + .model-selector { + justify-self: start; + } + + .model-selector select { + background: #1f2937; + color: white; + border: 1px solid #374151; + border-radius: 6px; + padding: 0.4rem 0.6rem; + font-size: 1.1rem; + font-weight: 600; + cursor: pointer; + appearance: auto; + } + + .model-selector select:disabled { + opacity: 0.5; + cursor: not-allowed; + } + + .model-selector select:hover:not(:disabled) { + border-color: #60a5fa; + } + .header-section h2 { justify-self: start; } diff --git a/ui/world_model/server/config.py b/ui/world_model/server/config.py index bdd4b061f..d274debd0 100644 --- a/ui/world_model/server/config.py +++ b/ui/world_model/server/config.py @@ -1,10 +1,34 @@ -# Model configuration -MODEL_CONFIG = { - "model_path": "FastVideo/Matrix-Game-2.0-Base-Diffusers", - "keyboard_dim": 4, - "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", +from pathlib import Path + +# Repo root: ui/world_model/server/config.py -> go up 3 levels +_REPO_ROOT = Path(__file__).resolve().parents[3] + +# Model registry - add new models here +MODEL_REGISTRY = { + "matrix-game-2.0-base": { + "name": "Matrix-Game 2.0 Base", + "model_path": "FastVideo/Matrix-Game-2.0-Base-Diffusers", + "keyboard_dim": 4, + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", + }, + "wangame-1.3b-mc-w-only-still-7k": { + "name": "WANGame 1.3B MC (W-Only, 7k)", + "model_path": "weizhou03/Wan2.1-Game-Fun-1.3B-InP-Diffusers", + "init_weights_from_safetensors": str(_REPO_ROOT / "checkpoints" / "wangame-1.3b-mc-w-only-still-7k" / "checkpoint-7000" / "transformer"), + "override_transformer_cls_name": "WanGameActionTransformer3DModel", + "override_pipeline_cls_name": "WanGameCausalDMDPipeline", + "keyboard_dim": 4, + #"image_url": "/server-assets/mc.png", + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", + # "image_path": str(Path(__file__).resolve().parent / "mc.png"), + }, } +DEFAULT_MODEL_ID = "matrix-game-2.0-base" + +# Active model configuration (set by server or user selection) +MODEL_CONFIG = MODEL_REGISTRY[DEFAULT_MODEL_ID] + # Keyboard mappings (WASD) KEYBOARD_MAP = { "w": [1, 0, 0, 0], @@ -24,12 +48,14 @@ CAMERA_MAP = { # Generation limits MAX_BLOCKS = 50 -SESSION_TIMEOUT_SECONDS = 90 +SESSION_TIMEOUT_SECONDS = 900 +MAX_USERS_PER_GPU = 16 +DISABLE_BATCHING = False # Process users sequentially (batch_size=1) instead of batching # Frame settings NUM_FRAMES = 597 FRAME_HEIGHT = 352 FRAME_WIDTH = 640 -NUM_INFERENCE_STEPS = 10 +NUM_INFERENCE_STEPS = 4 JPEG_QUALITY = 85 BATCH_SIZE = 12 diff --git a/ui/world_model/server/gpu_pool.py b/ui/world_model/server/gpu_pool.py index 154a906c9..05e130b07 100644 --- a/ui/world_model/server/gpu_pool.py +++ b/ui/world_model/server/gpu_pool.py @@ -12,7 +12,9 @@ import numpy as np import torch from config import ( - MODEL_CONFIG, NUM_FRAMES, FRAME_HEIGHT, FRAME_WIDTH, NUM_INFERENCE_STEPS + MODEL_CONFIG, MODEL_REGISTRY, DEFAULT_MODEL_ID, + NUM_FRAMES, FRAME_HEIGHT, FRAME_WIDTH, NUM_INFERENCE_STEPS, + MAX_USERS_PER_GPU, ) @@ -22,6 +24,11 @@ class CommandType(Enum): STEP = "step" RESET = "reset" SHUTDOWN = "shutdown" + # Multi-user commands + USER_JOIN = "user_join" + USER_STEP = "user_step" + USER_LEAVE = "user_leave" + RELOAD_MODEL = "reload_model" @dataclass @@ -29,6 +36,7 @@ class Command: """Command sent to GPU worker subprocess.""" type: CommandType data: dict = None + user_id: str = None def __post_init__(self): if self.data is None: @@ -41,6 +49,8 @@ class Response: success: bool frames: Optional[list[np.ndarray]] = None error: Optional[str] = None + timings: Optional[dict] = None + user_id: Optional[str] = None def gpu_worker_process( @@ -53,6 +63,7 @@ def gpu_worker_process( Worker process that runs on a single GPU. This function runs in a subprocess with CUDA_VISIBLE_DEVICES set to a single GPU. + Supports both single-user mode (legacy) and multi-user mode (ORCA batching). """ # Set CUDA_VISIBLE_DEVICES BEFORE importing torch or any CUDA code os.environ["CUDA_VISIBLE_DEVICES"] = cuda_device @@ -64,32 +75,78 @@ def gpu_worker_process( from fastvideo.configs.sample.base import SamplingParam from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.utils import align_to, shallow_asdict + from fastvideo.worker.multiproc_executor import ( + StreamingTask, StreamingTaskType, StreamingResult, + ) + + import gc + import time generator = None + current_model_config = dict(MODEL_CONFIG) - def initialize_generator(): - nonlocal generator - print(f"[GPU {gpu_id}] Loading model...") + def _gpu_mem(): + a = torch.cuda.memory_allocated() / 1024**3 + r = torch.cuda.memory_reserved() / 1024**3 + return f"alloc={a:.2f}GiB, reserved={r:.2f}GiB" - generator = StreamingVideoGenerator.from_pretrained( - MODEL_CONFIG["model_path"], + def initialize_generator(model_config=None): + nonlocal generator, current_model_config + if model_config is not None: + current_model_config = model_config + + # Free old generator if reloading + if generator is not None: + print(f"[GPU {gpu_id}] Freeing old model...") + try: + generator.shutdown() + except Exception: + pass + del generator + generator = None + gc.collect() + torch.cuda.empty_cache() + print(f"[GPU {gpu_id}] After cleanup: {_gpu_mem()}") + + print(f"[GPU {gpu_id}] Loading model: {current_model_config['model_path']}") + if current_model_config.get("init_weights_from_safetensors"): + print(f"[GPU {gpu_id}] With custom weights: {current_model_config['init_weights_from_safetensors']}") + print(f"[GPU {gpu_id}] Before model load: {_gpu_mem()}") + + load_kwargs = dict( num_gpus=1, - use_fsdp_inference=True, - dit_cpu_offload=True, + dit_layerwise_offload=False, + use_fsdp_inference=False, + dit_cpu_offload=False, vae_cpu_offload=False, text_encoder_cpu_offload=True, pin_cpu_memory=True, ) + if current_model_config.get("init_weights_from_safetensors"): + load_kwargs["init_weights_from_safetensors"] = current_model_config["init_weights_from_safetensors"] + if current_model_config.get("override_transformer_cls_name"): + load_kwargs["override_transformer_cls_name"] = current_model_config["override_transformer_cls_name"] + if current_model_config.get("override_pipeline_cls_name"): + load_kwargs["override_pipeline_cls_name"] = current_model_config["override_pipeline_cls_name"] - # Initial reset + generator = StreamingVideoGenerator.from_pretrained( + current_model_config["model_path"], + **load_kwargs, + ) + + print(f"[GPU {gpu_id}] After model load: {_gpu_mem()}") + + # Warmup: run initial reset to exercise the full pipeline once. + # This validates the model can encode 597 frames, warms up CUDA, + # and allocates KV caches. The caches will be freed when the first + # multi-user join arrives (add_user cleans up existing state). actions = { - "keyboard": torch.zeros((NUM_FRAMES, MODEL_CONFIG["keyboard_dim"])), + "keyboard": torch.zeros((NUM_FRAMES, current_model_config["keyboard_dim"])), "mouse": torch.zeros((NUM_FRAMES, 2)) } - generator.reset( prompt="", - image_path=MODEL_CONFIG["image_url"], + image_path=current_model_config.get("image_path", current_model_config["image_url"]), mouse_cond=actions["mouse"].unsqueeze(0), keyboard_cond=actions["keyboard"].unsqueeze(0), grid_sizes=torch.tensor([150, 44, 80]), @@ -99,17 +156,35 @@ def gpu_worker_process( num_inference_steps=NUM_INFERENCE_STEPS, ) + print(f"[GPU {gpu_id}] After warmup: {_gpu_mem()}") print(f"[GPU {gpu_id}] Model loaded and ready") - def do_step(keyboard_vector: list, mouse_vector: list) -> list[np.ndarray]: - """Execute a generation step.""" + def do_step(keyboard_vector: list, mouse_vector: list) -> tuple[list[np.ndarray], dict]: + """Execute a generation step with timing.""" + import time + timings = {} + + # Prepare action tensors + t0 = time.perf_counter() action = { "keyboard": torch.tensor(keyboard_vector).cuda(), "mouse": torch.tensor(mouse_vector).cuda() } keyboard_cond, mouse_cond = expand_action_to_frames(action, 12) - frames, _ = generator.step(keyboard_cond, mouse_cond) - return frames if frames is not None else [] + torch.cuda.synchronize() + timings["action_prep_ms"] = (time.perf_counter() - t0) * 1000 + + # Model step (DiT + VAE) + t0 = time.perf_counter() + frames, _, stage_timings = generator.step(keyboard_cond, mouse_cond) + torch.cuda.synchronize() + timings["model_step_ms"] = (time.perf_counter() - t0) * 1000 + + # Add stage timings if available + if stage_timings: + timings.update(stage_timings) + + return (frames if frames is not None else [], timings) def do_reset() -> list[np.ndarray]: """Reset the generator and return initial frames.""" @@ -130,7 +205,7 @@ def gpu_worker_process( # Set up sampling parameters actions = { - "keyboard": torch.zeros((NUM_FRAMES, MODEL_CONFIG["keyboard_dim"])), + "keyboard": torch.zeros((NUM_FRAMES, current_model_config["keyboard_dim"])), "mouse": torch.zeros((NUM_FRAMES, 2)) } @@ -140,7 +215,7 @@ def gpu_worker_process( generator.sampling_param.update({ "prompt": "", - "image_path": MODEL_CONFIG["image_url"], + "image_path": current_model_config.get("image_path", current_model_config["image_url"]), "mouse_cond": actions["mouse"].unsqueeze(0), "keyboard_cond": actions["keyboard"].unsqueeze(0), "grid_sizes": torch.tensor([150, 44, 80]), @@ -183,6 +258,249 @@ def gpu_worker_process( # Generate initial frame return do_step([0, 0, 0, 0], [0, 0]) + def prepare_user_batch() -> ForwardBatch: + """Create a fresh ForwardBatch for a new user joining.""" + actions = { + "keyboard": torch.zeros((NUM_FRAMES, current_model_config["keyboard_dim"])), + "mouse": torch.zeros((NUM_FRAMES, 2)) + } + + if generator.sampling_param is None: + generator.sampling_param = SamplingParam.from_pretrained( + generator.fastvideo_args.model_path) + + generator.sampling_param.update({ + "prompt": "", + "image_path": current_model_config.get("image_path", current_model_config["image_url"]), + "mouse_cond": actions["mouse"].unsqueeze(0), + "keyboard_cond": actions["keyboard"].unsqueeze(0), + "grid_sizes": torch.tensor([150, 44, 80]), + "num_frames": NUM_FRAMES, + "height": FRAME_HEIGHT, + "width": FRAME_WIDTH, + "num_inference_steps": NUM_INFERENCE_STEPS, + }) + + generator.sampling_param.height = align_to(generator.sampling_param.height, 16) + generator.sampling_param.width = align_to(generator.sampling_param.width, 16) + + latents_size = [ + (generator.sampling_param.num_frames - 1) // 4 + 1, + generator.sampling_param.height // 8, + generator.sampling_param.width // 8 + ] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + + generator.sampling_param.return_frames = True + generator.sampling_param.save_video = False + + batch = ForwardBatch( + **shallow_asdict(generator.sampling_param), + eta=0.0, + n_tokens=n_tokens, + VSA_sparsity=generator.fastvideo_args.VSA_sparsity, + ) + return batch + + def multi_user_event_loop(first_cmd: Command = None): + """Non-blocking event loop for multi-user mode. + + Drains commands from command_queue, submits them to the executor + without blocking, and collects results asynchronously. This enables + actual multi-user batching by allowing multiple USER_STEP commands + to be submitted before waiting for results. + """ + import queue as queue_module + + pending_results = 0 + + def reload_model(cmd: Command): + """Reload the model. Drains pending results first.""" + nonlocal pending_results + # Wait for all pending results to complete + while pending_results > 0: + result = generator.executor.get_result(timeout=5.0) + if result is not None: + pending_results -= 1 + # Route result to response queue + if result.error: + response_queue.put(Response( + success=False, error=str(result.error), + user_id=result.user_id)) + else: + response_queue.put(Response( + success=True, user_id=result.user_id)) + + model_config = cmd.data.get("model_config") + initialize_generator(model_config) + + # Re-enable streaming for the new executor + if not generator.executor._streaming_enabled: + generator.executor.enable_streaming() + + response_queue.put(Response(success=True, user_id=cmd.user_id)) + + def submit_command(cmd: Command): + """Submit a command to the executor. Returns True if a result is expected.""" + nonlocal pending_results + if cmd.type == CommandType.USER_JOIN: + batch = prepare_user_batch() + generator.executor.submit_user_join( + cmd.user_id, batch, generator.fastvideo_args) + pending_results += 1 + elif cmd.type == CommandType.USER_STEP: + action = { + "keyboard": torch.tensor( + cmd.data.get("keyboard", [0, 0, 0, 0])).cuda(), + "mouse": torch.tensor( + cmd.data.get("mouse", [0, 0])).cuda() + } + keyboard_cond, mouse_cond = expand_action_to_frames(action, 12) + torch.cuda.synchronize() + generator.executor.submit_user_step( + cmd.user_id, keyboard_cond, mouse_cond) + pending_results += 1 + elif cmd.type == CommandType.USER_LEAVE: + generator.executor.submit_user_leave(cmd.user_id) + pending_results += 1 + + def collect_results(): + """Collect available results from executor (non-blocking).""" + nonlocal pending_results + collected = False + while pending_results > 0: + result = generator.executor.get_result(timeout=0) + if result is None: + break + collected = True + pending_results -= 1 + + if result.error: + response_queue.put(Response( + success=False, error=str(result.error), + user_id=result.user_id)) + elif result.task_type == StreamingTaskType.USER_STEP: + frames = [] + timings = {} + if (result.output_batch is not None + and result.output_batch.output is not None): + frames = _process_output_to_frames(result.output_batch) + stage_timings = getattr( + result.output_batch, 'stage_timings', None) + if stage_timings: + timings.update(stage_timings) + response_queue.put(Response( + success=True, frames=frames, timings=timings, + user_id=result.user_id)) + elif result.task_type == StreamingTaskType.USER_JOIN: + print(f"[GPU {gpu_id}] User {result.user_id[:8]} joined") + response_queue.put(Response( + success=True, user_id=result.user_id)) + elif result.task_type == StreamingTaskType.USER_LEAVE: + print(f"[GPU {gpu_id}] User {result.user_id[:8]} left") + response_queue.put(Response( + success=True, user_id=result.user_id)) + return collected + + # Handle the first command that triggered the switch + if first_cmd is not None: + if first_cmd.type == CommandType.SHUTDOWN: + generator.shutdown() + response_queue.put(Response(success=True)) + return + try: + submit_command(first_cmd) + except Exception as e: + response_queue.put(Response( + success=False, error=str(e), + user_id=first_cmd.user_id)) + + print(f"[GPU {gpu_id}] Entered multi-user event loop") + + while True: + did_work = False + + # 1. Drain all available commands (non-blocking) + # If nothing pending, block for the first command + if pending_results == 0: + try: + cmd = command_queue.get(timeout=0.1) + if cmd.type == CommandType.SHUTDOWN: + generator.shutdown() + response_queue.put(Response(success=True)) + return + if cmd.type == CommandType.RELOAD_MODEL: + try: + reload_model(cmd) + except Exception as e: + response_queue.put(Response( + success=False, error=str(e), + user_id=cmd.user_id)) + did_work = True + continue + try: + submit_command(cmd) + did_work = True + except Exception as e: + response_queue.put(Response( + success=False, error=str(e), + user_id=cmd.user_id)) + except queue_module.Empty: + continue + + # Drain remaining commands (non-blocking) + while True: + try: + cmd = command_queue.get_nowait() + if cmd.type == CommandType.SHUTDOWN: + generator.shutdown() + response_queue.put(Response(success=True)) + return + if cmd.type == CommandType.RELOAD_MODEL: + try: + reload_model(cmd) + except Exception as e: + response_queue.put(Response( + success=False, error=str(e), + user_id=cmd.user_id)) + continue + try: + submit_command(cmd) + did_work = True + except Exception as e: + response_queue.put(Response( + success=False, error=str(e), + user_id=cmd.user_id)) + except queue_module.Empty: + break + + # 2. Collect results from executor (non-blocking) + if collect_results(): + did_work = True + + # 3. Brief sleep to avoid busy-waiting when results are pending + if not did_work and pending_results > 0: + time.sleep(0.001) + + def _process_output_to_frames(output_batch) -> list[np.ndarray]: + """Convert output batch tensor to list of numpy frames.""" + import torchvision + from einops import rearrange + + samples = output_batch.output + if samples is None: + return [] + if len(samples.shape) == 5: + videos = rearrange(samples, "b c t h w -> t b c h w") + else: + return [] + frames = [] + for x in videos: + x = torchvision.utils.make_grid(x, nrow=1) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + frames.append((x * 255).cpu().numpy().astype(np.uint8)) + return frames + # Main worker loop print(f"[GPU {gpu_id}] Worker process starting...") @@ -206,12 +524,23 @@ def gpu_worker_process( traceback.print_exc() response_queue.put(Response(success=False, error=str(e))) + elif cmd.type == CommandType.RELOAD_MODEL: + try: + model_config = cmd.data.get("model_config") + initialize_generator(model_config) + response_queue.put(Response(success=True, user_id=cmd.user_id)) + except Exception as e: + print(f"[GPU {gpu_id}] Reload error: {e}") + traceback.print_exc() + response_queue.put(Response( + success=False, error=str(e), user_id=cmd.user_id)) + elif cmd.type == CommandType.STEP: try: keyboard = cmd.data.get("keyboard", [0, 0, 0, 0]) mouse = cmd.data.get("mouse", [0, 0]) - frames = do_step(keyboard, mouse) - response_queue.put(Response(success=True, frames=frames)) + frames, timings = do_step(keyboard, mouse) + response_queue.put(Response(success=True, frames=frames, timings=timings)) except Exception as e: print(f"[GPU {gpu_id}] Step error: {e}") traceback.print_exc() @@ -226,6 +555,16 @@ def gpu_worker_process( traceback.print_exc() response_queue.put(Response(success=False, error=str(e))) + elif cmd.type in (CommandType.USER_JOIN, + CommandType.USER_STEP, + CommandType.USER_LEAVE): + # Enable streaming if needed (synchronous RPC) + if not generator.executor._streaming_enabled: + generator.executor.enable_streaming() + # Switch to non-blocking event loop for multi-user mode + multi_user_event_loop(first_cmd=cmd) + break # event loop handles everything until shutdown + except Exception as e: print(f"[GPU {gpu_id}] Worker crashed: {e}") traceback.print_exc() @@ -234,21 +573,43 @@ def gpu_worker_process( class GPUSlot: - """Manages a single GPU worker subprocess.""" + """Manages a single GPU worker subprocess with multi-client support.""" def __init__(self, gpu_id: int, cuda_device: str): self.gpu_id = gpu_id self.cuda_device = cuda_device - self.client_id: Optional[str] = None self.process: Optional[Process] = None self.command_queue: Optional[Queue] = None self.response_queue: Optional[Queue] = None self.ready: bool = False self._lock = asyncio.Lock() + # Multi-client state + self.connected_users: set[str] = set() # all connected user IDs + self._pending_futures: dict[str, asyncio.Future] = {} # user_id -> pending Future + self._response_reader_task: Optional[asyncio.Task] = None + self._multi_user_mode: bool = False + self._reader_lock: Optional[asyncio.Lock] = None + self.current_model_id: str = DEFAULT_MODEL_ID + + @property + def client_count(self) -> int: + return len(self.connected_users) + @property def is_available(self) -> bool: - return self.client_id is None and self.ready and self.process is not None and self.process.is_alive() + """A GPU is available if it has capacity for more users.""" + alive = self.ready and self.process is not None and self.process.is_alive() + if not alive: + return False + if self._multi_user_mode: + return len(self.connected_users) < MAX_USERS_PER_GPU + else: + return len(self.connected_users) == 0 + + @property + def is_empty(self) -> bool: + return len(self.connected_users) == 0 async def start(self): """Start the GPU worker subprocess.""" @@ -259,55 +620,225 @@ class GPUSlot: self.process = ctx.Process( target=gpu_worker_process, args=(self.gpu_id, self.cuda_device, self.command_queue, self.response_queue), - daemon=False, # Must be False so executor can spawn child workers + daemon=False, ) loop = asyncio.get_event_loop() await loop.run_in_executor(None, self.process.start) # Send init command and wait for response - response = await self._send_command(Command(CommandType.INIT), timeout=600.0) + response = await self._send_command( + Command(CommandType.INIT), timeout=600.0) if not response.success: - raise RuntimeError(f"GPU {self.gpu_id} failed to initialize: {response.error}") + raise RuntimeError( + f"GPU {self.gpu_id} failed to initialize: {response.error}") self.ready = True - async def _send_command(self, cmd: Command, timeout: float = 300.0) -> Response: - """Send a command to the worker and wait for response.""" + async def _send_command( + self, cmd: Command, timeout: float = 300.0 + ) -> Response: + """Send a command and wait for response (single-user mode).""" loop = asyncio.get_event_loop() - - # Send command in thread pool to not block await loop.run_in_executor(None, self.command_queue.put, cmd) - # Wait for response with timeout def get_response(): return self.response_queue.get(timeout=timeout) response = await loop.run_in_executor(None, get_response) return response - async def step(self, keyboard: list, mouse: list) -> list[np.ndarray]: - """Execute a generation step.""" + async def _send_command_tagged( + self, cmd: Command, timeout: float = 300.0 + ) -> Response: + """Send a tagged command and wait for the matching response (multi-user mode). + + Uses the response reader background task to route responses. + """ + loop = asyncio.get_event_loop() + future = loop.create_future() + + # Register this request's pending future + self._pending_futures[cmd.user_id] = future + + # Send command + await loop.run_in_executor(None, self.command_queue.put, cmd) + + # Ensure response reader is running + await self._ensure_response_reader() + + # Wait for the response + try: + response = await asyncio.wait_for(future, timeout=timeout) + except asyncio.TimeoutError: + # Clean up only our own future + if (cmd.user_id in self._pending_futures + and self._pending_futures[cmd.user_id] is future): + self._pending_futures.pop(cmd.user_id, None) + raise + + return response + + async def _ensure_response_reader(self): + """Start the response reader background task if not running.""" + if self._reader_lock is None: + self._reader_lock = asyncio.Lock() + + async with self._reader_lock: + if (self._response_reader_task is None + or self._response_reader_task.done()): + self._response_reader_task = asyncio.create_task( + self._response_reader()) + + async def _response_reader(self): + """Background task that reads response queue and routes to per-user futures.""" + loop = asyncio.get_event_loop() + + while self._multi_user_mode: + try: + def get_response_nonblocking(): + try: + return self.response_queue.get(timeout=0.5) + except Exception: + return None + + response = await loop.run_in_executor( + None, get_response_nonblocking) + + if response is None: + continue + + user_id = response.user_id + if user_id and user_id in self._pending_futures: + future = self._pending_futures.pop(user_id) + if not future.done(): + future.set_result(response) + else: + # No matching pending request - log and discard + print(f"[GPU {self.gpu_id}] Unmatched response for user " + f"{user_id[:8] if user_id else 'None'}") + except asyncio.CancelledError: + break + except Exception as e: + print(f"[GPU {self.gpu_id}] Response reader error: {e}") + await asyncio.sleep(0.01) + + # --- Single-user API (legacy) --- + + async def step(self, keyboard: list, mouse: list) -> tuple[list[np.ndarray], dict]: + """Execute a generation step (single-user mode).""" async with self._lock: response = await self._send_command( Command(CommandType.STEP, {"keyboard": keyboard, "mouse": mouse}) ) if not response.success: raise RuntimeError(f"Step failed: {response.error}") - return response.frames or [] + return response.frames or [], response.timings or {} async def reset(self) -> list[np.ndarray]: - """Reset the generator.""" + """Reset the generator (single-user mode).""" async with self._lock: response = await self._send_command(Command(CommandType.RESET)) if not response.success: raise RuntimeError(f"Reset failed: {response.error}") return response.frames or [] + # --- Multi-user API --- + + async def join_user(self, user_id: str, model_id: str = None) -> Response: + """Add a user to this GPU's multi-user engine. + + If model_id differs from the currently loaded model, triggers a full + model reload. This invalidates all existing users on this GPU — their + pending futures are cancelled with an error. + """ + if model_id is None: + model_id = DEFAULT_MODEL_ID + + # Reload model if a different one is requested + if model_id != self.current_model_id and model_id in MODEL_REGISTRY: + print(f"[GPU {self.gpu_id}] Model switch: " + f"{self.current_model_id} -> {model_id}") + + # Cancel pending futures for existing users (their steps will error) + for uid, future in list(self._pending_futures.items()): + if not future.done(): + future.set_exception( + RuntimeError("Model changed, session reset")) + self._pending_futures.clear() + self.connected_users.clear() + + # Send reload command to GPU worker. + # Use _send_command (not _send_command_tagged) because + # _multi_user_mode is still False, so the response reader + # would exit immediately and never route the response. + model_config = MODEL_REGISTRY[model_id] + response = await self._send_command( + Command(CommandType.RELOAD_MODEL, + data={"model_config": model_config}, + user_id="__reload__"), + timeout=600.0) + if not response.success: + raise RuntimeError( + f"Model reload failed: {response.error}") + + self.current_model_id = model_id + print(f"[GPU {self.gpu_id}] Model reloaded: {model_id}") + + self._multi_user_mode = True + self.connected_users.add(user_id) + try: + response = await self._send_command_tagged( + Command(CommandType.USER_JOIN, user_id=user_id), + timeout=600.0) + if not response.success: + self.connected_users.discard(user_id) + raise RuntimeError( + f"User join failed for {user_id[:8]}: {response.error}") + return response + except Exception: + self.connected_users.discard(user_id) + raise + + async def user_step( + self, user_id: str, keyboard: list, mouse: list + ) -> tuple[list[np.ndarray], dict]: + """Execute a generation step for a specific user.""" + response = await self._send_command_tagged( + Command(CommandType.USER_STEP, + {"keyboard": keyboard, "mouse": mouse}, + user_id=user_id), + timeout=300.0) + if not response.success: + raise RuntimeError( + f"User step failed for {user_id[:8]}: {response.error}") + return response.frames or [], response.timings or {} + + async def leave_user(self, user_id: str) -> None: + """Remove a user from this GPU's multi-user engine.""" + try: + response = await self._send_command_tagged( + Command(CommandType.USER_LEAVE, user_id=user_id), + timeout=30.0) + except Exception as e: + print(f"[GPU {self.gpu_id}] Leave user error: {e}") + finally: + self.connected_users.discard(user_id) + self._pending_futures.pop(user_id, None) + async def shutdown(self): """Shutdown the worker subprocess.""" + self._multi_user_mode = False # Stop the response reader loop + if self._response_reader_task and not self._response_reader_task.done(): + self._response_reader_task.cancel() + try: + await self._response_reader_task + except asyncio.CancelledError: + pass + if self.process and self.process.is_alive(): try: - await self._send_command(Command(CommandType.SHUTDOWN), timeout=30.0) + await self._send_command( + Command(CommandType.SHUTDOWN), timeout=30.0) except Exception: pass self.process.terminate() @@ -317,25 +848,20 @@ class GPUSlot: class GPUPool: - """Manages multiple GPU worker subprocesses.""" + """Manages multiple GPU worker subprocesses with multi-user support.""" def __init__(self, gpu_ids: list[int]): self.gpu_ids = gpu_ids self.slots: dict[int, GPUSlot] = { - gpu_id: GPUSlot(gpu_id, str(i)) for i, gpu_id in enumerate(gpu_ids) + gpu_id: GPUSlot(gpu_id, str(gpu_id)) for gpu_id in gpu_ids } self.waiting_list: list[tuple[str, asyncio.Event, "WebSocket"]] = [] self.client_gpu_map: dict[str, int] = {} self._pool_lock = asyncio.Lock() async def initialize(self): - """Start initializing all GPU workers in the background. - - Returns immediately so the server can start accepting connections. - GPUs become available as each one finishes loading. - """ + """Start initializing all GPU workers in the background.""" print(f"Initializing GPU pool with {len(self.gpu_ids)} GPUs: {self.gpu_ids}") - for gpu_id in self.gpu_ids: asyncio.create_task(self._init_gpu(gpu_id)) @@ -347,36 +873,45 @@ class GPUPool: print(f"GPU {gpu_id} failed to initialize: {e}") return - print(f"GPU pool: {gpu_id} ready ({sum(1 for s in self.slots.values() if s.ready)}/{len(self.gpu_ids)})") + print(f"GPU pool: {gpu_id} ready " + f"({sum(1 for s in self.slots.values() if s.ready)}/{len(self.gpu_ids)})") # Check if anyone is waiting for a GPU async with self._pool_lock: slot = self.slots[gpu_id] if slot.is_available and self.waiting_list: waiting_client_id, ready_event, _ = self.waiting_list.pop(0) - slot.client_id = waiting_client_id self.client_gpu_map[waiting_client_id] = gpu_id print(f"Client {waiting_client_id[:8]} assigned GPU {gpu_id} from queue") ready_event.set() - # Notify remaining clients of updated positions await self._send_queue_updates() async def acquire(self, client_id: str, websocket=None) -> tuple[int, GPUSlot]: - """ - Acquire a GPU slot for a client. - Returns (gpu_id, slot) when available. + """Acquire a GPU slot for a client. + + In multi-user mode, returns a GPU that has capacity for more users. """ async with self._pool_lock: + # First, try to find a GPU with existing users (better batching) + for gpu_id, slot in self.slots.items(): + if (slot.is_available and slot._multi_user_mode + and slot.client_count > 0): + self.client_gpu_map[client_id] = gpu_id + print(f"Client {client_id[:8]} acquired GPU {gpu_id} " + f"(co-located, {slot.client_count} users)") + return gpu_id, slot + + # Then, try any available GPU for gpu_id, slot in self.slots.items(): if slot.is_available: - slot.client_id = client_id self.client_gpu_map[client_id] = gpu_id print(f"Client {client_id[:8]} acquired GPU {gpu_id}") return gpu_id, slot # No slot available, wait in queue - print(f"Client {client_id[:8]} waiting in queue (all {len(self.gpu_ids)} GPUs busy)") + print(f"Client {client_id[:8]} waiting in queue " + f"(all {len(self.gpu_ids)} GPUs at capacity)") ready_event = asyncio.Event() async with self._pool_lock: self.waiting_list.append((client_id, ready_event, websocket)) @@ -385,12 +920,13 @@ class GPUPool: gpu_id = self.client_gpu_map.get(client_id) if gpu_id is None: - raise RuntimeError(f"Client {client_id} was signaled but has no GPU assigned") + raise RuntimeError( + f"Client {client_id} was signaled but has no GPU assigned") return gpu_id, self.slots[gpu_id] async def release(self, client_id: str): - """Release a GPU slot when a client disconnects.""" + """Release a client from its GPU slot.""" async with self._pool_lock: gpu_id = self.client_gpu_map.pop(client_id, None) if gpu_id is None: @@ -399,28 +935,31 @@ class GPUPool: slot = self.slots[gpu_id] print(f"Client {client_id[:8]} released GPU {gpu_id}") - # Reset the generator so it's fresh for the next client - try: - await slot.reset() - print(f"[GPU {gpu_id}] Reset complete, ready for next client") - except Exception as e: - print(f"[GPU {gpu_id}] Reset on release failed: {e}") + if slot._multi_user_mode: + # Multi-user: just remove this user, GPU stays active + try: + await slot.leave_user(client_id) + except Exception as e: + print(f"[GPU {gpu_id}] Leave user failed: {e}") + else: + # Single-user: reset the generator + try: + await slot.reset() + print(f"[GPU {gpu_id}] Reset complete, ready for next client") + except Exception as e: + print(f"[GPU {gpu_id}] Reset on release failed: {e}") - slot.client_id = None - - # Assign to next waiting client - if self.waiting_list: + # Assign to next waiting client if GPU has capacity + if slot.is_available and self.waiting_list: waiting_client_id, ready_event, _ = self.waiting_list.pop(0) - slot.client_id = waiting_client_id self.client_gpu_map[waiting_client_id] = gpu_id print(f"Client {waiting_client_id[:8]} assigned GPU {gpu_id} from queue") ready_event.set() - # Notify remaining clients of updated positions await self._send_queue_updates() async def _send_queue_updates(self): - """Send updated queue positions to all waiting clients. Must be called with _pool_lock held.""" + """Send updated queue positions to all waiting clients.""" for i, (cid, _, ws) in enumerate(self.waiting_list): if ws is not None: try: @@ -431,7 +970,7 @@ class GPUPool: "available_gpus": 0, }) except Exception: - pass # Client may have disconnected + pass async def shutdown(self): """Shutdown all GPU workers.""" @@ -444,13 +983,17 @@ class GPUPool: """Get the current status of the GPU pool.""" return { "total_gpus": len(self.gpu_ids), - "available_gpus": sum(1 for slot in self.slots.values() if slot.is_available), + "available_gpus": sum( + 1 for slot in self.slots.values() if slot.is_available), "queue_size": len(self.waiting_list), "gpu_status": { gpu_id: { "available": slot.is_available, - "client_id": slot.client_id[:8] if slot.client_id else None, - "process_alive": slot.process.is_alive() if slot.process else False, + "client_count": slot.client_count, + "multi_user": slot._multi_user_mode, + "current_model_id": slot.current_model_id, + "process_alive": ( + slot.process.is_alive() if slot.process else False), } for gpu_id, slot in self.slots.items() } diff --git a/ui/world_model/server/main.py b/ui/world_model/server/main.py index 5fb97017a..3bc3d77ae 100644 --- a/ui/world_model/server/main.py +++ b/ui/world_model/server/main.py @@ -11,7 +11,7 @@ import uuid from config import ( KEYBOARD_MAP, CAMERA_MAP, MAX_BLOCKS, JPEG_QUALITY, BATCH_SIZE, - SESSION_TIMEOUT_SECONDS, MODEL_CONFIG + SESSION_TIMEOUT_SECONDS, MODEL_CONFIG, MODEL_REGISTRY, DEFAULT_MODEL_ID ) from gpu_pool import GPUPool, GPUSlot, get_available_gpus @@ -60,6 +60,21 @@ async def get_status(): return gpu_pool.get_status() +@app.get("/models") +async def get_models(): + """Get available models and the currently active model.""" + models = [] + for model_id, config in MODEL_REGISTRY.items(): + models.append({ + "id": model_id, + "name": config["name"], + }) + return { + "models": models, + "default_model_id": DEFAULT_MODEL_ID, + } + + def encode_frames(frames: list) -> list[str]: """Encode frames to base64 JPEG.""" encoded = [] @@ -114,31 +129,50 @@ async def websocket_endpoint(websocket: WebSocket): pass # WebSocket may already be closed try: - # Acquire a GPU slot (may wait in queue) + # Wait for model selection from client + model_id = DEFAULT_MODEL_ID + model_config = MODEL_CONFIG + try: + init_data = await asyncio.wait_for(websocket.receive_json(), timeout=5.0) + if init_data.get("type") == "select_model": + model_id = init_data.get("model_id", DEFAULT_MODEL_ID) + if model_id in MODEL_REGISTRY: + model_config = MODEL_REGISTRY[model_id] + print(f"Client {client_id[:8]} selected model: {model_id}") + except asyncio.TimeoutError: + pass # Use default model config + + # Acquire a GPU slot (may wait in queue, may share with other users) gpu_id, slot = await gpu_pool.acquire(client_id, websocket) # Start session timeout timeout_task = asyncio.create_task(session_timeout()) + # Join the multi-user engine on this GPU (triggers reload if model differs) + await slot.join_user(client_id, model_id=model_id) + # Notify client they're connected to a GPU await websocket.send_json({ "type": "gpu_assigned", "gpu_id": gpu_id, "session_timeout": SESSION_TIMEOUT_SECONDS, - "image_url": MODEL_CONFIG.get("image_url") + "image_url": model_config.get("image_url") }) - # Generate and send initial frame (GPU is already reset and ready) + # Generate and send initial frame block_count = 0 await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS}) try: - frames = await slot.step([0, 0, 0, 0], [0, 0]) + frames, timings = await slot.user_step( + client_id, [0, 0, 0, 0], [0, 0]) block_count = 1 await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS}) await send_frames(websocket, frames) + if timings: + print(f"[GPU {gpu_id}] Initial ({client_id[:8]}): model={timings.get('model_step_ms', 0):.0f}ms") except Exception as e: - print(f"[GPU {gpu_id}] Initial frame error: {e}") + print(f"[GPU {gpu_id}] Initial frame error for {client_id[:8]}: {e}") # Main event loop while True: @@ -150,14 +184,19 @@ async def websocket_endpoint(websocket: WebSocket): await websocket.send_json({"type": "reset_started"}) try: - frames = await slot.reset() + # Leave and rejoin to reset user state + await slot.leave_user(client_id) + await slot.join_user(client_id, model_id=model_id) + + frames, timings = await slot.user_step( + client_id, [0, 0, 0, 0], [0, 0]) block_count = 1 await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS}) await send_frames(websocket, frames) await websocket.send_json({"type": "reset_complete"}) - print(f"[GPU {gpu_id}] Reset complete") + print(f"[GPU {gpu_id}] Reset complete for {client_id[:8]}") except Exception as e: - print(f"[GPU {gpu_id}] Reset error: {e}") + print(f"[GPU {gpu_id}] Reset error for {client_id[:8]}: {e}") await websocket.send_json({"type": "reset_complete"}) continue @@ -171,10 +210,11 @@ async def websocket_endpoint(websocket: WebSocket): keyboard_vector = KEYBOARD_MAP.get(key, [0, 0, 0, 0]) mouse_vector = [0, 0] - print(f"[GPU {gpu_id}] Key '{key}' pressed, generating block {block_count + 1}...") + print(f"[GPU {gpu_id}] Key '{key}' from {client_id[:8]}, generating block {block_count + 1}...") t_start = time.time() - frames = await slot.step(keyboard_vector, mouse_vector) + frames, timings = await slot.user_step( + client_id, keyboard_vector, mouse_vector) t_generation = time.time() - t_start block_count += 1 @@ -183,9 +223,11 @@ async def websocket_endpoint(websocket: WebSocket): if frames: t_encode_start = time.time() await send_frames(websocket, frames) - t_total = time.time() - t_encode_start + t_encode = time.time() - t_encode_start - print(f"[GPU {gpu_id}] Gen: {t_generation:.3f}s, Encode+Send: {t_total:.3f}s") + dit_ms = timings.get('dit_ms', 0) + vae_ms = timings.get('vae_ms', 0) + print(f"[GPU {gpu_id}] ({client_id[:8]}) DiT: {dit_ms:.0f}ms, VAE: {vae_ms:.0f}ms, Encode+Send: {t_encode*1000:.0f}ms, Total: {t_generation*1000:.0f}ms") except WebSocketDisconnect: print(f"Client {client_id[:8]} disconnected") @@ -205,12 +247,16 @@ async def websocket_endpoint(websocket: WebSocket): await gpu_pool.release(client_id) +# Serve server-side assets (images, etc.) +server_dir = os.path.dirname(os.path.abspath(__file__)) +app.mount("/server-assets", StaticFiles(directory=server_dir), name="server-assets") + # Serve built frontend (must be after API/WebSocket routes) -static_dir = os.path.join(os.path.dirname(__file__), "..", "client", "dist") +static_dir = os.path.join(server_dir, "..", "client", "dist") if os.path.isdir(static_dir): app.mount("/", StaticFiles(directory=static_dir, html=True), name="static") if __name__ == "__main__": import uvicorn - uvicorn.run(app, host="0.0.0.0", port=8000) + uvicorn.run(app, host="0.0.0.0", port=8001) diff --git a/ui/world_model/server/mock_server.py b/ui/world_model/server/mock_server.py new file mode 100644 index 000000000..d9b272692 --- /dev/null +++ b/ui/world_model/server/mock_server.py @@ -0,0 +1,332 @@ +""" +Mock server to test latency perception. +Generates frames with a moving object based on keyboard input. +Configurable latency to find the acceptable threshold. +""" + +import asyncio +import base64 +import io +import time +import uuid +from contextlib import asynccontextmanager + +import numpy as np +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from fastapi.middleware.cors import CORSMiddleware +from fastapi.staticfiles import StaticFiles +from PIL import Image, ImageDraw +import os + +# Configuration +LATENCY_MS = 200 # Simulated model latency - adjust this to test +FRAME_WIDTH = 640 +FRAME_HEIGHT = 352 +NUM_FRAMES = 12 # Frames per block +FPS = 24 +BATCH_SIZE = 4 +JPEG_QUALITY = 85 +SESSION_TIMEOUT_SECONDS = 90 +MAX_BLOCKS = 50 + +# Mock game state +class GameState: + def __init__(self): + self.reset() + + def reset(self): + self.x = FRAME_WIDTH // 2 + self.y = FRAME_HEIGHT // 2 + self.size = 40 + self.color = (100, 200, 255) # Light blue + self.trail = [] # Previous positions for motion blur effect + + def update(self, keyboard_vector, mouse_vector): + """Store the movement direction for this block.""" + self.keyboard_vector = keyboard_vector + self.mouse_vector = mouse_vector + + # Arrow keys for camera (just change color as visual feedback) + if mouse_vector[0] != 0 or mouse_vector[1] != 0: + r = int(100 + mouse_vector[0] * 50) % 256 + g = int(200 + mouse_vector[1] * 50) % 256 + self.color = (r, g, 255) + + def render_frame(self, frame_idx: int) -> np.ndarray: + """Render a single frame with incremental movement.""" + speed_per_frame = 8 + + # Apply movement for this frame + if hasattr(self, 'keyboard_vector'): + if self.keyboard_vector[0]: # W - up + self.y -= speed_per_frame + if self.keyboard_vector[1]: # A - left + self.x -= speed_per_frame + if self.keyboard_vector[2]: # S - down + self.y += speed_per_frame + if self.keyboard_vector[3]: # D - right + self.x += speed_per_frame + + # Wrap around screen + self.x = self.x % FRAME_WIDTH + self.y = self.y % FRAME_HEIGHT + + # Update trail + self.trail.append((self.x, self.y)) + if len(self.trail) > 5: + self.trail.pop(0) + """Render a single frame.""" + # Create frame with gradient background + img = Image.new('RGB', (FRAME_WIDTH, FRAME_HEIGHT), (20, 20, 30)) + draw = ImageDraw.Draw(img) + + # Draw grid for visual reference + for x in range(0, FRAME_WIDTH, 40): + draw.line([(x, 0), (x, FRAME_HEIGHT)], fill=(40, 40, 50), width=1) + for y in range(0, FRAME_HEIGHT, 40): + draw.line([(0, y), (FRAME_WIDTH, y)], fill=(40, 40, 50), width=1) + + # Interpolate position for smooth animation within block + t = frame_idx / NUM_FRAMES + + # Draw trail (motion blur effect) + for i, (tx, ty) in enumerate(self.trail): + alpha = (i + 1) / len(self.trail) * 0.3 + trail_size = int(self.size * (0.5 + alpha * 0.5)) + trail_color = tuple(int(c * alpha) for c in self.color) + draw.ellipse( + [tx - trail_size//2, ty - trail_size//2, + tx + trail_size//2, ty + trail_size//2], + fill=trail_color + ) + + # Draw main object (circle with glow) + # Outer glow + for r in range(3, 0, -1): + glow_size = self.size + r * 8 + glow_alpha = 0.2 / r + glow_color = tuple(int(c * glow_alpha) for c in self.color) + draw.ellipse( + [self.x - glow_size//2, self.y - glow_size//2, + self.x + glow_size//2, self.y + glow_size//2], + fill=glow_color + ) + + # Main circle + draw.ellipse( + [self.x - self.size//2, self.y - self.size//2, + self.x + self.size//2, self.y + self.size//2], + fill=self.color + ) + + # Inner highlight + highlight_size = self.size // 3 + highlight_offset = self.size // 6 + draw.ellipse( + [self.x - highlight_offset - highlight_size//2, + self.y - highlight_offset - highlight_size//2, + self.x - highlight_offset + highlight_size//2, + self.y - highlight_offset + highlight_size//2], + fill=(255, 255, 255) + ) + + # Add latency indicator text + draw.text((10, 10), f"Latency: {LATENCY_MS}ms", fill=(150, 150, 150)) + draw.text((10, 30), f"Frame: {frame_idx + 1}/{NUM_FRAMES}", fill=(100, 100, 100)) + + return np.array(img) + + +def encode_frame(frame: np.ndarray) -> str: + """Encode frame to base64 JPEG.""" + img = Image.fromarray(frame) + buffer = io.BytesIO() + img.save(buffer, format='JPEG', quality=JPEG_QUALITY) + return base64.b64encode(buffer.getvalue()).decode('utf-8') + + +def generate_frames(game_state: GameState) -> list[str]: + """Generate and encode a block of frames.""" + encoded = [] + for i in range(NUM_FRAMES): + frame = game_state.render_frame(i) + encoded.append(encode_frame(frame)) + return encoded + + +# Keyboard/camera mappings (same as real server) +KEYBOARD_MAP = { + 'w': [1, 0, 0, 0], + 'a': [0, 1, 0, 0], + 's': [0, 0, 1, 0], + 'd': [0, 0, 0, 1], +} + +CAMERA_MAP = { + 'ArrowUp': [0, -1], + 'ArrowDown': [0, 1], + 'ArrowLeft': [-1, 0], + 'ArrowRight': [1, 0], +} + + +@asynccontextmanager +async def lifespan(app: FastAPI): + print(f"Mock server starting with {LATENCY_MS}ms simulated latency...") + yield + print("Mock server shutting down...") + + +app = FastAPI(lifespan=lifespan) + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + + +async def send_frames(websocket: WebSocket, frames: list[str]): + """Send frames in batches.""" + for i in range(0, len(frames), BATCH_SIZE): + batch = frames[i:i + BATCH_SIZE] + await websocket.send_json({"type": "frame_batch", "frames": batch}) + + +@app.websocket("/ws") +async def websocket_endpoint(websocket: WebSocket): + await websocket.accept() + + client_id = str(uuid.uuid4()) + print(f"Client {client_id[:8]} connected") + + game_state = GameState() + block_count = 0 + timeout_task = None + + async def session_timeout(): + await asyncio.sleep(SESSION_TIMEOUT_SECONDS) + try: + await websocket.send_json({ + "type": "session_timeout", + "message": f"Session expired after {SESSION_TIMEOUT_SECONDS} seconds" + }) + await websocket.close() + except Exception: + pass + + try: + # Send initial status + await websocket.send_json({ + "type": "queue_status", + "position": 0, + "total_gpus": 1, + "available_gpus": 1 + }) + + # Simulate GPU assignment delay + await asyncio.sleep(0.1) + + # Start timeout + timeout_task = asyncio.create_task(session_timeout()) + + # Send GPU assigned + await websocket.send_json({ + "type": "gpu_assigned", + "gpu_id": 0, + "session_timeout": SESSION_TIMEOUT_SECONDS, + "image_url": None # No initial image for mock + }) + + # Generate initial frame + await websocket.send_json({"type": "block_count", "count": 0, "max": MAX_BLOCKS}) + + # Simulate latency + await asyncio.sleep(LATENCY_MS / 1000) + + frames = generate_frames(game_state) + block_count = 1 + await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS}) + await send_frames(websocket, frames) + + print(f"[Mock] Initial frames sent") + + # Main loop + while True: + data = await websocket.receive_json() + message_type = data.get("type", "key") + + if message_type == "reset": + print(f"[Mock] Reset requested") + await websocket.send_json({"type": "reset_started"}) + + game_state.reset() + await asyncio.sleep(LATENCY_MS / 1000) + + frames = generate_frames(game_state) + block_count = 1 + await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS}) + await send_frames(websocket, frames) + await websocket.send_json({"type": "reset_complete"}) + continue + + key = data.get("key") + if key and block_count < MAX_BLOCKS: + if key in CAMERA_MAP: + keyboard_vector = [0, 0, 0, 0] + mouse_vector = CAMERA_MAP[key] + else: + keyboard_vector = KEYBOARD_MAP.get(key, [0, 0, 0, 0]) + mouse_vector = [0, 0] + + print(f"[Mock] Key '{key}' pressed, generating block {block_count + 1}...") + + t_start = time.time() + + # Simulate model latency + await asyncio.sleep(LATENCY_MS / 1000) + + # Update game state and generate frames + game_state.update(keyboard_vector, mouse_vector) + frames = generate_frames(game_state) + + t_total = (time.time() - t_start) * 1000 + + block_count += 1 + await websocket.send_json({"type": "block_count", "count": block_count, "max": MAX_BLOCKS}) + await send_frames(websocket, frames) + + print(f"[Mock] Block generated in {t_total:.0f}ms (target: {LATENCY_MS}ms)") + + except WebSocketDisconnect: + print(f"Client {client_id[:8]} disconnected") + except Exception as e: + print(f"Client {client_id[:8]} error: {e}") + import traceback + traceback.print_exc() + finally: + if timeout_task and not timeout_task.done(): + timeout_task.cancel() + + +# Serve static files +static_dir = os.path.join(os.path.dirname(__file__), "..", "client", "dist") +if os.path.isdir(static_dir): + app.mount("/", StaticFiles(directory=static_dir, html=True), name="static") + + +if __name__ == "__main__": + import argparse + import uvicorn + + parser = argparse.ArgumentParser(description="Mock server for latency testing") + parser.add_argument("--latency", type=int, default=200, help="Simulated latency in ms") + parser.add_argument("--port", type=int, default=8001, help="Server port") + args = parser.parse_args() + + LATENCY_MS = args.latency + print(f"Starting mock server with {LATENCY_MS}ms latency on port {args.port}") + + uvicorn.run(app, host="0.0.0.0", port=args.port) diff --git a/ui/world_model/server/test_batching.py b/ui/world_model/server/test_batching.py new file mode 100644 index 000000000..ffadbbfb3 --- /dev/null +++ b/ui/world_model/server/test_batching.py @@ -0,0 +1,149 @@ +"""Test ORCA batching: synchronized vs random keypresses.""" +import argparse +import asyncio +import json +import random +import time +import websockets + +SERVER_URL = "ws://localhost:8000/ws" +KEYS = ["w", "a", "s", "d"] +# Think time range for random mode (ms) +MIN_THINK_MS = 200 +MAX_THINK_MS = 1500 + + +async def connect_and_wait(client_id, ws): + """Drain messages until initial frame_batch is received.""" + while True: + msg = json.loads(await ws.recv()) + if msg["type"] == "frame_batch": + print(f" Client {client_id}: got initial frame") + return + + +async def synced_client(client_id: int, barrier: asyncio.Barrier, results: list, + num_rounds: int): + """Synchronized keypresses — all clients press at the same time.""" + async with websockets.connect( + SERVER_URL, max_size=50_000_000, ping_interval=None, + ) as ws: + await connect_and_wait(client_id, ws) + + for round_idx in range(num_rounds): + await barrier.wait() + + t0 = time.perf_counter() + await ws.send(json.dumps({"type": "key", "key": "w"})) + + while True: + msg = json.loads(await ws.recv()) + if msg["type"] == "frame_batch": + break + + elapsed_ms = (time.perf_counter() - t0) * 1000 + results.append(elapsed_ms) + print(f" Client {client_id} round {round_idx}: {elapsed_ms:.0f}ms") + + +async def random_client(client_id: int, start_event: asyncio.Event, results: list, + num_rounds: int): + """Random keypresses with think time between presses.""" + async with websockets.connect( + SERVER_URL, max_size=50_000_000, ping_interval=None, + ) as ws: + await connect_and_wait(client_id, ws) + start_event.set() + + # Small stagger so clients don't all start at once + await asyncio.sleep(random.uniform(0, 1.0)) + + for _ in range(num_rounds): + think_ms = random.uniform(MIN_THINK_MS, MAX_THINK_MS) + await asyncio.sleep(think_ms / 1000) + + t0 = time.perf_counter() + await ws.send(json.dumps({"type": "key", "key": random.choice(KEYS)})) + + while True: + msg = json.loads(await ws.recv()) + if msg["type"] == "frame_batch": + break + + elapsed_ms = (time.perf_counter() - t0) * 1000 + results.append(elapsed_ms) + + print(f" Client {client_id}: done") + + +def print_stats(times: list[float], label: str): + """Print summary statistics and histogram.""" + times.sort() + n = len(times) + print(f"\n--- {label} ({n} keypresses) ---") + print(f" Min: {times[0]:.0f}ms") + print(f" p25: {times[n//4]:.0f}ms") + print(f" Median: {times[n//2]:.0f}ms") + print(f" p75: {times[3*n//4]:.0f}ms") + print(f" p95: {times[int(n*0.95)]:.0f}ms") + print(f" Max: {times[-1]:.0f}ms") + print(f" Avg: {sum(times)/n:.0f}ms") + + buckets = [0] * 10 + labels = ["<1s", "<1.5s", "<2s", "<2.5s", "<3s", + "<3.5s", "<4s", "<4.5s", "<5s", "5s+"] + for ms in times: + idx = min(int(ms / 500), 9) + buckets[idx] += 1 + print(f"\n Distribution:") + for lbl, count in zip(labels, buckets): + if count > 0: + bar = "#" * count + pct = count / n * 100 + print(f" {lbl:>6s}: {bar} ({count}, {pct:.0f}%)") + + +async def run_synced(num_clients: int, num_rounds: int): + print(f"=== SYNCHRONIZED MODE ({num_clients} clients, {num_rounds} rounds) ===") + barrier = asyncio.Barrier(num_clients) + results = [] + tasks = [synced_client(i, barrier, results, num_rounds) for i in range(num_clients)] + await asyncio.gather(*tasks) + print_stats(results, "Synchronized (worst case)") + + +async def run_random(num_clients: int, num_rounds: int): + print(f"=== RANDOM MODE ({num_clients} clients, {num_rounds} presses each) ===") + print(f" Think time: {MIN_THINK_MS}-{MAX_THINK_MS}ms") + results = [] + events = [asyncio.Event() for _ in range(num_clients)] + tasks = [random_client(i, events[i], results, num_rounds) for i in range(num_clients)] + await asyncio.gather(*tasks) + print_stats(results, "Random (realistic)") + + +def parse_args(): + parser = argparse.ArgumentParser(description="Test ORCA multi-user batching") + parser.add_argument("mode", nargs="?", default="both", + choices=["synced", "random", "both"]) + parser.add_argument("-c", "--clients", type=int, default=4, + help="Number of concurrent clients (default: 4)") + parser.add_argument("-r", "--rounds", type=int, default=10, + help="Number of rounds per client (default: 10)") + return parser.parse_args() + + +async def main(): + args = parse_args() + print(f"Connecting {args.clients} clients to {SERVER_URL}...\n") + + if args.mode in ("synced", "both"): + await run_synced(args.clients, args.rounds) + if args.mode in ("random", "both"): + if args.mode == "both": + print("\n" + "=" * 60 + "\n") + await run_random(args.clients, args.rounds) + + +if __name__ == "__main__": + asyncio.run(main())