Create download_model.py
Browse files- download_model.py +9 -0
download_model.py
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from diffusers import AutoencoderKLWan, WanPipeline
|
| 3 |
+
|
| 4 |
+
model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
|
| 5 |
+
print("Pre-downloading model weights on CPU...")
|
| 6 |
+
# Load the VAE and pipeline on CPU only
|
| 7 |
+
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
|
| 8 |
+
_ = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
|
| 9 |
+
print("Model weights pre-downloaded.")
|