This commit is contained in:
Artiprocher
2025-11-26 14:59:30 +08:00
parent 0336551544
commit 0eead33ed7
22 changed files with 71 additions and 42 deletions

View File

@@ -116,6 +116,7 @@ def wan_parser():
parser.add_argument("--audio_processor_path", type=str, default=None, help="Path to the audio processor. If provided, the processor will be used for Wan2.2-S2V model.")
parser.add_argument("--max_timestep_boundary", type=float, default=1.0, help="Max timestep boundary (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).")
parser.add_argument("--min_timestep_boundary", type=float, default=0.0, help="Min timestep boundary (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).")
parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.")
return parser
@@ -165,7 +166,7 @@ if __name__ == "__main__":
fp8_models=args.fp8_models,
offload_models=args.offload_models,
task=args.task,
device=accelerator.device,
device="cpu" if args.initialize_model_on_cpu else accelerator.device,
max_timestep_boundary=args.max_timestep_boundary,
min_timestep_boundary=args.min_timestep_boundary,
)