diff --git a/examples/inference/basic/basic_matrixgame_streaming.py b/examples/inference/basic/basic_matrixgame_streaming.py new file mode 100644 index 000000000..2454cda4a --- /dev/null +++ b/examples/inference/basic/basic_matrixgame_streaming.py @@ -0,0 +1,98 @@ +from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator +from fastvideo.models.dits.matrix_game.utils import get_current_action_async, expand_action_to_frames + +import torch +import asyncio + +# Available variants: "base_distilled_model", "gta_distilled_model", "templerun_distilled_model" +# Each variant has different keyboard_dim: +# - base_distilled_model: keyboard_dim=4 +# - gta_distilled_model: keyboard_dim=2 +# - templerun_distilled_model: keyboard_dim=7 (keyboard only, no mouse) +MODEL_VARIANT = "base_distilled_model" + +# Variant-specific settings +VARIANT_CONFIG = { + "base_distilled_model": { + "model_path": "FastVideo/Matrix-Game-2.0-Base-Diffusers", + "keyboard_dim": 4, + "mode": "universal", + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", + }, + "gta_distilled_model": { + "model_path": "FastVideo/Matrix-Game-2.0-GTA-Diffusers", + "keyboard_dim": 2, + "mode": "gta_drive", + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/gta_drive/0000.png", + }, + "templerun_distilled_model": { + "model_path": "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers", + "keyboard_dim": 7, + "mode": "templerun", + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/temple_run/0000.png", + }, +} + + +OUTPUT_PATH = "video_samples_matrixgame2" +async def main(): + # FastVideo will automatically use the optimal default arguments for the + # model. + # If a local path is provided, FastVideo will make a best effort + # attempt to identify the optimal arguments. + config = VARIANT_CONFIG[MODEL_VARIANT] + + generator = StreamingVideoGenerator.from_pretrained( + config["model_path"], + # FastVideo will automatically handle distributed setup + num_gpus=1, + use_fsdp_inference=True, + dit_cpu_offload=True, # DiT need to be offloaded for MoE + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer + pin_cpu_memory=True, + # image_encoder_cpu_offload=False, + ) + + max_blocks = 50 + num_frames = 597 + actions = { + "keyboard": torch.zeros((num_frames, config["keyboard_dim"])), + "mouse": torch.zeros((num_frames, 2)) + } + grid_sizes = torch.tensor([150, 44, 80]) + mode = config["mode"] + + generator.reset( + prompt="", + image_path=config["image_url"], + mouse_cond=actions["mouse"].unsqueeze(0), + keyboard_cond=actions["keyboard"].unsqueeze(0), + grid_sizes=grid_sizes, + num_frames=num_frames, + height=352, + width=640, + num_inference_steps=50, + output_path=OUTPUT_PATH, + save_video=True, + ) + print("Initialization complete.") + + for block_id in range(max_blocks): + print(f"\n=== Block {block_id + 1}/{max_blocks} ===") + + action = await get_current_action_async(mode) + keyboard_cond, mouse_cond = expand_action_to_frames(action, 12) + await generator.step_async(keyboard_cond, mouse_cond) + + if (await asyncio.to_thread(input, "\nContinue? (y/n): ")).lower() == 'n': + break + + # Save final video + generator.finalize() + generator.shutdown() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/inference/gradio/local/gradio_local_demo_matrixgame.py b/examples/inference/gradio/local/gradio_local_demo_matrixgame.py new file mode 100644 index 000000000..ae8333a04 --- /dev/null +++ b/examples/inference/gradio/local/gradio_local_demo_matrixgame.py @@ -0,0 +1,684 @@ +import argparse +import asyncio +import os +import time + +import gradio as gr +import torch +import uvicorn +from fastapi import FastAPI, Request, HTTPException +from fastapi.responses import HTMLResponse, FileResponse + +from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator +from fastvideo.models.dits.matrix_game.utils import expand_action_to_frames + + +VARIANT_CONFIG = { + "Matrix-Game-2.0-Base": { + "model_path": "FastVideo/Matrix-Game-2.0-Base-Diffusers", + "keyboard_dim": 4, + "mode": "universal", + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", + }, + "Matrix-Game-2.0-GTA": { + "model_path": "FastVideo/Matrix-Game-2.0-GTA-Diffusers", + "keyboard_dim": 2, + "mode": "gta_drive", + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/gta_drive/0000.png", + }, + "Matrix-Game-2.0-TempleRun": { + "model_path": "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers", + "keyboard_dim": 7, + "mode": "templerun", + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/temple_run/0000.png", + }, +} + +MODEL_PATH_MAPPING = { + name: config["model_path"] for name, config in VARIANT_CONFIG.items() +} + + +CAM_VALUE = 0.1 +KEYBOARD_MAP_UNIVERSAL = { + "W (Forward)": [1, 0, 0, 0], + "S (Back)": [0, 1, 0, 0], + "A (Left)": [0, 0, 1, 0], + "D (Right)": [0, 0, 0, 1], + "Q (Stop)": [0, 0, 0, 0], +} +KEYBOARD_MAP_GTA = { + "W (Forward)": [1, 0], + "S (Back)": [0, 1], + "Q (Stop)": [0, 0], +} +KEYBOARD_MAP_TEMPLERUN = { + "Q (Run)": [1, 0, 0, 0, 0, 0, 0], + "W (Jump)": [0, 1, 0, 0, 0, 0, 0], + "S (Slide)": [0, 0, 1, 0, 0, 0, 0], + "Z (Turn Left)": [0, 0, 0, 1, 0, 0, 0], + "C (Turn Right)": [0, 0, 0, 0, 1, 0, 0], + "A (Left)": [0, 0, 0, 0, 0, 1, 0], + "D (Right)": [0, 0, 0, 0, 0, 0, 1], +} + + +CAMERA_MAP_UNIVERSAL = { + "U (Center)": [0, 0], + "I (Up)": [CAM_VALUE, 0], + "K (Down)": [-CAM_VALUE, 0], + "J (Left)": [0, -CAM_VALUE], + "L (Right)": [0, CAM_VALUE], +} +CAMERA_MAP_GTA = { + "Q (Straight)": [0, 0], + "A (Steer Left)": [0, -CAM_VALUE], + "D (Steer Right)": [0, CAM_VALUE], +} + +def setup_model_environment(model_path: str) -> None: + # if "fullattn" in model_path.lower(): + # os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" + # else: + # os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN" + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" + os.environ["FASTVIDEO_STAGE_LOGGING"] = "1" + +def create_timing_display(inference_time, total_time, stage_execution_times, num_frames): + dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A" + + timing_html = f""" +
+

⏱️ Timing Breakdown

+
+
+
🚀
+
DiT Denoising
+
{dit_denoising_time}
+
+
+
🧠
+
E2E (w. vae/text encoder)
+
{inference_time:.2f}s
+
+
+
🎬
+
Video Encoding
+
N/A
+
+
+
🌐
+
Network Transfer
+
N/A
+
+
+
📊
+
Total Processing
+
{total_time:.2f}s
+
+
""" + + if inference_time > 0: + fps = num_frames / inference_time + timing_html += f""" +
+ Generation Speed: + {fps:.1f} frames/second +
""" + + return timing_html + "
" + +def get_action_tensors(mode: str, keyboard_key: str, mouse_key: str | None): + if mode == "universal": + keyboard = torch.tensor(KEYBOARD_MAP_UNIVERSAL.get(keyboard_key, [0, 0, 0, 0])).cuda() + mouse = torch.tensor(CAMERA_MAP_UNIVERSAL.get(mouse_key, [0, 0])).cuda() + elif mode == "gta_drive": + keyboard = torch.tensor(KEYBOARD_MAP_GTA.get(keyboard_key, [0, 0])).cuda() + mouse = torch.tensor(CAMERA_MAP_GTA.get(mouse_key, [0, 0])).cuda() + elif mode == "templerun": + keyboard = torch.tensor(KEYBOARD_MAP_TEMPLERUN.get(keyboard_key, [1, 0, 0, 0, 0, 0, 0])).cuda() + mouse = None + else: + raise ValueError(f"Unknown mode: {mode}") + + return {"keyboard": keyboard, "mouse": mouse} + +def create_gradio_interface(generators: dict[str, StreamingVideoGenerator], loaded_model_name: str): + initial_config = VARIANT_CONFIG.get(loaded_model_name, VARIANT_CONFIG["Matrix-Game-2.0-Base"]) + initial_mode = initial_config["mode"] + + if initial_mode == "universal": + initial_kb_choices = list(KEYBOARD_MAP_UNIVERSAL.keys()) + initial_mouse_choices = list(CAMERA_MAP_UNIVERSAL.keys()) + initial_mouse_visible = True + elif initial_mode == "gta_drive": + initial_kb_choices = list(KEYBOARD_MAP_GTA.keys()) + initial_mouse_choices = list(CAMERA_MAP_GTA.keys()) + initial_mouse_visible = True + else: # templerun + initial_kb_choices = list(KEYBOARD_MAP_TEMPLERUN.keys()) + initial_mouse_choices = [] + initial_mouse_visible = False + + theme = gr.themes.Base().set( + button_primary_background_fill="#2563eb", + button_primary_background_fill_hover="#1d4ed8", + button_primary_text_color="white", + slider_color="#2563eb", + checkbox_background_color_selected="#2563eb", + ) + + with gr.Blocks(title="FastVideo - Matrix Game 2.0", theme=theme) as demo: + game_state = gr.State({ + "initialized": False, + "current_model": None, + "block_idx": 0, + "max_blocks": 50, + }) + + # Header + gr.Image("assets/full.svg", show_label=False, container=False, height=80) + + gr.HTML(""" +
+

