File size: 5,589 Bytes
f2c0505 | 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 | from __future__ import annotations
import argparse
import json
import os
import sys
import hashlib
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
def main():
ap = argparse.ArgumentParser(description="Run official Wan-Animate-2 with packed OrbitQuant W4A4.")
ap.add_argument("--official-repo", default="/content/Wan-Animate-2")
ap.add_argument("--model-root", default="/content/Wan2_2_Animate_Repo")
ap.add_argument("--packed", default="/content/OrbitQuant_Animate2_W4A4_PACKED")
ap.add_argument("--image", default="/content/wan_reference/reference.jpg")
ap.add_argument("--video", default="/content/oAfghkYL_720p.mp4")
ap.add_argument("--output", default="/content/OrbitQuant_Animate2_W4A4_PACKED/output_smoke")
ap.add_argument("--prompt", default="A realistic person matching the reference image, performing the actions from the driving video, natural motion, detailed, high quality.")
ap.add_argument("--prompt-ref", default="人物动作的参考视频")
ap.add_argument("--negative-prompt", default="")
ap.add_argument("--width", type=int, default=256)
ap.add_argument("--height", type=int, default=320)
ap.add_argument("--fps", type=int, default=24)
ap.add_argument("--clip-len", type=int, default=17)
ap.add_argument("--steps", type=int, default=10)
ap.add_argument("--guidance", type=float, default=1.0)
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--placement", choices=["auto", "resident", "stream"], default="auto")
ap.add_argument("--attention", choices=["official", "sol", "para", "hybrid"], default="hybrid")
ap.add_argument("--sol-tau", type=float, default=1.0)
ap.add_argument("--kv-cache", choices=["cpu", "cpu-pinned", "gpu"], default="cpu")
ap.add_argument("--max-frames", type=int, default=17, help="Smoke input cap at target FPS; use 0 for the full driving video")
ap.add_argument("--kernel-gate", default="/content/orbitquant_w4a4_cuda_gate.json")
ap.add_argument("--skip-kernel-gate", action="store_true")
args = ap.parse_args()
if not args.skip_kernel_gate:
gate_path = Path(args.kernel_gate)
if not gate_path.is_file():
raise RuntimeError(
f"CUDA W4A4 kernel gate not found: {gate_path}. "
"Run scripts/kernel_selftest.py successfully before loading the 14B model."
)
gate = json.loads(gate_path.read_text())
if gate.get("status") != "PASS":
raise RuntimeError(f"CUDA W4A4 kernel gate is not PASS: {gate}")
import torch
if gate.get("gpu") != torch.cuda.get_device_name():
raise RuntimeError(
f"kernel gate was produced on {gate.get('gpu')!r}, current GPU is {torch.cuda.get_device_name()!r}"
)
if gate.get("torch") != torch.__version__ or gate.get("cuda") != torch.version.cuda:
raise RuntimeError(
"Torch/CUDA changed since the kernel gate; rerun scripts/kernel_selftest.py "
f"(gate torch={gate.get('torch')} cuda={gate.get('cuda')}, "
f"current torch={torch.__version__} cuda={torch.version.cuda})."
)
artifact_gate = gate.get("packed_artifact")
if not artifact_gate:
raise RuntimeError(
"CUDA gate did not validate the actual packed artifact. "
"Rerun scripts/kernel_selftest.py --packed-dir <packed-dir>."
)
pdir_gate = Path(args.packed)
manifest_path = pdir_gate / "packed_manifest.json"
if not manifest_path.is_file():
raise RuntimeError(f"packed manifest missing: {manifest_path}")
current_manifest_sha = hashlib.sha256(manifest_path.read_bytes()).hexdigest()
if current_manifest_sha != artifact_gate.get("manifest_sha256"):
raise RuntimeError("packed manifest changed since the CUDA artifact gate")
for name, size in artifact_gate.get("shard_sizes", {}).items():
p = pdir_gate / name
if not p.is_file() or p.stat().st_size != int(size):
raise RuntimeError(f"packed shard changed since CUDA gate: {p}")
print("✓ CUDA W4A4 kernel + actual packed-artifact gate:", gate_path)
official = Path(args.official_repo).resolve()
if not official.is_dir():
raise FileNotFoundError(f"official Wan-Animate-2 source not found: {official}; run COLAB_BOOTSTRAP.py first")
sys.path.insert(0, str(official))
from orbitquant_wan_a2.official_runtime import run_official_packed_w4a4
report = run_official_packed_w4a4(
official_repo=official,
model_root=args.model_root,
packed_dir=args.packed,
reference_image=args.image,
driving_video=args.video,
output_dir=args.output,
prompt=args.prompt,
prompt_ref=args.prompt_ref,
negative_prompt=args.negative_prompt,
width=args.width,
height=args.height,
fps=args.fps,
clip_len=args.clip_len,
steps=args.steps,
guidance_scale=args.guidance,
seed=args.seed,
transformer_placement=args.placement,
attention=args.attention,
sol_tau=args.sol_tau,
kv_cache_placement=args.kv_cache,
max_input_frames=(None if args.max_frames <= 0 else args.max_frames),
)
if int(os.environ.get("RANK", "0")) == 0:
print(json.dumps(report, indent=2, default=str))
if __name__ == "__main__":
main()
|