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

trackwan: 14B teacher pipeline, layer-wise track LR, and overfit tooling

Browse Source 
Training:
- finetune.py: optional track/base param groups (WANTRACK_TRACK_GROUP) so the
  track pathway can take a different LR (WANTRACK_TRACK_LR_MULT); patch_embedding
  gradient masked to the track slot [:, 36:]. Raises if nothing matches.
- track_warmup.py callback: hold the pretrained DiT at lr 0 while the track
  pathway warms up, then ramp it back in.

14B configs + held-alloc launchers:
- finetune_wantrack_{synth_sparse_fixed,synth_sparse_random,openvid_sparse_merged_bias,
  synth_stage2}_14b + openvid_14b_{warmup,lwlr} + 1.3B stage2 configs.
- run_wan14b_held.sh (held alloc + auto-restart + rank cleanup, per-stage WANTRACK
  env), acquire_nodes.sh, chain_stepAB.sh, run_stepB_seed.sh, run_stepC_merge.sh,
  watch_gate.sh, run_synth_stage2_slurm.sh.

Init / diagnostics:
- convert_trackwan_init_v2.py: --pe-src to lift patch_embedding[:, 36:] and
  --track-src to lift track_encoder together (co-adapted merge); sharded-source support.
- check_gate_growth.py, gradient_flow_analysis.py, gradient_flow_summarize.py.

Data pipeline: segment_tracks.py sharding; decode/segment/tracks slurm scripts;
precompute_null_text, two_track_val, cfg_ablation, lwlr_experiment.

Inference: app_multi_trace.py updates.
This commit is contained in:
shaoxiongduan
2026-07-23 03:52:56 +00:00
parent ac919b1c87
commit 92cc893cfc
31 changed files with 3458 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/inference/gradio/trackwan/app_multi_trace.py
+124 -174
View File
@@ -1,17 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
"""TrackWan multi-trace demo — pick a preprocessed openvid clip, draw multiple traces
with mouse (Space to record), generate.
"""TrackWan interactive rollout demo.
Two point-input modes on the same canvas:
- **Trace mode** — hover, press *Space* to record a 5-second (121-frame) drag.
- **Anchor mode** — click once to drop a static point (a track whose (x,y) is
constant across all 121 frames).
Design (matches training-time validation exactly):
- Loads conditioning (text_embedding, clip_feature, first_frame_latent) DIRECTLY from
the preprocessed openvid parquet — no on-the-fly encoding. Same path as
the preprocessed openvid parquet. Same path as
``fastvideo/train/callbacks/track_validation.py``.
- First frame image shown to the user is decoded from the VAE latent (``decode_reference``).
- Traces: press Space to record a 5s (121-frame) trajectory. Multi-trace stacks.
- Background points: sampled from SAM foreground-avoiding pixels of frame-0. Click a dot
to delete just it; buttons for resample / clear all.
- Generation: uses the sample's precomputed conditioning + user's track_points, then runs
the SAME motion-CFG denoise loop as ``TrackValidationCallback._sample``.
- Both traces and anchors are packed into a single ``track_points[T,N,2]`` tensor.
Anchors get their (x,y) broadcast across T=121; traces are linearly interpolated.
- Generation: the SAME motion-CFG denoise loop as ``TrackValidationCallback._sample``
(MotionStream Eq. 2, wt=3.0, wm=1.5 by default).
Launch::
@@ -42,53 +45,10 @@ import numpy as np
from PIL import Image
REPO = Path(__file__).resolve().parents[4]
SAM_WEIGHT_CANDIDATES = [
"/mnt/lustre/vlm-s4duan/models/seg/sam2.1_b.pt",
"/mnt/lustre/vlm-s4duan/FastVideo/sam2.1_b.pt",
"/mnt/lustre/vlm-s4duan/models/seg/sam2_b.pt",
]
NUM_FRAMES = 121
FPS = 24
# =============================================================================
# SAM
# =============================================================================
def load_sam(weight: str | None = None):
from ultralytics import SAM
if weight is None:
for w in SAM_WEIGHT_CANDIDATES:
if os.path.exists(w):
weight = w
break
if weight is None:
raise FileNotFoundError(f"No SAM weight found in {SAM_WEIGHT_CANDIDATES}")
return SAM(weight)
def sam_foreground_mask(sam, img: np.ndarray, imgsz: int = 640) -> np.ndarray:
res = sam(img, imgsz=imgsz, conf=0.4, iou=0.9, retina_masks=True, verbose=False)
if not res or res[0].masks is None:
return np.zeros(img.shape[:2], bool)
m = res[0].masks.data.cpu().numpy().astype(bool)
if m.shape[0] and m.shape[1:] != img.shape[:2]:
import torch
t = torch.from_numpy(m.astype(np.uint8))[None].float()
t = torch.nn.functional.interpolate(t, size=img.shape[:2], mode="nearest")[0]
m = t.numpy().astype(bool)
return m.any(0) if m.shape[0] > 0 else np.zeros(img.shape[:2], bool)
def sample_bg_points_normalized(fg_mask: np.ndarray, n: int, seed: int = 0) -> list:
ys, xs = np.where(~fg_mask)
if len(xs) == 0:
return []
rng = np.random.default_rng(seed)
idx = rng.choice(len(xs), size=min(n, len(xs)), replace=False)
H, W = fg_mask.shape
return [[float(xs[i]) / W, float(ys[i]) / H] for i in idx]
# =============================================================================
# frame b64 + track construction
# =============================================================================
@@ -98,15 +58,23 @@ def img_to_b64(arr: np.ndarray) -> str:
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
def build_track_points(traces_json: str, bg_pts: list) -> tuple[np.ndarray, np.ndarray, int, int]:
"""List-of-normalized-trajectories → (track_points[T,N,2], track_visibility[T,N], n_tr, n_bg)."""
def build_track_points(traces_json: str, anchors_json: str) -> tuple[np.ndarray, np.ndarray, int, int]:
"""(traces, anchors) → (track_points[T,N,2], track_visibility[T,N], n_tr, n_an).
Each trace is a list of normalized (x,y) samples at ~24 Hz over 5 s → interpolated to T=121.
Each anchor is a single normalized (x,y) → broadcast across all T frames.
"""
T = NUM_FRAMES
try:
traces = json.loads(traces_json) if traces_json else []
except Exception:
traces = []
n_tr, n_bg = len(traces), len(bg_pts)
N = n_tr + n_bg
try:
anchors = json.loads(anchors_json) if anchors_json else []
except Exception:
anchors = []
n_tr, n_an = len(traces), len(anchors)
N = n_tr + n_an
if N == 0:
return np.zeros((T, 0, 2), np.float32), np.zeros((T, 0), np.float32), 0, 0
tp = np.zeros((T, N, 2), np.float32)
@@ -122,14 +90,14 @@ def build_track_points(traces_json: str, bg_pts: list) -> tuple[np.ndarray, np.n
tgt = np.linspace(0, arr.shape[0] - 1, T)
tp[:, i, 0] = np.interp(tgt, src, arr[:, 0])
tp[:, i, 1] = np.interp(tgt, src, arr[:, 1])
for j, (bx, by) in enumerate(bg_pts):
tp[:, n_tr + j, 0] = bx
tp[:, n_tr + j, 1] = by
return tp, tv, n_tr, n_bg
for j, (ax, ay) in enumerate(anchors):
tp[:, n_tr + j, 0] = ax
tp[:, n_tr + j, 1] = ay
return tp, tv, n_tr, n_an
# =============================================================================
# canvas JS (multi-trace + click-to-delete bg points)
# canvas JS (traces via Space + anchors via click, mode radio decides)
# =============================================================================
def canvas_js(num_frames: int, fps: int) -> str:
palette = ["#ff3c3c", "#3cd23c", "#3c78ff", "#ffc83c", "#c83cc8", "#3cdcdc", "#ff821e", "#9696ff"]
@@ -138,32 +106,40 @@ def canvas_js(num_frames: int, fps: int) -> str:
" const ctx = cvs.getContext('2d');\n"
f" const PAL = {json.dumps(palette)};\n"
f" const N = {num_frames}; const DT = {int(1000 / fps)};\n"
" const st = { img:new Image(), mouse:[0.5,0.5], rec:false, current:[], traces:[], bg:[] };\n"
" const st = { img:new Image(), mouse:[0.5,0.5], rec:false, current:[], "
"traces:[], anchors:[], mode:'trace' };\n"
" window._st = st;\n"
" cvs.width = 832; cvs.height = 480;\n"
" function commit_traces(){\n"
" const tb = document.querySelector('#traces_json textarea');\n"
" if (tb) { tb.value = JSON.stringify(st.traces); tb.dispatchEvent(new Event('input',{bubbles:true})); }\n"
" }\n"
" function commit_bg(){\n"
" const tb = document.querySelector('#bg_pts_json textarea');\n"
" if (tb) { tb.value = JSON.stringify(st.bg); tb.dispatchEvent(new Event('input',{bubbles:true})); }\n"
" function commit_anchors(){\n"
" const tb = document.querySelector('#anchors_json textarea');\n"
" if (tb) { tb.value = JSON.stringify(st.anchors); tb.dispatchEvent(new Event('input',{bubbles:true})); }\n"
" }\n"
" function draw(){\n"
" ctx.clearRect(0,0,cvs.width,cvs.height);\n"
" if (st.img.complete && st.img.width) ctx.drawImage(st.img, 0, 0, cvs.width, cvs.height);\n"
" st.bg.forEach(p => { const x=p[0]*cvs.width, y=p[1]*cvs.height;\n"
" ctx.beginPath(); ctx.arc(x,y,3,0,7); ctx.fillStyle='#c8c8c8'; ctx.fill();\n"
" ctx.strokeStyle='#000'; ctx.stroke(); });\n"
" // anchors: solid gray dots with dark ring\n"
" st.anchors.forEach((p, i) => { const x=p[0]*cvs.width, y=p[1]*cvs.height;\n"
" ctx.beginPath(); ctx.arc(x,y,7,0,7); ctx.fillStyle='#888'; ctx.fill();\n"
" ctx.strokeStyle='#000'; ctx.lineWidth=2; ctx.stroke();\n"
" ctx.fillStyle='#fff'; ctx.font='bold 10px sans-serif'; ctx.textAlign='center'; ctx.textBaseline='middle';\n"
" ctx.fillText('A'+(i+1), x, y);\n"
" });\n"
" // traces: colored polyline with start/end markers\n"
" st.traces.forEach((traj, i) => {\n"
" const color = PAL[i % PAL.length];\n"
" ctx.strokeStyle=color; ctx.lineWidth=3; ctx.beginPath();\n"
" traj.forEach((p, j) => { const x=p[0]*cvs.width, y=p[1]*cvs.height; j?ctx.lineTo(x,y):ctx.moveTo(x,y); });\n"
" ctx.stroke();\n"
" const s = traj[0], e = traj[traj.length-1];\n"
" ctx.fillStyle=color; ctx.beginPath(); ctx.arc(s[0]*cvs.width, s[1]*cvs.height, 6, 0, 7); ctx.fill(); ctx.strokeStyle='#000'; ctx.stroke();\n"
" ctx.fillStyle=color; ctx.beginPath(); ctx.arc(s[0]*cvs.width, s[1]*cvs.height, 6, 0, 7); ctx.fill(); ctx.strokeStyle='#000'; ctx.lineWidth=1; ctx.stroke();\n"
" ctx.fillStyle='#ff3333'; ctx.beginPath(); ctx.arc(e[0]*cvs.width, e[1]*cvs.height, 6, 0, 7); ctx.fill(); ctx.strokeStyle='#000'; ctx.stroke();\n"
" ctx.fillStyle='#000'; ctx.font='13px sans-serif'; ctx.fillText('#'+(i+1), s[0]*cvs.width+8, s[1]*cvs.height-8);\n"
" ctx.fillStyle='#000'; ctx.font='13px sans-serif'; ctx.textAlign='left'; ctx.textBaseline='alphabetic'; ctx.fillText('T'+(i+1), s[0]*cvs.width+8, s[1]*cvs.height-8);\n"
" });\n"
" // live-drawing current trace\n"
" if (st.rec && st.current.length) {\n"
" const color = PAL[st.traces.length % PAL.length];\n"
" ctx.strokeStyle=color; ctx.lineWidth=3; ctx.beginPath();\n"
@@ -180,28 +156,34 @@ def canvas_js(num_frames: int, fps: int) -> str:
" const r = cvs.getBoundingClientRect();\n"
" const cx = (e.clientX - r.left) * (cvs.width / r.width);\n"
" const cy = (e.clientY - r.top) * (cvs.height / r.height);\n"
" let idx = -1, best = 12*12;\n"
" st.bg.forEach((p, i) => { const dx = p[0]*cvs.width - cx, dy = p[1]*cvs.height - cy;\n"
" const d2 = dx*dx + dy*dy; if (d2 < best) { best = d2; idx = i; } });\n"
" if (idx >= 0) { st.bg.splice(idx, 1); commit_bg(); draw(); }\n"
" if (st.mode === 'anchor') {\n"
" // toggle delete-if-near, else place\n"
" let idx = -1, best = 14*14;\n"
" st.anchors.forEach((p, i) => { const dx = p[0]*cvs.width - cx, dy = p[1]*cvs.height - cy;\n"
" const d2 = dx*dx + dy*dy; if (d2 < best) { best = d2; idx = i; } });\n"
" if (idx >= 0) st.anchors.splice(idx, 1);\n"
" else st.anchors.push([cx / cvs.width, cy / cvs.height]);\n"
" commit_anchors(); draw();\n"
" }\n"
" });\n"
" window._setmode = (m) => { st.mode = m || 'trace'; };\n"
" window._setbg = (b64) => {\n"
" if (!b64) return;\n"
" const im = new Image();\n"
" im.onload = () => { const ar = im.width / im.height; cvs.width = 832; cvs.height = Math.round(832 / ar); st.img = im; draw(); };\n"
" im.src = b64;\n"
" };\n"
" window._setbg_pts = (j) => { try { st.bg = JSON.parse(j) || []; } catch(e) { st.bg = []; } draw(); };\n"
" window._clear = () => {\n"
" window._clear_traces = () => {\n"
" st.traces = []; st.current = []; st.rec = false; commit_traces(); draw();\n"
" const stt = document.getElementById('recstatus'); if (stt) stt.textContent = 'cleared. Space for trace #1';\n"
" const stt = document.getElementById('recstatus'); if (stt) stt.textContent = 'traces cleared. Space for trace #1';\n"
" };\n"
" window._undo = () => {\n"
" window._undo_trace = () => {\n"
" if (st.traces.length) st.traces.pop();\n"
" commit_traces(); draw();\n"
" const stt = document.getElementById('recstatus'); if (stt) stt.textContent = st.traces.length + ' committed. Space for another';\n"
" const stt = document.getElementById('recstatus'); if (stt) stt.textContent = st.traces.length + ' trace(s). Space for another';\n"
" };\n"
" window._clearbg = () => { st.bg = []; commit_bg(); draw(); };\n"
" window._clear_anchors = () => { st.anchors = []; commit_anchors(); draw(); };\n"
" window._undo_anchor = () => { if (st.anchors.length) st.anchors.pop(); commit_anchors(); draw(); };\n"
" window._startRec = () => {\n"
" if (st.rec) return;\n"
" st.rec = true; st.current = []; let n = 0;\n"
@@ -221,6 +203,7 @@ def canvas_js(num_frames: int, fps: int) -> str:
" if (e.code !== 'Space' || e.repeat) return;\n"
" const tag = document.activeElement && document.activeElement.tagName;\n"
" if (tag === 'INPUT' || tag === 'TEXTAREA') return;\n"
" if (st.mode !== 'trace') return; // Space only records in trace mode\n"
" e.preventDefault(); window._startRec();\n"
" });\n"
" draw();\n"
@@ -234,35 +217,42 @@ def build_ui(state):
canvas_html = ("<div style='user-select:none'>"
"<canvas id='cvs' style='border:1px solid #888;cursor:crosshair;max-width:100%'></canvas>"
"<div id='recstatus' style='font-weight:bold;color:#0a0;height:1.4em;margin-top:6px'>"
"pick a clip, click <b>Load frame</b>, then hover and press <b>Space</b> to record trace #1</div>"
"pick a clip, click <b>Load frame</b>, then choose a mode below</div>"
"</div>")
labels = state["labels"]
with gr.Blocks(title="TrackWan multi-trace demo") as demo:
with gr.Blocks(title="TrackWan interactive rollout") as demo:
gr.Markdown(
"## TrackWan — multi-trace generation (parquet-conditioned)\n"
"1. Pick a preprocessed openvid clip → **Load frame**.\n"
"2. Hover over the frame and press **Space** to record a 5-second trace.\n"
"3. Repeat Space for more traces. Optionally include background points. **Generate**.")
"## TrackWan — interactive rollout demo\n"
"1. Pick a preprocessed clip → **Load frame**.\n"
"2. Choose a mode:\n"
" - **Trace** — hover on the frame and press **Space** to record a 5 s drag.\n"
" - **Anchor** — click on the frame to drop a *static* point. Click an existing anchor to remove it.\n"
"3. **Generate**. Denoise formula matches training's validation exactly (MotionStream Eq. 2).")
with gr.Row():
with gr.Column(scale=1):
clip = gr.Dropdown(labels, value=labels[0] if labels else None, label="Preprocessed clip")
caption = gr.Textbox(label="Caption (from parquet)", interactive=False, lines=2)
load_btn = gr.Button("Load frame (runs SAM)", variant="primary")
sam_on = gr.Checkbox(True, label="Run SAM (needed for background points)")
gr.Markdown("**Background points (SAM foreground-avoiding)**")
bg_on = gr.Checkbox(True, label="Include background points")
bg_n = gr.Slider(0, 50, 20, step=1, label="Num background points")
with gr.Row():
resample_bg_btn = gr.Button("Resample")
clear_bg_btn = gr.Button("Clear all bg pts")
gr.Markdown("_Tip: **click a bg dot** on the frame to delete just that one._")
gr.Markdown("**Trace drawing**")
load_btn = gr.Button("Load frame", variant="primary")
gr.Markdown("**Input mode**")
mode = gr.Radio(["trace", "anchor"], value="trace",
label="Mode",
info="Trace = Space to record a moving drag. Anchor = single click for a static point.")
gr.Markdown("**Traces** (moving)")
with gr.Row():
rec_btn = gr.Button("● Record (Space)")
undo_btn = gr.Button("Undo")
clear_btn = gr.Button("Clear")
trace_info = gr.Markdown("_no traces yet._")
undo_trace_btn = gr.Button("Undo trace")
clear_trace_btn = gr.Button("Clear traces")
gr.Markdown("**Anchors** (static)")
with gr.Row():
undo_anchor_btn = gr.Button("Undo anchor")
clear_anchor_btn = gr.Button("Clear anchors")
info_md = gr.Markdown("_no traces or anchors yet._")
gr.Markdown("**Generation** (denoise formula matches training's validation exactly)")
seed_in = gr.Slider(0, 9999, 1000, step=1, label="Seed")
steps_in = gr.Slider(10, 60, 30, step=1, label="Denoise steps")
@@ -278,13 +268,12 @@ def build_ui(state):
# hidden bridges
ff_b64 = gr.Textbox(elem_id="ff_b64", visible=False)
bg_pts_json = gr.Textbox(elem_id="bg_pts_json", visible=False, value="[]")
traces_json = gr.Textbox(elem_id="traces_json", visible=False, value="[]")
anchors_json = gr.Textbox(elem_id="anchors_json", visible=False, value="[]")
# server-side per-page state
st_clip_idx = gr.State(0)
st_frame = gr.State(None)
st_fg = gr.State(None)
# ---------------- caption on clip change ----------------
def _cap(clip_name):
@@ -292,35 +281,18 @@ def build_ui(state):
return i, state["samples"][i]["caption"][:400]
clip.change(_cap, [clip], [st_clip_idx, caption])
# ---------------- load frame (decode ref frame from VAE latent + SAM) ----------------
sam_cache: dict[int, np.ndarray] = {}
# ---------------- load frame (decode ref frame from VAE latent) ----------------
first_frame_cache: dict[int, np.ndarray] = {}
def _load(clip_name, sam_on_, bg_n_, bg_on_):
def _load(clip_name):
i = state["label_to_idx"].get(clip_name, 0)
if i not in first_frame_cache:
first_frame_cache[i] = state["decode_first_frame"](i)
frame = first_frame_cache[i]
if sam_on_:
if i in sam_cache:
fg = sam_cache[i]
msg = f"_clip loaded ({state['samples'][i]['caption'][:60]}...) — SAM cached, {int(fg.sum())} fg px_"
else:
t0 = time.time()
try:
fg = sam_foreground_mask(state["sam"], frame)
sam_cache[i] = fg
except Exception as e:
fg = np.zeros(frame.shape[:2], bool)
msg = f"_clip loaded — SAM {time.time()-t0:.1f}s, {int(fg.sum())} fg px_"
else:
fg = np.zeros(frame.shape[:2], bool)
msg = "_clip loaded (SAM skipped)_"
bg = sample_bg_points_normalized(fg, int(bg_n_)) if (bg_on_ and fg.any()) else []
return i, frame, fg, img_to_b64(frame), json.dumps(bg), msg
msg = f"_clip loaded ({state['samples'][i]['caption'][:60]}...)_"
return i, frame, img_to_b64(frame), msg
load_btn.click(_load, [clip, sam_on, bg_n, bg_on],
[st_clip_idx, st_frame, st_fg, ff_b64, bg_pts_json, status])
load_btn.click(_load, [clip], [st_clip_idx, st_frame, ff_b64, status])
# preload first clip on startup
def _preload():
@@ -328,58 +300,46 @@ def build_ui(state):
if i not in first_frame_cache:
first_frame_cache[i] = state["decode_first_frame"](i)
frame = first_frame_cache[i]
fg = np.zeros(frame.shape[:2], bool)
return (i, state["samples"][i]["caption"][:400], frame, fg, img_to_b64(frame),
"[]", "_default clip preloaded (SAM off). Click **Load frame** to enable bg points._")
demo.load(_preload, None,
[st_clip_idx, caption, st_frame, st_fg, ff_b64, bg_pts_json, status])
# ---------------- bg controls ----------------
state.setdefault("bg_seed", [0])
def _update_bg(fg, on_, n_, fresh: bool = False):
if fg is None or not on_ or not fg.any():
return "[]"
if fresh:
state["bg_seed"][0] += 1
return json.dumps(sample_bg_points_normalized(fg, int(n_), seed=state["bg_seed"][0]))
for w in (bg_on, bg_n):
w.change(_update_bg, [st_fg, bg_on, bg_n], bg_pts_json)
resample_bg_btn.click(lambda fg, on_, n_: _update_bg(fg, on_, n_, fresh=True),
[st_fg, bg_on, bg_n], bg_pts_json)
clear_bg_btn.click(None, None, None, js="() => window._clearbg && window._clearbg()")
return (i, state["samples"][i]["caption"][:400], frame, img_to_b64(frame),
"_default clip preloaded. Draw a trace (Space) or place anchors (click)._")
demo.load(_preload, None, [st_clip_idx, caption, st_frame, ff_b64, status])
# ---------------- JS bridges ----------------
ff_b64.change(None, ff_b64, None, js="(b) => window._setbg && window._setbg(b)")
bg_pts_json.change(None, bg_pts_json, None, js="(j) => window._setbg_pts && window._setbg_pts(j)")
mode.change(None, mode, None, js="(m) => window._setmode && window._setmode(m)")
rec_btn.click(None, None, None, js="() => window._startRec && window._startRec()")
undo_btn.click(None, None, None, js="() => window._undo && window._undo()")
clear_btn.click(None, None, None, js="() => window._clear && window._clear()")
undo_trace_btn.click(None, None, None, js="() => window._undo_trace && window._undo_trace()")
clear_trace_btn.click(None, None, None, js="() => window._clear_traces && window._clear_traces()")
undo_anchor_btn.click(None, None, None, js="() => window._undo_anchor && window._undo_anchor()")
clear_anchor_btn.click(None, None, None, js="() => window._clear_anchors && window._clear_anchors()")
def _tinfo(t):
def _tinfo(t, a):
try:
n = len(json.loads(t) if t else [])
n_t = len(json.loads(t) if t else [])
except Exception:
n = 0
return "_no traces yet. Press **Space** to record trace #1._" if n == 0 \
else f"_{n} trace(s) committed. **Space** for another, or **Generate**._"
traces_json.change(_tinfo, traces_json, trace_info)
n_t = 0
try:
n_a = len(json.loads(a) if a else [])
except Exception:
n_a = 0
if n_t == 0 and n_a == 0:
return "_no traces or anchors yet._"
return f"_{n_t} trace(s) + {n_a} anchor(s) → {n_t + n_a} total track(s)_"
traces_json.change(_tinfo, [traces_json, anchors_json], info_md)
anchors_json.change(_tinfo, [traces_json, anchors_json], info_md)
# ---------------- generate (MATCHES TrackValidationCallback._sample exactly) ----------------
def _generate(clip_idx, traces_str, bg_str, seed, steps, w_t, w_m):
def _generate(clip_idx, traces_str, anchors_str, seed, steps, w_t, w_m):
def logln(log, msg):
line = time.strftime("%H:%M:%S ") + msg
print(line, flush=True)
return (log + "\n" + line) if log else line
log = ""
try:
bg = json.loads(bg_str) if bg_str else []
except Exception:
bg = []
tp_np, tv_np, n_tr, n_bg = build_track_points(traces_str, bg)
log = logln(log, f"clip idx={clip_idx}: {n_tr} trace(s) + {n_bg} bg point(s) → N={tp_np.shape[1]}")
tp_np, tv_np, n_tr, n_an = build_track_points(traces_str, anchors_str)
log = logln(log, f"clip idx={clip_idx}: {n_tr} trace(s) + {n_an} anchor(s) → N={tp_np.shape[1]}")
if tp_np.shape[1] == 0:
return None, "_record at least one trace, or enable bg points_", log
return None, "_record at least one trace or place at least one anchor_", log
ts = time.strftime("%Y%m%d_%H%M%S")
req_dir = Path(state["out_dir"]) / f"req_{ts}"
@@ -388,7 +348,7 @@ def build_ui(state):
(req_dir / "meta.json").write_text(json.dumps({
"clip_idx": int(clip_idx), "seed": int(seed), "steps": int(steps),
"w_text": float(w_t), "w_motion": float(w_m),
"num_traces": n_tr, "num_bg": n_bg,
"num_traces": n_tr, "num_anchors": n_an,
}, indent=2))
log = logln(log, f"req -> {req_dir.name}, denoising ...")
@@ -406,7 +366,7 @@ def build_ui(state):
return None, f"_generation failed: {e}_", logln(log, f"EXC: {e}")
gen_btn.click(_generate,
[st_clip_idx, traces_json, bg_pts_json, seed_in, steps_in, w_text, w_motion],
[st_clip_idx, traces_json, anchors_json, seed_in, steps_in, w_text, w_motion],
[out_video, status, log_box])
demo.load(None, None, None, js=canvas_js(NUM_FRAMES, FPS))
@@ -490,10 +450,7 @@ def make_generator(model, tc, samples):
# Efficient parquet loader (reads only enough files to get N rows, NOT all 4494)
# =============================================================================
def _load_first_n_from_parquet(data_path: str, n: int, text_len: int) -> list:
"""Grab the first N rows from a preprocessed parquet dataset (like
``trackwan_infer.load_conditioning_from_parquet``, but stops after N rows so
259k-row datasets don't OOM the loader).
"""
"""Grab the first N rows from a preprocessed parquet dataset."""
import glob
import pyarrow.parquet as pq
from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v_track
@@ -508,7 +465,7 @@ def _load_first_n_from_parquet(data_path: str, n: int, text_len: int) -> list:
if len(rows) >= n:
break
sel = rows[:n]
print(f"[app] loaded {len(sel)} rows from {min(len(files), 1 + (n // max(1, (len(sel)//1))))} parquet file(s)", flush=True)
print(f"[app] loaded {len(sel)} rows from parquet", flush=True)
batch = collate_rows_from_parquet_schema(sel, pyarrow_schema_i2v_track,
text_padding_length=int(text_len), cfg_rate=0.0)
infos = batch.get("info_list") or [{} for _ in sel]
@@ -535,21 +492,16 @@ def main() -> None:
ap.add_argument("--model-dir", required=True, help="diffusers-exported ckpt dir")
ap.add_argument("--yaml", required=True, help="training yaml")
ap.add_argument("--data-path", required=True,
help="preprocessed openvid parquet root (combined_parquet_dataset)")
help="preprocessed parquet root (combined_parquet_dataset)")
ap.add_argument("--num-clips", type=int, default=30, help="how many parquet clips to preload")
ap.add_argument("--out-dir", default="/mnt/lustre/vlm-s4duan/multi_trace_out")
ap.add_argument("--host", default="0.0.0.0")
ap.add_argument("--port", type=int, default=7864)
ap.add_argument("--sam-weight", default=None)
args = ap.parse_args()
os.makedirs(args.out_dir, exist_ok=True)
state: dict = {"out_dir": args.out_dir}
# SAM (works on GPU when CUDA_VISIBLE_DEVICES set)
print(f"[app] loading SAM ...", flush=True)
state["sam"] = load_sam(args.sam_weight)
# Load model
print(f"[app] loading model from {args.model_dir} ...", flush=True)
sys.path.insert(0, str(REPO / "data_pipeline"))
@@ -559,14 +511,12 @@ def main() -> None:
state["model"] = model
state["tc"] = tc
# Load parquet conditioning EFFICIENTLY: stop as soon as we have enough rows.
# (trackwan_infer.load_conditioning_from_parquet reads ALL 4494 files → OOM on 259k rows.)
# Load parquet conditioning efficiently
print(f"[app] loading {args.num_clips} parquet clips ...", flush=True)
samples = _load_first_n_from_parquet(args.data_path, args.num_clips, text_len)
labels = []
label_to_idx = {}
for i, s in enumerate(samples):
# Trim caption to a short unique label for the dropdown
cap = (s["caption"] or f"clip_{i}").strip().replace("\n", " ")
lab = f"{i:03d} {cap[:60]}"
labels.append(lab)