Make Video Generation Go Blurrrrrrr

+

Code | Blog | Docs

+
+ """) + + with gr.Accordion("🎥 What Is FastVideo?", open=False): + gr.HTML(""" +
+

+ FastVideo is an inference and post-training framework for diffusion models. It features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered. +

+
+ """) + + # Model Selection + with gr.Row(): + model_selection = gr.Dropdown( + choices=[loaded_model_name], + value=loaded_model_name, + label="Select Model", + interactive=False + ) + + + # Main Layout + with gr.Row(equal_height=True, elem_classes="main-content-row"): + with gr.Column(scale=1, elem_classes="advanced-options-column"): + with gr.Group(): + gr.HTML("
Game Controls
") + + with gr.Group(): + gr.HTML("
🎮 Keyboard Control
") + keyboard_action = gr.Radio( + choices=initial_kb_choices, + value=initial_kb_choices[0] if initial_kb_choices else None, + label="Movement", + show_label=False, + interactive=True + ) + + with gr.Group(visible=initial_mouse_visible) as mouse_group: + gr.HTML("
🖱️ Mouse/Camera Control
") + mouse_action = gr.Radio( + choices=initial_mouse_choices if initial_mouse_visible else [], + value=initial_mouse_choices[0] if initial_mouse_choices else None, + label="Camera", + show_label=False, + interactive=True + ) + + with gr.Row(): + action_btn = gr.Button("Start", variant="primary") + stop_btn = gr.Button("Stop", variant="stop") + + gr.HTML("
") + + seed = gr.Slider( + label="Seed", + minimum=0, + maximum=1000000, + step=1, + value=1024, + ) + randomize_seed = gr.Checkbox(label="Randomize seed", value=False) + seed_output = gr.Number(label="Used Seed") + + block_counter = gr.Textbox(label="Progress", value="Block: 0 / 50", interactive=False, lines=1) + + + # Right Column: Video Output + with gr.Column(scale=1, elem_classes="video-column"): + video_output = gr.Video( + label="Generated Video", + show_label=True, + height=466, + width=600, + container=True, + elem_classes="video-component", + autoplay=True + ) + + # Styles + gr.HTML(""" + + """) + + # UI update based on model selection + def on_model_change(model_name): + config = VARIANT_CONFIG.get(model_name, VARIANT_CONFIG["Matrix-Game-2.0-Base"]) + mode = config["mode"] + + if mode == "universal": + kb_choices = list(KEYBOARD_MAP_UNIVERSAL.keys()) + mouse_choices = list(CAMERA_MAP_UNIVERSAL.keys()) + mouse_visible = True + elif mode == "gta_drive": + kb_choices = list(KEYBOARD_MAP_GTA.keys()) + mouse_choices = list(CAMERA_MAP_GTA.keys()) + mouse_visible = True + else: # templerun + kb_choices = list(KEYBOARD_MAP_TEMPLERUN.keys()) + mouse_choices = [] + mouse_visible = False + + return ( + gr.update(choices=kb_choices, value=kb_choices[0] if kb_choices else None), + gr.update(choices=mouse_choices, value=mouse_choices[0] if mouse_choices else None, visible=mouse_visible), + gr.update(visible=mouse_visible), + ) + + model_selection.change( + fn=on_model_change, + inputs=model_selection, + outputs=[keyboard_action, mouse_action, mouse_group] + ) + + def start_game(model_name, seed_val, randomize, state): + if randomize: + seed_val = torch.randint(0, 1000000, (1,)).item() + + config = VARIANT_CONFIG.get(model_name) + if not config: + return state, seed_val, "Block: 0 / 50", None, "", gr.update(), gr.update() + + generator = generators.get(config["model_path"]) + if not generator: + return state, seed_val, "Block: 0 / 50", None, "", gr.update(), gr.update() + + # If already initialized, clean up first + if state.get("initialized"): + try: + # Clear accumulated frames without saving + generator.accumulated_frames = [] + generator.executor.execute_streaming_clear() + except Exception as e: + print(f"Warning: cleanup error: {e}") + + # Streaming parameters + num_latent_frames_per_block = 3 + max_blocks = 50 + total_latent_frames = num_latent_frames_per_block * max_blocks + num_frames = (total_latent_frames - 1) * 4 + 1 + + actions = { + "keyboard": torch.zeros((num_frames, config["keyboard_dim"])), + "mouse": torch.zeros((num_frames, 2)) + } + grid_sizes = torch.tensor([150, 44, 80]) + + output_dir = os.path.abspath("outputs/matrixgame") + os.makedirs(output_dir, exist_ok=True) + video_path = os.path.join(output_dir, f"video_{int(time.time())}.mp4") + + generator.reset( + prompt="", + image_path=config["image_url"], + mouse_cond=actions["mouse"].unsqueeze(0), + keyboard_cond=actions["keyboard"].unsqueeze(0), + grid_sizes=grid_sizes, + num_frames=num_frames, + height=352, + width=640, + num_inference_steps=50, + output_path=video_path, + ) + + new_state = { + "initialized": True, + "current_model": model_name, + "block_idx": 0, + "max_blocks": max_blocks, + "video_path": video_path, + "frames_per_block": num_latent_frames_per_block * 4, + "mode": config["mode"], + "seed": seed_val, + } + + return new_state, seed_val, "Block: 0 / 50", None, gr.update(value="Step"), gr.update(interactive=True) + + async def step_game(keyboard_key, mouse_key, model_name, state): + if not state.get("initialized"): + return state, state.get("seed", 0), "Block: 0 / 50", None, gr.update(), gr.update() + + # total_start_time = time.time() + config = VARIANT_CONFIG.get(model_name) + generator = generators.get(config["model_path"]) + mode = state["mode"] + frames_per_block = state["frames_per_block"] + + # Parse inputs to tensors + action = get_action_tensors(mode, keyboard_key, mouse_key) + keyboard_cond, mouse_cond = expand_action_to_frames(action, frames_per_block) + + # run step async + # inference_start_time = time.time() + frames, block_future = await generator.step_async(keyboard_cond, mouse_cond) + # inference_time = time.time() - inference_start_time + + # wait for block file to be written + block_path = await asyncio.to_thread(block_future.result) if block_future else None + state["block_idx"] = generator.block_idx + block_str = f"Block: {state['block_idx']} / {state['max_blocks']}" + + # total_time = time.time() - total_start_time + + # Timing breakdown + # timing_html = create_timing_display(inference_time, total_time, [], frames_per_block) + + return state, state.get("seed", 0), block_str, block_path, gr.update(), gr.update() + + def stop_game(model_name, state): + if not state.get("initialized"): + return {"initialized": False}, 0, "Block: 0 / 50", None, gr.update(value="Start"), gr.update(interactive=False) + + config = VARIANT_CONFIG.get(model_name) + generator = generators.get(config["model_path"]) + + final_path = state.get("video_path") + generator.finalize(final_path) + + return {"initialized": False}, state.get("seed", 0), "Block: 0 / 50", final_path, gr.update(value="Start"), gr.update(interactive=False) + + async def handle_action(keyboard_key, mouse_key, model_name, seed_val, randomize, state): + if not state.get("initialized"): + return start_game(model_name, seed_val, randomize, state) + else: + return await step_game(keyboard_key, mouse_key, model_name, state) + + action_btn.click( + fn=handle_action, + inputs=[keyboard_action, mouse_action, model_selection, seed, randomize_seed, game_state], + outputs=[game_state, seed_output, block_counter, video_output, action_btn, stop_btn] + ) + + stop_btn.click( + fn=stop_game, + inputs=[model_selection, game_state], + outputs=[game_state, seed_output, block_counter, video_output, action_btn, stop_btn] + ) + + gr.HTML(""" +
+

Note that this demo is meant to showcase Matrix Game's quality and that under a large number of requests, generation speed may be affected.

