nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cleanup

Browse Source
This commit is contained in:
H1yori233
2026-03-24 15:58:35 -07:00
parent deefcde8db
commit 5d43534fe8
4 changed files with 227 additions and 558 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/models/schedulers/denoising_schedule.py
-124
View File
@@ -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,
)
fastvideo/pipelines/pipeline_batch_info.py
+6 -13
View File
@@ -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)
fastvideo/pipelines/stages/input_validation.py
+39 -82
View File
@@ -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
fastvideo/pipelines/stages/matrixgame_denoising.py
+182 -339
View File
@@ -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