Update generate.py

generate.py

@@ -1,7 +1,7 @@
 import argparse
 import os
+import subprocess  # Ensure subprocess is imported
 import torch
-import subprocess
 from huggingface_hub import snapshot_download
 
 # Arguments
@@ -11,7 +11,7 @@ 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-1.3B")
-parser.add_argument("--offload_model", type=str, default="True")
+parser.add_argument("--offload_model", type=str, default="True", choices=["True", "False"], help="Whether to offload the model")
 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")