Update generate.py

generate.py

@@ -1,36 +1,36 @@
 import argparse
-import torch
 import subprocess
 import os
+from huggingface_hub import snapshot_download
 
-# Define Arguments
+# Arguments
 parser = argparse.ArgumentParser()
-parser.add_argument("--task", type=str, default="t2v-1.3B")
+parser.add_argument("--task", type=str, default="t2v-14B")
 parser.add_argument("--size", type=str, default="832*480")
 parser.add_argument("--frame_num", type=int, default=60)
-parser.add_argument("--sample_steps", type=int, default=30)
-parser.add_argument("--ckpt_dir", type=str, default="./Wan2.1-T2V-1.3B")
+parser.add_argument("--sample_steps", type=int, default=25)
+parser.add_argument("--ckpt_dir", type=str, default="./Wan2.1-T2V-14B")
+parser.add_argument("--offload_model", type=str, default="True")
+parser.add_argument("--precision", type=str, default="bf16")
 parser.add_argument("--prompt", type=str, required=True)
 args = parser.parse_args()
 
-# Check GPU Availability
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using device: {device}")
+# Download model if not available
+if not os.path.exists(args.ckpt_dir):
+    print("🔄 Downloading WAN 2.1 model...")
+    snapshot_download(repo_id="Wan-AI/Wan2.1-T2V-14B", local_dir=args.ckpt_dir)
 
-# Run WAN 2.1 Inference
-command = f"python run_model.py --task {args.task} --size {args.size} --frame_num {args.frame_num} --sample_steps {args.sample_steps} --ckpt_dir {args.ckpt_dir} --prompt \"{args.prompt}\" --device {device}"
-
-subprocess.run(command, shell=True)
+# Run Model
+command = f"python run_model.py --task {args.task} --size {args.size} --frame_num {args.frame_num} --sample_steps {args.sample_steps} --ckpt_dir {args.ckpt_dir} --offload_model {args.offload_model} --precision {args.precision} --prompt \"{args.prompt}\""
 
 process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 stdout, stderr = process.communicate()
 
-print("Output:", stdout.decode())
-print("Error:", stderr.decode())
-
+print("🔹 Output:", stdout.decode())
+print("🔺 Error:", stderr.decode())
 
-# Save output
+# Check if video was created
 if os.path.exists("output.mp4"):
     print("✅ Video generated successfully: output.mp4")
 else:
-    print("❌ Error generating video.")
+    print("❌ Error: Video file not found!")