From d76dfa0df3e2bb82e3ffb499779849999f2cd363 Mon Sep 17 00:00:00 2001
From: RandNMR73 <notomatthew31@gmail.com>
Date: Tue, 13 Jan 2026 08:41:56 +0000
Subject: [PATCH] add inference script with custom weights

---
 ...geattn_inference_example_custom_weights.py | 54 +++++++++++++++++++
 1 file changed, 54 insertions(+)
 create mode 100644 examples/inference/sageattn_inference_example_custom_weights.py

diff --git a/examples/inference/sageattn_inference_example_custom_weights.py b/examples/inference/sageattn_inference_example_custom_weights.py
new file mode 100644
index 000000000..5457a7325
--- /dev/null
+++ b/examples/inference/sageattn_inference_example_custom_weights.py
@@ -0,0 +1,54 @@
+from fastvideo import VideoGenerator
+import os
+from pathlib import Path
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples"
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SAGE_ATTN_THREE"
+
+CHECKPOINT_PATH = Path(__file__).parent.parent.parent / "checkpoints" / "wan_t2v_finetune_qat" / "checkpoint-5000" / "transformer"
+
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        # image_encoder_cpu_offload=False,
+        # Load custom weights from checkpoint-5000
+        init_weights_from_safetensors=str(CHECKPOINT_PATH),
+    )
+
+    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    # sampling_param.num_frames = 45
+    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    # Generate videos with the same simple API, regardless of GPU count
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
+    # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+    prompt2 = (
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
+
+
+if __name__ == "__main__":
+    main()