Create generate.py

generate.py

@@ -0,0 +1,43 @@
+import argparse
+import subprocess
+import os
+import torch
+from huggingface_hub import snapshot_download
+
+# Arguments
+parser = argparse.ArgumentParser()
+parser.add_argument("--task", type=str, default="t2v-14B")
+parser.add_argument("--size", type=str, default="200*200")
+parser.add_argument("--frame_num", type=int, default=60)
+parser.add_argument("--sample_steps", type=int, default=20)
+parser.add_argument("--ckpt_dir", type=str, default="./Wan2.1-T2V-14B")
+parser.add_argument("--offload_model", type=str, default="True")
+parser.add_argument("--prompt", type=str, required=True)
+args = parser.parse_args()
+
+# Ensure the model is downloaded
+if not os.path.exists(args.ckpt_dir):
+    print("🔄 Downloading WAN 2.1 - 14B model from Hugging Face...")
+    snapshot_download(repo_id="Wan-AI/Wan2.1-T2V-14B", local_dir=args.ckpt_dir)
+
+# Free up GPU memory
+if torch.cuda.is_available():
+    torch.cuda.empty_cache()
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
+# Run WAN 2.1 - 14B Model
+command = f"python generate.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} --prompt \"{args.prompt}\""
+
+process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+stdout, stderr = process.communicate()
+
+# Print logs for debugging
+print("🔹 Output:", stdout.decode())
+print("🔺 Error:", stderr.decode())
+
+# Verify if video was created
+if os.path.exists("output.mp4"):
+    print("✅ Video generated successfully: output.mp4")
+else:
+    print("❌ Error: Video file not found!")