(args)
| 592 | |
| 593 | |
| 594 | def main(args): |
| 595 | logging_dir = Path(args.output_dir, args.logging_dir) |
| 596 | |
| 597 | accelerator = Accelerator( |
| 598 | gradient_accumulation_steps=args.gradient_accumulation_steps, |
| 599 | mixed_precision=args.mixed_precision, |
| 600 | log_with=args.report_to, |
| 601 | logging_dir=logging_dir, |
| 602 | ) |
| 603 | |
| 604 | if args.report_to == "wandb": |
| 605 | if not is_wandb_available(): |
| 606 | raise ImportError("Make sure to install wandb if you want to use it for logging during training.") |
| 607 | import wandb |
| 608 | |
| 609 | # Currently, it's not possible to do gradient accumulation when training two models with accelerate.accumulate |
| 610 | # This will be enabled soon in accelerate. For now, we don't allow gradient accumulation when training two models. |
| 611 | # TODO (patil-suraj): Remove this check when gradient accumulation with two models is enabled in accelerate. |
| 612 | # Make one log on every process with the configuration for debugging. |
| 613 | logging.basicConfig( |
| 614 | format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
| 615 | datefmt="%m/%d/%Y %H:%M:%S", |
| 616 | level=logging.INFO, |
| 617 | ) |
| 618 | logger.info(accelerator.state, main_process_only=False) |
| 619 | if accelerator.is_local_main_process: |
| 620 | transformers.utils.logging.set_verbosity_warning() |
| 621 | diffusers.utils.logging.set_verbosity_info() |
| 622 | else: |
| 623 | transformers.utils.logging.set_verbosity_error() |
| 624 | diffusers.utils.logging.set_verbosity_error() |
| 625 | |
| 626 | if args.seed is not None: |
| 627 | set_seed(args.seed) |
| 628 | if args.concepts_list is None: |
| 629 | args.concepts_list = [ |
| 630 | { |
| 631 | "instance_prompt": args.instance_prompt, |
| 632 | "class_prompt": args.class_prompt, |
| 633 | "instance_data_dir": args.instance_data_dir, |
| 634 | "class_data_dir": args.class_data_dir |
| 635 | } |
| 636 | ] |
| 637 | else: |
| 638 | with open(args.concepts_list, "r") as f: |
| 639 | args.concepts_list = json.load(f) |
| 640 | |
| 641 | if args.with_prior_preservation: |
| 642 | for i, concept in enumerate(args.concepts_list): |
| 643 | class_images_dir = Path(concept['class_data_dir']) |
| 644 | if not class_images_dir.exists(): |
| 645 | class_images_dir.mkdir(parents=True, exist_ok=True) |
| 646 | if args.real_prior: |
| 647 | if accelerator.is_main_process: |
| 648 | if not Path(os.path.join(class_images_dir, 'images')).exists() or len(list(Path(os.path.join(class_images_dir, 'images')).iterdir())) < args.num_class_images: |
| 649 | retrieve.retrieve(concept['class_prompt'], class_images_dir, args.num_class_images) |
| 650 | concept['class_prompt'] = os.path.join(class_images_dir, 'caption.txt') |
| 651 | concept['class_data_dir'] = os.path.join(class_images_dir, 'images.txt') |
no test coverage detected