| 59 | |
| 60 | |
| 61 | def _parse_args(): |
| 62 | parser = argparse.ArgumentParser( |
| 63 | description="Generate a image or video from a text prompt or image using Wan" |
| 64 | ) |
| 65 | parser.add_argument( |
| 66 | "--task", |
| 67 | type=str, |
| 68 | default="multitalk-14B", |
| 69 | choices=list(WAN_CONFIGS.keys()), |
| 70 | help="The task to run.") |
| 71 | parser.add_argument( |
| 72 | "--size", |
| 73 | type=str, |
| 74 | default="multitalk-480", |
| 75 | choices=list(SIZE_CONFIGS.keys()), |
| 76 | help="The buckget size of the generated video. The aspect ratio of the output video will follow that of the input image." |
| 77 | ) |
| 78 | parser.add_argument( |
| 79 | "--frame_num", |
| 80 | type=int, |
| 81 | default=81, |
| 82 | help="How many frames to be generated in one clip. The number should be 4n+1" |
| 83 | ) |
| 84 | parser.add_argument( |
| 85 | "--ckpt_dir", |
| 86 | type=str, |
| 87 | default='./weights/Wan2.1-I2V-14B-480P', |
| 88 | help="The path to the Wan checkpoint directory.") |
| 89 | parser.add_argument( |
| 90 | "--quant_dir", |
| 91 | type=str, |
| 92 | default=None, |
| 93 | help="The path to the Wan quant checkpoint directory.") |
| 94 | parser.add_argument( |
| 95 | "--wav2vec_dir", |
| 96 | type=str, |
| 97 | default='./weights/chinese-wav2vec2-base', |
| 98 | help="The path to the wav2vec checkpoint directory.") |
| 99 | parser.add_argument( |
| 100 | "--lora_dir", |
| 101 | type=str, |
| 102 | nargs='+', |
| 103 | default=None, |
| 104 | help="The path to the LoRA checkpoint directory.") |
| 105 | parser.add_argument( |
| 106 | "--lora_scale", |
| 107 | type=float, |
| 108 | nargs='+', |
| 109 | default=[1.2], |
| 110 | help="Controls how much to influence the outputs with the LoRA parameters. Accepts multiple float values." |
| 111 | ) |
| 112 | parser.add_argument( |
| 113 | "--offload_model", |
| 114 | type=str2bool, |
| 115 | default=None, |
| 116 | help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage." |
| 117 | ) |
| 118 | parser.add_argument( |