()
| 103 | |
| 104 | |
| 105 | def _parse_args(): |
| 106 | parser = argparse.ArgumentParser( |
| 107 | description="Generate a image or video from a text prompt or image using Wan" |
| 108 | ) |
| 109 | parser.add_argument( |
| 110 | "--task", |
| 111 | type=str, |
| 112 | default="t2v-A14B", |
| 113 | choices=list(WAN_CONFIGS.keys()), |
| 114 | help="The task to run.") |
| 115 | parser.add_argument( |
| 116 | "--size", |
| 117 | type=str, |
| 118 | default="1280*720", |
| 119 | choices=list(SIZE_CONFIGS.keys()), |
| 120 | help="The area (width*height) of the generated video. For the I2V task, the aspect ratio of the output video will follow that of the input image." |
| 121 | ) |
| 122 | parser.add_argument( |
| 123 | "--frame_num", |
| 124 | type=int, |
| 125 | default=None, |
| 126 | help="How many frames of video are generated. The number should be 4n+1" |
| 127 | ) |
| 128 | parser.add_argument( |
| 129 | "--ckpt_dir", |
| 130 | type=str, |
| 131 | default=None, |
| 132 | help="The path to the checkpoint directory.") |
| 133 | parser.add_argument( |
| 134 | "--offload_model", |
| 135 | type=str2bool, |
| 136 | default=None, |
| 137 | help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage." |
| 138 | ) |
| 139 | parser.add_argument( |
| 140 | "--ulysses_size", |
| 141 | type=int, |
| 142 | default=1, |
| 143 | help="The size of the ulysses parallelism in DiT.") |
| 144 | parser.add_argument( |
| 145 | "--t5_fsdp", |
| 146 | action="store_true", |
| 147 | default=False, |
| 148 | help="Whether to use FSDP for T5.") |
| 149 | parser.add_argument( |
| 150 | "--t5_cpu", |
| 151 | action="store_true", |
| 152 | default=False, |
| 153 | help="Whether to place T5 model on CPU.") |
| 154 | parser.add_argument( |
| 155 | "--dit_fsdp", |
| 156 | action="store_true", |
| 157 | default=False, |
| 158 | help="Whether to use FSDP for DiT.") |
| 159 | parser.add_argument( |
| 160 | "--save_file", |
| 161 | type=str, |
| 162 | default=None, |
no test coverage detected