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