File size: 9,044 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | """"""
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()
|