Spaces:
Running on Zero
Running on Zero
| from .data.dataset import ( | |
| PrecomputedLatentDataset, | |
| _default_blank_prompt_emb_path, | |
| _prompt_emb_to_cpu, | |
| _save_prompt_emb_atomic, | |
| ) | |
| import copy | |
| import glob | |
| import os | |
| import re | |
| import time | |
| import torch, os | |
| import lightning as pl | |
| from .pipelines.wan_video_Nplayers import WanVideoNPlayersPipeline | |
| from .models.model_manager import ModelManager | |
| from multics.pipelines.wan_video_Nplayers import model_fn_wan_video_causal | |
| import random | |
| import json | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import glob | |
| from lightning.pytorch.loggers import TensorBoardLogger | |
| from multics.models.wan_video_dit import ( | |
| prepare_hier_local_hub_state_dict_for_load, | |
| migrate_hier_no_local_hub_state_dict, | |
| ) | |
| _REPO_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| _BLANK_PROMPT_EMB_FILENAME = "blank_string_prompt_emb.pth" | |
| _CAMERA_PLAYER_PART = re.compile("(?:^|/)player=(\\d+)(?:/|$)") | |
| _GENERATED_HISTORY_NEGATIVE_PROMPT = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" | |
| def get_corrupt_noise_sigma( | |
| model_input, | |
| batch_size, | |
| corrupt_ratio=1 / 3, | |
| num_frames=None, | |
| is_frame_independent=False, | |
| ): | |
| """DF-style σ ~ U(0, corrupt_ratio) for (1-σ)*x + σ*noise. Broadcast to model_input.ndim.""" | |
| if is_frame_independent: | |
| noise_sigma_shape = (batch_size, 1, num_frames) | |
| else: | |
| noise_sigma_shape = (batch_size,) | |
| noise_sigma = ( | |
| torch.rand( | |
| size=noise_sigma_shape, device=model_input.device, dtype=model_input.dtype | |
| ) | |
| * corrupt_ratio | |
| ) | |
| while len(noise_sigma.shape) < model_input.ndim: | |
| noise_sigma = noise_sigma.unsqueeze(-1) | |
| return noise_sigma | |
| def _attach_hier_group_dit_kwargs(dit_kwargs: dict, batch: dict, device) -> None: | |
| """Copy grouping tensors from the dataloader batch into DiT kwargs. | |
| labels mode: group_labels [1,T,P], active_group [1,T,N] | |
| membership mode also: membership [1,T,N,P], global_only [1,T,P] | |
| """ | |
| if "group_labels" not in batch or "active_group" not in batch: | |
| raise KeyError( | |
| "hier_hub requires dataset batch keys group_labels and active_group" | |
| ) | |
| dit_kwargs["group_labels"] = batch["group_labels"].to( | |
| device=device, dtype=torch.long | |
| ) | |
| dit_kwargs["active_group"] = batch["active_group"].to( | |
| device=device, dtype=torch.bool | |
| ) | |
| if "membership" in batch: | |
| if "global_only" not in batch: | |
| raise KeyError("hier_group_format=membership batch requires global_only") | |
| dit_kwargs["membership"] = batch["membership"].to( | |
| device=device, dtype=torch.bool | |
| ) | |
| dit_kwargs["global_only"] = batch["global_only"].to( | |
| device=device, dtype=torch.bool | |
| ) | |
| class LightningModelForTrain(pl.LightningModule): | |
| def __init__( | |
| self, | |
| text_encoder_path, | |
| vae_path, | |
| image_encoder_path, | |
| dit_path, | |
| learning_rate=1e-05, | |
| use_gradient_checkpointing=True, | |
| use_gradient_checkpointing_offload=False, | |
| resume_ckpt_path=None, | |
| output_path="./", | |
| num_players=2, | |
| vae_latent_channels=48, | |
| perf_log_every_n_steps=100, | |
| perf_timing=True, | |
| target_rgb_frames=81, | |
| first_frame_dropout_p=0.0, | |
| first_frame_noise_p=0.0, | |
| first_frame_noise_corrupt_ratio=1 / 3, | |
| first_frame_noise_min_seg=1, | |
| ckpt_local_dir=None, | |
| run_version=None, | |
| blank_prompt_emb_path: str = None, | |
| text_prompt_dropout_p: float = 0.0, | |
| use_simplex_rope: bool = True, | |
| simplex_pool_size: int = 4, | |
| simplex_gamma: float = 1.0, | |
| simplex_train_permute: bool = True, | |
| simplex_train_sample_vertex_ids: bool = False, | |
| training_stage: str = "stage2", | |
| loss_type: str = "diffusion_forcing", | |
| attention_mode: str = "hier_hub", | |
| attn_impl: str = "sdpa", | |
| num_hub_tokens: int = 8, | |
| local_attn_size: int = 24, | |
| local_attn_size_after: int = None, | |
| local_attn_switch_step: int = None, | |
| train_step_offset: int = 0, | |
| num_frame_per_block: int = 7, | |
| num_groups: int = 1, | |
| num_local_hub_tokens: int = 0, | |
| num_global_hub_tokens: int = None, | |
| hier_no_local_hub: bool = False, | |
| hier_local_hub_mode: str = "shared", | |
| hier_direct_schedule: str = "none", | |
| hier_direct_frac: float = 0.0, | |
| hier_direct_layers=None, | |
| unfreeze_self_attn: bool = False, | |
| clean_prefix_prob: float = 0.0, | |
| clean_prefix_context_noise: int = 0, | |
| generated_history_prob: float = 0.0, | |
| df_per_frame_noise: bool = False, | |
| causal_forcing_training: bool = False, | |
| use_camera_pose: bool = False, | |
| camera_dim: int = 12, | |
| camera_dropout_p: float = 0.1, | |
| val_infer: bool = False, | |
| val_every_n_steps: int = 500, | |
| val_num_samples: int = 4, | |
| val_num_inference_steps: int = 50, | |
| val_causal_num_frames: int = None, | |
| val_output_dir: str = "outputs/validation", | |
| val_grid_fps: int = 15, | |
| ): | |
| super().__init__() | |
| if ( | |
| training_stage != "stage2" | |
| or loss_type != "diffusion_forcing" | |
| or attention_mode != "hier_hub" | |
| or (not use_simplex_rope) | |
| ): | |
| raise ValueError( | |
| "Multics supports hierarchical attention with diffusion forcing and simplex RoPE" | |
| ) | |
| model_manager = ModelManager(torch_dtype=torch.bfloat16, device="cpu") | |
| if "*" in dit_path: | |
| dit_files = sorted(glob.glob(dit_path)) | |
| model_manager.load_model(dit_files) | |
| elif os.path.isfile(dit_path): | |
| model_manager.load_models([dit_path]) | |
| else: | |
| dit_path = dit_path.split(",") | |
| model_manager.load_models([dit_path]) | |
| model_path = [text_encoder_path, vae_path] | |
| if image_encoder_path is not None: | |
| model_path.append(image_encoder_path) | |
| model_manager.load_models(model_path) | |
| self.pipe = WanVideoNPlayersPipeline.from_model_manager(model_manager) | |
| self.pipe.scheduler.set_timesteps(1000, training=True) | |
| self.pipe.num_players = num_players | |
| self.use_simplex_rope = bool(True) | |
| self.simplex_pool_size = int(simplex_pool_size) | |
| self.simplex_gamma = float(simplex_gamma) | |
| self.simplex_train_permute = bool(simplex_train_permute) | |
| self.simplex_train_sample_vertex_ids = bool(simplex_train_sample_vertex_ids) | |
| self.training_stage = str("stage2") | |
| self.loss_type = str("diffusion_forcing") | |
| self.attention_mode = str("hier_hub") | |
| self.attn_impl = str(attn_impl) | |
| self.num_hub_tokens = int(num_hub_tokens) | |
| self.local_attn_size_initial = int(local_attn_size) | |
| self.local_attn_size_after = ( | |
| None if local_attn_size_after is None else int(local_attn_size_after) | |
| ) | |
| self.local_attn_switch_step = ( | |
| None if local_attn_switch_step is None else int(local_attn_switch_step) | |
| ) | |
| self.train_step_offset = int(train_step_offset) | |
| if (self.local_attn_size_after is None) != ( | |
| self.local_attn_switch_step is None | |
| ): | |
| raise ValueError( | |
| "local_attn_size_after and local_attn_switch_step must be set together" | |
| ) | |
| if self.local_attn_switch_step is not None and self.local_attn_switch_step < 0: | |
| raise ValueError("local_attn_switch_step must be >= 0") | |
| if self.train_step_offset < 0: | |
| raise ValueError("train_step_offset must be >= 0") | |
| self._last_logged_local_attn_size = None | |
| self.num_frame_per_block = int(num_frame_per_block) | |
| for name, value in ( | |
| ("local_attn_size", self.local_attn_size_initial), | |
| ("local_attn_size_after", self.local_attn_size_after), | |
| ): | |
| if value is not None and value != -1 and (value < self.num_frame_per_block): | |
| raise ValueError( | |
| f"{name} must be -1 or >= num_frame_per_block; got {value} and {self.num_frame_per_block}" | |
| ) | |
| self.local_attn_size = self._local_attn_size_for_effective_step( | |
| self.train_step_offset | |
| ) | |
| self.num_groups = int(num_groups) | |
| self.num_local_hub_tokens = int(num_local_hub_tokens) | |
| self.num_global_hub_tokens = ( | |
| None if num_global_hub_tokens is None else int(num_global_hub_tokens) | |
| ) | |
| self.hier_no_local_hub = bool(hier_no_local_hub) | |
| self.hier_local_hub_mode = str(hier_local_hub_mode) | |
| self.hier_direct_schedule = str(hier_direct_schedule) | |
| self.hier_direct_frac = float(hier_direct_frac) | |
| self.hier_direct_layers = hier_direct_layers | |
| self.unfreeze_self_attn = bool(unfreeze_self_attn) | |
| self.target_rgb_frames = int(target_rgb_frames) | |
| self.val_infer = bool(val_infer) | |
| self.val_every_n_steps = int(val_every_n_steps) | |
| self.val_num_samples = int(val_num_samples) | |
| self.val_num_inference_steps = int(val_num_inference_steps) | |
| self.val_causal_num_frames = int(val_causal_num_frames or target_rgb_frames) | |
| self.val_output_dir = val_output_dir | |
| self.val_grid_fps = int(val_grid_fps) | |
| self.val_agent_vertex_ids = list(range(int(num_players))) | |
| self._manual_val_loader = None | |
| self._last_val_step = -1 | |
| self.clean_prefix_prob = float(clean_prefix_prob) | |
| self.clean_prefix_context_noise = int(clean_prefix_context_noise) | |
| self.generated_history_prob = float(generated_history_prob) | |
| if not 0.0 <= self.generated_history_prob <= 1.0: | |
| raise ValueError( | |
| f"generated_history_prob must be in [0,1], got {self.generated_history_prob}" | |
| ) | |
| if self.generated_history_prob > 0.0: | |
| if causal_forcing_training: | |
| raise ValueError( | |
| "Generated history uses the parallel target path and is incompatible with --causal_forcing_training" | |
| ) | |
| print( | |
| f"[generated-history] enabled prob={self.generated_history_prob} blocks=1..3 steps={self.val_num_inference_steps} cfg=5", | |
| flush=True, | |
| ) | |
| self._generated_history_negative_context_cpu = None | |
| self.df_per_frame_noise = bool(df_per_frame_noise) | |
| self.causal_forcing_training = bool(causal_forcing_training) | |
| self.use_camera_pose = bool(use_camera_pose) | |
| self.camera_dim = int(camera_dim) | |
| self.camera_dropout_p = float(camera_dropout_p) | |
| self.pipe.dit.keyboard_embedding = nn.Linear(23, 3072) | |
| nn.init.normal_(self.pipe.dit.keyboard_embedding.weight, mean=0, std=0.05) | |
| nn.init.zeros_(self.pipe.dit.keyboard_embedding.bias) | |
| dim = self.pipe.dit.blocks[0].self_attn.q.weight.shape[0] | |
| for block in self.pipe.dit.blocks: | |
| block.action_mlp = nn.Sequential( | |
| nn.LayerNorm(dim + 24, eps=1e-05, elementwise_affine=True), | |
| nn.Linear(dim + 24, dim + 24), | |
| nn.GELU(), | |
| nn.Linear(dim + 24, dim), | |
| nn.LayerNorm(dim, eps=1e-05, elementwise_affine=True), | |
| ) | |
| block.action_temporal_attn = copy.deepcopy(block.self_attn) | |
| block.action_cross_attn = copy.deepcopy(block.cross_attn) | |
| nn.init.zeros_(block.action_temporal_attn.o.weight) | |
| nn.init.zeros_(block.action_temporal_attn.o.bias) | |
| nn.init.zeros_(block.action_cross_attn.o.weight) | |
| nn.init.zeros_(block.action_cross_attn.o.bias) | |
| self.pipe.dit.enable_simplex_rope( | |
| pool_size=self.simplex_pool_size, gamma=self.simplex_gamma | |
| ) | |
| latent_t = (self.target_rgb_frames - 1) // 4 + 1 | |
| if latent_t % self.num_frame_per_block != 0: | |
| raise ValueError( | |
| f"Requires latent_T % num_frame_per_block == 0; got latent_T={latent_t}, n={self.num_frame_per_block}" | |
| ) | |
| stage2_attention_mode = "hier_hub" | |
| if self.simplex_train_permute: | |
| raise ValueError( | |
| "hier_hub requires --no_simplex_train_permute so group_labels stay aligned" | |
| ) | |
| self.pipe.enable_stage2( | |
| num_hub_tokens=self.num_hub_tokens, | |
| local_attn_size=self.local_attn_size, | |
| num_frame_per_block=self.num_frame_per_block, | |
| attention_mode="hier_hub", | |
| num_groups=self.num_groups, | |
| num_local_hub_tokens=self.num_local_hub_tokens, | |
| num_global_hub_tokens=self.num_global_hub_tokens, | |
| hier_direct_schedule=self.hier_direct_schedule, | |
| hier_direct_frac=self.hier_direct_frac, | |
| hier_direct_layers=self.hier_direct_layers, | |
| attn_impl=self.attn_impl, | |
| hier_no_local_hub=self.hier_no_local_hub, | |
| hier_local_hub_mode=self.hier_local_hub_mode, | |
| ) | |
| self.loss_type = "diffusion_forcing" | |
| self.attention_mode = "hier_hub" | |
| self.unfreeze_self_attn = True | |
| if self.use_camera_pose: | |
| self.pipe.dit.enable_camera_pose_alignment( | |
| cam_dim=self.camera_dim, dropout_p=self.camera_dropout_p | |
| ) | |
| print( | |
| f"[camera-pose] enabled bias-free Q/K-only camera conditioning (dim={self.camera_dim}, dropout_p={self.camera_dropout_p})" | |
| ) | |
| assert 0.0 <= float(first_frame_dropout_p) <= 1.0, first_frame_dropout_p | |
| self.first_frame_dropout_p = float(first_frame_dropout_p) | |
| assert 0.0 <= float(first_frame_noise_p) <= 1.0, first_frame_noise_p | |
| assert 0.0 <= float(first_frame_noise_corrupt_ratio) <= 1.0, ( | |
| first_frame_noise_corrupt_ratio | |
| ) | |
| self.first_frame_noise_p = float(first_frame_noise_p) | |
| self.first_frame_noise_corrupt_ratio = float(first_frame_noise_corrupt_ratio) | |
| self.first_frame_noise_min_seg = int(first_frame_noise_min_seg) | |
| if resume_ckpt_path is not None: | |
| state_dict = torch.load(resume_ckpt_path, map_location="cpu") | |
| if not self.hier_no_local_hub: | |
| state_dict, migration = prepare_hier_local_hub_state_dict_for_load( | |
| state_dict, | |
| target_mode=self.hier_local_hub_mode, | |
| target_num_groups=self.pipe.dit.num_groups, | |
| target_num_local_hub_tokens=self.pipe.dit.num_local_hub_tokens, | |
| model_dim=self.pipe.dit.dim, | |
| ) | |
| if migration is not None: | |
| print( | |
| f"[hier-local-hub] training warm-start migration: {migration['source_mode']} {migration['source_shape']} -> {migration['target_mode']} {migration['target_shape']}; optimizer state is not migrated", | |
| flush=True, | |
| ) | |
| old_hub = state_dict.get("hub_embed") | |
| new_hub = self.pipe.dit.hub_embed | |
| if ( | |
| torch.is_tensor(old_hub) | |
| and torch.is_tensor(new_hub) | |
| and (old_hub.ndim == 2) | |
| and (tuple(old_hub.shape) != tuple(new_hub.shape)) | |
| ): | |
| resized = new_hub.detach().cpu().clone() | |
| n = min(old_hub.shape[0], resized.shape[0]) | |
| resized[:n].copy_(old_hub[:n].to(dtype=resized.dtype)) | |
| state_dict["hub_embed"] = resized | |
| print( | |
| f"[checkpoint] resized hub_embed {tuple(old_hub.shape)} -> {tuple(resized.shape)}; copied {n} slots", | |
| flush=True, | |
| ) | |
| if self.hier_no_local_hub: | |
| state_dict, migration = migrate_hier_no_local_hub_state_dict( | |
| state_dict, | |
| target_num_global_hub_tokens=self.pipe.dit.num_global_hub_tokens, | |
| model_dim=self.pipe.dit.dim, | |
| ) | |
| if migration is not None: | |
| print( | |
| f"[hier-no-local] training warm-start migration: {migration['source']} {migration['source_shape']} -> global_hub_embed {migration['target_shape']} (replaced old global {migration['replaced_global_shape']}); optimizer state is not migrated", | |
| flush=True, | |
| ) | |
| missing, unexpected = self.pipe.dit.load_state_dict( | |
| state_dict, strict=False | |
| ) | |
| if missing: | |
| _hub_related = [k for k in missing if k.endswith("hub_embed")] | |
| _cam_keys = (".cam_q.", ".cam_k.") | |
| _cam_related = [k for k in missing if any((s in k for s in _cam_keys))] | |
| _other_missing = [ | |
| k | |
| for k in missing | |
| if k not in _hub_related and k not in _cam_related | |
| ] | |
| if _other_missing: | |
| raise RuntimeError( | |
| f"Unexpected missing keys when loading DiT: {_other_missing[:10]} ..." | |
| ) | |
| if _hub_related: | |
| print( | |
| f"[checkpoint] {len(_hub_related)} hub_embed param(s) absent in ckpt (zeros init, will train from seed)." | |
| ) | |
| if _cam_related: | |
| print( | |
| f"[camera-pose] {len(_cam_related)} camera adapter params absent in ckpt (zeros init)." | |
| ) | |
| if unexpected: | |
| _ff_keys = (".k_ff", ".v_ff", ".norm_k_ff", ".ff_gate") | |
| _ff_related = [k for k in unexpected if any((s in k for s in _ff_keys))] | |
| _handled_unexpected = list(_ff_related) | |
| _old_hub = [ | |
| k | |
| for k in unexpected | |
| if k == "hub_embed" or k.endswith(".hub_embed") | |
| ] | |
| if _old_hub: | |
| print( | |
| f"[checkpoint] ignoring {len(_old_hub)} legacy per-block hub_embed key(s) in ckpt." | |
| ) | |
| _handled_unexpected += _old_hub | |
| _legacy_cam_bias = [ | |
| k | |
| for k in unexpected | |
| if k.endswith(".cam_q.bias") or k.endswith(".cam_k.bias") | |
| ] | |
| if _legacy_cam_bias: | |
| print( | |
| f"[camera-pose] dropping {len(_legacy_cam_bias)} legacy camera projection bias key(s) for the bias-free adapter." | |
| ) | |
| _disabled_cam = [ | |
| k | |
| for k in unexpected | |
| if not self.use_camera_pose | |
| and (".cam_q." in k or ".cam_k." in k) | |
| and (k not in _legacy_cam_bias) | |
| ] | |
| if _disabled_cam: | |
| print( | |
| f"[camera-pose] ignoring {len(_disabled_cam)} camera adapter key(s) because camera conditioning is disabled." | |
| ) | |
| _handled_unexpected += _legacy_cam_bias + _disabled_cam | |
| if _ff_related: | |
| print( | |
| f"[WARN] DiT ckpt 中检测到 {len(_ff_related)} 个 first-frame cross-attn 分支参数 （k_ff/v_ff/norm_k_ff/ff_gate），当前架构已回退到 latent-overwrite 方案，这些参数会被忽略。如不需要可重新存一份干净的 ckpt。示例: {_ff_related[:3]}" | |
| ) | |
| _other_unexpected = [ | |
| k for k in unexpected if k not in _handled_unexpected | |
| ] | |
| if _other_unexpected: | |
| print( | |
| f"[WARN] unexpected keys in DiT ckpt: {_other_unexpected[:10]} ..." | |
| ) | |
| if "hub_embed" in state_dict and ( | |
| "local_hub_embed" not in state_dict | |
| or "global_hub_embed" not in state_dict | |
| ): | |
| old_hub = state_dict["hub_embed"] | |
| if ( | |
| torch.is_tensor(old_hub) | |
| and old_hub.dim() == 2 | |
| and (old_hub.shape[1] == self.pipe.dit.dim) | |
| ): | |
| with torch.no_grad(): | |
| kl = int(self.pipe.dit.num_local_hub_tokens) | |
| kg = int(self.pipe.dit.num_global_hub_tokens) | |
| if kl > 0: | |
| n = min(kl, old_hub.shape[0]) | |
| local_target = self.pipe.dit.local_hub_embed | |
| local_values = old_hub[:n].to(local_target) | |
| if local_target.dim() == 2: | |
| local_target[:n].copy_(local_values) | |
| elif local_target.dim() == 3: | |
| local_target[:, :n].copy_( | |
| local_values.unsqueeze(0).expand( | |
| local_target.shape[0], -1, -1 | |
| ) | |
| ) | |
| else: | |
| raise RuntimeError( | |
| f"local_hub_embed must be rank 2 or 3, got {tuple(local_target.shape)}" | |
| ) | |
| if kg > 0: | |
| n = min(kg, old_hub.shape[0]) | |
| self.pipe.dit.global_hub_embed[-n:].copy_( | |
| old_hub[-n:].to(self.pipe.dit.global_hub_embed) | |
| ) | |
| print( | |
| f"[hier_hub] initialized local/global hub seeds from central hub_embed {tuple(old_hub.shape)} -> local={(None if self.pipe.dit.local_hub_embed is None else tuple(self.pipe.dit.local_hub_embed.shape))} global={tuple(self.pipe.dit.global_hub_embed.shape)}" | |
| ) | |
| self.freeze_parameters() | |
| trainable_keywords = [ | |
| "keyboard_embedding", | |
| "action_mlp", | |
| "action_temporal_attn", | |
| "action_cross_attn", | |
| ] | |
| if self.unfreeze_self_attn: | |
| trainable_keywords.extend(["self_attn", "hub_embed"]) | |
| if self.use_camera_pose: | |
| trainable_keywords.extend(["cam_q", "cam_k"]) | |
| for name, module in self.pipe.denoising_model().named_modules(): | |
| if any((keyword in name for keyword in trainable_keywords)): | |
| print(f"Trainable: {name}") | |
| for param in module.parameters(): | |
| param.requires_grad = True | |
| if name.endswith(".cross_attn"): | |
| print(f"Trainable (text cross-attention): {name}") | |
| for param in module.parameters(): | |
| param.requires_grad = True | |
| if self.unfreeze_self_attn: | |
| _dit = self.pipe.denoising_model() | |
| if hasattr(_dit, "hub_embed") and isinstance(_dit.hub_embed, nn.Parameter): | |
| _dit.hub_embed.requires_grad_(True) | |
| print( | |
| f"Trainable (parameter, explicit): model-level hub_embed {tuple(_dit.hub_embed.shape)}" | |
| ) | |
| for _name in ("local_hub_embed", "global_hub_embed"): | |
| _p = getattr(_dit, _name, None) | |
| if isinstance(_p, nn.Parameter): | |
| _p.requires_grad_(True) | |
| print(f"Trainable (parameter, explicit): {_name} {tuple(_p.shape)}") | |
| os.makedirs(output_path, exist_ok=True) | |
| self.output_path = output_path | |
| self.ckpt_local_dir = ckpt_local_dir | |
| self.run_version = run_version | |
| self._prompt_encode_cache: dict = {} | |
| self._blank_prompt_ctx: torch.Tensor | None = None | |
| self._blank_prompt_emb_path = ( | |
| blank_prompt_emb_path or _default_blank_prompt_emb_path() | |
| ) | |
| self.text_prompt_dropout_p = float(text_prompt_dropout_p) | |
| if not 0.0 <= self.text_prompt_dropout_p <= 1.0: | |
| raise ValueError( | |
| f"--text_prompt_dropout_p must be in [0,1], got {self.text_prompt_dropout_p}" | |
| ) | |
| if self.text_prompt_dropout_p > 0: | |
| print( | |
| f"[text-cfg] text_prompt_dropout_p={self.text_prompt_dropout_p}", | |
| flush=True, | |
| ) | |
| trainable_params = 0 | |
| seen_params = set() | |
| for name, module in self.pipe.denoising_model().named_modules(): | |
| for param in module.parameters(): | |
| if param.requires_grad and param not in seen_params: | |
| trainable_params += param.numel() | |
| seen_params.add(param) | |
| print(f"Total number of trainable parameters: {trainable_params}") | |
| self.learning_rate = learning_rate | |
| self.use_gradient_checkpointing = use_gradient_checkpointing | |
| self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload | |
| tiled = True | |
| tile_size = (34, 34) | |
| tile_stride = (18, 16) | |
| self.tiler_kwargs = { | |
| "tiled": tiled, | |
| "tile_size": tile_size, | |
| "tile_stride": tile_stride, | |
| } | |
| self._perf_timing_enabled = bool(perf_timing) | |
| self.perf_log_every_n_steps = max(1, int(perf_log_every_n_steps)) | |
| self._iter_start_time = None | |
| self._prev_batch_end_time = None | |
| self._current_data_wait_time = 0.0 | |
| self._backward_start_time = None | |
| self._last_backward_time = 0.0 | |
| self._last_optimizer_time = 0.0 | |
| self._latest_step_part_times = {"encode_target": 0.0, "dit_forward_loss": 0.0} | |
| self._perf_window = { | |
| "count": 0, | |
| "iter_wall": 0.0, | |
| "data_wait": 0.0, | |
| "compute": 0.0, | |
| "backward": 0.0, | |
| "optimizer": 0.0, | |
| "encode_target": 0.0, | |
| "dit_forward_loss": 0.0, | |
| } | |
| def _synced_time(self): | |
| if self._perf_timing_enabled and torch.cuda.is_available(): | |
| try: | |
| torch.cuda.synchronize(self.device) | |
| except Exception: | |
| torch.cuda.synchronize() | |
| return time.perf_counter() | |
| def _classify_bottleneck(avg_data_wait, avg_compute, avg_backward): | |
| total_iter = avg_data_wait + avg_compute | |
| if total_iter <= 0: | |
| return "unknown" | |
| data_wait_ratio = avg_data_wait / total_iter | |
| compute_ratio = avg_compute / total_iter | |
| backward_ratio = avg_backward / max(avg_compute, 1e-08) | |
| if data_wait_ratio >= 0.4 and data_wait_ratio >= compute_ratio: | |
| return "data" | |
| if compute_ratio >= 0.6 and backward_ratio >= 0.55: | |
| return "communication" | |
| if compute_ratio >= 0.6: | |
| return "compute" | |
| if backward_ratio >= 0.65: | |
| return "communication" | |
| if data_wait_ratio >= 0.35: | |
| return "data" | |
| return "compute" | |
| def on_train_batch_start(self, batch, batch_idx): | |
| self._maybe_update_local_attn_size() | |
| if not self._perf_timing_enabled: | |
| return | |
| now = self._synced_time() | |
| if self._prev_batch_end_time is None: | |
| self._current_data_wait_time = 0.0 | |
| else: | |
| self._current_data_wait_time = max(0.0, now - self._prev_batch_end_time) | |
| self._iter_start_time = now | |
| self._backward_start_time = None | |
| self._last_backward_time = 0.0 | |
| self._last_optimizer_time = 0.0 | |
| def _local_attn_size_for_effective_step(self, effective_step: int) -> int: | |
| if ( | |
| self.local_attn_switch_step is not None | |
| and int(effective_step) >= self.local_attn_switch_step | |
| ): | |
| return int(self.local_attn_size_after) | |
| return int(self.local_attn_size_initial) | |
| def _maybe_update_local_attn_size(self): | |
| """Apply the attention window selected for the next optimizer update.""" | |
| if self.local_attn_switch_step is None: | |
| return | |
| local_step = int(self.global_step) | |
| effective_step = self.train_step_offset + local_step | |
| target = self._local_attn_size_for_effective_step(effective_step) | |
| if int(self.local_attn_size) != target: | |
| self.pipe.set_stage2_local_attn_size(target, reset_training_cache=True) | |
| self.local_attn_size = target | |
| values = { | |
| int(self.local_attn_size), | |
| int(self.pipe.local_attn_size), | |
| int(self.pipe.dit.local_attn_size), | |
| *(int(block.local_attn_size) for block in self.pipe.dit.blocks), | |
| *(int(block.self_attn.local_attn_size) for block in self.pipe.dit.blocks), | |
| } | |
| if values != {target}: | |
| raise RuntimeError( | |
| f"local-attn schedule propagation failed: expected {target}, observed {sorted(values)}" | |
| ) | |
| if self._last_logged_local_attn_size != target: | |
| if int(getattr(self, "global_rank", 0)) == 0: | |
| print( | |
| f"[local-attn-schedule] local_step={local_step} effective_step={effective_step} local_attn_size={target}", | |
| flush=True, | |
| ) | |
| self._last_logged_local_attn_size = target | |
| def on_before_backward(self, loss): | |
| if not self._perf_timing_enabled: | |
| return | |
| self._backward_start_time = self._synced_time() | |
| def on_after_backward(self): | |
| if not self._perf_timing_enabled: | |
| return | |
| if self._backward_start_time is None: | |
| self._last_backward_time = 0.0 | |
| return | |
| self._last_backward_time = max( | |
| 0.0, self._synced_time() - self._backward_start_time | |
| ) | |
| def optimizer_step(self, *args, **kwargs): | |
| if not self._perf_timing_enabled: | |
| return super().optimizer_step(*args, **kwargs) | |
| start_time = self._synced_time() | |
| out = super().optimizer_step(*args, **kwargs) | |
| self._last_optimizer_time = max(0.0, self._synced_time() - start_time) | |
| return out | |
| def on_train_batch_end(self, outputs, batch, batch_idx): | |
| end_time = self._synced_time() if self._perf_timing_enabled else None | |
| self._maybe_run_manual_validation() | |
| if not self._perf_timing_enabled: | |
| return | |
| compute_time = ( | |
| 0.0 | |
| if self._iter_start_time is None | |
| else max(0.0, end_time - self._iter_start_time) | |
| ) | |
| data_wait_time = max(0.0, self._current_data_wait_time) | |
| iter_wall_time = data_wait_time + compute_time | |
| self._prev_batch_end_time = end_time | |
| self._perf_window["count"] += 1 | |
| self._perf_window["iter_wall"] += iter_wall_time | |
| self._perf_window["data_wait"] += data_wait_time | |
| self._perf_window["compute"] += compute_time | |
| self._perf_window["backward"] += self._last_backward_time | |
| self._perf_window["optimizer"] += self._last_optimizer_time | |
| for key in ("encode_target", "dit_forward_loss"): | |
| self._perf_window[key] += self._latest_step_part_times.get(key, 0.0) | |
| if self._perf_window["count"] < self.perf_log_every_n_steps: | |
| return | |
| c = float(self._perf_window["count"]) | |
| avg_iter_wall = self._perf_window["iter_wall"] / c | |
| avg_data_wait = self._perf_window["data_wait"] / c | |
| avg_compute = self._perf_window["compute"] / c | |
| avg_backward = self._perf_window["backward"] / c | |
| avg_optimizer = self._perf_window["optimizer"] / c | |
| avg_encode_target = self._perf_window["encode_target"] / c | |
| avg_dit_forward_loss = self._perf_window["dit_forward_loss"] / c | |
| total_iter = max(avg_iter_wall, 1e-08) | |
| data_wait_ratio = avg_data_wait / total_iter | |
| compute_ratio = avg_compute / total_iter | |
| backward_ratio = avg_backward / max(avg_compute, 1e-08) | |
| bottleneck = self._classify_bottleneck(avg_data_wait, avg_compute, avg_backward) | |
| if self.trainer is not None and self.trainer.is_global_zero: | |
| print( | |
| f"[Perf][step={int(self.global_step)}][N={int(c)}] iter={avg_iter_wall:.4f}s data_wait={avg_data_wait:.4f}s({data_wait_ratio:.1%}) compute={avg_compute:.4f}s({compute_ratio:.1%}) backward={avg_backward:.4f}s(backward/compute={backward_ratio:.1%}) optimizer={avg_optimizer:.4f}s encode_target={avg_encode_target:.4f}s dit_forward_loss={avg_dit_forward_loss:.4f}s bottleneck={bottleneck}" | |
| ) | |
| self.log( | |
| "perf/iter_wall_s", | |
| avg_iter_wall, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/data_wait_s", | |
| avg_data_wait, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/compute_s", | |
| avg_compute, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/backward_s", | |
| avg_backward, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/optimizer_s", | |
| avg_optimizer, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/encode_target_s", | |
| avg_encode_target, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/dit_forward_loss_s", | |
| avg_dit_forward_loss, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/data_wait_ratio", | |
| data_wait_ratio, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/compute_ratio", | |
| compute_ratio, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/backward_in_compute_ratio", | |
| backward_ratio, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/bottleneck_data", | |
| 1.0 if bottleneck == "data" else 0.0, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/bottleneck_compute", | |
| 1.0 if bottleneck == "compute" else 0.0, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| self.log( | |
| "perf/bottleneck_communication", | |
| 1.0 if bottleneck == "communication" else 0.0, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=True, | |
| ) | |
| for key in self._perf_window.keys(): | |
| self._perf_window[key] = 0.0 if key != "count" else 0 | |
| def freeze_parameters(self): | |
| self.pipe.requires_grad_(False) | |
| self.pipe.eval() | |
| self.pipe.denoising_model().train() | |
| def _encode_prompt_with_cache(self, text: str) -> torch.Tensor: | |
| """将文本编码为 T5 context tensor (seq_len, 4096)，结果缓存在 CPU。 | |
| text 为空字符串时优先读取 repo-local blank cache；缺失则 encode 后写回。 | |
| 返回值为 CPU tensor；调用侧负责 .to(device)。 | |
| """ | |
| text = text.strip() if text else "" | |
| if not text: | |
| if self._blank_prompt_ctx is None: | |
| if os.path.isfile(self._blank_prompt_emb_path): | |
| blank = torch.load( | |
| self._blank_prompt_emb_path, | |
| map_location="cpu", | |
| weights_only=True, | |
| ) | |
| else: | |
| with torch.no_grad(): | |
| blank = self.pipe.encode_prompt("", positive=True) | |
| blank = _prompt_emb_to_cpu(blank) | |
| _save_prompt_emb_atomic(blank, self._blank_prompt_emb_path) | |
| print( | |
| f"[text-cfg] encoded blank prompt and cached to {self._blank_prompt_emb_path}", | |
| flush=True, | |
| ) | |
| ctx = blank["context"] | |
| if ctx.ndim == 3: | |
| ctx = ctx[0] | |
| self._blank_prompt_ctx = ctx.cpu() | |
| return self._blank_prompt_ctx | |
| if text not in self._prompt_encode_cache: | |
| with torch.no_grad(): | |
| emb = self.pipe.encode_prompt(text) | |
| ctx = emb["context"] | |
| if ctx.ndim == 3: | |
| ctx = ctx[0] | |
| self._prompt_encode_cache[text] = ctx.cpu() | |
| return self._prompt_encode_cache[text] | |
| def _get_context_for_seg(self, batch: dict, seg_idx: int) -> torch.Tensor: | |
| """返回 segment seg_idx 的 DiT context tensor。 | |
| - use_text_prompt=True：返回 (2, seq_len, 4096)，legacy 2-player 文本路径；空字符串 fallback 到 blank。 | |
| - use_text_prompt=False：返回 dataset 提供的 context，去掉 DataLoader batch 维； | |
| 通常为 per-player (P, seq_len, 4096)，没有 CSV prompt_emb_path 时 P 个 player 共用 blank。 | |
| 注意：返回的是 T5 原始 embedding（4096 dim），DiT 内部 text_embedding MLP 负责投影到 3072。 | |
| """ | |
| _dtype = self.pipe.torch_dtype | |
| _dev = self.device | |
| prompt_emb = batch.get("prompt_emb") or {} | |
| context = prompt_emb.get("context") if isinstance(prompt_emb, dict) else None | |
| if context is None: | |
| return self._blank_context_for_players( | |
| self.pipe.num_players, dtype=_dtype, device=_dev | |
| ) | |
| return context[0].to(dtype=_dtype, device=_dev) | |
| def _blank_context_for_players( | |
| self, num_players: int, *, dtype, device | |
| ) -> torch.Tensor: | |
| blank = self._encode_prompt_with_cache("").to(dtype=dtype, device=device) | |
| if blank.ndim != 2: | |
| raise ValueError( | |
| f"blank prompt context must be [S,D], got {tuple(blank.shape)}" | |
| ) | |
| return blank.unsqueeze(0).expand(int(num_players), -1, -1).contiguous() | |
| def _blank_context_like(self, context: torch.Tensor) -> torch.Tensor: | |
| blank = self._blank_context_for_players( | |
| int(context.shape[0]), dtype=context.dtype, device=context.device | |
| ) | |
| if tuple(blank.shape[1:]) != tuple(context.shape[1:]): | |
| raise ValueError( | |
| f"blank prompt context shape {tuple(blank.shape)} does not match training context {tuple(context.shape)}" | |
| ) | |
| return blank | |
| def _maybe_drop_text_prompt_context( | |
| self, context: torch.Tensor | |
| ) -> tuple[torch.Tensor, bool]: | |
| p = self.text_prompt_dropout_p | |
| if p <= 0.0 or not self.training: | |
| return (context, False) | |
| if bool(torch.rand((), device=context.device) < p): | |
| return (self._blank_context_like(context), True) | |
| return (context, False) | |
| def _encode_latents(pipe, video, tiler_kwargs): | |
| """VAE encode 可能返回 tensor 或 list/tuple，统一为 tensor。""" | |
| z = pipe.encode_video(video, **tiler_kwargs) | |
| if isinstance(z, (list, tuple)): | |
| z = z[0] | |
| return z | |
| def _should_corrupt_first_frame_anchor(self, seg_idx: int) -> bool: | |
| if self.first_frame_noise_p <= 0.0: | |
| return False | |
| if seg_idx < self.first_frame_noise_min_seg: | |
| return False | |
| return bool(torch.rand(1, device=self.device).item() < self.first_frame_noise_p) | |
| def _corrupt_origin_frame0(self, origin_latents: torch.Tensor) -> torch.Tensor: | |
| """(1-σ)*z + σ*ε on latent frame 0; σ ~ U(0, corrupt_ratio). Aligns with flow_match.add_noise.""" | |
| out = origin_latents.clone() | |
| z0 = out[:, :, 0:1] | |
| eps = torch.randn_like(z0) | |
| sigma = get_corrupt_noise_sigma( | |
| z0, | |
| batch_size=z0.shape[0], | |
| corrupt_ratio=self.first_frame_noise_corrupt_ratio, | |
| num_frames=1, | |
| is_frame_independent=False, | |
| ) | |
| out[:, :, 0:1] = (1.0 - sigma) * z0 + sigma * eps | |
| return out | |
| def _prepare_simplex_agents(self, latents, keyboard, mouse): | |
| """Select simplex vertices, then optionally permute aligned player streams.""" | |
| num_players = latents.shape[0] | |
| pool = self.simplex_pool_size | |
| if pool < num_players: | |
| raise ValueError(f"simplex_pool_size={pool} < num_players={num_players}") | |
| if self.simplex_train_sample_vertex_ids: | |
| vertex_ids = random.sample(range(pool), num_players) | |
| else: | |
| vertex_ids = list(range(num_players)) | |
| if self.simplex_train_permute: | |
| perm = torch.randperm(num_players).tolist() | |
| vertex_ids = [vertex_ids[i] for i in perm] | |
| latents = latents[perm] | |
| keyboard = keyboard[perm] | |
| mouse = mouse[perm] | |
| return (latents, keyboard, mouse, vertex_ids) | |
| def _apply_first_frame_conditioning( | |
| self, | |
| latents: torch.Tensor, | |
| noise: torch.Tensor, | |
| timestep: torch.Tensor, | |
| *, | |
| seg_idx: int = 0, | |
| allow_drop: bool = True, | |
| ): | |
| """Build noisy_latents with optional DF-style corrupt on anchor frame 0.""" | |
| origin_latents = latents.clone().detach() | |
| if timestep.ndim >= 1 and timestep.numel() > 1: | |
| t_bf = timestep.reshape(latents.shape[0], latents.shape[2]) | |
| sigma, _ = self.pipe.scheduler._sigma_from_timestep( | |
| t_bf, device=latents.device, dtype=latents.dtype | |
| ) | |
| sigma = sigma.view(latents.shape[0], 1, latents.shape[2], 1, 1) | |
| noisy_latents = (1.0 - sigma) * latents + sigma * noise | |
| else: | |
| noisy_latents = self.pipe.scheduler.add_noise(latents, noise, timestep) | |
| drop_first_frame = ( | |
| allow_drop | |
| and self.first_frame_dropout_p > 0.0 | |
| and bool(torch.rand(1).item() < self.first_frame_dropout_p) | |
| ) | |
| corrupt_anchor = False | |
| if not drop_first_frame: | |
| corrupt_anchor = self._should_corrupt_first_frame_anchor(seg_idx) | |
| if corrupt_anchor: | |
| origin_latents = self._corrupt_origin_frame0(origin_latents) | |
| noisy_latents[:, :, 0:1, ...] = origin_latents[:, :, 0:1, ...] | |
| return (noisy_latents, drop_first_frame, corrupt_anchor) | |
| def _stage2_causal_training_step( | |
| self, | |
| latents: torch.Tensor, | |
| keyboard: torch.Tensor, | |
| mouse: torch.Tensor, | |
| prompt_emb: dict, | |
| agent_vertex_ids, | |
| image_emb: dict, | |
| batch: dict, | |
| ) -> torch.Tensor: | |
| """Causal-Forcing teacher-forcing training for Stage 2. | |
| Uses a SINGLE parallel forward (WanModel.forward, same as standard training), | |
| NOT the KV-cache causal path. Using KV cache during the supervised block's | |
| forward breaks gradient flow to Wk, Wv, and hub_embed via in-place cache | |
| writes (cache tensors are requires_grad=False, so the autograd graph is severed | |
| at the slice-assignment boundary). | |
| Algorithm: | |
| - Sample supervised block k ∈ [0, num_blocks-1]. | |
| - Context blocks 0..k-1: timestep = 0 → sigma = 0 → noisy = clean latents. | |
| - Supervised block k: timestep = random t → noisy = latents + noise. | |
| - Frame 0: always clean (I2V anchor), timestep = 0. | |
| - Single WanModel.forward over the full sequence (block-causal mask handles | |
| attention restriction). Loss computed only on supervised block k. | |
| This gives full gradient flow to all trainable parameters (Wq, Wk, Wv, Wo, | |
| hub_embed, action layers) and matches inference semantics: block k is denoised | |
| while attending to clean context from earlier blocks. | |
| """ | |
| B, C, f_latent, H, W = latents.shape | |
| n = self.num_frame_per_block | |
| num_blocks = f_latent // n | |
| device = latents.device | |
| k = int(torch.randint(0, num_blocks, (1,)).item()) | |
| t_idx = torch.randint(0, self.pipe.scheduler.num_train_timesteps, (1,)) | |
| t_val = self.pipe.scheduler.timesteps[t_idx].to( | |
| dtype=self.pipe.torch_dtype, device=device | |
| ) | |
| timestep = torch.zeros( | |
| [B, f_latent], dtype=self.pipe.torch_dtype, device=device | |
| ) | |
| timestep[:, k * n : (k + 1) * n] = t_val | |
| timestep[:, 0] = 0.0 | |
| noise = torch.randn_like(latents) | |
| sigma, _ = self.pipe.scheduler._sigma_from_timestep( | |
| timestep, device=device, dtype=latents.dtype | |
| ) | |
| sigma_view = sigma.view(B, 1, f_latent, 1, 1) | |
| noisy_latents = (1.0 - sigma_view) * latents + sigma_view * noise | |
| noisy_latents[:, :, 0:1] = latents[:, :, 0:1] | |
| df_supervise_mask = torch.zeros(f_latent, dtype=torch.bool, device=device) | |
| df_supervise_mask[k * n : (k + 1) * n] = True | |
| df_supervise_mask[0] = False | |
| training_target = self.pipe.scheduler.training_target(latents, noise, timestep) | |
| extra_input = self.pipe.prepare_extra_input(latents) | |
| dit_kwargs = dict( | |
| mouse=mouse, | |
| keyboard=keyboard, | |
| loss_type="diffusion_forcing", | |
| **prompt_emb, | |
| **extra_input, | |
| **image_emb, | |
| use_gradient_checkpointing=self.use_gradient_checkpointing, | |
| use_gradient_checkpointing_offload=self.use_gradient_checkpointing_offload, | |
| ) | |
| if agent_vertex_ids is not None: | |
| dit_kwargs["agent_vertex_ids"] = agent_vertex_ids | |
| _attach_hier_group_dit_kwargs(dit_kwargs, batch, self.device) | |
| if self.use_camera_pose: | |
| if "camera_pose" not in batch: | |
| raise KeyError( | |
| "--use_camera_pose 需要 dataset 返回 camera_pose（CSV camera_pose_path* / camera12_path*）" | |
| ) | |
| dit_kwargs["camera_pose"] = batch["camera_pose"][0].to( | |
| device=self.device, dtype=self.pipe.torch_dtype | |
| ) | |
| noise_pred = self.pipe.denoising_model()( | |
| noisy_latents, timestep=timestep, **dit_kwargs | |
| ) | |
| sel = df_supervise_mask | |
| loss = F.mse_loss( | |
| noise_pred[:, :, sel, ...].float(), training_target[:, :, sel, ...].float() | |
| ) | |
| loss_weight = self.pipe.scheduler.training_weight(timestep[:, sel]) | |
| loss = loss * loss_weight | |
| if not torch.isfinite(loss.detach()).all(): | |
| raise RuntimeError( | |
| f"Non-finite training loss at step={int(self.global_step)} loss={loss.detach().float().cpu()}" | |
| ) | |
| return loss | |
| def _select_generated_history_plan(self, num_blocks: int): | |
| """Return a rank-synchronous (target k, history m, start g0) plan or None. | |
| The decision uses only global_step, so every DDP rank takes the same branch and | |
| uses the same block counts. It does not consume the training RNG. | |
| """ | |
| if self.generated_history_prob <= 0.0 or num_blocks < 2: | |
| return None | |
| step = int(self.global_step) | |
| mix_hash = step * 1103515245 + 12345 & 4294967295 | |
| if mix_hash / float(1 << 32) >= self.generated_history_prob: | |
| return None | |
| plan_hash = (step + 1) * 2654435761 + 1013904223 & 4294967295 | |
| k = 1 + plan_hash % (num_blocks - 1) | |
| max_m = min(3, k) | |
| m = 1 + (plan_hash >> 16) % max_m | |
| return (int(k), int(m), int(k - m)) | |
| def _generated_history_negative_context(self, num_players: int, *, dtype, device): | |
| """Encode the deployment negative prompt once and keep its cache on CPU.""" | |
| if self._generated_history_negative_context_cpu is None: | |
| self.pipe.text_encoder.to(device) | |
| encoded = self.pipe.encode_prompt( | |
| [_GENERATED_HISTORY_NEGATIVE_PROMPT] * int(num_players), positive=False | |
| )["context"] | |
| self._generated_history_negative_context_cpu = encoded.detach().cpu() | |
| self.pipe.text_encoder.to("cpu") | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| context = self._generated_history_negative_context_cpu | |
| if int(context.shape[0]) != int(num_players): | |
| raise ValueError( | |
| f"negative prompt cache P={context.shape[0]} does not match P={num_players}" | |
| ) | |
| return context.to(device=device, dtype=dtype) | |
| def _rollout_generated_history_online( | |
| self, | |
| *, | |
| latents, | |
| keyboard, | |
| mouse, | |
| positive_context, | |
| agent_vertex_ids, | |
| batch, | |
| target_block, | |
| history_blocks, | |
| batch_idx, | |
| ): | |
| """Generate the most recent 1--3 history blocks with the real cached CFG path. | |
| Input latents are [P,C,T,H,W]. Blocks [0,g0) are committed as clean GT, | |
| blocks [g0,k) are generated with the deployment sampler, and the returned | |
| [P,C,m*n,H,W] tensor is detached. The student target is intentionally not | |
| evaluated here; it stays on the parallel differentiable training path. | |
| """ | |
| P, C, T, H, W = latents.shape | |
| n = int(self.num_frame_per_block) | |
| k = int(target_block) | |
| m = int(history_blocks) | |
| g0 = k - m | |
| if not (1 <= k < T // n and 1 <= m <= min(3, k)): | |
| raise ValueError(f"invalid generated-history plan T={T} n={n} k={k} m={m}") | |
| rollout_scheduler = copy.deepcopy(self.pipe.scheduler) | |
| rollout_scheduler.set_timesteps( | |
| self.val_num_inference_steps, denoising_strength=1.0, shift=5.0 | |
| ) | |
| negative_context = self._generated_history_negative_context( | |
| P, dtype=latents.dtype, device=latents.device | |
| ) | |
| group_labels = batch["group_labels"][0].to( | |
| device=latents.device, dtype=torch.long | |
| ) | |
| active_group = batch["active_group"][0].to( | |
| device=latents.device, dtype=torch.bool | |
| ) | |
| membership = batch.get("membership") | |
| global_only = batch.get("global_only") | |
| if membership is not None: | |
| if global_only is None: | |
| raise KeyError( | |
| "membership generated-history rollout requires global_only" | |
| ) | |
| membership = membership[0].to(device=latents.device, dtype=torch.bool) | |
| global_only = global_only[0].to(device=latents.device, dtype=torch.bool) | |
| camera_pose = None | |
| if self.use_camera_pose: | |
| camera_pose = batch["camera_pose"][0].to( | |
| device=latents.device, dtype=latents.dtype | |
| ) | |
| def group_kwargs(lo, hi): | |
| out = { | |
| "group_labels": group_labels[lo:hi], | |
| "active_group": active_group[lo:hi], | |
| } | |
| if membership is not None: | |
| out.update( | |
| membership=membership, | |
| global_only=global_only, | |
| active_group=active_group, | |
| ) | |
| return out | |
| old_kv, old_cross = (self.pipe.kv_cache1, self.pipe.crossattn_cache) | |
| def new_cache_pair(): | |
| self.pipe.kv_cache1, self.pipe.crossattn_cache = (None, None) | |
| self.pipe.init_kv_cache_for_training(batch_size=P, device=latents.device) | |
| return (self.pipe.kv_cache1, self.pipe.crossattn_cache) | |
| kv_pos, cross_pos = new_cache_pair() | |
| kv_neg, cross_neg = new_cache_pair() | |
| if kv_pos is kv_neg or cross_pos is cross_neg: | |
| raise AssertionError("positive and negative CFG caches must be independent") | |
| token_stride = int(self.pipe.frame_seq_length) * P | |
| common = dict( | |
| mouse=mouse, | |
| keyboard=keyboard, | |
| loss_type="diffusion_forcing", | |
| agent_vertex_ids=agent_vertex_ids, | |
| camera_pose_full=camera_pose, | |
| ) | |
| pos_common = dict( | |
| common, context=positive_context, kv_cache=kv_pos, crossattn_cache=cross_pos | |
| ) | |
| neg_common = dict( | |
| common, context=negative_context, kv_cache=kv_neg, crossattn_cache=cross_neg | |
| ) | |
| module_states = [ | |
| (module, module.training) for module in self.pipe.dit.modules() | |
| ] | |
| self.pipe.dit.eval() | |
| generated = [] | |
| started = time.perf_counter() | |
| try: | |
| warm_t = torch.zeros((P, 1), device=latents.device, dtype=torch.int64) | |
| for branch in (pos_common, neg_common): | |
| model_fn_wan_video_causal( | |
| self.pipe.dit, | |
| latents[:, :, :1], | |
| timestep=warm_t, | |
| current_start=0, | |
| cache_start=0, | |
| **group_kwargs(0, 1), | |
| **branch, | |
| ) | |
| for block_idx in range(g0): | |
| lo, hi = (block_idx * n, (block_idx + 1) * n) | |
| context_timestep = torch.zeros( | |
| (P, n), device=latents.device, dtype=torch.int64 | |
| ) | |
| for branch in (pos_common, neg_common): | |
| model_fn_wan_video_causal( | |
| self.pipe.dit, | |
| latents[:, :, lo:hi], | |
| timestep=context_timestep, | |
| current_start=lo * token_stride, | |
| cache_start=lo * token_stride, | |
| **group_kwargs(lo, hi), | |
| **branch, | |
| ) | |
| generator = torch.Generator(device=latents.device) | |
| seed = ( | |
| 12648430 | |
| + 1000003 * int(self.global_rank) | |
| + 9176 * int(self.global_step) | |
| + int(batch_idx) | |
| & 9223372036854775807 | |
| ) | |
| generator.manual_seed(seed) | |
| for block_idx in range(g0, k): | |
| lo, hi = (block_idx * n, (block_idx + 1) * n) | |
| noisy = torch.randn( | |
| (P, C, n, H, W), | |
| generator=generator, | |
| device=latents.device, | |
| dtype=latents.dtype, | |
| ) | |
| if block_idx == 0: | |
| noisy[:, :, :1] = latents[:, :, :1] | |
| for current_timestep in rollout_scheduler.timesteps: | |
| timestep = torch.full( | |
| (P, n), | |
| current_timestep, | |
| device=latents.device, | |
| dtype=torch.int64, | |
| ) | |
| if block_idx == 0: | |
| timestep[:, 0] = 0 | |
| noisy[:, :, :1] = latents[:, :, :1] | |
| pred_pos = model_fn_wan_video_causal( | |
| self.pipe.dit, | |
| noisy, | |
| timestep=timestep, | |
| current_start=lo * token_stride, | |
| cache_start=lo * token_stride, | |
| **group_kwargs(lo, hi), | |
| **pos_common, | |
| ) | |
| pred_neg = model_fn_wan_video_causal( | |
| self.pipe.dit, | |
| noisy, | |
| timestep=timestep, | |
| current_start=lo * token_stride, | |
| cache_start=lo * token_stride, | |
| **group_kwargs(lo, hi), | |
| **neg_common, | |
| ) | |
| denoised_pred = pred_neg + 5.0 * (pred_pos - pred_neg) | |
| noisy = rollout_scheduler.step( | |
| denoised_pred, current_timestep, noisy | |
| ) | |
| if block_idx == 0: | |
| noisy[:, :, :1] = latents[:, :, :1] | |
| generated.append(noisy.detach().clone()) | |
| context_timestep = torch.zeros( | |
| (P, n), device=latents.device, dtype=torch.int64 | |
| ) | |
| for branch in (pos_common, neg_common): | |
| model_fn_wan_video_causal( | |
| self.pipe.dit, | |
| noisy, | |
| timestep=context_timestep, | |
| current_start=lo * token_stride, | |
| cache_start=lo * token_stride, | |
| **group_kwargs(lo, hi), | |
| **branch, | |
| ) | |
| finally: | |
| for module, state in module_states: | |
| module.training = state | |
| self.pipe.kv_cache1, self.pipe.crossattn_cache = (old_kv, old_cross) | |
| history = torch.cat(generated, dim=2).detach() | |
| expected_shape = (P, C, m * n, H, W) | |
| if tuple(history.shape) != expected_shape: | |
| raise AssertionError( | |
| f"generated history shape={tuple(history.shape)}, expected={expected_shape}" | |
| ) | |
| if history.requires_grad or history.grad_fn is not None: | |
| raise AssertionError("generated history must be detached") | |
| if not torch.isfinite(history).all(): | |
| raise RuntimeError("generated history contains non-finite values") | |
| if int(self.global_rank) == 0: | |
| print( | |
| f"[generated-history] step={int(self.global_step)} k={k} m={m} g0={g0} sampler_steps={self.val_num_inference_steps} cfg=5 rollout_s={time.perf_counter() - started:.3f} shape={tuple(history.shape)}", | |
| flush=True, | |
| ) | |
| return history | |
| def training_step(self, batch, batch_idx): | |
| self.pipe.device = self.device | |
| perf = self._perf_timing_enabled | |
| step_part_times = {"encode_target": 0.0, "dit_forward_loss": 0.0} | |
| if perf: | |
| t_encode_target_start = self._synced_time() | |
| if "latents" in batch: | |
| if batch["latents"].shape[0] != 1: | |
| raise ValueError( | |
| "N-player precomputed training currently expects DataLoader batch_size=1" | |
| ) | |
| latents = ( | |
| batch["latents"] | |
| .squeeze(0) | |
| .to(dtype=self.pipe.torch_dtype, device=self.pipe.device) | |
| ) | |
| if latents.shape[0] != self.pipe.num_players: | |
| raise ValueError( | |
| f"latents P={latents.shape[0]} does not match num_players={self.pipe.num_players}" | |
| ) | |
| elif "latents1" in batch: | |
| latents1 = batch["latents1"].to( | |
| dtype=self.pipe.torch_dtype, device=self.pipe.device | |
| ) | |
| latents2 = batch["latents2"].to( | |
| dtype=self.pipe.torch_dtype, device=self.pipe.device | |
| ) | |
| latents = torch.cat((latents1, latents2), dim=0) | |
| else: | |
| video1 = batch["video1"] | |
| video1 = video1.to(dtype=self.pipe.torch_dtype, device=self.pipe.device) | |
| latents1 = self._encode_latents(self.pipe, video1, self.tiler_kwargs) | |
| latents1 = latents1.to(dtype=self.pipe.torch_dtype, device=self.pipe.device) | |
| video2 = batch["video2"] | |
| video2 = video2.to(dtype=self.pipe.torch_dtype, device=self.pipe.device) | |
| latents2 = self._encode_latents(self.pipe, video2, self.tiler_kwargs) | |
| latents2 = latents2.to(dtype=self.pipe.torch_dtype, device=self.pipe.device) | |
| latents = torch.cat((latents1, latents2), dim=0) | |
| if perf: | |
| step_part_times["encode_target"] = max( | |
| 0.0, self._synced_time() - t_encode_target_start | |
| ) | |
| image_emb = {} | |
| keyboard = batch["keyboard"].squeeze(0)[:, :-1, :] | |
| mouse = batch["mouse"].squeeze(0)[:, :-1, :] | |
| keyboard = keyboard.to(device=self.pipe.device) | |
| mouse = mouse.to(device=self.pipe.device) | |
| if keyboard.shape[0] != latents.shape[0] or mouse.shape[0] != latents.shape[0]: | |
| raise ValueError( | |
| f"Player dim mismatch: latents P={latents.shape[0]}, keyboard P={keyboard.shape[0]}, mouse P={mouse.shape[0]}" | |
| ) | |
| latents, keyboard, mouse, agent_vertex_ids = self._prepare_simplex_agents( | |
| latents, keyboard, mouse | |
| ) | |
| generated_history_plan = None | |
| if self.generated_history_prob > 0.0: | |
| if latents.shape[2] % self.num_frame_per_block != 0: | |
| raise ValueError( | |
| f"generated-history latent T={latents.shape[2]} is not divisible by block size {self.num_frame_per_block}" | |
| ) | |
| generated_history_plan = self._select_generated_history_plan( | |
| int(latents.shape[2]) // self.num_frame_per_block | |
| ) | |
| prompt_emb = dict(batch.get("prompt_emb", {})) | |
| prompt_context = self._get_context_for_seg(batch, seg_idx=0) | |
| if generated_history_plan is None: | |
| prompt_context, text_prompt_dropped = self._maybe_drop_text_prompt_context( | |
| prompt_context | |
| ) | |
| else: | |
| text_prompt_dropped = False | |
| prompt_emb["context"] = prompt_context | |
| self.log( | |
| "train/text_prompt_dropped", | |
| 1.0 if text_prompt_dropped else 0.0, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=False, | |
| ) | |
| if perf: | |
| t_dit_forward_loss_start = self._synced_time() | |
| if self.causal_forcing_training: | |
| loss = self._stage2_causal_training_step( | |
| latents=latents, | |
| keyboard=keyboard, | |
| mouse=mouse, | |
| prompt_emb=prompt_emb, | |
| agent_vertex_ids=agent_vertex_ids, | |
| image_emb=image_emb, | |
| batch=batch, | |
| ) | |
| loss_weight = torch.tensor(1.0) | |
| self.log("train_loss", loss, prog_bar=True) | |
| self.log( | |
| "train/is_recon_warmup", | |
| 0.0, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=False, | |
| ) | |
| self._latest_step_part_times = step_part_times | |
| return loss | |
| noise = torch.randn_like(latents) | |
| f_latent = latents.shape[2] | |
| df_supervise_mask = None | |
| clean_prefix_block = None | |
| n_block = self.num_frame_per_block | |
| num_blocks = f_latent // n_block | |
| if self.df_per_frame_noise: | |
| timestep_per_frame = torch.randint( | |
| 0, self.pipe.scheduler.num_train_timesteps, (1, f_latent) | |
| ) | |
| else: | |
| timestep_id = torch.randint( | |
| 0, self.pipe.scheduler.num_train_timesteps, (1, num_blocks) | |
| ) | |
| timestep_per_frame = timestep_id.repeat_interleave(n_block, dim=1) | |
| if timestep_per_frame.shape[1] < f_latent: | |
| pad = timestep_id[:, -1:].expand( | |
| -1, f_latent - timestep_per_frame.shape[1] | |
| ) | |
| timestep_per_frame = torch.cat([timestep_per_frame, pad], dim=1) | |
| elif timestep_per_frame.shape[1] > f_latent: | |
| timestep_per_frame = timestep_per_frame[:, :f_latent] | |
| t_idx = self.pipe.scheduler.timesteps[timestep_per_frame[0]].to( | |
| dtype=self.pipe.torch_dtype, device=self.pipe.device | |
| ) | |
| timestep = t_idx.unsqueeze(0).expand(latents.shape[0], -1).clone() | |
| if generated_history_plan is not None: | |
| k, m, g0 = generated_history_plan | |
| prefix_end = k * n_block | |
| timestep[:, :prefix_end] = 0.0 | |
| cur_lo = k * n_block | |
| cur_hi = min((k + 1) * n_block, f_latent) | |
| df_supervise_mask = torch.zeros( | |
| f_latent, dtype=torch.bool, device=timestep.device | |
| ) | |
| df_supervise_mask[cur_lo:cur_hi] = True | |
| generated_history = self._rollout_generated_history_online( | |
| latents=latents, | |
| keyboard=keyboard, | |
| mouse=mouse, | |
| positive_context=prompt_emb["context"], | |
| agent_vertex_ids=agent_vertex_ids, | |
| batch=batch, | |
| target_block=k, | |
| history_blocks=m, | |
| batch_idx=batch_idx, | |
| ) | |
| mixed_latents = latents.clone() | |
| mixed_latents[:, :, g0 * n_block : k * n_block] = generated_history | |
| latents = mixed_latents | |
| elif ( | |
| self.clean_prefix_prob > 0.0 | |
| and num_blocks >= 2 | |
| and bool(torch.rand(1).item() < self.clean_prefix_prob) | |
| ): | |
| k = int(torch.randint(1, num_blocks, (1,)).item()) | |
| clean_prefix_block = k | |
| prefix_end = k * n_block | |
| timestep[:, :prefix_end] = float(self.clean_prefix_context_noise) | |
| cur_lo = k * n_block | |
| cur_hi = min((k + 1) * n_block, f_latent) | |
| df_supervise_mask = torch.zeros( | |
| f_latent, dtype=torch.bool, device=timestep.device | |
| ) | |
| df_supervise_mask[cur_lo:cur_hi] = True | |
| timestep[:, 0] = 0.0 | |
| extra_input = self.pipe.prepare_extra_input(latents) | |
| noisy_latents, drop_first_frame, corrupt_anchor = ( | |
| self._apply_first_frame_conditioning( | |
| latents, | |
| noise, | |
| timestep, | |
| seg_idx=0, | |
| allow_drop=df_supervise_mask is None, | |
| ) | |
| ) | |
| training_target = self.pipe.scheduler.training_target(latents, noise, timestep) | |
| dit_kwargs = dict( | |
| mouse=mouse, | |
| keyboard=keyboard, | |
| loss_type="diffusion_forcing", | |
| **prompt_emb, | |
| **extra_input, | |
| **image_emb, | |
| use_gradient_checkpointing=self.use_gradient_checkpointing, | |
| use_gradient_checkpointing_offload=self.use_gradient_checkpointing_offload, | |
| ) | |
| if agent_vertex_ids is not None: | |
| dit_kwargs["agent_vertex_ids"] = agent_vertex_ids | |
| _attach_hier_group_dit_kwargs(dit_kwargs, batch, self.device) | |
| if self.use_camera_pose: | |
| if "camera_pose" not in batch: | |
| raise KeyError( | |
| "--use_camera_pose 需要 dataset 返回 camera_pose（CSV camera_pose_path* / camera12_path*）" | |
| ) | |
| dit_kwargs["camera_pose"] = batch["camera_pose"][0].to( | |
| device=self.device, dtype=self.pipe.torch_dtype | |
| ) | |
| noise_pred = self.pipe.denoising_model()( | |
| noisy_latents, timestep=timestep, **dit_kwargs | |
| ) | |
| if df_supervise_mask is not None: | |
| sel = df_supervise_mask | |
| loss = torch.nn.functional.mse_loss( | |
| noise_pred[:, :, sel, ...].float(), | |
| training_target[:, :, sel, ...].float(), | |
| ) | |
| loss_weight = self.pipe.scheduler.training_weight(timestep[:, sel]) | |
| elif drop_first_frame: | |
| loss = torch.nn.functional.mse_loss( | |
| noise_pred.float(), training_target.float() | |
| ) | |
| loss_weight = self.pipe.scheduler.training_weight(timestep) | |
| else: | |
| loss = torch.nn.functional.mse_loss( | |
| noise_pred[:, :, 1:, ...].float(), | |
| training_target[:, :, 1:, ...].float(), | |
| ) | |
| loss_weight = self.pipe.scheduler.training_weight(timestep[:, 1:]) | |
| loss = loss * loss_weight | |
| if not torch.isfinite(loss.detach()).all(): | |
| raise RuntimeError( | |
| f"Non-finite training loss at step={int(self.global_step)} batch_idx={int(batch_idx)} loss={loss.detach().float().cpu()}" | |
| ) | |
| if perf: | |
| step_part_times["dit_forward_loss"] = max( | |
| 0.0, self._synced_time() - t_dit_forward_loss_start | |
| ) | |
| self._latest_step_part_times = step_part_times | |
| self.log("train_loss", loss, prog_bar=True) | |
| self.log( | |
| "train/is_recon_warmup", | |
| 0.0, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=False, | |
| ) | |
| self.log( | |
| "train/first_frame_dropped", | |
| 1.0 if drop_first_frame else 0.0, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=False, | |
| ) | |
| self.log( | |
| "train/first_frame_corrupted", | |
| 1.0 if corrupt_anchor else 0.0, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=False, | |
| ) | |
| self.log( | |
| "train/clean_prefix", | |
| 1.0 if clean_prefix_block is not None else 0.0, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=False, | |
| ) | |
| self.log( | |
| "train/generated_history", | |
| 1.0 if generated_history_plan is not None else 0.0, | |
| on_step=True, | |
| prog_bar=False, | |
| logger=True, | |
| sync_dist=False, | |
| ) | |
| return loss | |
| def _maybe_run_manual_validation(self): | |
| """方案二：按 global_step 触发验证，与 epoch 边界解耦。 | |
| 在 on_train_batch_end 里每步调用。所有 rank 依据同一 global_step 判定是否验证 | |
| （DDP 下 global_step 同步，判定一致），只有 rank0 真正跑生成/记录（无集合通信）， | |
| 其余 rank 到末尾 barrier 等待，避免跑太远 / logger 竞争。 | |
| """ | |
| if not self.val_infer or self._manual_val_loader is None: | |
| return | |
| gs = int(self.global_step) | |
| if gs <= 0 or gs % self.val_every_n_steps != 0: | |
| return | |
| if getattr(self, "_last_val_step", -1) == gs: | |
| return | |
| self._last_val_step = gs | |
| was_training = self.training | |
| self.eval() | |
| try: | |
| if self.trainer.is_global_zero: | |
| try: | |
| with torch.no_grad(): | |
| for i, batch in enumerate(self._manual_val_loader): | |
| self._run_val_sample(batch, i) | |
| except Exception as e: | |
| print( | |
| f"[val] manual validation failed @ step {gs}: {e}", flush=True | |
| ) | |
| finally: | |
| if was_training: | |
| self.train() | |
| self.pipe.scheduler.set_timesteps(1000, training=True) | |
| if ( | |
| getattr(self.trainer, "world_size", 1) > 1 | |
| and torch.distributed.is_initialized() | |
| ): | |
| torch.distributed.barrier() | |
| def _run_val_sample(self, batch, batch_idx): | |
| """训练集前 N 条的 causal_inference 可视化：只在 rank0 跑，生成/GT 视频记到 logger。 | |
| 依赖 PrecomputedLatentDataset 的 batch 结构（batch_size=1）： | |
| latents [1,P,48,T_lat,h,w]、keyboard [1,P,T,23]、mouse [1,P,T,2]。 | |
| KV cache 与 δ-mem state 在 causal_inference 内部逐次 reset，故每条样本独立。 | |
| """ | |
| if "latents" not in batch and "latents1" not in batch: | |
| return | |
| from .data.video import save_video | |
| from multics.data.utils import tensor_bchw_to_pil_list, pad_pil_list_to_length | |
| from PIL import Image as _PILImage | |
| gs = int(self.global_step) | |
| _dtype, _dev = (self.pipe.torch_dtype, self.device) | |
| if "latents" in batch: | |
| video_latents = batch["latents"].squeeze(0).to(dtype=_dtype, device=_dev) | |
| else: | |
| video_latents = torch.cat([batch["latents1"], batch["latents2"]], dim=0).to( | |
| dtype=_dtype, device=_dev | |
| ) | |
| P = int(video_latents.shape[0]) | |
| keyboard = batch["keyboard"].squeeze(0)[:, :-1, :].to(_dev) | |
| mouse = batch["mouse"].squeeze(0)[:, :-1, :].to(_dev) | |
| camera_pose = None | |
| if self.use_camera_pose: | |
| if "camera_pose" not in batch: | |
| raise KeyError("--use_camera_pose validation 需要 camera_pose") | |
| camera_pose = batch["camera_pose"].squeeze(0).to(dtype=_dtype, device=_dev) | |
| self.pipe.load_models_to_device(["vae"]) | |
| gt_all = self.pipe.decode_video(video_latents, **self.tiler_kwargs) | |
| gt_a = gt_all[0:1] | |
| gt_b = gt_all[1:2] | |
| prompt = "" | |
| negative_prompt = "" | |
| has_caption_prompt = False | |
| if "prompt" in batch: | |
| prompt = batch["prompt"] | |
| if isinstance(prompt, (list, tuple)): | |
| prompt = [ | |
| p[0] if isinstance(p, (list, tuple)) and len(p) > 0 else p | |
| for p in prompt | |
| ] | |
| else: | |
| prompt = [prompt] | |
| prompt = [str(p).strip() for p in prompt] | |
| has_caption_prompt = any(prompt) | |
| if has_caption_prompt: | |
| negative_prompt = [ | |
| "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" | |
| ] * len(prompt) | |
| else: | |
| prompt = "" | |
| val_cfg_scale = 5.0 if has_caption_prompt else 1.0 | |
| if batch_idx == 0: | |
| print( | |
| f"[val] prompt={('caption1/caption2' if has_caption_prompt else 'blank')} cfg_scale={val_cfg_scale}", | |
| flush=True, | |
| ) | |
| try: | |
| hier_infer_kwargs = {} | |
| _attach_hier_group_dit_kwargs(hier_infer_kwargs, batch, self.device) | |
| gen_players = self.pipe.causal_inference( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| video1=gt_a, | |
| video2=gt_b, | |
| mouse=mouse, | |
| keyboard=keyboard, | |
| num_frames=self.val_causal_num_frames, | |
| height=gt_a.shape[-2], | |
| width=gt_a.shape[-1], | |
| cfg_scale=val_cfg_scale, | |
| num_inference_steps=self.val_num_inference_steps, | |
| seed=0, | |
| tiled=self.tiler_kwargs["tiled"], | |
| tile_size=self.tiler_kwargs["tile_size"], | |
| tile_stride=self.tiler_kwargs["tile_stride"], | |
| loss_type="diffusion_forcing", | |
| return_both_players=True, | |
| agent_vertex_ids=self.val_agent_vertex_ids, | |
| video_latents=video_latents, | |
| camera_pose=camera_pose, | |
| overlap_path=None, | |
| **hier_infer_kwargs, | |
| ) | |
| finally: | |
| self.pipe.scheduler.set_timesteps(1000, training=True) | |
| if isinstance(gen_players, tuple): | |
| gen_players = list(gen_players) | |
| if len(gen_players) != P: | |
| raise RuntimeError( | |
| f"validation generated {len(gen_players)} players, expected P={P}" | |
| ) | |
| gen_a, gen_b = (gen_players[0], gen_players[1]) | |
| gt_a_pils = tensor_bchw_to_pil_list(gt_a) | |
| gt_b_pils = tensor_bchw_to_pil_list(gt_b) | |
| T = max(len(gen_a), len(gen_b), len(gt_a_pils), len(gt_b_pils)) | |
| ga = pad_pil_list_to_length(gen_a, T) | |
| gb = pad_pil_list_to_length(gen_b, T) | |
| qa = pad_pil_list_to_length(gt_a_pils, T) | |
| qb = pad_pil_list_to_length(gt_b_pils, T) | |
| w0, h0 = ga[0].size | |
| _cell = lambda im: im if im.size == (w0, h0) else im.resize((w0, h0)) | |
| bw = 2 | |
| total_w, total_h = (2 * w0 + bw, 2 * h0 + bw) | |
| grid = [] | |
| for t in range(T): | |
| canvas = _PILImage.new("RGB", (total_w, total_h), (0, 0, 0)) | |
| canvas.paste(_cell(ga[t]), (0, 0)) | |
| canvas.paste(_cell(gb[t]), (w0 + bw, 0)) | |
| canvas.paste(_cell(qa[t]), (0, h0 + bw)) | |
| canvas.paste(_cell(qb[t]), (w0 + bw, h0 + bw)) | |
| grid.append(canvas) | |
| os.makedirs(self.val_output_dir, exist_ok=True) | |
| grid_path = os.path.join(self.val_output_dir, f"step{gs}_s{batch_idx}_grid.mp4") | |
| save_video(grid, grid_path, fps=self.val_grid_fps, quality=5) | |
| def configure_optimizers(self): | |
| params = [ | |
| p for p in self.pipe.denoising_model().parameters() if p.requires_grad | |
| ] | |
| optimizer = torch.optim.AdamW(params, lr=self.learning_rate) | |
| return optimizer | |
| def _resolve_local_version_dir(self): | |
| """Resolve a local run directory for logs and checkpoints.""" | |
| try: | |
| ld = self.trainer.logger.log_dir | |
| except Exception: | |
| ld = None | |
| if ld: | |
| _out = os.path.abspath(self.output_path) | |
| if os.path.abspath(ld).startswith(_out + os.sep): | |
| return ld | |
| ver = self.run_version or "version_0" | |
| return os.path.join(self.output_path, ver) | |
| def on_train_start(self): | |
| """Record trainable parameter names in the local run directory.""" | |
| if not self.trainer.is_global_zero: | |
| return | |
| local_version_dir = self._resolve_local_version_dir() | |
| if not local_version_dir: | |
| print( | |
| "[on_train_start] cannot resolve local version dir; skip train_module.txt dump." | |
| ) | |
| return | |
| os.makedirs(local_version_dir, exist_ok=True) | |
| file_path = os.path.join(local_version_dir, "train_module.txt") | |
| with open(file_path, "w", encoding="utf-8") as f: | |
| for name, param in self.pipe.dit.named_parameters(): | |
| f.write(f"name: {name} grad: {param.requires_grad}\n") | |
| print(f"[on_train_start] wrote {file_path}") | |
| def _assert_dit_state_finite_before_save(self): | |
| bad = [] | |
| for name, param in self.pipe.denoising_model().named_parameters(): | |
| if not param.is_floating_point() and (not param.is_complex()): | |
| continue | |
| try: | |
| finite = bool(torch.isfinite(param.detach()).all().item()) | |
| except RuntimeError as exc: | |
| bad.append( | |
| f"{name}: shape={tuple(param.shape)} dtype={param.dtype} check_failed={exc}" | |
| ) | |
| continue | |
| if not finite: | |
| bad.append(f"{name}: shape={tuple(param.shape)} dtype={param.dtype}") | |
| if len(bad) >= 12: | |
| break | |
| if bad: | |
| raise RuntimeError( | |
| "DiT contains non-finite parameter(s) at checkpoint save time; refusing to save a checkpoint that will produce NaN/black videos.\n" | |
| + "\n".join((f" - {item}" for item in bad)) | |
| ) | |
| def on_save_checkpoint(self, checkpoint): | |
| checkpoint.clear() | |
| if not self.trainer.is_global_zero: | |
| return | |
| checkpoint_dir = self.trainer.checkpoint_callback.dirpath | |
| if self.ckpt_local_dir: | |
| save_dir = self.ckpt_local_dir | |
| else: | |
| save_dir = os.path.join(self._resolve_local_version_dir(), "checkpoints") | |
| os.makedirs(save_dir, exist_ok=True) | |
| print(f"Checkpoint save_dir: {save_dir}") | |
| current_step = self.global_step | |
| print(f"Current step: {current_step}") | |
| saved_files = [] | |
| self._assert_dit_state_finite_before_save() | |
| fname = f"step{current_step}.ckpt" | |
| torch.save( | |
| self.pipe.denoising_model().state_dict(), os.path.join(save_dir, fname) | |
| ) | |
| saved_files.append(fname) | |
| def _trainer_cluster_kwargs(): | |
| """Match Lightning num_nodes/devices to torchrun (WORLD_SIZE / LOCAL_WORLD_SIZE).""" | |
| ws = os.environ.get("WORLD_SIZE") | |
| lws = os.environ.get("LOCAL_WORLD_SIZE") | |
| if ws is not None and lws is not None: | |
| try: | |
| ws_i, lws_i = (int(ws), int(lws)) | |
| if lws_i > 0 and ws_i >= lws_i and (ws_i % lws_i == 0): | |
| return {"num_nodes": ws_i // lws_i, "devices": lws_i} | |
| except ValueError: | |
| pass | |
| return {"num_nodes": 1, "devices": "auto"} | |
| def _read_int_hparam(path: str, key: str, default: int = 0) -> int: | |
| if not path or not os.path.isfile(path): | |
| return int(default) | |
| with open(path, encoding="utf-8") as f: | |
| for raw in f: | |
| field, sep, value = raw.strip().partition(":") | |
| if not sep or field.strip() != key: | |
| continue | |
| value = value.strip() | |
| try: | |
| parsed = json.loads(value) | |
| except json.JSONDecodeError: | |
| parsed = value.strip("'\"") | |
| if parsed in (None, ""): | |
| return int(default) | |
| return int(parsed) | |
| return int(default) | |
| def _resolve_train_step_offset( | |
| resume_ckpt_path: str, explicit_offset, *, schedule_enabled: bool | |
| ) -> int: | |
| """Resolve cumulative completed updates for weight-only resume checkpoints.""" | |
| if explicit_offset is not None: | |
| offset = int(explicit_offset) | |
| if offset < 0: | |
| raise ValueError(f"train_step_offset must be >= 0, got {offset}") | |
| print(f"[train-step] using explicit train_step_offset={offset}", flush=True) | |
| return offset | |
| if not resume_ckpt_path: | |
| return 0 | |
| source_path = os.environ.get("WAN_RESUME_CKPT_SOURCE_PATH") or resume_ckpt_path | |
| match = re.fullmatch("step(\\d+)\\.ckpt", os.path.basename(source_path)) | |
| if match is None: | |
| message = f"cannot infer train_step_offset: resume checkpoint basename must be stepN.ckpt, got {source_path!r}" | |
| if schedule_enabled: | |
| raise ValueError(message + "; pass --train_step_offset explicitly") | |
| print(f"[train-step][warn] {message}; using 0", flush=True) | |
| return 0 | |
| checkpoint_step = int(match.group(1)) | |
| parent_hparams = os.path.join( | |
| os.path.dirname(os.path.dirname(source_path)), "hparams.yaml" | |
| ) | |
| parent_offset = _read_int_hparam(parent_hparams, "train_step_offset", default=0) | |
| offset = parent_offset + checkpoint_step | |
| print( | |
| f"[train-step] auto-resolved parent_offset={parent_offset} checkpoint_step={checkpoint_step} train_step_offset={offset} source={source_path}", | |
| flush=True, | |
| ) | |
| return offset | |
| def _write_hparams_yaml(path, args): | |
| """Write Lightning-style hparams.yaml next to checkpoints (TB logger does not).""" | |
| payload = {k: "" if v is None else str(v) for k, v in vars(args).items()} | |
| os.makedirs(os.path.dirname(path), exist_ok=True) | |
| with open(path, "w", encoding="utf-8") as f: | |
| for key, value in payload.items(): | |
| f.write(f"{key}: {json.dumps(value, ensure_ascii=False)}\n") | |
| def _compute_run_version(output_path): | |
| """Choose the next local lightning_logs/version_N directory.""" | |
| base = os.path.join(output_path, "lightning_logs") | |
| max_ver = -1 | |
| if os.path.isdir(base): | |
| for name in os.listdir(base): | |
| m = re.match("version_(\\d+)$", name) | |
| if m: | |
| max_ver = max(max_ver, int(m.group(1))) | |
| return f"version_{max_ver + 1}" | |
| def train(args): | |
| schedule_enabled = ( | |
| args.local_attn_size_after is not None | |
| or args.local_attn_switch_step is not None | |
| ) | |
| if (args.local_attn_size_after is None) != (args.local_attn_switch_step is None): | |
| raise ValueError( | |
| "--local_attn_size_after and --local_attn_switch_step must be set together" | |
| ) | |
| if args.local_attn_switch_step is not None and args.local_attn_switch_step < 0: | |
| raise ValueError("--local_attn_switch_step must be >= 0") | |
| args.train_step_offset = _resolve_train_step_offset( | |
| args.resume_ckpt_path, args.train_step_offset, schedule_enabled=schedule_enabled | |
| ) | |
| latent_t = (args.target_rgb_frames - 1) // 4 + 1 | |
| if latent_t % args.num_frame_per_block != 0: | |
| raise ValueError( | |
| f"latent_T={latent_t} must be divisible by num_frame_per_block={args.num_frame_per_block}" | |
| ) | |
| if args.local_attn_size != -1 and args.local_attn_size < args.num_frame_per_block: | |
| raise ValueError( | |
| f"local_attn_size must be -1 or >= num_frame_per_block; got local={args.local_attn_size}, n={args.num_frame_per_block}" | |
| ) | |
| if ( | |
| args.local_attn_size_after is not None | |
| and args.local_attn_size_after != -1 | |
| and (args.local_attn_size_after < args.num_frame_per_block) | |
| ): | |
| raise ValueError( | |
| f"local_attn_size_after must be -1 or >= num_frame_per_block; got local={args.local_attn_size_after}, n={args.num_frame_per_block}" | |
| ) | |
| if args.simplex_train_permute: | |
| raise ValueError("hier_hub requires --no_simplex_train_permute") | |
| kg = args.num_global_hub_tokens | |
| if kg is None: | |
| kg = args.num_hub_tokens - args.num_groups * args.num_local_hub_tokens | |
| if args.num_groups * args.num_local_hub_tokens + kg != args.num_hub_tokens: | |
| raise ValueError( | |
| f"hier_hub requires num_groups*num_local_hub_tokens + num_global_hub_tokens == num_hub_tokens; got {args.num_groups}*{args.num_local_hub_tokens}+{kg} != {args.num_hub_tokens}" | |
| ) | |
| if args.hier_no_local_hub: | |
| if args.hier_local_hub_mode != "shared": | |
| raise ValueError( | |
| "--hier_no_local_hub requires --hier_local_hub_mode shared" | |
| ) | |
| if args.num_groups <= 0 or args.num_local_hub_tokens != 0 or kg <= 0: | |
| raise ValueError( | |
| f"--hier_no_local_hub requires num_groups>0, K_L=0, and K_G>0; got N={args.num_groups}, K_L={args.num_local_hub_tokens}, K_G={kg}" | |
| ) | |
| if args.attn_impl == "hier_tiled_ste": | |
| raise ValueError("--hier_no_local_hub does not support hier_tiled_ste") | |
| elif args.num_local_hub_tokens <= 0: | |
| raise ValueError("ordinary hier_hub requires num_local_hub_tokens>0") | |
| print( | |
| "[train] 使用预压 latent 数据集（PrecomputedLatentDataset），跳过在线视频/VAE 编码。" | |
| ) | |
| dataset = PrecomputedLatentDataset( | |
| args.dataset_path, | |
| args.metadata_file_name | |
| if os.path.isabs(args.metadata_file_name) | |
| else os.path.join(args.dataset_path, args.metadata_file_name), | |
| steps_per_epoch=args.steps_per_epoch, | |
| num_players=args.num_players, | |
| target_rgb_frames=args.target_rgb_frames, | |
| blank_prompt_emb_path=args.blank_prompt_emb_path, | |
| vae_latent_channels=args.vae_latent_channels, | |
| num_frame_per_block=args.num_frame_per_block, | |
| local_attn_size=args.local_attn_size, | |
| use_camera_pose=args.use_camera_pose, | |
| use_group_labels=True, | |
| num_groups=args.num_groups, | |
| group_path_column=args.group_path_column, | |
| group_label_player_ids=args.group_label_player_ids, | |
| hier_group_mode=args.hier_group_mode, | |
| hier_group_seed=args.hier_group_seed, | |
| hier_random_preserve_sizes=bool(args.hier_random_preserve_sizes), | |
| hier_group_format=args.hier_group_format, | |
| grouping_cache_dir=args.grouping_cache_dir, | |
| grouping_model_profile=(args.num_local_hub_tokens, args.num_global_hub_tokens, | |
| args.hier_direct_frac, args.num_frame_per_block), | |
| ) | |
| dataloader = torch.utils.data.DataLoader( | |
| dataset, shuffle=True, batch_size=1, num_workers=args.dataloader_num_workers | |
| ) | |
| val_loader = None | |
| if args.val_infer: | |
| from torch.utils.data import Subset | |
| n_val = min(args.val_num_samples, len(dataset)) | |
| val_loader = torch.utils.data.DataLoader( | |
| Subset(dataset, list(range(n_val))), | |
| shuffle=False, | |
| batch_size=1, | |
| num_workers=2, | |
| ) | |
| print( | |
| f"[val] inline validation enabled: {n_val} samples, every {args.val_every_n_steps} steps." | |
| ) | |
| run_version = args.run_version or _compute_run_version(args.output_path) | |
| print( | |
| f"[run] version = {run_version}" | |
| + (" (来自 --run_version)" if args.run_version else " (扫描 max+1)") | |
| ) | |
| model = LightningModelForTrain( | |
| vae_path=args.vae_path, | |
| text_encoder_path=args.text_encoder_path, | |
| image_encoder_path=args.image_encoder_path, | |
| dit_path=args.dit_path, | |
| learning_rate=args.learning_rate, | |
| use_gradient_checkpointing=args.use_gradient_checkpointing, | |
| use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload, | |
| resume_ckpt_path=args.resume_ckpt_path, | |
| output_path=os.path.join(args.output_path, "lightning_logs"), | |
| num_players=args.num_players, | |
| vae_latent_channels=args.vae_latent_channels, | |
| perf_log_every_n_steps=args.perf_log_every_n_steps, | |
| perf_timing=not args.no_perf_timing, | |
| target_rgb_frames=args.target_rgb_frames, | |
| first_frame_dropout_p=args.first_frame_dropout_p, | |
| first_frame_noise_p=args.first_frame_noise_p, | |
| first_frame_noise_corrupt_ratio=args.first_frame_noise_corrupt_ratio, | |
| first_frame_noise_min_seg=args.first_frame_noise_min_seg, | |
| ckpt_local_dir=args.ckpt_local_dir, | |
| run_version=run_version, | |
| blank_prompt_emb_path=args.blank_prompt_emb_path, | |
| text_prompt_dropout_p=args.text_prompt_dropout_p, | |
| use_simplex_rope=True, | |
| simplex_pool_size=args.simplex_pool_size, | |
| simplex_gamma=args.simplex_gamma, | |
| simplex_train_permute=args.simplex_train_permute, | |
| simplex_train_sample_vertex_ids=args.simplex_train_sample_vertex_ids, | |
| training_stage="stage2", | |
| loss_type="diffusion_forcing", | |
| attention_mode="hier_hub", | |
| attn_impl=args.attn_impl, | |
| num_hub_tokens=args.num_hub_tokens, | |
| local_attn_size=args.local_attn_size, | |
| local_attn_size_after=args.local_attn_size_after, | |
| local_attn_switch_step=args.local_attn_switch_step, | |
| train_step_offset=args.train_step_offset, | |
| num_frame_per_block=args.num_frame_per_block, | |
| num_groups=args.num_groups, | |
| num_local_hub_tokens=args.num_local_hub_tokens, | |
| num_global_hub_tokens=args.num_global_hub_tokens, | |
| hier_no_local_hub=args.hier_no_local_hub, | |
| hier_local_hub_mode=args.hier_local_hub_mode, | |
| hier_direct_schedule=args.hier_direct_schedule, | |
| hier_direct_frac=args.hier_direct_frac, | |
| hier_direct_layers=args.hier_direct_layers, | |
| unfreeze_self_attn=args.unfreeze_self_attn or True, | |
| clean_prefix_prob=args.clean_prefix_prob, | |
| clean_prefix_context_noise=args.clean_prefix_context_noise, | |
| generated_history_prob=args.generated_history_prob, | |
| df_per_frame_noise=args.df_per_frame_noise, | |
| causal_forcing_training=args.causal_forcing_training, | |
| use_camera_pose=args.use_camera_pose, | |
| camera_dim=args.camera_dim, | |
| camera_dropout_p=args.camera_dropout_p, | |
| val_infer=args.val_infer, | |
| val_every_n_steps=args.val_every_n_steps, | |
| val_num_samples=args.val_num_samples, | |
| val_num_inference_steps=args.val_num_inference_steps, | |
| val_causal_num_frames=args.val_causal_num_frames, | |
| val_output_dir=args.val_output_dir, | |
| val_grid_fps=args.val_grid_fps, | |
| ) | |
| model._manual_val_loader = val_loader | |
| if args.use_swanlab: | |
| from swanlab.integration.pytorch_lightning import SwanLabLogger | |
| swanlab_config = {"UPPERFRAMEWORK": "DiffSynth-Studio"} | |
| swanlab_config.update(vars(args)) | |
| swanlab_logger = SwanLabLogger( | |
| project="wan", | |
| name=f"wan_{run_version}", | |
| config=swanlab_config, | |
| mode=args.swanlab_mode, | |
| logdir=os.path.join(args.output_path, "swanlog"), | |
| ) | |
| logger = [swanlab_logger] | |
| else: | |
| logger = TensorBoardLogger( | |
| save_dir=args.output_path, name="lightning_logs", version=run_version | |
| ) | |
| logger.log_hyperparams({k: str(v) for k, v in vars(args).items()}) | |
| checkpoint_callback = pl.pytorch.callbacks.ModelCheckpoint( | |
| every_n_train_steps=1000, save_top_k=2, monitor="step", mode="max" | |
| ) | |
| _val_trainer_kwargs = {} | |
| if args.val_infer: | |
| _val_trainer_kwargs = dict(use_distributed_sampler=False) | |
| trainer = pl.Trainer( | |
| max_epochs=args.max_epochs, | |
| max_steps=args.max_steps, | |
| accelerator="gpu", | |
| **_trainer_cluster_kwargs(), | |
| precision="bf16", | |
| strategy=args.training_strategy, | |
| default_root_dir=args.output_path, | |
| accumulate_grad_batches=args.accumulate_grad_batches, | |
| gradient_clip_val=args.gradient_clip_val | |
| if args.gradient_clip_val > 0 | |
| else None, | |
| callbacks=[checkpoint_callback], | |
| logger=logger, | |
| **_val_trainer_kwargs, | |
| ) | |
| if trainer.global_rank == 0: | |
| _primary_logger = logger[0] if isinstance(logger, (list, tuple)) else logger | |
| _log_dir = getattr(_primary_logger, "log_dir", None) or args.output_path | |
| checkpoints_dir = os.path.join(_log_dir, "checkpoints") | |
| os.makedirs(checkpoints_dir, exist_ok=True) | |
| print(f"Manually created checkpoint directory for rank 0: {checkpoints_dir}") | |
| version_dir = os.path.join(args.output_path, "lightning_logs", run_version) | |
| hparams_path = os.path.join(version_dir, "hparams.yaml") | |
| _write_hparams_yaml(hparams_path, args) | |
| print(f"[hparams] wrote {hparams_path}") | |
| trainer.fit(model, dataloader) | |