cleanup
This commit is contained in:
H1yori233
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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, )
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user