(args)
| 479 | |
| 480 | |
| 481 | def main(args): |
| 482 | logging_dir = Path(args.output_dir, args.logging_dir) |
| 483 | |
| 484 | accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) |
| 485 | kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) |
| 486 | accelerator = Accelerator( |
| 487 | gradient_accumulation_steps=args.gradient_accumulation_steps, |
| 488 | mixed_precision=args.mixed_precision, |
| 489 | log_with=args.report_to, |
| 490 | project_config=accelerator_project_config, |
| 491 | kwargs_handlers=[kwargs], |
| 492 | ) |
| 493 | |
| 494 | if args.report_to == "wandb": |
| 495 | if not is_wandb_available(): |
| 496 | raise ImportError("Make sure to install wandb if you want to use it for logging during training.") |
| 497 | import wandb |
| 498 | |
| 499 | # Make one log on every process with the configuration for debugging. |
| 500 | logging.basicConfig( |
| 501 | format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
| 502 | datefmt="%m/%d/%Y %H:%M:%S", |
| 503 | level=logging.INFO, |
| 504 | ) |
| 505 | logger.info(accelerator.state, main_process_only=False) |
| 506 | if accelerator.is_local_main_process: |
| 507 | datasets.utils.logging.set_verbosity_warning() |
| 508 | transformers.utils.logging.set_verbosity_warning() |
| 509 | diffusers.utils.logging.set_verbosity_info() |
| 510 | else: |
| 511 | datasets.utils.logging.set_verbosity_error() |
| 512 | transformers.utils.logging.set_verbosity_error() |
| 513 | diffusers.utils.logging.set_verbosity_error() |
| 514 | |
| 515 | # If passed along, set the training seed now. |
| 516 | if args.seed is not None: |
| 517 | set_seed(args.seed) |
| 518 | |
| 519 | # Handle the repository creation |
| 520 | if accelerator.is_main_process: |
| 521 | if args.output_dir is not None: |
| 522 | os.makedirs(args.output_dir, exist_ok=True) |
| 523 | |
| 524 | if args.push_to_hub: |
| 525 | repo_id = create_repo( |
| 526 | repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token |
| 527 | ).repo_id |
| 528 | |
| 529 | # Load the tokenizers |
| 530 | tokenizer_one = AutoTokenizer.from_pretrained( |
| 531 | args.pretrained_model_name_or_path, |
| 532 | subfolder="tokenizer", |
| 533 | revision=args.revision, |
| 534 | use_fast=False, |
| 535 | ) |
| 536 | tokenizer_two = AutoTokenizer.from_pretrained( |
| 537 | args.pretrained_model_name_or_path, |
| 538 | subfolder="tokenizer_2", |
no test coverage detected