diff --git a/fastvideo/models/schedulers/denoising_schedule.py b/fastvideo/models/schedulers/denoising_schedule.py deleted file mode 100644 index ab6c660dc..000000000 --- a/fastvideo/models/schedulers/denoising_schedule.py +++ /dev/null @@ -1,124 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from collections.abc import Iterable -from typing import Any - -import torch - - -def resolve_denoising_steps( - raw_steps: Iterable[int] | None, - scheduler_timesteps: torch.Tensor, - warp_denoising_step: bool, - *, - device: torch.device | None = None, -) -> torch.Tensor: - """Convert configured denoising steps into runtime scheduler timesteps.""" - if raw_steps is None: - raise ValueError("raw_steps must be provided") - - steps = torch.tensor(list(raw_steps), dtype=torch.long) - if warp_denoising_step: - timesteps = torch.cat( - ( - scheduler_timesteps.detach().cpu(), - torch.tensor([0], dtype=torch.float32), - ) - ) - steps = timesteps[1000 - steps] - - if device is not None: - steps = steps.to(device) - return steps - - -def build_block_denoising_steps( - base_steps: torch.Tensor, - block_index: int, - use_diagonal_denoising: bool, - warmup_mid_steps: torch.Tensor | None = None, -) -> torch.Tensor: - """ - Build the per-block denoising schedule for diagonal denoising. - - - block 0: first + 2 mids + last - - block 1: first + 1 mid + last - - block >= 2: first + last - """ - if (not use_diagonal_denoising) or base_steps.numel() <= 1: - return base_steps - - first = base_steps[0:1] - last = base_steps[-1:] - mids = base_steps[1:-1] - - if warmup_mid_steps is not None and warmup_mid_steps.numel() > 0: - mids = warmup_mid_steps.to( - device=base_steps.device, - dtype=base_steps.dtype, - ) - - if block_index <= 0: - if mids.numel() >= 2: - selected_mid = mids[:2] - elif mids.numel() == 1: - selected_mid = mids - else: - selected_mid = torch.empty( - 0, - dtype=base_steps.dtype, - device=base_steps.device, - ) - return torch.cat([first, selected_mid, last], dim=0) - - if block_index == 1: - if mids.numel() >= 1: - return torch.cat([first, mids[:1], last], dim=0) - return torch.cat([first, last], dim=0) - - return torch.cat([first, last], dim=0) - - -def resolve_base_denoising_timesteps( - pipeline_config: Any, - scheduler_timesteps: torch.Tensor, - *, - device: torch.device | None = None, -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Resolve the base denoising steps and optional diagonal warmup mids.""" - timesteps = resolve_denoising_steps( - pipeline_config.dmd_denoising_steps, - scheduler_timesteps, - getattr(pipeline_config, "warp_denoising_step", False), - device=device, - ) - warmup_mid_steps = None - raw_warmup_mid_steps = getattr( - pipeline_config, - "diagonal_warmup_mid_steps", - None, - ) - if raw_warmup_mid_steps: - warmup_mid_steps = resolve_denoising_steps( - raw_warmup_mid_steps, - scheduler_timesteps, - getattr(pipeline_config, "warp_denoising_step", False), - device=device, - ) - return timesteps, warmup_mid_steps - - -def resolve_block_denoising_timesteps( - base_timesteps: torch.Tensor, - block_idx: int, - use_diagonal_denoising: bool, - diagonal_warmup_mid_steps: torch.Tensor | None, -) -> torch.Tensor: - """Resolve the per-block denoising timesteps for causal rollout.""" - return build_block_denoising_steps( - base_steps=base_timesteps, - block_index=block_idx, - use_diagonal_denoising=use_diagonal_denoising, - warmup_mid_steps=diagonal_warmup_mid_steps, - ) diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 835b48ace..16beebfc1 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -53,8 +53,7 @@ class PipelineLoggingInfo: def get_total_execution_time(self) -> float: """Get total pipeline execution time.""" - return sum( - stage.get('execution_time', 0) for stage in self.stages.values()) + return sum(stage.get('execution_time', 0) for stage in self.stages.values()) @dataclass @@ -95,8 +94,7 @@ class ForwardBatch: t_thresh: float = 0.5 spatial_refine_only: bool = False num_cond_frames: int = 0 - stage1_video: list[ - PIL.Image.Image] | None = None # Loaded frames from refine_from + stage1_video: list[PIL.Image.Image] | None = None # Loaded frames from refine_from # Primary encoder embeddings prompt_embeds: list[torch.Tensor] = field(default_factory=list) @@ -140,8 +138,7 @@ class ForwardBatch: gt_latents: torch.Tensor | None = None # Ground truth latents for conditioning [B, 16, T, H, W] conditioning_mask: torch.Tensor | None = None # Mask for conditioning [B, 1, T, H, W] camera_trajectory: str | None = None # Camera trajectory file/identifier - action_list: list[ - str] | None = None # List of actions (e.g., ['forward', 'left']) + action_list: list[str] | None = None # List of actions (e.g., ['forward', 'left']) action_speed_list: list[float] | None = None # Speed for each action # Camera control inputs (LingBotWorld) c2ws_plucker_emb: torch.Tensor | None = None # Plucker embedding: [B, C, F_lat, H_lat, W_lat] @@ -214,15 +211,13 @@ class ForwardBatch: VSA_sparsity: float = 0.0 # Logging info - logging_info: PipelineLoggingInfo = field( - default_factory=PipelineLoggingInfo) + logging_info: PipelineLoggingInfo = field(default_factory=PipelineLoggingInfo) def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" # Enable CFG for standard guidance_scale and LTX-2 text CFG scales. - ltx2_text_cfg_enabled = (self.ltx2_cfg_scale_video != 1.0 - or self.ltx2_cfg_scale_audio != 1.0) + ltx2_text_cfg_enabled = (self.ltx2_cfg_scale_video != 1.0 or self.ltx2_cfg_scale_audio != 1.0) if self.guidance_scale > 1.0 or ltx2_text_cfg_enabled: self.do_classifier_free_guidance = True if self.negative_prompt_embeds is None: @@ -296,11 +291,9 @@ class TrainingBatch: dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict) latent_vis_dict: dict[str, Any] = field(default_factory=dict) fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict) - denoised_timestep_from: int | None = None - denoised_timestep_to: int | None = None @dataclass class PreprocessBatch(ForwardBatch): video_loader: list["VideoDecoder"] | list[str] = field(default_factory=list) - video_file_name: list[str] = field(default_factory=list) + video_file_name: list[str] = field(default_factory=list) \ No newline at end of file diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py index 26cc1de65..8547e8df9 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/fastvideo/pipelines/stages/input_validation.py @@ -12,8 +12,7 @@ from fastvideo.logger import init_logger from fastvideo.models.vision_utils import load_image, load_video, pil_to_numpy, numpy_to_pt, normalize, resize from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.stages.validators import (StageValidators, - VerificationResult) +from fastvideo.pipelines.stages.validators import (StageValidators, VerificationResult) from fastvideo.utils import best_output_size logger = init_logger(__name__) @@ -30,8 +29,7 @@ class InputValidationStage(PipelineStage): before proceeding with the diffusion process. """ - def _generate_seeds(self, batch: ForwardBatch, - fastvideo_args: FastVideoArgs): + def _generate_seeds(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs): """Generate seeds for the inference""" seed = batch.seed num_videos_per_prompt = batch.num_videos_per_prompt @@ -41,9 +39,7 @@ class InputValidationStage(PipelineStage): batch.seeds = seeds # Peiyuan: using GPU seed will cause A100 and H100 to generate different results... - batch.generator = [ - torch.Generator("cpu").manual_seed(seed) for seed in seeds - ] + batch.generator = [torch.Generator("cpu").manual_seed(seed) for seed in seeds] def forward( self, @@ -65,37 +61,27 @@ class InputValidationStage(PipelineStage): # Ensure prompt is properly formatted if batch.prompt is None and batch.prompt_embeds is None: - raise ValueError( - "Either `prompt` or `prompt_embeds` must be provided") + raise ValueError("Either `prompt` or `prompt_embeds` must be provided") # Ensure negative prompt is properly formatted if using classifier-free guidance if (batch.do_classifier_free_guidance and batch.negative_prompt is None and batch.negative_prompt_embeds is None): - raise ValueError( - "For classifier-free guidance, either `negative_prompt` or " - "`negative_prompt_embeds` must be provided") + raise ValueError("For classifier-free guidance, either `negative_prompt` or " + "`negative_prompt_embeds` must be provided") # Validate height and width if batch.height is None or batch.width is None: - raise ValueError( - "Height and width must be provided. Please set `height` and `width`." - ) + raise ValueError("Height and width must be provided. Please set `height` and `width`.") if batch.height % 8 != 0 or batch.width % 8 != 0: - raise ValueError( - f"Height and width must be divisible by 8 but are {batch.height} and {batch.width}." - ) + raise ValueError(f"Height and width must be divisible by 8 but are {batch.height} and {batch.width}.") # Validate number of inference steps if batch.num_inference_steps <= 0: - raise ValueError( - f"Number of inference steps must be positive, but got {batch.num_inference_steps}" - ) + raise ValueError(f"Number of inference steps must be positive, but got {batch.num_inference_steps}") # Validate guidance scale if using classifier-free guidance if batch.do_classifier_free_guidance and batch.guidance_scale <= 0: - raise ValueError( - f"Guidance scale must be positive, but got {batch.guidance_scale}" - ) + raise ValueError(f"Guidance scale must be positive, but got {batch.guidance_scale}") # for i2v, get image from image_path # @TODO(Wei) hard-coded for wan2.2 5b ti2v for now. Should put this in image_encoding stage @@ -108,8 +94,7 @@ class InputValidationStage(PipelineStage): # further processing for ti2v task if (fastvideo_args.pipeline_config.ti2v_task - or fastvideo_args.pipeline_config.is_causal - ) and batch.pil_image is not None: + or fastvideo_args.pipeline_config.is_causal) and batch.pil_image is not None: img = batch.pil_image ih, iw = img.height, img.width @@ -126,8 +111,7 @@ class InputValidationStage(PipelineStage): ow, oh = best_output_size(iw, ih, dw, dh, max_area) scale = max(ow / iw, oh / ih) - img = img.resize((round(iw * scale), round(ih * scale)), - Image.LANCZOS) + img = img.resize((round(iw * scale), round(ih * scale)), Image.LANCZOS) # center-crop x1 = (img.width - ow) // 2 @@ -135,12 +119,10 @@ class InputValidationStage(PipelineStage): img = img.crop((x1, y1, x1 + ow, y1 + oh)) assert img.width == ow and img.height == oh - logger.info("final processed img height: %s, img width: %s", - img.height, img.width) + logger.info("final processed img height: %s, img width: %s", img.height, img.width) # to tensor - img = TF.to_tensor(img).sub_(0.5).div_(0.5).to( - self.device).unsqueeze(1) + img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device).unsqueeze(1) img = img.unsqueeze(0) batch.height = oh batch.width = ow @@ -148,10 +130,8 @@ class InputValidationStage(PipelineStage): # for v2v, get control video from video path if batch.video_path is not None: - pil_images, original_fps = load_video(batch.video_path, - return_fps=True) - logger.info("Loaded video with %s frames, original FPS: %s", - len(pil_images), original_fps) + pil_images, original_fps = load_video(batch.video_path, return_fps=True) + logger.info("Loaded video with %s frames, original FPS: %s", len(pil_images), original_fps) # Get target parameters from batch target_fps = batch.fps @@ -164,25 +144,18 @@ class InputValidationStage(PipelineStage): if frame_skip > 1: pil_images = pil_images[::frame_skip] effective_fps = original_fps / frame_skip - logger.info( - "Resampled video from %.1f fps to %.1f fps (skip=%s)", - original_fps, effective_fps, frame_skip) + logger.info("Resampled video from %.1f fps to %.1f fps (skip=%s)", original_fps, effective_fps, + frame_skip) # Limit to target number of frames - if target_num_frames is not None and len( - pil_images) > target_num_frames: + if target_num_frames is not None and len(pil_images) > target_num_frames: pil_images = pil_images[:target_num_frames] - logger.info("Limited video to %s frames (from %s total)", - target_num_frames, len(pil_images)) + logger.info("Limited video to %s frames (from %s total)", target_num_frames, len(pil_images)) # Resize each PIL image to target dimensions resized_images = [] for pil_img in pil_images: - resized_img = resize(pil_img, - target_height, - target_width, - resize_mode="default", - resample="lanczos") + resized_img = resize(pil_img, target_height, target_width, resize_mode="default", resample="lanczos") resized_images.append(resized_img) # Convert PIL images to numpy array @@ -198,63 +171,47 @@ class InputValidationStage(PipelineStage): # Validate action control inputs (Matrix-Game) if batch.mouse_cond is not None: if batch.mouse_cond.dim() != 3 or batch.mouse_cond.shape[-1] != 2: - raise ValueError( - f"mouse_cond must have shape (B, T, 2), but got {batch.mouse_cond.shape}" - ) - logger.info("Action control: mouse_cond validated - shape %s", - batch.mouse_cond.shape) + raise ValueError(f"mouse_cond must have shape (B, T, 2), but got {batch.mouse_cond.shape}") + logger.info("Action control: mouse_cond validated - shape %s", batch.mouse_cond.shape) if batch.keyboard_cond is not None: if batch.keyboard_cond.dim() != 3: - raise ValueError( - f"keyboard_cond must have 3 dimensions (B, T, K), but got {batch.keyboard_cond.dim()}" - ) + raise ValueError(f"keyboard_cond must have 3 dimensions (B, T, K), but got {batch.keyboard_cond.dim()}") keyboard_dim = batch.keyboard_cond.shape[-1] - logger.info( - "Action control: keyboard_cond validated - shape %s (dim=%d)", - batch.keyboard_cond.shape, keyboard_dim) + logger.info("Action control: keyboard_cond validated - shape %s (dim=%d)", batch.keyboard_cond.shape, + keyboard_dim) if batch.grid_sizes is not None: if not isinstance(batch.grid_sizes, list | tuple | torch.Tensor): raise ValueError("grid_sizes must be a list, tuple, or tensor") if isinstance(batch.grid_sizes, torch.Tensor): if batch.grid_sizes.numel() != 3: - raise ValueError( - "grid_sizes must have 3 elements [F, H, W]") + raise ValueError("grid_sizes must have 3 elements [F, H, W]") else: if len(batch.grid_sizes) != 3: - raise ValueError( - "grid_sizes must have 3 elements [F, H, W]") - logger.info("Action control: grid_sizes validated - %s", - batch.grid_sizes) + raise ValueError("grid_sizes must have 3 elements [F, H, W]") + logger.info("Action control: grid_sizes validated - %s", batch.grid_sizes) return batch - def verify_input(self, batch: ForwardBatch, - fastvideo_args: FastVideoArgs) -> VerificationResult: + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: """Verify input validation stage inputs.""" result = VerificationResult() # Cosmos-Predict2.5 default seed is 0; allow non-negative seeds here. result.add_check("seed", batch.seed, [V.not_none, V.non_negative_int]) - result.add_check("num_videos_per_prompt", batch.num_videos_per_prompt, - V.positive_int) - result.add_check( - "prompt_or_embeds", None, lambda _: V.string_or_list_strings( - batch.prompt) or V.list_not_empty(batch.prompt_embeds)) + result.add_check("num_videos_per_prompt", batch.num_videos_per_prompt, V.positive_int) + result.add_check("prompt_or_embeds", None, + lambda _: V.string_or_list_strings(batch.prompt) or V.list_not_empty(batch.prompt_embeds)) result.add_check("height", batch.height, V.positive_int) result.add_check("width", batch.width, V.positive_int) - result.add_check("num_inference_steps", batch.num_inference_steps, - V.positive_int) - result.add_check( - "guidance_scale", batch.guidance_scale, lambda x: not batch. - do_classifier_free_guidance or V.positive_float(x)) + result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int) + result.add_check("guidance_scale", batch.guidance_scale, + lambda x: not batch.do_classifier_free_guidance or V.positive_float(x)) return result - def verify_output(self, batch: ForwardBatch, - fastvideo_args: FastVideoArgs) -> VerificationResult: + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: """Verify input validation stage outputs.""" result = VerificationResult() result.add_check("seeds", batch.seeds, V.list_not_empty) - result.add_check("generator", batch.generator, - V.generator_or_list_generators) - return result + result.add_check("generator", batch.generator, V.generator_or_list_generators) + return result \ No newline at end of file diff --git a/fastvideo/pipelines/stages/matrixgame_denoising.py b/fastvideo/pipelines/stages/matrixgame_denoising.py index 72ecfe72a..7e1496857 100644 --- a/fastvideo/pipelines/stages/matrixgame_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame_denoising.py @@ -9,8 +9,6 @@ from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger -from fastvideo.models.schedulers.denoising_schedule import ( - resolve_base_denoising_timesteps, resolve_block_denoising_timesteps) from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.denoising import DenoisingStage @@ -18,15 +16,12 @@ from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult try: - from fastvideo.attention.backends.video_sparse_attn import ( - VideoSparseAttentionBackend) + from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionBackend) vsa_available = True except ImportError: vsa_available = False VideoSparseAttentionBackend = None # type: ignore -from fastvideo.models.dits.matrixgame.kv_cache import KVCache - logger = init_logger(__name__) @@ -39,10 +34,10 @@ class BlockProcessingContext: block_idx: int start_index: int - kv_cache1: list[KVCache] - kv_cache2: list[KVCache] | None - kv_cache_mouse: list[KVCache] | None - kv_cache_keyboard: list[KVCache] | None + 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 @@ -53,14 +48,13 @@ class BlockProcessingContext: target_dtype: torch.dtype autocast_enabled: bool boundary_timestep: float | None + high_noise_timesteps: torch.Tensor | None context_noise: float - use_diagonal_denoising: bool - diagonal_warmup_mid_steps: torch.Tensor | None image_kwargs: dict[str, Any] pos_cond_kwargs: dict[str, Any] - def get_kv_cache(self, timestep_val: float) -> list[KVCache]: + 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 @@ -72,48 +66,33 @@ class BlockProcessingContext: class MatrixGameCausalDenoisingStage(DenoisingStage): - def __init__(self, - transformer, - scheduler, - pipeline=None, - transformer_2=None, - vae=None) -> None: + def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None) -> None: super().__init__(transformer, scheduler, pipeline, transformer_2, vae) self.transformer = transformer self.transformer_2 = transformer_2 self.vae = vae self.num_transformer_blocks = len(self.transformer.blocks) - if hasattr(self.transformer, 'config') and hasattr( - self.transformer.config, 'arch_config'): - self.num_frame_per_block = getattr( - self.transformer.config.arch_config, 'num_frames_per_block', - getattr(self.transformer, 'num_frame_per_block', 1)) - self.sliding_window_num_frames = getattr( - self.transformer.config.arch_config, - 'sliding_window_num_frames', 15) + if hasattr(self.transformer, 'config') and hasattr(self.transformer.config, 'arch_config'): + self.num_frame_per_block = getattr(self.transformer.config.arch_config, 'num_frames_per_block', + getattr(self.transformer, 'num_frame_per_block', 1)) + self.sliding_window_num_frames = getattr(self.transformer.config.arch_config, 'sliding_window_num_frames', + 15) else: - self.num_frame_per_block = getattr(self.transformer, - 'num_frame_per_block', 1) + self.num_frame_per_block = getattr(self.transformer, 'num_frame_per_block', 1) self.sliding_window_num_frames = 15 try: - self.local_attn_size = getattr(self.transformer, "local_attn_size", - -1) + self.local_attn_size = getattr(self.transformer, "local_attn_size", -1) except Exception: self.local_attn_size = -1 - assert self.local_attn_size != -1, ( - f"local_attn_size must be set for Matrix-Game causal inference, " - f"got {self.local_attn_size}. Check MatrixGameWanVideoArchConfig.") - assert self.num_frame_per_block > 0, ( - f"num_frame_per_block must be positive, got {self.num_frame_per_block}" - ) + assert self.local_attn_size != -1, (f"local_attn_size must be set for Matrix-Game causal inference, " + f"got {self.local_attn_size}. Check MatrixGameWanVideoArchConfig.") + assert self.num_frame_per_block > 0, (f"num_frame_per_block must be positive, got {self.num_frame_per_block}") - logger.info( - "MatrixGame causal inference initialized: " - "local_attn_size=%s, num_frame_per_block=%s", self.local_attn_size, - self.num_frame_per_block) + logger.info("MatrixGame causal inference initialized: " + "local_attn_size=%s, num_frame_per_block=%s", self.local_attn_size, self.num_frame_per_block) self.action_config = getattr(self.transformer, 'action_config', {}) self.use_action_module = len(self.action_config) > 0 @@ -127,40 +106,31 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): fastvideo_args: FastVideoArgs, ) -> ForwardBatch: target_dtype = torch.bfloat16 - autocast_enabled = (target_dtype != torch.float32 - ) and not fastvideo_args.disable_autocast + 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, diagonal_warmup_mid_steps = ( - resolve_base_denoising_timesteps( - fastvideo_args.pipeline_config, - self.scheduler.timesteps, - device=get_local_torch_device(), - ) - ) - use_diagonal_denoising = getattr( - fastvideo_args.pipeline_config, - "use_diagonal_denoising", - False, - ) + 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) + 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 - ] + 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} @@ -172,9 +142,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): prompt_embeds = batch.prompt_embeds assert torch.isnan(prompt_embeds[0]).sum() == 0 - kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], - dtype=target_dtype, - device=latents.device) + 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], @@ -184,10 +152,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): 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) + 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], @@ -196,9 +163,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): device=latents.device) if t % self.num_frame_per_block != 0: - raise ValueError( - "num_frames must be divisible by num_frame_per_block for causal denoising" - ) + 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 start_index = 0 @@ -226,57 +191,34 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): target_dtype=target_dtype, autocast_enabled=autocast_enabled, boundary_timestep=boundary_timestep, - context_noise=getattr(fastvideo_args.pipeline_config, - "context_noise", 0), - use_diagonal_denoising=use_diagonal_denoising, - diagonal_warmup_mid_steps=diagonal_warmup_mid_steps, + 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) + context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", 0) - total_steps = sum( - len( - resolve_block_denoising_timesteps( - timesteps, - block_idx, - use_diagonal_denoising, - diagonal_warmup_mid_steps, - ) - ) - for block_idx in range(len(block_sizes)) - ) - with self.progress_bar(total=total_steps) as progress_bar: + 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, :, :] - block_timesteps = resolve_block_denoising_timesteps( - timesteps, - block_idx, - use_diagonal_denoising, - diagonal_warmup_mid_steps, - ) + current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :] - action_kwargs = self._prepare_action_kwargs( - batch, start_index, current_num_frames) + action_kwargs = self._prepare_action_kwargs(batch, start_index, current_num_frames) current_latents = self._process_single_block( current_latents=current_latents, batch=batch, start_index=start_index, current_num_frames=current_num_frames, - timesteps=block_timesteps, + timesteps=timesteps, ctx=ctx, action_kwargs=action_kwargs, progress_bar=progress_bar, ) - latents[:, :, start_index:start_index + - current_num_frames, :, :] = current_latents + latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents # Update KV caches with clean context self._update_context_cache( @@ -299,110 +241,102 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): batch.latents = latents return batch - def _prepare_action_kwargs(self, batch: ForwardBatch, start_index: int, - num_frames: int) -> dict[str, Any]: + def _prepare_action_kwargs(self, batch: ForwardBatch, start_index: int, num_frames: int) -> dict[str, Any]: action_kwargs: dict[str, Any] = {} if not self.use_action_module: return action_kwargs vae_time_compression_ratio = 4 - end_frame_idx = 1 + vae_time_compression_ratio * (start_index + - num_frames - 1) + end_frame_idx = 1 + vae_time_compression_ratio * (start_index + num_frames - 1) if hasattr(batch, 'mouse_cond') and batch.mouse_cond is not None: action_kwargs['mouse_cond'] = batch.mouse_cond[:, :end_frame_idx] if hasattr(batch, 'keyboard_cond') and batch.keyboard_cond is not None: - action_kwargs[ - 'keyboard_cond'] = batch.keyboard_cond[:, :end_frame_idx] + action_kwargs['keyboard_cond'] = batch.keyboard_cond[:, :end_frame_idx] # CRITICAL: Pass num_frame_per_block to model - this should be num_frames (current block size) action_kwargs['num_frame_per_block'] = num_frames return action_kwargs - def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, - device: torch.device) -> list[KVCache]: + def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device) -> list[dict]: kv_cache = [] num_attention_heads = self.transformer.num_attention_heads - attention_head_dim = getattr( - self.transformer, 'attention_head_dim', - self.transformer.hidden_size // num_attention_heads) + attention_head_dim = getattr(self.transformer, 'attention_head_dim', + self.transformer.hidden_size // num_attention_heads) 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 for _ in range(self.num_transformer_blocks): - kv_cache.append( - KVCache.zeros( - batch_size=batch_size, - cache_size=kv_cache_size, - num_heads=num_attention_heads, - head_dim=attention_head_dim, - dtype=dtype, - device=device, - )) + 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 - def _initialize_action_kv_cache(self, batch_size: int, dtype: torch.dtype, - device: torch.device): + def _initialize_action_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device): kv_cache_mouse = [] kv_cache_keyboard = [] action_heads = self.action_config.get('heads_num', 16) - mouse_head_dim = self.action_config.get('mouse_hidden_dim', - 1024) // action_heads - keyboard_head_dim = self.action_config.get('keyboard_hidden_dim', - 1024) // action_heads + mouse_head_dim = self.action_config.get('mouse_hidden_dim', 1024) // action_heads + keyboard_head_dim = self.action_config.get('keyboard_hidden_dim', 1024) // action_heads - if self.local_attn_size != -1: - kv_cache_size = self.local_attn_size - else: - kv_cache_size = 15 + kv_cache_size = self.local_attn_size if self.local_attn_size != -1 else 15 for _ in range(self.num_transformer_blocks): - kv_cache_keyboard.append( - KVCache.zeros( - batch_size=batch_size, - cache_size=kv_cache_size, - num_heads=action_heads, - head_dim=keyboard_head_dim, - dtype=dtype, - device=device, - )) - kv_cache_mouse.append( - KVCache.zeros( - batch_size=batch_size * self.frame_seq_length, - cache_size=kv_cache_size, - num_heads=action_heads, - head_dim=mouse_head_dim, - dtype=dtype, - device=device, - )) - - return kv_cache_mouse, kv_cache_keyboard - - def _initialize_crossattn_cache(self, batch_size: int, max_text_len: int, - dtype: torch.dtype, - device: torch.device) -> list[dict]: - crossattn_cache = [] - num_attention_heads = self.transformer.num_attention_heads - attention_head_dim = getattr( - self.transformer, 'attention_head_dim', - self.transformer.hidden_size // num_attention_heads) - for _ in range(self.num_transformer_blocks): - crossattn_cache.append({ + kv_cache_keyboard.append({ "k": - torch.zeros([ - batch_size, max_text_len, num_attention_heads, - attention_head_dim - ], + torch.zeros([batch_size, kv_cache_size, action_heads, keyboard_head_dim], dtype=dtype, device=device), + "v": + torch.zeros([batch_size, kv_cache_size, action_heads, keyboard_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), + }) + kv_cache_mouse.append({ + "k": + torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, action_heads, mouse_head_dim], dtype=dtype, device=device), "v": - torch.zeros([ - batch_size, max_text_len, num_attention_heads, - attention_head_dim - ], + torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, action_heads, mouse_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_mouse, kv_cache_keyboard + + def _initialize_crossattn_cache(self, batch_size: int, max_text_len: int, dtype: torch.dtype, + device: torch.device) -> list[dict]: + crossattn_cache = [] + num_attention_heads = self.transformer.num_attention_heads + attention_head_dim = getattr(self.transformer, 'attention_head_dim', + self.transformer.hidden_size // num_attention_heads) + for _ in range(self.num_transformer_blocks): + crossattn_cache.append({ + "k": + torch.zeros([batch_size, max_text_len, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size, max_text_len, num_attention_heads, attention_head_dim], dtype=dtype, device=device), "is_init": @@ -426,12 +360,6 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): prompt_embeds = batch.prompt_embeds noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) - local_high_noise_timesteps = None - if ctx.boundary_timestep is not None: - local_high_noise_timesteps = timesteps[ - timesteps >= ctx.boundary_timestep - ] - 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 @@ -441,32 +369,23 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): 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) + 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) + 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) + 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( - ) + 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( - ) + 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, + patch_size=ctx.fastvideo_args.pipeline_config.dit_config.patch_size, VSA_sparsity=ctx.fastvideo_args.VSA_sparsity, device=get_local_torch_device(), ) @@ -483,25 +402,21 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): 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) + 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, - "update_kv_cache": False, } 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, + "kv_cache_mouse": ctx.kv_cache_mouse, + "kv_cache_keyboard": ctx.kv_cache_keyboard, }) model_kwargs.update(action_kwargs) @@ -519,57 +434,42 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): 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]) + 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]) + 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) + 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) + 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) + 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 local_high_noise_timesteps is not None - and i < len(local_high_noise_timesteps) - 1 - ): + 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( + 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 local_high_noise_timesteps is not None - and i == len(local_high_noise_timesteps) - 1 - ): + 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]) + 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) @@ -593,8 +493,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): 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, + 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) @@ -610,15 +509,12 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): "crossattn_cache": ctx.crossattn_cache, "current_start": start_index * self.frame_seq_length, "start_frame": start_index, - "update_kv_cache": True, } if self.use_action_module: context_model_kwargs.update({ - "kv_cache_mouse": - ctx.kv_cache_mouse, - "kv_cache_keyboard": - ctx.kv_cache_keyboard, + "kv_cache_mouse": ctx.kv_cache_mouse, + "kv_cache_keyboard": ctx.kv_cache_keyboard, }) context_model_kwargs.update(action_kwargs) @@ -631,7 +527,6 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): crossattn_cache=ctx.crossattn_cache, current_start=start_index * self.frame_seq_length, start_frame=start_index, - update_kv_cache=True, **ctx.image_kwargs, **ctx.pos_cond_kwargs, ) @@ -645,43 +540,33 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): **context_model_kwargs, ) - def streaming_reset(self, batch: ForwardBatch, - fastvideo_args: FastVideoArgs) -> ForwardBatch: + 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 + 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, diagonal_warmup_mid_steps = ( - resolve_base_denoising_timesteps( - fastvideo_args.pipeline_config, - self.scheduler.timesteps, - device=get_local_torch_device(), - ) - ) - use_diagonal_denoising = getattr( - fastvideo_args.pipeline_config, - "use_diagonal_denoising", - False, - ) + 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) + 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 - ] + 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} @@ -694,9 +579,7 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): 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_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], @@ -706,10 +589,9 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): 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) + 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], @@ -719,33 +601,21 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): # 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" - ) + 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 - max_block_step_count = max( - len( - resolve_block_denoising_timesteps( - timesteps, - block_idx, - use_diagonal_denoising, - diagonal_warmup_mid_steps, - ) - ) - for block_idx in range(len(block_sizes)) - ) + 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(max_block_step_count - 1) + ) for _ in range(num_denoising_steps - 1) ] # Create and store context @@ -765,10 +635,8 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): target_dtype=target_dtype, autocast_enabled=autocast_enabled, boundary_timestep=boundary_timestep, - context_noise=getattr(fastvideo_args.pipeline_config, - "context_noise", 0), - use_diagonal_denoising=use_diagonal_denoising, - diagonal_warmup_mid_steps=diagonal_warmup_mid_steps, + 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, ) @@ -776,13 +644,11 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): self._streaming_initialized = True return batch - def streaming_step( - self, - keyboard_action: torch.Tensor | None = None, - mouse_action: torch.Tensor | None = None) -> ForwardBatch: + 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.") + raise RuntimeError("Streaming not initialized! Call streaming_reset first.") ctx = self._streaming_ctx if ctx.block_idx >= len(ctx.block_sizes): @@ -795,61 +661,46 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): current_num_frames = ctx.block_sizes[ctx.block_idx] start_index = ctx.start_index - current_latents = latents[:, :, start_index:start_index + - current_num_frames, :, :] + 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) + 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) + 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) + 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) + 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: + 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) + 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) + 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=resolve_block_denoising_timesteps( - ctx.timesteps, - ctx.block_idx, - ctx.use_diagonal_denoising, - ctx.diagonal_warmup_mid_steps, - ), + timesteps=ctx.timesteps, ctx=ctx, action_kwargs=action_kwargs, noise_generator=streaming_noise_generator, ) - latents[:, :, start_index:start_index + - current_num_frames, :, :] = current_latents + latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents # Update KV caches with clean context self._update_context_cache( @@ -872,25 +723,17 @@ class MatrixGameCausalDenoisingStage(DenoisingStage): self._streaming_initialized = False self._streaming_ctx = None - def verify_input(self, batch: ForwardBatch, - fastvideo_args: FastVideoArgs) -> VerificationResult: + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: result = VerificationResult() - result.add_check("latents", batch.latents, - [V.is_tensor, V.with_dims(5)]) + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) result.add_check("image_embeds", batch.image_embeds, V.is_list) - result.add_check("image_latent", batch.image_latent, - V.none_or_tensor_with_dims(5)) - result.add_check("num_inference_steps", batch.num_inference_steps, - V.positive_int) - result.add_check("guidance_scale", batch.guidance_scale, - V.positive_float) + result.add_check("image_latent", batch.image_latent, V.none_or_tensor_with_dims(5)) + result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int) + result.add_check("guidance_scale", batch.guidance_scale, V.positive_float) result.add_check("eta", batch.eta, V.non_negative_float) - result.add_check("generator", batch.generator, - V.generator_or_list_generators) - result.add_check("do_classifier_free_guidance", - batch.do_classifier_free_guidance, V.bool_value) - result.add_check( - "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: - not batch.do_classifier_free_guidance or V.list_not_empty(x)) + result.add_check("generator", batch.generator, V.generator_or_list_generators) + result.add_check("do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value) + result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds, + lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x)) return result