Update generate.py

generate.py

@@ -1,24 +1,37 @@
 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("--task", type=str, default="t2v-1.3B")
+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=20)
-parser.add_argument("--ckpt_dir", type=str, default="./Wan2.1-T2V-14B")
+parser.add_argument("--ckpt_dir", type=str, default="./Wan2.1-T2V-1.3B")
 parser.add_argument("--offload_model", type=str, default="True")
+parser.add_argument("--t5_cpu", action="store_true", help="Use CPU for T5 model (optional)")
+parser.add_argument("--sample_shift", type=int, default=8, help="Sampling shift for generation")
+parser.add_argument("--sample_guide_scale", type=int, default=6, help="Sampling guide scale for generation")
 parser.add_argument("--prompt", type=str, required=True)
 args = parser.parse_args()
-print("🔄 Downloading WAN 2.1 start..")
+
+# Log input parameters
+print(f"Generating video with the following settings:\n"
+      f"Task: {args.task}\n"
+      f"Resolution: {args.size}\n"
+      f"Frames: {args.frame_num}\n"
+      f"Sample Steps: {args.sample_steps}\n"
+      f"Prompt: {args.prompt}\n"
+      f"Sample Shift: {args.sample_shift}\n"
+      f"Sample Guide Scale: {args.sample_guide_scale}\n"
+      f"Using T5 on CPU: {args.t5_cpu}")
+
 # 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)
+    print("🔄 Downloading WAN 2.1 - 1.3B model from Hugging Face...")
+    snapshot_download(repo_id="Wan-AI/Wan2.1-T2V-1.3B", local_dir=args.ckpt_dir)
 
 # Free up GPU memory
 if torch.cuda.is_available():
@@ -26,9 +39,8 @@ if torch.cuda.is_available():
     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}\""
-command = f"python generate.py  --task t2v-1.3B --size 832*480 --ckpt_dir ./Wan2.1-T2V-1.3B --offload_model True --t5_cpu --sample_shift 8 --sample_guide_scale 6 --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
+# Run the model (Ensure that `generate.py` includes these new params in its model call)
+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} --t5_cpu {args.t5_cpu} --sample_shift {args.sample_shift} --sample_guide_scale {args.sample_guide_scale} --prompt \"{args.prompt}\""
 
 process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 stdout, stderr = process.communicate()
@@ -41,4 +53,4 @@ print("🔺 Error:", stderr.decode())
 if os.path.exists("output.mp4"):
     print("✅ Video generated successfully: output.mp4")
 else:
-    print("❌ Error: Video file not found!")
+    print("❌ Error: Video file not found!")