| 109 | |
| 110 | |
| 111 | def parse_args(): |
| 112 | parser = argparse.ArgumentParser(description="training script of jodiffusion.") |
| 113 | parser.add_argument( |
| 114 | "--pretrained_model_name_or_path", |
| 115 | type=str, |
| 116 | default='inference/saved_pipeline/jodiffusion', |
| 117 | help="Path to pretrained model or model identifier from huggingface.co/models.", |
| 118 | ) |
| 119 | parser.add_argument( |
| 120 | "--pretrained_label_vae_path", |
| 121 | type=str, |
| 122 | default=None, |
| 123 | help="Path to pretrained label vae model.", |
| 124 | ) |
| 125 | parser.add_argument( |
| 126 | "--dataset_name", |
| 127 | type=str, |
| 128 | default="ade20k_semantic", |
| 129 | help="The name of the dataset to use for training.", |
| 130 | ) |
| 131 | parser.add_argument( |
| 132 | "--caption_column", |
| 133 | type=str, |
| 134 | default="blip2", |
| 135 | choices=["category_name", "blip2", "coco"], |
| 136 | help="The name of the column in the dataset that contains the captions.", |
| 137 | ) |
| 138 | parser.add_argument( |
| 139 | "--lightweight_label_vae", |
| 140 | action="store_true", |
| 141 | help="Whether or not to use lightweight vae.", |
| 142 | ) |
| 143 | parser.add_argument( |
| 144 | "--noise_type", |
| 145 | type=str, |
| 146 | default="image_only", |
| 147 | choices=["image_only", "joint"], |
| 148 | help="The type of noise to use for training.", |
| 149 | ) |
| 150 | parser.add_argument( |
| 151 | "--output_dir", |
| 152 | type=str, |
| 153 | default=None, |
| 154 | help="The output directory where the model predictions and checkpoints will be written.", |
| 155 | ) |
| 156 | parser.add_argument( |
| 157 | "--seed", |
| 158 | type=int, |
| 159 | default=42, |
| 160 | help="A seed for reproducible training." |
| 161 | ) |
| 162 | parser.add_argument( |
| 163 | "--resolution", |
| 164 | type=int, |
| 165 | default=512, |
| 166 | help="The resolution for input images, all the images in the train/validation dataset will be resized to this" |
| 167 | " resolution", |
| 168 | ) |