| """""" |
| 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 根目录：<root>/<clip_id>/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() |
|
|
