Skip to content

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