+
+ """) + + return demo + + +def main(): + parser = argparse.ArgumentParser(description="Matrix Game Gradio Demo") + parser.add_argument("--model", type=str, default="Matrix-Game-2.0-Base", + choices=list(VARIANT_CONFIG.keys()), + help="Model variant to load") + parser.add_argument("--host", type=str, default="0.0.0.0") + parser.add_argument("--port", type=int, default=7860) + args = parser.parse_args() + + # Load the selected model + config = VARIANT_CONFIG[args.model] + model_path = config["model_path"] + + print(f"Loading model: {model_path}") + setup_model_environment(model_path) + generator = StreamingVideoGenerator.from_pretrained( + model_path, + num_gpus=1, + use_fsdp_inference=True, + dit_cpu_offload=True, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + pin_cpu_memory=True, + ) + + generators = {model_path: generator} + + demo = create_gradio_interface(generators, args.model) + + print(f"Starting Gradio at http://{args.host}:{args.port}") + + # FastAPI Wrapper + app = FastAPI() + + @app.get("/logo.png") + def get_logo(): + return FileResponse( + "assets/full.svg", + media_type="image/svg+xml", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + } + ) + + @app.get("/favicon.ico") + def get_favicon(): + favicon_path = "assets/icon-simple.svg" + + if os.path.exists(favicon_path): + return FileResponse( + favicon_path, + media_type="image/svg+xml", + headers={ + "Cache-Control": "public, max-age=3600", + "Access-Control-Allow-Origin": "*" + } + ) + else: + raise HTTPException(status_code=404, detail="Favicon not found") + + @app.get("/", response_class=HTMLResponse) + def index(request: Request): + base_url = str(request.base_url).rstrip('/') + return f""" + + + + + + + FastVideo - Matrix Game 2.0 + + + + + + + + + + + + + + + + + + + + + + + + + + + + """ + + app = gr.mount_gradio_app( + app, + demo, + path="/gradio", + allowed_paths=[os.path.abspath("outputs"), os.path.abspath("fastvideo-logos")] + ) + + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/entrypoints/streaming_generator.py b/fastvideo/entrypoints/streaming_generator.py new file mode 100644 index 000000000..fd52dc9dc --- /dev/null +++ b/fastvideo/entrypoints/streaming_generator.py @@ -0,0 +1,288 @@ +import asyncio +import os +from concurrent.futures import Future, ThreadPoolExecutor + +import imageio +import numpy as np +import torch +import torchvision +from einops import rearrange + +from fastvideo.configs.sample import SamplingParam +from fastvideo.entrypoints.video_generator import VideoGenerator +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines import ForwardBatch +from fastvideo.utils import align_to, shallow_asdict +from fastvideo.worker.executor import Executor +from fastvideo.worker.multiproc_executor import MultiprocExecutor + +logger = init_logger(__name__) + + +class IncrementalVideoWriter: + + def __init__(self, path: str, fps: int = 24, block_dir: str | None = None): + self._executor = ThreadPoolExecutor(max_workers=2, + thread_name_prefix="video_write_") + self._path = path + self._writer = imageio.get_writer(path, fps=fps, format="mp4") + self._pending_main: Future | None = None + self._block_dir = block_dir + self._block_idx = 0 + self._fps = fps + + @property + def path(self) -> str: + return self._path + + def add_frames(self, frames: list[np.ndarray]) -> Future | None: + # Wait for previous main video write to complete + if self._pending_main is not None: + self._pending_main.result() + + # Copy frames to avoid race conditions + frames_copy = [f.copy() for f in frames] + self._pending_main = self._executor.submit(self._write_frames, + frames_copy) + + # Write block file if block_dir is set + block_future = None + if self._block_dir: + self._block_idx += 1 + block_path = os.path.join(self._block_dir, + f"b{self._block_idx}.mp4") + block_future = self._executor.submit(self._write_block, frames_copy, + block_path) + return block_future + + def _write_frames(self, frames: list[np.ndarray]) -> None: + for frame in frames: + self._writer.append_data(frame) + + def _write_block(self, frames: list[np.ndarray], path: str) -> str: + imageio.mimsave(path, frames, fps=self._fps) + return path + + def close(self) -> None: + if self._pending_main is not None: + self._pending_main.result() + self._pending_main = None + if self._writer: + self._writer.close() + self._writer = None + self._executor.shutdown(wait=True) + + +class StreamingVideoGenerator(VideoGenerator): + """ + This class extends VideoGenerator with streaming capabilities, + allowing incremental video generation with step-by-step control. + """ + + def __init__(self, + fastvideo_args: FastVideoArgs, + executor_class: type[Executor], + log_stats: bool, + use_queue_mode: bool = True): + super().__init__(fastvideo_args, executor_class, log_stats) + self.accumulated_frames: list[np.ndarray] = [] + self.sampling_param: SamplingParam | None = None + self.batch: ForwardBatch | None = None + self._use_queue_mode = use_queue_mode and isinstance( + self.executor, MultiprocExecutor) + self.writer: IncrementalVideoWriter | None = None + self.block_dir: str | None = None + self.block_idx: int = 0 + + @classmethod + def from_fastvideo_args( + cls, fastvideo_args: FastVideoArgs) -> "StreamingVideoGenerator": + executor_class = Executor.get_class(fastvideo_args) + return cls( + fastvideo_args=fastvideo_args, + executor_class=executor_class, + log_stats=False, + ) + + def reset( + self, + prompt: str = "A gameplay video of a cyberpunk city", + image_path: str | None = None, + num_frames: int = 120, # Default max frames + **kwargs): + self.accumulated_frames = [] + self.block_idx = 0 + self.block_dir = None + if self.writer: + self.writer.close() + self.writer = None + self.executor.execute_streaming_clear() + + # Handle batch processing from text file + if self.sampling_param is None: + self.sampling_param = SamplingParam.from_pretrained( + self.fastvideo_args.model_path) + + self.sampling_param.update(kwargs) + self.sampling_param.prompt = prompt + if image_path: + self.sampling_param.image_path = image_path + self.sampling_param.num_frames = num_frames + + if "output_path" in kwargs: + output_path = self._prepare_output_path(kwargs["output_path"], + prompt) + # Create block directory for individual block files + block_dir = output_path.replace(".mp4", "") + os.makedirs(block_dir, exist_ok=True) + self.block_dir = block_dir + self.writer = IncrementalVideoWriter(output_path, + fps=24, + block_dir=block_dir) + + fastvideo_args = self.fastvideo_args + + self.sampling_param.height = align_to(self.sampling_param.height, 16) + self.sampling_param.width = align_to(self.sampling_param.width, 16) + + latents_size = [(self.sampling_param.num_frames - 1) // 4 + 1, + self.sampling_param.height // 8, + self.sampling_param.width // 8] + n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + + self.sampling_param.return_frames = True + self.sampling_param.save_video = False + + self.batch = ForwardBatch( + **shallow_asdict(self.sampling_param), + eta=0.0, + n_tokens=n_tokens, + VSA_sparsity=fastvideo_args.VSA_sparsity, + ) + + if self._use_queue_mode: + self.executor.submit_reset(self.batch, fastvideo_args) + result = self.executor.wait_result() + if result.error: + raise result.error + else: + self.executor.execute_streaming_reset(self.batch, fastvideo_args) + + def step( + self, keyboard_cond: torch.Tensor, + mouse_cond: torch.Tensor) -> tuple[list[np.ndarray], Future | None]: + if self.batch is None: + raise RuntimeError("Call reset() before step()") + + if self._use_queue_mode and self.executor._streaming_enabled: + self.executor.submit_step(keyboard_cond, mouse_cond) + result = self.executor.wait_result() + if result.error: + raise result.error + output_batch = result.output_batch + else: + # Fallback to RPC-based + output_batch = self.executor.execute_streaming_step( + keyboard_action=keyboard_cond, mouse_action=mouse_cond) + + frames = self._process_output_batch(output_batch) + block_future = None + if len(frames) > 0: + self.accumulated_frames.extend(frames) + self.block_idx += 1 + if self.writer: + # Returns Future for block file, or None if no block_dir + block_future = self.writer.add_frames(frames) + + return frames, block_future + + async def step_async( + self, keyboard_cond: torch.Tensor, + mouse_cond: torch.Tensor) -> tuple[list[np.ndarray], Future | None]: + if self.batch is None: + raise RuntimeError("Call reset() before step_async()") + + if self._use_queue_mode and self.executor._streaming_enabled: + self.executor.submit_step(keyboard_cond, mouse_cond) + + loop = asyncio.get_running_loop() + result = await loop.run_in_executor(None, self.executor.wait_result) + + if result.error: + raise result.error + output_batch = result.output_batch + else: + # Fallback to RPC-based + output_batch = await self.executor.execute_streaming_step_async( + keyboard_action=keyboard_cond, + mouse_action=mouse_cond, + ) + + frames = self._process_output_batch(output_batch) + block_future = None + if len(frames) > 0: + self.accumulated_frames.extend(frames) + self.block_idx += 1 + if self.writer: + block_future = self.writer.add_frames(frames) + + return frames, block_future + + def finalize(self, + output_path: str = "streaming_output.mp4", + fps: int = 24) -> str: + if not self.accumulated_frames: + logger.warning("No frames to save.") + return "" + + if self.writer: + output_path = self.writer.path + self.writer.close() + self.writer = None + logger.info("Saved video to %s", output_path) + else: + imageio.mimsave(output_path, + self.accumulated_frames, + fps=fps, + format="mp4") + logger.info("Saved video to %s", output_path) + + if self._use_queue_mode and self.executor._streaming_enabled: + self.executor.submit_clear() + else: + self.executor.execute_streaming_clear() + self.accumulated_frames = [] + return output_path + + def _process_output_batch(self, + output_batch: ForwardBatch) -> list[np.ndarray]: + if output_batch.output is None: + return [] + + samples = output_batch.output + # [B, C, T, H, W] or [1, C, T, H, W] + if len(samples.shape) == 5: + # Rearrange to [T, B, C, H, W] for processing loop + videos = rearrange(samples, "b c t h w -> t b c h w") + else: + logger.warning("Unexpected output shape: %s", samples.shape) + 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 + + def shutdown(self): + if self.writer: + self.writer.close() + self.writer = None + + if self._use_queue_mode and self.executor._streaming_enabled: + self.executor.disable_streaming() + + super().shutdown() diff --git a/fastvideo/models/dits/matrix_game/utils.py b/fastvideo/models/dits/matrix_game/utils.py index d3ef3667a..7f365eed7 100644 --- a/fastvideo/models/dits/matrix_game/utils.py +++ b/fastvideo/models/dits/matrix_game/utils.py @@ -1,14 +1,13 @@ from __future__ import annotations +import asyncio import os import random -# import cv2 -import numpy as np import torch -from diffusers.utils import export_to_video from PIL import Image +from fastvideo.distributed.parallel_state import get_local_torch_device from fastvideo.utils import logger @@ -36,6 +35,120 @@ KEYBOARD_MAP_7 = { # templerun_distilled_model: still/w/s/left/right/a/d } KEYBOARD_MAP = KEYBOARD_MAP_4 # Default for backward compatibility +def expand_action_to_frames(action: dict, num_frames: int) -> tuple[torch.Tensor, torch.Tensor]: + result = {} + for key, tensor in action.items(): + if tensor is not None: + # Expand to [num_frames, D] then unsqueeze to [1, num_frames, D] + result[key] = tensor.unsqueeze(0).repeat(num_frames, 1).unsqueeze(0) + else: + result[key] = None + + if "mouse" not in result or result["mouse"] is None: + # keyboard device if available, otherwise default + device = result.get("keyboard", torch.tensor([])).device if result.get("keyboard") is not None else get_local_torch_device() + result["mouse"] = torch.zeros(1, num_frames, 2, device=device) + + return result["keyboard"], result["mouse"] + +def get_current_action(mode="universal"): + + CAM_VALUE = 0.1 + if mode == 'universal': + logger.info("") + logger.info('-'*30) + logger.info("PRESS [I, K, J, L, U] FOR CAMERA TRANSFORM\n (I: up, K: down, J: left, L: right, U: no move)") + logger.info("PRESS [W, S, A, D, Q] FOR MOVEMENT\n (W: forward, S: back, A: left, D: right, Q: no move)") + logger.info('-'*30) + CAMERA_VALUE_MAP = { + "i": [CAM_VALUE, 0], + "k": [-CAM_VALUE, 0], + "j": [0, -CAM_VALUE], + "l": [0, CAM_VALUE], + "u": [0, 0] + } + KEYBOARD_IDX = { + "w": [1, 0, 0, 0], "s": [0, 1, 0, 0], "a": [0, 0, 1, 0], "d": [0, 0, 0, 1], + "q": [0, 0, 0, 0] + } + flag = 0 + while flag != 1: + try: + idx_mouse = input('Please input the mouse action (e.g. `U`):\n').strip().lower() + idx_keyboard = input('Please input the keyboard action (e.g. `W`):\n').strip().lower() + if idx_mouse in CAMERA_VALUE_MAP and idx_keyboard in KEYBOARD_IDX: + flag = 1 + except Exception: + pass + mouse_cond = torch.tensor(CAMERA_VALUE_MAP[idx_mouse]).cuda() + keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard]).cuda() + elif mode == 'gta_drive': + logger.info("") + logger.info('-'*30) + logger.info("PRESS [W, S, A, D, Q] FOR MOVEMENT\n (W: forward, S: back, A: left, D: right, Q: no move)") + logger.info('-'*30) + CAMERA_VALUE_MAP = { + "a": [0, -CAM_VALUE], + "d": [0, CAM_VALUE], + "q": [0, 0] + } + KEYBOARD_IDX = { + "w": [1, 0], "s": [0, 1], + "q": [0, 0] + } + flag = 0 + while flag != 1: + try: + indexes = input('Please input the actions (split with ` `):\n(e.g. `W` for forward, `W A` for forward and left)\n').strip().lower().split(' ') + idx_mouse = [] + idx_keyboard = [] + for i in indexes: + if i in CAMERA_VALUE_MAP.keys(): + idx_mouse += [i] + elif i in KEYBOARD_IDX.keys(): + idx_keyboard += [i] + if len(idx_mouse) == 0: + idx_mouse += ['q'] + if len(idx_keyboard) == 0: + idx_keyboard += ['q'] + assert idx_mouse in [['a'], ['d'], ['q']] and idx_keyboard in [['q'], ['w'], ['s']] + flag = 1 + except Exception: + pass + mouse_cond = torch.tensor(CAMERA_VALUE_MAP[idx_mouse[0]]).cuda() + keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard[0]]).cuda() + elif mode == 'templerun': + logger.info("") + logger.info('-'*30) + logger.info("PRESS [W, S, A, D, Z, C, Q] FOR ACTIONS\n (W: jump, S: slide, A: left side, D: right side, Z: turn left, C: turn right, Q: no move)") + logger.info('-'*30) + KEYBOARD_IDX = { + "w": [0, 1, 0, 0, 0, 0, 0], "s": [0, 0, 1, 0, 0, 0, 0], + "a": [0, 0, 0, 0, 0, 1, 0], "d": [0, 0, 0, 0, 0, 0, 1], + "z": [0, 0, 0, 1, 0, 0, 0], "c": [0, 0, 0, 0, 1, 0, 0], + "q": [1, 0, 0, 0, 0, 0, 0] + } + flag = 0 + while flag != 1: + try: + idx_keyboard = input('Please input the action: \n(e.g. `W` for forward, `Z` for turning left)\n').strip().lower() + if idx_keyboard in KEYBOARD_IDX.keys(): + flag = 1 + except Exception: + pass + keyboard_cond = torch.tensor(KEYBOARD_IDX[idx_keyboard]).cuda() + + if mode != 'templerun': + return { + "mouse": mouse_cond, + "keyboard": keyboard_cond + } + return { + "keyboard": keyboard_cond + } + +async def get_current_action_async(mode="universal"): + return await asyncio.to_thread(get_current_action, mode) def load_initial_image(image_path: str = None) -> Image.Image: if image_path and os.path.exists(image_path): @@ -43,7 +156,6 @@ def load_initial_image(image_path: str = None) -> Image.Image: logger.warning("No image provided, creating placeholder...") return Image.new("RGB", (640, 352), (128, 128, 128)) - def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None): if keyboard_dim not in (2, 4, 7): raise ValueError(f"keyboard_dim must be 2, 4, or 7, got {keyboard_dim}") @@ -145,7 +257,6 @@ def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = No return {"keyboard": keyboard_condition, "mouse": mouse_condition} - def parse_config(config, mode="universal"): assert mode in ['universal', 'gta_drive', 'templerun'] key_data = {} @@ -183,7 +294,6 @@ def parse_config(config, mode="universal"): ) return key_data, mouse_data - # NOTE: drawing functions are commented out to avoid cv2/libGL dependency. # # def draw_rounded_rectangle(image, top_left, bottom_right, color, radius=10, alpha=0.5): @@ -199,7 +309,6 @@ def parse_config(config, mode="universal"): # cv2.ellipse(overlay, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, -1) # cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image) # -# # def draw_keys_on_frame(frame, keys, key_size=(80, 50), spacing=20, bottom_margin=30, mode='universal'): # h, w, _ = frame.shape # horison_shift = 90 @@ -244,7 +353,6 @@ def parse_config(config, mode="universal"): # text_y = y + (key_size[1] + text_size[1]) // 2 # cv2.putText(frame, key_icon[key], (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2) # -# # def overlay_icon(frame, icon, position, scale=1.0, rotation=0): # x, y = position # h, w, _ = icon.shape @@ -282,7 +390,6 @@ def parse_config(config, mode="universal"): # frame_region[:, :, c] = (1 - alpha) * frame_region[:, :, c] + alpha * icon_rgb[:, :, c] # frame[top_left_y:bottom_right_y, top_left_x:bottom_right_x] = frame_region # -# # def process_video(input_video, output_video, config, mouse_icon_path, # mouse_scale=1.0, mouse_rotation=0, process_icon=True, mode='universal'): # key_data, mouse_data = parse_config(config, mode=mode) diff --git a/fastvideo/models/vaes/wanvae.py b/fastvideo/models/vaes/wanvae.py index 4e18dc00d..2d1502be8 100644 --- a/fastvideo/models/vaes/wanvae.py +++ b/fastvideo/models/vaes/wanvae.py @@ -1252,6 +1252,52 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): dec = dec[:, :, start_frame_idx:] return dec + def get_streaming_cache(self) -> list[torch.Tensor | None]: + def _count_conv3d(model) -> int: + count = 0 + for m in model.modules(): + if isinstance(m, WanCausalConv3d): + count += 1 + return count + + conv_num = _count_conv3d(self.decoder) + return [None] * conv_num + + def streaming_decode( + self, + z: torch.Tensor, + cache: list[torch.Tensor | None], + is_first_chunk: bool = False, + ) -> tuple[torch.Tensor, list[torch.Tensor | None]]: + """ + Args: + z (`torch.Tensor`): Latent tensor of shape [B, C, T, H, W]. + cache (`list[torch.Tensor | None]`): The VAE cache. + is_first_chunk (`bool`): Whether this is the first chunk in the sequence. + + Returns: + A tuple of (decoded_frames, updated_cache). + """ + iter_ = z.shape[2] + x = self.post_quant_conv(z) + + with forward_context(feat_cache_arg=cache, feat_idx_arg=0): + outputs = [] + for i in range(iter_): + feat_idx.set(0) + first_chunk.set(is_first_chunk and i == 0) + decoded_chunk = self.decoder(x[:, :, i:i + 1, :, :]) + outputs.append(decoded_chunk) + out = torch.cat(outputs, dim=2) + + if self.config.patch_size is not None: + out = unpatchify(out, patch_size=self.config.patch_size) + + out = out.float() + out = torch.clamp(out, min=-1.0, max=1.0) + + return out, cache + def forward( self, sample: torch.Tensor, @@ -1272,3 +1318,4 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): z = posterior.mode() dec = self.decode(z) return dec + diff --git a/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py index e2f8ce6bb..837c47988 100644 --- a/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py +++ b/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py @@ -2,8 +2,9 @@ """Matrix-Game causal DMD pipeline implementation.""" from fastvideo.fastvideo_args import FastVideoArgs +import torch from fastvideo.logger import init_logger -from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline +from fastvideo.pipelines import ComposedPipelineBase, ForwardBatch, LoRAPipeline from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, InputValidationStage, @@ -69,5 +70,61 @@ class MatrixGameCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): logger.info( "MatrixGameCausalDMDPipeline initialized with action support") + @torch.no_grad() + def streaming_reset(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs): + if not self.post_init_called: + self.post_init() + + # 1. Run Pre-processing stages + 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) + + # 2. Reset Denoising Stage + denoiser = self._stage_name_mapping["denoising_stage"] + denoiser.streaming_reset(batch, fastvideo_args) + + # 3. Initialize VAE cache + 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 + + # Decode only the new generated block + 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 = [MatrixGameCausalDMDPipeline] diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 264419958..7de5d8960 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -50,32 +50,7 @@ class DecodingStage(PipelineStage): result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)]) return result - @torch.no_grad() - def decode(self, latents: torch.Tensor, - fastvideo_args: FastVideoArgs) -> torch.Tensor: - """ - Decode latent representations into pixel space using VAE. - - Args: - latents: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents) - fastvideo_args: Configuration containing: - - disable_autocast: Whether to disable automatic mixed precision (default: False) - - pipeline_config.vae_precision: VAE computation precision ("fp32", "fp16", "bf16") - - pipeline_config.vae_tiling: Whether to enable VAE tiling for memory efficiency - - Returns: - Decoded video tensor with shape (batch, channels, frames, height, width), - normalized to [0, 1] range and moved to CPU as float32 - """ - self.vae = self.vae.to(get_local_torch_device()) - latents = latents.to(get_local_torch_device()) - - # Setup VAE precision - vae_dtype = PRECISION_TO_TYPE[ - fastvideo_args.pipeline_config.vae_precision] - vae_autocast_enabled = ( - vae_dtype != torch.float32) and not fastvideo_args.disable_autocast - + def _denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor: # denormalization for MatrixGame VAE # z = z * std + mean during decode if (hasattr(self.vae.config, 'latents_mean') @@ -109,6 +84,35 @@ class DecodingStage(PipelineStage): latents.dtype) else: latents += self.vae.shift_factor + return latents + + @torch.no_grad() + def decode(self, latents: torch.Tensor, + fastvideo_args: FastVideoArgs) -> torch.Tensor: + """ + Decode latent representations into pixel space using VAE. + + Args: + latents: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents) + fastvideo_args: Configuration containing: + - disable_autocast: Whether to disable automatic mixed precision (default: False) + - pipeline_config.vae_precision: VAE computation precision ("fp32", "fp16", "bf16") + - pipeline_config.vae_tiling: Whether to enable VAE tiling for memory efficiency + + Returns: + Decoded video tensor with shape (batch, channels, frames, height, width), + normalized to [0, 1] range and moved to CPU as float32 + """ + self.vae = self.vae.to(get_local_torch_device()) + latents = latents.to(get_local_torch_device()) + + # Setup VAE precision + vae_dtype = PRECISION_TO_TYPE[ + fastvideo_args.pipeline_config.vae_precision] + vae_autocast_enabled = ( + vae_dtype != torch.float32) and not fastvideo_args.disable_autocast + + latents = self._denormalize_latents(latents) # Decode latents with torch.autocast(device_type="cuda", @@ -126,6 +130,57 @@ class DecodingStage(PipelineStage): image = (image / 2 + 0.5).clamp(0, 1) return image + @torch.no_grad() + def streaming_decode( + self, + latents: torch.Tensor, + fastvideo_args: FastVideoArgs, + cache: list[torch.Tensor | None] | None = None, + is_first_chunk: bool = False, + ) -> tuple[torch.Tensor, list[torch.Tensor | None]]: + """ + Decode latent representations into pixel space using VAE with streaming cache. + + Args: + latents: Input latent tensor with shape (batch, channels, frames, height_latents, width_latents) + fastvideo_args: Configuration object. + cache: VAE cache from previous call, or None to initialize a new cache. + is_first_chunk: Whether this is the first chunk. + + Returns: + A tuple of (decoded_frames, updated_cache). + """ + self.vae = self.vae.to(get_local_torch_device()) + latents = latents.to(get_local_torch_device()) + + # Setup VAE precision + vae_dtype = PRECISION_TO_TYPE[ + fastvideo_args.pipeline_config.vae_precision] + vae_autocast_enabled = ( + vae_dtype != torch.float32) and not fastvideo_args.disable_autocast + + latents = self._denormalize_latents(latents) + + # Initialize cache if needed + if cache is None: + cache = self.vae.get_streaming_cache() + + # Decode latents with streaming + with 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: + latents = latents.to(vae_dtype) + image, cache = self.vae.streaming_decode(latents, cache, + is_first_chunk) + + # Normalize image to [0, 1] range + image = (image / 2 + 0.5).clamp(0, 1) + assert cache is not None, "cache should not be None after streaming_decode" + return image, cache + @torch.no_grad() def forward( self, diff --git a/fastvideo/pipelines/stages/matrixgame_denoising.py b/fastvideo/pipelines/stages/matrixgame_denoising.py index 70e17127d..3b442b068 100644 --- a/fastvideo/pipelines/stages/matrixgame_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame_denoising.py @@ -1,4 +1,6 @@ from __future__ import annotations +from collections.abc import Callable +from dataclasses import dataclass from typing import Any import torch # type: ignore @@ -32,6 +34,45 @@ except ImportError: logger = init_logger(__name__) +@dataclass +class BlockProcessingContext: + """Dataclass contains for block processing.""" + + batch: ForwardBatch + + block_idx: int + start_index: int + + kv_cache1: list[dict[Any, Any]] + kv_cache2: list[dict[Any, Any]] | None + kv_cache_mouse: list[dict[Any, Any]] | None + kv_cache_keyboard: list[dict[Any, Any]] | None + crossattn_cache: list[dict[Any, Any]] + + timesteps: torch.Tensor + block_sizes: list[int] + noise_pool: list[torch.Tensor] | None + + fastvideo_args: FastVideoArgs + target_dtype: torch.dtype + autocast_enabled: bool + boundary_timestep: float | None + high_noise_timesteps: torch.Tensor | None + context_noise: float + + image_kwargs: dict[str, Any] + pos_cond_kwargs: dict[str, Any] + + def get_kv_cache(self, timestep_val: float) -> list[dict[Any, Any]]: + if self.boundary_timestep is not None: + if timestep_val >= self.boundary_timestep: + return self.kv_cache1 + else: + assert self.kv_cache2 is not None, "kv_cache2 is not initialized" + return self.kv_cache2 + return self.kv_cache1 + + class MatrixGameCausalDenoisingStage(DenoisingStage): def __init__(self, @@ -80,6 +121,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): self.action_config = getattr(self.transformer, 'action_config', {}) self.use_action_module = len(self.action_config) > 0 + self._streaming_initialized: bool = False + self._streaming_ctx: BlockProcessingContext | None = None + def forward( self, batch: ForwardBatch, @@ -94,8 +138,6 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): patch_ratio = patch_size[-1] * patch_size[-2] self.frame_seq_length = latent_seq_length // patch_ratio - independent_first_frame = getattr(self.transformer, - 'independent_first_frame', False) timesteps = torch.tensor( fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu() @@ -152,23 +194,12 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): dtype=target_dtype, device=latents.device) - def _get_kv_cache(timestep_val: float) -> list[dict]: - if boundary_timestep is not None: - if timestep_val >= boundary_timestep: - return kv_cache1 - else: - assert kv_cache2 is not None, "kv_cache2 is not initialized" - return kv_cache2 - return kv_cache1 - crossattn_cache = self._initialize_crossattn_cache( batch_size=latents.shape[0], max_text_len=257, # 1 CLS + 256 patch tokens dtype=target_dtype, device=latents.device) - pos_start_base = 0 - if t % self.num_frame_per_block != 0: raise ValueError( "num_frames must be divisible by num_frame_per_block for causal denoising" @@ -184,211 +215,67 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): # The first frame information is already encoded in batch.image_latent (cond_concat) # and will be used by the model via channel concatenation: torch.cat([x, cond_concat], dim=1) + ctx = BlockProcessingContext( + batch=batch, + block_idx=0, + start_index=0, + kv_cache1=kv_cache1, + kv_cache2=kv_cache2, + kv_cache_mouse=kv_cache_mouse, + kv_cache_keyboard=kv_cache_keyboard, + crossattn_cache=crossattn_cache, + timesteps=timesteps, + block_sizes=block_sizes, + noise_pool=None, + fastvideo_args=fastvideo_args, + target_dtype=target_dtype, + autocast_enabled=autocast_enabled, + boundary_timestep=boundary_timestep, + high_noise_timesteps=high_noise_timesteps, + context_noise=getattr(fastvideo_args.pipeline_config, + "context_noise", 0), + image_kwargs=image_kwargs, + pos_cond_kwargs=pos_cond_kwargs, + ) + + context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", + 0) + with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar: for block_idx, current_num_frames in enumerate(block_sizes): + ctx.block_idx = block_idx + ctx.start_index = start_index current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :] - noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) - _video_raw_latent_shape = noise_latents_btchw.shape # noqa: F841 - - # NOTE: crossattn_cache should NOT be reset between blocks! action_kwargs = self._prepare_action_kwargs( batch, start_index, current_num_frames) - for i, t_cur in enumerate(timesteps): - if boundary_timestep is not None and t_cur < 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(target_dtype) - - 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(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) - - 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( - ) - attn_metadata = self.attn_metadata_builder.build( - current_timestep=i, - raw_latent_shape=(current_num_frames, h, w), - patch_size=fastvideo_args.pipeline_config. - dit_config.patch_size, - STA_param=batch.STA_param, - VSA_sparsity=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=target_dtype, - enabled=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": _get_kv_cache(t_cur), - "crossattn_cache": crossattn_cache, - "current_start": (pos_start_base + 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": - kv_cache_mouse, - "kv_cache_keyboard": - kv_cache_keyboard, - }) - model_kwargs.update(action_kwargs) - - pred_noise_btchw = current_model( - latent_model_input, - prompt_embeds, - t_expanded_noise, - **image_kwargs, - **pos_cond_kwargs, - **model_kwargs, - ).permute(0, 2, 1, 3, 4) - - if boundary_timestep is not None and t_cur >= 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) * - 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) - 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 boundary_timestep is not None and high_noise_timesteps is not None and i < len( - 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) * - boundary_timestep).unflatten( - 0, pred_video_btchw.shape[:2]) - elif boundary_timestep is not None and high_noise_timesteps is not None and i == len( - 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) - - if progress_bar is not None: - progress_bar.update() + current_latents = self._process_single_block( + current_latents=current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + timesteps=timesteps, + ctx=ctx, + action_kwargs=action_kwargs, + progress_bar=progress_bar, + ) latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents - context_noise = getattr(fastvideo_args.pipeline_config, - "context_noise", 0) - # Expand context timestep to per-frame format [batch, num_frames] for causal model - t_context = torch.ones([latents.shape[0], current_num_frames], - device=latents.device, - dtype=torch.long) * int(context_noise) - context_bcthw = current_latents.to(target_dtype) - - with torch.autocast(device_type="cuda", - dtype=target_dtype, - enabled=autocast_enabled), \ - set_forward_context(current_timestep=0, - attn_metadata=attn_metadata, - forward_batch=batch): - - context_model_kwargs = { - "kv_cache": kv_cache1, - "crossattn_cache": crossattn_cache, - "current_start": - (pos_start_base + start_index) * self.frame_seq_length, - "start_frame": start_index, - } - - if self.use_action_module: - context_model_kwargs.update({ - "kv_cache_mouse": - kv_cache_mouse, - "kv_cache_keyboard": - kv_cache_keyboard, - }) - context_model_kwargs.update(action_kwargs) - - if boundary_timestep is not None and self.transformer_2 is not None: - self.transformer_2( - context_bcthw, - prompt_embeds, - t_context, - kv_cache=kv_cache2, - crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index) * - self.frame_seq_length, - start_frame=start_index, - **image_kwargs, - **pos_cond_kwargs, - ) - - self.transformer( - context_bcthw, - prompt_embeds, - t_context, - **image_kwargs, - **pos_cond_kwargs, - **context_model_kwargs, - ) + # Update KV caches with clean context + self._update_context_cache( + current_latents=current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + ctx=ctx, + action_kwargs=action_kwargs, + context_noise=context_noise, + ) start_index += current_num_frames @@ -542,6 +429,440 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): }) return crossattn_cache + def _process_single_block( + self, + current_latents: torch.Tensor, + batch: ForwardBatch, + start_index: int, + current_num_frames: int, + timesteps: torch.Tensor, + ctx: BlockProcessingContext, + action_kwargs: dict[str, Any], + noise_generator: Callable[[tuple, torch.dtype, int], torch.Tensor] + | 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) + + if progress_bar is not None: + progress_bar.update() + + return current_latents + + def _update_context_cache( + self, + current_latents: torch.Tensor, + batch: ForwardBatch, + start_index: int, + current_num_frames: int, + ctx: BlockProcessingContext, + action_kwargs: dict[str, Any], + context_noise: float, + ) -> None: + prompt_embeds = batch.prompt_embeds + latents_device = current_latents.device + + # Expand context timestep to per-frame format [batch, num_frames] for causal model + t_context = torch.ones([current_latents.shape[0], current_num_frames], + device=latents_device, + dtype=torch.long) * int(context_noise) + context_bcthw = current_latents.to(ctx.target_dtype) + + with torch.autocast(device_type="cuda", + dtype=ctx.target_dtype, + enabled=ctx.autocast_enabled), \ + set_forward_context(current_timestep=0, + attn_metadata=None, + forward_batch=batch): + + context_model_kwargs = { + "kv_cache": ctx.kv_cache1, + "crossattn_cache": ctx.crossattn_cache, + "current_start": start_index * self.frame_seq_length, + "start_frame": start_index, + } + + if self.use_action_module: + context_model_kwargs.update({ + "kv_cache_mouse": + ctx.kv_cache_mouse, + "kv_cache_keyboard": + ctx.kv_cache_keyboard, + }) + context_model_kwargs.update(action_kwargs) + + if ctx.boundary_timestep is not None and self.transformer_2 is not None: + self.transformer_2( + context_bcthw, + prompt_embeds, + t_context, + kv_cache=ctx.kv_cache2, + crossattn_cache=ctx.crossattn_cache, + current_start=start_index * self.frame_seq_length, + start_frame=start_index, + **ctx.image_kwargs, + **ctx.pos_cond_kwargs, + ) + + self.transformer( + context_bcthw, + prompt_embeds, + t_context, + **ctx.image_kwargs, + **ctx.pos_cond_kwargs, + **context_model_kwargs, + ) + + def streaming_reset(self, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> ForwardBatch: + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32 + ) and not fastvideo_args.disable_autocast + + latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] + patch_size = self.transformer.patch_size + patch_ratio = patch_size[-1] * patch_size[-2] + self.frame_seq_length = latent_seq_length // patch_ratio + + timesteps = torch.tensor( + fastvideo_args.pipeline_config.dmd_denoising_steps, + dtype=torch.long).cpu() + if fastvideo_args.pipeline_config.warp_denoising_step: + scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), + torch.tensor([0], + dtype=torch.float32))) + timesteps = scheduler_timesteps[1000 - timesteps] + timesteps = timesteps.to(get_local_torch_device()) + + boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, + 'boundary_ratio', None) + if boundary_ratio is not None: + boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps + high_noise_timesteps = timesteps[timesteps >= boundary_timestep] + else: + boundary_timestep = None + high_noise_timesteps = None + + image_embeds = batch.image_embeds + if len(image_embeds) > 0: + assert torch.isnan(image_embeds[0]).sum() == 0 + image_embeds = [ + image_embed.to(target_dtype) for image_embed in image_embeds + ] + + # directly set the kwarg. + image_kwargs = {"encoder_hidden_states_image": image_embeds} + pos_cond_kwargs: dict[str, Any] = {} + + if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: + self.prepare_sta_param(batch, fastvideo_args) + + assert batch.latents is not None, "latents must be provided" + latents = batch.latents + b, c, t, h, w = latents.shape + prompt_embeds = batch.prompt_embeds + assert torch.isnan(prompt_embeds[0]).sum() == 0 + + # Initialize caches + kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + kv_cache2 = None + if boundary_timestep is not None: + kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + kv_cache_mouse = None + kv_cache_keyboard = None + if self.use_action_module: + kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache( + batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + crossattn_cache = self._initialize_crossattn_cache( + batch_size=latents.shape[0], + max_text_len=257, # 1 CLS + 256 patch tokens + dtype=target_dtype, + device=latents.device) + + # Calculate block sizes + if t % self.num_frame_per_block != 0: + raise ValueError( + "num_frames must be divisible by num_frame_per_block for causal denoising" + ) + num_blocks = t // self.num_frame_per_block + block_sizes = [self.num_frame_per_block] * num_blocks + if boundary_timestep is not None: + block_sizes[0] = 1 + + # Pre-allocate noise pool + num_denoising_steps = len(timesteps) + noise_shape = (b, self.num_frame_per_block, c, h, w) + noise_pool = [ + torch.randn( + noise_shape, + dtype=target_dtype, + device=latents.device, + ) for _ in range(num_denoising_steps - 1) + ] + + # Create and store context + self._streaming_ctx = BlockProcessingContext( + batch=batch, + block_idx=0, + start_index=0, + kv_cache1=kv_cache1, + kv_cache2=kv_cache2, + kv_cache_mouse=kv_cache_mouse, + kv_cache_keyboard=kv_cache_keyboard, + crossattn_cache=crossattn_cache, + timesteps=timesteps, + block_sizes=block_sizes, + noise_pool=noise_pool, + fastvideo_args=fastvideo_args, + target_dtype=target_dtype, + autocast_enabled=autocast_enabled, + boundary_timestep=boundary_timestep, + high_noise_timesteps=high_noise_timesteps, + context_noise=getattr(fastvideo_args.pipeline_config, + "context_noise", 0), + image_kwargs=image_kwargs, + pos_cond_kwargs=pos_cond_kwargs, + ) + + self._streaming_initialized = True + return batch + + def streaming_step( + self, + keyboard_action: torch.Tensor | None = None, + mouse_action: torch.Tensor | None = None) -> ForwardBatch: + if not self._streaming_initialized or self._streaming_ctx is None: + raise RuntimeError( + "Streaming not initialized! Call streaming_reset first.") + + ctx = self._streaming_ctx + if ctx.block_idx >= len(ctx.block_sizes): + return ctx.batch + + batch = ctx.batch + latents = batch.latents + assert latents is not None, "latents must be set in batch" + + 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 for this block + 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._prepare_action_kwargs(batch, start_index, + current_num_frames) + + # Create noise generator that uses pre-allocated noise pool + def streaming_noise_generator(shape: tuple, dtype: torch.dtype, + step_idx: int) -> torch.Tensor: + 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: + # Fallback to dynamic allocation if pool not available + return torch.randn( + shape, + dtype=dtype, + generator=(batch.generator[0] if isinstance( + batch.generator, list) else batch.generator)).to( + latents.device) + + current_latents = self._process_single_block( + current_latents=current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + timesteps=ctx.timesteps, + ctx=ctx, + action_kwargs=action_kwargs, + noise_generator=streaming_noise_generator, + ) + + latents[:, :, start_index:start_index + + current_num_frames, :, :] = current_latents + + # Update KV caches with clean context + self._update_context_cache( + current_latents=current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + ctx=ctx, + action_kwargs=action_kwargs, + context_noise=ctx.context_noise, + ) + + # Advance streaming state + ctx.start_index += current_num_frames + ctx.block_idx += 1 + + return batch + + def streaming_clear(self) -> None: + self._streaming_initialized = False + self._streaming_ctx = None + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() diff --git a/fastvideo/worker/gpu_worker.py b/fastvideo/worker/gpu_worker.py index fe5832aac..0460a97e5 100644 --- a/fastvideo/worker/gpu_worker.py +++ b/fastvideo/worker/gpu_worker.py @@ -114,6 +114,20 @@ class Worker: return {"status": "lora_adapter_unmerged"} return {"status": "failed: pipeline is not a LoRAPipeline"} + def execute_streaming_reset( + self, forward_batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> dict[str, Any]: + self.pipeline.streaming_reset(forward_batch, self.fastvideo_args) + return {"status": "reset_complete"} + + def execute_streaming_step(self, keyboard_action: torch.Tensor, + mouse_action: torch.Tensor) -> ForwardBatch: + return self.pipeline.streaming_step(keyboard_action, mouse_action) + + def execute_streaming_clear(self) -> dict[str, Any]: + self.pipeline.streaming_clear() + return {"status": "cleared"} + def merge_lora_weights(self) -> dict[str, Any]: if isinstance(self.pipeline, LoRAPipeline): self.pipeline.merge_lora_weights() diff --git a/fastvideo/worker/multiproc_executor.py b/fastvideo/worker/multiproc_executor.py index 5c062e3ed..2e3268334 100644 --- a/fastvideo/worker/multiproc_executor.py +++ b/fastvideo/worker/multiproc_executor.py @@ -1,11 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import asyncio import atexit import contextlib from dataclasses import dataclass +from enum import Enum import faulthandler import multiprocessing as mp from multiprocessing.connection import Connection +from multiprocessing.queues import Queue import os +import queue import signal import time from collections.abc import Callable @@ -13,19 +19,55 @@ from multiprocessing.process import BaseProcess from typing import Any, cast import psutil +import torch from fastvideo.distributed.parallel_state import get_dp_group, get_tp_group import fastvideo.envs as envs from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger from fastvideo.pipelines.pipeline_batch_info import ForwardBatch -from fastvideo.utils import decorate_logs, get_distributed_init_method, get_exception_traceback, get_loopback_ip, get_mp_context, get_open_port, kill_itself_when_parent_died, force_spawn +from fastvideo.utils import (decorate_logs, get_distributed_init_method, + get_exception_traceback, get_loopback_ip, + get_mp_context, get_open_port, + kill_itself_when_parent_died, force_spawn) from fastvideo.worker.executor import Executor from fastvideo.worker.worker_base import WorkerWrapperBase 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" + + +@dataclass +class StreamingTask: + """Task submitted to worker via input queue.""" + task_type: StreamingTaskType + # For STEP tasks: + keyboard_action: torch.Tensor | None = None + mouse_action: torch.Tensor | None = None + # For RESET tasks: + batch: ForwardBatch | None = None + fastvideo_args: FastVideoArgs | None = None + + +@dataclass +class StreamingResult: + """Result returned from worker via output queue.""" + task_type: StreamingTaskType + output_batch: ForwardBatch | None = None + error: Exception | None = None + + class MultiprocExecutor(Executor): def _init_executor(self) -> None: @@ -40,6 +82,12 @@ class MultiprocExecutor(Executor): get_loopback_ip(), master_port) logger.info("Use master port: %s", master_port) + # Create streaming queues BEFORE spawning workers + ctx = get_mp_context() + self._streaming_input_queue: Queue | None = ctx.Queue() + self._streaming_output_queue: Queue | None = ctx.Queue() + self._streaming_enabled = False + unready_workers: list[UnreadyWorkerProcHandle] = [] success = False try: @@ -50,6 +98,8 @@ class MultiprocExecutor(Executor): local_rank=rank, rank=rank, distributed_init_method=distributed_init_method, + streaming_input_queue=self._streaming_input_queue, + streaming_output_queue=self._streaming_output_queue, )) # Workers must be created before wait_for_ready to avoid @@ -88,6 +138,108 @@ class MultiprocExecutor(Executor): return result_batch + def execute_streaming_reset( + self, forward_batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> dict[str, Any]: + responses = self.collective_rpc("execute_streaming_reset", + kwargs={ + "forward_batch": forward_batch, + "fastvideo_args": fastvideo_args, + }) + return responses[0] + + def execute_streaming_step(self, keyboard_action: Any, + mouse_action: Any) -> ForwardBatch: + responses = self.collective_rpc("execute_streaming_step", + kwargs={ + "keyboard_action": keyboard_action, + "mouse_action": mouse_action, + }) + return responses[0] + + async def execute_streaming_step_async(self, keyboard_action: Any, + mouse_action: Any) -> ForwardBatch: + responses = await self.collective_rpc_async("execute_streaming_step", + kwargs={ + "keyboard_action": + keyboard_action, + "mouse_action": + mouse_action, + }) + return responses[0] + + def execute_streaming_clear(self) -> dict[str, Any]: + responses = self.collective_rpc("execute_streaming_clear") + return responses[0] + + def enable_streaming(self) -> None: + if self._streaming_enabled: + return + + self.collective_rpc("start_streaming_queue_loop") + self._streaming_enabled = True + + def disable_streaming(self) -> None: + if not self._streaming_enabled: + return + + if self._streaming_input_queue is not None: + self._streaming_input_queue.put( + StreamingTask(task_type=StreamingTaskType.EXIT)) + self._streaming_enabled = False + + self._streaming_input_queue = None + self._streaming_output_queue = None + + def submit_reset(self, forward_batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> None: + if not self._streaming_enabled: + self.enable_streaming() + + self._streaming_input_queue.put( + StreamingTask( + task_type=StreamingTaskType.RESET, + batch=forward_batch, + fastvideo_args=fastvideo_args, + )) + + def submit_step(self, 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.STEP, + keyboard_action=keyboard_action, + mouse_action=mouse_action, + )) + + def submit_clear(self) -> None: + if self._streaming_enabled and self._streaming_input_queue is not None: + self._streaming_input_queue.put( + StreamingTask(task_type=StreamingTaskType.CLEAR)) + + def get_result(self, + timeout: float | None = None) -> StreamingResult | None: + if not self._streaming_enabled or self._streaming_output_queue is None: + return None + + try: + if timeout == 0: + return self._streaming_output_queue.get_nowait() + else: + return self._streaming_output_queue.get(timeout=timeout) + except queue.Empty: + return None + + def wait_result(self) -> StreamingResult: + if not self._streaming_enabled or self._streaming_output_queue is None: + raise RuntimeError("Streaming mode not enabled.") + + return self._streaming_output_queue.get() + def set_lora_adapter(self, lora_nickname: str, lora_path: str | None = None) -> None: @@ -147,6 +299,24 @@ class MultiprocExecutor(Executor): except Exception as e: raise e + async def collective_rpc_async(self, + method: str | Callable, + timeout: float | None = None, + args: tuple = (), + kwargs: dict | None = None) -> list[Any]: + kwargs = kwargs or {} + loop = asyncio.get_running_loop() + + for worker in self.workers: + worker.pipe.send({"method": method, "args": args, "kwargs": kwargs}) + + async def recv_from_worker(worker: WorkerProcHandle) -> Any: + return await loop.run_in_executor(None, worker.pipe.recv) + + responses = await asyncio.gather( + *[recv_from_worker(worker) for worker in self.workers]) + return list(responses) + def shutdown(self) -> None: """Properly shut down the executor and its workers""" if hasattr(self, 'shutting_down') and self.shutting_down: @@ -266,7 +436,7 @@ class WorkerProcHandle: @classmethod def from_unready_handle( - cls, unready_handle: UnreadyWorkerProcHandle) -> "WorkerProcHandle": + cls, unready_handle: UnreadyWorkerProcHandle) -> WorkerProcHandle: return cls( proc=unready_handle.proc, rank=unready_handle.rank, @@ -286,9 +456,13 @@ class WorkerMultiprocProc: rank: int, distributed_init_method: str, pipe: Connection, + streaming_input_queue: Queue | None = None, + streaming_output_queue: Queue | None = None, ): self.rank = rank self.pipe = pipe + self.streaming_input_queue = streaming_input_queue + self.streaming_output_queue = streaming_output_queue wrapper = WorkerWrapperBase(fastvideo_args=fastvideo_args, rpc_rank=rank) @@ -314,6 +488,8 @@ class WorkerMultiprocProc: local_rank: int, rank: int, distributed_init_method: str, + streaming_input_queue: Queue | None = None, + streaming_output_queue: Queue | None = None, ) -> UnreadyWorkerProcHandle: context = get_mp_context() executor_pipe, worker_pipe = context.Pipe(duplex=True) @@ -326,6 +502,8 @@ class WorkerMultiprocProc: "distributed_init_method": distributed_init_method, "pipe": worker_pipe, "ready_pipe": writer, + "streaming_input_queue": streaming_input_queue, + "streaming_output_queue": streaming_output_queue, } # Run EngineCore busy loop in background process. proc = context.Process(target=WorkerMultiprocProc.worker_main, @@ -455,6 +633,11 @@ class WorkerMultiprocProc: with contextlib.suppress(Exception): self.pipe.send(response) break + if method == "start_streaming_queue_loop": + self.pipe.send( + {"status": "streaming_queue_loop_started"}) + self.streaming_queue_loop() + continue if method == 'execute_forward': forward_batch = kwargs['forward_batch'] fastvideo_args = kwargs['fastvideo_args'] @@ -488,6 +671,50 @@ class WorkerMultiprocProc: self.rank, str(e)) continue + def streaming_queue_loop(self) -> None: + 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 + + while True: + try: + task: StreamingTask = self.streaming_input_queue.get() + + if task.task_type == StreamingTaskType.EXIT: + break + elif task.task_type == StreamingTaskType.RESET: + try: + self.worker.execute_streaming_reset( + task.batch, task.fastvideo_args) + self.streaming_output_queue.put( + StreamingResult(task_type=StreamingTaskType.RESET)) + except Exception as e: + logger.error("Worker %d reset error: %s", self.rank, e) + self.streaming_output_queue.put( + StreamingResult(task_type=StreamingTaskType.RESET, + error=e)) + elif task.task_type == StreamingTaskType.STEP: + try: + batch = self.worker.execute_streaming_step( + task.keyboard_action, task.mouse_action) + self.streaming_output_queue.put( + StreamingResult(task_type=StreamingTaskType.STEP, + output_batch=batch)) + except Exception as e: + logger.error("Worker %d step error: %s", self.rank, e) + self.streaming_output_queue.put( + StreamingResult(task_type=StreamingTaskType.STEP, + error=e)) + elif task.task_type == StreamingTaskType.CLEAR: + self.worker.execute_streaming_clear() + self.streaming_output_queue.put( + StreamingResult(task_type=StreamingTaskType.CLEAR)) + 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)) + @staticmethod def setup_proc_title_and_log_prefix() -> None: dp_size = get_dp_group().world_size diff --git a/fastvideo/worker/ray_distributed_executor.py b/fastvideo/worker/ray_distributed_executor.py index 8468edfba..daf31d4b7 100644 --- a/fastvideo/worker/ray_distributed_executor.py +++ b/fastvideo/worker/ray_distributed_executor.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Adapt from https://github.com/vllm-project/vllm/blob/releases/v0.11.0/vllm/executor/ray_distributed_executor.py +import asyncio from collections import defaultdict import os import cloudpickle @@ -269,6 +270,47 @@ class RayDistributedExecutor(Executor): else: self.non_driver_workers.append(worker) + def execute_streaming_reset( + self, forward_batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> dict[str, Any]: + responses: list[dict[str, Any]] = self.collective_rpc( + "execute_streaming_reset", + kwargs={ + "forward_batch": forward_batch, + "fastvideo_args": fastvideo_args, + }, + ) + return responses[0] + + def execute_streaming_step(self, + keyboard_action=None, + mouse_action=None) -> ForwardBatch: + responses: list[ForwardBatch] = self.collective_rpc( + "execute_streaming_step", + kwargs={ + "keyboard_action": keyboard_action, + "mouse_action": mouse_action, + }, + ) + return responses[0] + + async def execute_streaming_step_async(self, + keyboard_action=None, + mouse_action=None) -> ForwardBatch: + kwargs = { + "keyboard_action": keyboard_action, + "mouse_action": mouse_action, + } + futures = [ + w.execute_method.remote("execute_streaming_step", **kwargs) + for w in self.workers + ] + responses = await asyncio.gather(*futures) + return responses[0] + + def execute_streaming_clear(self) -> None: + self.collective_rpc("execute_streaming_clear") + def execute_forward(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: responses: list[ForwardBatch] = self.collective_rpc(