()
| 840 | |
| 841 | |
| 842 | def main(): |
| 843 | |
| 844 | warnings.filterwarnings('ignore', category=DeprecationWarning) |
| 845 | warnings.filterwarnings('ignore', category=FutureWarning) |
| 846 | torch.multiprocessing.set_start_method('spawn') |
| 847 | |
| 848 | args = parse_args() |
| 849 | |
| 850 | if args.non_ema_revision is not None: |
| 851 | deprecate( |
| 852 | "non_ema_revision!=None", |
| 853 | "0.15.0", |
| 854 | message=( |
| 855 | "Downloading 'non_ema' weights from revision branches of the Hub is deprecated. Please make sure to" |
| 856 | " use `--variant=non_ema` instead." |
| 857 | ), |
| 858 | ) |
| 859 | logging_dir = os.path.join(args.output_dir, args.logging_dir) |
| 860 | accelerator_project_config = ProjectConfiguration( |
| 861 | project_dir=args.output_dir, logging_dir=logging_dir) |
| 862 | # ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) |
| 863 | accelerator = Accelerator( |
| 864 | gradient_accumulation_steps=args.gradient_accumulation_steps, |
| 865 | mixed_precision=args.mixed_precision, |
| 866 | project_config=accelerator_project_config, |
| 867 | ) |
| 868 | |
| 869 | generator = torch.Generator( |
| 870 | device=accelerator.device).manual_seed(23123134) |
| 871 | |
| 872 | if args.report_to == "wandb": |
| 873 | if not is_wandb_available(): |
| 874 | raise ImportError( |
| 875 | "Make sure to install wandb if you want to use it for logging during training.") |
| 876 | import wandb |
| 877 | |
| 878 | # Make one log on every process with the configuration for debugging. |
| 879 | logging.basicConfig( |
| 880 | format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", |
| 881 | datefmt="%m/%d/%Y %H:%M:%S", |
| 882 | level=logging.INFO, |
| 883 | ) |
| 884 | logger.info(accelerator.state, main_process_only=False) |
| 885 | if accelerator.is_local_main_process: |
| 886 | transformers.utils.logging.set_verbosity_warning() |
| 887 | diffusers.utils.logging.set_verbosity_info() |
| 888 | else: |
| 889 | transformers.utils.logging.set_verbosity_error() |
| 890 | diffusers.utils.logging.set_verbosity_error() |
| 891 | |
| 892 | # If passed along, set the training seed now. |
| 893 | if args.seed is not None: |
| 894 | set_seed(args.seed) |
| 895 | |
| 896 | # Handle the repository creation |
| 897 | if accelerator.is_main_process: |
| 898 | if args.output_dir is not None: |
| 899 | os.makedirs(args.output_dir, exist_ok=True) |
no test coverage detected