Implementation:AUTOMATIC1111 Stable diffusion webui Train embedding

Knowledge Sources	stable-diffusion-webui
Domains	Textual Inversion, Training Loop, Optimization, Stable Diffusion
Last Updated	2026-02-08 00:00 GMT

Overview

Concrete tool for executing the full textual inversion training loop including dataset loading, mixed-precision optimization with AdamW, gradient accumulation, periodic checkpointing, and preview image generation, provided by the AUTOMATIC1111 stable-diffusion-webui repository.

Description

The train_embedding function orchestrates the complete textual inversion training workflow. It performs the following steps:

1. Validation: Validates all input parameters (model name, learn rate, batch size, paths, etc.)
2. Setup: Creates log directories for embeddings, images, and image embeddings; sets up TensorBoard if enabled
3. Dataset construction: Instantiates a PersonalizedBase dataset with all specified augmentation and sampling options, then wraps it in a PersonalizedDataLoader
4. Optimizer initialization: Creates an AdamW optimizer targeting only embedding.vec with weight_decay=0.0; optionally restores optimizer state from a prior checkpoint
5. Mixed precision setup: Initializes a torch.cuda.amp.GradScaler for float16 training
6. Training loop: Iterates through the dataset, performing forward passes through the U-Net, accumulating gradients, clipping them, and updating the embedding
7. Checkpointing: Saves the embedding (and optionally optimizer state) at regular intervals
8. Preview generation: Generates sample images at regular intervals using either the training prompt or custom txt2img settings
9. Steganographic embedding: Optionally embeds the checkpoint data into preview images for portable distribution
10. Loss logging: Writes loss and learning rate values to CSV and TensorBoard

The training loop handles both standard and inpainting models, offloads the VAE to CPU during training when low VRAM mode is active, and supports graceful interruption via shared.state.interrupted.

Usage

Use this function when:

• You want to train a textual inversion embedding from the WebUI or API
• You need a complete training pipeline with checkpointing, preview generation, and logging
• You want mixed-precision training with gradient accumulation and clipping
• You need support for resuming interrupted training runs

Code Reference

Source Location

• Repository: stable-diffusion-webui
• File: modules/textual_inversion/textual_inversion.py
• Lines: L400-688

Signature

def train_embedding(
    id_task,
    embedding_name,
    learn_rate,
    batch_size,
    gradient_step,
    data_root,
    log_directory,
    training_width,
    training_height,
    varsize,
    steps,
    clip_grad_mode,
    clip_grad_value,
    shuffle_tags,
    tag_drop_out,
    latent_sampling_method,
    use_weight,
    create_image_every,
    save_embedding_every,
    template_filename,
    save_image_with_stored_embedding,
    preview_from_txt2img,
    preview_prompt,
    preview_negative_prompt,
    preview_steps,
    preview_sampler_name,
    preview_cfg_scale,
    preview_seed,
    preview_width,
    preview_height
):
    ...
    return embedding, filename

Import

from modules.textual_inversion.textual_inversion import train_embedding

I/O Contract

Inputs

Name	Type	Required	Description
id_task	str	Yes	Unique task identifier for progress tracking
embedding_name	str	Yes	Name of the embedding to train (must already exist in the embeddings database)
learn_rate	str	Yes	Learning rate schedule string, e.g. "0.005:100, 0.0001:1000"
batch_size	int	Yes	Number of images per training batch
gradient_step	int	Yes	Number of batches to accumulate gradients over before an optimizer step
data_root	str	Yes	Path to directory containing training images
log_directory	str	Yes	Base directory for saving checkpoints, images, and logs
training_width	int	Yes	Width in pixels for training images
training_height	int	Yes	Height in pixels for training images
varsize	bool	Yes	If True, preserve original image aspect ratios with resolution bucketing
steps	int	Yes	Maximum number of training steps
clip_grad_mode	str	Yes	Gradient clipping mode: "value", "norm", or "disabled"
clip_grad_value	str	Yes	Gradient clipping threshold schedule string (same format as learn_rate)
shuffle_tags	bool	Yes	If True, randomly shuffle comma-separated tags in captions each step
tag_drop_out	float	Yes	Probability of dropping each tag from captions (0 to disable)
latent_sampling_method	str	Yes	One of "once", "deterministic", or "random"
use_weight	bool	Yes	If True, use alpha channel as per-pixel loss weight
create_image_every	int	Yes	Generate a preview image every N steps (0 to disable)
save_embedding_every	int	Yes	Save an embedding checkpoint every N steps (0 to disable)
template_filename	str	Yes	Filename of the prompt template file from the textual inversion templates directory
save_image_with_stored_embedding	bool	Yes	If True, embed checkpoint data steganographically into preview images
preview_from_txt2img	bool	Yes	If True, use custom txt2img settings for preview generation instead of training prompts
preview_prompt	str	Yes	Custom prompt for preview generation (used when preview_from_txt2img=True)
preview_negative_prompt	str	Yes	Custom negative prompt for preview generation
preview_steps	int	Yes	Number of sampling steps for preview generation
preview_sampler_name	str	Yes	Sampler name for preview generation
preview_cfg_scale	float	Yes	CFG scale for preview generation
preview_seed	int	Yes	Seed for preview generation
preview_width	int	Yes	Width in pixels for preview images
preview_height	int	Yes	Height in pixels for preview images

Outputs

Name	Type	Description
embedding	Embedding	The trained Embedding object with updated vec, step, and checkpoint metadata
filename	str	Path to the final saved .pt embedding file

Usage Examples

Basic Usage

from modules.textual_inversion.textual_inversion import train_embedding

embedding, filepath = train_embedding(
    id_task="task-001",
    embedding_name="my-cat",
    learn_rate="0.005:100, 0.001:500, 1e-4:2000",
    batch_size=1,
    gradient_step=1,
    data_root="/path/to/cat_images",
    log_directory="/path/to/logs",
    training_width=512,
    training_height=512,
    varsize=False,
    steps=2000,
    clip_grad_mode="norm",
    clip_grad_value="1.0",
    shuffle_tags=True,
    tag_drop_out=0.1,
    latent_sampling_method="once",
    use_weight=False,
    create_image_every=100,
    save_embedding_every=500,
    template_filename="style_filewords.txt",
    save_image_with_stored_embedding=True,
    preview_from_txt2img=False,
    preview_prompt="",
    preview_negative_prompt="",
    preview_steps=20,
    preview_sampler_name="Euler a",
    preview_cfg_scale=7.0,
    preview_seed=-1,
    preview_width=512,
    preview_height=512
)

print(f"Training completed at step {embedding.step}")
print(f"Embedding saved to: {filepath}")

Related Pages

Implements Principle

• Principle:AUTOMATIC1111_Stable_diffusion_webui_Textual_inversion_training_loop

Requires Environment

• Environment:AUTOMATIC1111_Stable_diffusion_webui_Python_And_PyTorch_Runtime
• Environment:AUTOMATIC1111_Stable_diffusion_webui_GPU_Compute_Backend