| """""" |
| 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 |
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| from model_resources import wan_base_root, wan_model_paths, wan_tokenizer_path |
| VACE = wan_base_root() |
| DS = _project_path('third_party/DiffSynth-Studio') |
| EVAL_MANIFEST = _project_path('manifests/eval.json') |
| MASK_ROOT = seg_mask_root() |
| PALETTE = np.array([[0, 0, 0], [220, 40, 40], [40, 200, 80]], dtype=np.uint8) |
| NEG = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景' |
|
|
| def build_pipe(device, lora_path): |
| sys.path.insert(0, DS) |
| sys.path.insert(0, HERE) |
| from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig |
| 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'已加载 LoRA: {lora_path}') |
| return pipe |
|
|
| 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(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--lora', default=_project_path('checkpoints/generator_hoi/ckpt.safetensors')) |
| ap.add_argument('--out', default=_project_path('outputs/pred_hoi')) |
| 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('--manifest', default=EVAL_MANIFEST) |
| 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() |
| manifest = json.load(open(args.manifest)) |
| if args.limit > 0: |
| manifest = manifest[:args.limit] |
| out = os.path.abspath(args.out) |
| os.makedirs(out, exist_ok=True) |
| pipe = build_pipe(args.device, args.lora) |
| size = (args.width, args.height) |
| print(f'开始 M2 推理, {len(manifest)} 条') |
| for k, s in enumerate(manifest): |
| if os.path.isfile(os.path.join(out, s['take_uid'], 'pred.mp4')): |
| print(f' [skip] 已存在: {s['take_uid']}') |
| continue |
| mask_dir = os.path.join(args.mask_root, s['take_uid'], 'ego') |
| if not os.path.isdir(mask_dir): |
| print(f' [skip] 无mask: {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) |
| print(f"[{k + 1}/{len(manifest)}] {s['take_name']} | {nf}帧 | '{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, s['take_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 |
| print(f' -> {sd}') |
| json.dump({'model': 'Wan2.1-VACE-1.3B + HOI-LoRA', 'lora': args.lora, 'height': args.height, 'width': args.width, 'rotate': False, 'vace_scale': args.vace_scale, 'note': 'seg 预测的 ego HOI mask 引导生成'}, open(os.path.join(out, 'infer_config.json'), 'w'), indent=2) |
| print('完成 ->', out) |
| if __name__ == '__main__': |
| main() |
|
|
