From ab8610b6718262d9068960e50788390964c22760 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Mon, 13 Jul 2026 08:23:48 +0000 Subject: [PATCH 01/10] fastwam-libero training --- .../train/fastwam/libero_uncond_2cam224.yaml | 63 ++ .../lightx2v_train/data/__init__.py | 2 + .../lightx2v_train/data/libero/__init__.py | 3 + .../lightx2v_train/data/libero/dataset.py | 119 +++ .../data/libero/lerobot_dataset.py | 197 ++++ .../lightx2v_train/data/libero/normalizer.py | 73 ++ .../lightx2v_train/data/libero/processor.py | 115 +++ .../data/libero/robot_video_dataset.py | 128 +++ .../data/libero/video_decoder.py | 69 ++ .../lightx2v_train/model_zoo/__init__.py | 3 +- .../model_zoo/native/wan/fastwam/__init__.py | 3 + .../native/wan/fastwam/action_dit.py | 367 ++++++++ .../model_zoo/native/wan/fastwam/layers.py | 110 +++ .../model_zoo/native/wan/fastwam/loader.py | 127 +++ .../model_zoo/native/wan/fastwam/model.py | 847 ++++++++++++++++++ .../model_zoo/native/wan/fastwam/mot.py | 547 +++++++++++ .../model_zoo/native/wan/fastwam/video_dit.py | 237 +++++ .../lightx2v_train/model_zoo/wan_fastwam.py | 120 +++ .../schedulers/flow_matching.py | 66 ++ .../lightx2v_train/trainers/__init__.py | 3 +- .../lightx2v_train/trainers/fastwam.py | 517 +++++++++++ lightx2v_train/lightx2v_train/utils/video.py | 120 +++ .../scripts/run_fastwam_libero_train.sh | 11 + 23 files changed, 3845 insertions(+), 2 deletions(-) create mode 100644 lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml create mode 100644 lightx2v_train/lightx2v_train/data/libero/__init__.py create mode 100644 lightx2v_train/lightx2v_train/data/libero/dataset.py create mode 100644 lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py create mode 100644 lightx2v_train/lightx2v_train/data/libero/normalizer.py create mode 100644 lightx2v_train/lightx2v_train/data/libero/processor.py create mode 100644 lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py create mode 100644 lightx2v_train/lightx2v_train/data/libero/video_decoder.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/__init__.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py create mode 100644 lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py create mode 100644 lightx2v_train/lightx2v_train/trainers/fastwam.py create mode 100644 lightx2v_train/lightx2v_train/utils/video.py create mode 100755 lightx2v_train/scripts/run_fastwam_libero_train.sh diff --git a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml new file mode 100644 index 000000000..51d35982c --- /dev/null +++ b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml @@ -0,0 +1,63 @@ +model: + name: wan_fastwam + running_dtype: bf16 + model_path: /data/nvme0/models/Wan-AI/Wan2.2-TI2V-5B + action_dit_pretrained_path: /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt + +data: + train: &libero_data + name: libero_fastwam_dataset + dataset_dirs: + - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_spatial_no_noops_lerobot + - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_object_no_noops_lerobot + - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_goal_no_noops_lerobot + - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_10_no_noops_lerobot + text_embedding_cache_dir: /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/text_embeds_cache/libero + pretrained_norm_stats: /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/dataset_stats.json + batch_size: 1 + num_workers: 0 + val: + <<: *libero_data + +training: + method: fastwam + output_dir: /data/nvme0/chendingyu/LightX2V/lightx2v_train/runs/fastwam_libero + # Smoke default: one optimizer step proves the LightX2V integration path. + # Increase this for a real training run. + max_train_iters: 1 + gradient_accumulation_iters: 1 + max_grad_norm: 1.0 + save_every_iters: 0 + save_total_limit: 2 + save_final: true + lr_scheduler: constant + lr_warmup_iters: 0 + optimizer: + learning_rate: 1.0e-4 + weight_decay: 1.0e-2 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_epsilon: 1.0e-8 + +evaluation: + eval_every_iters: 1 + eval_num_inference_steps: 5 + num_samples: 1 + run_inference: true + save_video: true + save_preview: true + fps: 8 + tiled: false + seed: 42 + +logging: + train_log_every_iters: 1 + +resume: + auto_resume: false + resume_ckpt_path: null + +distributed: + backend: nccl + sequence_parallel: + enabled: false diff --git a/lightx2v_train/lightx2v_train/data/__init__.py b/lightx2v_train/lightx2v_train/data/__init__.py index 7df8a3c00..460d520ff 100644 --- a/lightx2v_train/lightx2v_train/data/__init__.py +++ b/lightx2v_train/lightx2v_train/data/__init__.py @@ -1,11 +1,13 @@ from lightx2v_train.utils.registry import build_data from .image_dataset import build_image_dataset +from .libero.dataset import build_libero_fastwam_dataset from .video_dataset import build_causal_forcing_lmdb_dataset, build_prompt_dataset, build_wan_t2v_cached_dataset, build_wan_t2v_video_dataset __all__ = [ "build_data", "build_image_dataset", + "build_libero_fastwam_dataset", "build_prompt_dataset", "build_wan_t2v_video_dataset", "build_wan_t2v_cached_dataset", diff --git a/lightx2v_train/lightx2v_train/data/libero/__init__.py b/lightx2v_train/lightx2v_train/data/libero/__init__.py new file mode 100644 index 000000000..41a716905 --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/libero/__init__.py @@ -0,0 +1,3 @@ +from .dataset import build_libero_fastwam_dataset + +__all__ = ["build_libero_fastwam_dataset"] diff --git a/lightx2v_train/lightx2v_train/data/libero/dataset.py b/lightx2v_train/lightx2v_train/data/libero/dataset.py new file mode 100644 index 000000000..7c5f498fb --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/libero/dataset.py @@ -0,0 +1,119 @@ +from copy import deepcopy +from pathlib import Path + +import torch +from loguru import logger +from torch.utils.data import DataLoader, DistributedSampler + +from lightx2v_train.runtime.distributed import get_data_parallel_rank, get_data_parallel_world_size +from lightx2v_train.utils.registry import DATA_REGISTER + +from .processor import FastWAMProcessor +from .robot_video_dataset import RobotVideoDataset + +ASSET_ROOT = Path(__file__).resolve().parents[3] / "assets" / "libero_fastwam" +DEFAULT_DATASET_DIRS = [ + ASSET_ROOT / "datasets" / name + for name in ( + "libero_spatial_no_noops_lerobot", + "libero_object_no_noops_lerobot", + "libero_goal_no_noops_lerobot", + "libero_10_no_noops_lerobot", + ) +] + + +def _default_shape_meta(): + return { + "images": [ + {"key": "image", "raw_shape": [3, 512, 512], "shape": [3, 224, 224]}, + {"key": "wrist_image", "raw_shape": [3, 512, 512], "shape": [3, 224, 224]}, + ], + "action": [{"key": "default", "raw_shape": 7, "shape": 7}], + "state": [{"key": "default", "raw_shape": 8, "shape": 8}], + } + + +class DatasetSliceRepeat(torch.utils.data.Dataset): + def __init__(self, dataset, max_samples=None, dataset_repeat=1): + self.dataset = dataset + self.base_len = len(dataset) if max_samples is None else min(int(max_samples), len(dataset)) + self.dataset_repeat = max(1, int(dataset_repeat)) + if self.base_len <= 0: + raise RuntimeError("LIBERO dataset is empty after applying max_samples.") + + def __len__(self): + return self.base_len * self.dataset_repeat + + def __getitem__(self, index): + return self.dataset[index % self.base_len] + + +def _path(value): + return str(Path(value).expanduser().resolve()) + + +def _build_dataset(config, split): + shape_meta = deepcopy(config.get("shape_meta") or _default_shape_meta()) + num_frames = int(config.get("num_frames", 33)) + processor = FastWAMProcessor(shape_meta, num_frames) + dataset_dirs = config.get("dataset_dirs") or DEFAULT_DATASET_DIRS + if isinstance(dataset_dirs, (str, Path)): + dataset_dirs = [dataset_dirs] + + dataset = RobotVideoDataset( + dataset_dirs=[_path(item) for item in dataset_dirs], + shape_meta=shape_meta, + processor=processor, + text_embedding_cache_dir=_path( + config.get("text_embedding_cache_dir") + or ASSET_ROOT / "text_embeds_cache" / "libero" + ), + pretrained_norm_stats=_path( + config.get("pretrained_norm_stats") or ASSET_ROOT / "dataset_stats.json" + ), + num_frames=num_frames, + context_len=int(config.get("context_len", 128)), + val_set_proportion=float(config.get("val_set_proportion", 0.0)), + is_training_set=bool(config.get("is_training_set", split == "train")), + global_sample_stride=int(config.get("global_sample_stride", 1)), + action_video_freq_ratio=int(config.get("action_video_freq_ratio", 4)), + skip_padding_as_possible=bool(config.get("skip_padding_as_possible", False)), + max_padding_retry=int(config.get("max_padding_retry", 3)), + video_backend=config.get("video_backend"), + ) + logger.info("[data] built LIBERO FastWAM {} dataset size={}", split, len(dataset)) + return DatasetSliceRepeat( + dataset, + max_samples=config.get("max_samples"), + dataset_repeat=config.get("dataset_repeat", 1), + ) + + +@DATA_REGISTER("libero_fastwam_dataset") +def build_libero_fastwam_dataset(data_config, train_or_val="train"): + dataset = _build_dataset(data_config, train_or_val) + if data_config.get("return_dataset", False): + return dataset + + sampler = None + shuffle = bool(data_config.get("shuffle", train_or_val == "train")) + world_size = get_data_parallel_world_size() + if train_or_val == "train" and world_size > 1: + sampler = DistributedSampler( + dataset, + num_replicas=world_size, + rank=get_data_parallel_rank(), + shuffle=shuffle, + drop_last=bool(data_config.get("drop_last", False)), + ) + shuffle = False + return DataLoader( + dataset, + batch_size=int(data_config.get("batch_size", 1)), + shuffle=shuffle, + sampler=sampler, + num_workers=int(data_config.get("num_workers", 4)), + pin_memory=bool(data_config.get("pin_memory", torch.cuda.is_available())), + drop_last=bool(data_config.get("drop_last", False)), + ) diff --git a/lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py b/lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py new file mode 100644 index 000000000..4af8c3312 --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py @@ -0,0 +1,197 @@ +import bisect +import json +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import torch +from torch.utils.data import Dataset + +from .video_decoder import decode_video_frames + + +@dataclass(frozen=True) +class Episode: + root: Path + index: int + length: int + data_path: str + video_path: str + chunks_size: int + tasks: dict[int, str] + + def parquet_path(self) -> Path: + return self.root / self.data_path.format( + episode_chunk=self.index // self.chunks_size, + episode_index=self.index, + ) + + def camera_path(self, key: str) -> Path: + return self.root / self.video_path.format( + episode_chunk=self.index // self.chunks_size, + episode_index=self.index, + video_key=key, + ) + + +def _read_json(path: Path): + with path.open(encoding="utf-8") as handle: + return json.load(handle) + + +def _read_jsonl(path: Path): + with path.open(encoding="utf-8") as handle: + return [json.loads(line) for line in handle if line.strip()] + + +def _column_to_tensor(column: pa.ChunkedArray) -> torch.Tensor: + column = column.combine_chunks() + if pa.types.is_list(column.type) or pa.types.is_fixed_size_list(column.type): + array = np.asarray(column.to_pylist()) + else: + array = column.to_numpy(zero_copy_only=False) + return torch.from_numpy(np.array(array, copy=True)) + + +class LiberoLeRobotDataset(Dataset): + """Minimal local reader for FastWAM's processed LIBERO LeRobot v2.1 data.""" + + def __init__( + self, + dataset_dirs, + image_keys, + state_key, + action_key, + num_frames, + global_sample_stride=1, + val_set_proportion=0.0, + is_training_set=True, + seed=42, + video_backend=None, + ): + if num_frames < 2: + raise ValueError(f"num_frames must be at least 2, got {num_frames}") + self.image_keys = list(image_keys) + self.state_key = state_key + self.action_key = action_key + self.num_frames = int(num_frames) + self.global_sample_stride = int(global_sample_stride) + self.video_backend = video_backend + + episodes = [] + fps_values = set() + for dataset_dir in dataset_dirs: + root = Path(dataset_dir).expanduser().resolve() + info = _read_json(root / "meta" / "info.json") + tasks = { + int(item["task_index"]): item["task"] + for item in _read_jsonl(root / "meta" / "tasks.jsonl") + } + episode_meta = _read_jsonl(root / "meta" / "episodes.jsonl") + fps_values.add(int(info["fps"])) + self._validate_features(root, info["features"]) + + indices = list(range(len(episode_meta))) + if val_set_proportion >= 1e-6: + rng = np.random.default_rng(seed) + rng.shuffle(indices) + split = int(len(indices) * (1.0 - val_set_proportion)) + indices = indices[:split] if is_training_set else indices[split:] + + by_index = {int(item["episode_index"]): item for item in episode_meta} + for episode_index in indices: + item = by_index[episode_index] + episodes.append( + Episode( + root=root, + index=episode_index, + length=int(item["length"]), + data_path=info["data_path"], + video_path=info["video_path"], + chunks_size=int(info["chunks_size"]), + tasks=tasks, + ) + ) + + if not episodes: + raise RuntimeError("No LIBERO episodes were selected.") + if len(fps_values) != 1: + raise ValueError(f"All LIBERO datasets must have the same fps, got {sorted(fps_values)}") + self.fps = fps_values.pop() + self.episodes = episodes + self._episode_ends = np.cumsum([episode.length for episode in episodes]).tolist() + + def _validate_features(self, root, features): + required = [*self.image_keys, self.state_key, self.action_key, "timestamp", "task_index"] + missing = [key for key in required if key not in features] + if missing: + raise ValueError(f"LIBERO dataset {root} is missing features: {missing}") + non_video = [key for key in self.image_keys if features[key].get("dtype") != "video"] + if non_video: + raise ValueError(f"LIBERO camera features must be videos, got: {non_video}") + + def __len__(self): + return self._episode_ends[-1] + + @lru_cache(maxsize=32) + def _load_episode(self, episode_position): + episode = self.episodes[episode_position] + table = pq.read_table(episode.parquet_path()) + return {key: _column_to_tensor(table[key]) for key in table.column_names} + + def _locate(self, index): + if index < 0: + index += len(self) + if index < 0 or index >= len(self): + raise IndexError(f"Index {index} out of bounds for dataset of size {len(self)}") + episode_position = bisect.bisect_right(self._episode_ends, index) + episode_start = 0 if episode_position == 0 else self._episode_ends[episode_position - 1] + return episode_position, index - episode_start + + def _window(self, values, frame_index, length): + offsets = torch.arange(length, dtype=torch.long) * self.global_sample_stride + indices = frame_index + offsets + is_pad = indices >= values.shape[0] + indices.clamp_(max=values.shape[0] - 1) + return values[indices], is_pad, indices + + def __getitem__(self, index): + episode_position, frame_index = self._locate(index) + episode = self.episodes[episode_position] + data = self._load_episode(episode_position) + + state, state_is_pad, observation_indices = self._window( + data[self.state_key], frame_index, self.num_frames + ) + action, action_is_pad, _ = self._window( + data[self.action_key], frame_index, self.num_frames - 1 + ) + timestamps = data["timestamp"][observation_indices].float().tolist() + tolerance = max(1e-4, 1.0 / self.fps - 1e-4) + images = { + key: decode_video_frames( + episode.camera_path(key), + timestamps, + tolerance_s=tolerance, + backend=self.video_backend, + ) + for key in self.image_keys + } + + task_index = int(data["task_index"][frame_index].item()) + return { + "idx": int(index), + "task": episode.tasks[task_index], + "action": {"default": action.float()}, + "state": {"default": state.float()}, + "images": { + key.removeprefix("observation.images."): value + for key, value in images.items() + }, + "action_is_pad": action_is_pad, + "state_is_pad": state_is_pad, + "image_is_pad": state_is_pad.clone(), + } diff --git a/lightx2v_train/lightx2v_train/data/libero/normalizer.py b/lightx2v_train/lightx2v_train/data/libero/normalizer.py new file mode 100644 index 000000000..a18718991 --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/libero/normalizer.py @@ -0,0 +1,73 @@ +import json + +import torch + + +def _to_tensor_tree(value): + if isinstance(value, dict): + return {key: _to_tensor_tree(item) for key, item in value.items()} + if isinstance(value, list): + try: + return torch.tensor(value, dtype=torch.float32) + except (TypeError, ValueError): + return [_to_tensor_tree(item) for item in value] + return value + + +def load_dataset_stats(path): + with open(path, encoding="utf-8") as handle: + return _to_tensor_tree(json.load(handle)) + + +class FieldNormalizer: + def __init__(self, stats, mode="min/max"): + if mode == "z-score": + scale = 1.0 / (stats["std"] + 1e-8) + offset = -stats["mean"] * scale + else: + low_key, high_key = ("min", "max") if mode == "min/max" else ("q01", "q99") + lower = stats[low_key] + upper = stats[high_key] + width = upper - lower + constant = width < 1e-4 + width = width.clone() + width[constant] = 2.0 + scale = 2.0 / width + offset = -1.0 - scale * lower + offset[constant] = -lower[constant] + self.scale = scale + self.offset = offset + + def forward(self, value): + return (value * self.scale + self.offset).clamp_(-5.0, 5.0) + + def backward(self, value): + return (value - self.offset) / self.scale + + +class LinearNormalizer: + def __init__(self, shape_meta, stats, mode="min/max"): + self.normalizers = {"action": {}, "state": {}} + for group in self.normalizers: + for item in shape_meta[group]: + key = item["key"] + global_stats = { + name.removeprefix("global_"): value + for name, value in stats[group][key].items() + if name.startswith("global_") + } + self.normalizers[group][key] = FieldNormalizer(global_stats, mode=mode) + + def forward(self, batch): + for group, fields in self.normalizers.items(): + if group not in batch: + continue + for key, normalizer in fields.items(): + batch[group][key] = normalizer.forward(batch[group][key]) + return batch + + def backward(self, batch): + for group, fields in self.normalizers.items(): + for key, normalizer in fields.items(): + batch[group][key] = normalizer.backward(batch[group][key]) + return batch diff --git a/lightx2v_train/lightx2v_train/data/libero/processor.py b/lightx2v_train/lightx2v_train/data/libero/processor.py new file mode 100644 index 000000000..2f8359590 --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/libero/processor.py @@ -0,0 +1,115 @@ +import torch +import torchvision.transforms.functional as vision + +from .normalizer import LinearNormalizer + + +class ConcatLeftAlign: + def __init__(self, shape_meta): + self.action_meta = shape_meta["action"] + self.state_meta = shape_meta["state"] + + @staticmethod + def _merge(values, metadata): + return torch.cat([values[item["key"]] for item in metadata], dim=-1) + + @staticmethod + def _split(value, metadata): + result = {} + offset = 0 + for item in metadata: + width = int(item["shape"]) + result[item["key"]] = value[..., offset : offset + width] + offset += width + return result + + def forward(self, batch): + if "action" in batch: + batch["action"] = self._merge(batch["action"], self.action_meta) + batch["action_dim_is_pad"] = torch.zeros(batch["action"].shape[-1], dtype=torch.bool) + batch["state"] = self._merge(batch["state"], self.state_meta) + batch["state_dim_is_pad"] = torch.zeros(batch["state"].shape[-1], dtype=torch.bool) + return batch + + def backward(self, batch): + batch["action"] = self._split(batch["action"], self.action_meta) + batch["state"] = self._split(batch["state"], self.state_meta) + return batch + + +class FastWAMProcessor: + def __init__( + self, + shape_meta, + num_obs_steps, + image_size=(224, 224), + delta_action_dim_mask=(True, True, True, True, True, True, False), + ): + self.shape_meta = shape_meta + self.num_obs_steps = int(num_obs_steps) + self.image_size = tuple(image_size) + self.action_state_merger = ConcatLeftAlign(shape_meta) + self.delta_action_dim_mask = torch.as_tensor(delta_action_dim_mask, dtype=torch.bool) + self._normalizer = None + + @property + def normalizer(self): + if self._normalizer is None: + raise RuntimeError("FastWAM processor normalization statistics have not been loaded.") + return self._normalizer + + def set_normalizer_from_stats(self, stats): + self._normalizer = LinearNormalizer(self.shape_meta, stats) + + def _validate_action_state(self, data): + for group in ("action", "state"): + for item in self.shape_meta[group]: + actual = data[group][item["key"]].shape[-1] + expected = int(item["raw_shape"]) + if actual != expected: + raise ValueError( + f"{group}.{item['key']} width mismatch: expected {expected}, got {actual}" + ) + + def _process_images(self, images): + processed = [] + for item in self.shape_meta["images"]: + image = images[item["key"]] + if image.ndim != 4: + raise ValueError(f"Image sequence must be [T,C,H,W], got {tuple(image.shape)}") + image = vision.resize( + image.to(torch.float32), + self.image_size, + interpolation=vision.InterpolationMode.BILINEAR, + antialias=True, + ) + expected = (self.num_obs_steps, *tuple(item["shape"])) + if tuple(image.shape) != expected: + raise ValueError(f"Processed image shape mismatch: expected {expected}, got {tuple(image.shape)}") + processed.append(image) + return torch.stack(processed) + + def preprocess(self, data): + self._validate_action_state(data) + action_is_pad = torch.as_tensor(data["action_is_pad"], dtype=torch.bool) + action = data["action"]["default"] + if bool(action_is_pad.any()): + mask = action_is_pad[:, None] & self.delta_action_dim_mask[None, :] + action[mask] = 0.0 + + transformed = self.normalizer.forward( + {"action": data["action"], "state": data["state"]} + ) + transformed = self.action_state_merger.forward(transformed) + return { + "idx": data["idx"], + "instruction": str(data["task"]), + "pixel_values": self._process_images(data["images"]), + "image_is_pad": data["image_is_pad"], + "action": transformed["action"], + "action_is_pad": action_is_pad, + "action_dim_is_pad": transformed["action_dim_is_pad"], + "proprio": transformed["state"], + "proprio_is_pad": data["state_is_pad"], + "proprio_dim_is_pad": transformed["state_dim_is_pad"], + } diff --git a/lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py b/lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py new file mode 100644 index 000000000..ee088937f --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py @@ -0,0 +1,128 @@ +import hashlib +import os +import traceback + +import numpy as np +import torch +from torch.utils.data import Dataset + +from .lerobot_dataset import LiberoLeRobotDataset +from .normalizer import load_dataset_stats + +PROMPT_TEMPLATE = "A video recorded from a robot's point of view executing the following instruction: {task}" + + +class RobotVideoDataset(Dataset): + def __init__( + self, + dataset_dirs, + shape_meta, + processor, + text_embedding_cache_dir, + pretrained_norm_stats, + num_frames=33, + context_len=128, + val_set_proportion=0.0, + is_training_set=True, + global_sample_stride=1, + action_video_freq_ratio=4, + skip_padding_as_possible=False, + max_padding_retry=3, + video_backend=None, + ): + image_keys = [f"observation.images.{item['key']}" for item in shape_meta["images"]] + self.lerobot_dataset = LiberoLeRobotDataset( + dataset_dirs=dataset_dirs, + image_keys=image_keys, + state_key="observation.state", + action_key="action", + num_frames=num_frames, + global_sample_stride=global_sample_stride, + val_set_proportion=val_set_proportion, + is_training_set=is_training_set, + video_backend=video_backend, + ) + self.num_frames = int(num_frames) + self.action_video_freq_ratio = int(action_video_freq_ratio) + if (self.num_frames - 1) % self.action_video_freq_ratio: + raise ValueError("num_frames - 1 must be divisible by action_video_freq_ratio") + if ((self.num_frames - 1) // self.action_video_freq_ratio) % 4: + raise ValueError("The number of future video frames must be divisible by the VAE temporal factor 4") + self.video_sample_indices = list(range(0, self.num_frames, self.action_video_freq_ratio)) + self.context_len = int(context_len) + self.text_embedding_cache_dir = os.path.abspath(os.path.expanduser(text_embedding_cache_dir)) + self.skip_padding_as_possible = bool(skip_padding_as_possible) + self.max_padding_retry = int(max_padding_retry) + + stats_path = os.path.abspath(os.path.expanduser(pretrained_norm_stats)) + if not os.path.isfile(stats_path): + raise FileNotFoundError(f"LIBERO normalization stats do not exist: {stats_path}") + processor.set_normalizer_from_stats(load_dataset_stats(stats_path)) + self.processor = processor + + def __len__(self): + return len(self.lerobot_dataset) + + def _sample_without_padding(self, index): + for attempt in range(self.max_padding_retry + 1): + sample = self.processor.preprocess(self.lerobot_dataset[index]) + if not self.skip_padding_as_possible: + return sample + has_padding = any( + bool(sample[key].any()) + for key in ("action_is_pad", "image_is_pad", "proprio_is_pad") + ) + if not has_padding or attempt == self.max_padding_retry: + return sample + index = np.random.randint(len(self)) + raise AssertionError("unreachable") + + def _cached_context(self, prompt): + digest = hashlib.sha256(prompt.encode("utf-8")).hexdigest() + path = os.path.join( + self.text_embedding_cache_dir, + f"{digest}.t5_len{self.context_len}.wan22ti2v5b.pt", + ) + if not os.path.isfile(path): + raise FileNotFoundError(f"Missing cached text embedding: {path}") + payload = torch.load(path, map_location="cpu", weights_only=True) + context = payload["context"] + mask = payload["mask"].bool() + if context.ndim != 2 or context.shape[0] != self.context_len: + raise ValueError(f"Invalid cached context shape {tuple(context.shape)} in {path}") + if mask.shape != (self.context_len,): + raise ValueError(f"Invalid cached context mask shape {tuple(mask.shape)} in {path}") + context[~mask] = 0.0 + return context, torch.ones_like(mask) + + def _get(self, index): + sample = self._sample_without_padding(index) + video = sample["pixel_values"][:, self.video_sample_indices] + if video.shape[0] != 2: + raise ValueError(f"FastWAM LIBERO expects two cameras, got {video.shape[0]}") + video = torch.cat((video[0], video[1]), dim=-1) + video = video.mul(2.0).sub(1.0).permute(1, 0, 2, 3).contiguous() + + action = sample["action"] + proprio = sample["proprio"][:-1] + prompt = PROMPT_TEMPLATE.format(task=sample["instruction"]) + context, context_mask = self._cached_context(prompt) + return { + "video": video, + "action": action, + "proprio": proprio, + "prompt": prompt, + "context": context, + "context_mask": context_mask, + "image_is_pad": sample["image_is_pad"][self.video_sample_indices], + "action_is_pad": sample["action_is_pad"], + "proprio_is_pad": sample["proprio_is_pad"], + } + + def __getitem__(self, index): + try: + return self._get(index) + except Exception: + traceback.print_exc() + fallback = np.random.randint(len(self)) + return self._get(fallback) diff --git a/lightx2v_train/lightx2v_train/data/libero/video_decoder.py b/lightx2v_train/lightx2v_train/data/libero/video_decoder.py new file mode 100644 index 000000000..1a2e576c5 --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/libero/video_decoder.py @@ -0,0 +1,69 @@ +import importlib.util +import logging +import warnings +from pathlib import Path + +import torch +import torchvision + + +def _default_backend(): + if importlib.util.find_spec("torchcodec") is not None: + return "torchcodec" + return "pyav" + + +def decode_video_frames(video_path, timestamps, tolerance_s, backend=None): + backend = backend or _default_backend() + if backend == "torchcodec": + try: + return _decode_torchcodec(video_path, timestamps, tolerance_s) + except Exception as error: + warnings.warn( + f"torchcodec decode failed ({type(error).__name__}: {error}); falling back to pyav." + ) + backend = "pyav" + if backend not in {"pyav", "video_reader"}: + raise ValueError(f"Unsupported video backend: {backend}") + return _decode_torchvision(video_path, timestamps, tolerance_s, backend) + + +def _decode_torchvision(video_path, timestamps, tolerance_s, backend): + torchvision.set_video_backend(backend) + reader = torchvision.io.VideoReader(str(video_path), "video") + reader.seek(min(timestamps), keyframes_only=backend == "pyav") + frames = [] + frame_timestamps = [] + for frame in reader: + frames.append(frame["data"]) + frame_timestamps.append(frame["pts"]) + if frame["pts"] >= max(timestamps): + break + if backend == "pyav": + reader.container.close() + if not frames: + raise RuntimeError(f"No frames decoded from {video_path}") + return _select_frames(frames, frame_timestamps, timestamps, tolerance_s, video_path) + + +def _decode_torchcodec(video_path, timestamps, tolerance_s): + from torchcodec.decoders import VideoDecoder + + decoder = VideoDecoder(str(video_path), device="cpu", seek_mode="approximate") + fps = decoder.metadata.average_fps + batch = decoder.get_frames_at(indices=[round(timestamp * fps) for timestamp in timestamps]) + frames = list(batch.data) + frame_timestamps = batch.pts_seconds.tolist() + return _select_frames(frames, frame_timestamps, timestamps, tolerance_s, video_path) + + +def _select_frames(frames, frame_timestamps, query_timestamps, tolerance_s, video_path): + query = torch.as_tensor(query_timestamps, dtype=torch.float32) + available = torch.as_tensor(frame_timestamps, dtype=torch.float32) + distances = torch.cdist(query[:, None], available[:, None], p=1) + minimum, indices = distances.min(dim=1) + if not bool((minimum <= tolerance_s).all()): + logging.warning("Video timestamp mismatch for %s: max error %.6fs", video_path, minimum.max()) + raise RuntimeError(f"Unable to decode requested timestamps from {Path(video_path)}") + selected = torch.stack([frames[index] for index in indices.tolist()]) + return selected.to(torch.float32).div_(255.0) diff --git a/lightx2v_train/lightx2v_train/model_zoo/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/__init__.py index 623651261..151173cba 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/__init__.py +++ b/lightx2v_train/lightx2v_train/model_zoo/__init__.py @@ -5,7 +5,8 @@ from .longcat_image import LongCatImageModel from .qwen_image import QwenImageModel from .qwen_image_edit import QwenImageEditModel +from .wan_fastwam import WanFastWAMModel from .wan_t2v import WanT2VModel from .wan_ti2v_5b import WanTI2V5BModel -__all__ = ["build_model", "QwenImageModel", "QwenImageEditModel", "LongCatImageModel", "Flux2KleinModel", "WanT2VModel", "WanTI2V5BModel"] +__all__ = ["build_model", "QwenImageModel", "QwenImageEditModel", "LongCatImageModel", "Flux2KleinModel", "WanT2VModel", "WanTI2V5BModel", "WanFastWAMModel"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/__init__.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/__init__.py new file mode 100644 index 000000000..080d4c67e --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/__init__.py @@ -0,0 +1,3 @@ +from .model import FastWAM + +__all__ = ["FastWAM"] diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py new file mode 100644 index 000000000..1ac68dd88 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py @@ -0,0 +1,367 @@ +import logging +import os +from typing import Any, Dict, Optional + +import torch +import torch.nn as nn + +from ..modules.model2_2 import rope_params +from .layers import expert_block_forward, timestep_embedding + +logger = logging.getLogger(__name__) + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, value): + dtype = value.dtype + value = value.float() + value = value * torch.rsqrt(value.pow(2).mean(dim=-1, keepdim=True) + self.eps) + return value.to(dtype) * self.weight + + +class WideAttention(nn.Module): + """Attention projection whose width may differ from the expert hidden width.""" + + def __init__(self, hidden_dim, num_heads, head_dim, eps): + super().__init__() + attention_dim = num_heads * head_dim + self.num_heads = num_heads + self.head_dim = head_dim + self.q = nn.Linear(hidden_dim, attention_dim) + self.k = nn.Linear(hidden_dim, attention_dim) + self.v = nn.Linear(hidden_dim, attention_dim) + self.o = nn.Linear(attention_dim, hidden_dim) + self.norm_q = RMSNorm(attention_dim, eps) + self.norm_k = RMSNorm(attention_dim, eps) + + +class ActionBlock(nn.Module): + def __init__(self, hidden_dim, ffn_dim, num_heads, head_dim, eps): + super().__init__() + self.num_heads = num_heads + self.self_attn = WideAttention(hidden_dim, num_heads, head_dim, eps) + self.cross_attn = WideAttention(hidden_dim, num_heads, head_dim, eps) + self.norm1 = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False) + self.norm2 = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False) + self.norm3 = nn.LayerNorm(hidden_dim, eps=eps) + self.ffn = nn.Sequential( + nn.Linear(hidden_dim, ffn_dim), + nn.GELU(approximate="tanh"), + nn.Linear(ffn_dim, hidden_dim), + ) + self.modulation = nn.Parameter(torch.randn(1, 6, hidden_dim) / hidden_dim**0.5) + + +class ActionDiT(nn.Module): + ACTION_BACKBONE_SKIP_PREFIXES = ("action_encoder.", "head.") + ACTION_BACKBONE_META_KEYS = ( + "hidden_dim", + "ffn_dim", + "num_layers", + "num_heads", + "attn_head_dim", + "text_dim", + "freq_dim", + "eps", + ) + + def __init__( + self, + hidden_dim: int, + action_dim: int, + ffn_dim: int, + text_dim: int, + freq_dim: int, + eps: float, + num_heads: int, + attn_head_dim: int, + num_layers: int, + use_gradient_checkpointing: bool = False, + ): + super().__init__() + self.hidden_dim = hidden_dim + self.action_dim = action_dim + self.ffn_dim = ffn_dim + self.text_dim = text_dim + self.freq_dim = freq_dim + self.num_heads = num_heads + self.attn_head_dim = attn_head_dim + + if num_heads <= 0: + raise ValueError(f"`num_heads` must be > 0, got {num_heads}") + if attn_head_dim <= 0: + raise ValueError(f"`attn_head_dim` must be > 0, got {attn_head_dim}") + if attn_head_dim % 2 != 0: + raise ValueError(f"`attn_head_dim` must be even for RoPE, got {attn_head_dim}") + + self.action_encoder = nn.Linear(action_dim, hidden_dim) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, hidden_dim), + nn.GELU(approximate="tanh"), + nn.Linear(hidden_dim, hidden_dim), + ) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, hidden_dim), + nn.SiLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 6)) + self.blocks = nn.ModuleList( + [ + ActionBlock( + hidden_dim=hidden_dim, + ffn_dim=ffn_dim, + num_heads=num_heads, + head_dim=attn_head_dim, + eps=eps, + ) + for _ in range(num_layers) + ] + ) + self.head = nn.Linear(hidden_dim, action_dim) + self.freqs = rope_params(1024, attn_head_dim) + + self.use_gradient_checkpointing = use_gradient_checkpointing + + @classmethod + def backbone_key_set(cls, keys) -> set[str]: + return { + key + for key in keys + if not any(key.startswith(prefix) for prefix in cls.ACTION_BACKBONE_SKIP_PREFIXES) + } + + @classmethod + def from_pretrained( + cls, + action_dit_config: dict[str, Any], + action_dit_pretrained_path: str | None = None, + skip_dit_load_from_pretrain: bool = False, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + ) -> "ActionDiT": + if action_dit_config is None: + raise ValueError("`action_dit_config` is required for ActionDiT.from_pretrained().") + if skip_dit_load_from_pretrain: + logger.info( + "Skipping ActionDiT pretrained load (`skip_dit_load_from_pretrain=True`); " + "initializing action expert randomly and expecting checkpoint override." + ) + return cls(**action_dit_config).to(device=device, dtype=torch_dtype) + if not action_dit_pretrained_path: + logger.info("No `action_dit_pretrained_path` provided, initializing ActionDiT with random weights.") + return cls(**action_dit_config).to(device=device, dtype=torch_dtype) + from pathlib import Path + p = Path(action_dit_pretrained_path) + if not p.is_absolute(): + p = Path(__file__).resolve().parents[4] / p + action_dit_pretrained_path = str(p) + if not os.path.isfile(action_dit_pretrained_path): + raise FileNotFoundError( + f"`action_dit_pretrained_path` does not exist: {action_dit_pretrained_path}" + ) + + action_cfg = dict(action_dit_config) + action_expert = cls(**action_cfg).to(device=device, dtype=torch_dtype) + action_state = action_expert.state_dict() + expected_backbone_keys = cls.backbone_key_set(action_state.keys()) + + payload = torch.load(action_dit_pretrained_path, map_location="cpu", weights_only=True) + if not isinstance(payload, dict): + raise ValueError( + f"Invalid action backbone payload type from {action_dit_pretrained_path}: {type(payload)}" + ) + + policy = payload.get("policy", {}) + if policy: + logger.info(f"ActionDiT backbone payload policy: {policy}") + + meta = payload.get("meta") + expected_meta = { + "hidden_dim": int(action_cfg["hidden_dim"]), + "ffn_dim": int(action_cfg["ffn_dim"]), + "num_layers": int(action_cfg["num_layers"]), + "num_heads": int(action_cfg["num_heads"]), + "attn_head_dim": int(action_cfg["attn_head_dim"]), + "text_dim": int(action_cfg["text_dim"]), + "freq_dim": int(action_cfg["freq_dim"]), + "eps": float(action_cfg["eps"]), + } + for key in cls.ACTION_BACKBONE_META_KEYS: + if key not in meta: + raise ValueError(f"`meta.{key}` missing in {action_dit_pretrained_path}") + expected_value = expected_meta[key] + got_value = meta[key] + if key == "eps": + if abs(float(got_value) - float(expected_value)) > 1e-12: + raise ValueError( + f"`meta.{key}` mismatch in {action_dit_pretrained_path}: " + f"expected {expected_value}, got {got_value}" + ) + elif int(got_value) != int(expected_value): + raise ValueError( + f"`meta.{key}` mismatch in {action_dit_pretrained_path}: " + f"expected {expected_value}, got {got_value}" + ) + + backbone_state_dict = payload.get("backbone_state_dict") + if not isinstance(backbone_state_dict, dict): + raise ValueError( + f"`backbone_state_dict` must be a dict in {action_dit_pretrained_path}, " + f"got {type(backbone_state_dict)}" + ) + + provided_keys = set(backbone_state_dict.keys()) + missing_keys = sorted(expected_backbone_keys - provided_keys) + unexpected_keys = sorted(provided_keys - expected_backbone_keys) + if missing_keys or unexpected_keys: + raise ValueError( + "Action backbone key mismatch in preprocessed payload. " + f"missing={missing_keys[:10]}{'...' if len(missing_keys) > 10 else ''}, " + f"unexpected={unexpected_keys[:10]}{'...' if len(unexpected_keys) > 10 else ''}" + ) + + merged_state = dict(action_state) + for key in expected_backbone_keys: + value = backbone_state_dict[key] + if not isinstance(value, torch.Tensor): + raise ValueError( + f"`backbone_state_dict[{key}]` must be torch.Tensor in {action_dit_pretrained_path}, " + f"got {type(value)}" + ) + target = merged_state[key] + if tuple(value.shape) != tuple(target.shape): + raise ValueError( + f"Shape mismatch for `{key}` in {action_dit_pretrained_path}: " + f"expected {tuple(target.shape)}, got {tuple(value.shape)}" + ) + merged_state[key] = value.to(device=target.device, dtype=target.dtype) + + action_expert.load_state_dict(merged_state, strict=True) + logger.info( + "Loaded ActionDiT backbone from %s (keys=%d; random_kept_prefixes=%s).", + action_dit_pretrained_path, + len(expected_backbone_keys), + list(cls.ACTION_BACKBONE_SKIP_PREFIXES), + ) + return action_expert.to(device=device, dtype=torch_dtype) + + def pre_dit( + self, + action_tokens: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: Optional[torch.Tensor] = None, + ) -> Dict[str, Any]: + if action_tokens.ndim != 3: + raise ValueError( + f"`action_tokens` must be 3D [B, T, action_dim], got shape {tuple(action_tokens.shape)}" + ) + if action_tokens.shape[2] != self.action_dim: + raise ValueError( + f"`action_tokens` last dim must be {self.action_dim}, got {action_tokens.shape[2]}" + ) + if timestep.ndim != 1: + raise ValueError(f"`timestep` must be 1D [B] or [1], got shape {tuple(timestep.shape)}") + if context.ndim != 3: + raise ValueError( + f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}" + ) + + batch_size = action_tokens.shape[0] + if context.shape[0] != batch_size: + raise ValueError( + f"Batch mismatch between action tokens and text context: {batch_size} vs {context.shape[0]}" + ) + if timestep.shape[0] not in (1, batch_size): + raise ValueError( + f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}" + ) + if timestep.shape[0] == 1 and batch_size > 1: + if self.training: + raise ValueError("During training, action timestep length must match batch_size.") + timestep = timestep.expand(batch_size) + + if context_mask is None: + context_mask = torch.ones( + (batch_size, context.shape[1]), dtype=torch.bool, device=context.device + ) + else: + if context_mask.ndim != 2: + raise ValueError(f"`context_mask` must be 2D [B, L], got shape {tuple(context_mask.shape)}") + if context_mask.shape[0] != batch_size or context_mask.shape[1] != context.shape[1]: + raise ValueError( + f"`context_mask` shape must match `context` shape [B, L], got {tuple(context_mask.shape)} vs {tuple(context.shape)}" + ) + + seq_len = action_tokens.shape[1] + if seq_len > self.freqs.shape[0]: + raise ValueError( + f"Action token length {seq_len} exceeds RoPE cache {self.freqs.shape[0]}." + ) + + t = self.time_embedding(timestep_embedding(self.freq_dim, timestep)) + t_mod = self.time_projection(t).unflatten(1, (6, self.hidden_dim)) + + tokens = self.action_encoder(action_tokens) + context_emb = self.text_embedding(context) + context_attn_mask = context_mask.unsqueeze(1).expand(-1, seq_len, -1) + freqs = self.freqs[:seq_len].view(seq_len, 1, -1).to(tokens.device) + + return { + "tokens": tokens, + "freqs": freqs, + "t": t, + "t_mod": t_mod, + "context": context_emb, + "context_mask": context_attn_mask, + "meta": { + "batch_size": batch_size, + "seq_len": seq_len, + }, + } + + def post_dit(self, tokens: torch.Tensor, pre_state: Dict[str, Any]) -> torch.Tensor: + return self.head(tokens) + + def forward( + self, + action_tokens: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + pre_state = self.pre_dit( + action_tokens=action_tokens, + timestep=timestep, + context=context, + context_mask=context_mask, + ) + x = pre_state["tokens"] + context = pre_state["context"] + t_mod = pre_state["t_mod"] + freqs = pre_state["freqs"] + context_mask = pre_state["context_mask"] + + for block in self.blocks: + def run(value, current_block=block): + return expert_block_forward( + current_block, + value, + context, + t_mod, + freqs, + context_mask, + ) + + if self.use_gradient_checkpointing and self.training: + x = torch.utils.checkpoint.checkpoint(run, x, use_reentrant=False) + else: + x = run(x) + + return self.post_dit(x, pre_state) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py new file mode 100644 index 000000000..338431aa3 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py @@ -0,0 +1,110 @@ +import torch +import torch.nn.functional as F + + +def timestep_embedding(dim: int, position: torch.Tensor) -> torch.Tensor: + half = dim // 2 + frequencies = torch.pow( + 10000, + -torch.arange(half, dtype=torch.float64, device=position.device) / half, + ) + embedding = torch.outer(position.to(torch.float64), frequencies) + return torch.cat([torch.cos(embedding), torch.sin(embedding)], dim=1).to(position.dtype) + + +def apply_rope(x: torch.Tensor, frequencies: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + complex_x = torch.view_as_complex(x.to(torch.float64).reshape(*x.shape[:-1], -1, 2)) + complex_frequencies = frequencies.to(device=x.device, dtype=torch.complex128) + return torch.view_as_real(complex_x * complex_frequencies).flatten(-2).to(dtype) + + +def scaled_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor | None = None, +) -> torch.Tensor: + attention_mask = None + if mask is not None: + attention_mask = mask.to(device=query.device, dtype=torch.bool) + if attention_mask.ndim == 2: + attention_mask = attention_mask.unsqueeze(0).unsqueeze(0) + elif attention_mask.ndim == 3: + attention_mask = attention_mask.unsqueeze(1) + elif attention_mask.ndim != 4: + raise ValueError(f"Attention mask must be 2D, 3D, or 4D, got {attention_mask.ndim}D") + + output = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=attention_mask, + ) + return output.transpose(1, 2).contiguous() + + +def split_modulation(block, time_modulation: torch.Tensor): + has_sequence_dimension = time_modulation.ndim == 4 + chunk_dimension = 2 if has_sequence_dimension else 1 + values = (block.modulation.to(time_modulation) + time_modulation).chunk(6, dim=chunk_dimension) + if has_sequence_dimension: + values = tuple(value.squeeze(2) for value in values) + return values + + +def layer_norm(module, value: torch.Tensor) -> torch.Tensor: + if not isinstance(module, torch.nn.LayerNorm): + return module(value) + weight = None if module.weight is None else module.weight.float() + bias = None if module.bias is None else module.bias.float() + return F.layer_norm( + value.float(), module.normalized_shape, weight, bias, module.eps + ).to(value.dtype) + + +def masked_cross_attention(block, x, context, context_mask=None): + batch_size, _, num_heads = x.shape[0], x.shape[1], block.cross_attn.num_heads + head_dim = block.cross_attn.head_dim + query = block.cross_attn.norm_q(block.cross_attn.q(x)).view(batch_size, -1, num_heads, head_dim) + key = block.cross_attn.norm_k(block.cross_attn.k(context)).view(batch_size, -1, num_heads, head_dim) + value = block.cross_attn.v(context).view(batch_size, -1, num_heads, head_dim) + output = scaled_attention(query, key, value, context_mask) + return block.cross_attn.o(output.flatten(2)) + + +def expert_block_forward( + block, + x: torch.Tensor, + context: torch.Tensor, + time_modulation: torch.Tensor, + frequencies: torch.Tensor, + context_mask: torch.Tensor | None = None, + self_attention_mask: torch.Tensor | None = None, +) -> torch.Tensor: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = split_modulation( + block, time_modulation + ) + attention_input = layer_norm(block.norm1, x) * (1 + scale_msa) + shift_msa + batch_size, sequence_length = attention_input.shape[:2] + num_heads = block.self_attn.num_heads + head_dim = block.self_attn.head_dim + query = block.self_attn.norm_q(block.self_attn.q(attention_input)).view( + batch_size, sequence_length, num_heads, head_dim + ) + key = block.self_attn.norm_k(block.self_attn.k(attention_input)).view( + batch_size, sequence_length, num_heads, head_dim + ) + value = block.self_attn.v(attention_input).view( + batch_size, sequence_length, num_heads, head_dim + ) + attention = scaled_attention( + apply_rope(query, frequencies), + apply_rope(key, frequencies), + value, + self_attention_mask, + ) + x = x + gate_msa * block.self_attn.o(attention.flatten(2)) + x = x + masked_cross_attention(block, layer_norm(block.norm3, x), context, context_mask) + mlp_input = layer_norm(block.norm2, x) * (1 + scale_mlp) + shift_mlp + return x + gate_mlp * block.ffn(mlp_input) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py new file mode 100644 index 000000000..9da059a6e --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py @@ -0,0 +1,127 @@ +import inspect +import logging +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from safetensors.torch import load_file + +from ..modules.t5 import T5EncoderModel +from ..modules.vae2_2 import Wan2_2_VAE +from .video_dit import FastWAMVideoDiT + +logger = logging.getLogger(__name__) +@dataclass +class Wan22LoadedComponents: + dit: FastWAMVideoDiT + vae: Wan2_2_VAE + text_encoder: T5EncoderModel | None + tokenizer: Any | None + + +def _validate_dit_config(dit_config: dict[str, Any]) -> dict[str, Any]: + if not isinstance(dit_config, dict): + raise ValueError(f"video_dit_config must be a dict, got {type(dit_config)}") + signature = inspect.signature(FastWAMVideoDiT.__init__) + allowed = {name for name in signature.parameters if name != "self"} + required = { + name + for name, parameter in signature.parameters.items() + if name != "self" and parameter.default is inspect.Signature.empty + } + unknown = sorted(set(dit_config) - allowed) + missing = sorted(required - set(dit_config)) + if unknown: + raise ValueError(f"Unknown video_dit_config keys: {unknown}") + if missing: + raise ValueError(f"Missing video_dit_config keys: {missing}") + return dict(dit_config) + + +def _require_path(path: Path, description: str, directory: bool = False) -> str: + exists = path.is_dir() if directory else path.is_file() + if not exists: + kind = "directory" if directory else "file" + raise FileNotFoundError(f"Wan model {description} {kind} does not exist: {path}") + return str(path) + + +def _resolve_model_paths(model_path: str): + root = Path(model_path).expanduser().resolve() + _require_path(root, "root", directory=True) + single_dit = root / "diffusion_pytorch_model.safetensors" + if single_dit.is_file(): + dit_path: str | list[str] = str(single_dit) + else: + shards = sorted(root.glob("diffusion_pytorch_model-*-of-*.safetensors")) + if not shards: + raise FileNotFoundError(f"Wan diffusion weights do not exist under: {root}") + dit_path = [str(path) for path in shards] + return ( + dit_path, + _require_path(root / "models_t5_umt5-xxl-enc-bf16.pth", "text encoder"), + _require_path(root / "Wan2.2_VAE.pth", "VAE"), + _require_path(root / "google" / "umt5-xxl", "tokenizer", directory=True), + ) + + +def _load_video_expert(path, config, torch_dtype, device): + model = FastWAMVideoDiT(**config) + paths = path if isinstance(path, list) else [path] + state_dict = {} + for shard_path in paths: + shard = load_file(shard_path, device="cpu") + state_dict.update({key: value.to(torch_dtype) for key, value in shard.items()}) + model.load_state_dict(state_dict, strict=True) + return model.to(device=device, dtype=torch_dtype) + + +def load_wan22_ti2v_5b_components( + model_path: str, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + tokenizer_max_len: int = 512, + dit_config: dict[str, Any] | None = None, + skip_dit_load_from_pretrain: bool = False, + load_text_encoder: bool = True, +): + start = time.time() + if dit_config is None: + raise ValueError("video_dit_config is required for Wan2.2-TI2V-5B loading.") + config = _validate_dit_config(dit_config) + dit_path, text_encoder_path, vae_path, tokenizer_path = _resolve_model_paths(model_path) + + if skip_dit_load_from_pretrain: + dit = FastWAMVideoDiT(**config).to(device=device, dtype=torch_dtype) + else: + dit = _load_video_expert(dit_path, config, torch_dtype, device) + + text_encoder = None + tokenizer = None + if load_text_encoder: + text_encoder = T5EncoderModel( + text_len=int(tokenizer_max_len), + dtype=torch_dtype, + device=torch.device(device), + checkpoint_path=text_encoder_path, + tokenizer_path=tokenizer_path, + ) + tokenizer = text_encoder.tokenizer + + vae = Wan2_2_VAE( + vae_pth=vae_path, + dtype=torch_dtype, + device=device, + ) + vae.model.to(device=device, dtype=torch_dtype) + vae.temporal_downsample_factor = 4 + vae.upsampling_factor = 16 + logger.info("Loaded Wan2.2 FastWAM components in %.2f seconds", time.time() - start) + return Wan22LoadedComponents( + dit=dit, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + ) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py new file mode 100644 index 000000000..3426eefb7 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py @@ -0,0 +1,847 @@ +import logging +from typing import Any, Optional, Sequence, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from PIL import Image + +from lightx2v_train.schedulers.flow_matching import WanContinuousFlowMatchScheduler + +from .action_dit import ActionDiT +from .loader import load_wan22_ti2v_5b_components +from .mot import MoT + +logger = logging.getLogger(__name__) + +WAN22_TI2V_5B_VIDEO_CONFIG = { + "has_image_input": False, + "patch_size": (1, 2, 2), + "in_dim": 48, + "hidden_dim": 3072, + "ffn_dim": 14336, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 48, + "num_heads": 24, + "attn_head_dim": 128, + "num_layers": 30, + "eps": 1e-6, + "seperated_timestep": True, + "fuse_vae_embedding_in_latents": True, + "video_attention_mask_mode": "first_frame_causal", +} + +FASTWAM_ACTION_CONFIG = { + "action_dim": 7, + "hidden_dim": 1024, + "ffn_dim": 4096, + "num_heads": 24, + "attn_head_dim": 128, + "num_layers": 30, + "text_dim": 4096, + "freq_dim": 256, + "eps": 1e-6, +} + + +class FastWAM(torch.nn.Module): + """MoT world model with video/action experts.""" + + def __init__( + self, + video_expert, + action_expert: ActionDiT, + mot: MoT, + vae, + text_encoder=None, + tokenizer=None, + text_dim: Optional[int] = None, + proprio_dim: Optional[int] = None, + device: str = "cpu", + torch_dtype: torch.dtype = torch.float32, + video_train_shift: float = 5.0, + video_infer_shift: float = 5.0, + video_num_train_timesteps: int = 1000, + action_train_shift: float = 5.0, + action_infer_shift: float = 5.0, + action_num_train_timesteps: int = 1000, + loss_lambda_video: float = 1.0, + loss_lambda_action: float = 1.0, + ): + super().__init__() + self.video_expert = video_expert + self.action_expert = action_expert + self.mot = mot + # Keep trainer compatibility: optimizer and freeze logic use `model.dit`. + self.dit = self.mot + + self.vae = vae + self.text_encoder = text_encoder + self.tokenizer = tokenizer + if text_dim is None: + if self.text_encoder is None: + raise ValueError("`text_dim` is required when `text_encoder` is not loaded.") + text_dim = int(self.text_encoder.model.dim) + self.text_dim = int(text_dim) + self.proprio_dim = None if proprio_dim is None else int(proprio_dim) + if self.proprio_dim is not None: + self.proprio_encoder = nn.Linear(self.proprio_dim, self.text_dim).to(torch_dtype) + else: + self.proprio_encoder = None + + self.train_video_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=video_num_train_timesteps, + shift=video_train_shift, + ) + self.infer_video_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=video_num_train_timesteps, + shift=video_infer_shift, + ) + self.train_action_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=action_num_train_timesteps, + shift=action_train_shift, + ) + self.infer_action_scheduler = WanContinuousFlowMatchScheduler( + num_train_timesteps=action_num_train_timesteps, + shift=action_infer_shift, + ) + self.device = torch.device(device) + self.torch_dtype = torch_dtype + self.loss_lambda_video = float(loss_lambda_video) + self.loss_lambda_action = float(loss_lambda_action) + + self.to(self.device) + + @classmethod + def from_wan22_pretrained( + cls, + model_path: str, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + tokenizer_max_len: int = 512, + load_text_encoder: bool = True, + proprio_dim: Optional[int] = None, + video_dit_config: dict[str, Any] | None = None, + action_dit_config: dict[str, Any] | None = None, + action_dit_pretrained_path: str | None = None, + skip_dit_load_from_pretrain: bool = False, + mot_checkpoint_mixed_attn: bool = True, + video_train_shift: float = 5.0, + video_infer_shift: float = 5.0, + video_num_train_timesteps: int = 1000, + action_train_shift: float = 5.0, + action_infer_shift: float = 5.0, + action_num_train_timesteps: int = 1000, + loss_lambda_video: float = 1.0, + loss_lambda_action: float = 1.0, + ): + video_dit_config = {**WAN22_TI2V_5B_VIDEO_CONFIG, **(video_dit_config or {})} + action_dit_config = {**FASTWAM_ACTION_CONFIG, **(action_dit_config or {})} + + components = load_wan22_ti2v_5b_components( + model_path=model_path, + device=device, + torch_dtype=torch_dtype, + tokenizer_max_len=tokenizer_max_len, + dit_config=video_dit_config, + skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, + load_text_encoder=load_text_encoder, + ) + + video_expert = components.dit + action_expert = ActionDiT.from_pretrained( + action_dit_config=action_dit_config, + action_dit_pretrained_path=action_dit_pretrained_path, + skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, + device=device, + torch_dtype=torch_dtype, + ) + if int(action_expert.num_heads) != int(video_expert.num_heads): + raise ValueError("ActionDiT `num_heads` must match video expert for MoT mixed attention.") + if int(action_expert.attn_head_dim) != int(video_expert.attn_head_dim): + raise ValueError("ActionDiT `attn_head_dim` must match video expert for MoT mixed attention.") + if int(len(action_expert.blocks)) != int(len(video_expert.blocks)): + raise ValueError("ActionDiT `num_layers` must match video expert.") + + mot = MoT( + mixtures={"video": video_expert, "action": action_expert}, + mot_checkpoint_mixed_attn=mot_checkpoint_mixed_attn, + ) + + model = cls( + video_expert=video_expert, + action_expert=action_expert, + mot=mot, + vae=components.vae, + text_encoder=components.text_encoder, + tokenizer=components.tokenizer, + text_dim=int(video_dit_config["text_dim"]), + proprio_dim=proprio_dim, + device=device, + torch_dtype=torch_dtype, + video_train_shift=video_train_shift, + video_infer_shift=video_infer_shift, + video_num_train_timesteps=video_num_train_timesteps, + action_train_shift=action_train_shift, + action_infer_shift=action_infer_shift, + action_num_train_timesteps=action_num_train_timesteps, + loss_lambda_video=loss_lambda_video, + loss_lambda_action=loss_lambda_action, + ) + return model + + def to(self, *args, **kwargs): + super().to(*args, **kwargs) + self.mot.to(*args, **kwargs) + if self.text_encoder is not None: + self.text_encoder.model.to(*args, **kwargs) + self.vae.model.to(*args, **kwargs) + return self + + @staticmethod + def _check_resize_height_width(height, width, num_frames): + if height % 16 != 0: + height = (height + 15) // 16 * 16 + if width % 16 != 0: + width = (width + 15) // 16 * 16 + if num_frames % 4 != 1: + num_frames = (num_frames + 3) // 4 * 4 + 1 + return height, width, num_frames + + @torch.no_grad() + def encode_prompt(self, prompt: Union[str, Sequence[str]]): + if self.text_encoder is None or self.tokenizer is None: + raise ValueError( + "Prompt encoding requires loaded text encoder/tokenizer. " + "Set `load_text_encoder=true` or provide precomputed `context/context_mask`." + ) + ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True) + ids = ids.to(self.device) + mask = mask.to(self.device, dtype=torch.bool) + prompt_emb = self.text_encoder.model(ids, mask) + # FIXME: original implementation's zero padding is visible in cross-attn. + seq_lens = mask.gt(0).sum(dim=1).long() + for i, v in enumerate(seq_lens): + prompt_emb[i, v:] = 0 + mask = torch.ones_like(mask) + return prompt_emb.to(device=self.device), mask + + def _append_proprio_to_context( + self, + context: torch.Tensor, + context_mask: torch.Tensor, + proprio: Optional[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.proprio_encoder is None or proprio is None: + return context, context_mask + if proprio.ndim != 2: + raise ValueError(f"`proprio` must be 2D [B, D], got shape {tuple(proprio.shape)}") + if self.proprio_dim is None or proprio.shape[1] != self.proprio_dim: + raise ValueError( + f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}" + ) + proprio_token = self.proprio_encoder( + proprio.to(device=self.device, dtype=context.dtype).unsqueeze(1) + ).to(dtype=context.dtype) # [B, 1, D] + proprio_mask = torch.ones((context_mask.shape[0], 1), dtype=torch.bool, device=context_mask.device) + return ( + torch.cat([context, proprio_token], dim=1), + torch.cat([context_mask, proprio_mask], dim=1), + ) + + @torch.no_grad() + def _encode_video_latents(self, video_tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): + del tile_size, tile_stride + if tiled: + raise ValueError("Wan2.2 training VAE does not support tiled encoding.") + latents = self.vae.encode([video for video in video_tensor]) + return torch.stack(latents).to(device=self.device, dtype=self.torch_dtype) + + @torch.no_grad() + def _encode_input_image_latents_tensor(self, input_image: torch.Tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): + if input_image.ndim == 3: + input_image = input_image.unsqueeze(0) + if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: + raise ValueError( + f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" + ) + image = input_image.to(device=self.device)[0].unsqueeze(1) + del tile_size, tile_stride + if tiled: + raise ValueError("Wan2.2 training VAE does not support tiled encoding.") + return torch.stack(self.vae.encode([image])).to(device=self.device, dtype=self.torch_dtype) + + def _decode_latents(self, latents, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): + del tile_size, tile_stride + if tiled: + raise ValueError("Wan2.2 training VAE does not support tiled decoding.") + video_tensor = torch.stack(self.vae.decode([latent for latent in latents])) + video_tensor = video_tensor.squeeze(0).detach().float().clamp(-1, 1) + video_tensor = ((video_tensor + 1.0) * 127.5).to(torch.uint8).cpu() + frames = [] + for t in range(video_tensor.shape[1]): + frame = video_tensor[:, t].permute(1, 2, 0).numpy() + frames.append(Image.fromarray(frame)) + return frames + + def build_inputs(self, sample, tiled: bool = False): + video = sample["video"] + if "context" not in sample or "context_mask" not in sample: + raise ValueError( + "FastWAM training requires `sample['context']` and `sample['context_mask']`." + ) + context = sample["context"] + context_mask = sample["context_mask"] + proprio = sample.get("proprio", None) + if video.ndim != 5: + raise ValueError(f"`sample['video']` must be 5D [B, 3, T, H, W], got shape {tuple(video.shape)}") + if video.shape[1] != 3: + raise ValueError(f"`sample['video']` channel dimension must be 3, got shape {tuple(video.shape)}") + + batch_size, _, num_frames, height, width = video.shape + if height % 16 != 0 or width % 16 != 0: + raise ValueError( + f"Video spatial dims must be multiples of 16, got H={height}, W={width}" + ) + if num_frames % 4 != 1: + raise ValueError(f"Video T must satisfy T % 4 == 1, got T={num_frames}") + if num_frames <= 1: + raise ValueError(f"Video T must be > 1 for action-conditioned training, got T={num_frames}") + + if "action" not in sample: + raise ValueError("`sample['action']` is required for FastWAM training.") + + action = sample["action"] + if action.ndim != 3: + raise ValueError(f"`sample['action']` must be 3D [B, T, a_dim], got shape {tuple(action.shape)}") + action_horizon = int(action.shape[1]) + if action_horizon % (num_frames - 1) != 0: + raise ValueError( + f"`sample['action']` temporal dimension must be divisible by video transitions ({num_frames - 1}), got {action_horizon}" + ) + + action_is_pad = sample.get("action_is_pad", None) + if action_is_pad is not None: + if action_is_pad.ndim != 2: + raise ValueError( + f"`sample['action_is_pad']` must be 2D [B, T], got shape {tuple(action_is_pad.shape)}" + ) + if action_is_pad.shape[0] != batch_size or action_is_pad.shape[1] != action_horizon: + raise ValueError( + "`sample['action_is_pad']` shape mismatch: " + f"got {tuple(action_is_pad.shape)} vs expected ({batch_size}, {action_horizon})" + ) + + image_is_pad = sample.get("image_is_pad", None) + if image_is_pad is not None: + if image_is_pad.ndim != 2: + raise ValueError( + f"`sample['image_is_pad']` must be 2D [B, T], got shape {tuple(image_is_pad.shape)}" + ) + if image_is_pad.shape[0] != batch_size or image_is_pad.shape[1] != num_frames: + raise ValueError( + "`sample['image_is_pad']` shape mismatch: " + f"got {tuple(image_is_pad.shape)} vs expected ({batch_size}, {num_frames})" + ) + + input_video = video.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + input_latents = self._encode_video_latents(input_video, tiled=tiled) + + first_frame_latents = None + fuse_flag = False + if getattr(self.video_expert, "fuse_vae_embedding_in_latents", False): + first_frame_latents = input_latents[:, :, 0:1] + fuse_flag = True + + if context.ndim != 3 or context_mask.ndim != 2: + raise ValueError( + f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" + ) + context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) + if self.proprio_encoder is not None: + if proprio is None: + raise ValueError("`sample['proprio']` is required when `proprio_dim` is enabled.") + if proprio.ndim != 3: + raise ValueError(f"`sample['proprio']` must be 3D [B, T, d], got shape {tuple(proprio.shape)}") + if proprio.shape[2] != self.proprio_dim: + raise ValueError( + f"`sample['proprio']` last dim must be {self.proprio_dim}, got {proprio.shape[2]}" + ) + proprio = proprio[:, 0, :] # [B, D] + context, context_mask = self._append_proprio_to_context( + context=context, + context_mask=context_mask, + proprio=proprio.to(device=self.device, dtype=self.torch_dtype), + ) + action = action.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + + if action_is_pad is not None: + action_is_pad = action_is_pad.to(device=self.device, dtype=torch.bool, non_blocking=True) + if image_is_pad is not None: + image_is_pad = image_is_pad.to(device=self.device, dtype=torch.bool, non_blocking=True) + + return { + "context": context, + "context_mask": context_mask, + "input_latents": input_latents, + "first_frame_latents": first_frame_latents, + "fuse_vae_embedding_in_latents": fuse_flag, + "action": action, + "action_is_pad": action_is_pad, + "image_is_pad": image_is_pad, + } + + @torch.no_grad() + def _build_mot_attention_mask( + self, + video_seq_len: int, + action_seq_len: int, + video_tokens_per_frame: int, + device: torch.device, + ) -> torch.Tensor: + total_seq_len = video_seq_len + action_seq_len + mask = torch.zeros((total_seq_len, total_seq_len), dtype=torch.bool, device=device) + + # video -> video + mask[:video_seq_len, :video_seq_len] = self.video_expert.build_video_to_video_mask( + video_seq_len=video_seq_len, + video_tokens_per_frame=video_tokens_per_frame, + device=device, + ) + # action -> action + mask[video_seq_len:, video_seq_len:] = True + # action -> first-frame video only + first_frame_tokens = min(video_tokens_per_frame, video_seq_len) + mask[video_seq_len:, :first_frame_tokens] = True + return mask + + def _compute_video_loss_per_sample( + self, + pred_video: torch.Tensor, + target_video: torch.Tensor, + image_is_pad: Optional[torch.Tensor], + include_initial_video_step: bool, + ) -> torch.Tensor: + video_loss_token = F.mse_loss(pred_video.float(), target_video.float(), reduction="none").mean(dim=(1, 3, 4)) + if image_is_pad is None: + return video_loss_token.mean(dim=1) + + temporal_factor = int(self.vae.temporal_downsample_factor) + if temporal_factor <= 0: + raise ValueError(f"`vae.temporal_downsample_factor` must be positive, got {temporal_factor}.") + if image_is_pad.shape[1] < 1: + raise ValueError("`image_is_pad` must contain at least one frame.") + if (image_is_pad.shape[1] - 1) % temporal_factor != 0: + raise ValueError( + "Cannot align `image_is_pad` with video latent steps: " + f"num_frames={image_is_pad.shape[1]}, temporal_downsample_factor={temporal_factor}." + ) + + tail_is_pad = image_is_pad[:, 1:] + latent_tail_is_pad = tail_is_pad.view(image_is_pad.shape[0], -1, temporal_factor).all(dim=2) + if include_initial_video_step: + video_is_pad = torch.cat([image_is_pad[:, :1], latent_tail_is_pad], dim=1) + else: + video_is_pad = latent_tail_is_pad + + if video_is_pad.shape[1] != video_loss_token.shape[1]: + raise ValueError( + "Video-loss mask shape mismatch: " + f"mask steps={video_is_pad.shape[1]}, loss steps={video_loss_token.shape[1]}." + ) + + valid = (~video_is_pad).to(device=video_loss_token.device, dtype=video_loss_token.dtype) + valid_sum = valid.sum(dim=1).clamp(min=1.0) + return (video_loss_token * valid).sum(dim=1) / valid_sum + + def training_loss(self, sample, tiled: bool = False): + inputs = self.build_inputs(sample, tiled=tiled) + input_latents = inputs["input_latents"] + batch_size = input_latents.shape[0] + context = inputs["context"] + context_mask = inputs["context_mask"] + action = inputs["action"] + action_is_pad = inputs["action_is_pad"] + image_is_pad = inputs["image_is_pad"] + + noise_video = torch.randn_like(input_latents) + timestep_video = self.train_video_scheduler.sample_training_t( + batch_size=batch_size, + device=self.device, + dtype=input_latents.dtype, + ) + latents = self.train_video_scheduler.add_noise(input_latents, noise_video, timestep_video) + target_video = self.train_video_scheduler.training_target(input_latents, noise_video, timestep_video) + + if inputs["first_frame_latents"] is not None: + latents[:, :, 0:1] = inputs["first_frame_latents"] + + noise_action = torch.randn_like(action) + timestep_action = self.train_action_scheduler.sample_training_t( + batch_size=batch_size, + device=self.device, + dtype=action.dtype, + ) + noisy_action = self.train_action_scheduler.add_noise(action, noise_action, timestep_action) + target_action = self.train_action_scheduler.training_target(action, noise_action, timestep_action) + + video_pre = self.video_expert.pre_dit( + x=latents, + timestep=timestep_video, + context=context, + context_mask=context_mask, + action=action, + fuse_vae_embedding_in_latents=inputs["fuse_vae_embedding_in_latents"], + ) + + action_pre = self.action_expert.pre_dit( + action_tokens=noisy_action, + timestep=timestep_action, + context=context, + context_mask=context_mask, + ) + + video_tokens = video_pre["tokens"] + action_tokens = action_pre["tokens"] + + attention_mask = self._build_mot_attention_mask( + video_seq_len=video_tokens.shape[1], + action_seq_len=action_tokens.shape[1], + video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), + device=video_tokens.device, + ) + tokens_out = self.mot( + embeds_all={ + "video": video_tokens, + "action": action_tokens, + }, + attention_mask=attention_mask, + freqs_all={ + "video": video_pre["freqs"], + "action": action_pre["freqs"], + }, + context_all={ + "video": { + "context": video_pre["context"], + "mask": video_pre["context_mask"], + }, + "action": { + "context": action_pre["context"], + "mask": action_pre["context_mask"], + }, + }, + t_mod_all={ + "video": video_pre["t_mod"], + "action": action_pre["t_mod"], + }, + ) + + pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) + + pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) + + include_initial_video_step = inputs["first_frame_latents"] is None + if inputs["first_frame_latents"] is not None: + pred_video = pred_video[:, :, 1:] + target_video = target_video[:, :, 1:] + + loss_video_per_sample = self._compute_video_loss_per_sample( + pred_video=pred_video, + target_video=target_video, + image_is_pad=image_is_pad, + include_initial_video_step=include_initial_video_step, + ) + video_weight = self.train_video_scheduler.training_weight(timestep_video).to( + loss_video_per_sample.device, dtype=loss_video_per_sample.dtype + ) + loss_video = (loss_video_per_sample * video_weight).mean() + + action_loss_token = F.mse_loss(pred_action.float(), target_action.float(), reduction="none").mean(dim=2) # [B, T] + if action_is_pad is not None: + valid = (~action_is_pad).to(device=action_loss_token.device, dtype=action_loss_token.dtype) + valid_sum = valid.sum(dim=1).clamp(min=1.0) + action_loss_per_sample = (action_loss_token * valid).sum(dim=1) / valid_sum + else: + action_loss_per_sample = action_loss_token.mean(dim=1) + + action_weight = self.train_action_scheduler.training_weight(timestep_action).to( + action_loss_per_sample.device, dtype=action_loss_per_sample.dtype + ) + loss_action = (action_loss_per_sample * action_weight).mean() + + loss_total = self.loss_lambda_video * loss_video + self.loss_lambda_action * loss_action + loss_dict = { + "loss_video": self.loss_lambda_video * float(loss_video.detach().item()), + "loss_action": self.loss_lambda_action * float(loss_action.detach().item()), + } + return loss_total, loss_dict + + @torch.no_grad() + def _predict_joint_noise( + self, + latents_video: torch.Tensor, + latents_action: torch.Tensor, + timestep_video: torch.Tensor, + timestep_action: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor, + fuse_vae_embedding_in_latents: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + video_pre = self.video_expert.pre_dit( + x=latents_video, + timestep=timestep_video, + context=context, + context_mask=context_mask, + action=None, + fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, + ) + action_pre = self.action_expert.pre_dit( + action_tokens=latents_action, + timestep=timestep_action, + context=context, + context_mask=context_mask, + ) + + attention_mask = self._build_mot_attention_mask( + video_seq_len=video_pre["tokens"].shape[1], + action_seq_len=action_pre["tokens"].shape[1], + video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), + device=video_pre["tokens"].device, + ) + + tokens_out = self.mot( + embeds_all={ + "video": video_pre["tokens"], + "action": action_pre["tokens"], + }, + attention_mask=attention_mask, + freqs_all={ + "video": video_pre["freqs"], + "action": action_pre["freqs"], + }, + context_all={ + "video": { + "context": video_pre["context"], + "mask": video_pre["context_mask"], + }, + "action": { + "context": action_pre["context"], + "mask": action_pre["context_mask"], + }, + }, + t_mod_all={ + "video": video_pre["t_mod"], + "action": action_pre["t_mod"], + }, + ) + + pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) + pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) + return pred_video, pred_action + + @torch.no_grad() + def infer_joint( + self, + prompt: Optional[str], + input_image: torch.Tensor, + num_video_frames: int, + action_horizon: int, + proprio: Optional[torch.Tensor] = None, + context: Optional[torch.Tensor] = None, + context_mask: Optional[torch.Tensor] = None, + num_inference_steps: int = 20, + sigma_shift: Optional[float] = None, + seed: Optional[int] = None, + rand_device: str = "cpu", + tiled: bool = False, + ) -> dict[str, Any]: + self.eval() + if input_image.ndim == 3: + input_image = input_image.unsqueeze(0) + if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: + raise ValueError( + f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" + ) + _, _, height, width = input_image.shape + checked_h, checked_w, checked_t = self._check_resize_height_width(height, width, num_video_frames) + if (checked_h, checked_w) != (height, width): + raise ValueError( + f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})" + ) + if checked_t != num_video_frames: + raise ValueError( + f"`num_video_frames` must satisfy T % 4 == 1, got {num_video_frames}" + ) + if proprio is not None: + if self.proprio_dim is None: + raise ValueError("`proprio` was provided but `proprio_dim=None` so `proprio_encoder` is disabled.") + if proprio.ndim == 1: + proprio = proprio.unsqueeze(0) + elif proprio.ndim == 2 and proprio.shape[0] == 1: + pass + else: + raise ValueError(f"`proprio` must be [D] or [1,D], got shape {tuple(proprio.shape)}") + if proprio.shape[1] != self.proprio_dim: + raise ValueError(f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}") + proprio = proprio.to(device=self.device, dtype=self.torch_dtype) + + latent_t = (num_video_frames - 1) // self.vae.temporal_downsample_factor + 1 + latent_h = height // self.vae.upsampling_factor + latent_w = width // self.vae.upsampling_factor + + video_generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) + action_generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) + latents_video = torch.randn( + (1, self.vae.model.z_dim, latent_t, latent_h, latent_w), + generator=video_generator, + device=rand_device, + dtype=torch.float32, + ).to(device=self.device, dtype=self.torch_dtype) + latents_action = torch.randn( + (1, action_horizon, self.action_expert.action_dim), + generator=action_generator, + device=rand_device, + dtype=torch.float32, + ).to(device=self.device, dtype=self.torch_dtype) + + input_image = input_image.to(device=self.device, dtype=self.torch_dtype) + first_frame_latents = self._encode_input_image_latents_tensor(input_image=input_image, tiled=tiled) + latents_video[:, :, 0:1] = first_frame_latents.clone() + fuse_flag = bool(getattr(self.video_expert, "fuse_vae_embedding_in_latents", False)) + + use_prompt = prompt is not None + use_context = context is not None or context_mask is not None + if use_prompt and use_context: + raise ValueError("`prompt` and `context/context_mask` are mutually exclusive.") + if not use_prompt and not use_context: + raise ValueError("Either `prompt` or both `context/context_mask` must be provided.") + + if use_prompt: + context, context_mask = self.encode_prompt(prompt) + else: + if context is None or context_mask is None: + raise ValueError("`context` and `context_mask` must be both provided together.") + if context.ndim == 2: + context = context.unsqueeze(0) + if context_mask.ndim == 1: + context_mask = context_mask.unsqueeze(0) + if context.ndim != 3 or context_mask.ndim != 2: + raise ValueError( + f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" + ) + context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) + context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) + if proprio is not None: + context, context_mask = self._append_proprio_to_context( + context=context, + context_mask=context_mask, + proprio=proprio, + ) + + infer_timesteps_video, infer_deltas_video = self.infer_video_scheduler.build_inference_schedule( + num_inference_steps=num_inference_steps, + device=self.device, + dtype=latents_video.dtype, + shift_override=sigma_shift, + ) + infer_timesteps_action, infer_deltas_action = self.infer_action_scheduler.build_inference_schedule( + num_inference_steps=num_inference_steps, + device=self.device, + dtype=latents_action.dtype, + shift_override=sigma_shift, + ) + for step_t_video, step_delta_video, step_t_action, step_delta_action in zip( + infer_timesteps_video, + infer_deltas_video, + infer_timesteps_action, + infer_deltas_action, + ): + timestep_video = step_t_video.unsqueeze(0).to(dtype=latents_video.dtype, device=self.device) + timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device) + + pred_video_posi, pred_action_posi = self._predict_joint_noise( + latents_video=latents_video, + latents_action=latents_action, + timestep_video=timestep_video, + timestep_action=timestep_action, + context=context, + context_mask=context_mask, + fuse_vae_embedding_in_latents=fuse_flag, + ) + pred_video = pred_video_posi + pred_action = pred_action_posi + + latents_video = self.infer_video_scheduler.step(pred_video, step_delta_video, latents_video) + latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action) + latents_video[:, :, 0:1] = first_frame_latents.clone() + + action_out = latents_action[0].detach().to(device="cpu", dtype=torch.float32) + return { + "video": self._decode_latents(latents_video, tiled=tiled), + "action": action_out, + } + + @torch.no_grad() + def infer( + self, + prompt: Optional[str], + input_image: torch.Tensor, + num_frames: int, + action_horizon: int, + proprio: Optional[torch.Tensor] = None, + context: Optional[torch.Tensor] = None, + context_mask: Optional[torch.Tensor] = None, + num_inference_steps: int = 20, + sigma_shift: Optional[float] = None, + seed: Optional[int] = None, + rand_device: str = "cpu", + tiled: bool = False, + ): + return self.infer_joint( + prompt=prompt, + input_image=input_image, + num_video_frames=num_frames, + action_horizon=action_horizon, + proprio=proprio, + context=context, + context_mask=context_mask, + num_inference_steps=num_inference_steps, + sigma_shift=sigma_shift, + seed=seed, + rand_device=rand_device, + tiled=tiled, + ) + + def save_checkpoint(self, path, step=None): + payload = { + "mot": self.mot.state_dict(), + "step": step, + "torch_dtype": str(self.torch_dtype), + } + if self.proprio_encoder is not None: + payload["proprio_encoder"] = self.proprio_encoder.state_dict() + torch.save(payload, path) + + def load_checkpoint(self, path): + payload = torch.load(path, map_location="cpu", weights_only=True) + if "mot" in payload: + self.mot.load_state_dict(payload["mot"], strict=False) + elif "dit" in payload: + logger.warning("Loading legacy `dit` checkpoint into video expert only.") + self.video_expert.load_state_dict(payload["dit"], strict=False) + else: + raise ValueError(f"Checkpoint missing both `mot` and `dit` keys: {path}") + if self.proprio_encoder is not None: + if "proprio_encoder" in payload: + self.proprio_encoder.load_state_dict(payload["proprio_encoder"], strict=True) + else: + logger.warning("Checkpoint has no `proprio_encoder` weights; keeping current `proprio_encoder` params.") + elif "proprio_encoder" in payload: + logger.warning("Checkpoint contains `proprio_encoder` weights but current model has `proprio_dim=None`; ignoring.") + + return payload + + def forward(self, *args, **kwargs): + return self.training_loss(*args, **kwargs) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py new file mode 100644 index 000000000..389cbe0e2 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py @@ -0,0 +1,547 @@ +from __future__ import annotations + +import logging +from typing import Dict, Optional + +import torch +import torch.nn as nn + +from .layers import apply_rope, layer_norm, masked_cross_attention, scaled_attention, split_modulation + +logger = logging.getLogger(__name__) + + +class MoT(nn.Module): + def __init__( + self, + mixtures: Dict[str, nn.Module], + mot_checkpoint_mixed_attn: bool = True, + ): + super().__init__() + if not mixtures: + raise ValueError("`mixtures` cannot be empty.") + if "video" not in mixtures or "action" not in mixtures: + raise ValueError("`mixtures` must include both 'video' and 'action' experts.") + + self.mixtures = nn.ModuleDict(mixtures) + self.expert_order = list(self.mixtures.keys()) + self.mot_checkpoint_mixed_attn = mot_checkpoint_mixed_attn + if mot_checkpoint_mixed_attn: + logger.info("Using gradient checkpointing for mixture attention. This will save memory but use more computation.") + + first_expert = self.mixtures[self.expert_order[0]] + self.num_layers = len(first_expert.blocks) + self.num_heads = first_expert.num_heads + self.attn_head_dim = first_expert.attn_head_dim + + for name in self.expert_order[1:]: + expert = self.mixtures[name] + if len(expert.blocks) != self.num_layers: + raise ValueError( + f"All experts must have same number of layers; got {self.num_layers} and {len(expert.blocks)}" + ) + if expert.num_heads != self.num_heads: + raise ValueError( + f"All experts must have same num_heads; got {self.num_heads} and {expert.num_heads}" + ) + if expert.attn_head_dim != self.attn_head_dim: + raise ValueError( + "All experts must have same attn_head_dim; " + f"got {self.attn_head_dim} and {expert.attn_head_dim}" + ) + + logger.info(f"Initialized MoT with experts: {self.expert_order}, num_layers={self.num_layers}") + for name in self.expert_order: + expert = self.mixtures[name] + logger.info(f" Expert '{name}': num_params={sum(p.numel() for p in expert.parameters()) / 1e9:.2f} B") + + @staticmethod + def _split_modulation(block, t_mod: torch.Tensor): + return split_modulation(block, t_mod) + + def _mixed_attention( + self, + q_cat: torch.Tensor, + k_cat: torch.Tensor, + v_cat: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + def _forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return scaled_attention(q, k, v, attention_mask) + + if self.mot_checkpoint_mixed_attn and self.training: + return torch.utils.checkpoint.checkpoint( + _forward, + q_cat, + k_cat, + v_cat, + use_reentrant=False, + ) + return _forward(q_cat, k_cat, v_cat) + + @staticmethod + def _apply_expert_post_block( + block, + residual_x: torch.Tensor, + mixed_attn_out: torch.Tensor, + gate_msa: torch.Tensor, + shift_mlp: torch.Tensor, + scale_mlp: torch.Tensor, + gate_mlp: torch.Tensor, + context_payload: Optional[dict], + ) -> torch.Tensor: + x = residual_x + gate_msa * block.self_attn.o(mixed_attn_out.flatten(2)) + + if context_payload is not None: + context = context_payload.get("context") + if context is not None: + x = x + masked_cross_attention( + block, + layer_norm(block.norm3, x), + context, + context_payload.get("mask"), + ) + + mlp_input = layer_norm(block.norm2, x) * (1 + scale_mlp) + shift_mlp + return x + gate_mlp * block.ffn(mlp_input) + + def _build_expert_attention_io( + self, + expert, + block, + x: torch.Tensor, + freqs: torch.Tensor, + t_mod: torch.Tensor, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + bool, + ]: + """Build per-expert attention tensors and post-block states. + + Args: + expert: Expert module that owns this `block`; only used to read + `use_gradient_checkpointing`. + block: Transformer block for current layer (`expert.blocks[layer_idx]`). + x: Current expert tokens, shape [B, S, D]. + freqs: RoPE frequencies aligned with token sequence, shape [S, 1, rope_dim]. + t_mod: Time modulation tensor for this expert/layer. + + Returns: + q: Query after q-proj, RMSNorm, and RoPE, shape [B, S, H*Dh]. + k: Key after k-proj, RMSNorm, and RoPE, shape [B, S, H*Dh]. + v: Value after v-proj, shape [B, S, H*Dh]. + residual_x: Original input `x` for residual path in post block. + gate_msa: Gating tensor for self-attention residual branch. + shift_mlp: Shift tensor for MLP modulation. + scale_mlp: Scale tensor for MLP modulation. + gate_mlp: Gating tensor for MLP residual branch. + use_gradient_checkpointing: Whether this expert enables checkpointing. + """ + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self._split_modulation(block, t_mod) + attn_input = layer_norm(block.norm1, x) * (1 + scale_msa) + shift_msa + + batch_size, sequence_length = attn_input.shape[:2] + q = block.self_attn.norm_q(block.self_attn.q(attn_input)).view( + batch_size, sequence_length, self.num_heads, self.attn_head_dim + ) + k = block.self_attn.norm_k(block.self_attn.k(attn_input)).view( + batch_size, sequence_length, self.num_heads, self.attn_head_dim + ) + v = block.self_attn.v(attn_input).view( + batch_size, sequence_length, self.num_heads, self.attn_head_dim + ) + + q = apply_rope(q, freqs) + k = apply_rope(k, freqs) + + use_gradient_checkpointing = bool(getattr(expert, "use_gradient_checkpointing", False)) + return ( + q, + k, + v, + x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) + + def _apply_post_with_optional_checkpoint( + self, + block, + residual_x: torch.Tensor, + gate_msa: torch.Tensor, + shift_mlp: torch.Tensor, + scale_mlp: torch.Tensor, + gate_mlp: torch.Tensor, + use_gradient_checkpointing: bool, + mixed_slice: torch.Tensor, + context_payload: Optional[dict], + ) -> torch.Tensor: + """Apply post-attention computations, with optional checkpointing. + + Args: + block: Transformer block for current layer. + residual_x: Residual input tokens before attention update, shape [B, S, D]. + gate_msa: Gating tensor used after mixed self-attention. + shift_mlp: Shift tensor for MLP input modulation. + scale_mlp: Scale tensor for MLP input modulation. + gate_mlp: Gating tensor used after MLP. + use_gradient_checkpointing: If True and training, checkpoint this post block. + mixed_slice: Mixed-attention output for this expert, shape [B, S, H*Dh]. + context_payload: Optional dict for cross-attention. + - `context`: encoder states [B, L, D] + - `mask`: attention mask [B, S, L] or [B, 1, S, L] + + Returns: + Updated expert tokens after self-attn residual, optional cross-attn, and MLP. + """ + def _post_fn( + _mixed_slice: torch.Tensor, + _x: torch.Tensor, + _gate_msa: torch.Tensor, + _shift_mlp: torch.Tensor, + _scale_mlp: torch.Tensor, + _gate_mlp: torch.Tensor, + _block=block, + _context_payload=context_payload, + ) -> torch.Tensor: + return self._apply_expert_post_block( + block=_block, + residual_x=_x, + mixed_attn_out=_mixed_slice, + gate_msa=_gate_msa, + shift_mlp=_shift_mlp, + scale_mlp=_scale_mlp, + gate_mlp=_gate_mlp, + context_payload=_context_payload, + ) + + if use_gradient_checkpointing and self.training: + return torch.utils.checkpoint.checkpoint( + _post_fn, + mixed_slice, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_reentrant=False, + ) + return _post_fn( + mixed_slice, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + ) + + def prefill_video_cache( + self, + video_tokens: torch.Tensor, + video_freqs: torch.Tensor, + video_t_mod: torch.Tensor, + video_context_payload: Optional[dict], + video_attention_mask: torch.Tensor, + ) -> list[dict[str, torch.Tensor]]: + """Prefill video branch once and cache per-layer K/V for action denoising. + + Args: + video_tokens: Video tokens before layer 0, shape [B, Sv, D]. + video_freqs: Video RoPE frequencies, shape [Sv, 1, rope_dim]. + video_t_mod: Video time modulation tensor. + video_context_payload: Optional dict for video cross-attention. + - `context`: encoder states [B, L, D] + - `mask`: attention mask [B, Sv, L] or [B, 1, Sv, L] + video_attention_mask: Video self-attention mask, shape [Sv, Sv]. + + Returns: + Layer-wise cache list with length `num_layers`. + Each entry contains: + - `k`: video key tensor [B, Sv, H*Dh] + - `v`: video value tensor [B, Sv, H*Dh] + """ + if "video" not in self.mixtures: + raise ValueError("MoT requires `video` expert for `prefill_video_cache`.") + if video_attention_mask.ndim != 2: + raise ValueError( + f"`video_attention_mask` must be 2D [S,S], got shape {tuple(video_attention_mask.shape)}" + ) + if video_attention_mask.shape[0] != video_attention_mask.shape[1]: + raise ValueError( + f"`video_attention_mask` must be square, got shape {tuple(video_attention_mask.shape)}" + ) + if video_attention_mask.shape[0] != video_tokens.shape[1]: + raise ValueError( + "`video_attention_mask` seq length mismatch: " + f"mask={video_attention_mask.shape[0]} vs tokens={video_tokens.shape[1]}" + ) + + expert = self.mixtures["video"] + x = video_tokens + kv_cache: list[dict[str, torch.Tensor]] = [] + for layer_idx in range(self.num_layers): + block = expert.blocks[layer_idx] + # Build video Q/K/V from current layer input tokens. + ( + q, + k, + v, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) = self._build_expert_attention_io( + expert=expert, + block=block, + x=x, + freqs=video_freqs, + t_mod=video_t_mod, + ) + # Video prefill uses only video self-attention mask. + mixed = self._mixed_attention( + q_cat=q, + k_cat=k, + v_cat=v, + attention_mask=video_attention_mask, + ) + # Update video tokens for the next layer and persist current layer K/V. + x = self._apply_post_with_optional_checkpoint( + block=block, + residual_x=residual_x, + gate_msa=gate_msa, + shift_mlp=shift_mlp, + scale_mlp=scale_mlp, + gate_mlp=gate_mlp, + use_gradient_checkpointing=use_gradient_checkpointing, + mixed_slice=mixed, + context_payload=video_context_payload, + ) + kv_cache.append({"k": k, "v": v}) + return kv_cache + + def forward_action_with_video_cache( + self, + action_tokens: torch.Tensor, + action_freqs: torch.Tensor, + action_t_mod: torch.Tensor, + action_context_payload: Optional[dict], + video_kv_cache: list[dict[str, torch.Tensor]], + attention_mask: torch.Tensor, + video_seq_len: int, + ) -> torch.Tensor: + """Run action branch with cached video K/V instead of recomputing video tokens. + + Args: + action_tokens: Action tokens before layer 0, shape [B, Sa, D]. + action_freqs: Action RoPE frequencies, shape [Sa, 1, rope_dim]. + action_t_mod: Action time modulation tensor. + action_context_payload: Optional dict for action cross-attention. + - `context`: encoder states [B, L, D] + - `mask`: attention mask [B, Sa, L] or [B, 1, Sa, L] + video_kv_cache: Layer-wise cached video K/V from `prefill_video_cache`. + attention_mask: Joint [video+action] mask, shape [Sv+Sa, Sv+Sa]. + video_seq_len: Video token count `Sv` in the joint sequence prefix. + + Returns: + Updated action tokens after all layers, shape [B, Sa, D]. + """ + if "action" not in self.mixtures: + raise ValueError("MoT requires `action` expert for `forward_action_with_video_cache`.") + if len(video_kv_cache) != self.num_layers: + raise ValueError( + f"`video_kv_cache` must contain {self.num_layers} layers, got {len(video_kv_cache)}." + ) + if attention_mask.ndim != 2: + raise ValueError(f"`attention_mask` must be 2D [S,S], got shape {tuple(attention_mask.shape)}") + if attention_mask.shape[0] != attention_mask.shape[1]: + raise ValueError(f"`attention_mask` must be square, got shape {tuple(attention_mask.shape)}") + + action_seq_len = int(action_tokens.shape[1]) + total_seq_len = int(video_seq_len) + action_seq_len + if attention_mask.shape[0] != total_seq_len: + raise ValueError( + "`attention_mask` seq length mismatch: " + f"mask={attention_mask.shape[0]} vs expected_total={total_seq_len}" + ) + # Use the action query rows from the joint [video+action] mask. + action_attention_mask = attention_mask[video_seq_len:total_seq_len, :total_seq_len] + + expert = self.mixtures["action"] + x = action_tokens + for layer_idx in range(self.num_layers): + block = expert.blocks[layer_idx] + # Action query/key/value are still step-dependent and must be recomputed each step. + ( + q_action, + k_action, + v_action, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) = self._build_expert_attention_io( + expert=expert, + block=block, + x=x, + freqs=action_freqs, + t_mod=action_t_mod, + ) + layer_cache = video_kv_cache[layer_idx] + if "k" not in layer_cache or "v" not in layer_cache: + raise ValueError( + f"`video_kv_cache[{layer_idx}]` must contain `k` and `v`." + ) + + k_video = layer_cache["k"] + v_video = layer_cache["v"] + if k_video.shape[1] != video_seq_len or v_video.shape[1] != video_seq_len: + raise ValueError( + f"`video_kv_cache[{layer_idx}]` seq len mismatch, expected {video_seq_len}." + ) + + # Mixed attention: action queries attend to cached video K/V plus current action K/V. + k_cat = torch.cat([k_video, k_action], dim=1) + v_cat = torch.cat([v_video, v_action], dim=1) + mixed = self._mixed_attention( + q_cat=q_action, + k_cat=k_cat, + v_cat=v_cat, + attention_mask=action_attention_mask, + ) + x = self._apply_post_with_optional_checkpoint( + block=block, + residual_x=residual_x, + gate_msa=gate_msa, + shift_mlp=shift_mlp, + scale_mlp=scale_mlp, + gate_mlp=gate_mlp, + use_gradient_checkpointing=use_gradient_checkpointing, + mixed_slice=mixed, + context_payload=action_context_payload, + ) + return x + + def forward( + self, + embeds_all: Dict[str, torch.Tensor], + attention_mask: torch.Tensor, + freqs_all: Dict[str, torch.Tensor], + context_all: Dict[str, Optional[dict]], + t_mod_all: Dict[str, torch.Tensor], + ): + missing = [k for k in self.expert_order if k not in embeds_all] + if missing: + raise ValueError(f"Missing expert tokens for {missing}") + missing = [k for k in self.expert_order if k not in freqs_all] + if missing: + raise ValueError(f"Missing expert freqs for {missing}") + missing = [k for k in self.expert_order if k not in t_mod_all] + if missing: + raise ValueError(f"Missing expert t_mod for {missing}") + + if attention_mask.ndim != 2: + raise ValueError(f"`attention_mask` must be 2D [S, S], got shape {tuple(attention_mask.shape)}") + if attention_mask.shape[0] != attention_mask.shape[1]: + raise ValueError(f"`attention_mask` must be square, got shape {tuple(attention_mask.shape)}") + + tokens_all = {k: v for k, v in embeds_all.items()} + + for layer_idx in range(self.num_layers): + q_chunks = [] + k_chunks = [] + v_chunks = [] + cached = {} + seq_lens = [] + + for name in self.expert_order: + expert = self.mixtures[name] + block = expert.blocks[layer_idx] + x = tokens_all[name] + freqs = freqs_all[name] + t_mod = t_mod_all[name] + + ( + q, + k, + v, + residual_x, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + use_gradient_checkpointing, + ) = self._build_expert_attention_io( + expert=expert, + block=block, + x=x, + freqs=freqs, + t_mod=t_mod, + ) + + q_chunks.append(q) + k_chunks.append(k) + v_chunks.append(v) + seq_lens.append(x.shape[1]) + cached[name] = { + "block": block, + "residual_x": residual_x, + "gate_msa": gate_msa, + "shift_mlp": shift_mlp, + "scale_mlp": scale_mlp, + "gate_mlp": gate_mlp, + "use_gradient_checkpointing": use_gradient_checkpointing, + } + + # 3. concat all tokens for mixed attention + q_cat = torch.cat(q_chunks, dim=1) + k_cat = torch.cat(k_chunks, dim=1) + v_cat = torch.cat(v_chunks, dim=1) + + total_seq = q_cat.shape[1] + if attention_mask.shape[0] != total_seq: + raise ValueError( + "Attention mask seq length mismatch: " + f"mask={attention_mask.shape[0]} vs tokens={total_seq}" + ) + + mixed = self._mixed_attention(q_cat=q_cat, k_cat=k_cat, v_cat=v_cat, attention_mask=attention_mask) + + start = 0 + for name, seq_len in zip(self.expert_order, seq_lens): + # 4. split mixed attention output and apply post-attention blocks for each expert + end = start + seq_len + mixed_slice = mixed[:, start:end, :] + cached_expert = cached[name] + block = cached_expert["block"] + context_payload = context_all.get(name) + + updated_tokens = self._apply_post_with_optional_checkpoint( + block=block, + residual_x=cached_expert["residual_x"], + gate_msa=cached_expert["gate_msa"], + shift_mlp=cached_expert["shift_mlp"], + scale_mlp=cached_expert["scale_mlp"], + gate_mlp=cached_expert["gate_mlp"], + use_gradient_checkpointing=cached_expert["use_gradient_checkpointing"], + mixed_slice=mixed_slice, + context_payload=context_payload, + ) + + tokens_all[name] = updated_tokens + start = end + + return tokens_all diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py new file mode 100644 index 000000000..aaf91d024 --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py @@ -0,0 +1,237 @@ +from typing import Any + +import torch +from einops import rearrange + +from ..modules.model2_2 import WanModel +from .layers import expert_block_forward, timestep_embedding + + +class FastWAMVideoDiT(WanModel): + """Wan2.2 video expert with FastWAM's split pre/post interface.""" + + def __init__( + self, + hidden_dim: int, + in_dim: int, + ffn_dim: int, + out_dim: int, + text_dim: int, + freq_dim: int, + eps: float, + patch_size: tuple[int, int, int], + num_heads: int, + attn_head_dim: int, + num_layers: int, + has_image_input: bool, + has_image_pos_emb: bool = False, + has_ref_conv: bool = False, + add_control_adapter: bool = False, + in_dim_control_adapter: int = 24, + seperated_timestep: bool = False, + require_vae_embedding: bool = False, + require_clip_embedding: bool = False, + fuse_vae_embedding_in_latents: bool = True, + action_conditioned: bool = False, + action_dim: int = 7, + action_group_causal_mask_mode: str = "causal", + video_attention_mask_mode: str = "bidirectional", + use_gradient_checkpointing: bool = False, + ): + del has_image_pos_emb, in_dim_control_adapter, action_dim, action_group_causal_mask_mode + if has_image_input or has_ref_conv or add_control_adapter: + raise ValueError("FastWAM training only supports the Wan2.2 latent-input video expert.") + if require_vae_embedding or require_clip_embedding or not fuse_vae_embedding_in_latents: + raise ValueError("FastWAM requires the first-frame VAE latent to be fused into the video input.") + if action_conditioned: + raise ValueError("FastWAM conditions actions through the MoT action expert, not the video expert.") + if hidden_dim != num_heads * attn_head_dim: + raise ValueError( + f"hidden_dim must equal num_heads * attn_head_dim, got {hidden_dim} and " + f"{num_heads} * {attn_head_dim}" + ) + + super().__init__( + model_type="ti2v", + patch_size=tuple(patch_size), + text_len=512, + in_dim=in_dim, + dim=hidden_dim, + ffn_dim=ffn_dim, + freq_dim=freq_dim, + text_dim=text_dim, + out_dim=out_dim, + num_heads=num_heads, + num_layers=num_layers, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=eps, + ) + self.hidden_dim = hidden_dim + self.attn_head_dim = attn_head_dim + self.seperated_timestep = seperated_timestep + self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents + self.video_attention_mask_mode = str(video_attention_mask_mode) + self.use_gradient_checkpointing = bool(use_gradient_checkpointing) + + def _validate_inputs(self, x, timestep, context, context_mask): + if x.ndim != 5: + raise ValueError(f"latents must be [B,C,T,H,W], got {tuple(x.shape)}") + if context.ndim != 3: + raise ValueError(f"context must be [B,L,D], got {tuple(context.shape)}") + if timestep.ndim != 1: + raise ValueError(f"timestep must be [B] or [1], got {tuple(timestep.shape)}") + batch_size = x.shape[0] + if context.shape[0] != batch_size: + if not self.training and batch_size == 1: + x = x.expand(context.shape[0], -1, -1, -1, -1) + batch_size = context.shape[0] + else: + raise ValueError(f"Batch mismatch between latents and context: {batch_size} vs {context.shape[0]}") + if timestep.shape[0] not in (1, batch_size): + raise ValueError(f"timestep length must be 1 or {batch_size}, got {timestep.shape[0]}") + if timestep.shape[0] == 1 and batch_size > 1: + if self.training: + raise ValueError("Training timestep length must match batch size.") + timestep = timestep.expand(batch_size) + if context_mask is None: + context_mask = torch.ones(context.shape[:2], dtype=torch.bool, device=context.device) + elif context_mask.shape != context.shape[:2]: + raise ValueError( + f"context_mask must have shape {tuple(context.shape[:2])}, got {tuple(context_mask.shape)}" + ) + return x, timestep, context_mask + + def build_video_to_video_mask(self, video_seq_len, video_tokens_per_frame, device): + if self.video_attention_mask_mode == "bidirectional": + return torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) + if video_seq_len % video_tokens_per_frame != 0: + raise ValueError("Video token count must be divisible by tokens per frame.") + if self.video_attention_mask_mode == "per_frame_causal": + frame_count = video_seq_len // video_tokens_per_frame + mask = torch.tril(torch.ones((frame_count, frame_count), dtype=torch.bool, device=device)) + return mask.repeat_interleave(video_tokens_per_frame, 0).repeat_interleave( + video_tokens_per_frame, 1 + ) + if self.video_attention_mask_mode == "first_frame_causal": + mask = torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) + mask[:video_tokens_per_frame, video_tokens_per_frame:] = False + return mask + raise ValueError(f"Unsupported video_attention_mask_mode: {self.video_attention_mask_mode}") + + def _video_frequencies(self, grid_size, device): + frames, height, width = grid_size + half_dim = self.attn_head_dim // 2 + split_sizes = [half_dim - 2 * (half_dim // 3), half_dim // 3, half_dim // 3] + time_freqs, height_freqs, width_freqs = self.freqs.split(split_sizes, dim=1) + frequencies = torch.cat( + [ + time_freqs[:frames].view(frames, 1, 1, -1).expand(frames, height, width, -1), + height_freqs[:height].view(1, height, 1, -1).expand(frames, height, width, -1), + width_freqs[:width].view(1, 1, width, -1).expand(frames, height, width, -1), + ], + dim=-1, + ) + return frequencies.reshape(frames * height * width, 1, -1).to(device) + + def pre_dit( + self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + context_mask: torch.Tensor | None = None, + action: torch.Tensor | None = None, + fuse_vae_embedding_in_latents: bool = False, + control_camera_latents_input: torch.Tensor | None = None, + ) -> dict[str, Any]: + del action, control_camera_latents_input + x, timestep, context_mask = self._validate_inputs(x, timestep, context, context_mask) + if not self.seperated_timestep or not fuse_vae_embedding_in_latents: + raise ValueError("FastWAM video expert requires separated timestep with fused first-frame latents.") + + batch_size = x.shape[0] + patch_h, patch_w = self.patch_size[1:] + if x.shape[3] % patch_h or x.shape[4] % patch_w: + raise ValueError("Latent spatial dimensions must be divisible by the Wan patch size.") + tokens_per_frame = (x.shape[3] // patch_h) * (x.shape[4] // patch_w) + token_timesteps = timestep.view(batch_size, 1, 1).expand( + batch_size, x.shape[2], tokens_per_frame + ).clone() + token_timesteps[:, 0] = 0 + token_timesteps = token_timesteps.reshape(batch_size, -1) + time_embedding = self.time_embedding( + timestep_embedding(self.freq_dim, token_timesteps.reshape(-1)) + ).reshape(batch_size, -1, self.hidden_dim) + time_modulation = self.time_projection(time_embedding).unflatten(2, (6, self.hidden_dim)) + + patches = self.patch_embedding(x) + frames, height, width = patches.shape[2:] + tokens = rearrange(patches, "b c f h w -> b (f h w) c").contiguous() + context = self.text_embedding(context) + context_mask = context_mask.unsqueeze(1).expand(-1, tokens.shape[1], -1) + return { + "tokens": tokens, + "freqs": self._video_frequencies((frames, height, width), tokens.device), + "t": time_embedding, + "t_mod": time_modulation, + "context": context, + "context_mask": context_mask, + "meta": { + "grid_size": (frames, height, width), + "tokens_per_frame": tokens_per_frame, + "batch_size": batch_size, + }, + } + + def post_dit(self, tokens: torch.Tensor, pre_state: dict[str, Any]) -> torch.Tensor: + time_embedding = pre_state["t"] + shift, scale = ( + self.head.modulation.unsqueeze(0).to(time_embedding) + time_embedding.unsqueeze(2) + ).chunk(2, dim=2) + tokens = self.head.head( + self.head.norm(tokens) * (1 + scale.squeeze(2)) + shift.squeeze(2) + ) + frames, height, width = pre_state["meta"]["grid_size"] + return rearrange( + tokens, + "b (f h w) (pf ph pw c) -> b c (f pf) (h ph) (w pw)", + f=frames, + h=height, + w=width, + pf=self.patch_size[0], + ph=self.patch_size[1], + pw=self.patch_size[2], + c=self.out_dim, + ) + + def forward(self, x, timestep, context, context_mask=None, action=None, fuse_vae_embedding_in_latents=False): + state = self.pre_dit( + x=x, + timestep=timestep, + context=context, + context_mask=context_mask, + action=action, + fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, + ) + tokens = state["tokens"] + self_mask = self.build_video_to_video_mask( + tokens.shape[1], state["meta"]["tokens_per_frame"], tokens.device + ) + for block in self.blocks: + def run(value, current_block=block): + return expert_block_forward( + current_block, + value, + state["context"], + state["t_mod"], + state["freqs"], + state["context_mask"], + self_mask, + ) + + if self.use_gradient_checkpointing and self.training: + tokens = torch.utils.checkpoint.checkpoint(run, tokens, use_reentrant=False) + else: + tokens = run(tokens) + return self.post_dit(tokens, state) diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py b/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py new file mode 100644 index 000000000..016059a5e --- /dev/null +++ b/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py @@ -0,0 +1,120 @@ +import os +from contextlib import nullcontext +from pathlib import Path + +import torch +from loguru import logger + +from lightx2v_train.model_zoo.native.wan.fastwam import FastWAM +from lightx2v_train.utils.registry import MODEL_REGISTER +from lightx2v_train.utils.utils import get_running_dtype + +DEFAULT_ACTION_DIT_PATH = ( + Path(__file__).resolve().parents[2] + / "assets" + / "libero_fastwam" + / "checkpoints" + / "ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt" +) + + +def _resolve_local_path(path, name, directory=False): + if not path: + raise ValueError(f"model.{name} must be an explicit local path.") + resolved = os.path.abspath(os.path.expanduser(str(path))) + exists = os.path.isdir(resolved) if directory else os.path.isfile(resolved) + if not exists: + expected = "directory" if directory else "file" + raise FileNotFoundError(f"model.{name} {expected} does not exist: {resolved}") + return resolved + + +@MODEL_REGISTER("wan_fastwam") +class WanFastWAMModel: + def __init__(self, config): + self.config = config + self.model_config = config["model"] + self.running_dtype = get_running_dtype(self.model_config.get("running_dtype", "bf16")) + self.device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") + self.module = None + + def load_components(self, transformer_only=False, reference_model=None): + if transformer_only: + raise ValueError("wan_fastwam does not support transformer_only loading.") + + model_path = _resolve_local_path(self.model_config.get("model_path"), "model_path", directory=True) + action_dit_pretrained_path = _resolve_local_path( + self.model_config.get("action_dit_pretrained_path", DEFAULT_ACTION_DIT_PATH), + "action_dit_pretrained_path", + ) + load_text_encoder = bool(self.model_config.get("load_text_encoder", False)) + + video_scheduler = self.model_config.get("video_scheduler", {}) + action_scheduler = self.model_config.get("action_scheduler", {}) + loss_config = self.model_config.get("loss", {}) + self.module = FastWAM.from_wan22_pretrained( + model_path=model_path, + device=str(self.device), + torch_dtype=self.running_dtype, + tokenizer_max_len=int(self.model_config.get("tokenizer_max_len", 128)), + load_text_encoder=load_text_encoder, + proprio_dim=self.model_config.get("proprio_dim", 8), + video_dit_config=dict(self.model_config.get("video_dit_config", {})), + action_dit_config=dict(self.model_config.get("action_dit_config", {})), + action_dit_pretrained_path=action_dit_pretrained_path, + skip_dit_load_from_pretrain=bool(self.model_config.get("skip_dit_load_from_pretrain", False)), + mot_checkpoint_mixed_attn=bool(self.model_config.get("mot_checkpoint_mixed_attn", False)), + video_train_shift=float(video_scheduler.get("train_shift", 5.0)), + video_infer_shift=float(video_scheduler.get("infer_shift", 5.0)), + video_num_train_timesteps=int(video_scheduler.get("num_train_timesteps", 1000)), + action_train_shift=float(action_scheduler.get("train_shift", 5.0)), + action_infer_shift=float(action_scheduler.get("infer_shift", 5.0)), + action_num_train_timesteps=int(action_scheduler.get("num_train_timesteps", 1000)), + loss_lambda_video=float(loss_config.get("lambda_video", 1.0)), + loss_lambda_action=float(loss_config.get("lambda_action", 1.0)), + ) + + checkpoint_path = self.model_config.get("checkpoint_path") + if checkpoint_path: + logger.info("[model] loading FastWAM checkpoint {}", checkpoint_path) + self.unwrap_module().load_checkpoint(checkpoint_path) + + def unwrap_module(self): + if self.module is None: + raise RuntimeError("FastWAM module has not been loaded.") + return self.module + + def autocast_context(self): + if self.device.type == "cuda" and self.running_dtype in {torch.float16, torch.bfloat16}: + return torch.autocast(device_type="cuda", dtype=self.running_dtype) + return nullcontext() + + def set_dit_only_trainable(self): + module = self.unwrap_module() + module.eval() + module.requires_grad_(False) + module.dit.train() + module.dit.requires_grad_(True) + proprio_encoder = getattr(module, "proprio_encoder", None) + if proprio_encoder is not None: + proprio_encoder.train() + proprio_encoder.requires_grad_(True) + + def trainable_parameters(self): + module = self.unwrap_module() + params = list(module.dit.parameters()) + proprio_encoder = getattr(module, "proprio_encoder", None) + if proprio_encoder is not None: + params.extend(list(proprio_encoder.parameters())) + return (param for param in params if param.requires_grad) + + def save_checkpoint(self, path, step=None): + self.unwrap_module().save_checkpoint(path, step=step) + + def load_checkpoint(self, path): + return self.unwrap_module().load_checkpoint(path) + + def log_model_structure(self): + module = self.unwrap_module() + logger.info("[model] class=WanFastWAMModel fastwam_class={}", module.__class__.__name__) + logger.info("[model] trainable root={}", module.dit.__class__.__name__) diff --git a/lightx2v_train/lightx2v_train/schedulers/flow_matching.py b/lightx2v_train/lightx2v_train/schedulers/flow_matching.py index 31d26752c..0d9d4e4c2 100644 --- a/lightx2v_train/lightx2v_train/schedulers/flow_matching.py +++ b/lightx2v_train/lightx2v_train/schedulers/flow_matching.py @@ -177,3 +177,69 @@ def sample_clean_augmentation(self, batch_size, num_frames, num_frame_per_chunk, def add_noise(self, latent, noise, sigmas): sigmas = sigmas.reshape(sigmas.shape[0], 1, sigmas.shape[1], 1, 1) return (1.0 - sigmas) * latent + sigmas * noise + + +class WanContinuousFlowMatchScheduler: + """Continuous flow matching used by FastWAM video and action experts.""" + + def __init__(self, num_train_timesteps=1000, shift=5.0, eps=1e-10): + if num_train_timesteps <= 0: + raise ValueError("num_train_timesteps must be positive.") + if shift <= 0: + raise ValueError("shift must be positive.") + self.num_train_timesteps = int(num_train_timesteps) + self.shift = float(shift) + self.eps = float(eps) + self._y_min, self._weight_norm_const = self._precompute_training_weight_stats() + + @staticmethod + def _phi(value, shift): + return shift * value / (1.0 + (shift - 1.0) * value) + + def _precompute_training_weight_stats(self): + steps = self.num_train_timesteps + grid = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64)[:-1] + timestep = self._phi(grid, self.shift) * float(steps) + weights = torch.exp(-2.0 * ((timestep - steps / 2.0) / steps) ** 2) + minimum = float(weights.min().item()) + return minimum, float((weights - minimum).mean().item()) + + def sample_training_t(self, batch_size, device, dtype): + value = torch.rand((batch_size,), device=device, dtype=torch.float32) + return (self._phi(value, self.shift) * self.num_train_timesteps).to(dtype=dtype) + + def training_weight(self, timestep): + timestep = timestep.to(dtype=torch.float32) + steps = float(self.num_train_timesteps) + weight = torch.exp(-2.0 * ((timestep - steps / 2.0) / steps) ** 2) + weight = (weight - self._y_min) / (self._weight_norm_const + self.eps) + return weight.reshape(()) if weight.numel() == 1 else weight + + def add_noise(self, original_samples, noise, timestep): + sigma = (timestep / self.num_train_timesteps).to( + device=original_samples.device, dtype=original_samples.dtype + ) + if sigma.ndim: + sigma = sigma.view(-1, *([1] * (original_samples.ndim - 1))) + return (1 - sigma) * original_samples + sigma * noise + + @staticmethod + def training_target(sample, noise, timestep): + del timestep + return noise - sample + + def build_inference_schedule(self, num_inference_steps, device, dtype, shift_override=None): + if num_inference_steps <= 0: + raise ValueError("num_inference_steps must be positive.") + shift = self.shift if shift_override is None else float(shift_override) + values = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device) + sigmas = self._phi(values, shift) + timesteps = sigmas[:-1] * self.num_train_timesteps + return timesteps.to(dtype=dtype), (sigmas[1:] - sigmas[:-1]).to(dtype=dtype) + + @staticmethod + def step(model_output, delta, sample): + delta = delta.to(device=sample.device, dtype=sample.dtype) + if delta.ndim: + delta = delta.view(-1, *([1] * (sample.ndim - 1))) + return sample + model_output * delta diff --git a/lightx2v_train/lightx2v_train/trainers/__init__.py b/lightx2v_train/lightx2v_train/trainers/__init__.py index 31678a443..0ff03510d 100644 --- a/lightx2v_train/lightx2v_train/trainers/__init__.py +++ b/lightx2v_train/lightx2v_train/trainers/__init__.py @@ -2,9 +2,10 @@ from .dmd import DmdTrainer, VideoArDmdTrainer, VideoDmdTrainer from .dopsd import DopsdTrainer +from .fastwam import FastWAMTrainer from .flow import FlowMatchingTrainer from .tf import TFTrainer ARDmdTrainer = VideoArDmdTrainer -__all__ = ["build_trainer", "ARDmdTrainer", "DmdTrainer", "FlowMatchingTrainer", "TFTrainer", "VideoArDmdTrainer", "VideoDmdTrainer", "DopsdTrainer"] +__all__ = ["build_trainer", "ARDmdTrainer", "DmdTrainer", "FlowMatchingTrainer", "TFTrainer", "VideoArDmdTrainer", "VideoDmdTrainer", "DopsdTrainer", "FastWAMTrainer"] diff --git a/lightx2v_train/lightx2v_train/trainers/fastwam.py b/lightx2v_train/lightx2v_train/trainers/fastwam.py new file mode 100644 index 000000000..fb6252e67 --- /dev/null +++ b/lightx2v_train/lightx2v_train/trainers/fastwam.py @@ -0,0 +1,517 @@ +import json +import os +import shutil +import time +from contextlib import nullcontext + +import numpy as np +import torch +from PIL import Image, ImageDraw +from diffusers.optimization import get_scheduler +from loguru import logger +from torch.nn.parallel import DistributedDataParallel + +from lightx2v_train.runtime.checkpoint import find_latest_checkpoint, parse_checkpoint_iteration, prune_checkpoints +from lightx2v_train.runtime.distributed import barrier, get_rank, get_world_size, is_distributed, is_main_process, reduce_mean +from lightx2v_train.utils.registry import TRAINER_REGISTER +from lightx2v_train.utils.video import pil_frames_to_video_tensor, save_mp4, video_psnr, video_ssim + + +def _unwrap_dataset(dataset): + while hasattr(dataset, "dataset"): + dataset = dataset.dataset + return dataset + + +@TRAINER_REGISTER("fastwam") +class FastWAMTrainer: + def __init__(self, config): + self.config = config + self.model_config = config["model"] + self.training_config = config["training"] + self.evaluation_config = config.get("evaluation", config.get("inference", {})) + self.logging_config = config.get("logging", {}) + + optimizer_config = self.training_config.get("optimizer", {}) + self.learning_rate = float(optimizer_config.get("learning_rate", self.training_config.get("learning_rate", 1e-4))) + self.weight_decay = float(optimizer_config.get("weight_decay", self.training_config.get("weight_decay", 0.0))) + self.adam_beta1 = float(optimizer_config.get("adam_beta1", 0.9)) + self.adam_beta2 = float(optimizer_config.get("adam_beta2", 0.95)) + self.adam_epsilon = float(optimizer_config.get("adam_epsilon", 1e-8)) + + self.output_train_dir = self.training_config["output_dir"] + self.max_train_iters = self._resolve_max_train_iters() + self.gradient_accumulation_iters = int(self.training_config.get("gradient_accumulation_iters", 1)) + self.max_grad_norm = float(self.training_config.get("max_grad_norm", 1.0)) + self.save_every_iters = int(self.training_config.get("save_every_iters", 0) or 0) + self.save_total_limit = int(self.training_config.get("save_total_limit", 3)) + self.save_final = bool(self.training_config.get("save_final", True)) + self.lr_scheduler_name = self.training_config.get("lr_scheduler", "constant") + self.lr_warmup_iters = int(self.training_config.get("lr_warmup_iters", 0)) + self.train_log_every_iters = max(1, int(self.logging_config.get("train_log_every_iters", 10))) + + self.eval_every_iters = int(self.evaluation_config.get("eval_every_iters", self.evaluation_config.get("infer_every_iters", 0)) or 0) + self.eval_num_inference_steps = int(self.evaluation_config.get("eval_num_inference_steps", 5)) + self.eval_num_samples = max(1, int(self.evaluation_config.get("num_samples", 1))) + self.eval_run_inference = bool(self.evaluation_config.get("run_inference", True)) + self.eval_save_video = bool(self.evaluation_config.get("save_video", True)) + self.eval_save_preview = bool(self.evaluation_config.get("save_preview", True)) + self.eval_fps = max(1, int(self.evaluation_config.get("fps", 8))) + self.eval_tiled = bool(self.evaluation_config.get("tiled", False)) + self.eval_output_dir = self.evaluation_config.get("output_dir", os.path.join(self.output_train_dir, "eval")) + self.eval_seed = int(self.evaluation_config.get("seed", 42)) + + resume_config = self.config.get("resume", {}) + self.auto_resume = bool(resume_config.get("auto_resume", False)) + self.resume_ckpt_path = resume_config.get("resume_ckpt_path") + + def _resolve_max_train_iters(self): + max_train_iters = self.training_config.get("max_train_iters", self.training_config.get("max_steps")) + if max_train_iters is not None: + return int(max_train_iters) + num_epochs = int(self.training_config.get("num_epochs", 1)) + return max(1, num_epochs) + + def set_model(self, model): + self.model = model + + def set_data(self, dataloader_train, dataloader_eval=None): + self.dataloader_train = dataloader_train + self.dataloader_eval = dataloader_eval + + def setup(self, resume_ckpt_path=None): + self.model.set_dit_only_trainable() + self.model.log_model_structure() + self.trainable_params = list(self.model.trainable_parameters()) + if not self.trainable_params: + raise RuntimeError("FastWAM has no trainable parameters.") + + self.optimizer = torch.optim.AdamW( + self.trainable_params, + lr=self.learning_rate, + weight_decay=self.weight_decay, + betas=(self.adam_beta1, self.adam_beta2), + eps=self.adam_epsilon, + ) + self.lr_scheduler = get_scheduler( + self.lr_scheduler_name, + optimizer=self.optimizer, + num_warmup_steps=self.lr_warmup_iters, + num_training_steps=self.max_train_iters, + ) + + self.train_module = self.model.unwrap_module() + self.ddp_enabled = False + if is_distributed(): + self.train_module = DistributedDataParallel( + self.train_module, + device_ids=[torch.cuda.current_device()] if torch.cuda.is_available() else None, + find_unused_parameters=bool(self.training_config.get("ddp_find_unused_parameters", False)), + ) + self.ddp_enabled = True + + if resume_ckpt_path is not None: + self._load_resume_state(resume_ckpt_path) + + def _resolve_resume(self): + if self.resume_ckpt_path: + current_iter = parse_checkpoint_iteration(self.resume_ckpt_path) + return self.resume_ckpt_path, current_iter + if not self.auto_resume: + return None, 0 + ckpt_path, current_iter = find_latest_checkpoint(self.output_train_dir) + if ckpt_path is None: + logger.info("Auto-resume enabled but no checkpoint found in '{}'. Starting from scratch.", self.output_train_dir) + return None, 0 + logger.info("Auto-resuming from checkpoint: {} (iteration {})", ckpt_path, current_iter) + return ckpt_path, current_iter + + def _load_resume_state(self, resume_ckpt_path): + weights_path = os.path.join(resume_ckpt_path, "fastwam.pt") + state_path = os.path.join(resume_ckpt_path, "training_state.pt") + if not os.path.exists(weights_path): + raise RuntimeError(f"fastwam.pt not found in {resume_ckpt_path}") + if not os.path.exists(state_path): + raise RuntimeError(f"training_state.pt not found in {resume_ckpt_path}") + self.model.load_checkpoint(weights_path) + state = torch.load(state_path, map_location="cpu", weights_only=False) + if state.get("world_size") != get_world_size(): + raise RuntimeError(f"Cannot resume world_size={state.get('world_size')} checkpoint with world_size={get_world_size()}.") + self.optimizer.load_state_dict(state["optimizer"]) + self.lr_scheduler.load_state_dict(state["lr_scheduler"]) + logger.info("[resume] restored FastWAM training state from {}", resume_ckpt_path) + + def _save_checkpoint(self, iteration): + if is_main_process(): + prune_checkpoints(self.output_train_dir, self.save_total_limit) + save_dir = os.path.join(self.output_train_dir, f"checkpoint-{iteration:09d}") + logger.info("[checkpoint] saving FastWAM iter={} path={}", iteration, save_dir) + if is_main_process(): + os.makedirs(save_dir, exist_ok=True) + self.model.save_checkpoint(os.path.join(save_dir, "fastwam.pt"), step=iteration) + torch.save( + { + "iteration": iteration, + "world_size": get_world_size(), + "optimizer": self.optimizer.state_dict(), + "lr_scheduler": self.lr_scheduler.state_dict(), + }, + os.path.join(save_dir, "training_state.pt"), + ) + config_path = self.config.get("config_path") + if config_path is not None: + shutil.copy2(config_path, os.path.join(save_dir, "config.yaml")) + barrier() + logger.info("[checkpoint] saved FastWAM iter={} path={}", iteration, save_dir) + + def _forward_loss(self, sample): + with self.model.autocast_context(): + loss, loss_dict = self.train_module(sample) + return loss, loss_dict + + def _ddp_sync_context(self, sync_grad): + if self.ddp_enabled and not sync_grad: + return self.train_module.no_sync() + return nullcontext() + + def train(self): + resume_ckpt_path, current_iter = self._resolve_resume() + self.setup(resume_ckpt_path=resume_ckpt_path) + if is_main_process(): + os.makedirs(self.output_train_dir, exist_ok=True) + barrier() + + grad_accum_counter = 0 + running_loss = 0.0 + running_loss_dict = {} + epoch = 0 + start_time = time.perf_counter() + + logger.info( + "[train] start method=fastwam iter={}/{} world_size={} grad_accum={} log_every={}", + current_iter, + self.max_train_iters, + get_world_size(), + self.gradient_accumulation_iters, + self.train_log_every_iters, + ) + + while current_iter < self.max_train_iters: + sampler = getattr(self.dataloader_train, "sampler", None) + if hasattr(sampler, "set_epoch"): + sampler.set_epoch(epoch) + + for sample in self.dataloader_train: + sync_grad = (grad_accum_counter + 1) % self.gradient_accumulation_iters == 0 + with self._ddp_sync_context(sync_grad): + loss, loss_dict = self._forward_loss(sample) + (loss / self.gradient_accumulation_iters).backward() + + running_loss += float(loss.detach().item()) / self.gradient_accumulation_iters + for key, value in loss_dict.items(): + running_loss_dict[key] = running_loss_dict.get(key, 0.0) + float(value) / self.gradient_accumulation_iters + + grad_accum_counter += 1 + if grad_accum_counter % self.gradient_accumulation_iters != 0: + continue + + grad_norm = torch.nn.utils.clip_grad_norm_(self.trainable_params, self.max_grad_norm) + self.optimizer.step() + self.lr_scheduler.step() + self.optimizer.zero_grad(set_to_none=True) + + current_iter += 1 + display_loss = reduce_mean(running_loss) + current_lr = self.lr_scheduler.get_last_lr()[0] + if current_iter == 1 or current_iter % self.train_log_every_iters == 0 or current_iter >= self.max_train_iters: + elapsed = max(time.perf_counter() - start_time, 1e-6) + detail = " ".join(f"{key}={reduce_mean(value):.6f}" for key, value in sorted(running_loss_dict.items())) + logger.info( + "[train] iter={}/{} loss={:.6f} {} lr={:.8f} grad_norm={:.4f} speed={:.3f} it/s", + current_iter, + self.max_train_iters, + display_loss, + detail, + current_lr, + float(grad_norm), + current_iter / elapsed, + ) + running_loss = 0.0 + running_loss_dict = {} + + if self.eval_every_iters > 0 and self.dataloader_eval is not None and current_iter % self.eval_every_iters == 0: + self.evaluate(current_iter) + + if self.save_every_iters > 0 and current_iter % self.save_every_iters == 0: + self._save_checkpoint(current_iter) + + if current_iter >= self.max_train_iters: + break + + epoch += 1 + + if self.save_final and (self.save_every_iters <= 0 or current_iter % self.save_every_iters != 0): + self._save_checkpoint(current_iter) + logger.info("[train] finished FastWAM iter={}/{}", current_iter, self.max_train_iters) + + @torch.no_grad() + def evaluate(self, current_iter): + if self.dataloader_eval is None: + return + + if torch.cuda.is_available(): + torch.cuda.synchronize() + eval_start_time = time.perf_counter() + eval_dataset = self.dataloader_eval.dataset + model_dataset = _unwrap_dataset(eval_dataset) + if len(eval_dataset) == 0: + return + + model = self.model.unwrap_module() + model.eval() + step_dir = os.path.join(self.eval_output_dir, f"step_{current_iter:09d}") + os.makedirs(step_dir, exist_ok=True) + + eval_indices = self._select_eval_indices(len(eval_dataset)) + local_results = [] + for sample_number in range(get_rank(), len(eval_indices), get_world_size()): + eval_index = eval_indices[sample_number] + sample = self._to_batched_eval_sample(eval_dataset[eval_index]) + local_results.append( + self._evaluate_sample( + model=model, + dataset=model_dataset, + sample=sample, + sample_number=sample_number, + eval_index=eval_index, + step_dir=step_dir, + ) + ) + + results = self._gather_eval_results(local_results) + summary = self._summarize_eval_results(results) + if torch.cuda.is_available(): + torch.cuda.synchronize() + duration_seconds = time.perf_counter() - eval_start_time + if is_main_process(): + metrics_path = os.path.join(step_dir, "metrics.json") + with open(metrics_path, "w", encoding="utf-8") as handle: + json.dump( + { + "iteration": int(current_iter), + "num_samples": len(results), + "duration_seconds": duration_seconds, + "mean": summary, + "samples": results, + }, + handle, + ensure_ascii=True, + indent=2, + ) + details = " ".join(f"{key}={value:.6f}" for key, value in sorted(summary.items())) + logger.info( + "[eval] iter={} samples={} duration={:.3f}s {} output={}", + current_iter, + len(results), + duration_seconds, + details, + step_dir, + ) + + barrier() + self.model.set_dit_only_trainable() + + def _to_batched_eval_sample(self, sample): + result = {} + for key, value in sample.items(): + if isinstance(value, torch.Tensor): + result[key] = value.unsqueeze(0) if value.ndim >= 1 else value.reshape(1) + elif isinstance(value, str): + result[key] = [value] + else: + result[key] = value + return result + + def _select_eval_indices(self, dataset_size): + count = min(self.eval_num_samples, dataset_size) + generator = torch.Generator(device="cpu").manual_seed(self.eval_seed) + return torch.randperm(dataset_size, generator=generator)[:count].tolist() + + def _evaluate_sample(self, model, dataset, sample, sample_number, eval_index, step_dir): + sample_seed = self.eval_seed + sample_number + rng_devices = [torch.cuda.current_device()] if torch.cuda.is_available() else [] + with torch.random.fork_rng(devices=rng_devices): + torch.manual_seed(sample_seed) + with self.model.autocast_context(): + val_loss, val_loss_dict = model.training_loss(sample) + result = { + "sample_number": int(sample_number), + "dataset_index": int(eval_index), + "val_loss": float(val_loss.detach().float().item()), + } + result.update({f"val_{key}": float(value) for key, value in val_loss_dict.items()}) + if not self.eval_run_inference: + return result + + video0 = sample["video"][0] + action = sample.get("action") + proprio = sample.get("proprio") + context = sample.get("context") + context_mask = sample.get("context_mask") + prompt = sample.get("prompt", [None])[0] + input_image = video0[:, 0].unsqueeze(0) + action_horizon = ( + int(action.shape[1]) + if action is not None + else int((video0.shape[1] - 1) * dataset.action_video_freq_ratio) + ) + infer_kwargs = { + "input_image": input_image, + "num_frames": int(video0.shape[1]), + "action_horizon": action_horizon, + "proprio": None if proprio is None else proprio[0, 0], + "num_inference_steps": self.eval_num_inference_steps, + "seed": sample_seed, + "rand_device": "cpu", + "tiled": self.eval_tiled, + } + if context is not None: + infer_kwargs.update(prompt=None, context=context[0], context_mask=context_mask[0]) + else: + infer_kwargs["prompt"] = prompt + prediction = model.infer(**infer_kwargs) + + pred_video_tensor = pil_frames_to_video_tensor(prediction["video"]) + gt_video_tensor = ((video0.detach().float().cpu().clamp(-1.0, 1.0) + 1.0) * 0.5).contiguous() + self._validate_video_shapes(pred_video_tensor, gt_video_tensor, "prediction", "ground truth") + + gt_video_batch = video0.unsqueeze(0).to(device=model.device, dtype=model.torch_dtype) + vae_latents = model._encode_video_latents(gt_video_batch, tiled=self.eval_tiled) + vae_video_tensor = pil_frames_to_video_tensor(model._decode_latents(vae_latents, tiled=self.eval_tiled)) + self._validate_video_shapes(vae_video_tensor, gt_video_tensor, "VAE reconstruction", "ground truth") + + pred_future = pred_video_tensor[:, 1:] + vae_future = vae_video_tensor[:, 1:] + gt_future = gt_video_tensor[:, 1:] + result.update( + { + "future_video_psnr_pred_gt": video_psnr(pred_future, gt_future), + "future_video_ssim_pred_gt": video_ssim(pred_future, gt_future), + "future_video_psnr_pred_vae": video_psnr(pred_future, vae_future), + "future_video_ssim_pred_vae": video_ssim(pred_future, vae_future), + "future_video_psnr_vae_gt": video_psnr(vae_future, gt_future), + "future_video_ssim_vae_gt": video_ssim(vae_future, gt_future), + } + ) + + pred_action = prediction.get("action") + if action is not None and pred_action is not None: + result.update(self._action_metrics(dataset, sample, pred_action, action)) + + if self.eval_save_video or self.eval_save_preview: + comparison_frames = self._comparison_frames(pred_video_tensor, vae_video_tensor, gt_video_tensor) + artifact_stem = f"sample_{sample_number:03d}_index_{eval_index:06d}" + if self.eval_save_video: + video_path = os.path.join(step_dir, f"{artifact_stem}.mp4") + save_mp4(comparison_frames, video_path, fps=self.eval_fps) + result["video_path"] = video_path + if self.eval_save_preview: + preview_path = os.path.join(step_dir, f"{artifact_stem}_preview.png") + comparison_frames[len(comparison_frames) // 2].save(preview_path) + result["preview_path"] = preview_path + return result + + def _action_metrics(self, dataset, sample, pred_action, target_action): + processor = dataset.processor + proprio = sample.get("proprio") + if proprio is None: + raise ValueError("Eval action denormalization requires sample['proprio'].") + + denormalized = {} + action_meta = processor.shape_meta["action"] + state_meta = processor.shape_meta["state"] + for name, raw_action in (("pred", pred_action), ("gt", target_action)): + action_btd = raw_action.unsqueeze(0) if raw_action.ndim == 2 else raw_action + action_btd = action_btd.detach().to(device="cpu", dtype=torch.float32) + batch = { + "action": action_btd, + "state": proprio.detach().to(device="cpu", dtype=torch.float32), + } + batch = processor.action_state_merger.backward(batch) + batch = processor.normalizer.backward(batch) + merged_batch = { + "action": {meta["key"]: batch["action"][meta["key"]].squeeze(0) for meta in action_meta}, + "state": {meta["key"]: batch["state"][meta["key"]].squeeze(0) for meta in state_meta}, + } + denormalized[name] = processor.action_state_merger.forward(merged_batch)["action"].float() + + diff = denormalized["pred"] - denormalized["gt"] + action_is_pad = sample.get("action_is_pad") + if action_is_pad is not None: + valid = ~action_is_pad[0].detach().cpu().bool() + if valid.any(): + diff = diff[valid] + return { + "action_l1": float(diff.abs().mean().item()), + "action_l2": float(diff.pow(2).mean().item()), + } + + @staticmethod + def _validate_video_shapes(left, right, left_name, right_name): + if left.shape != right.shape: + raise ValueError( + f"Eval {left_name}/{right_name} shape mismatch: {tuple(left.shape)} vs {tuple(right.shape)}" + ) + if left.shape[1] <= 1: + raise ValueError(f"Eval video must contain at least one future frame, got {tuple(left.shape)}") + + @staticmethod + def _comparison_frames(pred, vae, gt): + labels = ("Prediction", "VAE reconstruction", "Ground truth") + videos = (pred, vae, gt) + frame_height = int(gt.shape[2]) + frame_width = int(gt.shape[3]) + label_height = 32 + frames = [] + for frame_index in range(gt.shape[1]): + canvas = Image.new("RGB", (frame_width * 3, frame_height + label_height), color=(18, 18, 18)) + draw = ImageDraw.Draw(canvas) + for column, (label, video) in enumerate(zip(labels, videos)): + array = ( + video[:, frame_index] + .permute(1, 2, 0) + .clamp(0.0, 1.0) + .mul(255.0) + .round() + .to(torch.uint8) + .numpy() + ) + canvas.paste(Image.fromarray(np.ascontiguousarray(array)), (column * frame_width, label_height)) + draw.text((column * frame_width + 8, 9), label, fill=(245, 245, 245)) + frames.append(canvas) + return frames + + @staticmethod + def _gather_eval_results(local_results): + if not is_distributed(): + return local_results + gathered = [None for _ in range(get_world_size())] + torch.distributed.all_gather_object(gathered, local_results) + return sorted( + [result for rank_results in gathered for result in rank_results], + key=lambda result: result["sample_number"], + ) + + @staticmethod + def _summarize_eval_results(results): + if not results: + return {} + metric_keys = sorted( + key + for key, value in results[0].items() + if isinstance(value, float) + ) + return { + key: float(sum(result[key] for result in results if key in result) / sum(key in result for result in results)) + for key in metric_keys + } diff --git a/lightx2v_train/lightx2v_train/utils/video.py b/lightx2v_train/lightx2v_train/utils/video.py new file mode 100644 index 000000000..284665f9f --- /dev/null +++ b/lightx2v_train/lightx2v_train/utils/video.py @@ -0,0 +1,120 @@ +import os +from typing import Iterable, Sequence + +import imageio +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + + +def _to_even_frame(frame: np.ndarray) -> np.ndarray: + height, width = frame.shape[:2] + if height % 2 == 0 and width % 2 == 0: + return frame + return np.pad(frame, ((0, height % 2), (0, width % 2), (0, 0)), mode="edge") + + +def save_mp4(frames: Iterable[Image.Image], path: str, fps: int = 8): + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + writer = imageio.get_writer( + path, + fps=max(fps, 1), + codec="libx264", + format="FFMPEG", + pixelformat="yuv420p", + ) + try: + for frame in frames: + writer.append_data(_to_even_frame(np.array(frame.convert("RGB")))) + finally: + writer.close() + + +def pil_frames_to_video_tensor(frames: Sequence[Image.Image]) -> torch.Tensor: + if len(frames) == 0: + raise ValueError("`frames` must be non-empty.") + + frame_tensors = [] + for frame in frames: + arr = np.array(frame.convert("RGB"), dtype=np.float32) / 255.0 + x = torch.from_numpy(arr).permute(2, 0, 1).contiguous() # [3, H, W] + frame_tensors.append(x) + return torch.stack(frame_tensors, dim=1) # [3, T, H, W] + + +def _gaussian_kernel_2d(kernel_size: int, sigma: float, channels: int, device: torch.device, dtype: torch.dtype): + coords = torch.arange(kernel_size, device=device, dtype=dtype) - (kernel_size - 1) / 2.0 + g = torch.exp(-(coords**2) / (2.0 * sigma * sigma)) + g = g / g.sum() + kernel_2d = torch.outer(g, g) + kernel_2d = kernel_2d / kernel_2d.sum() + kernel_2d = kernel_2d.view(1, 1, kernel_size, kernel_size) + return kernel_2d.repeat(channels, 1, 1, 1) + + +def video_psnr(pred: torch.Tensor, target: torch.Tensor, data_range: float = 1.0, eps: float = 1e-8) -> float: + """ + Compute average PSNR over all frames. + Expects `pred` and `target` in shape [3, T, H, W] with values in [0, 1]. + """ + if pred.shape != target.shape: + raise ValueError(f"Shape mismatch: pred={tuple(pred.shape)} target={tuple(target.shape)}") + + pred = pred.float() + target = target.float() + mse = (pred - target).pow(2).mean(dim=(0, 2, 3)) # [T] + psnr = 10.0 * torch.log10((data_range * data_range) / (mse + eps)) + return float(psnr.mean().item()) + + +def video_ssim( + pred: torch.Tensor, + target: torch.Tensor, + data_range: float = 1.0, + kernel_size: int = 11, + sigma: float = 1.5, + k1: float = 0.01, + k2: float = 0.03, +) -> float: + """ + Compute average SSIM over all frames. + Expects `pred` and `target` in shape [3, T, H, W] with values in [0, 1]. + """ + if pred.shape != target.shape: + raise ValueError(f"Shape mismatch: pred={tuple(pred.shape)} target={tuple(target.shape)}") + if pred.ndim != 4 or pred.shape[0] != 3: + raise ValueError(f"Expected [3, T, H, W], got {tuple(pred.shape)}") + if kernel_size % 2 == 0: + raise ValueError("`kernel_size` must be odd.") + + pred = pred.float().permute(1, 0, 2, 3).contiguous() # [T, 3, H, W] + target = target.float().permute(1, 0, 2, 3).contiguous() # [T, 3, H, W] + channels = pred.shape[1] + kernel = _gaussian_kernel_2d( + kernel_size=kernel_size, + sigma=sigma, + channels=channels, + device=pred.device, + dtype=pred.dtype, + ) + + pad = kernel_size // 2 + mu_x = F.conv2d(pred, kernel, padding=pad, groups=channels) + mu_y = F.conv2d(target, kernel, padding=pad, groups=channels) + + mu_x2 = mu_x * mu_x + mu_y2 = mu_y * mu_y + mu_xy = mu_x * mu_y + + sigma_x2 = F.conv2d(pred * pred, kernel, padding=pad, groups=channels) - mu_x2 + sigma_y2 = F.conv2d(target * target, kernel, padding=pad, groups=channels) - mu_y2 + sigma_xy = F.conv2d(pred * target, kernel, padding=pad, groups=channels) - mu_xy + + c1 = (k1 * data_range) ** 2 + c2 = (k2 * data_range) ** 2 + + numerator = (2.0 * mu_xy + c1) * (2.0 * sigma_xy + c2) + denominator = (mu_x2 + mu_y2 + c1) * (sigma_x2 + sigma_y2 + c2) + ssim_map = numerator / (denominator + 1e-12) + return float(ssim_map.mean().item()) diff --git a/lightx2v_train/scripts/run_fastwam_libero_train.sh b/lightx2v_train/scripts/run_fastwam_libero_train.sh new file mode 100755 index 000000000..a823d629f --- /dev/null +++ b/lightx2v_train/scripts/run_fastwam_libero_train.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +set -euo pipefail + +LIGHTX2V_ROOT="${LIGHTX2V_ROOT:-/data/nvme0/chendingyu/LightX2V}" +CONFIG_PATH="${CONFIG_PATH:-${LIGHTX2V_ROOT}/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml}" + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-1}" +export PYTHONPATH="${LIGHTX2V_ROOT}/lightx2v_train:${LIGHTX2V_ROOT}:${PYTHONPATH:-}" + +cd "${LIGHTX2V_ROOT}" +python lightx2v_train/train.py --config "${CONFIG_PATH}" "$@" From ca2d5840b6c75198b7e7e07e0d0cd3efda2ce545 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Mon, 13 Jul 2026 08:41:00 +0000 Subject: [PATCH 02/10] u --- lightx2v_train/scripts/run_fastwam_libero_train.sh | 1 - 1 file changed, 1 deletion(-) diff --git a/lightx2v_train/scripts/run_fastwam_libero_train.sh b/lightx2v_train/scripts/run_fastwam_libero_train.sh index a823d629f..233ddbe8e 100755 --- a/lightx2v_train/scripts/run_fastwam_libero_train.sh +++ b/lightx2v_train/scripts/run_fastwam_libero_train.sh @@ -5,7 +5,6 @@ LIGHTX2V_ROOT="${LIGHTX2V_ROOT:-/data/nvme0/chendingyu/LightX2V}" CONFIG_PATH="${CONFIG_PATH:-${LIGHTX2V_ROOT}/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml}" export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-1}" -export PYTHONPATH="${LIGHTX2V_ROOT}/lightx2v_train:${LIGHTX2V_ROOT}:${PYTHONPATH:-}" cd "${LIGHTX2V_ROOT}" python lightx2v_train/train.py --config "${CONFIG_PATH}" "$@" From a43e4100962f24edff890d2f8078e0b2429c60c1 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Mon, 13 Jul 2026 10:10:09 +0000 Subject: [PATCH 03/10] update yaml --- .../train/fastwam/libero_uncond_2cam224.yaml | 23 +++++++++---------- 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml index 51d35982c..5950b6be4 100644 --- a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml +++ b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml @@ -15,23 +15,22 @@ data: text_embedding_cache_dir: /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/text_embeds_cache/libero pretrained_norm_stats: /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/dataset_stats.json batch_size: 1 - num_workers: 0 + num_workers: 4 val: <<: *libero_data training: method: fastwam - output_dir: /data/nvme0/chendingyu/LightX2V/lightx2v_train/runs/fastwam_libero - # Smoke default: one optimizer step proves the LightX2V integration path. - # Increase this for a real training run. - max_train_iters: 1 - gradient_accumulation_iters: 1 + output_dir: /data/nvme0/chendingyu/LightX2V/lightx2v_train/runs/fastwam_libero_full + # Equivalent to the official 10-epoch, global-batch-128 LIBERO schedule. + max_train_iters: 21700 + gradient_accumulation_iters: 32 max_grad_norm: 1.0 - save_every_iters: 0 + save_every_iters: 2000 save_total_limit: 2 save_final: true - lr_scheduler: constant - lr_warmup_iters: 0 + lr_scheduler: cosine + lr_warmup_iters: 1085 optimizer: learning_rate: 1.0e-4 weight_decay: 1.0e-2 @@ -40,7 +39,7 @@ training: adam_epsilon: 1.0e-8 evaluation: - eval_every_iters: 1 + eval_every_iters: 200 eval_num_inference_steps: 5 num_samples: 1 run_inference: true @@ -51,10 +50,10 @@ evaluation: seed: 42 logging: - train_log_every_iters: 1 + train_log_every_iters: 10 resume: - auto_resume: false + auto_resume: true resume_ckpt_path: null distributed: From a0541ecba54c7ce90c7a8e2900fd6980b0096e13 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Mon, 13 Jul 2026 10:48:24 +0000 Subject: [PATCH 04/10] zero1 --- .../train/fastwam/libero_uncond_2cam224.yaml | 8 ++- .../lightx2v_train/trainers/fastwam.py | 64 ++++++++++++++++--- 2 files changed, 60 insertions(+), 12 deletions(-) diff --git a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml index 5950b6be4..12c682061 100644 --- a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml +++ b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml @@ -14,8 +14,8 @@ data: - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_10_no_noops_lerobot text_embedding_cache_dir: /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/text_embeds_cache/libero pretrained_norm_stats: /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/dataset_stats.json - batch_size: 1 - num_workers: 4 + batch_size: 8 + num_workers: 8 val: <<: *libero_data @@ -24,7 +24,7 @@ training: output_dir: /data/nvme0/chendingyu/LightX2V/lightx2v_train/runs/fastwam_libero_full # Equivalent to the official 10-epoch, global-batch-128 LIBERO schedule. max_train_iters: 21700 - gradient_accumulation_iters: 32 + gradient_accumulation_iters: 4 max_grad_norm: 1.0 save_every_iters: 2000 save_total_limit: 2 @@ -58,5 +58,7 @@ resume: distributed: backend: nccl + zero1: + enabled: true sequence_parallel: enabled: false diff --git a/lightx2v_train/lightx2v_train/trainers/fastwam.py b/lightx2v_train/lightx2v_train/trainers/fastwam.py index fb6252e67..ea1508ec3 100644 --- a/lightx2v_train/lightx2v_train/trainers/fastwam.py +++ b/lightx2v_train/lightx2v_train/trainers/fastwam.py @@ -12,7 +12,16 @@ from torch.nn.parallel import DistributedDataParallel from lightx2v_train.runtime.checkpoint import find_latest_checkpoint, parse_checkpoint_iteration, prune_checkpoints -from lightx2v_train.runtime.distributed import barrier, get_rank, get_world_size, is_distributed, is_main_process, reduce_mean +from lightx2v_train.runtime.distributed import ( + barrier, + get_data_parallel_group, + get_data_parallel_world_size, + get_rank, + get_world_size, + is_distributed, + is_main_process, + reduce_mean, +) from lightx2v_train.utils.registry import TRAINER_REGISTER from lightx2v_train.utils.video import pil_frames_to_video_tensor, save_mp4, video_psnr, video_ssim @@ -50,6 +59,10 @@ def __init__(self, config): self.lr_warmup_iters = int(self.training_config.get("lr_warmup_iters", 0)) self.train_log_every_iters = max(1, int(self.logging_config.get("train_log_every_iters", 10))) + zero1_config = self.config.get("distributed", {}).get("zero1", {}) + self.zero1_requested = bool(zero1_config.get("enabled", False)) if isinstance(zero1_config, dict) else bool(zero1_config) + self.zero1_enabled = False + self.eval_every_iters = int(self.evaluation_config.get("eval_every_iters", self.evaluation_config.get("infer_every_iters", 0)) or 0) self.eval_num_inference_steps = int(self.evaluation_config.get("eval_num_inference_steps", 5)) self.eval_num_samples = max(1, int(self.evaluation_config.get("num_samples", 1))) @@ -79,6 +92,35 @@ def set_data(self, dataloader_train, dataloader_eval=None): self.dataloader_train = dataloader_train self.dataloader_eval = dataloader_eval + def _build_optimizer(self): + optimizer_kwargs = { + "lr": self.learning_rate, + "weight_decay": self.weight_decay, + "betas": (self.adam_beta1, self.adam_beta2), + "eps": self.adam_epsilon, + } + data_parallel_size = get_data_parallel_world_size() + if self.zero1_requested and data_parallel_size > 1: + from torch.distributed.optim import ZeroRedundancyOptimizer + + self.zero1_enabled = True + logger.info("[optimizer] using ZeroRedundancyOptimizer with AdamW across {} data-parallel ranks", data_parallel_size) + return ZeroRedundancyOptimizer( + self.trainable_params, + optimizer_class=torch.optim.AdamW, + process_group=get_data_parallel_group(), + parameters_as_bucket_view=False, + overlap_with_ddp=False, + **optimizer_kwargs, + ) + + self.zero1_enabled = False + if self.zero1_requested: + logger.info("[optimizer] ZeroRedundancyOptimizer requested but data-parallel size is 1; using AdamW") + else: + logger.info("[optimizer] using AdamW without optimizer-state sharding") + return torch.optim.AdamW(self.trainable_params, **optimizer_kwargs) + def setup(self, resume_ckpt_path=None): self.model.set_dit_only_trainable() self.model.log_model_structure() @@ -86,13 +128,7 @@ def setup(self, resume_ckpt_path=None): if not self.trainable_params: raise RuntimeError("FastWAM has no trainable parameters.") - self.optimizer = torch.optim.AdamW( - self.trainable_params, - lr=self.learning_rate, - weight_decay=self.weight_decay, - betas=(self.adam_beta1, self.adam_beta2), - eps=self.adam_epsilon, - ) + self.optimizer = self._build_optimizer() self.lr_scheduler = get_scheduler( self.lr_scheduler_name, optimizer=self.optimizer, @@ -139,13 +175,22 @@ def _load_resume_state(self, resume_ckpt_path): raise RuntimeError(f"Cannot resume world_size={state.get('world_size')} checkpoint with world_size={get_world_size()}.") self.optimizer.load_state_dict(state["optimizer"]) self.lr_scheduler.load_state_dict(state["lr_scheduler"]) - logger.info("[resume] restored FastWAM training state from {}", resume_ckpt_path) + saved_optimizer_sharding = state.get("optimizer_sharding", "none") + active_optimizer_sharding = "zero1" if self.zero1_enabled else "none" + logger.info( + "[resume] restored FastWAM training state from {} (optimizer sharding: {} -> {})", + resume_ckpt_path, + saved_optimizer_sharding, + active_optimizer_sharding, + ) def _save_checkpoint(self, iteration): if is_main_process(): prune_checkpoints(self.output_train_dir, self.save_total_limit) save_dir = os.path.join(self.output_train_dir, f"checkpoint-{iteration:09d}") logger.info("[checkpoint] saving FastWAM iter={} path={}", iteration, save_dir) + if self.zero1_enabled: + self.optimizer.consolidate_state_dict(to=0) if is_main_process(): os.makedirs(save_dir, exist_ok=True) self.model.save_checkpoint(os.path.join(save_dir, "fastwam.pt"), step=iteration) @@ -153,6 +198,7 @@ def _save_checkpoint(self, iteration): { "iteration": iteration, "world_size": get_world_size(), + "optimizer_sharding": "zero1" if self.zero1_enabled else "none", "optimizer": self.optimizer.state_dict(), "lr_scheduler": self.lr_scheduler.state_dict(), }, From b73d58e63c7b3a0e8b73fc256dcece3da3daad25 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Mon, 13 Jul 2026 13:00:19 +0000 Subject: [PATCH 05/10] update save-resume --- .../train/fastwam/libero_uncond_2cam224.yaml | 3 - .../lightx2v_train/data/__init__.py | 2 + .../lightx2v_train/data/libero/dataset.py | 9 +- .../lightx2v_train/data/libero/preparation.py | 321 ++++++++++++++++++ .../lightx2v_train/data/preparation.py | 11 + .../native/wan/fastwam/action_dit.py | 78 +++++ .../model_zoo/native/wan/fastwam/model.py | 22 +- .../lightx2v_train/model_zoo/wan_fastwam.py | 17 +- .../lightx2v_train/trainers/fastwam.py | 140 +++++++- lightx2v_train/train.py | 4 +- 10 files changed, 560 insertions(+), 47 deletions(-) create mode 100644 lightx2v_train/lightx2v_train/data/libero/preparation.py create mode 100644 lightx2v_train/lightx2v_train/data/preparation.py diff --git a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml index 12c682061..79ce41ab1 100644 --- a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml +++ b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml @@ -2,7 +2,6 @@ model: name: wan_fastwam running_dtype: bf16 model_path: /data/nvme0/models/Wan-AI/Wan2.2-TI2V-5B - action_dit_pretrained_path: /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt data: train: &libero_data @@ -12,8 +11,6 @@ data: - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_object_no_noops_lerobot - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_goal_no_noops_lerobot - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_10_no_noops_lerobot - text_embedding_cache_dir: /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/text_embeds_cache/libero - pretrained_norm_stats: /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/dataset_stats.json batch_size: 8 num_workers: 8 val: diff --git a/lightx2v_train/lightx2v_train/data/__init__.py b/lightx2v_train/lightx2v_train/data/__init__.py index 460d520ff..e8a8ab3c9 100644 --- a/lightx2v_train/lightx2v_train/data/__init__.py +++ b/lightx2v_train/lightx2v_train/data/__init__.py @@ -2,6 +2,7 @@ from .image_dataset import build_image_dataset from .libero.dataset import build_libero_fastwam_dataset +from .preparation import prepare_data from .video_dataset import build_causal_forcing_lmdb_dataset, build_prompt_dataset, build_wan_t2v_cached_dataset, build_wan_t2v_video_dataset __all__ = [ @@ -12,4 +13,5 @@ "build_wan_t2v_video_dataset", "build_wan_t2v_cached_dataset", "build_causal_forcing_lmdb_dataset", + "prepare_data", ] diff --git a/lightx2v_train/lightx2v_train/data/libero/dataset.py b/lightx2v_train/lightx2v_train/data/libero/dataset.py index 7c5f498fb..6abc84afa 100644 --- a/lightx2v_train/lightx2v_train/data/libero/dataset.py +++ b/lightx2v_train/lightx2v_train/data/libero/dataset.py @@ -65,13 +65,8 @@ def _build_dataset(config, split): dataset_dirs=[_path(item) for item in dataset_dirs], shape_meta=shape_meta, processor=processor, - text_embedding_cache_dir=_path( - config.get("text_embedding_cache_dir") - or ASSET_ROOT / "text_embeds_cache" / "libero" - ), - pretrained_norm_stats=_path( - config.get("pretrained_norm_stats") or ASSET_ROOT / "dataset_stats.json" - ), + text_embedding_cache_dir=_path(config["text_embedding_cache_dir"]), + pretrained_norm_stats=_path(config["pretrained_norm_stats"]), num_frames=num_frames, context_len=int(config.get("context_len", 128)), val_set_proportion=float(config.get("val_set_proportion", 0.0)), diff --git a/lightx2v_train/lightx2v_train/data/libero/preparation.py b/lightx2v_train/lightx2v_train/data/libero/preparation.py new file mode 100644 index 000000000..b4741ae7f --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/libero/preparation.py @@ -0,0 +1,321 @@ +import gc +import hashlib +import json +import os +import shutil +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from pathlib import Path + +import torch +from loguru import logger + +from lightx2v_train.runtime.distributed import barrier, is_main_process + +from .dataset import DEFAULT_DATASET_DIRS, _default_shape_meta +from .lerobot_dataset import LiberoLeRobotDataset +from .robot_video_dataset import PROMPT_TEMPLATE + +TEXT_ENCODER_ID = "wan22ti2v5b" +TEXT_EMBEDDING_DIM = 4096 + + +def _resolve_path(value): + return Path(value).expanduser().resolve() + + +def _dataset_dirs(split_config): + values = split_config.get("dataset_dirs") or DEFAULT_DATASET_DIRS + if isinstance(values, (str, Path)): + values = [values] + return [_resolve_path(value) for value in values] + + +def _atomic_json_dump(payload, path): + path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") + with temporary_path.open("w", encoding="utf-8") as handle: + json.dump(_to_serializable(payload), handle, ensure_ascii=True, indent=2) + os.replace(temporary_path, path) + + +def _atomic_torch_save(payload, path): + path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") + torch.save(payload, temporary_path) + os.replace(temporary_path, path) + + +def _to_serializable(value): + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + if isinstance(value, dict): + return {key: _to_serializable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_to_serializable(item) for item in value] + return value + + +def _sliding_window_with_replication(value, window_size): + offsets = torch.arange(window_size, dtype=torch.long) + indices = torch.arange(value.shape[0], dtype=torch.long)[:, None] + offsets[None, :] + indices.clamp_(max=value.shape[0] - 1) + return value[indices] + + +def _episode_field_stats(value): + return { + "min": value.amin(0), + "max": value.amax(0), + "q01": torch.quantile(value, 0.01, dim=0, keepdim=False), + "q99": torch.quantile(value, 0.99, dim=0, keepdim=False), + "mean": value.mean(0), + "var": value.var(0), + } + + +def _aggregate_field_stats(episode_stats): + values = { + name: torch.stack([stats[name] for stats in episode_stats]) + for name in ("min", "max", "q01", "q99", "mean", "var") + } + stepwise_mean = values["mean"].mean(0) + stepwise_std = (values["var"] + (values["mean"] - stepwise_mean) ** 2).mean(0).sqrt() + global_mean = values["mean"].mean((0, 1)) + global_std = (values["var"] + (values["mean"] - global_mean) ** 2).mean((0, 1)).sqrt() + stepwise_min = values["min"].amin(0) + stepwise_max = values["max"].amax(0) + stepwise_q01 = values["q01"].amin(0) + stepwise_q99 = values["q99"].amax(0) + return { + "stepwise_min": stepwise_min, + "stepwise_max": stepwise_max, + "global_min": stepwise_min.amin(0), + "global_max": stepwise_max.amax(0), + "stepwise_q01": stepwise_q01, + "stepwise_q99": stepwise_q99, + "global_q01": stepwise_q01.amin(0), + "global_q99": stepwise_q99.amax(0), + "stepwise_mean": stepwise_mean, + "stepwise_std": stepwise_std, + "global_mean": global_mean, + "global_std": global_std, + } + + +def calculate_dataset_stats(split_config): + shape_meta = deepcopy(split_config.get("shape_meta") or _default_shape_meta()) + if len(shape_meta["action"]) != 1 or len(shape_meta["state"]) != 1: + raise ValueError("Automatic LIBERO stats currently require one action field and one state field.") + + image_keys = [f"observation.images.{item['key']}" for item in shape_meta["images"]] + dataset = LiberoLeRobotDataset( + dataset_dirs=[str(path) for path in _dataset_dirs(split_config)], + image_keys=image_keys, + state_key="observation.state", + action_key="action", + num_frames=int(split_config.get("num_frames", 33)), + global_sample_stride=int(split_config.get("global_sample_stride", 1)), + val_set_proportion=float(split_config.get("val_set_proportion", 0.0)), + is_training_set=True, + video_backend=split_config.get("video_backend"), + ) + action_horizon = int(split_config.get("num_frames", 33)) - 1 + + def process_episode(episode_position): + raw = dataset._load_episode(episode_position) + state = raw[dataset.state_key].float().unsqueeze(1) + action = _sliding_window_with_replication(raw[dataset.action_key].float(), action_horizon) + return _episode_field_stats(state), _episode_field_stats(action) + + worker_count = min(16, len(dataset.episodes), os.cpu_count() or 1) + with ThreadPoolExecutor(max_workers=worker_count) as executor: + episode_results = list(executor.map(process_episode, range(len(dataset.episodes)))) + + state_key = shape_meta["state"][0]["key"] + action_key = shape_meta["action"][0]["key"] + return { + "state": {state_key: _aggregate_field_stats([result[0] for result in episode_results])}, + "action": {action_key: _aggregate_field_stats([result[1] for result in episode_results])}, + "num_episodes": len(dataset.episodes), + "num_transition": len(dataset), + } + + +def _collect_prompts(split_configs): + prompts = [] + seen = set() + for split_config in split_configs: + for dataset_dir in _dataset_dirs(split_config): + tasks_path = dataset_dir / "meta" / "tasks.jsonl" + if not tasks_path.is_file(): + raise FileNotFoundError(f"LIBERO tasks file does not exist: {tasks_path}") + with tasks_path.open(encoding="utf-8") as handle: + for line in handle: + if not line.strip(): + continue + task = str(json.loads(line)["task"]) + prompt = PROMPT_TEMPLATE.format(task=task) + if prompt not in seen: + seen.add(prompt) + prompts.append(prompt) + return prompts + + +def _text_cache_path(cache_dir, prompt, context_len): + digest = hashlib.sha256(prompt.encode("utf-8")).hexdigest() + return cache_dir / f"{digest}.t5_len{context_len}.{TEXT_ENCODER_ID}.pt" + + +def _valid_text_cache(path, context_len): + if not path.is_file(): + return False + try: + payload = torch.load(path, map_location="cpu", weights_only=True) + return tuple(payload["context"].shape) == (context_len, TEXT_EMBEDDING_DIM) and tuple(payload["mask"].shape) == (context_len,) + except Exception: + return False + + +def _validate_text_caches(split_configs): + for split_config in split_configs: + cache_dir = _resolve_path(split_config["text_embedding_cache_dir"]) + context_len = int(split_config.get("context_len", 128)) + prompts = _collect_prompts([split_config]) + missing = [ + _text_cache_path(cache_dir, prompt, context_len) + for prompt in prompts + if not _valid_text_cache(_text_cache_path(cache_dir, prompt, context_len), context_len) + ] + if missing: + raise FileNotFoundError( + f"LIBERO text embedding cache is incomplete: missing={len(missing)}/{len(prompts)} first={missing[0]}" + ) + + +def precompute_text_embeddings(model_path, cache_dir, context_len, prompts): + missing_prompts = [ + prompt + for prompt in prompts + if not _valid_text_cache(_text_cache_path(cache_dir, prompt, context_len), context_len) + ] + if not missing_prompts: + logger.info("[data-preflight] text embeddings already cached: prompts={} path={}", len(prompts), cache_dir) + return + if not torch.cuda.is_available(): + raise RuntimeError("Automatic FastWAM text embedding generation requires CUDA.") + + from lightx2v_train.model_zoo.native.wan.modules.t5 import T5EncoderModel + + model_root = _resolve_path(model_path) + checkpoint_path = model_root / "models_t5_umt5-xxl-enc-bf16.pth" + tokenizer_path = model_root / "google" / "umt5-xxl" + if not checkpoint_path.is_file(): + raise FileNotFoundError(f"Wan text encoder checkpoint does not exist: {checkpoint_path}") + if not tokenizer_path.is_dir(): + raise FileNotFoundError(f"Wan tokenizer directory does not exist: {tokenizer_path}") + + device = torch.device("cuda", torch.cuda.current_device()) + logger.info( + "[data-preflight] generating text embeddings: missing={}/{} context_len={} path={}", + len(missing_prompts), + len(prompts), + context_len, + cache_dir, + ) + with torch.random.fork_rng(devices=[torch.cuda.current_device()]): + encoder = T5EncoderModel( + text_len=context_len, + dtype=torch.bfloat16, + device=device, + checkpoint_path=str(checkpoint_path), + tokenizer_path=str(tokenizer_path), + ) + encoder.model.eval() + batch_size = 16 + with torch.no_grad(): + for start in range(0, len(missing_prompts), batch_size): + batch_prompts = missing_prompts[start : start + batch_size] + input_ids, mask = encoder.tokenizer(batch_prompts, return_mask=True, add_special_tokens=True) + input_ids = input_ids.to(device) + mask = mask.to(device=device, dtype=torch.bool) + context = encoder.model(input_ids, mask) + for index, prompt in enumerate(batch_prompts): + _atomic_torch_save( + { + "context": context[index].detach().to(device="cpu", dtype=torch.bfloat16).contiguous(), + "mask": mask[index].detach().to(device="cpu", dtype=torch.bool).contiguous(), + }, + _text_cache_path(cache_dir, prompt, context_len), + ) + + del context, input_ids, mask, encoder + gc.collect() + torch.cuda.empty_cache() + logger.info("[data-preflight] generated {} text embeddings in {}", len(missing_prompts), cache_dir) + + +def _is_missing(value): + return value is None or not str(value).strip() + + +def prepare_libero_fastwam_assets(config): + data_config = config["data"] + train_config = data_config["train"] + val_config = data_config.get("val") + output_dir = _resolve_path(config["training"]["output_dir"]) + + configured_stats = train_config.get("pretrained_norm_stats") + if _is_missing(configured_stats): + stats_path = output_dir / "dataset_stats.json" + if is_main_process() and not stats_path.is_file(): + logger.info("[data-preflight] calculating LIBERO normalization statistics") + stats = calculate_dataset_stats(train_config) + _atomic_json_dump(stats, stats_path) + logger.info( + "[data-preflight] saved normalization statistics: episodes={} transitions={} path={}", + stats["num_episodes"], + stats["num_transition"], + stats_path, + ) + barrier() + else: + stats_path = _resolve_path(configured_stats) + if not stats_path.is_file(): + raise FileNotFoundError(f"Configured LIBERO normalization stats do not exist: {stats_path}") + eval_stats_path = output_dir / "dataset_stats.json" + if is_main_process() and stats_path != eval_stats_path: + eval_stats_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(stats_path, eval_stats_path) + barrier() + + train_config["pretrained_norm_stats"] = str(stats_path) + if val_config is not None and _is_missing(val_config.get("pretrained_norm_stats")): + val_config["pretrained_norm_stats"] = str(stats_path) + + configured_cache = train_config.get("text_embedding_cache_dir") + auto_text_cache = _is_missing(configured_cache) + cache_dir = output_dir / "text_embeds_cache" if auto_text_cache else _resolve_path(configured_cache) + train_config["text_embedding_cache_dir"] = str(cache_dir) + if val_config is not None and _is_missing(val_config.get("text_embedding_cache_dir")): + val_config["text_embedding_cache_dir"] = str(cache_dir) + + if auto_text_cache: + generated_splits = [train_config] + if val_config is not None and _resolve_path(val_config["text_embedding_cache_dir"]) == cache_dir: + generated_splits.append(val_config) + context_lengths = {int(split.get("context_len", 128)) for split in generated_splits} + if len(context_lengths) != 1: + raise ValueError(f"Automatic text cache generation requires one context_len, got {sorted(context_lengths)}") + context_len = context_lengths.pop() + prompts = _collect_prompts(generated_splits) + if is_main_process(): + precompute_text_embeddings( + model_path=config["model"]["model_path"], + cache_dir=cache_dir, + context_len=context_len, + prompts=prompts, + ) + barrier() + + _validate_text_caches([split for split in (train_config, val_config) if split is not None]) diff --git a/lightx2v_train/lightx2v_train/data/preparation.py b/lightx2v_train/lightx2v_train/data/preparation.py new file mode 100644 index 000000000..6fb13217f --- /dev/null +++ b/lightx2v_train/lightx2v_train/data/preparation.py @@ -0,0 +1,11 @@ +def prepare_data(config): + data_config = config.get("data", {}) + data_names = { + split_config.get("name") + for split_config in data_config.values() + if isinstance(split_config, dict) + } + if "libero_fastwam_dataset" in data_names: + from .libero.preparation import prepare_libero_fastwam_assets + + prepare_libero_fastwam_assets(config) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py index 1ac68dd88..658c39e44 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn +import torch.nn.functional as F from ..modules.model2_2 import rope_params from .layers import expert_block_forward, timestep_embedding @@ -11,6 +12,42 @@ logger = logging.getLogger(__name__) +def _interpolate_last_dim(tensor: torch.Tensor, new_size: int) -> torch.Tensor: + if tensor.shape[-1] == new_size: + return tensor + flat = tensor.reshape(-1, 1, tensor.shape[-1]).to(torch.float32) + flat = F.interpolate(flat, size=new_size, mode="linear", align_corners=True) + return flat.reshape(*tensor.shape[:-1], new_size) + + +def _resize_tensor_to_shape(source: torch.Tensor, target_shape: tuple[int, ...]) -> torch.Tensor: + if tuple(source.shape) == target_shape: + return source + + output = source.to(torch.float32) + while output.ndim < len(target_shape): + output = output.unsqueeze(0) + while output.ndim > len(target_shape): + if output.shape[0] != 1: + raise ValueError(f"Cannot reduce tensor rank for resize: source={tuple(source.shape)} target={target_shape}") + output = output.squeeze(0) + + for dimension, new_size in enumerate(target_shape): + if output.shape[dimension] == new_size: + continue + permutation = [index for index in range(output.ndim) if index != dimension] + [dimension] + inverse_permutation = [0] * output.ndim + for index, original_dimension in enumerate(permutation): + inverse_permutation[original_dimension] = index + permuted = output.permute(*permutation).contiguous() + permuted = _interpolate_last_dim(permuted, new_size) + output = permuted.permute(*inverse_permutation).contiguous() + + if tuple(output.shape) != target_shape: + raise ValueError(f"Resize produced wrong shape: source={tuple(source.shape)} target={target_shape} got={tuple(output.shape)}") + return output.to(dtype=source.dtype) + + class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() @@ -136,6 +173,47 @@ def backbone_key_set(cls, keys) -> set[str]: if not any(key.startswith(prefix) for prefix in cls.ACTION_BACKBONE_SKIP_PREFIXES) } + @classmethod + def from_video_expert( + cls, + action_dit_config: dict[str, Any], + video_expert: nn.Module, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + apply_alpha_scaling: bool = True, + ) -> "ActionDiT": + action_expert = cls(**action_dit_config).to(device=device, dtype=torch_dtype) + action_state = action_expert.state_dict() + video_state = video_expert.state_dict() + copied = 0 + interpolated = 0 + + with torch.no_grad(): + for key in sorted(cls.backbone_key_set(action_state.keys())): + if key not in video_state: + raise ValueError(f"ActionDiT backbone key `{key}` does not exist in the Wan video expert.") + source = video_state[key] + target = action_state[key] + if tuple(source.shape) == tuple(target.shape): + value = source + copied += 1 + else: + value = _resize_tensor_to_shape(source, tuple(target.shape)) + if apply_alpha_scaling and source.ndim >= 2 and source.shape[-1] != target.shape[-1]: + alpha = (float(source.shape[-1]) / float(target.shape[-1])) ** 0.5 + value = value.to(torch.float32).mul_(alpha) + interpolated += 1 + target.copy_(value.to(device=target.device, dtype=target.dtype)) + + logger.info( + "Initialized ActionDiT backbone from Wan video expert " + "(copied=%d, interpolated=%d, random_kept_prefixes=%s).", + copied, + interpolated, + list(cls.ACTION_BACKBONE_SKIP_PREFIXES), + ) + return action_expert + @classmethod def from_pretrained( cls, diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py index 3426eefb7..d9d616373 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py @@ -150,13 +150,21 @@ def from_wan22_pretrained( ) video_expert = components.dit - action_expert = ActionDiT.from_pretrained( - action_dit_config=action_dit_config, - action_dit_pretrained_path=action_dit_pretrained_path, - skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, - device=device, - torch_dtype=torch_dtype, - ) + if action_dit_pretrained_path or skip_dit_load_from_pretrain: + action_expert = ActionDiT.from_pretrained( + action_dit_config=action_dit_config, + action_dit_pretrained_path=action_dit_pretrained_path, + skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, + device=device, + torch_dtype=torch_dtype, + ) + else: + action_expert = ActionDiT.from_video_expert( + action_dit_config=action_dit_config, + video_expert=video_expert, + device=device, + torch_dtype=torch_dtype, + ) if int(action_expert.num_heads) != int(video_expert.num_heads): raise ValueError("ActionDiT `num_heads` must match video expert for MoT mixed attention.") if int(action_expert.attn_head_dim) != int(video_expert.attn_head_dim): diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py b/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py index 016059a5e..ec978d6b2 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py @@ -1,6 +1,5 @@ import os from contextlib import nullcontext -from pathlib import Path import torch from loguru import logger @@ -9,14 +8,6 @@ from lightx2v_train.utils.registry import MODEL_REGISTER from lightx2v_train.utils.utils import get_running_dtype -DEFAULT_ACTION_DIT_PATH = ( - Path(__file__).resolve().parents[2] - / "assets" - / "libero_fastwam" - / "checkpoints" - / "ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt" -) - def _resolve_local_path(path, name, directory=False): if not path: @@ -43,9 +34,11 @@ def load_components(self, transformer_only=False, reference_model=None): raise ValueError("wan_fastwam does not support transformer_only loading.") model_path = _resolve_local_path(self.model_config.get("model_path"), "model_path", directory=True) - action_dit_pretrained_path = _resolve_local_path( - self.model_config.get("action_dit_pretrained_path", DEFAULT_ACTION_DIT_PATH), - "action_dit_pretrained_path", + configured_action_dit_path = self.model_config.get("action_dit_pretrained_path") + action_dit_pretrained_path = ( + _resolve_local_path(configured_action_dit_path, "action_dit_pretrained_path") + if configured_action_dit_path + else None ) load_text_encoder = bool(self.model_config.get("load_text_encoder", False)) diff --git a/lightx2v_train/lightx2v_train/trainers/fastwam.py b/lightx2v_train/lightx2v_train/trainers/fastwam.py index ea1508ec3..7171e5bb3 100644 --- a/lightx2v_train/lightx2v_train/trainers/fastwam.py +++ b/lightx2v_train/lightx2v_train/trainers/fastwam.py @@ -25,6 +25,38 @@ from lightx2v_train.utils.registry import TRAINER_REGISTER from lightx2v_train.utils.video import pil_frames_to_video_tensor, save_mp4, video_psnr, video_ssim +ZERO1_PER_RANK_FORMAT = "zero1_per_rank_v1" +FULL_OPTIMIZER_FORMAT = "full_v1" + + +def _optimizer_shard_filename(rank, world_size): + return f"optimizer-rank-{rank:05d}-of-{world_size:05d}.pt" + + +def _temporary_path(path): + return f"{path}.tmp-{os.getpid()}-{get_rank()}" + + +def _atomic_torch_save(payload, path): + temporary_path = _temporary_path(path) + try: + torch.save(payload, temporary_path) + os.replace(temporary_path, path) + finally: + if os.path.exists(temporary_path): + os.remove(temporary_path) + + +def _sync_optimizer_group_options(source_groups, target_groups): + if len(source_groups) != len(target_groups): + raise RuntimeError( + f"Optimizer parameter-group count mismatch: source={len(source_groups)} target={len(target_groups)}" + ) + for source, target in zip(source_groups, target_groups): + for key, value in source.items(): + if key != "params": + target[key] = value + def _unwrap_dataset(dataset): while hasattr(dataset, "dataset"): @@ -173,42 +205,118 @@ def _load_resume_state(self, resume_ckpt_path): state = torch.load(state_path, map_location="cpu", weights_only=False) if state.get("world_size") != get_world_size(): raise RuntimeError(f"Cannot resume world_size={state.get('world_size')} checkpoint with world_size={get_world_size()}.") - self.optimizer.load_state_dict(state["optimizer"]) + optimizer_state_format = state.get("optimizer_state_format") + expected_optimizer_sharding = "zero1" if self.zero1_enabled else "none" + if state.get("optimizer_sharding") != expected_optimizer_sharding: + raise RuntimeError( + "Optimizer sharding mode mismatch: " + f"checkpoint={state.get('optimizer_sharding')!r} active={expected_optimizer_sharding!r}." + ) + if optimizer_state_format == ZERO1_PER_RANK_FORMAT: + if not self.zero1_enabled: + raise RuntimeError("A per-rank ZeRO-1 optimizer checkpoint requires distributed.zero1.enabled=true.") + shard_files = state.get("optimizer_shard_files") + if not isinstance(shard_files, list) or len(shard_files) != get_world_size(): + raise RuntimeError( + f"Invalid ZeRO-1 optimizer shard manifest: expected {get_world_size()} files, got {shard_files!r}." + ) + shard_path = os.path.join(resume_ckpt_path, shard_files[get_rank()]) + if not os.path.isfile(shard_path): + raise RuntimeError(f"ZeRO-1 optimizer shard not found for rank {get_rank()}: {shard_path}") + shard = torch.load(shard_path, map_location="cpu", weights_only=False) + if shard.get("rank") != get_rank() or shard.get("world_size") != get_world_size(): + raise RuntimeError( + "ZeRO-1 optimizer shard metadata mismatch: " + f"rank={shard.get('rank')} world_size={shard.get('world_size')} " + f"expected_rank={get_rank()} expected_world_size={get_world_size()}." + ) + self.optimizer.optim.load_state_dict(shard["optimizer"]) + _sync_optimizer_group_options(self.optimizer.optim.param_groups, self.optimizer.param_groups) + elif optimizer_state_format == FULL_OPTIMIZER_FORMAT: + if self.zero1_enabled: + raise RuntimeError("A full optimizer checkpoint requires distributed.zero1.enabled=false.") + self.optimizer.load_state_dict(state["optimizer"]) + else: + raise RuntimeError( + f"Unsupported optimizer checkpoint format {optimizer_state_format!r} in {state_path}." + ) self.lr_scheduler.load_state_dict(state["lr_scheduler"]) - saved_optimizer_sharding = state.get("optimizer_sharding", "none") - active_optimizer_sharding = "zero1" if self.zero1_enabled else "none" logger.info( - "[resume] restored FastWAM training state from {} (optimizer sharding: {} -> {})", + "[resume] restored FastWAM training state from {} (optimizer sharding: {}, state format: {})", resume_ckpt_path, - saved_optimizer_sharding, - active_optimizer_sharding, + expected_optimizer_sharding, + optimizer_state_format, ) def _save_checkpoint(self, iteration): + save_start_time = time.perf_counter() if is_main_process(): prune_checkpoints(self.output_train_dir, self.save_total_limit) save_dir = os.path.join(self.output_train_dir, f"checkpoint-{iteration:09d}") logger.info("[checkpoint] saving FastWAM iter={} path={}", iteration, save_dir) - if self.zero1_enabled: - self.optimizer.consolidate_state_dict(to=0) if is_main_process(): os.makedirs(save_dir, exist_ok=True) - self.model.save_checkpoint(os.path.join(save_dir, "fastwam.pt"), step=iteration) - torch.save( + barrier() + + optimizer_state_format = FULL_OPTIMIZER_FORMAT + optimizer_shard_files = None + if self.zero1_enabled: + optimizer_state_format = ZERO1_PER_RANK_FORMAT + optimizer_shard_files = [ + _optimizer_shard_filename(rank, get_world_size()) + for rank in range(get_world_size()) + ] + _sync_optimizer_group_options(self.optimizer.param_groups, self.optimizer.optim.param_groups) + shard_path = os.path.join(save_dir, optimizer_shard_files[get_rank()]) + shard_start_time = time.perf_counter() + _atomic_torch_save( { - "iteration": iteration, + "rank": get_rank(), "world_size": get_world_size(), - "optimizer_sharding": "zero1" if self.zero1_enabled else "none", - "optimizer": self.optimizer.state_dict(), - "lr_scheduler": self.lr_scheduler.state_dict(), + "optimizer": self.optimizer.optim.state_dict(), }, - os.path.join(save_dir, "training_state.pt"), + shard_path, + ) + logger.info( + "[checkpoint] saved optimizer shard rank={} duration={:.3f}s path={}", + get_rank(), + time.perf_counter() - shard_start_time, + shard_path, ) + barrier() + + if is_main_process(): + weights_path = os.path.join(save_dir, "fastwam.pt") + temporary_weights_path = _temporary_path(weights_path) + try: + self.model.save_checkpoint(temporary_weights_path, step=iteration) + os.replace(temporary_weights_path, weights_path) + finally: + if os.path.exists(temporary_weights_path): + os.remove(temporary_weights_path) + + training_state = { + "iteration": iteration, + "world_size": get_world_size(), + "optimizer_sharding": "zero1" if self.zero1_enabled else "none", + "optimizer_state_format": optimizer_state_format, + "lr_scheduler": self.lr_scheduler.state_dict(), + } + if self.zero1_enabled: + training_state["optimizer_shard_files"] = optimizer_shard_files + else: + training_state["optimizer"] = self.optimizer.state_dict() + _atomic_torch_save(training_state, os.path.join(save_dir, "training_state.pt")) config_path = self.config.get("config_path") if config_path is not None: shutil.copy2(config_path, os.path.join(save_dir, "config.yaml")) barrier() - logger.info("[checkpoint] saved FastWAM iter={} path={}", iteration, save_dir) + logger.info( + "[checkpoint] saved FastWAM iter={} duration={:.3f}s path={}", + iteration, + time.perf_counter() - save_start_time, + save_dir, + ) def _forward_loss(self, sample): with self.model.autocast_context(): diff --git a/lightx2v_train/train.py b/lightx2v_train/train.py index 6cddb2d70..37e0c7b16 100644 --- a/lightx2v_train/train.py +++ b/lightx2v_train/train.py @@ -1,8 +1,7 @@ import argparse import torch - -from lightx2v_train.data import build_data +from lightx2v_train.data import build_data, prepare_data from lightx2v_train.model_zoo import build_model from lightx2v_train.runtime import cleanup_distributed, init_distributed, load_config, setup_logger from lightx2v_train.trainers import build_trainer @@ -23,6 +22,7 @@ def main(): setup_logger(config) try: + prepare_data(config) model = build_model(config) model.load_components() From eaa715e5669681108022ef519685a98c9c322fa2 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Mon, 13 Jul 2026 13:22:49 +0000 Subject: [PATCH 06/10] update --- .../train/fastwam/libero_uncond_2cam224.yaml | 11 ++++++----- lightx2v_train/scripts/run_fastwam_libero_train.sh | 14 +++++++------- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml index 79ce41ab1..b9a32b45a 100644 --- a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml +++ b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml @@ -7,10 +7,11 @@ data: train: &libero_data name: libero_fastwam_dataset dataset_dirs: - - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_spatial_no_noops_lerobot - - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_object_no_noops_lerobot - - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_goal_no_noops_lerobot - - /data/nvme0/chendingyu/LightX2V/lightx2v_train/assets/libero_fastwam/datasets/libero_10_no_noops_lerobot + # https://huggingface.co/datasets/yuanty/LIBERO-fastwam + - /path/to/LIBERO-fastwam/libero_spatial_no_noops_lerobot + - /path/to/LIBERO-fastwam/libero_object_no_noops_lerobot + - /path/to/LIBERO-fastwam/libero_goal_no_noops_lerobot + - /path/to/LIBERO-fastwam/libero_10_no_noops_lerobot batch_size: 8 num_workers: 8 val: @@ -18,7 +19,7 @@ data: training: method: fastwam - output_dir: /data/nvme0/chendingyu/LightX2V/lightx2v_train/runs/fastwam_libero_full + output_dir: runs/fastwam_libero_full # Equivalent to the official 10-epoch, global-batch-128 LIBERO schedule. max_train_iters: 21700 gradient_accumulation_iters: 4 diff --git a/lightx2v_train/scripts/run_fastwam_libero_train.sh b/lightx2v_train/scripts/run_fastwam_libero_train.sh index 233ddbe8e..6d93f44f3 100755 --- a/lightx2v_train/scripts/run_fastwam_libero_train.sh +++ b/lightx2v_train/scripts/run_fastwam_libero_train.sh @@ -1,10 +1,10 @@ -#!/usr/bin/env bash -set -euo pipefail +#!/bin/bash -LIGHTX2V_ROOT="${LIGHTX2V_ROOT:-/data/nvme0/chendingyu/LightX2V}" -CONFIG_PATH="${CONFIG_PATH:-${LIGHTX2V_ROOT}/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml}" +export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3} -export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-1}" +NPROC_PER_NODE=${NPROC_PER_NODE:-4} -cd "${LIGHTX2V_ROOT}" -python lightx2v_train/train.py --config "${CONFIG_PATH}" "$@" +torchrun \ +--standalone \ +--nproc_per_node="${NPROC_PER_NODE}" \ +train.py --config configs/train/fastwam/libero_uncond_2cam224.yaml From fdef47a18ab4a45659c8aca673d882350ce684bb Mon Sep 17 00:00:00 2001 From: Musisoul Date: Mon, 13 Jul 2026 13:26:55 +0000 Subject: [PATCH 07/10] lint --- lightx2v_train/train.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lightx2v_train/train.py b/lightx2v_train/train.py index 37e0c7b16..6c93b4e61 100644 --- a/lightx2v_train/train.py +++ b/lightx2v_train/train.py @@ -1,6 +1,7 @@ import argparse import torch + from lightx2v_train.data import build_data, prepare_data from lightx2v_train.model_zoo import build_model from lightx2v_train.runtime import cleanup_distributed, init_distributed, load_config, setup_logger From 84627552bdf851f60f925cb96642ac30da6cb75b Mon Sep 17 00:00:00 2001 From: Musisoul Date: Mon, 13 Jul 2026 13:27:38 +0000 Subject: [PATCH 08/10] lint --- .../data/libero/lerobot_dataset.py | 18 +--- .../lightx2v_train/data/libero/normalizer.py | 6 +- .../lightx2v_train/data/libero/preparation.py | 21 +--- .../lightx2v_train/data/libero/processor.py | 8 +- .../data/libero/robot_video_dataset.py | 5 +- .../data/libero/video_decoder.py | 4 +- .../lightx2v_train/data/preparation.py | 6 +- .../native/wan/fastwam/action_dit.py | 83 ++++------------ .../model_zoo/native/wan/fastwam/layers.py | 20 +--- .../model_zoo/native/wan/fastwam/loader.py | 8 +- .../model_zoo/native/wan/fastwam/model.py | 95 +++++-------------- .../model_zoo/native/wan/fastwam/mot.py | 63 ++++-------- .../model_zoo/native/wan/fastwam/video_dit.py | 34 ++----- .../lightx2v_train/model_zoo/wan_fastwam.py | 6 +- .../schedulers/flow_matching.py | 4 +- .../lightx2v_train/trainers/fastwam.py | 57 +++-------- 16 files changed, 105 insertions(+), 333 deletions(-) diff --git a/lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py b/lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py index 4af8c3312..a213441af 100644 --- a/lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py +++ b/lightx2v_train/lightx2v_train/data/libero/lerobot_dataset.py @@ -86,10 +86,7 @@ def __init__( for dataset_dir in dataset_dirs: root = Path(dataset_dir).expanduser().resolve() info = _read_json(root / "meta" / "info.json") - tasks = { - int(item["task_index"]): item["task"] - for item in _read_jsonl(root / "meta" / "tasks.jsonl") - } + tasks = {int(item["task_index"]): item["task"] for item in _read_jsonl(root / "meta" / "tasks.jsonl")} episode_meta = _read_jsonl(root / "meta" / "episodes.jsonl") fps_values.add(int(info["fps"])) self._validate_features(root, info["features"]) @@ -163,12 +160,8 @@ def __getitem__(self, index): episode = self.episodes[episode_position] data = self._load_episode(episode_position) - state, state_is_pad, observation_indices = self._window( - data[self.state_key], frame_index, self.num_frames - ) - action, action_is_pad, _ = self._window( - data[self.action_key], frame_index, self.num_frames - 1 - ) + state, state_is_pad, observation_indices = self._window(data[self.state_key], frame_index, self.num_frames) + action, action_is_pad, _ = self._window(data[self.action_key], frame_index, self.num_frames - 1) timestamps = data["timestamp"][observation_indices].float().tolist() tolerance = max(1e-4, 1.0 / self.fps - 1e-4) images = { @@ -187,10 +180,7 @@ def __getitem__(self, index): "task": episode.tasks[task_index], "action": {"default": action.float()}, "state": {"default": state.float()}, - "images": { - key.removeprefix("observation.images."): value - for key, value in images.items() - }, + "images": {key.removeprefix("observation.images."): value for key, value in images.items()}, "action_is_pad": action_is_pad, "state_is_pad": state_is_pad, "image_is_pad": state_is_pad.clone(), diff --git a/lightx2v_train/lightx2v_train/data/libero/normalizer.py b/lightx2v_train/lightx2v_train/data/libero/normalizer.py index a18718991..91cbbac46 100644 --- a/lightx2v_train/lightx2v_train/data/libero/normalizer.py +++ b/lightx2v_train/lightx2v_train/data/libero/normalizer.py @@ -51,11 +51,7 @@ def __init__(self, shape_meta, stats, mode="min/max"): for group in self.normalizers: for item in shape_meta[group]: key = item["key"] - global_stats = { - name.removeprefix("global_"): value - for name, value in stats[group][key].items() - if name.startswith("global_") - } + global_stats = {name.removeprefix("global_"): value for name, value in stats[group][key].items() if name.startswith("global_")} self.normalizers[group][key] = FieldNormalizer(global_stats, mode=mode) def forward(self, batch): diff --git a/lightx2v_train/lightx2v_train/data/libero/preparation.py b/lightx2v_train/lightx2v_train/data/libero/preparation.py index b4741ae7f..8f2a06629 100644 --- a/lightx2v_train/lightx2v_train/data/libero/preparation.py +++ b/lightx2v_train/lightx2v_train/data/libero/preparation.py @@ -75,10 +75,7 @@ def _episode_field_stats(value): def _aggregate_field_stats(episode_stats): - values = { - name: torch.stack([stats[name] for stats in episode_stats]) - for name in ("min", "max", "q01", "q99", "mean", "var") - } + values = {name: torch.stack([stats[name] for stats in episode_stats]) for name in ("min", "max", "q01", "q99", "mean", "var")} stepwise_mean = values["mean"].mean(0) stepwise_std = (values["var"] + (values["mean"] - stepwise_mean) ** 2).mean(0).sqrt() global_mean = values["mean"].mean((0, 1)) @@ -182,23 +179,13 @@ def _validate_text_caches(split_configs): cache_dir = _resolve_path(split_config["text_embedding_cache_dir"]) context_len = int(split_config.get("context_len", 128)) prompts = _collect_prompts([split_config]) - missing = [ - _text_cache_path(cache_dir, prompt, context_len) - for prompt in prompts - if not _valid_text_cache(_text_cache_path(cache_dir, prompt, context_len), context_len) - ] + missing = [_text_cache_path(cache_dir, prompt, context_len) for prompt in prompts if not _valid_text_cache(_text_cache_path(cache_dir, prompt, context_len), context_len)] if missing: - raise FileNotFoundError( - f"LIBERO text embedding cache is incomplete: missing={len(missing)}/{len(prompts)} first={missing[0]}" - ) + raise FileNotFoundError(f"LIBERO text embedding cache is incomplete: missing={len(missing)}/{len(prompts)} first={missing[0]}") def precompute_text_embeddings(model_path, cache_dir, context_len, prompts): - missing_prompts = [ - prompt - for prompt in prompts - if not _valid_text_cache(_text_cache_path(cache_dir, prompt, context_len), context_len) - ] + missing_prompts = [prompt for prompt in prompts if not _valid_text_cache(_text_cache_path(cache_dir, prompt, context_len), context_len)] if not missing_prompts: logger.info("[data-preflight] text embeddings already cached: prompts={} path={}", len(prompts), cache_dir) return diff --git a/lightx2v_train/lightx2v_train/data/libero/processor.py b/lightx2v_train/lightx2v_train/data/libero/processor.py index 2f8359590..cd2747ba9 100644 --- a/lightx2v_train/lightx2v_train/data/libero/processor.py +++ b/lightx2v_train/lightx2v_train/data/libero/processor.py @@ -67,9 +67,7 @@ def _validate_action_state(self, data): actual = data[group][item["key"]].shape[-1] expected = int(item["raw_shape"]) if actual != expected: - raise ValueError( - f"{group}.{item['key']} width mismatch: expected {expected}, got {actual}" - ) + raise ValueError(f"{group}.{item['key']} width mismatch: expected {expected}, got {actual}") def _process_images(self, images): processed = [] @@ -97,9 +95,7 @@ def preprocess(self, data): mask = action_is_pad[:, None] & self.delta_action_dim_mask[None, :] action[mask] = 0.0 - transformed = self.normalizer.forward( - {"action": data["action"], "state": data["state"]} - ) + transformed = self.normalizer.forward({"action": data["action"], "state": data["state"]}) transformed = self.action_state_merger.forward(transformed) return { "idx": data["idx"], diff --git a/lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py b/lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py index ee088937f..76ff3d147 100644 --- a/lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py +++ b/lightx2v_train/lightx2v_train/data/libero/robot_video_dataset.py @@ -68,10 +68,7 @@ def _sample_without_padding(self, index): sample = self.processor.preprocess(self.lerobot_dataset[index]) if not self.skip_padding_as_possible: return sample - has_padding = any( - bool(sample[key].any()) - for key in ("action_is_pad", "image_is_pad", "proprio_is_pad") - ) + has_padding = any(bool(sample[key].any()) for key in ("action_is_pad", "image_is_pad", "proprio_is_pad")) if not has_padding or attempt == self.max_padding_retry: return sample index = np.random.randint(len(self)) diff --git a/lightx2v_train/lightx2v_train/data/libero/video_decoder.py b/lightx2v_train/lightx2v_train/data/libero/video_decoder.py index 1a2e576c5..748be766b 100644 --- a/lightx2v_train/lightx2v_train/data/libero/video_decoder.py +++ b/lightx2v_train/lightx2v_train/data/libero/video_decoder.py @@ -19,9 +19,7 @@ def decode_video_frames(video_path, timestamps, tolerance_s, backend=None): try: return _decode_torchcodec(video_path, timestamps, tolerance_s) except Exception as error: - warnings.warn( - f"torchcodec decode failed ({type(error).__name__}: {error}); falling back to pyav." - ) + warnings.warn(f"torchcodec decode failed ({type(error).__name__}: {error}); falling back to pyav.") backend = "pyav" if backend not in {"pyav", "video_reader"}: raise ValueError(f"Unsupported video backend: {backend}") diff --git a/lightx2v_train/lightx2v_train/data/preparation.py b/lightx2v_train/lightx2v_train/data/preparation.py index 6fb13217f..fc3bb0ea9 100644 --- a/lightx2v_train/lightx2v_train/data/preparation.py +++ b/lightx2v_train/lightx2v_train/data/preparation.py @@ -1,10 +1,6 @@ def prepare_data(config): data_config = config.get("data", {}) - data_names = { - split_config.get("name") - for split_config in data_config.values() - if isinstance(split_config, dict) - } + data_names = {split_config.get("name") for split_config in data_config.values() if isinstance(split_config, dict)} if "libero_fastwam_dataset" in data_names: from .libero.preparation import prepare_libero_fastwam_assets diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py index 658c39e44..d05773b44 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/action_dit.py @@ -167,11 +167,7 @@ def __init__( @classmethod def backbone_key_set(cls, keys) -> set[str]: - return { - key - for key in keys - if not any(key.startswith(prefix) for prefix in cls.ACTION_BACKBONE_SKIP_PREFIXES) - } + return {key for key in keys if not any(key.startswith(prefix) for prefix in cls.ACTION_BACKBONE_SKIP_PREFIXES)} @classmethod def from_video_expert( @@ -206,8 +202,7 @@ def from_video_expert( target.copy_(value.to(device=target.device, dtype=target.dtype)) logger.info( - "Initialized ActionDiT backbone from Wan video expert " - "(copied=%d, interpolated=%d, random_kept_prefixes=%s).", + "Initialized ActionDiT backbone from Wan video expert (copied=%d, interpolated=%d, random_kept_prefixes=%s).", copied, interpolated, list(cls.ACTION_BACKBONE_SKIP_PREFIXES), @@ -226,23 +221,19 @@ def from_pretrained( if action_dit_config is None: raise ValueError("`action_dit_config` is required for ActionDiT.from_pretrained().") if skip_dit_load_from_pretrain: - logger.info( - "Skipping ActionDiT pretrained load (`skip_dit_load_from_pretrain=True`); " - "initializing action expert randomly and expecting checkpoint override." - ) + logger.info("Skipping ActionDiT pretrained load (`skip_dit_load_from_pretrain=True`); initializing action expert randomly and expecting checkpoint override.") return cls(**action_dit_config).to(device=device, dtype=torch_dtype) if not action_dit_pretrained_path: logger.info("No `action_dit_pretrained_path` provided, initializing ActionDiT with random weights.") return cls(**action_dit_config).to(device=device, dtype=torch_dtype) from pathlib import Path + p = Path(action_dit_pretrained_path) if not p.is_absolute(): p = Path(__file__).resolve().parents[4] / p action_dit_pretrained_path = str(p) if not os.path.isfile(action_dit_pretrained_path): - raise FileNotFoundError( - f"`action_dit_pretrained_path` does not exist: {action_dit_pretrained_path}" - ) + raise FileNotFoundError(f"`action_dit_pretrained_path` does not exist: {action_dit_pretrained_path}") action_cfg = dict(action_dit_config) action_expert = cls(**action_cfg).to(device=device, dtype=torch_dtype) @@ -251,10 +242,8 @@ def from_pretrained( payload = torch.load(action_dit_pretrained_path, map_location="cpu", weights_only=True) if not isinstance(payload, dict): - raise ValueError( - f"Invalid action backbone payload type from {action_dit_pretrained_path}: {type(payload)}" - ) - + raise ValueError(f"Invalid action backbone payload type from {action_dit_pretrained_path}: {type(payload)}") + policy = payload.get("policy", {}) if policy: logger.info(f"ActionDiT backbone payload policy: {policy}") @@ -277,22 +266,13 @@ def from_pretrained( got_value = meta[key] if key == "eps": if abs(float(got_value) - float(expected_value)) > 1e-12: - raise ValueError( - f"`meta.{key}` mismatch in {action_dit_pretrained_path}: " - f"expected {expected_value}, got {got_value}" - ) + raise ValueError(f"`meta.{key}` mismatch in {action_dit_pretrained_path}: expected {expected_value}, got {got_value}") elif int(got_value) != int(expected_value): - raise ValueError( - f"`meta.{key}` mismatch in {action_dit_pretrained_path}: " - f"expected {expected_value}, got {got_value}" - ) + raise ValueError(f"`meta.{key}` mismatch in {action_dit_pretrained_path}: expected {expected_value}, got {got_value}") backbone_state_dict = payload.get("backbone_state_dict") if not isinstance(backbone_state_dict, dict): - raise ValueError( - f"`backbone_state_dict` must be a dict in {action_dit_pretrained_path}, " - f"got {type(backbone_state_dict)}" - ) + raise ValueError(f"`backbone_state_dict` must be a dict in {action_dit_pretrained_path}, got {type(backbone_state_dict)}") provided_keys = set(backbone_state_dict.keys()) missing_keys = sorted(expected_backbone_keys - provided_keys) @@ -308,16 +288,10 @@ def from_pretrained( for key in expected_backbone_keys: value = backbone_state_dict[key] if not isinstance(value, torch.Tensor): - raise ValueError( - f"`backbone_state_dict[{key}]` must be torch.Tensor in {action_dit_pretrained_path}, " - f"got {type(value)}" - ) + raise ValueError(f"`backbone_state_dict[{key}]` must be torch.Tensor in {action_dit_pretrained_path}, got {type(value)}") target = merged_state[key] if tuple(value.shape) != tuple(target.shape): - raise ValueError( - f"Shape mismatch for `{key}` in {action_dit_pretrained_path}: " - f"expected {tuple(target.shape)}, got {tuple(value.shape)}" - ) + raise ValueError(f"Shape mismatch for `{key}` in {action_dit_pretrained_path}: expected {tuple(target.shape)}, got {tuple(value.shape)}") merged_state[key] = value.to(device=target.device, dtype=target.dtype) action_expert.load_state_dict(merged_state, strict=True) @@ -337,51 +311,35 @@ def pre_dit( context_mask: Optional[torch.Tensor] = None, ) -> Dict[str, Any]: if action_tokens.ndim != 3: - raise ValueError( - f"`action_tokens` must be 3D [B, T, action_dim], got shape {tuple(action_tokens.shape)}" - ) + raise ValueError(f"`action_tokens` must be 3D [B, T, action_dim], got shape {tuple(action_tokens.shape)}") if action_tokens.shape[2] != self.action_dim: - raise ValueError( - f"`action_tokens` last dim must be {self.action_dim}, got {action_tokens.shape[2]}" - ) + raise ValueError(f"`action_tokens` last dim must be {self.action_dim}, got {action_tokens.shape[2]}") if timestep.ndim != 1: raise ValueError(f"`timestep` must be 1D [B] or [1], got shape {tuple(timestep.shape)}") if context.ndim != 3: - raise ValueError( - f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}" - ) + raise ValueError(f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}") batch_size = action_tokens.shape[0] if context.shape[0] != batch_size: - raise ValueError( - f"Batch mismatch between action tokens and text context: {batch_size} vs {context.shape[0]}" - ) + raise ValueError(f"Batch mismatch between action tokens and text context: {batch_size} vs {context.shape[0]}") if timestep.shape[0] not in (1, batch_size): - raise ValueError( - f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}" - ) + raise ValueError(f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}") if timestep.shape[0] == 1 and batch_size > 1: if self.training: raise ValueError("During training, action timestep length must match batch_size.") timestep = timestep.expand(batch_size) if context_mask is None: - context_mask = torch.ones( - (batch_size, context.shape[1]), dtype=torch.bool, device=context.device - ) + context_mask = torch.ones((batch_size, context.shape[1]), dtype=torch.bool, device=context.device) else: if context_mask.ndim != 2: raise ValueError(f"`context_mask` must be 2D [B, L], got shape {tuple(context_mask.shape)}") if context_mask.shape[0] != batch_size or context_mask.shape[1] != context.shape[1]: - raise ValueError( - f"`context_mask` shape must match `context` shape [B, L], got {tuple(context_mask.shape)} vs {tuple(context.shape)}" - ) + raise ValueError(f"`context_mask` shape must match `context` shape [B, L], got {tuple(context_mask.shape)} vs {tuple(context.shape)}") seq_len = action_tokens.shape[1] if seq_len > self.freqs.shape[0]: - raise ValueError( - f"Action token length {seq_len} exceeds RoPE cache {self.freqs.shape[0]}." - ) + raise ValueError(f"Action token length {seq_len} exceeds RoPE cache {self.freqs.shape[0]}.") t = self.time_embedding(timestep_embedding(self.freq_dim, timestep)) t_mod = self.time_projection(t).unflatten(1, (6, self.hidden_dim)) @@ -427,6 +385,7 @@ def forward( context_mask = pre_state["context_mask"] for block in self.blocks: + def run(value, current_block=block): return expert_block_forward( current_block, diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py index 338431aa3..6309546bb 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/layers.py @@ -58,9 +58,7 @@ def layer_norm(module, value: torch.Tensor) -> torch.Tensor: return module(value) weight = None if module.weight is None else module.weight.float() bias = None if module.bias is None else module.bias.float() - return F.layer_norm( - value.float(), module.normalized_shape, weight, bias, module.eps - ).to(value.dtype) + return F.layer_norm(value.float(), module.normalized_shape, weight, bias, module.eps).to(value.dtype) def masked_cross_attention(block, x, context, context_mask=None): @@ -82,22 +80,14 @@ def expert_block_forward( context_mask: torch.Tensor | None = None, self_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = split_modulation( - block, time_modulation - ) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = split_modulation(block, time_modulation) attention_input = layer_norm(block.norm1, x) * (1 + scale_msa) + shift_msa batch_size, sequence_length = attention_input.shape[:2] num_heads = block.self_attn.num_heads head_dim = block.self_attn.head_dim - query = block.self_attn.norm_q(block.self_attn.q(attention_input)).view( - batch_size, sequence_length, num_heads, head_dim - ) - key = block.self_attn.norm_k(block.self_attn.k(attention_input)).view( - batch_size, sequence_length, num_heads, head_dim - ) - value = block.self_attn.v(attention_input).view( - batch_size, sequence_length, num_heads, head_dim - ) + query = block.self_attn.norm_q(block.self_attn.q(attention_input)).view(batch_size, sequence_length, num_heads, head_dim) + key = block.self_attn.norm_k(block.self_attn.k(attention_input)).view(batch_size, sequence_length, num_heads, head_dim) + value = block.self_attn.v(attention_input).view(batch_size, sequence_length, num_heads, head_dim) attention = scaled_attention( apply_rope(query, frequencies), apply_rope(key, frequencies), diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py index 9da059a6e..e7048066d 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/loader.py @@ -13,6 +13,8 @@ from .video_dit import FastWAMVideoDiT logger = logging.getLogger(__name__) + + @dataclass class Wan22LoadedComponents: dit: FastWAMVideoDiT @@ -26,11 +28,7 @@ def _validate_dit_config(dit_config: dict[str, Any]) -> dict[str, Any]: raise ValueError(f"video_dit_config must be a dict, got {type(dit_config)}") signature = inspect.signature(FastWAMVideoDiT.__init__) allowed = {name for name in signature.parameters if name != "self"} - required = { - name - for name, parameter in signature.parameters.items() - if name != "self" and parameter.default is inspect.Signature.empty - } + required = {name for name, parameter in signature.parameters.items() if name != "self" and parameter.default is inspect.Signature.empty} unknown = sorted(set(dit_config) - allowed) missing = sorted(required - set(dit_config)) if unknown: diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py index d9d616373..804eef573 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/model.py @@ -220,10 +220,7 @@ def _check_resize_height_width(height, width, num_frames): @torch.no_grad() def encode_prompt(self, prompt: Union[str, Sequence[str]]): if self.text_encoder is None or self.tokenizer is None: - raise ValueError( - "Prompt encoding requires loaded text encoder/tokenizer. " - "Set `load_text_encoder=true` or provide precomputed `context/context_mask`." - ) + raise ValueError("Prompt encoding requires loaded text encoder/tokenizer. Set `load_text_encoder=true` or provide precomputed `context/context_mask`.") ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True) ids = ids.to(self.device) mask = mask.to(self.device, dtype=torch.bool) @@ -246,12 +243,8 @@ def _append_proprio_to_context( if proprio.ndim != 2: raise ValueError(f"`proprio` must be 2D [B, D], got shape {tuple(proprio.shape)}") if self.proprio_dim is None or proprio.shape[1] != self.proprio_dim: - raise ValueError( - f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}" - ) - proprio_token = self.proprio_encoder( - proprio.to(device=self.device, dtype=context.dtype).unsqueeze(1) - ).to(dtype=context.dtype) # [B, 1, D] + raise ValueError(f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}") + proprio_token = self.proprio_encoder(proprio.to(device=self.device, dtype=context.dtype).unsqueeze(1)).to(dtype=context.dtype) # [B, 1, D] proprio_mask = torch.ones((context_mask.shape[0], 1), dtype=torch.bool, device=context_mask.device) return ( torch.cat([context, proprio_token], dim=1), @@ -271,9 +264,7 @@ def _encode_input_image_latents_tensor(self, input_image: torch.Tensor, tiled=Fa if input_image.ndim == 3: input_image = input_image.unsqueeze(0) if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: - raise ValueError( - f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" - ) + raise ValueError(f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}") image = input_image.to(device=self.device)[0].unsqueeze(1) del tile_size, tile_stride if tiled: @@ -296,9 +287,7 @@ def _decode_latents(self, latents, tiled=False, tile_size=(30, 52), tile_stride= def build_inputs(self, sample, tiled: bool = False): video = sample["video"] if "context" not in sample or "context_mask" not in sample: - raise ValueError( - "FastWAM training requires `sample['context']` and `sample['context_mask']`." - ) + raise ValueError("FastWAM training requires `sample['context']` and `sample['context_mask']`.") context = sample["context"] context_mask = sample["context_mask"] proprio = sample.get("proprio", None) @@ -309,9 +298,7 @@ def build_inputs(self, sample, tiled: bool = False): batch_size, _, num_frames, height, width = video.shape if height % 16 != 0 or width % 16 != 0: - raise ValueError( - f"Video spatial dims must be multiples of 16, got H={height}, W={width}" - ) + raise ValueError(f"Video spatial dims must be multiples of 16, got H={height}, W={width}") if num_frames % 4 != 1: raise ValueError(f"Video T must satisfy T % 4 == 1, got T={num_frames}") if num_frames <= 1: @@ -325,34 +312,22 @@ def build_inputs(self, sample, tiled: bool = False): raise ValueError(f"`sample['action']` must be 3D [B, T, a_dim], got shape {tuple(action.shape)}") action_horizon = int(action.shape[1]) if action_horizon % (num_frames - 1) != 0: - raise ValueError( - f"`sample['action']` temporal dimension must be divisible by video transitions ({num_frames - 1}), got {action_horizon}" - ) + raise ValueError(f"`sample['action']` temporal dimension must be divisible by video transitions ({num_frames - 1}), got {action_horizon}") action_is_pad = sample.get("action_is_pad", None) if action_is_pad is not None: if action_is_pad.ndim != 2: - raise ValueError( - f"`sample['action_is_pad']` must be 2D [B, T], got shape {tuple(action_is_pad.shape)}" - ) + raise ValueError(f"`sample['action_is_pad']` must be 2D [B, T], got shape {tuple(action_is_pad.shape)}") if action_is_pad.shape[0] != batch_size or action_is_pad.shape[1] != action_horizon: - raise ValueError( - "`sample['action_is_pad']` shape mismatch: " - f"got {tuple(action_is_pad.shape)} vs expected ({batch_size}, {action_horizon})" - ) + raise ValueError(f"`sample['action_is_pad']` shape mismatch: got {tuple(action_is_pad.shape)} vs expected ({batch_size}, {action_horizon})") image_is_pad = sample.get("image_is_pad", None) if image_is_pad is not None: if image_is_pad.ndim != 2: - raise ValueError( - f"`sample['image_is_pad']` must be 2D [B, T], got shape {tuple(image_is_pad.shape)}" - ) + raise ValueError(f"`sample['image_is_pad']` must be 2D [B, T], got shape {tuple(image_is_pad.shape)}") if image_is_pad.shape[0] != batch_size or image_is_pad.shape[1] != num_frames: - raise ValueError( - "`sample['image_is_pad']` shape mismatch: " - f"got {tuple(image_is_pad.shape)} vs expected ({batch_size}, {num_frames})" - ) - + raise ValueError(f"`sample['image_is_pad']` shape mismatch: got {tuple(image_is_pad.shape)} vs expected ({batch_size}, {num_frames})") + input_video = video.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) input_latents = self._encode_video_latents(input_video, tiled=tiled) @@ -363,9 +338,7 @@ def build_inputs(self, sample, tiled: bool = False): fuse_flag = True if context.ndim != 3 or context_mask.ndim != 2: - raise ValueError( - f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" - ) + raise ValueError(f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}") context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) if self.proprio_encoder is not None: @@ -374,10 +347,8 @@ def build_inputs(self, sample, tiled: bool = False): if proprio.ndim != 3: raise ValueError(f"`sample['proprio']` must be 3D [B, T, d], got shape {tuple(proprio.shape)}") if proprio.shape[2] != self.proprio_dim: - raise ValueError( - f"`sample['proprio']` last dim must be {self.proprio_dim}, got {proprio.shape[2]}" - ) - proprio = proprio[:, 0, :] # [B, D] + raise ValueError(f"`sample['proprio']` last dim must be {self.proprio_dim}, got {proprio.shape[2]}") + proprio = proprio[:, 0, :] # [B, D] context, context_mask = self._append_proprio_to_context( context=context, context_mask=context_mask, @@ -442,10 +413,7 @@ def _compute_video_loss_per_sample( if image_is_pad.shape[1] < 1: raise ValueError("`image_is_pad` must contain at least one frame.") if (image_is_pad.shape[1] - 1) % temporal_factor != 0: - raise ValueError( - "Cannot align `image_is_pad` with video latent steps: " - f"num_frames={image_is_pad.shape[1]}, temporal_downsample_factor={temporal_factor}." - ) + raise ValueError(f"Cannot align `image_is_pad` with video latent steps: num_frames={image_is_pad.shape[1]}, temporal_downsample_factor={temporal_factor}.") tail_is_pad = image_is_pad[:, 1:] latent_tail_is_pad = tail_is_pad.view(image_is_pad.shape[0], -1, temporal_factor).all(dim=2) @@ -455,10 +423,7 @@ def _compute_video_loss_per_sample( video_is_pad = latent_tail_is_pad if video_is_pad.shape[1] != video_loss_token.shape[1]: - raise ValueError( - "Video-loss mask shape mismatch: " - f"mask steps={video_is_pad.shape[1]}, loss steps={video_loss_token.shape[1]}." - ) + raise ValueError(f"Video-loss mask shape mismatch: mask steps={video_is_pad.shape[1]}, loss steps={video_loss_token.shape[1]}.") valid = (~video_is_pad).to(device=video_loss_token.device, dtype=video_loss_token.dtype) valid_sum = valid.sum(dim=1).clamp(min=1.0) @@ -561,12 +526,10 @@ def training_loss(self, sample, tiled: bool = False): image_is_pad=image_is_pad, include_initial_video_step=include_initial_video_step, ) - video_weight = self.train_video_scheduler.training_weight(timestep_video).to( - loss_video_per_sample.device, dtype=loss_video_per_sample.dtype - ) + video_weight = self.train_video_scheduler.training_weight(timestep_video).to(loss_video_per_sample.device, dtype=loss_video_per_sample.dtype) loss_video = (loss_video_per_sample * video_weight).mean() - action_loss_token = F.mse_loss(pred_action.float(), target_action.float(), reduction="none").mean(dim=2) # [B, T] + action_loss_token = F.mse_loss(pred_action.float(), target_action.float(), reduction="none").mean(dim=2) # [B, T] if action_is_pad is not None: valid = (~action_is_pad).to(device=action_loss_token.device, dtype=action_loss_token.dtype) valid_sum = valid.sum(dim=1).clamp(min=1.0) @@ -574,9 +537,7 @@ def training_loss(self, sample, tiled: bool = False): else: action_loss_per_sample = action_loss_token.mean(dim=1) - action_weight = self.train_action_scheduler.training_weight(timestep_action).to( - action_loss_per_sample.device, dtype=action_loss_per_sample.dtype - ) + action_weight = self.train_action_scheduler.training_weight(timestep_action).to(action_loss_per_sample.device, dtype=action_loss_per_sample.dtype) loss_action = (action_loss_per_sample * action_weight).mean() loss_total = self.loss_lambda_video * loss_video + self.loss_lambda_action * loss_action @@ -669,19 +630,13 @@ def infer_joint( if input_image.ndim == 3: input_image = input_image.unsqueeze(0) if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: - raise ValueError( - f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" - ) + raise ValueError(f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}") _, _, height, width = input_image.shape checked_h, checked_w, checked_t = self._check_resize_height_width(height, width, num_video_frames) if (checked_h, checked_w) != (height, width): - raise ValueError( - f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})" - ) + raise ValueError(f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})") if checked_t != num_video_frames: - raise ValueError( - f"`num_video_frames` must satisfy T % 4 == 1, got {num_video_frames}" - ) + raise ValueError(f"`num_video_frames` must satisfy T % 4 == 1, got {num_video_frames}") if proprio is not None: if self.proprio_dim is None: raise ValueError("`proprio` was provided but `proprio_dim=None` so `proprio_encoder` is disabled.") @@ -736,9 +691,7 @@ def infer_joint( if context_mask.ndim == 1: context_mask = context_mask.unsqueeze(0) if context.ndim != 3 or context_mask.ndim != 2: - raise ValueError( - f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" - ) + raise ValueError(f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}") context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) if proprio is not None: diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py index 389cbe0e2..9d2e8495d 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/mot.py @@ -37,19 +37,12 @@ def __init__( for name in self.expert_order[1:]: expert = self.mixtures[name] if len(expert.blocks) != self.num_layers: - raise ValueError( - f"All experts must have same number of layers; got {self.num_layers} and {len(expert.blocks)}" - ) + raise ValueError(f"All experts must have same number of layers; got {self.num_layers} and {len(expert.blocks)}") if expert.num_heads != self.num_heads: - raise ValueError( - f"All experts must have same num_heads; got {self.num_heads} and {expert.num_heads}" - ) + raise ValueError(f"All experts must have same num_heads; got {self.num_heads} and {expert.num_heads}") if expert.attn_head_dim != self.attn_head_dim: - raise ValueError( - "All experts must have same attn_head_dim; " - f"got {self.attn_head_dim} and {expert.attn_head_dim}" - ) - + raise ValueError(f"All experts must have same attn_head_dim; got {self.attn_head_dim} and {expert.attn_head_dim}") + logger.info(f"Initialized MoT with experts: {self.expert_order}, num_layers={self.num_layers}") for name in self.expert_order: expert = self.mixtures[name] @@ -148,15 +141,9 @@ def _build_expert_attention_io( attn_input = layer_norm(block.norm1, x) * (1 + scale_msa) + shift_msa batch_size, sequence_length = attn_input.shape[:2] - q = block.self_attn.norm_q(block.self_attn.q(attn_input)).view( - batch_size, sequence_length, self.num_heads, self.attn_head_dim - ) - k = block.self_attn.norm_k(block.self_attn.k(attn_input)).view( - batch_size, sequence_length, self.num_heads, self.attn_head_dim - ) - v = block.self_attn.v(attn_input).view( - batch_size, sequence_length, self.num_heads, self.attn_head_dim - ) + q = block.self_attn.norm_q(block.self_attn.q(attn_input)).view(batch_size, sequence_length, self.num_heads, self.attn_head_dim) + k = block.self_attn.norm_k(block.self_attn.k(attn_input)).view(batch_size, sequence_length, self.num_heads, self.attn_head_dim) + v = block.self_attn.v(attn_input).view(batch_size, sequence_length, self.num_heads, self.attn_head_dim) q = apply_rope(q, freqs) k = apply_rope(k, freqs) @@ -204,6 +191,7 @@ def _apply_post_with_optional_checkpoint( Returns: Updated expert tokens after self-attn residual, optional cross-attn, and MLP. """ + def _post_fn( _mixed_slice: torch.Tensor, _x: torch.Tensor, @@ -273,18 +261,11 @@ def prefill_video_cache( if "video" not in self.mixtures: raise ValueError("MoT requires `video` expert for `prefill_video_cache`.") if video_attention_mask.ndim != 2: - raise ValueError( - f"`video_attention_mask` must be 2D [S,S], got shape {tuple(video_attention_mask.shape)}" - ) + raise ValueError(f"`video_attention_mask` must be 2D [S,S], got shape {tuple(video_attention_mask.shape)}") if video_attention_mask.shape[0] != video_attention_mask.shape[1]: - raise ValueError( - f"`video_attention_mask` must be square, got shape {tuple(video_attention_mask.shape)}" - ) + raise ValueError(f"`video_attention_mask` must be square, got shape {tuple(video_attention_mask.shape)}") if video_attention_mask.shape[0] != video_tokens.shape[1]: - raise ValueError( - "`video_attention_mask` seq length mismatch: " - f"mask={video_attention_mask.shape[0]} vs tokens={video_tokens.shape[1]}" - ) + raise ValueError(f"`video_attention_mask` seq length mismatch: mask={video_attention_mask.shape[0]} vs tokens={video_tokens.shape[1]}") expert = self.mixtures["video"] x = video_tokens @@ -360,9 +341,7 @@ def forward_action_with_video_cache( if "action" not in self.mixtures: raise ValueError("MoT requires `action` expert for `forward_action_with_video_cache`.") if len(video_kv_cache) != self.num_layers: - raise ValueError( - f"`video_kv_cache` must contain {self.num_layers} layers, got {len(video_kv_cache)}." - ) + raise ValueError(f"`video_kv_cache` must contain {self.num_layers} layers, got {len(video_kv_cache)}.") if attention_mask.ndim != 2: raise ValueError(f"`attention_mask` must be 2D [S,S], got shape {tuple(attention_mask.shape)}") if attention_mask.shape[0] != attention_mask.shape[1]: @@ -371,10 +350,7 @@ def forward_action_with_video_cache( action_seq_len = int(action_tokens.shape[1]) total_seq_len = int(video_seq_len) + action_seq_len if attention_mask.shape[0] != total_seq_len: - raise ValueError( - "`attention_mask` seq length mismatch: " - f"mask={attention_mask.shape[0]} vs expected_total={total_seq_len}" - ) + raise ValueError(f"`attention_mask` seq length mismatch: mask={attention_mask.shape[0]} vs expected_total={total_seq_len}") # Use the action query rows from the joint [video+action] mask. action_attention_mask = attention_mask[video_seq_len:total_seq_len, :total_seq_len] @@ -402,16 +378,12 @@ def forward_action_with_video_cache( ) layer_cache = video_kv_cache[layer_idx] if "k" not in layer_cache or "v" not in layer_cache: - raise ValueError( - f"`video_kv_cache[{layer_idx}]` must contain `k` and `v`." - ) + raise ValueError(f"`video_kv_cache[{layer_idx}]` must contain `k` and `v`.") k_video = layer_cache["k"] v_video = layer_cache["v"] if k_video.shape[1] != video_seq_len or v_video.shape[1] != video_seq_len: - raise ValueError( - f"`video_kv_cache[{layer_idx}]` seq len mismatch, expected {video_seq_len}." - ) + raise ValueError(f"`video_kv_cache[{layer_idx}]` seq len mismatch, expected {video_seq_len}.") # Mixed attention: action queries attend to cached video K/V plus current action K/V. k_cat = torch.cat([k_video, k_action], dim=1) @@ -513,10 +485,7 @@ def forward( total_seq = q_cat.shape[1] if attention_mask.shape[0] != total_seq: - raise ValueError( - "Attention mask seq length mismatch: " - f"mask={attention_mask.shape[0]} vs tokens={total_seq}" - ) + raise ValueError(f"Attention mask seq length mismatch: mask={attention_mask.shape[0]} vs tokens={total_seq}") mixed = self._mixed_attention(q_cat=q_cat, k_cat=k_cat, v_cat=v_cat, attention_mask=attention_mask) diff --git a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py index aaf91d024..97f3aadd7 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py +++ b/lightx2v_train/lightx2v_train/model_zoo/native/wan/fastwam/video_dit.py @@ -46,10 +46,7 @@ def __init__( if action_conditioned: raise ValueError("FastWAM conditions actions through the MoT action expert, not the video expert.") if hidden_dim != num_heads * attn_head_dim: - raise ValueError( - f"hidden_dim must equal num_heads * attn_head_dim, got {hidden_dim} and " - f"{num_heads} * {attn_head_dim}" - ) + raise ValueError(f"hidden_dim must equal num_heads * attn_head_dim, got {hidden_dim} and {num_heads} * {attn_head_dim}") super().__init__( model_type="ti2v", @@ -98,9 +95,7 @@ def _validate_inputs(self, x, timestep, context, context_mask): if context_mask is None: context_mask = torch.ones(context.shape[:2], dtype=torch.bool, device=context.device) elif context_mask.shape != context.shape[:2]: - raise ValueError( - f"context_mask must have shape {tuple(context.shape[:2])}, got {tuple(context_mask.shape)}" - ) + raise ValueError(f"context_mask must have shape {tuple(context.shape[:2])}, got {tuple(context_mask.shape)}") return x, timestep, context_mask def build_video_to_video_mask(self, video_seq_len, video_tokens_per_frame, device): @@ -111,9 +106,7 @@ def build_video_to_video_mask(self, video_seq_len, video_tokens_per_frame, devic if self.video_attention_mask_mode == "per_frame_causal": frame_count = video_seq_len // video_tokens_per_frame mask = torch.tril(torch.ones((frame_count, frame_count), dtype=torch.bool, device=device)) - return mask.repeat_interleave(video_tokens_per_frame, 0).repeat_interleave( - video_tokens_per_frame, 1 - ) + return mask.repeat_interleave(video_tokens_per_frame, 0).repeat_interleave(video_tokens_per_frame, 1) if self.video_attention_mask_mode == "first_frame_causal": mask = torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) mask[:video_tokens_per_frame, video_tokens_per_frame:] = False @@ -155,14 +148,10 @@ def pre_dit( if x.shape[3] % patch_h or x.shape[4] % patch_w: raise ValueError("Latent spatial dimensions must be divisible by the Wan patch size.") tokens_per_frame = (x.shape[3] // patch_h) * (x.shape[4] // patch_w) - token_timesteps = timestep.view(batch_size, 1, 1).expand( - batch_size, x.shape[2], tokens_per_frame - ).clone() + token_timesteps = timestep.view(batch_size, 1, 1).expand(batch_size, x.shape[2], tokens_per_frame).clone() token_timesteps[:, 0] = 0 token_timesteps = token_timesteps.reshape(batch_size, -1) - time_embedding = self.time_embedding( - timestep_embedding(self.freq_dim, token_timesteps.reshape(-1)) - ).reshape(batch_size, -1, self.hidden_dim) + time_embedding = self.time_embedding(timestep_embedding(self.freq_dim, token_timesteps.reshape(-1))).reshape(batch_size, -1, self.hidden_dim) time_modulation = self.time_projection(time_embedding).unflatten(2, (6, self.hidden_dim)) patches = self.patch_embedding(x) @@ -186,12 +175,8 @@ def pre_dit( def post_dit(self, tokens: torch.Tensor, pre_state: dict[str, Any]) -> torch.Tensor: time_embedding = pre_state["t"] - shift, scale = ( - self.head.modulation.unsqueeze(0).to(time_embedding) + time_embedding.unsqueeze(2) - ).chunk(2, dim=2) - tokens = self.head.head( - self.head.norm(tokens) * (1 + scale.squeeze(2)) + shift.squeeze(2) - ) + shift, scale = (self.head.modulation.unsqueeze(0).to(time_embedding) + time_embedding.unsqueeze(2)).chunk(2, dim=2) + tokens = self.head.head(self.head.norm(tokens) * (1 + scale.squeeze(2)) + shift.squeeze(2)) frames, height, width = pre_state["meta"]["grid_size"] return rearrange( tokens, @@ -215,10 +200,9 @@ def forward(self, x, timestep, context, context_mask=None, action=None, fuse_vae fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, ) tokens = state["tokens"] - self_mask = self.build_video_to_video_mask( - tokens.shape[1], state["meta"]["tokens_per_frame"], tokens.device - ) + self_mask = self.build_video_to_video_mask(tokens.shape[1], state["meta"]["tokens_per_frame"], tokens.device) for block in self.blocks: + def run(value, current_block=block): return expert_block_forward( current_block, diff --git a/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py b/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py index ec978d6b2..6b667bef7 100644 --- a/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py +++ b/lightx2v_train/lightx2v_train/model_zoo/wan_fastwam.py @@ -35,11 +35,7 @@ def load_components(self, transformer_only=False, reference_model=None): model_path = _resolve_local_path(self.model_config.get("model_path"), "model_path", directory=True) configured_action_dit_path = self.model_config.get("action_dit_pretrained_path") - action_dit_pretrained_path = ( - _resolve_local_path(configured_action_dit_path, "action_dit_pretrained_path") - if configured_action_dit_path - else None - ) + action_dit_pretrained_path = _resolve_local_path(configured_action_dit_path, "action_dit_pretrained_path") if configured_action_dit_path else None load_text_encoder = bool(self.model_config.get("load_text_encoder", False)) video_scheduler = self.model_config.get("video_scheduler", {}) diff --git a/lightx2v_train/lightx2v_train/schedulers/flow_matching.py b/lightx2v_train/lightx2v_train/schedulers/flow_matching.py index 0d9d4e4c2..d50b5bb2f 100644 --- a/lightx2v_train/lightx2v_train/schedulers/flow_matching.py +++ b/lightx2v_train/lightx2v_train/schedulers/flow_matching.py @@ -216,9 +216,7 @@ def training_weight(self, timestep): return weight.reshape(()) if weight.numel() == 1 else weight def add_noise(self, original_samples, noise, timestep): - sigma = (timestep / self.num_train_timesteps).to( - device=original_samples.device, dtype=original_samples.dtype - ) + sigma = (timestep / self.num_train_timesteps).to(device=original_samples.device, dtype=original_samples.dtype) if sigma.ndim: sigma = sigma.view(-1, *([1] * (original_samples.ndim - 1))) return (1 - sigma) * original_samples + sigma * noise diff --git a/lightx2v_train/lightx2v_train/trainers/fastwam.py b/lightx2v_train/lightx2v_train/trainers/fastwam.py index 7171e5bb3..0e204ee90 100644 --- a/lightx2v_train/lightx2v_train/trainers/fastwam.py +++ b/lightx2v_train/lightx2v_train/trainers/fastwam.py @@ -49,9 +49,7 @@ def _atomic_torch_save(payload, path): def _sync_optimizer_group_options(source_groups, target_groups): if len(source_groups) != len(target_groups): - raise RuntimeError( - f"Optimizer parameter-group count mismatch: source={len(source_groups)} target={len(target_groups)}" - ) + raise RuntimeError(f"Optimizer parameter-group count mismatch: source={len(source_groups)} target={len(target_groups)}") for source, target in zip(source_groups, target_groups): for key, value in source.items(): if key != "params": @@ -208,27 +206,20 @@ def _load_resume_state(self, resume_ckpt_path): optimizer_state_format = state.get("optimizer_state_format") expected_optimizer_sharding = "zero1" if self.zero1_enabled else "none" if state.get("optimizer_sharding") != expected_optimizer_sharding: - raise RuntimeError( - "Optimizer sharding mode mismatch: " - f"checkpoint={state.get('optimizer_sharding')!r} active={expected_optimizer_sharding!r}." - ) + raise RuntimeError(f"Optimizer sharding mode mismatch: checkpoint={state.get('optimizer_sharding')!r} active={expected_optimizer_sharding!r}.") if optimizer_state_format == ZERO1_PER_RANK_FORMAT: if not self.zero1_enabled: raise RuntimeError("A per-rank ZeRO-1 optimizer checkpoint requires distributed.zero1.enabled=true.") shard_files = state.get("optimizer_shard_files") if not isinstance(shard_files, list) or len(shard_files) != get_world_size(): - raise RuntimeError( - f"Invalid ZeRO-1 optimizer shard manifest: expected {get_world_size()} files, got {shard_files!r}." - ) + raise RuntimeError(f"Invalid ZeRO-1 optimizer shard manifest: expected {get_world_size()} files, got {shard_files!r}.") shard_path = os.path.join(resume_ckpt_path, shard_files[get_rank()]) if not os.path.isfile(shard_path): raise RuntimeError(f"ZeRO-1 optimizer shard not found for rank {get_rank()}: {shard_path}") shard = torch.load(shard_path, map_location="cpu", weights_only=False) if shard.get("rank") != get_rank() or shard.get("world_size") != get_world_size(): raise RuntimeError( - "ZeRO-1 optimizer shard metadata mismatch: " - f"rank={shard.get('rank')} world_size={shard.get('world_size')} " - f"expected_rank={get_rank()} expected_world_size={get_world_size()}." + f"ZeRO-1 optimizer shard metadata mismatch: rank={shard.get('rank')} world_size={shard.get('world_size')} expected_rank={get_rank()} expected_world_size={get_world_size()}." ) self.optimizer.optim.load_state_dict(shard["optimizer"]) _sync_optimizer_group_options(self.optimizer.optim.param_groups, self.optimizer.param_groups) @@ -237,9 +228,7 @@ def _load_resume_state(self, resume_ckpt_path): raise RuntimeError("A full optimizer checkpoint requires distributed.zero1.enabled=false.") self.optimizer.load_state_dict(state["optimizer"]) else: - raise RuntimeError( - f"Unsupported optimizer checkpoint format {optimizer_state_format!r} in {state_path}." - ) + raise RuntimeError(f"Unsupported optimizer checkpoint format {optimizer_state_format!r} in {state_path}.") self.lr_scheduler.load_state_dict(state["lr_scheduler"]) logger.info( "[resume] restored FastWAM training state from {} (optimizer sharding: {}, state format: {})", @@ -262,10 +251,7 @@ def _save_checkpoint(self, iteration): optimizer_shard_files = None if self.zero1_enabled: optimizer_state_format = ZERO1_PER_RANK_FORMAT - optimizer_shard_files = [ - _optimizer_shard_filename(rank, get_world_size()) - for rank in range(get_world_size()) - ] + optimizer_shard_files = [_optimizer_shard_filename(rank, get_world_size()) for rank in range(get_world_size())] _sync_optimizer_group_options(self.optimizer.param_groups, self.optimizer.optim.param_groups) shard_path = os.path.join(save_dir, optimizer_shard_files[get_rank()]) shard_start_time = time.perf_counter() @@ -514,11 +500,7 @@ def _evaluate_sample(self, model, dataset, sample, sample_number, eval_index, st context_mask = sample.get("context_mask") prompt = sample.get("prompt", [None])[0] input_image = video0[:, 0].unsqueeze(0) - action_horizon = ( - int(action.shape[1]) - if action is not None - else int((video0.shape[1] - 1) * dataset.action_video_freq_ratio) - ) + action_horizon = int(action.shape[1]) if action is not None else int((video0.shape[1] - 1) * dataset.action_video_freq_ratio) infer_kwargs = { "input_image": input_image, "num_frames": int(video0.shape[1]), @@ -613,9 +595,7 @@ def _action_metrics(self, dataset, sample, pred_action, target_action): @staticmethod def _validate_video_shapes(left, right, left_name, right_name): if left.shape != right.shape: - raise ValueError( - f"Eval {left_name}/{right_name} shape mismatch: {tuple(left.shape)} vs {tuple(right.shape)}" - ) + raise ValueError(f"Eval {left_name}/{right_name} shape mismatch: {tuple(left.shape)} vs {tuple(right.shape)}") if left.shape[1] <= 1: raise ValueError(f"Eval video must contain at least one future frame, got {tuple(left.shape)}") @@ -631,15 +611,7 @@ def _comparison_frames(pred, vae, gt): canvas = Image.new("RGB", (frame_width * 3, frame_height + label_height), color=(18, 18, 18)) draw = ImageDraw.Draw(canvas) for column, (label, video) in enumerate(zip(labels, videos)): - array = ( - video[:, frame_index] - .permute(1, 2, 0) - .clamp(0.0, 1.0) - .mul(255.0) - .round() - .to(torch.uint8) - .numpy() - ) + array = video[:, frame_index].permute(1, 2, 0).clamp(0.0, 1.0).mul(255.0).round().to(torch.uint8).numpy() canvas.paste(Image.fromarray(np.ascontiguousarray(array)), (column * frame_width, label_height)) draw.text((column * frame_width + 8, 9), label, fill=(245, 245, 245)) frames.append(canvas) @@ -660,12 +632,5 @@ def _gather_eval_results(local_results): def _summarize_eval_results(results): if not results: return {} - metric_keys = sorted( - key - for key, value in results[0].items() - if isinstance(value, float) - ) - return { - key: float(sum(result[key] for result in results if key in result) / sum(key in result for result in results)) - for key in metric_keys - } + metric_keys = sorted(key for key, value in results[0].items() if isinstance(value, float)) + return {key: float(sum(result[key] for result in results if key in result) / sum(key in result for result in results)) for key in metric_keys} From fcd64944ef614ee9a9c6748cf7a0284931046b65 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Mon, 13 Jul 2026 13:34:12 +0000 Subject: [PATCH 09/10] u --- .../lightx2v_train/data/libero/dataset.py | 13 +------------ .../lightx2v_train/data/libero/preparation.py | 4 ++-- 2 files changed, 3 insertions(+), 14 deletions(-) diff --git a/lightx2v_train/lightx2v_train/data/libero/dataset.py b/lightx2v_train/lightx2v_train/data/libero/dataset.py index 6abc84afa..8e9da4b3d 100644 --- a/lightx2v_train/lightx2v_train/data/libero/dataset.py +++ b/lightx2v_train/lightx2v_train/data/libero/dataset.py @@ -11,17 +11,6 @@ from .processor import FastWAMProcessor from .robot_video_dataset import RobotVideoDataset -ASSET_ROOT = Path(__file__).resolve().parents[3] / "assets" / "libero_fastwam" -DEFAULT_DATASET_DIRS = [ - ASSET_ROOT / "datasets" / name - for name in ( - "libero_spatial_no_noops_lerobot", - "libero_object_no_noops_lerobot", - "libero_goal_no_noops_lerobot", - "libero_10_no_noops_lerobot", - ) -] - def _default_shape_meta(): return { @@ -57,7 +46,7 @@ def _build_dataset(config, split): shape_meta = deepcopy(config.get("shape_meta") or _default_shape_meta()) num_frames = int(config.get("num_frames", 33)) processor = FastWAMProcessor(shape_meta, num_frames) - dataset_dirs = config.get("dataset_dirs") or DEFAULT_DATASET_DIRS + dataset_dirs = config.get("dataset_dirs") if isinstance(dataset_dirs, (str, Path)): dataset_dirs = [dataset_dirs] diff --git a/lightx2v_train/lightx2v_train/data/libero/preparation.py b/lightx2v_train/lightx2v_train/data/libero/preparation.py index 8f2a06629..a662709a1 100644 --- a/lightx2v_train/lightx2v_train/data/libero/preparation.py +++ b/lightx2v_train/lightx2v_train/data/libero/preparation.py @@ -12,7 +12,7 @@ from lightx2v_train.runtime.distributed import barrier, is_main_process -from .dataset import DEFAULT_DATASET_DIRS, _default_shape_meta +from .dataset import _default_shape_meta from .lerobot_dataset import LiberoLeRobotDataset from .robot_video_dataset import PROMPT_TEMPLATE @@ -25,7 +25,7 @@ def _resolve_path(value): def _dataset_dirs(split_config): - values = split_config.get("dataset_dirs") or DEFAULT_DATASET_DIRS + values = split_config.get("dataset_dirs") if isinstance(values, (str, Path)): values = [values] return [_resolve_path(value) for value in values] From daa58cf5d74ca68d97ab98671b936590b96136fe Mon Sep 17 00:00:00 2001 From: Musisoul Date: Mon, 13 Jul 2026 13:35:56 +0000 Subject: [PATCH 10/10] u --- lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml index b9a32b45a..9fa39df88 100644 --- a/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml +++ b/lightx2v_train/configs/train/fastwam/libero_uncond_2cam224.yaml @@ -1,7 +1,7 @@ model: name: wan_fastwam running_dtype: bf16 - model_path: /data/nvme0/models/Wan-AI/Wan2.2-TI2V-5B + model_path: /path/to/Wan-AI/Wan2.2-TI2V-5B data: train: &libero_data