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