""""""
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()