()
| 97 | |
| 98 | |
| 99 | def _parse_args(): |
| 100 | parser = argparse.ArgumentParser( |
| 101 | description="Generate a image or video from a text prompt or image using Wan" |
| 102 | ) |
| 103 | parser.add_argument( |
| 104 | "--task", |
| 105 | type=str, |
| 106 | default="t2v-14B", |
| 107 | choices=list(WAN_CONFIGS.keys()), |
| 108 | help="The task to run.") |
| 109 | parser.add_argument( |
| 110 | "--size", |
| 111 | type=str, |
| 112 | default="1280*720", |
| 113 | choices=list(SIZE_CONFIGS.keys()), |
| 114 | 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." |
| 115 | ) |
| 116 | parser.add_argument( |
| 117 | "--frame_num", |
| 118 | type=int, |
| 119 | default=None, |
| 120 | help="How many frames to sample from a image or video. The number should be 4n+1" |
| 121 | ) |
| 122 | parser.add_argument( |
| 123 | "--ckpt_dir", |
| 124 | type=str, |
| 125 | default=None, |
| 126 | help="The path to the checkpoint directory.") |
| 127 | parser.add_argument( |
| 128 | "--offload_model", |
| 129 | type=str2bool, |
| 130 | default=None, |
| 131 | help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage." |
| 132 | ) |
| 133 | parser.add_argument( |
| 134 | "--ulysses_size", |
| 135 | type=int, |
| 136 | default=1, |
| 137 | help="The size of the ulysses parallelism in DiT.") |
| 138 | parser.add_argument( |
| 139 | "--ring_size", |
| 140 | type=int, |
| 141 | default=1, |
| 142 | help="The size of the ring attention parallelism in DiT.") |
| 143 | parser.add_argument( |
| 144 | "--t5_fsdp", |
| 145 | action="store_true", |
| 146 | default=False, |
| 147 | help="Whether to use FSDP for T5.") |
| 148 | parser.add_argument( |
| 149 | "--t5_cpu", |
| 150 | action="store_true", |
| 151 | default=False, |
| 152 | help="Whether to place T5 model on CPU.") |
| 153 | parser.add_argument( |
| 154 | "--dit_fsdp", |
| 155 | action="store_true", |
| 156 | default=False, |
no test coverage detected