()
| 64 | |
| 65 | |
| 66 | def main(): |
| 67 | cfgs = parse_cfgs() |
| 68 | if cfgs.non_ema_revision is not None: |
| 69 | deprecate( |
| 70 | "non_ema_revision!=None", |
| 71 | "0.15.0", |
| 72 | message=( |
| 73 | "Downloading 'non_ema' weights from revision branches of the Hub is deprecated. Please make sure to" |
| 74 | " use `--variant=non_ema` instead." |
| 75 | ), |
| 76 | ) |
| 77 | |
| 78 | logging_dir = os.path.join(cfgs.output_dir, cfgs.logging_dir) |
| 79 | |
| 80 | accelerator_project_config = ProjectConfiguration(total_limit=cfgs.checkpoints_total_limit) |
| 81 | |
| 82 | accelerator = Accelerator( |
| 83 | gradient_accumulation_steps=cfgs.gradient_accumulation_steps, |
| 84 | mixed_precision=cfgs.mixed_precision, |
| 85 | log_with=cfgs.report_to, |
| 86 | project_dir=logging_dir, |
| 87 | project_config=accelerator_project_config, |
| 88 | ) |
| 89 | |
| 90 | # Make one log on every process with the configuration for debugging. |
| 91 | logging.basicConfig( |
| 92 | format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
| 93 | datefmt="%m/%d/%Y %H:%M:%S", |
| 94 | level=logging.INFO, |
| 95 | ) |
| 96 | |
| 97 | logger.parent = None |
| 98 | logger.info(accelerator.state, main_process_only=False) |
| 99 | if accelerator.is_local_main_process: |
| 100 | datasets.utils.logging.set_verbosity_warning() |
| 101 | transformers.utils.logging.set_verbosity_warning() |
| 102 | diffusers.utils.logging.set_verbosity_info() |
| 103 | else: |
| 104 | datasets.utils.logging.set_verbosity_error() |
| 105 | transformers.utils.logging.set_verbosity_error() |
| 106 | diffusers.utils.logging.set_verbosity_error() |
| 107 | |
| 108 | # If passed along, set the training seed now. |
| 109 | if cfgs.seed is not None: |
| 110 | set_seed(cfgs.seed) |
| 111 | |
| 112 | # Handle the repository creation |
| 113 | if accelerator.is_main_process: |
| 114 | if cfgs.push_to_hub: |
| 115 | if cfgs.hub_model_id is None: |
| 116 | repo_name = get_full_repo_name(Path(cfgs.output_dir).name, token=cfgs.hub_token) |
| 117 | else: |
| 118 | repo_name = cfgs.hub_model_id |
| 119 | create_repo(repo_name, exist_ok=True, token=cfgs.hub_token) |
| 120 | repo = Repository(cfgs.output_dir, clone_from=repo_name, token=cfgs.hub_token) |
| 121 | |
| 122 | with open(os.path.join(cfgs.output_dir, ".gitignore"), "w+") as gitignore: |
| 123 | if "step_*" not in gitignore: |
no test coverage detected