yitongl/sparse_quant_exp

1#!/usr/bin/env python32"""Run Wan T2V inference with the sparse FP4 checkpoint-700 transformer."""3 4from __future__ import annotations5 6import argparse7import os8from pathlib import Path9 10 11DEFAULT_PROMPT = (12    "In the video, a woman is elegantly showcasing her earrings, bringing "13    "attention to their intricate design with a gentle touch of her fingers. "14    "She is bathed in ambient purple and pink lighting, which casts a soft "15    "glow on her delicate features and enhances the vivid tones of her lipstick "16    "and eye makeup. Her hair is styled to frame her face smoothly, emphasizing "17    "the contours of her jawline and cheekbones. The background features a "18    "blurred neon light, adding an artistic and modern touch to the overall "19    "aesthetic."20)21 22DEFAULT_NEGATIVE_PROMPT = (23    "Bright tones, overexposed, static, blurred details, subtitles, style, "24    "works, paintings, images, static, overall gray, worst quality, low quality, "25    "JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn "26    "hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused "27    "fingers, still picture, messy background, three legs, many people in the "28    "background, walking backwards"29)30 31 32def _resolve_weights(repo_id: str, weights: str | None, local_dir: str) -> str:33    if weights:34        path = Path(weights).expanduser()35        if path.exists():36            return str(path.resolve())37        raise FileNotFoundError(f"--weights does not exist: {path}")38 39    from huggingface_hub import hf_hub_download40 41    path = hf_hub_download(42        repo_id=repo_id,43        filename="transformer/diffusion_pytorch_model.safetensors",44        local_dir=local_dir,45        repo_type="model",46    )47    return str(Path(path).resolve())48 49 50def main() -> int:51    parser = argparse.ArgumentParser()52    parser.add_argument("--repo-id", default="yitongl/sparse_quant_exp")53    parser.add_argument(54        "--model-path",55        default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",56        help="Base Wan Diffusers model repo/path.",57    )58    parser.add_argument("--weights", default=None)59    parser.add_argument(60        "--local-dir",61        default="checkpoints/hf_download/sparse_quant_exp",62        help="Local Hugging Face download directory for the uploaded weights.",63    )64    parser.add_argument("--prompt", default=DEFAULT_PROMPT)65    parser.add_argument("--negative-prompt", default=DEFAULT_NEGATIVE_PROMPT)66    parser.add_argument("--output-path", default="outputs/sfp4_checkpoint_700")67    parser.add_argument("--height", type=int, default=448)68    parser.add_argument("--width", type=int, default=832)69    parser.add_argument("--num-frames", type=int, default=77)70    parser.add_argument("--num-inference-steps", type=int, default=50)71    parser.add_argument("--fps", type=int, default=16)72    parser.add_argument("--guidance-scale", type=float, default=5.0)73    parser.add_argument("--flow-shift", type=float, default=1.0)74    parser.add_argument("--seed", type=int, default=1000)75    parser.add_argument("--vsa-sparsity", type=float, default=0.9)76    parser.add_argument("--num-gpus", type=int, default=1)77    parser.add_argument("--sp-size", type=int, default=1)78    parser.add_argument("--tp-size", type=int, default=1)79    parser.add_argument("--text-encoder-cpu-offload", action="store_true", default=True)80    parser.add_argument("--pin-cpu-memory", action="store_true", default=False)81    args = parser.parse_args()82 83    os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "SPARSE_FP4_OURS_P_ATTN")84    os.environ.setdefault("FASTVIDEO_SPARSE_FP4_USE_HIGH_PREC_O", "1")85 86    weights_path = _resolve_weights(args.repo_id, args.weights, args.local_dir)87 88    from fastvideo import VideoGenerator89 90    generator = VideoGenerator.from_pretrained(91        model_path=args.model_path,92        num_gpus=args.num_gpus,93        sp_size=args.sp_size,94        tp_size=args.tp_size,95        init_weights_from_safetensors=weights_path,96        dit_cpu_offload=False,97        vae_cpu_offload=False,98        text_encoder_cpu_offload=args.text_encoder_cpu_offload,99        pin_cpu_memory=args.pin_cpu_memory,100        flow_shift=args.flow_shift,101        VSA_sparsity=args.vsa_sparsity,102    )103 104    result = generator.generate_video(105        prompt=args.prompt,106        negative_prompt=args.negative_prompt,107        output_path=args.output_path,108        save_video=True,109        return_frames=False,110        height=args.height,111        width=args.width,112        num_frames=args.num_frames,113        num_inference_steps=args.num_inference_steps,114        fps=args.fps,115        guidance_scale=args.guidance_scale,116        seed=args.seed,117    )118    print(result)119    return 0120 121 122if __name__ == "__main__":123    raise SystemExit(main())124