nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[bugfix] WanTrack: track subsampling was a no-op when max_points == N

Browse Source 
_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.
This commit is contained in:
shaoxiongduan
2026-07-03 08:18:44 +00:00
parent de8e9e20b4
commit c6d0233eab
3 changed files with 174 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/train/scenario/worldmodel/finetune_wantrack_golf_overfit.yaml
+77
View File
@@ -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
fastvideo/train/callbacks/track_validation.py
+58 -16
View File
@@ -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()
fastvideo/train/models/wantrack/wantrack.py
+39 -2
View File
@@ -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: