# Written by yams_any4d agent, 2026-07-10. """ Any4D + V4Head inference with RERUN panels. Like infer_anydata.py, but after each generation it also runs the V4Head action head on the DiT's stashed pre-final-layer features (the LAST denoise step = fixed low-sigma readout) and logs everything to rerun: gen/scene, gen/wrist generated video per view (frame timeline) gt/scene ground-truth video pred/kp GT (white circles + trail) vs pred (rainbow) keypoints on the GENERATED scene frames pred/heatmap_scene volume Z-marginal softmax JET-blended on gen frames pred/pca_scene PCA of per-latent-frame DiT features (shared basis) pred/grip_grid, pred/rot_grid bin grids (whole window, once per sample) status per-sample losses + metadata Output: a .rrd file (open at https://app.rerun.io or `rerun file.rrd`), or --serve to host the web viewer directly (rerun_helpers.init_web pattern from yam_local_train; port = web UI, port+1 = grpc). LoRA checkpoints are auto-merged in memory (peft base_layer + lora_A/B -> plain weights) before loading — loading unmerged LoRA ckpts into a plain model produces pure noise (hard-won lesson, 2026-07-10). V4Head keys pass through. Usage (yukon): PYTHONPATH=$PWD:$PWD/externals/AnyData python custom/eval/infer_v4head_rerun.py \ --exp_cfg=custom/experiment/basile/rwm4_pp70_v4head_yukon.py \ --ckpt_path=/model/iter_*.pt \ --dset_cfg=custom/config/anydata/web/debug/yamyukon_pp70.yaml \ --output_dir=$HOME/any4d_work/output/v4head_rerun \ --num_steps=25 --stop_after=3 """ import argparse import copy import os import cv2 import numpy as np import torch import custom.eval.infer_anydata as base from custom.eval import infer_utils from custom.eval.visuals import cached_decode_latent_video_auto from custom.any4d.v4_head_viz import bin_grid_panel, _pca_rgb # ---------------------------------------------------------------------------- def merge_lora_state_dict(sd): """In-memory peft-LoRA merge: W += B @ A (scale = alpha/r = 1 for our r16/a16), rename base_layer.* -> *, drop lora keys. No-op for plain checkpoints.""" la = [k for k in sd if k.endswith('lora_A.default.weight')] if not la: return sd, False for ka in la: stem = ka[: -len('lora_A.default.weight')] A = sd.pop(ka).float() B = sd.pop(stem + 'lora_B.default.weight').float() kw = stem + 'base_layer.weight' sd[kw] = (sd[kw].float() + B @ A).to(sd[kw].dtype) for k in list(sd.keys()): if '.base_layer.' in k: sd[k.replace('.base_layer.', '.')] = sd.pop(k) return sd, True def to_u8(frame_chw, bgr=True): """(3, H, W) float in [-1,1] or [0,1] -> HxWx3 uint8 (BGR for cv2).""" f = frame_chw.detach().float().cpu().numpy().transpose(1, 2, 0) if f.min() < -0.01: f = (f + 1.0) / 2.0 f = np.clip(f, 0, 1) img = (f * 255).astype(np.uint8) return cv2.cvtColor(img, cv2.COLOR_RGB2BGR) if bgr else img def jpeg(img_bgr, quality=80): import rerun as rr ok, buf = cv2.imencode('.jpg', img_bgr, [cv2.IMWRITE_JPEG_QUALITY, quality]) if not ok: return rr.Image(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)) return rr.EncodedImage(contents=bytes(buf), media_type='image/jpeg') def gt_action_from_batch(data_batch): """(T, 20) float tensor from the parsed anydata action entries (mask-free).""" act = data_batch['anydata']['action'] keys = sorted(act.keys(), key=lambda k: (k[0], k[1]) if isinstance(k, tuple) else k) seen, rows = set(), [] for k in keys: t = k[0] if isinstance(k, tuple) else k if t in seen: continue seen.add(t) a = act[k]['action'] a = a[0] if getattr(a, 'ndim', 1) == 2 else a # collate may add B dim rows.append(torch.as_tensor(np.asarray(a), dtype=torch.float32)) return torch.stack(rows) # (T, 20) def kp_frame(bg_bgr, gt_uv, pred_uv, t): """Per-frame keypoint overlay: full GT trail (thin) + trails up to t + big current markers.""" img = bg_bgr.copy() gt = np.round(gt_uv).astype(np.int32) pr = np.round(pred_uv).astype(np.int32) for i in range(len(gt) - 1): cv2.line(img, tuple(gt[i]), tuple(gt[i + 1]), (180, 180, 180), 1, cv2.LINE_AA) for i in range(max(t, 1) - 1): cv2.line(img, tuple(pr[i]), tuple(pr[i + 1]), (255, 255, 255), 1, cv2.LINE_AA) cv2.circle(img, tuple(gt[t]), 10, (255, 255, 255), 2, cv2.LINE_AA) cv2.circle(img, tuple(pr[t]), 6, (0, 0, 255), -1, cv2.LINE_AA) return img def heatmap_frame(bg_bgr, vol_t, alpha=0.55): """(Z, P, P) logits -> Z-marginal softmax JET blended on bg (v4 marginal style).""" x = vol_t.detach().float().cpu().numpy() flat = x.reshape(-1) e = np.exp(flat - flat.max()) p = (e / e.sum()).reshape(x.shape).sum(0) # (P, P) marginal over Z p = (p - p.min()) / (p.max() - p.min() + 1e-8) hm = cv2.applyColorMap((p * 255).astype(np.uint8), cv2.COLORMAP_JET) hm = cv2.resize(hm, (bg_bgr.shape[1], bg_bgr.shape[0]), interpolation=cv2.INTER_LINEAR) return cv2.addWeighted(hm, alpha, bg_bgr, 1 - alpha, 0) # ---------------------------------------------------------------------------- def main(args): import rerun as rr (model, config, lpips_loss, device, exp_cfg) = base.setup_runtime(args) # --- checkpoint: auto-merge LoRA then load --- print(f'[cyan]Loading ckpt {args.ckpt_path} (auto-merge LoRA if present)...') sd = torch.load(args.ckpt_path, map_location='cpu', weights_only=False) sd, merged = merge_lora_state_dict(sd) print(f'[cyan] lora merged: {merged} | keys: {len(sd)}') missing, unexpected = model.pipe.dit.load_state_dict(sd, strict=False) print(f'[cyan] load: {len(missing)} missing, {len(unexpected)} unexpected') v4h_missing = [k for k in missing if 'v4head' in k] if v4h_missing: print(f'[red] WARNING: v4head keys missing from ckpt: {v4h_missing[:4]} ...') dataloader = base.load_dataset(args, exp_cfg) # --- rerun init --- rr.init('v4head_infer', spawn=False) rrd_path = os.path.join(args.output_dir, 'v4head_infer.rrd') os.makedirs(args.output_dir, exist_ok=True) if args.serve: grpc_port = args.port + 1 rr.serve_grpc(grpc_port=grpc_port) rr.serve_web_viewer(web_port=args.port, open_browser=False) print(f'[green]rerun web viewer: http://:{args.port}' f'?url=rerun%2Bhttp%3A%2F%2F127.0.0.1%3A{grpc_port}%2Fproxy') else: rr.save(rrd_path) print(f'[green]streaming to {rrd_path}') # Head: either embedded in the DiT ckpt (joint runs) or a standalone state_dict # from the cached-training pipeline (--head_ckpt, scene-only). if args.head_ckpt: from custom.any4d.v4_head import V4Head # Head config MUST match training (else state_dict shapes mismatch / feature # standardization is silently off). Canonical recipe: per_step, pred 128, t_out 17, # scene-only, input_channel_norm ON. Override via CLI for other heads. head = V4Head(prep_path=args.prep, use_wrist=not args.scene_only, temporal_mode=args.temporal_mode, pred_size=args.pred_size, t_out=args.t_out, n_lat_frames=args.n_lat_frames, input_channel_norm=args.input_channel_norm).to(device).float() hmiss, hunexp = head.load_state_dict( torch.load(args.head_ckpt, map_location='cpu', weights_only=False), strict=False) assert not hunexp and all('K_' in k or 'T_w2c' in k or 'hand_eye' in k or 'rot_centroids' in k or 'range' in k or 'save_wh' in k for k in hmiss), \ f'head config mismatch — missing {hmiss[:4]} unexpected {hunexp[:4]}' head.eval() # attach a stub so a4d_network stashes features during sampling if getattr(model.pipe.dit, 'v4head', None) is None: model.pipe.dit.v4head = torch.nn.Identity() model.pipe.dit._v4head_feats = None print(f'[cyan]standalone head loaded from {args.head_ckpt} (use_wrist={head.use_wrist})') else: head = model.pipe.dit.v4head assert head is not None, 'exp_cfg must have v4head_enabled=True (or pass --head_ckpt)' directives = dict(config.val_directives) if isinstance(getattr(config, 'val_directives', None), dict) else {} if args.cond_frames_raw is not None: directives['cond_frames_raw'] = args.cond_frames_raw n_done = 0 for batch_idx, data_batch in enumerate(dataloader): if args.stop_after > 0 and n_done >= args.stop_after: break print(f'[cyan]=== sample {n_done} (batch {batch_idx}) ===') gt_action = gt_action_from_batch(data_batch) # (T, 20) with torch.no_grad(): val_dict, _ = model.validation_step( data_batch=copy.deepcopy(data_batch), iteration=-1, dataloader_key='infer', local_path=args.output_dir, directives=directives, val_iter=n_done, seed=args.seed + n_done) feats = model.pipe.dit._v4head_feats # last denoise step assert feats is not None and 0 in feats, 'no stashed features — v4head_enabled off?' (Hp, Wp) = feats[0].shape[2], feats[0].shape[3] pixel_hw = (Hp * 16, Wp * 16) if args.scene_only: feats = {0: feats[0]} with torch.no_grad(): head_out = head({k: v.float() for k, v in feats.items()}, pixel_hw, gt_action[None].to(feats[0].device)) # decoded generated + GT videos: (3, Tp, H, W) each gen = {} for (v, key) in [(0, 'rgb0'), (1, 'rgb1')]: if args.scene_only and v == 1: continue if key in val_dict.get('y0_pred_entries', {}): dec = cached_decode_latent_video_auto(val_dict['y0_pred_entries'], key, model.vae)[0] if dec.shape[1] > 1: # skip 1-frame stubs of absent views gen[v] = dec gt_vid = None if 'rgb0' in val_dict.get('x0_entries', {}): gt_vid = cached_decode_latent_video_auto(val_dict['x0_entries'], 'rgb0', model.vae)[0] vol = head_out['volume_logits'][0].float() # (T, NZ, P, P) vpos = head_out['voxel_positions'] if vpos.ndim == 4: # legacy (NZ,P,P,3) vpos = vpos.unsqueeze(0) pred_xyz = head.decode_pred_xyz(head_out['volume_logits'].float(), vpos.float())[0] # (T, 3) gt_xyz = head_out['gt_xyz'][0] pred_uv = head.project_world_to_pix(pred_xyz, pixel_hw).cpu().numpy() gt_uv = head.project_world_to_pix(gt_xyz, pixel_hw).cpu().numpy() T = vol.shape[0] Tp = gen[0].shape[1] if 0 in gen else T rr.set_time('sample', sequence=n_done) # once-per-sample panels rr.set_time('frame', sequence=0) rr.log('pred/grip_grid', rr.Image(cv2.cvtColor(bin_grid_panel( head_out['grip_logits'][0].float().cpu().numpy(), head_out['tgt_grip'][0].cpu().numpy()), cv2.COLOR_BGR2RGB))) rr.log('pred/rot_grid', rr.Image(cv2.cvtColor(bin_grid_panel( head_out['rot_logits'][0].float().cpu().numpy(), head_out['tgt_rot'][0].cpu().numpy()), cv2.COLOR_BGR2RGB))) losses = {k: round(float(v), 4) for k, v in head_out['losses'].items()} rr.log('status', rr.TextDocument( f"sample {n_done} | id={data_batch.get('sample_id', ['?'])[0]}\n" f"v4head losses: {losses}\nckpt: {os.path.basename(args.ckpt_path)} (merged={merged})")) # PCA per latent frame (shared basis), logged at raw t = lat * 4 p0 = head_out['scene_feats_per_frame'][0].detach().float().cpu().numpy() # (Tl, C, Hp, Wp) pca_imgs = _pca_rgb(p0, (Wp * 16, Hp * 16)) for lat in range(len(pca_imgs)): rr.set_time('frame', sequence=min(lat * 4, T - 1)) rr.log('pred/pca_scene', rr.Image(pca_imgs[lat])) for t in range(min(T, Tp)): rr.set_time('frame', sequence=t) bg = to_u8(gen[0][:, t]) if 0 in gen else np.full( (pixel_hw[0], pixel_hw[1], 3), 80, np.uint8) rr.log('gen/scene', jpeg(bg)) if 1 in gen: rr.log('gen/wrist', jpeg(to_u8(gen[1][:, t]))) if gt_vid is not None: rr.log('gt/scene', jpeg(to_u8(gt_vid[:, t]))) rr.log('pred/kp', jpeg(kp_frame(bg, gt_uv, pred_uv, t))) rr.log('pred/heatmap_scene', jpeg(heatmap_frame(bg, vol[t]))) n_done += 1 if not args.serve: print(f'[green]DONE — rrd at {rrd_path}') else: print('[green]DONE — viewer stays up; Ctrl-C to stop') import time while True: time.sleep(60) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--exp_cfg', type=str, required=True) parser.add_argument('--ckpt_path', type=str, required=True) parser.add_argument('--dset_cfg', type=str, required=True) parser.add_argument('--output_dir', type=str, required=True) parser.add_argument('--num_steps', type=int, default=25) parser.add_argument('--stop_after', type=int, default=3) parser.add_argument('--batch_size', type=int, default=1) parser.add_argument('--seed', type=int, default=1234) parser.add_argument('--cond_frames_raw', type=int, default=None) parser.add_argument('--head_ckpt', type=str, default=None) parser.add_argument('--prep', type=str, default='/home/robot-lab/any4d_work/v4head_prep.npz') parser.add_argument('--scene_only', action='store_true') # V4Head config — MUST match how --head_ckpt was trained (defaults = canonical 2026-07-13 recipe) parser.add_argument('--temporal_mode', type=str, default='per_step', choices=['per_step', 'stack_mlp']) parser.add_argument('--pred_size', type=int, default=128) parser.add_argument('--t_out', type=int, default=17) parser.add_argument('--n_lat_frames', type=int, default=5) parser.add_argument('--input_channel_norm', action='store_true', help='online per-channel feature standardization (ON in the trained head)') parser.add_argument('--serve', action='store_true') parser.add_argument('--port', type=int, default=9092) parser.add_argument('--device', type=str, default='cuda') parser.add_argument('--gpu_id', type=int, default=0) parser.add_argument('--debug', action='store_true') parser.add_argument('--data_overrides', type=str, default=None) parser.add_argument('--data_defaults', type=str, default=None) parser.add_argument('--any4d_overrides', type=str, default=None) parser.add_argument('--donor_rgb', type=str, default=None) parser.add_argument('--donor_rgb_cam', type=int, default=0) _args = parser.parse_args() # setup_runtime/load_dataset (borrowed from infer_anydata) touch more args than we # expose — fill safe defaults for any that are absent. _defaults = dict(exp_overrides=None, cond_aug_sigma=None, guidance=None, infer_metrics=False, metric_resolution=None, num_samples=1, num_segments=4, perturb_action=None, perturb_per_sample=False, perturb_traj=None, run_autoregressive=False, shard=None, visual_detail=1, viz_extra_modes=None, donor_rgb_path=None, extrapolation_strategy='backtrack') for (_k, _v) in _defaults.items(): if not hasattr(_args, _k): setattr(_args, _k, _v) main(_args)