()
| 377 | |
| 378 | |
| 379 | def main(): |
| 380 | args = parse_args() |
| 381 | logging_dir = Path(args.output_dir, args.logging_dir) |
| 382 | |
| 383 | accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) |
| 384 | |
| 385 | accelerator = Accelerator( |
| 386 | gradient_accumulation_steps=args.gradient_accumulation_steps, |
| 387 | mixed_precision=args.mixed_precision, |
| 388 | log_with=args.report_to, |
| 389 | project_config=accelerator_project_config, |
| 390 | ) |
| 391 | if args.report_to == "wandb": |
| 392 | if not is_wandb_available(): |
| 393 | raise ImportError("Make sure to install wandb if you want to use it for logging during training.") |
| 394 | import wandb |
| 395 | |
| 396 | # Make one log on every process with the configuration for debugging. |
| 397 | logging.basicConfig( |
| 398 | format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
| 399 | datefmt="%m/%d/%Y %H:%M:%S", |
| 400 | level=logging.INFO, |
| 401 | ) |
| 402 | logger.info(accelerator.state, main_process_only=False) |
| 403 | if accelerator.is_local_main_process: |
| 404 | datasets.utils.logging.set_verbosity_warning() |
| 405 | transformers.utils.logging.set_verbosity_warning() |
| 406 | diffusers.utils.logging.set_verbosity_info() |
| 407 | else: |
| 408 | datasets.utils.logging.set_verbosity_error() |
| 409 | transformers.utils.logging.set_verbosity_error() |
| 410 | diffusers.utils.logging.set_verbosity_error() |
| 411 | |
| 412 | # If passed along, set the training seed now. |
| 413 | if args.seed is not None: |
| 414 | set_seed(args.seed) |
| 415 | |
| 416 | # Handle the repository creation |
| 417 | if accelerator.is_main_process: |
| 418 | if args.output_dir is not None: |
| 419 | os.makedirs(args.output_dir, exist_ok=True) |
| 420 | |
| 421 | if args.push_to_hub: |
| 422 | repo_id = create_repo( |
| 423 | repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token |
| 424 | ).repo_id |
| 425 | # Load scheduler, tokenizer and models. |
| 426 | noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") |
| 427 | tokenizer = CLIPTokenizer.from_pretrained( |
| 428 | args.pretrained_model_name_or_path, subfolder="tokenizer", revision=args.revision |
| 429 | ) |
| 430 | text_encoder = CLIPTextModel.from_pretrained( |
| 431 | args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision |
| 432 | ) |
| 433 | vae = AutoencoderKL.from_pretrained( |
| 434 | args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision, variant=args.variant |
| 435 | ) |
| 436 | unet = UNet2DConditionModel.from_pretrained( |
no test coverage detected