"""""" from pathlib import Path as _Path import sys as _sys _bundle_root = _Path(__file__).resolve().parents[1] _sys.path.insert(0, str(_bundle_root)) from project_paths import project_path as _project_path, env as _env, initialize as _initialize _initialize(__name__ == '__main__') from project_paths import seg_mask_root, seg_mask_dir import os import sys import json import argparse import numpy as np import torch from PIL import Image from model_resources import wan_base_root, wan_model_paths, wan_tokenizer_path VACE = wan_base_root() DS = _project_path('third_party/DiffSynth-Studio') HERE = os.path.dirname(os.path.abspath(__file__)) EVAL_MANIFEST = _project_path('manifests/eval.json') MASK_ROOT = seg_mask_root() CAPTION_JSON = _project_path('data/captions/grpo.json') PALETTE = np.array([[0, 0, 0], [220, 40, 40], [40, 200, 80]], dtype=np.uint8) NEG = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，鱼眼畸变，鱼眼镜头，广角畸变，桶形畸变，镜头畸变，边缘拉伸变形，画面弯曲，fisheye distortion，wide-angle distortion，lens distortion，barrel distortion，可见的相机，三脚架，GoPro，运动相机' def build_pipe(device, lora_path, branch_ckpt, caption_dim=4096): sys.path.insert(0, DS) sys.path.insert(0, HERE) from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig from caption_cross_attn import inject_caption_cross_attn pipe = WanVideoPipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=[ModelConfig(path=path) for path in wan_model_paths(VACE, strict=True)], tokenizer_config=ModelConfig(path=wan_tokenizer_path(VACE, strict=True))) if lora_path: pipe.load_lora(pipe.vace, lora_path, alpha=1.0) print(f'已加载 Stage1 底座 LoRA: {lora_path}') dit = pipe.dit dev = next(dit.parameters()).device dtype = next(dit.parameters()).dtype from safetensors.torch import load_file sd = load_file(branch_ckpt) ckpt_has_route = any(('spatial_route' in k for k in sd.keys())) env_route = os.environ.get('USE_SPATIAL_ROUTE') use_spatial_route = env_route != '0' if env_route is not None else ckpt_has_route inject_caption_cross_attn(dit, caption_dim=caption_dim, vace_module=pipe.vace, use_spatial_route=use_spatial_route) for block in dit.blocks: block.cross_attn._cap_branch.to(device=dev, dtype=dtype) dit_sd = dit.state_dict() n_cap = 0 miss = [] for k, v in sd.items(): kk = k for pre in ('pipe.dit.', 'dit.', 'model.'): if kk.startswith(pre): kk = kk[len(pre):] if kk in dit_sd: dit_sd[kk].copy_(v.to(dit_sd[kk].dtype).to(dit_sd[kk].device)) if '_cap_branch' in kk: n_cap += 1 else: miss.append(k) print(f'[caption ckpt] 加载 caption {n_cap}/{len(sd)} 个参数 <- {branch_ckpt}') if miss: print(f'[caption ckpt] 未匹配 {len(miss)} 个 key, 样本: {miss[:3]}') for block in dit.blocks: cb = block.cross_attn._cap_branch cb.gate_cap.data = cb.gate_cap.data.float() print(f'[caption ckpt] block0 gate_cap={dit.blocks[0].cross_attn._cap_branch.gate_cap.item():+.4e}') return pipe @torch.no_grad() def encode_caption(pipe, cap): """""" if cap is None: return None cap = str(cap).strip() if len(cap) == 0: return None pipe.load_models_to_device(('text_encoder',)) ids, mask = pipe.tokenizer(cap, return_mask=True, add_special_tokens=True) ids = ids.to(pipe.device) mask = mask.to(pipe.device) seq_lens = mask.gt(0).sum(dim=1).long() feat = pipe.text_encoder(ids, mask) for i, v in enumerate(seq_lens): feat[:, v:] = 0 return feat.to(device=pipe.device, dtype=pipe.torch_dtype) def load_mask_video(mask_dir, size): """""" fs = sorted((f for f in os.listdir(mask_dir) if f.endswith('.png'))) out = [] for f in fs: m = np.asarray(Image.open(os.path.join(mask_dir, f))) if m.ndim == 3: m = m[..., 0] col = PALETTE[np.clip(m, 0, 2)] out.append(Image.fromarray(col).resize(size)) return out def main(): mode_parser = argparse.ArgumentParser(add_help=False) modes = ['global', 'spatial'] mode_parser.add_argument('--variant', choices=modes, default='spatial') mode, _ = mode_parser.parse_known_args() branch = 'global' if mode.variant == 'global' else 'spatial' ap = argparse.ArgumentParser() ap.add_argument('--variant', choices=modes, default=mode.variant) ap.add_argument('--lora', default=_project_path('checkpoints/generator_hoi/ckpt.safetensors'), help='Stage 1 HOI 底座，与分支训练的 preset_lora_path 一致') ap.add_argument('--branch_ckpt', default=_project_path(f'checkpoints/generator_{branch}/ckpt.safetensors')) ap.add_argument('--out', default=_project_path(f'outputs/pred_{mode.variant}')) ap.add_argument('--manifest', default=EVAL_MANIFEST) ap.add_argument('--caption_json', default=CAPTION_JSON) ap.add_argument('--caption_dim', type=int, default=4096) ap.add_argument('--height', type=int, default=480) ap.add_argument('--width', type=int, default=480) ap.add_argument('--steps', type=int, default=50) ap.add_argument('--cfg_scale', type=float, default=5.0) ap.add_argument('--vace_scale', type=float, default=1.0) ap.add_argument('--fps', type=float, default=8.0) ap.add_argument('--seed', type=int, default=0) ap.add_argument('--limit', type=int, default=-1) ap.add_argument('--mask_root', default=MASK_ROOT, help='外部 seg 预测 mask 根目录：//ego/*.png') ap.add_argument('--device', default='cuda') args = ap.parse_args() os.environ['USE_SPATIAL_ROUTE'] = '0' if args.variant == 'global' else '1' sys.path.insert(0, HERE) from caption_cross_attn import set_caption_context manifest = json.load(open(args.manifest)) if args.limit > 0: manifest = manifest[:args.limit] captions = json.load(open(args.caption_json)) out = os.path.abspath(args.out) os.makedirs(out, exist_ok=True) pipe = build_pipe(args.device, args.lora, args.branch_ckpt, args.caption_dim) size = (args.width, args.height) print(f'开始 caption 旁路推理, {len(manifest)} 条 | cfg={args.cfg_scale}') n_done = 0 for k, s in enumerate(manifest): uid = s['take_uid'] if os.path.isfile(os.path.join(out, uid, 'pred.mp4')): print(f' [skip] 已存在: {uid}') n_done += 1 continue mask_dir = os.path.join(args.mask_root, uid, 'ego') if not os.path.isdir(mask_dir): print(f' [skip] 无mask: {s['take_name']}') continue cap = captions.get(uid) if cap is None: print(f' [skip] 无caption: {s['take_name']}') continue mask_vid = load_mask_video(mask_dir, size) nf = len(mask_vid) ref = Image.open(s['first_frame']).convert('RGB').resize(size) cap_feat = encode_caption(pipe, cap) set_caption_context(pipe.dit, cap_feat) cap_shape = tuple(cap_feat.shape) if cap_feat is not None else 'None(旁路关)' print(f"[{k + 1}/{len(manifest)}] {s['take_name']} | {nf}帧 | variant={args.variant} | cap{cap_shape} | '{s['prompt'][:40]}'") video = pipe(prompt=s['prompt'], negative_prompt=NEG, vace_video=mask_vid, vace_reference_image=ref, vace_scale=args.vace_scale, height=args.height, width=args.width, num_frames=nf, num_inference_steps=args.steps, cfg_scale=args.cfg_scale, seed=args.seed, tiled=True) sd = os.path.join(out, uid) os.makedirs(sd, exist_ok=True) for i, fr in enumerate(video): Image.fromarray(np.asarray(fr)).save(os.path.join(sd, f'{i:03d}.png')) try: from diffsynth.utils.data import save_video save_video(video, os.path.join(sd, 'pred.mp4'), fps=args.fps, quality=5) except Exception: pass n_done += 1 print(f' -> {sd}') json.dump({'model': 'Wan2.1-VACE-1.3B + HOI-LoRA + caption 旁路', 'lora': args.lora, 'branch_ckpt': args.branch_ckpt, 'caption_json': args.caption_json, 'height': args.height, 'width': args.width, 'rotate': False, 'vace_scale': args.vace_scale, 'steps': args.steps, 'cfg_scale': args.cfg_scale, 'variant': args.variant, 'seed': args.seed, 'num_samples': n_done, 'note': 'ego HOI mask + narration(原生) + ego dense caption(旁路) 引导生成'}, open(os.path.join(out, 'infer_config.json'), 'w'), ensure_ascii=False, indent=2) print(f'完成 {n_done} 条 ->', out) if __name__ == '__main__': main()