(args)
| 135 | return prompt_embeds |
| 136 | |
| 137 | def main(args): |
| 138 | args.pretrained_model_name_or_path = "ashawkey/imagedream-ipmv-diffusers" |
| 139 | |
| 140 | args.max_train_steps = None |
| 141 | args.num_train_epochs = 100 |
| 142 | args.learning_rate = 1e-4 |
| 143 | args.mixed_precision = 'bf16' |
| 144 | |
| 145 | args.output_dir = "mvd_pretrain" |
| 146 | args.tracker_project_name = "train_mvd_pretrain" |
| 147 | |
| 148 | args.num_gpu = 1 |
| 149 | args.train_batch_size = 1 |
| 150 | args.gradient_accumulation_steps = 1 |
| 151 | args.enable_xformers_memory_efficient_attention = True |
| 152 | |
| 153 | args.resolution = 256 |
| 154 | args.output_dir = args.output_dir +"_bs_"+str(args.train_batch_size * args.num_gpu * args.gradient_accumulation_steps) |
| 155 | args.tracker_project_name = args.tracker_project_name + "_bs_" + str(args.train_batch_size * args.num_gpu * args.gradient_accumulation_steps) |
| 156 | |
| 157 | logging_dir = Path(args.output_dir, args.logging_dir) |
| 158 | |
| 159 | accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) |
| 160 | accelerator = Accelerator( |
| 161 | gradient_accumulation_steps=args.gradient_accumulation_steps, |
| 162 | mixed_precision=args.mixed_precision, |
| 163 | log_with=args.report_to, |
| 164 | project_config=accelerator_project_config, |
| 165 | ) |
| 166 | |
| 167 | logging.basicConfig( |
| 168 | format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
| 169 | datefmt="%m/%d/%Y %H:%M:%S", |
| 170 | level=logging.INFO, |
| 171 | ) |
| 172 | logger.info(accelerator.state, main_process_only=False) |
| 173 | if accelerator.is_local_main_process: |
| 174 | transformers.utils.logging.set_verbosity_warning() |
| 175 | diffusers.utils.logging.set_verbosity_info() |
| 176 | else: |
| 177 | transformers.utils.logging.set_verbosity_error() |
| 178 | diffusers.utils.logging.set_verbosity_error() |
| 179 | |
| 180 | if args.seed is not None: |
| 181 | set_seed(args.seed) |
| 182 | |
| 183 | if accelerator.is_main_process: |
| 184 | if args.output_dir is not None: |
| 185 | os.makedirs(args.output_dir, exist_ok=True) |
| 186 | |
| 187 | if args.push_to_hub: |
| 188 | repo_id = create_repo( |
| 189 | repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token, private=True |
| 190 | ).repo_id |
| 191 | |
| 192 | noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler", revision=None) |
| 193 | image_encoder = CLIPVisionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="image_encoder", revision=None) |
| 194 | text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder", revision=None) |
no test coverage detected