diff --git a/fastvideo/configs/pipelines/base.py b/fastvideo/configs/pipelines/base.py index 87baeb1cf..80ddc04b3 100644 --- a/fastvideo/configs/pipelines/base.py +++ b/fastvideo/configs/pipelines/base.py @@ -63,8 +63,6 @@ class PipelineConfig: # DMD parameters dmd_denoising_steps: list[int] | None = field(default=None) - use_diagonal_denoising: bool = False - diagonal_warmup_mid_steps: list[int] | None = field(default=None) # Wan2.2 TI2V parameters ti2v_task: bool = False @@ -170,25 +168,6 @@ class PipelineConfig: default=PipelineConfig.dmd_denoising_steps, help="Comma-separated list of denoising steps (e.g., '1000,757,522')", ) - parser.add_argument( - f"--{prefix_with_dot}use-diagonal-denoising", - action=StoreBoolean, - dest= - f"{prefix_with_dot.replace('-', '_')}use_diagonal_denoising", - default=PipelineConfig.use_diagonal_denoising, - help="Enable block-wise diagonal denoising schedule.", - ) - parser.add_argument( - f"--{prefix_with_dot}diagonal-warmup-mid-steps", - type=parse_int_list, - dest= - f"{prefix_with_dot.replace('-', '_')}diagonal_warmup_mid_steps", - default=PipelineConfig.diagonal_warmup_mid_steps, - help=( - "Optional comma-separated warmup mid steps for diagonal " - "denoising." - ), - ) # Add VAE configuration arguments from fastvideo.configs.models.vaes.base import VAEConfig diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index cf21c948b..10a5fe0b3 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -858,16 +858,12 @@ class TrainingArgs(FastVideoArgs): # Self-forcing specific arguments num_frame_per_block: int = 3 - chunkwise_timestep: bool = True independent_first_frame: bool = False enable_gradient_masking: bool = True gradient_mask_last_n_frames: int = 21 same_step_across_blocks: bool = False # Use same exit timestep for all blocks last_step_only: bool = False # Only use the last timestep for training context_noise: int = 0 # Context noise level for cache updates - - reinit_action_module: bool = False - override_keyboard_dim: int | None = None @classmethod def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs": @@ -1157,13 +1153,6 @@ class TrainingArgs(FastVideoArgs): type=int, default=TrainingArgs.context_noise, help="Context noise level for cache updates") - parser.add_argument("--reinit-action-module", - action=StoreBoolean, - help="Whether to reinitialize the action module during training") - parser.add_argument("--override-keyboard-dim", - type=int, - default=TrainingArgs.override_keyboard_dim, - help="Override keyboard dimension") return parser diff --git a/fastvideo/train/methods/fine_tuning/dfsft.py b/fastvideo/train/methods/fine_tuning/dfsft.py index f6a5a8bf8..861482fc2 100644 --- a/fastvideo/train/methods/fine_tuning/dfsft.py +++ b/fastvideo/train/methods/fine_tuning/dfsft.py @@ -34,9 +34,6 @@ class DiffusionForcingSFTMethod(TrainingMethod): self._attn_kind: Literal["dense", "vsa"] = (self._infer_attn_kind()) self._chunk_size = self._parse_chunk_size(self.method_config.get("chunk_size", None)) - self._share_timestep_within_chunk = self._parse_share_timestep_within_chunk( - self.method_config.get("share_timestep_within_chunk", True) - ) self._timestep_index_range = (self._parse_timestep_index_range()) # Initialize preprocessors on student. @@ -278,19 +275,6 @@ class DiffusionForcingSFTMethod(TrainingMethod): return min_index, max_index + 1 - def _parse_share_timestep_within_chunk( - self, - raw: Any, - ) -> bool: - if raw is None: - return True - if not isinstance(raw, bool): - raise ValueError( - "method_config.share_timestep_within_chunk must be a bool, " - f"got {type(raw).__name__}" - ) - return raw - def _init_optimizers_and_schedulers(self) -> None: tc = self.training_config student_lr = float(tc.optimizer.learning_rate) @@ -320,35 +304,16 @@ class DiffusionForcingSFTMethod(TrainingMethod): num_latents: int, device: torch.device, ) -> torch.Tensor: + chunk_size = self._chunk_size + num_chunks = ((num_latents + chunk_size - 1) // chunk_size) low, high = self._timestep_index_range - timestep_indices = torch.randint( + chunk_indices = torch.randint( low=low, high=high, - size=(batch_size, num_latents), + size=(batch_size, num_chunks), device=device, dtype=torch.long, generator=self.cuda_generator, ) - - if self._share_timestep_within_chunk: - chunk_size = self._chunk_size - num_chunks = (num_latents + chunk_size - 1) // chunk_size - padded_num_latents = num_chunks * chunk_size - if padded_num_latents != num_latents: - pad = timestep_indices[:, -1:].expand( - batch_size, - padded_num_latents - num_latents, - ) - timestep_indices = torch.cat([timestep_indices, pad], dim=1) - timestep_indices = timestep_indices.reshape( - batch_size, - num_chunks, - chunk_size, - ) - timestep_indices[:, :, 1:] = timestep_indices[:, :, :1] - timestep_indices = timestep_indices.reshape( - batch_size, - padded_num_latents, - )[:, :num_latents] - - return timestep_indices + expanded = chunk_indices.repeat_interleave(chunk_size, dim=1) + return expanded[:, :num_latents] diff --git a/fastvideo/train/models/matrixgame/__init__.py b/fastvideo/train/models/matrixgame/__init__.py deleted file mode 100644 index 7a71811a6..000000000 --- a/fastvideo/train/models/matrixgame/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""MatrixGame model plugin package.""" - -from fastvideo.train.models.matrixgame.matrixgame import ( - MatrixGameModel as MatrixGameModel, ) -from fastvideo.train.models.matrixgame.matrixgame_causal import ( - MatrixGameCausalModel as MatrixGameCausalModel, ) diff --git a/fastvideo/train/models/matrixgame/matrixgame.py b/fastvideo/train/models/matrixgame/matrixgame.py deleted file mode 100644 index 63fc09a4b..000000000 --- a/fastvideo/train/models/matrixgame/matrixgame.py +++ /dev/null @@ -1,288 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""MatrixGame training model plugin.""" - -from __future__ import annotations - -import copy -from typing import Any, Literal - -import torch - -from fastvideo.dataset.dataloader.schema import pyarrow_schema_matrixgame -from fastvideo.pipelines import TrainingBatch -from fastvideo.training.training_utils import normalize_dit_input - -from fastvideo.train.models.wan.wan import WanModel -from fastvideo.train.utils.dataloader import ( - build_parquet_matrixgame_train_dataloader, -) - - -class MatrixGameModel(WanModel): - """MatrixGame per-role model for finetuning in the new trainer.""" - - _transformer_cls_name: str = "MatrixGameWanModel" - - def init_preprocessors(self, training_config: Any) -> None: - self.vae = self._load_vae(training_config) - self.world_group = self._get_world_group() - self.sp_group = self._get_sp_group() - self._init_timestep_mechanics() - self.dataloader = build_parquet_matrixgame_train_dataloader( - training_config.data, - parquet_schema=pyarrow_schema_matrixgame, - ) - self.start_step = 0 - - def on_train_start(self) -> None: - # MatrixGame finetuning does not use negative text conditioning. - return - - def prepare_batch( - self, - raw_batch: dict[str, Any], - *, - generator: torch.Generator, - latents_source: Literal["data", "zeros"] = "data", - ) -> TrainingBatch: - assert self.training_config is not None - tc = self.training_config - dtype = self._get_training_dtype() - device = self.device - - training_batch = TrainingBatch() - infos = raw_batch.get("info_list") - batch_size = self._infer_batch_size(raw_batch) - - if latents_source == "zeros": - latents = self._make_zero_latents(batch_size=batch_size) - elif latents_source == "data": - latents = raw_batch["vae_latent"][:, :, :tc.data.num_latent_t] - latents = latents.to(device=device, dtype=dtype) - else: - raise ValueError(f"Unknown latents_source: {latents_source!r}") - - clip_feature = raw_batch["clip_feature"].to(device=device, dtype=dtype) - first_frame_latent = raw_batch["first_frame_latent"] - first_frame_latent = first_frame_latent[:, :, :tc.data.num_latent_t] - first_frame_latent = first_frame_latent.to(device=device, dtype=dtype) - - pil_image = raw_batch.get("pil_image") - if pil_image is not None: - pil_image = pil_image.to(device=device) - - keyboard_cond = self._get_optional_action( - raw_batch, - key="keyboard_cond", - expected_frames=self._expected_action_frames(tc.data.num_latent_t), - dtype=dtype, - ) - mouse_cond = self._get_optional_action( - raw_batch, - key="mouse_cond", - expected_frames=self._expected_action_frames(tc.data.num_latent_t), - dtype=dtype, - ) - - training_batch.latents = normalize_dit_input("wan", latents, self.vae) - training_batch.encoder_hidden_states = None - training_batch.encoder_attention_mask = None - training_batch.preprocessed_image = pil_image - training_batch.image_embeds = clip_feature - training_batch.image_latents = first_frame_latent - training_batch.keyboard_cond = keyboard_cond - training_batch.mouse_cond = mouse_cond - training_batch.infos = infos - - return self._finalize_prepared_batch(training_batch, generator) - - def _prepare_dit_inputs( - self, - training_batch: TrainingBatch, - generator: torch.Generator, - ) -> TrainingBatch: - training_batch = super()._prepare_dit_inputs(training_batch, generator) - - image_latents = training_batch.image_latents - image_embeds = training_batch.image_embeds - if image_latents is None or image_embeds is None: - raise RuntimeError("MatrixGame requires image_latents and image_embeds") - - cond_latents = self._build_matrixgame_cond_concat(image_latents) - training_batch.image_latents = cond_latents - training_batch.noisy_model_input = torch.cat( - [training_batch.noisy_model_input, cond_latents], - dim=1, - ) - training_batch.conditional_dict = { - "encoder_hidden_states": None, - "encoder_attention_mask": None, - "encoder_hidden_states_image": image_embeds, - "keyboard_cond": training_batch.keyboard_cond, - "mouse_cond": training_batch.mouse_cond, - } - training_batch.unconditional_dict = dict(training_batch.conditional_dict) - return training_batch - - def _get_uncond_text_dict( - self, - batch: TrainingBatch, - *, - cfg_uncond: dict[str, Any] | None, - ) -> dict[str, Any]: - del cfg_uncond - cond_dict = batch.conditional_dict - if cond_dict is None: - raise RuntimeError("Missing conditional_dict in TrainingBatch") - return batch.unconditional_dict or cond_dict - - def _build_distill_input_kwargs( - self, - noise_input: torch.Tensor, - timestep: torch.Tensor, - text_dict: dict[str, Any] | None, - ) -> dict[str, Any]: - if text_dict is None: - raise ValueError("text_dict cannot be None for MatrixGame") - return { - "hidden_states": noise_input.permute(0, 2, 1, 3, 4), - "encoder_hidden_states": None, - "timestep": timestep.to(device=self.device, dtype=torch.bfloat16), - "encoder_hidden_states_image": text_dict["encoder_hidden_states_image"], - "keyboard_cond": text_dict["keyboard_cond"], - "mouse_cond": text_dict["mouse_cond"], - "return_dict": False, - } - - def _build_matrixgame_cond_concat( - self, - image_latents: torch.Tensor, - ) -> torch.Tensor: - if image_latents.ndim != 5: - raise ValueError( - "first_frame_latent must have shape [B, C, T, H, W], " - f"got {tuple(image_latents.shape)}" - ) - if image_latents.shape[1] == 20: - return image_latents - if image_latents.shape[1] != 16: - raise ValueError( - "MatrixGame expects first_frame_latent with 16 or 20 channels, " - f"got {image_latents.shape[1]}" - ) - - temporal_compression_ratio = self._temporal_compression_ratio() - batch_size, _, num_latent_t, latent_height, latent_width = image_latents.shape - num_frames = (num_latent_t - 1) * temporal_compression_ratio + 1 - - mask_lat_size = torch.ones( - batch_size, - 1, - num_frames, - latent_height, - latent_width, - device=image_latents.device, - dtype=image_latents.dtype, - ) - mask_lat_size[:, :, 1:] = 0 - first_frame_mask = mask_lat_size[:, :, :1] - first_frame_mask = torch.repeat_interleave( - first_frame_mask, - dim=2, - repeats=temporal_compression_ratio, - ) - mask_lat_size = torch.cat( - [first_frame_mask, mask_lat_size[:, :, 1:]], - dim=2, - ) - mask_lat_size = mask_lat_size.view( - batch_size, - -1, - temporal_compression_ratio, - latent_height, - latent_width, - ).transpose(1, 2) - return torch.cat([mask_lat_size, image_latents], dim=1) - - def _get_optional_action( - self, - raw_batch: dict[str, Any], - *, - key: str, - expected_frames: int, - dtype: torch.dtype, - ) -> torch.Tensor | None: - value = raw_batch.get(key) - if value is None or value.numel() == 0: - return None - if value.shape[1] < expected_frames: - raise ValueError( - f"{key} has {value.shape[1]} frames but requires at least {expected_frames}" - ) - return value[:, :expected_frames].to(device=self.device, dtype=dtype) - - def _finalize_prepared_batch( - self, - training_batch: TrainingBatch, - generator: torch.Generator, - ) -> TrainingBatch: - training_batch = self._prepare_dit_inputs(training_batch, generator) - training_batch = self._build_attention_metadata(training_batch) - - training_batch.attn_metadata_vsa = copy.deepcopy(training_batch.attn_metadata) - if training_batch.attn_metadata is not None: - training_batch.attn_metadata.VSA_sparsity = 0.0 # type: ignore[attr-defined] - - return training_batch - - def _expected_action_frames(self, num_latent_t: int) -> int: - return (num_latent_t - 1) * self._temporal_compression_ratio() + 1 - - def _temporal_compression_ratio(self) -> int: - assert self.training_config is not None - return int( - self.training_config.pipeline_config.vae_config.arch_config.temporal_compression_ratio # type: ignore[union-attr] - ) - - def _infer_batch_size(self, raw_batch: dict[str, Any]) -> int: - if "vae_latent" in raw_batch: - return int(raw_batch["vae_latent"].shape[0]) - if "clip_feature" in raw_batch: - return int(raw_batch["clip_feature"].shape[0]) - raise ValueError("Unable to infer batch size from MatrixGame batch") - - def _make_zero_latents(self, *, batch_size: int) -> torch.Tensor: - assert self.training_config is not None - vae_config = self.training_config.pipeline_config.vae_config.arch_config # type: ignore[union-attr] - num_channels = vae_config.z_dim - spatial_compression_ratio = vae_config.spatial_compression_ratio - latent_height = self.training_config.data.num_height // spatial_compression_ratio - latent_width = self.training_config.data.num_width // spatial_compression_ratio - return torch.zeros( - batch_size, - num_channels, - self.training_config.data.num_latent_t, - latent_height, - latent_width, - device=self.device, - dtype=self._get_training_dtype(), - ) - - def _load_vae(self, training_config: Any) -> Any: - from fastvideo.train.utils.moduleloader import load_module_from_path - - return load_module_from_path( - model_path=str(training_config.model_path), - module_type="vae", - training_config=training_config, - ) - - def _get_world_group(self) -> Any: - from fastvideo.distributed import get_world_group - - return get_world_group() - - def _get_sp_group(self) -> Any: - from fastvideo.distributed import get_sp_group - - return get_sp_group() diff --git a/fastvideo/train/models/matrixgame/matrixgame_causal.py b/fastvideo/train/models/matrixgame/matrixgame_causal.py deleted file mode 100644 index c7f998678..000000000 --- a/fastvideo/train/models/matrixgame/matrixgame_causal.py +++ /dev/null @@ -1,476 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""MatrixGame causal training model plugin.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Literal - -import torch - -from fastvideo.forward_context import set_forward_context -from fastvideo.models.dits.matrixgame.kv_cache import KVCache - -from fastvideo.train.models.matrixgame.matrixgame import MatrixGameModel -from fastvideo.train.models.wan.wan_causal import WanCausalModel - - -@dataclass(slots=True) -class _MatrixGameStreamingCaches: - kv_cache: list[dict[str, Any]] - kv_cache_mouse: list[KVCache | None] - kv_cache_keyboard: list[KVCache | None] - crossattn_cache: list[dict[str, Any]] | None - frame_seq_length: int - local_attn_size: int - sliding_window_num_frames: int - batch_size: int - dtype: torch.dtype - device: torch.device - - -class MatrixGameCausalModel(MatrixGameModel, WanCausalModel): - """MatrixGame per-role model with causal/streaming primitives.""" - - _transformer_cls_name: str = "CausalMatrixGameWanModel" - - def __init__(self, *args: Any, **kwargs: Any) -> None: - super().__init__(*args, **kwargs) - self._streaming_caches: ( - dict[tuple[int, str], _MatrixGameStreamingCaches] - ) = {} - - def clear_caches(self, *, cache_tag: str = "pos") -> None: - self._streaming_caches.pop((id(self), str(cache_tag)), None) - - def predict_noise_streaming( - self, - noisy_latents: torch.Tensor, - timestep: torch.Tensor, - batch: Any, - *, - conditional: bool, - cache_tag: str = "pos", - store_kv: bool = False, - cur_start_frame: int = 0, - cfg_uncond: dict[str, Any] | None = None, - attn_kind: Literal["dense", "vsa"] = "dense", - ) -> torch.Tensor | None: - del cfg_uncond - attn_metadata = self._get_attn_metadata(batch, attn_kind=attn_kind) - - cache_tag = str(cache_tag) - cur_start_frame = int(cur_start_frame) - if cur_start_frame < 0: - raise ValueError("cur_start_frame must be >= 0") - - batch_size = int(noisy_latents.shape[0]) - num_frames = int(noisy_latents.shape[1]) - timestep_full = self._ensure_per_frame_timestep( - timestep=timestep, - batch_size=batch_size, - num_frames=num_frames, - device=noisy_latents.device, - ) - - transformer = self._get_transformer(timestep_full) - caches = self._get_or_init_streaming_caches( - cache_tag=cache_tag, - transformer=transformer, - noisy_latents=noisy_latents, - ) - - kv_cache = caches.kv_cache - kv_cache_mouse = caches.kv_cache_mouse - kv_cache_keyboard = caches.kv_cache_keyboard - crossattn_cache = caches.crossattn_cache - - ( - kv_cache, - kv_cache_mouse, - kv_cache_keyboard, - ) = self._maybe_snapshot_streaming_caches( - kv_cache=kv_cache, - kv_cache_mouse=kv_cache_mouse, - kv_cache_keyboard=kv_cache_keyboard, - ) - model_kwargs = self._build_streaming_model_kwargs( - kv_cache=kv_cache, - kv_cache_mouse=kv_cache_mouse, - kv_cache_keyboard=kv_cache_keyboard, - crossattn_cache=crossattn_cache, - cur_start_frame=cur_start_frame, - frame_seq_length=int(caches.frame_seq_length), - store_kv=store_kv, - num_frames=num_frames, - ) - cond_dict = self._get_streaming_conditioning( - batch, - conditional=conditional, - ) - - device_type = self.device.type - dtype = noisy_latents.dtype - with ( - torch.autocast(device_type, dtype=dtype), - set_forward_context( - current_timestep=batch.timesteps, - attn_metadata=attn_metadata, - ), - ): - input_kwargs = self._build_streaming_input_kwargs( - noisy_latents=noisy_latents, - timestep=timestep_full, - cond_dict=cond_dict, - batch=batch, - cur_start_frame=cur_start_frame, - ) - input_kwargs["timestep"] = timestep_full.to( - device=self.device, - dtype=torch.long, - ) - input_kwargs.update(model_kwargs) - - if store_kv: - with torch.no_grad(): - _ = transformer(**input_kwargs) - return None - - pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4) - return pred_noise - - def _get_attn_metadata( - self, - batch: Any, - *, - attn_kind: Literal["dense", "vsa"], - ) -> Any: - if attn_kind == "dense": - return batch.attn_metadata - if attn_kind == "vsa": - return batch.attn_metadata_vsa - raise ValueError(f"Unknown attn_kind: {attn_kind!r}") - - def _maybe_snapshot_streaming_caches( - self, - *, - kv_cache: list[dict[str, Any]], - kv_cache_mouse: list[KVCache | None], - kv_cache_keyboard: list[KVCache | None], - ) -> tuple[list[dict[str, Any]], list[KVCache | None], list[KVCache | None]]: - if self._should_snapshot_streaming_cache() and torch.is_grad_enabled(): - kv_cache = self._snapshot_kv_cache_indices(kv_cache) - kv_cache_mouse = self._snapshot_action_kv_cache_lengths(kv_cache_mouse) - kv_cache_keyboard = self._snapshot_action_kv_cache_lengths( - kv_cache_keyboard - ) - return kv_cache, kv_cache_mouse, kv_cache_keyboard - - def _build_streaming_model_kwargs( - self, - *, - kv_cache: list[dict[str, Any]], - kv_cache_mouse: list[KVCache | None], - kv_cache_keyboard: list[KVCache | None], - crossattn_cache: list[dict[str, Any]] | None, - cur_start_frame: int, - frame_seq_length: int, - store_kv: bool, - num_frames: int, - ) -> dict[str, Any]: - return { - "kv_cache": kv_cache, - "kv_cache_mouse": kv_cache_mouse, - "kv_cache_keyboard": kv_cache_keyboard, - "crossattn_cache": crossattn_cache, - "current_start": cur_start_frame * frame_seq_length, - "start_frame": cur_start_frame, - "update_kv_cache": bool(store_kv), - "num_frame_per_block": num_frames, - } - - def _get_streaming_conditioning( - self, - batch: Any, - *, - conditional: bool, - ) -> dict[str, Any]: - cond_dict = batch.conditional_dict - if cond_dict is None: - raise RuntimeError("Missing conditional_dict in TrainingBatch") - if conditional: - return cond_dict - return batch.unconditional_dict or cond_dict - - def _get_or_init_streaming_caches( - self, - *, - cache_tag: str, - transformer: torch.nn.Module, - noisy_latents: torch.Tensor, - ) -> _MatrixGameStreamingCaches: - key = (id(self), cache_tag) - cached = self._streaming_caches.get(key) - - batch_size = int(noisy_latents.shape[0]) - dtype = noisy_latents.dtype - device = noisy_latents.device - - frame_seq_length = self._compute_frame_seq_length(transformer, noisy_latents) - local_attn_size = self._get_local_attn_size(transformer) - sliding_window_num_frames = self._get_sliding_window_num_frames(transformer) - - meta = ( - frame_seq_length, - local_attn_size, - sliding_window_num_frames, - batch_size, - dtype, - device, - ) - - if cached is not None: - cached_meta = ( - cached.frame_seq_length, - cached.local_attn_size, - cached.sliding_window_num_frames, - cached.batch_size, - cached.dtype, - cached.device, - ) - if cached_meta == meta: - return cached - - kv_cache = self._initialize_kv_cache( - transformer=transformer, - batch_size=batch_size, - dtype=dtype, - device=device, - frame_seq_length=frame_seq_length, - local_attn_size=local_attn_size, - sliding_window_num_frames=sliding_window_num_frames, - checkpoint_safe=self._should_use_checkpoint_safe_kv_cache(), - ) - kv_cache_mouse = self._initialize_action_kv_cache( - transformer=transformer, - batch_size=batch_size, - dtype=dtype, - device=device, - local_attn_size=local_attn_size, - sliding_window_num_frames=sliding_window_num_frames, - channel="mouse", - ) - kv_cache_keyboard = self._initialize_action_kv_cache( - transformer=transformer, - batch_size=batch_size, - dtype=dtype, - device=device, - local_attn_size=local_attn_size, - sliding_window_num_frames=sliding_window_num_frames, - channel="keyboard", - ) - crossattn_cache = self._initialize_crossattn_cache( - transformer=transformer, - device=device, - ) - - caches = _MatrixGameStreamingCaches( - kv_cache=kv_cache, - kv_cache_mouse=kv_cache_mouse, - kv_cache_keyboard=kv_cache_keyboard, - crossattn_cache=crossattn_cache, - frame_seq_length=frame_seq_length, - local_attn_size=local_attn_size, - sliding_window_num_frames=sliding_window_num_frames, - batch_size=batch_size, - dtype=dtype, - device=device, - ) - self._streaming_caches[key] = caches - return caches - - def _build_streaming_input_kwargs( - self, - *, - noisy_latents: torch.Tensor, - timestep: torch.Tensor, - cond_dict: dict[str, Any], - batch: Any, - cur_start_frame: int, - ) -> dict[str, Any]: - if batch.image_latents is None: - raise RuntimeError("MatrixGame causal rollout requires image_latents") - if batch.image_embeds is None: - raise RuntimeError("MatrixGame causal rollout requires image_embeds") - - num_frames = int(noisy_latents.shape[1]) - frame_end = cur_start_frame + num_frames - cond_latents = self._slice_image_latents_for_chunk( - batch.image_latents, - start=cur_start_frame, - end=frame_end, - ) - hidden_states = torch.cat( - [noisy_latents.permute(0, 2, 1, 3, 4), cond_latents], - dim=1, - ) - - return { - "hidden_states": hidden_states, - "encoder_hidden_states": None, - "timestep": timestep, - "encoder_hidden_states_image": cond_dict["encoder_hidden_states_image"], - "keyboard_cond": self._slice_action_prefix( - batch.keyboard_cond, - frame_end=frame_end, - ), - "mouse_cond": self._slice_action_prefix( - batch.mouse_cond, - frame_end=frame_end, - ), - "return_dict": False, - } - - def _slice_image_latents_for_chunk( - self, - image_latents: torch.Tensor, - *, - start: int, - end: int, - ) -> torch.Tensor: - num_frames = end - start - if image_latents.ndim != 5: - raise ValueError( - "image_latents must have shape [B, C, T, H, W], " - f"got {tuple(image_latents.shape)}" - ) - if image_latents.shape[2] >= end: - return image_latents[:, :, start:end] - if image_latents.shape[2] > start: - cond = image_latents[:, :, start:] - pad_frames = num_frames - int(cond.shape[2]) - if pad_frames <= 0: - return cond - pad = torch.zeros( - cond.shape[0], - cond.shape[1], - pad_frames, - cond.shape[3], - cond.shape[4], - device=cond.device, - dtype=cond.dtype, - ) - return torch.cat([cond, pad], dim=2) - return torch.zeros( - image_latents.shape[0], - image_latents.shape[1], - num_frames, - image_latents.shape[3], - image_latents.shape[4], - device=image_latents.device, - dtype=image_latents.dtype, - ) - - def _slice_action_prefix( - self, - action: torch.Tensor | None, - *, - frame_end: int, - ) -> torch.Tensor | None: - if action is None: - return None - action_frame_end = ((frame_end - 1) * self._temporal_compression_ratio()) + 1 - if action.shape[1] < action_frame_end: - raise ValueError( - "Action tensor is shorter than required for causal rollout: " - f"got={action.shape[1]}, required>={action_frame_end}" - ) - return action[:, :action_frame_end] - - def _initialize_action_kv_cache( - self, - *, - transformer: torch.nn.Module, - batch_size: int, - dtype: torch.dtype, - device: torch.device, - local_attn_size: int, - sliding_window_num_frames: int, - channel: Literal["mouse", "keyboard"], - ) -> list[KVCache | None]: - blocks = list(getattr(transformer, "blocks", [])) - if not blocks: - return [] - - cache_size = ( - int(local_attn_size) - if int(local_attn_size) != -1 - else int(sliding_window_num_frames) - ) - if cache_size <= 0: - raise ValueError( - f"Invalid action cache size for channel {channel!r}: {cache_size}" - ) - - caches: list[KVCache | None] = [] - for block in blocks: - action_model = getattr(block, "action_model", None) - if action_model is None: - caches.append(None) - continue - - if channel == "mouse": - enabled = bool(getattr(action_model, "enable_mouse", False)) - hidden_dim = getattr( - getattr(action_model, "proj_mouse", None), - "in_features", - None, - ) - else: - enabled = bool(getattr(action_model, "enable_keyboard", False)) - hidden_dim = getattr( - getattr(action_model, "proj_keyboard", None), - "in_features", - None, - ) - - if not enabled: - caches.append(None) - continue - if hidden_dim is None: - raise ValueError( - f"Unable to determine {channel} action hidden size for " - "MatrixGame causal cache initialization" - ) - - num_heads = int(getattr(action_model, "heads_num")) - head_dim = int(hidden_dim) // max(1, num_heads) - caches.append( - KVCache.zeros( - batch_size=batch_size, - cache_size=cache_size, - num_heads=num_heads, - head_dim=head_dim, - dtype=dtype, - device=device, - ) - ) - return caches - - def _snapshot_action_kv_cache_lengths( - self, - caches: list[KVCache | None], - ) -> list[KVCache | None]: - snapshot: list[KVCache | None] = [] - for cache in caches: - if cache is None: - snapshot.append(None) - continue - snapshot.append( - KVCache( - k=cache.k, - v=cache.v, - length=cache.length.detach().clone(), - ) - ) - return snapshot diff --git a/fastvideo/train/utils/dataloader.py b/fastvideo/train/utils/dataloader.py index 5c2ce6794..a1b22d3ba 100644 --- a/fastvideo/train/utils/dataloader.py +++ b/fastvideo/train/utils/dataloader.py @@ -31,26 +31,3 @@ def build_parquet_t2v_train_dataloader( seed=int(data_config.seed or 0), )) return dataloader - - -def build_parquet_matrixgame_train_dataloader( - data_config: DataConfig, - *, - parquet_schema: Any, -) -> Any: - """Build a parquet dataloader for MatrixGame datasets.""" - - from fastvideo.dataset import ( - build_parquet_map_style_dataloader, ) - - _dataset, dataloader = (build_parquet_map_style_dataloader( - data_config.data_path, - data_config.train_batch_size, - num_data_workers=(data_config.dataloader_num_workers), - parquet_schema=parquet_schema, - cfg_rate=float(data_config.training_cfg_rate or 0.0), - drop_last=True, - text_padding_length=512, - seed=int(data_config.seed or 0), - )) - return dataloader \ No newline at end of file