372 lines
11 KiB
Python
372 lines
11 KiB
Python
import argparse
|
|
import os
|
|
|
|
def parse_args(input_args=None):
|
|
parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
|
parser.add_argument(
|
|
"--pretrained_model_name_or_path",
|
|
type=str,
|
|
default=None,
|
|
required=True,
|
|
help="Path to pretrained model or model identifier from huggingface.co/models.",
|
|
)
|
|
parser.add_argument(
|
|
"--pretrained_vae_name_or_path",
|
|
type=str,
|
|
default=None,
|
|
help="Path to pretrained vae or vae identifier from huggingface.co/models.",
|
|
)
|
|
parser.add_argument(
|
|
"--revision",
|
|
type=str,
|
|
default=None,
|
|
required=False,
|
|
help="Revision of pretrained model identifier from huggingface.co/models.",
|
|
)
|
|
parser.add_argument(
|
|
"--tokenizer_name",
|
|
type=str,
|
|
default=None,
|
|
help="Pretrained tokenizer name or path if not the same as model_name",
|
|
)
|
|
parser.add_argument(
|
|
"--instance_data_dir",
|
|
type=str,
|
|
default=None,
|
|
required=True,
|
|
help="A folder containing the training data of instance images.",
|
|
)
|
|
parser.add_argument(
|
|
"--class_data_dir",
|
|
type=str,
|
|
default=None,
|
|
required=False,
|
|
help="A folder containing the training data of class images.",
|
|
)
|
|
parser.add_argument(
|
|
"--json_path",
|
|
type=str,
|
|
default="",
|
|
required=True,
|
|
help="A JSON file with the same args as argparse (instance_data_dir, class_data_dir, etc.)",
|
|
)
|
|
parser.add_argument(
|
|
"--instance_prompt",
|
|
type=str,
|
|
default=None,
|
|
required=True,
|
|
help="The prompt with identifier specifying the instance",
|
|
)
|
|
parser.add_argument(
|
|
"--preview_prompt",
|
|
type=str,
|
|
default=None,
|
|
help="The prompt to use when generating preview images",
|
|
)
|
|
parser.add_argument(
|
|
"--class_prompt",
|
|
type=str,
|
|
default=None,
|
|
help="The prompt to specify images in the same class as provided instance images.",
|
|
)
|
|
parser.add_argument(
|
|
"--with_prior_preservation",
|
|
default=False,
|
|
action="store_true",
|
|
help="Flag to add prior preservation loss.",
|
|
)
|
|
parser.add_argument(
|
|
"--prior_loss_weight",
|
|
type=float,
|
|
default=1.0,
|
|
help="The weight of prior preservation loss.",
|
|
)
|
|
parser.add_argument(
|
|
"--prior_preservation_mode",
|
|
type=str,
|
|
choices=["additive", "multiply", "single_pass", "text"],
|
|
default="multiply",
|
|
help=("The prior preservation loss mode."
|
|
"Additive: loss + (prior_loss * loss_weight)"
|
|
"Multiply: loss + loss_weight * prior_loss",
|
|
"Text: The class prompt is used as the initializer"
|
|
"Single Pass:" "Compute the losses together"
|
|
)
|
|
)
|
|
parser.add_argument(
|
|
"--num_class_images",
|
|
type=int,
|
|
default=800,
|
|
help=(
|
|
"Minimal class images for prior preservation loss. If not have enough images, additional images will be"
|
|
" sampled with class_prompt."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--output_dir",
|
|
type=str,
|
|
default="text-inversion-model",
|
|
help="The output directory where the model predictions and checkpoints will be written.",
|
|
)
|
|
parser.add_argument(
|
|
"--seed", type=int, default=None, help="A seed for reproducible training."
|
|
)
|
|
parser.add_argument(
|
|
"--resolution",
|
|
type=int,
|
|
default=512,
|
|
help=(
|
|
"The resolution for input images, all the images in the train/validation dataset will be resized to this"
|
|
" resolution"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--center_crop",
|
|
action="store_true",
|
|
help="Whether to center crop images before resizing to resolution",
|
|
)
|
|
parser.add_argument(
|
|
"--color_jitter",
|
|
action="store_true",
|
|
help="Whether to apply color jitter to images",
|
|
)
|
|
parser.add_argument(
|
|
"--train_text_encoder",
|
|
action="store_true",
|
|
help="Whether to train the text encoder",
|
|
)
|
|
parser.add_argument(
|
|
"--clip_layers",
|
|
type=int,
|
|
default=12,
|
|
help="Amount of hidden layers to include in CLIP (Also known as CLIP Skip / Penultimate)",
|
|
)
|
|
parser.add_argument(
|
|
"--train_batch_size",
|
|
type=int,
|
|
default=1,
|
|
help="Batch size (per device) for the training dataloader.",
|
|
)
|
|
parser.add_argument(
|
|
"--sample_batch_size",
|
|
type=int,
|
|
default=1,
|
|
help="Batch size (per device) for sampling images.",
|
|
)
|
|
parser.add_argument("--num_train_epochs", type=int, default=1)
|
|
parser.add_argument(
|
|
"--max_train_steps",
|
|
type=int,
|
|
default=None,
|
|
help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
|
|
)
|
|
parser.add_argument(
|
|
"--save_steps",
|
|
type=int,
|
|
default=500,
|
|
help="Save checkpoint every X updates steps.",
|
|
)
|
|
parser.add_argument(
|
|
"--save_for_webui",
|
|
action="store_true",
|
|
default=True,
|
|
help="Save a LoRA model for usage in the AUTOMATIC1111 webui.",
|
|
)
|
|
parser.add_argument(
|
|
"--preview_steps",
|
|
type=int,
|
|
default=100,
|
|
help="Save preview every X updates steps.",
|
|
)
|
|
parser.add_argument(
|
|
"--gradient_accumulation_steps",
|
|
type=int,
|
|
default=1,
|
|
help="Number of updates steps to accumulate before performing a backward/update pass.",
|
|
)
|
|
parser.add_argument(
|
|
"--gradient_checkpointing",
|
|
action="store_true",
|
|
help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
|
|
)
|
|
parser.add_argument(
|
|
"--lora_rank",
|
|
type=int,
|
|
default=4,
|
|
help="Rank of LoRA approximation.",
|
|
)
|
|
parser.add_argument(
|
|
"--lora_bias",
|
|
type=str,
|
|
default="none",
|
|
help="Whether or not to use bias when training LoRA.",
|
|
choices=["none", "lora_only", "all"]
|
|
)
|
|
parser.add_argument(
|
|
"--learning_rate",
|
|
type=float,
|
|
default=None,
|
|
help="Initial learning rate (after the potential warmup period) to use.",
|
|
)
|
|
parser.add_argument(
|
|
"--learning_rate_text",
|
|
type=float,
|
|
default=5e-6,
|
|
help="Initial learning rate for text encoder (after the potential warmup period) to use.",
|
|
)
|
|
parser.add_argument(
|
|
"--dropout",
|
|
type=float,
|
|
default=0,
|
|
help="Dropout for both UNET and Text Encoder"
|
|
)
|
|
parser.add_argument(
|
|
"--scale_lr",
|
|
action="store_true",
|
|
default=False,
|
|
help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
|
|
)
|
|
parser.add_argument(
|
|
"--dataset_norm",
|
|
action="store_true",
|
|
default=False,
|
|
help="Normalizes the entire dataset by calculating the mean and standard deviation of all elements.",
|
|
)
|
|
parser.add_argument(
|
|
"--save_preview",
|
|
action="store_true",
|
|
default=False,
|
|
help="Save preview images during training.",
|
|
)
|
|
parser.add_argument(
|
|
"--lr_scheduler",
|
|
type=str,
|
|
default="constant",
|
|
help=(
|
|
'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
|
|
' "constant", "constant_with_warmup"]'
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--lr_warmup_steps",
|
|
type=int,
|
|
default=500,
|
|
help="Number of steps for the warmup in the lr scheduler.",
|
|
)
|
|
parser.add_argument(
|
|
"--use_8bit_adam",
|
|
action="store_true",
|
|
help="Whether or not to use 8-bit Adam from bitsandbytes.",
|
|
)
|
|
parser.add_argument(
|
|
"--adam_beta1",
|
|
type=float,
|
|
default=0.9,
|
|
help="The beta1 parameter for the Adam optimizer.",
|
|
)
|
|
parser.add_argument(
|
|
"--adam_beta2",
|
|
type=float,
|
|
default=0.999,
|
|
help="The beta2 parameter for the Adam optimizer.",
|
|
)
|
|
parser.add_argument(
|
|
"--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use."
|
|
)
|
|
parser.add_argument(
|
|
"--adam_epsilon",
|
|
type=float,
|
|
default=1e-08,
|
|
help="Epsilon value for the Adam optimizer",
|
|
)
|
|
parser.add_argument(
|
|
"--max_grad_norm", default=1.0, type=float, help="Max gradient norm."
|
|
)
|
|
parser.add_argument(
|
|
"--push_to_hub",
|
|
action="store_true",
|
|
help="Whether or not to push the model to the Hub.",
|
|
)
|
|
parser.add_argument(
|
|
"--hub_token",
|
|
type=str,
|
|
default=None,
|
|
help="The token to use to push to the Model Hub.",
|
|
)
|
|
parser.add_argument(
|
|
"--logging_dir",
|
|
type=str,
|
|
default="logs",
|
|
help=(
|
|
"[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
|
|
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--mixed_precision",
|
|
type=str,
|
|
default=None,
|
|
choices=["no", "fp16", "bf16"],
|
|
help=(
|
|
"Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
|
|
" 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the"
|
|
" flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--local_rank",
|
|
type=int,
|
|
default=-1,
|
|
help="For distributed training: local_rank. Not to be confused with LoRA rank.",
|
|
)
|
|
parser.add_argument(
|
|
"--resume_unet",
|
|
type=str,
|
|
default=None,
|
|
help=("File path for unet lora to resume training."),
|
|
)
|
|
parser.add_argument(
|
|
"--resume_text_encoder",
|
|
type=str,
|
|
default=None,
|
|
help=("File path for text encoder lora to resume training."),
|
|
)
|
|
parser.add_argument(
|
|
"--resize",
|
|
type=bool,
|
|
default=True,
|
|
required=False,
|
|
help="Should images be resized to --resolution before training?",
|
|
)
|
|
parser.add_argument(
|
|
"--only_attn", action="store_true", help="Only finetune attention layers."
|
|
)
|
|
parser.add_argument(
|
|
"--only_webui", action="store_true", help="Only save weights for webui."
|
|
)
|
|
parser.add_argument(
|
|
"--use_xformers", action="store_true", help="Whether or not to use xformers"
|
|
)
|
|
parser.add_argument(
|
|
"--lora_name",
|
|
type=str,
|
|
default="stable_lora",
|
|
help="The name of your project. Will get saved as LoRA metadata."
|
|
)
|
|
|
|
if input_args is not None:
|
|
args = parser.parse_args(input_args)
|
|
else:
|
|
args = parser.parse_args()
|
|
|
|
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
|
|
if env_local_rank != -1 and env_local_rank != args.local_rank:
|
|
args.local_rank = env_local_rank
|
|
|
|
if args.with_prior_preservation:
|
|
if args.class_data_dir is None:
|
|
raise ValueError("You must specify a data directory for class images.")
|
|
if args.class_prompt is None:
|
|
raise ValueError("You must specify prompt for class images.")
|
|
|
|
return args |