()
| 494 | |
| 495 | |
| 496 | def main(): |
| 497 | args = parse_args() |
| 498 | |
| 499 | if args.non_ema_revision is not None: |
| 500 | deprecate( |
| 501 | "non_ema_revision!=None", |
| 502 | "0.15.0", |
| 503 | message=( |
| 504 | "Downloading 'non_ema' weights from revision branches of the Hub is deprecated. Please make sure to" |
| 505 | " use `--variant=non_ema` instead." |
| 506 | ), |
| 507 | ) |
| 508 | logging_dir = os.path.join(args.output_dir, args.logging_dir) |
| 509 | |
| 510 | accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) |
| 511 | |
| 512 | accelerator = Accelerator( |
| 513 | gradient_accumulation_steps=args.gradient_accumulation_steps, |
| 514 | mixed_precision=args.mixed_precision, |
| 515 | log_with=args.report_to, |
| 516 | project_config=accelerator_project_config, |
| 517 | ) |
| 518 | |
| 519 | # Make one log on every process with the configuration for debugging. |
| 520 | logging.basicConfig( |
| 521 | format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
| 522 | datefmt="%m/%d/%Y %H:%M:%S", |
| 523 | level=logging.INFO, |
| 524 | ) |
| 525 | logger.info(accelerator.state, main_process_only=False) |
| 526 | if accelerator.is_local_main_process: |
| 527 | datasets.utils.logging.set_verbosity_warning() |
| 528 | transformers.utils.logging.set_verbosity_warning() |
| 529 | diffusers.utils.logging.set_verbosity_info() |
| 530 | else: |
| 531 | datasets.utils.logging.set_verbosity_error() |
| 532 | transformers.utils.logging.set_verbosity_error() |
| 533 | diffusers.utils.logging.set_verbosity_error() |
| 534 | |
| 535 | # If passed along, set the training seed now. |
| 536 | if args.seed is not None: |
| 537 | set_seed(args.seed) |
| 538 | |
| 539 | # Handle the repository creation |
| 540 | if accelerator.is_main_process: |
| 541 | if args.output_dir is not None: |
| 542 | os.makedirs(args.output_dir, exist_ok=True) |
| 543 | |
| 544 | if args.push_to_hub: |
| 545 | repo_id = create_repo( |
| 546 | repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token |
| 547 | ).repo_id |
| 548 | |
| 549 | # Load scheduler, tokenizer and models. |
| 550 | noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler") |
| 551 | |
| 552 | def deepspeed_zero_init_disabled_context_manager(): |
| 553 | """ |
no test coverage detected