From c6d0233eaba990461b2d3b010fbcd5db64b3acb2 Mon Sep 17 00:00:00 2001 From: shaoxiongduan Date: Fri, 3 Jul 2026 08:18:44 +0000 Subject: [PATCH] [bugfix] WanTrack: track subsampling was a no-op when max_points == N _augment_tracks gated the object-coverage + low-rank-weighted + uniform draw on `0 < max_points < N`, so with the default WANTRACK_MAX_POINTS=2500 and the full 2500-track grid (max_points == N) the entire subsample block was skipped -- the model saw all 2500 tracks every step and only chunk-masking ran. Gate on `min_points < N` instead so K~U[1000,2500] actually samples (verified: kept 1179/1437/1573 across steps, was a constant 2500). Also: TrackValidation now applies the same sampler (loads object_ids/track_weights, runs _augment_tracks deterministically) so validation is conditioned on -- and overlays -- the sampled subset, not the full grid. Env-gated WANTRACK_DEBUG=1 logs per-step sampling stats (kept count, object coverage, coord ranges). Adds a single-video (golf) overfit config for debugging. --- .../finetune_wantrack_golf_overfit.yaml | 77 +++++++++++++++++++ fastvideo/train/callbacks/track_validation.py | 74 ++++++++++++++---- fastvideo/train/models/wantrack/wantrack.py | 41 +++++++++- 3 files changed, 174 insertions(+), 18 deletions(-) create mode 100644 examples/train/scenario/worldmodel/finetune_wantrack_golf_overfit.yaml diff --git a/examples/train/scenario/worldmodel/finetune_wantrack_golf_overfit.yaml b/examples/train/scenario/worldmodel/finetune_wantrack_golf_overfit.yaml new file mode 100644 index 000000000..96fbb004b --- /dev/null +++ b/examples/train/scenario/worldmodel/finetune_wantrack_golf_overfit.yaml @@ -0,0 +1,77 @@ +# WanTrack Wan-200 finetune: SAME params as the 50-video (control50) run, scaled to +# 200 Wan2.2-A14B clips. Fun-Control init, WANTRACK_AUG=0 (launch env), unique prompts +# (standard funinp preprocess). Isolates the data-scale variable (50 -> 200). + +models: + student: + _target_: fastvideo.train.models.wantrack.wantrack.WanTrackModel + init_from: /mnt/weka/home/hao.zhang/shao/data/models/trackwan_1.3b_i2v_control_init + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod + +training: + model_path: /mnt/weka/home/hao.zhang/shao/data/models/trackwan_1.3b_i2v_control_init + + distributed: + num_gpus: 1 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 1 + + data: + data_path: /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/wan_golf_overfit/preprocessed_i2v_track/combined_parquet_dataset + dataloader_num_workers: 2 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 31 + num_height: 480 + num_width: 832 + num_frames: 121 + + optimizer: + learning_rate: 1.0e-4 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 3000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: /mnt/weka/home/hao.zhang/shao/data/motion_pipeline/wantrack_golf_overfit_out + training_state_checkpointing_steps: 1000 + checkpoints_total_limit: 2 + resume_from_checkpoint: latest + + tracker: + trackers: [wandb] + project_name: wantrack-bidir + run_name: golf-overfit-debug + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + track_validation: + _target_: fastvideo.train.callbacks.track_validation.TrackValidationCallback + every_steps: 200 + num_val_samples: 1 + num_inference_steps: 30 + guidance_scale: 1.0 + fps: 24 + grid_stride: 3 + tail: 12 + validate_at_start: true + seed: 1000 + +pipeline: + flow_shift: 3 diff --git a/fastvideo/train/callbacks/track_validation.py b/fastvideo/train/callbacks/track_validation.py index 810a325e5..25da08bdb 100644 --- a/fastvideo/train/callbacks/track_validation.py +++ b/fastvideo/train/callbacks/track_validation.py @@ -194,6 +194,15 @@ class TrackValidationCallback(Callback): n = batch["text_embedding"].shape[0] infos = batch.get("info_list") or [{} for _ in range(n)] + + def _opt(key: str, i: int) -> Any: + # optional per-sample tensor (object_ids / track_weights); None if absent/empty + t = batch.get(key) + if t is None or not torch.is_tensor(t) or t.numel() == 0 or t.shape[0] <= i: + return None + row = t[i:i + 1].clone() + return row if row.numel() > 0 else None + for i in range(n): self._samples.append({ "text_embedding": batch["text_embedding"][i:i + 1].clone(), @@ -203,6 +212,8 @@ class TrackValidationCallback(Callback): "clip_feature": batch["clip_feature"][i:i + 1].clone(), "track_points": batch["track_points"][i:i + 1].clone(), "track_visibility": batch["track_visibility"][i:i + 1].clone(), + "object_ids": _opt("object_ids", i), + "track_weights": _opt("track_weights", i), "caption": str(infos[i].get("caption", "") if i < len(infos) else ""), }) @@ -222,9 +233,12 @@ class TrackValidationCallback(Callback): gen_logs: list[Any] = [] ref_logs: list[Any] = [] for i, s in enumerate(self._samples): - gen_latents = self._sample(student, transformer, s) # [1,16,T,H,W] normalized + # Apply the SAME training-time sampler so the model sees the sampled subset + # (not the full 2500 grid) and the overlay shows those sampled points. + tp_s, tv_s = self._sampled_tracks(student, s) + gen_latents = self._sample(student, transformer, s, tp_s, tv_s) # [1,16,T,H,W] normalized gen_px = student.decode_latents(gen_latents.permute(0, 2, 1, 3, 4))[0] # [3,T,H,W] in [0,1] - gen_frames = self._overlay_tracks(gen_px, s) + gen_frames = self._overlay_tracks(gen_px, s, tp_s, tv_s) if self._is_main: fn = os.path.join(out_dir, f"step{step:06d}_sample{i}_gen.mp4") imageio.mimsave(fn, gen_frames, fps=self.fps, macro_block_size=1) @@ -232,10 +246,10 @@ class TrackValidationCallback(Callback): if art is not None: gen_logs.append(art) - # One-time ground-truth reference (VAE round-trip + same tracks). + # One-time ground-truth reference (VAE round-trip + same sampled tracks). if not self._ref_logged and self._is_main: ref_px = self._decode_reference(student, s) - ref_frames = self._overlay_tracks(ref_px, s) + ref_frames = self._overlay_tracks(ref_px, s, tp_s, tv_s) fn_ref = os.path.join(out_dir, f"reference_sample{i}_gt.mp4") imageio.mimsave(fn_ref, ref_frames, fps=self.fps, macro_block_size=1) art = self.tracker.video(fn_ref, caption=f"GT {s['caption'][:120]}") @@ -253,7 +267,26 @@ class TrackValidationCallback(Callback): if was_training: transformer.train() - def _sample(self, student: Any, transformer: torch.nn.Module, s: dict[str, Any]) -> torch.Tensor: + def _sampled_tracks(self, student: Any, s: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + """Run the training-time track sampler (_augment_tracks) on this val sample, deterministically + (fixed seed), so the model is conditioned on the sampled subset exactly like training. Falls + back to the full tracks if the sampler is disabled (WANTRACK_AUG=0) or drops motion.""" + device = student.device + tp = s["track_points"].to(device) + tv = s["track_visibility"].to(device) + n = tp.shape[2] + oid = s.get("object_ids") + oid = oid.to(device).long() if (oid is not None and oid.shape[-1] == n) else None + tw = s.get("track_weights") + tw = tw.to(device).float() if (tw is not None and tw.shape[-1] == n) else None + gen = torch.Generator(device=device).manual_seed(self.seed) + tp_s, tv_s = student._augment_tracks(tp.float(), tv.float(), gen, object_ids=oid, track_weights=tw) + if tp_s is None or tv_s is None: # motion-drop (off in val) -> keep full tracks + return tp, tv + return tp_s, tv_s + + def _sample(self, student: Any, transformer: torch.nn.Module, s: dict[str, Any], tp_in: torch.Tensor, + tv_in: torch.Tensor) -> torch.Tensor: device = student.device dtype = torch.bfloat16 flow_shift = float(student.timestep_shift) @@ -262,8 +295,8 @@ class TrackValidationCallback(Callback): cond20 = student._build_i2v_cond_concat(ff) # [1,20,T,H,W] txt = s["text_embedding"].to(device, dtype) mask = s["text_attention_mask"].to(device, dtype) - tp = s["track_points"].to(device, dtype) # [1,T,N,2] normalized - tv = s["track_visibility"].to(device, dtype) # [1,T,N] + tp = tp_in.to(device, dtype) # [1,T,N,2] normalized (sampled subset) + tv = tv_in.to(device, dtype) # [1,T,N] (sampled visibility) img = s["clip_feature"].to(device, dtype) # [1,SeqLen,Dim] CLIP frame-0 _, _, T, H, W = ff.shape @@ -298,22 +331,31 @@ class TrackValidationCallback(Callback): norm = normalize_dit_input("wan", raw, student.vae) return student.decode_latents(norm.permute(0, 2, 1, 3, 4))[0] # [3,T,H,W] - def _overlay_tracks(self, px: torch.Tensor, s: dict[str, Any]) -> list[np.ndarray]: - # px: [3,T,H,W] in [0,1] + def _overlay_tracks(self, px: torch.Tensor, s: dict[str, Any], tp_in: torch.Tensor, + tv_in: torch.Tensor) -> list[np.ndarray]: + # px: [3,T,H,W] in [0,1]. Draws ONLY the sampled tracks (those the sampler kept), gated by + # the sampled visibility, so the overlay matches what the model was conditioned on. video = (px.clamp(0, 1).float().cpu().numpy() * 255.0).astype(np.uint8) frames = np.transpose(video, (1, 2, 3, 0)) # [T,H,W,3] T, H, W, _ = frames.shape - tp = s["track_points"][0].float().cpu().numpy() # [Tt,N,2] normalized - tv = s["track_visibility"][0].float().cpu().numpy() # [Tt,N] + tp = tp_in[0].float().cpu().numpy() # [Tt,N,2] normalized + tv = tv_in[0].float().cpu().numpy() # [Tt,N] sampled visibility tt = min(T, tp.shape[0]) frames, tp, tv = frames[:tt], tp[:tt], tv[:tt] + n = tp.shape[1] + grid = int(round(n**0.5)) - grid = int(round(tp.shape[1]**0.5)) - tp, tv = _subsample(tp, tv, grid, self.grid_stride) - colors = _grid_colors(grid, self.grid_stride) - if colors.shape[0] != tp.shape[1]: - colors = _grid_colors(int(round(tp.shape[1]**0.5)) or 1, 1)[:tp.shape[1]] + colors = _grid_colors(grid, 1) # one color per grid cell (row-major) + if colors.shape[0] < n: + colors = np.resize(colors, (n, 3)) + colors = colors[:n] + + # keep only tracks that are ever visible after sampling (the sampled subset) + ever = tv.max(axis=0) >= self.vis_thresh + if not ever.any(): + ever = np.ones(n, dtype=bool) + tp, tv, colors = tp[:, ever], tv[:, ever], colors[ever] # Normalized [0,1] -> pixel coords of the generated frame. tp_px = tp.copy() diff --git a/fastvideo/train/models/wantrack/wantrack.py b/fastvideo/train/models/wantrack/wantrack.py index 33fe5e400..8ee9f9113 100644 --- a/fastvideo/train/models/wantrack/wantrack.py +++ b/fastvideo/train/models/wantrack/wantrack.py @@ -22,13 +22,15 @@ import torch from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track from fastvideo.distributed import get_sp_group, get_world_group +from fastvideo.logger import init_logger from fastvideo.pipelines import TrainingBatch from fastvideo.training.training_utils import normalize_dit_input - from fastvideo.train.models.wan.wan import WanModel from fastvideo.train.utils.dataloader import build_parquet_t2v_train_dataloader from fastvideo.train.utils.moduleloader import load_module_from_path +logger = init_logger(__name__) + class WanTrackModel(WanModel): """Wan2.2 + point-track I2V model for finetuning in the new trainer.""" @@ -89,6 +91,10 @@ class WanTrackModel(WanModel): # CFG dropout (ours, not MotionStream) -> default OFF; enables motion/text CFG when >0. self._aug_motion_drop = float(os.getenv("WANTRACK_MOTION_DROP", "0.0")) self._aug_text_drop = float(os.getenv("WANTRACK_TEXT_DROP", "0.0")) + # TEMP debug: WANTRACK_DEBUG=1 logs per-step sampling stats (throttled) so we can verify + # the sampler keeps the right number of tracks / covers objects / uses the weights. + self._aug_debug = int(os.getenv("WANTRACK_DEBUG", "0")) + self._aug_calls = 0 def _augment_tracks( self, @@ -126,7 +132,10 @@ class WanTrackModel(WanModel): disp = ((tp - tp[:, :1])**2).sum(-1).sqrt() # [B,T,N] motion = (disp * (track_visibility > 0.5).float()).amax(dim=1) # [B,N] - if 0 < self._aug_max_points < N: + # Subsample whenever there's room to draw fewer than N (min_points < N). The old + # `max_points < N` gate silently disabled ALL sampling when max_points == N (== 2500), + # so coverage/weighted/uniform never ran. max_points == 0 still means "keep all". + if self._aug_max_points > 0 and self._aug_min_points < N: oids = object_ids if oids is not None and oids.dim() == 1: oids = oids.unsqueeze(0).expand(B, -1) @@ -184,6 +193,25 @@ class WanTrackModel(WanModel): frame_keep[b, c * L:(c + 1) * L] = 0.0 vis = vis * frame_keep.to(device=device, dtype=vis.dtype).unsqueeze(-1) + # TEMP debug: verify the sampling is doing what we think (throttled). + if self._aug_debug and (self._aug_calls < 8 or self._aug_calls % 200 == 0): + ever0 = (vis[0].amax(dim=0) > 0.5) # [N] tracks kept (any visible frame after sampling) + kept = int(ever0.sum().item()) + cov = "n/a" + if object_ids is not None: + ob = object_ids[0] if object_ids.dim() > 1 else object_ids + present = int(torch.unique(ob[ob >= 0]).numel()) + kept_obj = int(torch.unique(ob[(ever0) & (ob >= 0)]).numel()) + cov = f"{kept_obj}/{present}" + tp0 = track_points[0] + logger.info( + "[WANTRACK_DEBUG] aug call=%d B=%d N=%d sample_range=[%d,%d] kept=%d obj_cov=%s " + "tw=%s oid=%s pmask=%.2f | coords x[%.3f,%.3f] y[%.3f,%.3f] vis_sum(f0)=%.0f", self._aug_calls, B, N, + self._aug_min_points, self._aug_max_points, kept, cov, "yes" if track_weights is not None else "NO", + "yes" if object_ids is not None else "NO", self._aug_pmask, float(tp0[..., 0].min()), + float(tp0[..., 0].max()), float(tp0[..., 1].min()), float(tp0[..., 1].max()), float(vis[0, 0].sum())) + self._aug_calls += 1 + # Motion CFG dropout (ours; default off): drop ALL motion for this step (-> cm = zeros). if (self._aug_motion_drop > 0.0 and float(torch.rand(1, generator=generator, device=gdev).item()) < self._aug_motion_drop): @@ -261,6 +289,15 @@ class WanTrackModel(WanModel): object_ids=object_ids, track_weights=track_weights) + # TEMP debug: confirm what actually flows into the forward (tracks not silently dropped). + if self._aug_debug and self._aug_calls <= 8: + logger.info( + "[WANTRACK_DEBUG] prepare_batch: latents=%s track_points=%s track_vis=%s object_ids=%s track_weights=%s", + tuple(latents.shape), (tuple(track_points.shape) if track_points is not None else None), + (tuple(track_visibility.shape) if track_visibility is not None else None), + (tuple(object_ids.shape) if object_ids is not None else None), + (tuple(track_weights.shape) if track_weights is not None else None)) + encoder_hidden_states = encoder_hidden_states.to(device, dtype=dtype) # Text CFG dropout: zero the text embedding for this step with prob. if self._aug_enabled and self._aug_text_drop > 0.0: