UCPE/scripts/predict_one_sample.py

"""Run UCPE's WanVideoPipeline on ONE PanShot sample and write the output mp4.

Standalone (no Lightning trainer); reuses UCPE's `PanShotTrainModule.__init__` for
the pipeline + camera-patch setup, then loads the DeepSpeed checkpoint manually,
fetches a single sample by index from PanShotDataset, and calls `pipe(...)` the
same way `PanShotTrainModule.forward` does in src/main.py.

Used by ../cf_ucpe/scripts/compare_inference.py via subprocess.
"""
import argparse
import os
import sys
from pathlib import Path

import torch

# Make UCPE imports available regardless of cwd at invocation time.
HERE = Path(__file__).resolve().parent
UCPE_ROOT = HERE.parent
sys.path.insert(0, str(UCPE_ROOT))

from diffsynth import save_video                                   # noqa: E402
from src.main import PanShotTrainModule                                 # noqa: E402
from src.dataset import PanShotDataset                             # noqa: E402

NEGATIVE_PROMPT = (
    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，"
    "最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，"
    "畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
)


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--sample_idx", type=int, default=None,
                   help="Index into the test split (after filtering). "
                        "Mutually exclusive with --video_id.")
    p.add_argument("--video_id", default=None,
                   help="Pick the exact PanShot video by id (recommended for "
                        "cross-codebase reproducibility).")
    p.add_argument("--ckpt_path", required=True,
                   help="Path to UCPE checkpoint. Either a .ckpt file or a "
                        "DeepSpeed folder (last.ckpt/checkpoint/mp_rank_00_model_states.pt).")
    p.add_argument("--output_path", required=True, help="Where to write the mp4.")
    p.add_argument("--data_root", default=str(UCPE_ROOT / "data" / "UCPE"),
                   help="UCPE data root (parent of PanShot/).")
    p.add_argument("--video_subdir", default="videos_704")
    p.add_argument("--model_id", default=str(UCPE_ROOT / "Wan2.2-TI2V-5B"),
                   help="Local Wan2.2-TI2V-5B model dir, or HF repo id.")
    p.add_argument("--height", type=int, default=704)
    p.add_argument("--width", type=int, default=1280)
    p.add_argument("--num_frames", type=int, default=81)
    p.add_argument("--num_inference_steps", type=int, default=50)
    p.add_argument("--camera_condition", default="relray_absmap")
    p.add_argument("--attn_compress", type=int, default=8)
    p.add_argument("--adaptation_method", default="parallel",
                   choices=["before", "after", "parallel"])
    p.add_argument("--split", default="test")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--fps", type=int, default=16)
    return p.parse_args()


def resolve_ckpt(path_str):
    """Accept either a .pt file or a DeepSpeed folder (returns the .pt inside)."""
    p = Path(path_str)
    if p.is_dir():
        cand = p / "checkpoint" / "mp_rank_00_model_states.pt"
        if not cand.exists():
            sys.exit(f"DeepSpeed folder {p} missing checkpoint/mp_rank_00_model_states.pt")
        return str(cand)
    return str(p)


def main():
    args = parse_args()
    torch.set_grad_enabled(False)

    # -------- 1. Build UCPE model (pipe + camera patch) --------
    # PanShotTrainModule handles: from_pretrained, patch_dit, enable_grad. ckpt_path=None
    # because we load it manually below to support DeepSpeed format.
    print(f"[predict_one_sample] building model (model_id={args.model_id})")
    model = PanShotTrainModule(
        model_id=args.model_id,
        ckpt_path=None,
        height=args.height,
        width=args.width,
        num_frames=args.num_frames,
        num_inference_steps=args.num_inference_steps,
        camera_condition=args.camera_condition,
        attn_compress=args.attn_compress,
        adaptation_method=args.adaptation_method,
    )
    model = model.to("cuda")
    model.pipe.device = torch.device("cuda")
    model.eval()

    # -------- 2. Load checkpoint (DeepSpeed or Lightning .ckpt) --------
    ckpt_file = resolve_ckpt(args.ckpt_path)
    print(f"[predict_one_sample] loading ckpt: {ckpt_file}")
    sd = torch.load(ckpt_file, map_location="cpu", weights_only=False)
    if "module" in sd:                  # DeepSpeed
        sd = sd["module"]
    elif "state_dict" in sd:            # Lightning
        sd = sd["state_dict"]
    missing, unexpected = model.load_state_dict(sd, strict=False)
    print(f"[predict_one_sample] loaded; missing={len(missing)} unexpected={len(unexpected)}")
    if unexpected:
        print(f"[predict_one_sample] (first unexpected) {unexpected[0]}")

    # patch_dit adds cam_self_attn modules in default fp32; pipe.from_pretrained
    # already loaded the rest of the DiT in bf16. Mixed dtypes blow up at
    # bf16-input × fp32-weight matmuls. Cast the whole DiT to bf16 after load
    # so cam modules align with the rest.
    model.pipe.dit = model.pipe.dit.to(torch.bfloat16)

    # -------- 3. Build a small DataModule-equivalent args object --------
    # PanShotDataset uses both attribute access (args.data_root) AND
    # membership tests (`"model_id" in args`), so we need a Namespace-with-
    # __contains__. omegaconf.DictConfig satisfies both.
    from omegaconf import DictConfig
    hp = DictConfig({
        "data_root": args.data_root,
        "video_subdir": args.video_subdir,
        "zero_first_yaw": True,
    })

    if (args.video_id is None) == (args.sample_idx is None):
        sys.exit("specify exactly one of --video_id or --sample_idx")
    dataset_kwargs = dict(load_keys=["video", "pose", "input_image"])
    if args.video_id is not None:
        dataset_kwargs["video_ids"] = [args.video_id]
    dataset = PanShotDataset(hp, split=args.split, **dataset_kwargs)
    if args.video_id is not None:
        if len(dataset) == 0 or dataset.metas[0]["video_id"] != args.video_id:
            sys.exit(f"video_id {args.video_id!r} not found in {args.split} split")
        idx = 0
    else:
        if not (0 <= args.sample_idx < len(dataset)):
            sys.exit(f"sample_idx {args.sample_idx} out of range [0, {len(dataset)})")
        idx = args.sample_idx
    sample = dataset[idx]
    video_id = sample["video_id"]
    print(f"[predict_one_sample] picked video_id={video_id}")

    # -------- 4. Build batch (single-element 'collated') --------
    # The pipe was built with torch_dtype=bf16, so float tensors must be bf16
    # to match the VAE / DiT / camera-encoder Linear weights. (Lightning's
    # precision machinery handles this in normal training; we have to do it
    # manually here.)
    def _to_batch(v, cast_float=True):
        import numpy as np
        if isinstance(v, np.ndarray):
            v = torch.from_numpy(v)
        if isinstance(v, torch.Tensor):
            t = v.unsqueeze(0).to("cuda")
            if cast_float and t.is_floating_point():
                t = t.to(dtype=torch.bfloat16)
            return t
        return [v]

    batch = {
        "caption": [sample["caption"]],
        "input_image": _to_batch(sample["input_image"]),
        "pose": _to_batch(sample["pose"]),
        "xi": _to_batch(torch.tensor(float(sample["xi"]))),
        "x_fov": _to_batch(torch.tensor(float(sample["x_fov"]))),
    }

    # -------- 5. Run inference (mirrors PanShotTrainModule.forward) --------
    print(f"[predict_one_sample] running pipe ({args.num_inference_steps} steps)")
    video = model.pipe(
        prompt=batch["caption"][0],
        input_image=batch.get("input_image", None),
        camera_control_panshot={k: batch[k] for k in ["pose", "xi", "x_fov"]},
        negative_prompt=NEGATIVE_PROMPT,
        num_inference_steps=args.num_inference_steps,
        tiled=False,
        seed=args.seed,
        height=args.height,
        width=args.width,
        num_frames=args.num_frames,
    )

    out_path = Path(args.output_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    save_video(video, str(out_path), fps=args.fps, quality=8)
    print(f"[predict_one_sample] wrote: {out_path}")


if __name__ == "__main__":
    main()