Update generate.py

generate.py

@@ -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", choices=["True", "False"], help="Whether to offload the model")
+parser.add_argument("--offload_model", type=bool, default=True, help="Whether to offload the model (True/False)")
 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")
@@ -27,7 +27,8 @@ print(f"Generating video with the following settings:\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}")
+      f"Using T5 on CPU: {args.t5_cpu}\n"
+      f"Offload Model: {args.offload_model}")
 
 # Ensure the model is downloaded
 if not os.path.exists(args.ckpt_dir):
@@ -41,7 +42,9 @@ if torch.cuda.is_available():
     torch.backends.cudnn.deterministic = True
 
 # 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}\""
+# Pass offload_model correctly as either True or False
+offload_model_value = "True" if args.offload_model else "False"
+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 {offload_model_value} --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()