File size: 5,005 Bytes
c28ffff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """"""
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()
|