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