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.api.sampling_param 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()