acmyu/KeyframesAI

1import argparse2parser = argparse.ArgumentParser(description="Simple example of a ControlNet training script.")3parser.add_argument(4    "--pretrained_model_name_or_path",5    type=str,6    default=None,7    required=True,8    help="Path to pretrained model or model identifier from huggingface.co/models.",)9 10parser.add_argument("--revision",type=str,default=None,required=False,help=(11        "Revision of pretrained model identifier from huggingface.co/models. Trainable model components should be"12        " float32 precision."),)13parser.add_argument("--json_path", type=str, default="./datasets/deepfashing/test_data.json", help="json path", )14parser.add_argument("--img_path", type=str, default="./datasets/deepfashing/train_all_png/", help="image path", )15parser.add_argument("--gen_t_img_path", type=str,default="./save_data/stage2/guidancescale2_seed42_numsteps20/",help="gen target image path", )16parser.add_argument("--image_encoder_path", type=str, default="./dinov2-giant",17                    help="Path to pretrained model or model identifier from huggingface.co/models.", )18parser.add_argument("--output_dir",type=str,default="controlnet-model",help="The output directory where the model predictions and checkpoints will be written.",)19 20 21parser.add_argument(22    "--seed", type=int, default=None, help="A seed for reproducible training."23)24parser.add_argument(25    "--resolution",26    type=int,27    default=512,28    help=(29        "The resolution for input images, all the images in the train/validation dataset will be resized to this"30        " resolution"31    ),32)33parser.add_argument(34    "--train_batch_size",35    type=int,36    default=4,37    help="Batch size (per device) for the training dataloader.",38)39parser.add_argument("--num_train_epochs", type=int, default=1)40parser.add_argument("--noise_level", type=int, default=250)41 42parser.add_argument(43    "--max_train_steps",44    type=int,45    default=None,46    help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",47)48parser.add_argument(49    "--checkpointing_steps",50    type=int,51    default=500,52    help=(53        "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "54        "In the case that the checkpoint is better than the final trained model, the checkpoint can also be used for inference."55        "Using a checkpoint for inference requires separate loading of the original pipeline and the individual checkpointed model components."56        "See https://huggingface.co/docs/diffusers/main/en/training/dreambooth#performing-inference-using-a-saved-checkpoint for step by step"57        "instructions."58    ),59)60 61parser.add_argument(62    "--resume_from_checkpoint",63    type=str,64    default=None,65    help=(66        "Whether training should be resumed from a previous checkpoint. Use a path saved by"67        ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'68    ),69)70parser.add_argument(71    "--gradient_accumulation_steps",72    type=int,73    default=1,74    help="Number of updates steps to accumulate before performing a backward/update pass.",75)76parser.add_argument(77    "--gradient_checkpointing",78    action="store_true",79    help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",80)81parser.add_argument(82    "--learning_rate",83    type=float,84    default=5e-6,85    help="Initial learning rate (after the potential warmup period) to use.",86)87 88parser.add_argument(89    "--scale_lr",90    action="store_true",91    default=False,92    help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",93)94parser.add_argument(95    "--lr_scheduler",96    type=str,97    default="constant",98    help=(99        'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'100        ' "constant", "constant_with_warmup"]'101    ),102)103parser.add_argument(104    "--lr_warmup_steps",105    type=int,106    default=500,107    help="Number of steps for the warmup in the lr scheduler.",108)109parser.add_argument(110    "--lr_num_cycles",111    type=int,112    default=1,113    help="Number of hard resets of the lr in cosine_with_restarts scheduler.",114)115parser.add_argument(116    "--lr_power",117    type=float,118    default=1.0,119    help="Power factor of the polynomial scheduler.",120)121 122 123parser.add_argument(124    "--adam_beta1",125    type=float,126    default=0.9,127    help="The beta1 parameter for the Adam optimizer.",128)129parser.add_argument(130    "--adam_beta2",131    type=float,132    default=0.999,133    help="The beta2 parameter for the Adam optimizer.",134)135parser.add_argument(136    "--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use."137)138parser.add_argument(139    "--adam_epsilon",140    type=float,141    default=1e-08,142    help="Epsilon value for the Adam optimizer",143)144parser.add_argument(145    "--max_grad_norm", default=1.0, type=float, help="Max gradient norm."146)147 148 149parser.add_argument(150    "--logging_dir",151    type=str,152    default="logs",153    help=(154        "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"155        " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."156    ),157)158parser.add_argument(159    "--allow_tf32",160    action="store_true",161    help=(162        "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"163        " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"164    ),165)166parser.add_argument(167    "--report_to",168    type=str,169    default="tensorboard",170    help=(171        'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'172        ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'173    ),174)175parser.add_argument(176    "--mixed_precision",177    type=str,178    default=None,179    choices=["no", "fp16", "bf16"],180    help=(181        "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="182        " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"183        " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."184    ),185)186 187parser.add_argument(188    "--set_grads_to_none",189    action="store_true",190    help=(191        "Save more memory by using setting grads to None instead of zero. Be aware, that this changes certain"192        " behaviors, so disable this argument if it causes any problems. More info:"193        " https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html"194    ),195)196 197parser.add_argument("--noise_offset", type=float, default=0, help="The scale of noise offset.")198 199parser.add_argument(200    "--tracker_project_name",201    type=str,202    default="train_baseline",203    help=(204        "The `project_name` argument passed to Accelerator.init_trackers for"205        " more information see https://huggingface.co/docs/accelerate/v0.17.0/en/package_reference/accelerator#accelerate.Accelerator"206    ),207)208 209 210args = parser.parse_args()211print(args)212if args.resolution % 8 != 0:213    raise ValueError(214        "`--resolution` must be divisible by 8 for consistently sized encoded images between the VAE and the controlnet encoder."215    )216 217 218