File size: 4,992 Bytes
368274e | 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 | #!/usr/bin/env python
"""Run the official Wan2.2-Animate-14B example against this quantized repo.
Usage (from the repository root):
python examples/run_example.py --backend native
python examples/run_example.py --backend sage_v2 --steps 20
"""
from __future__ import annotations
import argparse
from pathlib import Path
import imageio.v3 as iio
import numpy as np
import torch
from diffusers import AutoencoderKLWan, WanAnimatePipeline
from diffusers.models.attention_dispatch import (
AttentionBackendName,
_AttentionBackendRegistry,
_check_device_cuda_atleast_smXY,
_check_shape,
)
from diffusers.utils import load_image, load_video
MODEL_DIR = Path(__file__).resolve().parents[1]
EXAMPLE_DIR = Path(__file__).resolve().parent / "animate"
PROMPT = "People in the video are doing actions."
def register_sage_backends() -> None:
from sageattention import sageattn, sageattn_qk_int8_pv_fp16_triton
from sageattn3 import sageattn3_blackwell
def sage_v1(query, key, value, attn_mask=None, is_causal=False, scale=None, **kw):
if attn_mask is not None:
raise ValueError("SAGE_ATTN_1 does not support attn_mask")
return sageattn_qk_int8_pv_fp16_triton(
query, key, value, tensor_layout="NHD", is_causal=is_causal, sm_scale=scale, smooth_k=True
)
def sage_v2(query, key, value, attn_mask=None, is_causal=False, scale=None, **kw):
if attn_mask is not None:
raise ValueError("SAGE_ATTN does not support attn_mask")
return sageattn(query, key, value, tensor_layout="NHD", is_causal=is_causal, sm_scale=scale, smooth_k=True)
def sage_v3(query, key, value, attn_mask=None, is_causal=False, **kw):
if attn_mask is not None:
raise ValueError("SAGE_ATTN_3 does not support attn_mask")
return sageattn3_blackwell(
query.transpose(1, 2).contiguous(),
key.transpose(1, 2).contiguous(),
value.transpose(1, 2).contiguous(),
is_causal=is_causal,
).transpose(1, 2).contiguous()
_AttentionBackendRegistry.register(
AttentionBackendName._SAGE_QK_INT8_PV_FP16_TRITON,
constraints=[_check_device_cuda_atleast_smXY(8, 0), _check_shape],
)(sage_v1)
_AttentionBackendRegistry.register(
AttentionBackendName._SAGE_QK_INT8_PV_FP8_CUDA,
constraints=[_check_device_cuda_atleast_smXY(8, 0), _check_shape],
)(sage_v2)
_AttentionBackendRegistry.register(
AttentionBackendName._SAGE_QK_INT8_PV_FP8_CUDA_SM90,
constraints=[_check_device_cuda_atleast_smXY(10, 0), _check_shape],
)(sage_v3)
BACKENDS = {
"native": "native",
"sage_v1": "_sage_qk_int8_pv_fp16_triton",
"sage_v2": "_sage_qk_int8_pv_fp8_cuda",
"sage_v3": "_sage_qk_int8_pv_fp8_cuda_sm90",
}
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser()
p.add_argument("--backend", default="native", choices=sorted(BACKENDS))
p.add_argument("--steps", type=int, default=20)
p.add_argument("--segment-frame-length", type=int, default=77)
p.add_argument("--limit-frames", type=int, default=0)
p.add_argument("--output", default="")
p.add_argument("--no-offload", action="store_true", help="Keep all components on GPU (needs ~68 GB).")
return p.parse_args()
def main() -> None:
args = parse_args()
backend = BACKENDS[args.backend]
if args.backend.startswith("sage"):
register_sage_backends()
vae = AutoencoderKLWan.from_pretrained(MODEL_DIR, subfolder="vae", torch_dtype=torch.float32)
pipe = WanAnimatePipeline.from_pretrained(
MODEL_DIR,
torch_dtype=torch.bfloat16,
vae=vae,
device_map="balanced",
)
if args.no_offload:
pipe.to("cuda")
else:
pipe.reset_device_map()
pipe.enable_model_cpu_offload(device="cuda")
pipe.transformer.set_attention_backend(backend)
image = load_image(str(EXAMPLE_DIR / "src_ref.png"))
pose = load_video(str(EXAMPLE_DIR / "src_pose.mp4"))
face = load_video(str(EXAMPLE_DIR / "src_face.mp4"))
if args.limit_frames:
pose = pose[: args.limit_frames]
face = face[: args.limit_frames]
result = pipe(
image=image,
pose_video=pose,
face_video=face,
prompt=PROMPT,
height=720,
width=1280,
segment_frame_length=args.segment_frame_length,
prev_segment_conditioning_frames=1,
num_inference_steps=args.steps,
mode="animate",
guidance_scale=1.0,
generator=torch.Generator(device="cuda").manual_seed(42),
output_type="np",
)
frames = np.stack([np.asarray(f) for f in result.frames[0]], axis=0)
output = Path(args.output or f"example_{args.backend}.mp4")
iio.imwrite(output, frames, fps=30, codec="libx264", quality=8)
print(f"saved {output} ({frames.shape[0]} frames, {frames.shape[2]}x{frames.shape[1]})")
if __name__ == "__main__":
main()
|