FullStableDiffusionConfig
FullStableDiffusionConfig
dataclass
Bases: FinetuningConfig
Fine-tuning config for full fine-tuning with Stable Diffusion.
Examples:
from baseten.training import Dataset, FullStableDiffusionConfig
config = FullStableDiffusionConfig(
input_dataset=Dataset("DATASET_ID"),
train_text_encoder=False,
num_train_epochs=10
)
Parameters:
Name | Type | Description | Default |
---|---|---|---|
input_dataset |
DatasetIdentifier
|
An identifier, either an ID or a public URL, for the Dataset that Dreambooth should use |
required |
wandb_api_key |
Optional[str]
|
API key for Weights & Biases to monitor your model training |
None
|
pretrained_model_name_or_path |
str
|
Path to pretrained model or model identifier from huggingface.co/models. |
'CompVis/stable-diffusion-v1-4'
|
revision |
Optional[str]
|
Revision of pretrained model identifier from huggingface.co/models. |
None
|
seed |
Optional[int]
|
The random seed to use for fine-tuning |
None
|
gradient_accumulation_steps |
int
|
Number of gradient accumulation steps to use. Defaults to 1. |
1
|
resolution |
int
|
The resolution of your input images. Images will be resized to this value. |
512
|
num_train_epochs |
int
|
The number of epochs to train for. If you set max_train_steps, this will be ignored. |
10
|
train_text_encoder |
bool
|
Whether to train the text encoder alongside the UNet |
False
|
center_crop |
bool
|
Whether to center crop the images to the resolution |
False
|
random_flip |
bool
|
Whether to randomly flip the images horizontally |
False
|
train_batch_size |
int
|
The batch size to use for training. This value can cause OOMs if too large. We recommend using between 1-4 based on the resolution of your images. |
1
|
max_train_steps |
Optional[int]
|
The number of training steps to train for. If you set this, num_train_epochs will be ignored. |
None
|
learning_rate |
float
|
Initial learning rate (after the potential warmup period) to use. |
1e-05
|
lr_scheduler |
str
|
The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"]. |
'cosine'
|
lr_warmup_steps |
int
|
The number of steps to warmup for the learning rate schedule. |
500
|
use_ema |
bool
|
Whether to use exponential moving average for the model weights. |
False
|
non_ema_revision |
Optional[str]
|
Revision of pretrained non-ema model identifier. Must be remote repository specified with --pretrained_model_name_or_path. |
None
|
adam_beta1 |
float
|
The beta1 parameter for the Adam optimizer. |
0.9
|
adam_beta2 |
float
|
The beta2 parameter for the Adam optimizer. |
0.999
|
adam_weight_decay |
float
|
The weight decay value for the Adam optimizer. |
0.01
|
adam_epsilon |
float
|
The epsilon value for the Adam optimizer. |
1e-08
|
max_grad_norm |
float
|
The max gradient normalization to clip gradients to. Prevents exploding gradients. |
1.0
|
image_log_steps |
int
|
The number of steps to log sample images to Weights and Biases. This allows you to visually assess your model during training. Only relevant if wandb_api_key is set. |
50
|