LoRA Fine-Tuning Dual Pipeline

Introduction

30 Second Summary

AI image generators can produce incredible visuals, but teaching them to generate images of a specific subject or style requires fine-tuning. The challenge is building training infrastructure that works reliably on both a local GPU and a cloud environment.

In this project, you will build a dual-execution LoRA fine-tuning pipeline for Stable Diffusion XL that runs identically on your local GPU and on Amazon SageMaker.

What You'll Build

A complete, modular MLOps repository called diffusers-lora-dual-pipeline that fine-tunes Stable Diffusion XL using LoRA adapters, supporting both local GPU training via Hugging Face Accelerate and cloud training via Amazon SageMaker.

By the end of this project, you'll have:

• A production-ready training script that fine-tunes Stable Diffusion XL with LoRA adapters using your own images.
• A local execution pipeline using Hugging Face Accelerate for single-GPU training with mixed-precision support.
• A SageMaker submission pipeline that packages data to S3, provisions a cloud GPU instance, and retrieves trained LoRA weights.
• An inference script that loads trained LoRA weights into the base SDXL model and generates evaluation images.
• Secret Mission: Automatic LoRA weight upload to Hugging Face Hub with generated model cards.

Are there any prerequisites?

This project assumes familiarity with Python and basic PyTorch concepts. You will need an NVIDIA GPU with 16GB+ VRAM for local training, or an AWS account with SageMaker access for the cloud path.

You will also need a Hugging Face account to download models.

Set Up Repository Structure and Dependencies

Every ML training pipeline starts with a clean, organized repository. A well-structured project means your training script can find data, save checkpoints, and work identically whether it runs on your local GPU or a SageMaker instance.

In this step, you'll scaffold the project directory, pin your Python dependencies, and configure all training hyperparameters in a single YAML file.

In this step, get ready to:

• Create the diffusers-lora-dual-pipeline project structure with all required subdirectories.
• Install pinned Python dependencies in an isolated virtual environment.
• Configure training hyperparameters in config/default.yaml.

Create the repository directory structure

Your project needs dedicated folders for source code, configuration, training data, and outputs. Separating these concerns keeps the repository navigable and makes it straightforward to package different parts for local or cloud execution.

• Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
• Type Terminal (macOS) or Git Bash (Windows) and press Enter to open it.

Why Git Bash on Windows?

This project uses bash scripts for local training (local_run.sh). Git Bash comes bundled with Git for Windows and gives you a Unix-compatible shell. If you prefer, WSL (Windows Subsystem for Linux) also works.

• Move to your Desktop by running this command:

cd ~/Desktop

• Create the project folder and all subdirectories by running this command:

mkdir -p diffusers-lora-dual-pipeline/{src,config,data/instance_images,outputs}
cd diffusers-lora-dual-pipeline

What do these directories hold?

• src/ holds your training and inference Python scripts.
• config/ stores YAML hyperparameter files.
• data/instance_images/ is where you'll place training images later.
• outputs/ collects trained LoRA weights, checkpoints, and logs.

• Confirm the directories exist by running:

ls

You should see config, data, outputs, and src listed in the output.

Don't see the four folders?

Make sure you ran the mkdir -p command from your Desktop, not from inside another project. Run pwd to check your current location. It should end with diffusers-lora-dual-pipeline.

Help me debug my directory structure.

Create the requirements file and install dependencies

Your project depends on several libraries: diffusers for the Stable Diffusion XL pipeline, PEFT for LoRA adapters, and Accelerate for distributed training.

Pinning versions ensures reproducible training runs across your local machine and SageMaker.

• Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
• Type VS Code (or Cursor) and press Enter to open your code editor.
• Click File in the top menu bar.
• Click Open Folder.
• Select the diffusers-lora-dual-pipeline folder on your Desktop, then click Open.
• Create a new file called requirements.txt in the project root by pressing Cmd+N (macOS) or Ctrl+N (Windows), then Cmd+S (macOS) or Ctrl+S (Windows) and name it requirements.txt.
• Paste the following dependency list into requirements.txt:

diffusers==0.38.0
transformers>=4.41.2
accelerate>=1.0.0
peft>=0.17.0
sagemaker>=2.257.0
boto3>=1.35.0
safetensors>=0.4.3
torchvision
Pillow
wandb
pyyaml
huggingface_hub>=0.21.0

What are these packages for?

• diffusers==0.38.0 is pinned exactly because the training API can change between releases.
• peft provides the LoRA adapter implementation you'll inject into the UNet.
• sagemaker and boto3 handle cloud job submission and S3 data transfers.
• wandb enables metric logging so you can track loss curves during training.

• Save requirements.txt.

✔️ Awesome, I've got everything!

Great. Double check you've saved your file.

ⓧ I'd like to double check the full code

diffusers==0.38.0
transformers>=4.41.2
accelerate>=1.0.0
peft>=0.17.0
sagemaker>=2.257.0
boto3>=1.35.0
safetensors>=0.4.3
torchvision
Pillow
wandb
pyyaml
huggingface_hub>=0.21.0

Now create a virtual environment and install these packages.

macOS

• In your terminal (still inside diffusers-lora-dual-pipeline), create and activate a virtual environment by running these commands:

python3 -m venv venv
source venv/bin/activate

Windows

• In your terminal (still inside diffusers-lora-dual-pipeline), create and activate a virtual environment by running these commands:

python -m venv venv
source venv/Scripts/activate

Your terminal prompt should now show (venv) at the beginning, confirming the virtual environment is active.

• Install all project dependencies by running:

pip install -r requirements.txt

This install takes 2-3 minutes depending on your network speed. You should see packages downloading and a final confirmation line.

Seeing install errors?

If you see a permission error, make sure your (venv) prefix is visible in your prompt. If it's missing, the virtual environment isn't active. Re-run the activation command above.

If a specific package fails to build, check that you have Python 3.10+ by running python3 --version.

Help me fix my pip install errors.

Create the hyperparameter configuration file

A centralized config file lets you change training parameters without editing Python code. This is especially important when the same script runs in two different environments.

• In your code editor, create a new file at config/default.yaml inside the config/ folder.
• Paste the following configuration into config/default.yaml:

# Model Configuration
model:
  pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
  pretrained_vae_model_name_or_path: "madebyollin/sdxl-vae-fp16-fix"

# LoRA Configuration
lora:
  rank: 16
  alpha: 16
  target_modules_unet:
    - "to_k"
    - "to_q"
    - "to_v"
    - "to_out.0"
  target_modules_text_encoder:
    - "q_proj"
    - "k_proj"
    - "v_proj"
    - "out_proj"
  train_text_encoder: false

# Training Configuration
training:
  instance_prompt: "a photo of sks subject"
  resolution: 1024
  train_batch_size: 1
  gradient_accumulation_steps: 4
  learning_rate: 1.0e-4
  lr_scheduler: "constant"
  lr_warmup_steps: 0
  max_train_steps: 1000
  mixed_precision: "fp16"
  gradient_checkpointing: true
  seed: 42
  checkpointing_steps: 250

# Data Configuration
data:
  instance_data_dir: "./data/instance_images"
  caption_column: "text"

# Output Configuration
output:
  output_dir: "./outputs/lora_weights"
  logging_dir: "./outputs/logs"
  report_to: "wandb"

# Hub Configuration
hub:
  push_to_hub: false
  hub_model_id: null

What do these parameters control?

• rank: 16 sets the LoRA rank. Higher rank captures more detail but uses more VRAM.
• target_modules_unet lists the UNet attention layers that get LoRA adapters injected.
• resolution: 1024 matches SDXL's native training resolution.
• gradient_checkpointing: true trades compute time for VRAM savings, critical for fitting SDXL on a 16GB GPU.
• mixed_precision: "fp16" halves memory usage for activations and gradients.

Why a separate VAE model?

SDXL's default VAE can produce NaN values during fp16 training. The madebyollin/sdxl-vae-fp16-fix is a community fix that resolves this numerical instability.

• Save config/default.yaml.
• Confirm the file appears in the config/ folder in your editor's sidebar.

File not appearing in the config folder?

Make sure you saved the file inside the config/ subdirectory, not in the project root. The full path should be diffusers-lora-dual-pipeline/config/default.yaml.

Help me find where my config file ended up.

✔️ Awesome, I've got everything!

Great. Double check you've saved your file.

ⓧ I'd like to double check the full code

# Model Configuration
model:
  pretrained_model_name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
  pretrained_vae_model_name_or_path: "madebyollin/sdxl-vae-fp16-fix"

# LoRA Configuration
lora:
  rank: 16
  alpha: 16
  target_modules_unet:
    - "to_k"
    - "to_q"
    - "to_v"
    - "to_out.0"
  target_modules_text_encoder:
    - "q_proj"
    - "k_proj"
    - "v_proj"
    - "out_proj"
  train_text_encoder: false

# Training Configuration
training:
  instance_prompt: "a photo of sks subject"
  resolution: 1024
  train_batch_size: 1
  gradient_accumulation_steps: 4
  learning_rate: 1.0e-4
  lr_scheduler: "constant"
  lr_warmup_steps: 0
  max_train_steps: 1000
  mixed_precision: "fp16"
  gradient_checkpointing: true
  seed: 42
  checkpointing_steps: 250

# Data Configuration
data:
  instance_data_dir: "./data/instance_images"
  caption_column: "text"

# Output Configuration
output:
  output_dir: "./outputs/lora_weights"
  logging_dir: "./outputs/logs"
  report_to: "wandb"

# Hub Configuration
hub:
  push_to_hub: false
  hub_model_id: null

Your project skeleton is ready. Next up, you'll build the unified training script that uses this configuration to fine-tune Stable Diffusion XL with LoRA adapters.

Build the Unified Training Script

Your repository structure and dependencies are ready. Now it's time to build the core of the pipeline: a single training script that works identically on your local GPU and on Amazon SageMaker.

The script needs to handle several responsibilities: loading the massive SDXL model, injecting lightweight LoRA adapters into the UNet, running a diffusion-based training loop, and saving checkpoints. By keeping all of this in one file, you avoid the common MLOps headache of maintaining separate local and cloud training codebases.

In this step, get ready to:

• Write the script foundation with argument parsing, config loading, and SDXL model loading.
• Implement LoRA adapter injection using PEFT's LoraConfig targeting the UNet's attention layers.
• Build the full training loop with Accelerate, including noise scheduling, loss computation, and checkpoint saving.

Write the script foundation with imports, dataset class, and config utilities

The training script starts with all the imports, a custom dataset class for loading instance images, and utility functions for parsing arguments and YAML configuration.

• Create the file src/train_lora.py and add the following imports, dataset class, and utility functions:

#!/usr/bin/env python3
"""Unified LoRA fine-tuning script for Stable Diffusion XL.

Supports both local GPU training (via Accelerate) and Amazon SageMaker execution.
Reads hyperparameters from config/default.yaml or CLI arguments.
"""

import argparse
import logging
import math
import os
from pathlib import Path

import torch
import torch.nn.functional as F
import yaml
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration, set_seed
from diffusers import (
    AutoencoderKL,
    DDPMScheduler,
    StableDiffusionXLPipeline,
    UNet2DConditionModel,
)
from diffusers.optimization import get_scheduler
from huggingface_hub import upload_folder
from peft import LoraConfig
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from transformers import AutoTokenizer, PretrainedConfig

logger = get_logger(__name__, log_level="INFO")


class InstanceDataset(Dataset):
    """Dataset for instance images with captions."""

    def __init__(self, instance_data_dir, instance_prompt, tokenizer_one, tokenizer_two, size=1024):
        self.size = size
        self.instance_prompt = instance_prompt
        self.tokenizer_one = tokenizer_one
        self.tokenizer_two = tokenizer_two

        self.instance_images_path = []
        self.instance_captions = []

        inst_dir = Path(instance_data_dir)
        supported_ext = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}

        metadata_file = inst_dir / "metadata.jsonl"
        if metadata_file.exists():
            import json
            with open(metadata_file, "r") as f:
                for line in f:
                    entry = json.loads(line.strip())
                    img_path = inst_dir / entry["file_name"]
                    if img_path.exists():
                        self.instance_images_path.append(img_path)
                        self.instance_captions.append(entry.get("text", instance_prompt))
        else:
            for img_file in sorted(inst_dir.iterdir()):
                if img_file.suffix.lower() in supported_ext:
                    self.instance_images_path.append(img_file)
                    self.instance_captions.append(instance_prompt)

        self.num_instance_images = len(self.instance_images_path)
        self.image_transforms = transforms.Compose([
            transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.CenterCrop(size),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ])

    def __len__(self):
        return self.num_instance_images

    def __getitem__(self, index):
        image = Image.open(self.instance_images_path[index % self.num_instance_images])
        image = image.convert("RGB")
        image = self.image_transforms(image)
        caption = self.instance_captions[index % self.num_instance_images]

        tokens_one = self.tokenizer_one(
            caption, truncation=True, padding="max_length",
            max_length=self.tokenizer_one.model_max_length, return_tensors="pt"
        ).input_ids[0]

        tokens_two = self.tokenizer_two(
            caption, truncation=True, padding="max_length",
            max_length=self.tokenizer_two.model_max_length, return_tensors="pt"
        ).input_ids[0]

        return {"pixel_values": image, "input_ids_one": tokens_one, "input_ids_two": tokens_two}


def load_config(config_path):
    """Load YAML configuration file."""
    with open(config_path, "r") as f:
        return yaml.safe_load(f)


def parse_args():
    parser = argparse.ArgumentParser(description="LoRA fine-tuning for SDXL")
    parser.add_argument("--config", type=str, default="config/default.yaml")
    parser.add_argument("--pretrained_model_name_or_path", type=str, default=None)
    parser.add_argument("--instance_data_dir", type=str, default=None)
    parser.add_argument("--instance_prompt", type=str, default=None)
    parser.add_argument("--output_dir", type=str, default=None)
    parser.add_argument("--resolution", type=int, default=None)
    parser.add_argument("--train_batch_size", type=int, default=None)
    parser.add_argument("--max_train_steps", type=int, default=None)
    parser.add_argument("--learning_rate", type=float, default=None)
    parser.add_argument("--lora_rank", type=int, default=None)
    parser.add_argument("--lora_alpha", type=int, default=None)
    parser.add_argument("--mixed_precision", type=str, default=None)
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--checkpointing_steps", type=int, default=None)
    parser.add_argument("--report_to", type=str, default=None)
    parser.add_argument("--push_to_hub", action="store_true", default=False)
    parser.add_argument("--hub_model_id", type=str, default=None)
    parser.add_argument("--validation_prompt", type=str, default=None)
    parser.add_argument("--gradient_checkpointing", action="store_true", default=None)
    parser.add_argument("--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "./outputs"))
    args = parser.parse_args()
    return args

What does this code do?

The imports bring in the core libraries: diffusers for the SDXL model components, peft for LoRA configuration, accelerate for distributed training, and transformers for tokenizers.

The InstanceDataset class handles loading training images. It checks for a metadata.jsonl file with per-image captions. If none exists, it uses the global instance_prompt for every image. SDXL requires dual tokenization, so each image produces tokens for both CLIP text encoders.

The parse_args function accepts CLI arguments that can override the YAML config. The --model_dir argument reads from SM_MODEL_DIR so SageMaker can pass its output path automatically.

• Save src/train_lora.py.

You won't see output yet. The next chunk adds the main() function that ties everything together.

Seeing import errors in your editor?

Make sure your virtual environment is activated and selected as the Python interpreter in your editor. The packages from requirements.txt must be installed in the active environment.

Help me fix import errors in my LoRA training script

Now add the main() function. This is where the script loads configuration, initializes Accelerate, and brings all four SDXL model components into memory.

• In src/train_lora.py, add the following main() function and entry point below parse_args:

def main():
    args = parse_args()
    config = {}
    if os.path.exists(args.config):
        config = load_config(args.config)

    # Merge config with CLI args (CLI takes precedence)
    model_name = args.pretrained_model_name_or_path or config.get("model", {}).get("pretrained_model_name_or_path", "stabilityai/stable-diffusion-xl-base-1.0")
    vae_name = config.get("model", {}).get("pretrained_vae_model_name_or_path", "madebyollin/sdxl-vae-fp16-fix")
    instance_data_dir = args.instance_data_dir or config.get("data", {}).get("instance_data_dir", "./data/instance_images")
    instance_prompt = args.instance_prompt or config.get("training", {}).get("instance_prompt", "a photo of sks subject")
    output_dir = args.output_dir or config.get("output", {}).get("output_dir", "./outputs/lora_weights")
    resolution = args.resolution or config.get("training", {}).get("resolution", 1024)
    train_batch_size = args.train_batch_size or config.get("training", {}).get("train_batch_size", 1)
    max_train_steps = args.max_train_steps or config.get("training", {}).get("max_train_steps", 1000)
    learning_rate = args.learning_rate or config.get("training", {}).get("learning_rate", 1e-4)
    lora_rank = args.lora_rank or config.get("lora", {}).get("rank", 16)
    lora_alpha = args.lora_alpha or config.get("lora", {}).get("alpha", 16)
    mixed_precision = args.mixed_precision or config.get("training", {}).get("mixed_precision", "fp16")
    seed = args.seed if args.seed is not None else config.get("training", {}).get("seed", 42)
    checkpointing_steps = args.checkpointing_steps or config.get("training", {}).get("checkpointing_steps", 250)
    report_to = args.report_to or config.get("output", {}).get("report_to", "wandb")
    gradient_accumulation_steps = config.get("training", {}).get("gradient_accumulation_steps", 4)
    lr_scheduler_type = config.get("training", {}).get("lr_scheduler", "constant")
    lr_warmup_steps = config.get("training", {}).get("lr_warmup_steps", 0)
    gradient_checkpointing = args.gradient_checkpointing if args.gradient_checkpointing is not None else config.get("training", {}).get("gradient_checkpointing", True)
    train_text_encoder = config.get("lora", {}).get("train_text_encoder", False)
    target_modules_unet = config.get("lora", {}).get("target_modules_unet", ["to_k", "to_q", "to_v", "to_out.0"])
    push_to_hub = args.push_to_hub or config.get("hub", {}).get("push_to_hub", False)
    hub_model_id = args.hub_model_id or config.get("hub", {}).get("hub_model_id")

    # Handle SageMaker data channel
    sm_channel_training = os.environ.get("SM_CHANNEL_TRAINING")
    if sm_channel_training:
        instance_data_dir = sm_channel_training
        output_dir = args.model_dir

    # Setup accelerator
    logging_dir = Path(output_dir) / "logs"
    project_config = ProjectConfiguration(project_dir=output_dir, logging_dir=str(logging_dir))

    accelerator = Accelerator(
        gradient_accumulation_steps=gradient_accumulation_steps,
        mixed_precision=mixed_precision,
        log_with=report_to if report_to != "none" else None,
        project_config=project_config,
    )

    logging.basicConfig(format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO)
    logger.info(accelerator.state, main_process_only=False)

    if seed is not None:
        set_seed(seed)

    os.makedirs(output_dir, exist_ok=True)

    # Load tokenizers
    tokenizer_one = AutoTokenizer.from_pretrained(model_name, subfolder="tokenizer", use_fast=False)
    tokenizer_two = AutoTokenizer.from_pretrained(model_name, subfolder="tokenizer_2", use_fast=False)

    # Load models
    noise_scheduler = DDPMScheduler.from_pretrained(model_name, subfolder="scheduler")
    vae = AutoencoderKL.from_pretrained(vae_name, torch_dtype=torch.float16)
    unet = UNet2DConditionModel.from_pretrained(model_name, subfolder="unet")

    from transformers import CLIPTextModel, CLIPTextModelWithProjection
    text_encoder_one = CLIPTextModel.from_pretrained(model_name, subfolder="text_encoder")
    text_encoder_two = CLIPTextModelWithProjection.from_pretrained(model_name, subfolder="text_encoder_2")

    # Freeze base model
    vae.requires_grad_(False)
    text_encoder_one.requires_grad_(False)
    text_encoder_two.requires_grad_(False)
    unet.requires_grad_(False)

    # Move to device
    weight_dtype = torch.float16 if mixed_precision == "fp16" else torch.float32
    vae.to(accelerator.device, dtype=weight_dtype)
    text_encoder_one.to(accelerator.device, dtype=weight_dtype)
    text_encoder_two.to(accelerator.device, dtype=weight_dtype)

What does this code do?

The config merge pattern gives CLI arguments priority over config/default.yaml values. This lets SageMaker pass hyperparameters as CLI arguments while local runs use the YAML file.

The SageMaker detection block checks for SM_CHANNEL_TRAINING. When SageMaker runs the script, it sets this environment variable to the path where training data was downloaded from S3. On a local machine, this variable won't exist.

SDXL has four main components: two CLIP text encoders (one standard, one with projection), a UNet for denoising, and a VAE for encoding/decoding images. All base model weights are frozen because only the LoRA adapters will be trained.

• Save src/train_lora.py.
• Verify the argument parser works by running:

python src/train_lora.py --help

You should see a list of all the arguments you defined (--config, --lora_rank, --max_train_steps, etc.) printed in the terminal.

Getting a ModuleNotFoundError?

Confirm your virtual environment is active. You should see (venv) at the start of your terminal prompt. If not, activate it with source venv/bin/activate (macOS/Linux) or venv\Scripts\activate (Windows).

Help me fix ModuleNotFoundError when running my training script

Implement LoRA adapter injection using PEFT

LoRA works by injecting small, trainable matrices into specific layers of a frozen model. Instead of fine-tuning all 3.5 billion UNet parameters, you'll train roughly 4-8 million parameters in the LoRA adapters. This makes training feasible on a single GPU.

• In src/train_lora.py, add the LoRA injection and optimizer setup at the end of main(), directly below the line text_encoder_two.to(accelerator.device, dtype=weight_dtype):

    # Setup LoRA
    unet_lora_config = LoraConfig(
        r=lora_rank,
        lora_alpha=lora_alpha,
        init_lora_weights="gaussian",
        target_modules=target_modules_unet,
    )
    unet.add_adapter(unet_lora_config)

    if gradient_checkpointing:
        unet.enable_gradient_checkpointing()

    # Collect trainable parameters
    lora_layers = list(filter(lambda p: p.requires_grad, unet.parameters()))
    logger.info(f"Trainable LoRA parameters: {sum(p.numel() for p in lora_layers):,}")

    # Optimizer
    optimizer = torch.optim.AdamW(lora_layers, lr=learning_rate, betas=(0.9, 0.999), weight_decay=1e-2, eps=1e-8)

What does this code do?

The LoraConfig defines adapter properties. r=lora_rank sets the rank (bottleneck dimension) of the adapter matrices. target_modules specifies which UNet layers get adapters: to_k, to_q, to_v, and to_out.0 are the attention projection layers.

After calling unet.add_adapter(), only the LoRA parameters have requires_grad=True. The filter collects these into lora_layers for the AdamW optimizer. Gradient checkpointing trades compute for VRAM by recomputing activations during the backward pass instead of storing them.

Why target attention layers specifically?

The UNet's cross-attention layers are where text conditioning meets visual features. By injecting LoRA into the key, query, value, and output projections, you're teaching the model a new mapping between text tokens and visual patterns. This is the highest-leverage injection point for subject-driven fine-tuning.

• Save src/train_lora.py.

You won't see a change yet. The next chunk uses these trainable parameters in the training loop.

Implement the training loop with Accelerate

The training loop is where the actual learning happens. Each iteration samples noise, predicts it using the UNet, and computes the MSE loss between predicted and actual noise. Hugging Face Accelerate handles device placement, gradient accumulation, and mixed-precision under the hood.

• In src/train_lora.py, add the dataset setup, dataloader, learning rate scheduler, and Accelerate preparation below the optimizer line:

    # Dataset
    train_dataset = InstanceDataset(
        instance_data_dir=instance_data_dir,
        instance_prompt=instance_prompt,
        tokenizer_one=tokenizer_one,
        tokenizer_two=tokenizer_two,
        size=resolution,
    )
    train_dataloader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True, num_workers=0)

    # LR Scheduler
    lr_scheduler = get_scheduler(
        lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps,
    )

    # Prepare with accelerator
    unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        unet, optimizer, train_dataloader, lr_scheduler
    )

    # Initialize trackers
    if accelerator.is_main_process:
        tracker_config = {"lora_rank": lora_rank, "lora_alpha": lora_alpha, "learning_rate": learning_rate, "max_train_steps": max_train_steps, "resolution": resolution, "model": model_name}
        if report_to and report_to != "none":
            accelerator.init_trackers("lora-sdxl-training", config=tracker_config)

What does this code do?

The accelerator.prepare() call wraps the model, optimizer, dataloader, and scheduler for distributed training and mixed-precision. Even on a single GPU, this handles fp16 casting and gradient scaling automatically.

The tracker initialization sends hyperparameters to W&B (Weights & Biases) for experiment tracking. If report_to is "none", logging is skipped entirely.

• Save src/train_lora.py.

You won't see a change yet. The next chunk adds the core training loop.

• Now add the training loop, final weight saving, and Hub upload logic below the tracker initialization. Also add the if __name__ == "__main__": entry point at the very end of the file:

    # Training loop
    global_step = 0
    logger.info("***** Starting training *****")

    progress_bar = None
    try:
        from tqdm.auto import tqdm
        progress_bar = tqdm(range(max_train_steps), desc="Training")
    except ImportError:
        pass

    unet.train()
    while global_step < max_train_steps:
        for batch in train_dataloader:
            with accelerator.accumulate(unet):
                pixel_values = batch["pixel_values"].to(dtype=weight_dtype)
                latents = vae.encode(pixel_values).latent_dist.sample()
                latents = latents * vae.config.scaling_factor
                noise = torch.randn_like(latents)
                bsz = latents.shape[0]
                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device).long()
                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
                prompt_embeds_one = text_encoder_one(batch["input_ids_one"].to(accelerator.device), output_hidden_states=True)
                prompt_embeds_two = text_encoder_two(batch["input_ids_two"].to(accelerator.device), output_hidden_states=True)
                pooled_prompt_embeds = prompt_embeds_two[0]
                prompt_embeds = torch.cat([prompt_embeds_one.hidden_states[-2], prompt_embeds_two.hidden_states[-2]], dim=-1)
                add_time_ids = torch.tensor([[resolution, resolution, 0, 0, resolution, resolution]], device=latents.device).repeat(bsz, 1)
                added_cond_kwargs = {"text_embeds": pooled_prompt_embeds, "time_ids": add_time_ids}
                model_pred = unet(noisy_latents, timesteps, prompt_embeds, added_cond_kwargs=added_cond_kwargs).sample
                loss = F.mse_loss(model_pred.float(), noise.float(), reduction="mean")
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.sync_gradients:
                global_step += 1
                if progress_bar:
                    progress_bar.update(1)
                logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
                if progress_bar:
                    progress_bar.set_postfix(**logs)
                accelerator.log(logs, step=global_step)
                if global_step % checkpointing_steps == 0 and accelerator.is_main_process:
                    save_path = os.path.join(output_dir, f"checkpoint-{global_step}")
                    accelerator.save_state(save_path)
                    logger.info(f"Saved checkpoint to {save_path}")

            if global_step >= max_train_steps:
                break

    # Save final LoRA weights
    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        unet_model = accelerator.unwrap_model(unet)
        unet_model.save_pretrained(output_dir)
        logger.info(f"Saved LoRA weights to {output_dir}")
        if push_to_hub and hub_model_id:
            logger.info(f"Uploading to Hugging Face Hub: {hub_model_id}")
            upload_folder(
                repo_id=hub_model_id,
                folder_path=output_dir,
                commit_message="End of LoRA training",
                ignore_patterns=["checkpoint-*", "logs"],
            )

    accelerator.end_training()
    logger.info("Training complete!")


if __name__ == "__main__":
    main()

What does this code do?

Each training iteration follows the diffusion training pattern: encode images to latents with the VAE, add random noise at a random timestep, predict the noise with the UNet, and compute MSE loss between the prediction and actual noise.

SDXL's UNet expects concatenated text embeddings from both CLIP encoders plus add_time_ids (original resolution and crop coordinates). The accelerator.accumulate() context manager handles gradient accumulation over 4 steps before updating weights.

Checkpoints are saved every checkpointing_steps (250 by default). The final LoRA weights are saved as pytorch_lora_weights.safetensors via save_pretrained(). If --push_to_hub is set, upload_folder() pushes the weights to Hugging Face Hub.

• Save src/train_lora.py.
• Verify the script parses correctly by running:

python -c "import ast; ast.parse(open('src/train_lora.py').read()); print('Syntax OK')"

You should see Syntax OK in the terminal. This confirms your entire script is valid Python without running it.

Seeing a SyntaxError?

Check the line number in the error message. Common issues: missing closing parentheses on the LoraConfig() call, unmatched brackets in the training loop, or incorrect indentation on the if __name__ block (it should be at the top level, not indented inside main()).

Help me fix a SyntaxError in my LoRA training script

✔️ Awesome, I've got everything!

Great. Double check you've saved src/train_lora.py and that the syntax check passed.

ⓧ I'd like to double check the full code

Here is the complete src/train_lora.py file. Compare it against yours to find any differences:

#!/usr/bin/env python3
"""Unified LoRA fine-tuning script for Stable Diffusion XL.

Supports both local GPU training (via Accelerate) and Amazon SageMaker execution.
Reads hyperparameters from config/default.yaml or CLI arguments.
"""

import argparse
import logging
import math
import os
from pathlib import Path

import torch
import torch.nn.functional as F
import yaml
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration, set_seed
from diffusers import (
    AutoencoderKL,
    DDPMScheduler,
    StableDiffusionXLPipeline,
    UNet2DConditionModel,
)
from diffusers.optimization import get_scheduler
from huggingface_hub import upload_folder
from peft import LoraConfig
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from transformers import AutoTokenizer, PretrainedConfig

logger = get_logger(__name__, log_level="INFO")


class InstanceDataset(Dataset):
    """Dataset for instance images with captions."""

    def __init__(self, instance_data_dir, instance_prompt, tokenizer_one, tokenizer_two, size=1024):
        self.size = size
        self.instance_prompt = instance_prompt
        self.tokenizer_one = tokenizer_one
        self.tokenizer_two = tokenizer_two

        self.instance_images_path = []
        self.instance_captions = []

        inst_dir = Path(instance_data_dir)
        supported_ext = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}

        metadata_file = inst_dir / "metadata.jsonl"
        if metadata_file.exists():
            import json
            with open(metadata_file, "r") as f:
                for line in f:
                    entry = json.loads(line.strip())
                    img_path = inst_dir / entry["file_name"]
                    if img_path.exists():
                        self.instance_images_path.append(img_path)
                        self.instance_captions.append(entry.get("text", instance_prompt))
        else:
            for img_file in sorted(inst_dir.iterdir()):
                if img_file.suffix.lower() in supported_ext:
                    self.instance_images_path.append(img_file)
                    self.instance_captions.append(instance_prompt)

        self.num_instance_images = len(self.instance_images_path)
        self.image_transforms = transforms.Compose([
            transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.CenterCrop(size),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ])

    def __len__(self):
        return self.num_instance_images

    def __getitem__(self, index):
        image = Image.open(self.instance_images_path[index % self.num_instance_images])
        image = image.convert("RGB")
        image = self.image_transforms(image)
        caption = self.instance_captions[index % self.num_instance_images]

        tokens_one = self.tokenizer_one(
            caption, truncation=True, padding="max_length",
            max_length=self.tokenizer_one.model_max_length, return_tensors="pt"
        ).input_ids[0]

        tokens_two = self.tokenizer_two(
            caption, truncation=True, padding="max_length",
            max_length=self.tokenizer_two.model_max_length, return_tensors="pt"
        ).input_ids[0]

        return {"pixel_values": image, "input_ids_one": tokens_one, "input_ids_two": tokens_two}


def load_config(config_path):
    """Load YAML configuration file."""
    with open(config_path, "r") as f:
        return yaml.safe_load(f)


def parse_args():
    parser = argparse.ArgumentParser(description="LoRA fine-tuning for SDXL")
    parser.add_argument("--config", type=str, default="config/default.yaml")
    parser.add_argument("--pretrained_model_name_or_path", type=str, default=None)
    parser.add_argument("--instance_data_dir", type=str, default=None)
    parser.add_argument("--instance_prompt", type=str, default=None)
    parser.add_argument("--output_dir", type=str, default=None)
    parser.add_argument("--resolution", type=int, default=None)
    parser.add_argument("--train_batch_size", type=int, default=None)
    parser.add_argument("--max_train_steps", type=int, default=None)
    parser.add_argument("--learning_rate", type=float, default=None)
    parser.add_argument("--lora_rank", type=int, default=None)
    parser.add_argument("--lora_alpha", type=int, default=None)
    parser.add_argument("--mixed_precision", type=str, default=None)
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--checkpointing_steps", type=int, default=None)
    parser.add_argument("--report_to", type=str, default=None)
    parser.add_argument("--push_to_hub", action="store_true", default=False)
    parser.add_argument("--hub_model_id", type=str, default=None)
    parser.add_argument("--validation_prompt", type=str, default=None)
    parser.add_argument("--gradient_checkpointing", action="store_true", default=None)
    parser.add_argument("--model_dir", type=str, default=os.environ.get("SM_MODEL_DIR", "./outputs"))
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    config = {}
    if os.path.exists(args.config):
        config = load_config(args.config)

    model_name = args.pretrained_model_name_or_path or config.get("model", {}).get("pretrained_model_name_or_path", "stabilityai/stable-diffusion-xl-base-1.0")
    vae_name = config.get("model", {}).get("pretrained_vae_model_name_or_path", "madebyollin/sdxl-vae-fp16-fix")
    instance_data_dir = args.instance_data_dir or config.get("data", {}).get("instance_data_dir", "./data/instance_images")
    instance_prompt = args.instance_prompt or config.get("training", {}).get("instance_prompt", "a photo of sks subject")
    output_dir = args.output_dir or config.get("output", {}).get("output_dir", "./outputs/lora_weights")
    resolution = args.resolution or config.get("training", {}).get("resolution", 1024)
    train_batch_size = args.train_batch_size or config.get("training", {}).get("train_batch_size", 1)
    max_train_steps = args.max_train_steps or config.get("training", {}).get("max_train_steps", 1000)
    learning_rate = args.learning_rate or config.get("training", {}).get("learning_rate", 1e-4)
    lora_rank = args.lora_rank or config.get("lora", {}).get("rank", 16)
    lora_alpha = args.lora_alpha or config.get("lora", {}).get("alpha", 16)
    mixed_precision = args.mixed_precision or config.get("training", {}).get("mixed_precision", "fp16")
    seed = args.seed if args.seed is not None else config.get("training", {}).get("seed", 42)
    checkpointing_steps = args.checkpointing_steps or config.get("training", {}).get("checkpointing_steps", 250)
    report_to = args.report_to or config.get("output", {}).get("report_to", "wandb")
    gradient_accumulation_steps = config.get("training", {}).get("gradient_accumulation_steps", 4)
    lr_scheduler_type = config.get("training", {}).get("lr_scheduler", "constant")
    lr_warmup_steps = config.get("training", {}).get("lr_warmup_steps", 0)
    gradient_checkpointing = args.gradient_checkpointing if args.gradient_checkpointing is not None else config.get("training", {}).get("gradient_checkpointing", True)
    train_text_encoder = config.get("lora", {}).get("train_text_encoder", False)
    target_modules_unet = config.get("lora", {}).get("target_modules_unet", ["to_k", "to_q", "to_v", "to_out.0"])
    push_to_hub = args.push_to_hub or config.get("hub", {}).get("push_to_hub", False)
    hub_model_id = args.hub_model_id or config.get("hub", {}).get("hub_model_id")

    sm_channel_training = os.environ.get("SM_CHANNEL_TRAINING")
    if sm_channel_training:
        instance_data_dir = sm_channel_training
        output_dir = args.model_dir

    logging_dir = Path(output_dir) / "logs"
    project_config = ProjectConfiguration(project_dir=output_dir, logging_dir=str(logging_dir))

    accelerator = Accelerator(
        gradient_accumulation_steps=gradient_accumulation_steps,
        mixed_precision=mixed_precision,
        log_with=report_to if report_to != "none" else None,
        project_config=project_config,
    )

    logging.basicConfig(format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO)
    logger.info(accelerator.state, main_process_only=False)

    if seed is not None:
        set_seed(seed)

    os.makedirs(output_dir, exist_ok=True)

    tokenizer_one = AutoTokenizer.from_pretrained(model_name, subfolder="tokenizer", use_fast=False)
    tokenizer_two = AutoTokenizer.from_pretrained(model_name, subfolder="tokenizer_2", use_fast=False)

    noise_scheduler = DDPMScheduler.from_pretrained(model_name, subfolder="scheduler")
    vae = AutoencoderKL.from_pretrained(vae_name, torch_dtype=torch.float16)
    unet = UNet2DConditionModel.from_pretrained(model_name, subfolder="unet")

    from transformers import CLIPTextModel, CLIPTextModelWithProjection
    text_encoder_one = CLIPTextModel.from_pretrained(model_name, subfolder="text_encoder")
    text_encoder_two = CLIPTextModelWithProjection.from_pretrained(model_name, subfolder="text_encoder_2")

    vae.requires_grad_(False)
    text_encoder_one.requires_grad_(False)
    text_encoder_two.requires_grad_(False)
    unet.requires_grad_(False)

    weight_dtype = torch.float16 if mixed_precision == "fp16" else torch.float32
    vae.to(accelerator.device, dtype=weight_dtype)
    text_encoder_one.to(accelerator.device, dtype=weight_dtype)
    text_encoder_two.to(accelerator.device, dtype=weight_dtype)

    unet_lora_config = LoraConfig(
        r=lora_rank,
        lora_alpha=lora_alpha,
        init_lora_weights="gaussian",
        target_modules=target_modules_unet,
    )
    unet.add_adapter(unet_lora_config)

    if gradient_checkpointing:
        unet.enable_gradient_checkpointing()

    lora_layers = list(filter(lambda p: p.requires_grad, unet.parameters()))
    logger.info(f"Trainable LoRA parameters: {sum(p.numel() for p in lora_layers):,}")

    optimizer = torch.optim.AdamW(lora_layers, lr=learning_rate, betas=(0.9, 0.999), weight_decay=1e-2, eps=1e-8)

    train_dataset = InstanceDataset(
        instance_data_dir=instance_data_dir,
        instance_prompt=instance_prompt,
        tokenizer_one=tokenizer_one,
        tokenizer_two=tokenizer_two,
        size=resolution,
    )
    train_dataloader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True, num_workers=0)

    lr_scheduler = get_scheduler(
        lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps,
    )

    unet, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        unet, optimizer, train_dataloader, lr_scheduler
    )

    if accelerator.is_main_process:
        tracker_config = {"lora_rank": lora_rank, "lora_alpha": lora_alpha, "learning_rate": learning_rate, "max_train_steps": max_train_steps, "resolution": resolution, "model": model_name}
        if report_to and report_to != "none":
            accelerator.init_trackers("lora-sdxl-training", config=tracker_config)

    global_step = 0
    logger.info("***** Starting training *****")

    progress_bar = None
    try:
        from tqdm.auto import tqdm
        progress_bar = tqdm(range(max_train_steps), desc="Training")
    except ImportError:
        pass

    unet.train()
    while global_step < max_train_steps:
        for batch in train_dataloader:
            with accelerator.accumulate(unet):
                pixel_values = batch["pixel_values"].to(dtype=weight_dtype)
                latents = vae.encode(pixel_values).latent_dist.sample()
                latents = latents * vae.config.scaling_factor
                noise = torch.randn_like(latents)
                bsz = latents.shape[0]
                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device).long()
                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
                prompt_embeds_one = text_encoder_one(batch["input_ids_one"].to(accelerator.device), output_hidden_states=True)
                prompt_embeds_two = text_encoder_two(batch["input_ids_two"].to(accelerator.device), output_hidden_states=True)
                pooled_prompt_embeds = prompt_embeds_two[0]
                prompt_embeds = torch.cat([prompt_embeds_one.hidden_states[-2], prompt_embeds_two.hidden_states[-2]], dim=-1)
                add_time_ids = torch.tensor([[resolution, resolution, 0, 0, resolution, resolution]], device=latents.device).repeat(bsz, 1)
                added_cond_kwargs = {"text_embeds": pooled_prompt_embeds, "time_ids": add_time_ids}
                model_pred = unet(noisy_latents, timesteps, prompt_embeds, added_cond_kwargs=added_cond_kwargs).sample
                loss = F.mse_loss(model_pred.float(), noise.float(), reduction="mean")
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            if accelerator.sync_gradients:
                global_step += 1
                if progress_bar:
                    progress_bar.update(1)
                logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
                if progress_bar:
                    progress_bar.set_postfix(**logs)
                accelerator.log(logs, step=global_step)
                if global_step % checkpointing_steps == 0 and accelerator.is_main_process:
                    save_path = os.path.join(output_dir, f"checkpoint-{global_step}")
                    accelerator.save_state(save_path)
                    logger.info(f"Saved checkpoint to {save_path}")

            if global_step >= max_train_steps:
                break

    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        unet_model = accelerator.unwrap_model(unet)
        unet_model.save_pretrained(output_dir)
        logger.info(f"Saved LoRA weights to {output_dir}")
        if push_to_hub and hub_model_id:
            logger.info(f"Uploading to Hugging Face Hub: {hub_model_id}")
            upload_folder(
                repo_id=hub_model_id,
                folder_path=output_dir,
                commit_message="End of LoRA training",
                ignore_patterns=["checkpoint-*", "logs"],
            )

    accelerator.end_training()
    logger.info("Training complete!")


if __name__ == "__main__":
    main()

Your unified training script is complete. It handles configuration from both YAML and CLI, loads the full SDXL model stack, injects LoRA adapters into the attention layers, and runs a noise-prediction training loop with checkpointing. Next, you'll build the local execution pipeline to actually run this script on your GPU.

Build the Local Execution Pipeline

Your training script is complete and handles everything from loading SDXL to saving LoRA checkpoints. Now you need a way to actually launch it on your local GPU with the right settings.

In this step, you will create a bash launcher script that orchestrates the full local training workflow. You will also prepare a small sample dataset and run a quick 50-step test to confirm the entire pipeline works end-to-end before committing to a longer training run.

In this step, get ready to:

• Write a bash script that automates environment setup and launches training via Hugging Face Accelerate.
• Prepare a sample dataset of images with captions in the expected directory format.
• Validate the full pipeline with a short 50-step test run and verify LoRA weight output.

Write the local training launcher script

A shell script wraps the repetitive commands you would otherwise type every time you want to train. It ensures the virtual environment is active, dependencies are current, and Accelerate launches with the correct mixed-precision flag.

• Create a file called local_run.sh in the root of your diffusers-lora-dual-pipeline folder, then paste the following code:

#!/bin/bash
# Local GPU training script for diffusers-lora-dual-pipeline
# Requires: NVIDIA GPU with 16GB+ VRAM, Python 3.10+

set -e

echo "=== diffusers-lora-dual-pipeline: Local Training ==="

# Create virtual environment if it doesn't exist
if [ ! -d "venv" ]; then
    echo "[1/4] Creating virtual environment..."
    python -m venv venv
fi

# Activate virtual environment
echo "[2/4] Activating virtual environment..."
source venv/bin/activate

# Install dependencies
echo "[3/4] Installing dependencies..."
pip install --quiet --upgrade pip
pip install --quiet -r requirements.txt

# Configure accelerate (single GPU, fp16)
echo "[4/4] Configuring Accelerate and launching training..."
accelerate launch --mixed_precision=fp16 src/train_lora.py \
    --config config/default.yaml \
    --output_dir ./outputs/lora_weights \
    --instance_data_dir ./data/instance_images \
    --report_to wandb

echo "=== Training Complete ==="
echo "LoRA weights saved to: ./outputs/lora_weights/"

What does this script do?

• set -e tells bash to stop immediately if any command fails, so errors don't cascade silently.
• The if [ ! -d "venv" ] block creates a virtual environment only when one does not already exist. Since you already have one from the previous step, this check will skip creation.
• accelerate launch --mixed_precision=fp16 starts the training script through Accelerate, which handles device placement and fp16 mixed-precision automatically.
• The trailing backslashes (\) split a single long command across multiple lines for readability.

• Save local_run.sh.
• Make the script executable by running:

chmod +x local_run.sh

• Confirm the file exists by running:

ls local_run.sh

You should see local_run.sh listed in the output.

Don't see local_run.sh?

Make sure you saved the file in the root of your diffusers-lora-dual-pipeline folder, not inside src/ or config/.

Help me find my local_run.sh file

✔️ Awesome, I've got everything!

Great. Make sure you have saved local_run.sh and made it executable before continuing.

ⓧ I'd like to double check the full code

#!/bin/bash
# Local GPU training script for diffusers-lora-dual-pipeline
# Requires: NVIDIA GPU with 16GB+ VRAM, Python 3.10+

set -e

echo "=== diffusers-lora-dual-pipeline: Local Training ==="

# Create virtual environment if it doesn't exist
if [ ! -d "venv" ]; then
    echo "[1/4] Creating virtual environment..."
    python -m venv venv
fi

# Activate virtual environment
echo "[2/4] Activating virtual environment..."
source venv/bin/activate

# Install dependencies
echo "[3/4] Installing dependencies..."
pip install --quiet --upgrade pip
pip install --quiet -r requirements.txt

# Configure accelerate (single GPU, fp16)
echo "[4/4] Configuring Accelerate and launching training..."
accelerate launch --mixed_precision=fp16 src/train_lora.py \
    --config config/default.yaml \
    --output_dir ./outputs/lora_weights \
    --instance_data_dir ./data/instance_images \
    --report_to wandb

echo "=== Training Complete ==="
echo "LoRA weights saved to: ./outputs/lora_weights/"

Prepare a sample dataset and run a training test

The training script expects a folder of images, with an optional metadata.jsonl file that pairs each image with a caption. Without the metadata file, the script falls back to using the instance_prompt from your config for every image.

• Create the dataset directory by running:

mkdir -p data/instance_images

• Copy 5 to 10 images of your training subject into data/instance_images/. Use .jpg or .png format.
• Optionally, create a metadata.jsonl file inside data/instance_images/ with one JSON line per image:

{"file_name": "image_01.jpg", "text": "a photo of sks subject sitting on a bench"}
{"file_name": "image_02.jpg", "text": "a photo of sks subject standing outdoors"}
{"file_name": "image_03.jpg", "text": "a close-up photo of sks subject"}

What is the sks token?

The token sks is a rare identifier that the model learns to associate with your specific subject during training. By using a word that has no strong pre-existing meaning in the model's vocabulary, you avoid conflicting with concepts SDXL already knows.

At inference time, you include sks in your prompt to activate the learned LoRA weights for that subject.

• Confirm your images are in the right place by running:

ls data/instance_images/

You should see your image files listed, along with metadata.jsonl if you created one.

Now run a short 50-step training test to validate the pipeline without waiting for a full 1000-step run.

• Launch a short test run by running this command:

accelerate launch --mixed_precision=fp16 src/train_lora.py \
    --config config/default.yaml \
    --output_dir ./outputs/lora_weights \
    --instance_data_dir ./data/instance_images \
    --max_train_steps 50 \
    --report_to none

Why override max_train_steps and report_to?

Passing --max_train_steps 50 overrides the 1000-step default in your config so the test finishes in a few minutes. Setting --report_to none skips W&B logging, which means you do not need a wandb account configured for this initial validation.

Training takes a few minutes depending on your GPU. You should see a progress bar counting from 0 to 50, with loss values printed at each step.

Seeing a CUDA out of memory error?

SDXL LoRA training at 1024 resolution requires approximately 16 GB of VRAM. If your GPU has less, try reducing resolution by adding --resolution 512 to the command above.

Also confirm that no other processes are using your GPU. Run nvidia-smi to check current VRAM usage.

Help me troubleshoot my CUDA memory error

Verify the checkpoint output

After the 50-step run completes, the training script saves the final LoRA adapter weights as a safetensors file. This small file (typically 3 to 8 MB) contains only the trained LoRA parameters, not the full SDXL model.

• Check that the LoRA weights file exists by running:

ls outputs/lora_weights/

You should see pytorch_lora_weights.safetensors in the output. This is the trained LoRA adapter file that you will load during inference.

• Confirm the file size is reasonable (a few MB) by running:

ls -lh outputs/lora_weights/

You should see a file size in the range of 3 to 8 MB. If you see 0 bytes, training likely failed silently.

Why are LoRA weights so small?

LoRA works by injecting small trainable matrices into the attention layers of the frozen base model. Only these low-rank matrices are saved, not the billions of SDXL parameters.

This is what makes LoRA practical for fine-tuning large models on consumer GPUs. You get personalized outputs without needing to store or distribute a full model copy.

Output directory is empty?

If outputs/lora_weights/ is empty or does not exist, scroll up in your terminal and check for errors during training. Common causes include an empty data/instance_images/ folder or images that cannot be opened.

Verify you have at least 5 valid images by running ls data/instance_images/ and confirming you see .jpg or .png files.

Help me debug why no weights were saved

Your local pipeline is validated. Next up, you will build the SageMaker submission script that packages your data to S3 and launches the same training script on a cloud GPU.

Build the SageMaker Submission Pipeline

Your local training pipeline works. You proved LoRA fine-tuning runs end-to-end on your own GPU.

Production ML teams rarely train on a single developer machine. They submit jobs to managed GPU clusters that spin up on demand, run training, save artifacts to cloud storage, and shut down automatically. In this step, you will build a submission script that packages your data to S3, provisions a SageMaker GPU instance, runs training with the same script you tested locally, and retrieves the trained LoRA weights.

In this step, get ready to:

• Write the SageMaker submission script with S3 data upload and IAM role configuration.
• Configure the PyTorch Estimator with hyperparameters, environment variables, and automatic dependency installation.
• Add post-training artifact download from S3 to your local machine.

Write the submission script with S3 upload and IAM role setup

The submission script handles three jobs: uploading your training images to S3, configuring the GPU instance with the right framework and dependencies, and downloading the trained weights when the job finishes.

• Create a new file called sagemaker_submit.py in the root of your diffusers-lora-dual-pipeline folder.
• Add the file header, imports, and argument parsing function by pasting this code:

#!/usr/bin/env python3
"""Submit LoRA fine-tuning job to Amazon SageMaker."""

import argparse
import os
import time
from pathlib import Path

import boto3
import sagemaker
from sagemaker.pytorch import PyTorch


def parse_args():
    parser = argparse.ArgumentParser(description="Submit LoRA training to SageMaker")
    parser.add_argument("--instance_data_dir", type=str, default="./data/instance_images")
    parser.add_argument("--s3_bucket", type=str, default=None)
    parser.add_argument("--s3_prefix", type=str, default="lora-training")
    parser.add_argument("--instance_type", type=str, default="ml.g5.2xlarge")
    parser.add_argument("--role", type=str, default=None)
    parser.add_argument("--max_train_steps", type=int, default=1000)
    parser.add_argument("--lora_rank", type=int, default=16)
    parser.add_argument("--learning_rate", type=float, default=1e-4)
    parser.add_argument("--resolution", type=int, default=1024)
    parser.add_argument("--instance_prompt", type=str, default="a photo of sks subject")
    parser.add_argument("--output_dir", type=str, default="./outputs/sagemaker_results")
    parser.add_argument("--hf_token", type=str, default=None)
    parser.add_argument("--wandb_api_key", type=str, default=None)
    return parser.parse_args()

What does this code do?

The imports bring in boto3 (the AWS SDK for Python) and the sagemaker library. The PyTorch import from sagemaker.pytorch gives you a pre-built estimator class that knows how to provision GPU instances with PyTorch Deep Learning Containers.

The argument parser defines all the knobs you can tweak at submission time: where your data lives, which instance type to use, the LoRA hyperparameters, and optional API keys for logging.

• Below the parse_args function, add the start of main() with session initialization and the S3 data upload:

def main():
    args = parse_args()
    # Initialize the SageMaker session and resolve bucket/role
    sagemaker_session = sagemaker.Session()
    bucket = args.s3_bucket or sagemaker_session.default_bucket()
    role = args.role or sagemaker.get_execution_role()

    # Upload training images to S3
    s3_data_uri = sagemaker_session.upload_data(
        path=args.instance_data_dir,
        bucket=bucket,
        key_prefix=f"{args.s3_prefix}/data"
    )

What does this code do?

sagemaker.Session() creates a session object tied to your AWS credentials. It handles all API calls to the SageMaker service.

sagemaker.get_execution_role() retrieves the IAM role that gives SageMaker permission to access your S3 data, provision instances, and write outputs. If you are running this from a SageMaker notebook, it auto-detects the role. From your local machine, you pass it explicitly via the --role argument.

upload_data() pushes your local training images to S3, returning an s3:// URI that SageMaker can feed directly into the training container.

SageMaker default bucket

If you do not pass --s3_bucket, the SDK creates a bucket named sagemaker-{region}-{account_id} automatically. This keeps your training artifacts organized without manual bucket creation.

Configure the estimator with hyperparameters and environment variables

The PyTorch Estimator needs to know which framework image to use, what hyperparameters to pass to your training script, and which secrets to inject as environment variables. Because you set source_dir=".", SageMaker automatically finds and installs your requirements.txt inside the Deep Learning Container before training starts.

• Add the environment variable configuration and hyperparameters dictionary below the S3 upload code:

    # Pass API tokens as environment variables (not logged)
    environment = {}
    if args.hf_token:
        environment["HF_TOKEN"] = args.hf_token
    elif os.environ.get("HF_TOKEN"):
        environment["HF_TOKEN"] = os.environ["HF_TOKEN"]
    if args.wandb_api_key:
        environment["WANDB_API_KEY"] = args.wandb_api_key
    elif os.environ.get("WANDB_API_KEY"):
        environment["WANDB_API_KEY"] = os.environ["WANDB_API_KEY"]

    # Hyperparameters passed as CLI args to train_lora.py
    hyperparameters = {
        "max_train_steps": args.max_train_steps,
        "lora_rank": args.lora_rank,
        "learning_rate": args.learning_rate,
        "resolution": args.resolution,
        "instance_prompt": args.instance_prompt,
        "mixed_precision": "fp16",
        "checkpointing_steps": 250,
        "report_to": "wandb" if environment.get("WANDB_API_KEY") else "none",
    }

What does this code do?

Environment variables are injected into the training container at runtime. HF_TOKEN lets the container download gated models from Hugging Face Hub. WANDB_API_KEY enables metric logging to Weights & Biases.

The script checks both CLI arguments and local environment variables as fallbacks. This way you can either export the keys in your shell or pass them explicitly at submission time.

The hyperparameters dictionary gets serialized and passed as CLI arguments to src/train_lora.py inside the container. SageMaker converts each key-value pair into --key value format automatically.

• Add the PyTorch Estimator creation and the fit() call that launches training:

    # Create the PyTorch Estimator targeting a GPU DLC
    estimator = PyTorch(
        entry_point="src/train_lora.py",
        source_dir=".",
        role=role,
        instance_count=1,
        instance_type=args.instance_type,
        framework_version="2.4.0",
        py_version="py311",
        hyperparameters=hyperparameters,
        environment=environment,
        output_path=f"s3://{bucket}/{args.s3_prefix}/output",
        max_run=7200,
        sagemaker_session=sagemaker_session,
    )

    # Launch the training job and wait for completion
    estimator.fit({"training": s3_data_uri}, wait=True)

What does this code do?

entry_point tells SageMaker which script to execute. source_dir="." packages your entire project directory (including requirements.txt) and uploads it to the training container. SageMaker installs the packages from requirements.txt automatically before running your script.

framework_version="2.4.0" and py_version="py311" select a pre-built Deep Learning Container with PyTorch 2.4.0 and Python 3.11. The max_run=7200 parameter caps the job at 2 hours to prevent runaway costs.

estimator.fit() starts the job. The {"training": s3_data_uri} dictionary maps the S3 path to the SM_CHANNEL_TRAINING environment variable inside the container. Your train_lora.py already checks for this variable and uses it as the data directory.

Cost warning

The default instance type ml.g5.2xlarge costs approximately $1.52/hr on-demand. A 1000-step training run on SDXL typically takes 30-60 minutes. The max_run=7200 safety cap ensures the job stops after 2 hours even if something goes wrong.

Add post-training artifact download from S3

When training finishes, SageMaker compresses your model artifacts into a model.tar.gz file and stores it in S3. The final section of the script downloads and extracts these weights to your local machine so you can use them for inference.

• Add the artifact download logic and the if __name__ block below the estimator.fit() call:

    # Download trained LoRA weights from S3
    os.makedirs(args.output_dir, exist_ok=True)
    model_data = estimator.model_data
    s3_client = boto3.client("s3")
    output_path = Path(args.output_dir)

    # Parse S3 URI and download the tar archive
    if model_data.startswith("s3://"):
        parts = model_data.replace("s3://", "").split("/", 1)
        s3_bucket_name = parts[0]
        s3_key = parts[1]
        local_tar = output_path / "model.tar.gz"
        s3_client.download_file(s3_bucket_name, s3_key, str(local_tar))
        # Extract the LoRA weights from the archive
        import tarfile
        with tarfile.open(local_tar, "r:gz") as tar:
            tar.extractall(path=str(output_path))
        local_tar.unlink()

    print(f"LoRA weights saved to: {args.output_dir}")


if __name__ == "__main__":
    main()

What does this code do?

estimator.model_data returns the S3 URI of the compressed model output. The code parses this URI to extract the bucket name and key, then uses boto3 to download the model.tar.gz archive.

After extraction, your outputs/sagemaker_results/ directory contains the same pytorch_lora_weights.safetensors file that local training produces. The .unlink() call deletes the tar archive to save disk space.

• Save sagemaker_submit.py.
• Verify the file exists by running:

ls sagemaker_submit.py

• You should see sagemaker_submit.py listed in the output.

File not showing up?

Make sure you saved the file in the root of your diffusers-lora-dual-pipeline folder, not inside src/. The estimator uses source_dir="." to package the entire project root, so the script must be at the top level.

help me troubleshoot why my sagemaker_submit.py file is not in the right location

• Test that the script parses without syntax errors by running:

python sagemaker_submit.py --help

• You should see the argument list printed with defaults for each flag.

Seeing a ModuleNotFoundError?

If Python cannot find sagemaker or boto3, make sure your virtual environment is active. Run source venv/bin/activate (macOS/Linux) or venv\Scripts\activate (Windows) and try again.

help me debug a ModuleNotFoundError when running sagemaker_submit.py

✔️ Awesome, I've got everything!

Great. Make sure you have saved sagemaker_submit.py before moving on.

ⓧ I'd like to double check the full code

Here is the complete sagemaker_submit.py file. Compare your version against it:

#!/usr/bin/env python3
"""Submit LoRA fine-tuning job to Amazon SageMaker."""

import argparse
import os
import time
from pathlib import Path

import boto3
import sagemaker
from sagemaker.pytorch import PyTorch


def parse_args():
    parser = argparse.ArgumentParser(description="Submit LoRA training to SageMaker")
    parser.add_argument("--instance_data_dir", type=str, default="./data/instance_images")
    parser.add_argument("--s3_bucket", type=str, default=None)
    parser.add_argument("--s3_prefix", type=str, default="lora-training")
    parser.add_argument("--instance_type", type=str, default="ml.g5.2xlarge")
    parser.add_argument("--role", type=str, default=None)
    parser.add_argument("--max_train_steps", type=int, default=1000)
    parser.add_argument("--lora_rank", type=int, default=16)
    parser.add_argument("--learning_rate", type=float, default=1e-4)
    parser.add_argument("--resolution", type=int, default=1024)
    parser.add_argument("--instance_prompt", type=str, default="a photo of sks subject")
    parser.add_argument("--output_dir", type=str, default="./outputs/sagemaker_results")
    parser.add_argument("--hf_token", type=str, default=None)
    parser.add_argument("--wandb_api_key", type=str, default=None)
    return parser.parse_args()


def main():
    args = parse_args()
    sagemaker_session = sagemaker.Session()
    bucket = args.s3_bucket or sagemaker_session.default_bucket()
    role = args.role or sagemaker.get_execution_role()

    s3_data_uri = sagemaker_session.upload_data(
        path=args.instance_data_dir,
        bucket=bucket,
        key_prefix=f"{args.s3_prefix}/data"
    )

    environment = {}
    if args.hf_token:
        environment["HF_TOKEN"] = args.hf_token
    elif os.environ.get("HF_TOKEN"):
        environment["HF_TOKEN"] = os.environ["HF_TOKEN"]
    if args.wandb_api_key:
        environment["WANDB_API_KEY"] = args.wandb_api_key
    elif os.environ.get("WANDB_API_KEY"):
        environment["WANDB_API_KEY"] = os.environ["WANDB_API_KEY"]

    hyperparameters = {
        "max_train_steps": args.max_train_steps,
        "lora_rank": args.lora_rank,
        "learning_rate": args.learning_rate,
        "resolution": args.resolution,
        "instance_prompt": args.instance_prompt,
        "mixed_precision": "fp16",
        "checkpointing_steps": 250,
        "report_to": "wandb" if environment.get("WANDB_API_KEY") else "none",
    }

    estimator = PyTorch(
        entry_point="src/train_lora.py",
        source_dir=".",
        role=role,
        instance_count=1,
        instance_type=args.instance_type,
        framework_version="2.4.0",
        py_version="py311",
        hyperparameters=hyperparameters,
        environment=environment,
        output_path=f"s3://{bucket}/{args.s3_prefix}/output",
        max_run=7200,
        sagemaker_session=sagemaker_session,
    )

    estimator.fit({"training": s3_data_uri}, wait=True)

    os.makedirs(args.output_dir, exist_ok=True)
    model_data = estimator.model_data
    s3_client = boto3.client("s3")
    output_path = Path(args.output_dir)

    if model_data.startswith("s3://"):
        parts = model_data.replace("s3://", "").split("/", 1)
        s3_bucket_name = parts[0]
        s3_key = parts[1]
        local_tar = output_path / "model.tar.gz"
        s3_client.download_file(s3_bucket_name, s3_key, str(local_tar))
        import tarfile
        with tarfile.open(local_tar, "r:gz") as tar:
            tar.extractall(path=str(output_path))
        local_tar.unlink()

    print(f"LoRA weights saved to: {args.output_dir}")


if __name__ == "__main__":
    main()

Your cloud submission pipeline is ready. Next up, you will write the inference script that loads your trained LoRA weights and generates evaluation images.

Build the Inference and Evaluation Pipeline

Your training pipeline is complete. You have LoRA weights sitting in outputs/lora_weights/ ready to be used. The next challenge is turning those raw weight files into actual images.

In this step, you will build an inference script that loads the base Stable Diffusion XL model, injects your trained LoRA adapters on top, and generates evaluation images to verify training quality.

In this step, get ready to:

• Write an inference script that loads SDXL and injects LoRA weights.
• Implement batch image generation with configurable prompts and seeds.
• Test inference with your trained weights and generate evaluation images.

Write the inference script with argument parsing

The inference script needs to be flexible. Different prompts, different LoRA scales, and different output locations should all be configurable from the command line without editing code.

• Create a new file called src/inference.py and add the imports and argument parsing by pasting this code:

#!/usr/bin/env python3
"""Inference script for loading LoRA weights and generating images."""

import argparse
import os
from pathlib import Path

import torch
from diffusers import DiffusionPipeline, AutoencoderKL


def parse_args():
    # Configure all inference parameters as CLI arguments
    parser = argparse.ArgumentParser(description="Generate images with trained LoRA")
    parser.add_argument("--lora_weights_path", type=str, required=True)
    parser.add_argument("--base_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0")
    parser.add_argument("--vae_model", type=str, default="madebyollin/sdxl-vae-fp16-fix")
    parser.add_argument("--prompt", type=str, nargs="+", default=["a photo of sks subject in a garden"])
    parser.add_argument("--negative_prompt", type=str, default="blurry, low quality, distorted")
    parser.add_argument("--num_images_per_prompt", type=int, default=2)
    parser.add_argument("--guidance_scale", type=float, default=7.5)
    parser.add_argument("--num_inference_steps", type=int, default=30)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output_dir", type=str, default="./outputs/generated_images")
    parser.add_argument("--lora_scale", type=float, default=1.0)
    return parser.parse_args()

What does this code do?

The --lora_weights_path argument is required because the script needs to know which trained adapter to load. All other arguments have sensible defaults.

The --lora_scale parameter controls how strongly the LoRA adapter influences generation. A value of 1.0 means full influence, while 0.5 blends halfway between the base model and your fine-tune.

The --prompt argument accepts multiple values via nargs="+", so you can generate images for several prompts in a single run.

Implement pipeline loading and batch image generation

The core of inference is a three-step process: load the base SDXL pipeline, inject the LoRA weights, then run generation. The pipeline.load_lora_weights() method handles merging your adapter into the UNet's attention layers.

• In src/inference.py, add the main() function and entry point below the parse_args() function by pasting this code:

def main():
    args = parse_args()
    # Load the fp16-fix VAE to avoid numerical instability
    vae = AutoencoderKL.from_pretrained(args.vae_model, torch_dtype=torch.float16)
    # Load base SDXL pipeline with the fixed VAE
    pipeline = DiffusionPipeline.from_pretrained(
        args.base_model,
        vae=vae,
        torch_dtype=torch.float16,
    ).to("cuda")
    # Inject the trained LoRA adapter into the pipeline
    pipeline.load_lora_weights(args.lora_weights_path)
    os.makedirs(args.output_dir, exist_ok=True)
    # Seed the generator for reproducible outputs
    generator = torch.Generator(device="cuda").manual_seed(args.seed)
    # Generate images for each prompt in the batch
    for i, prompt in enumerate(args.prompt):
        images = pipeline(
            prompt=prompt,
            negative_prompt=args.negative_prompt,
            num_images_per_prompt=args.num_images_per_prompt,
            guidance_scale=args.guidance_scale,
            num_inference_steps=args.num_inference_steps,
            generator=generator,
            cross_attention_kwargs={"scale": args.lora_scale},
        ).images
        # Save each image with a structured filename
        for j, image in enumerate(images):
            filename = f"prompt_{i:02d}_image_{j:02d}.png"
            filepath = Path(args.output_dir) / filename
            image.save(filepath)


if __name__ == "__main__":
    main()

What does this code do?

The script loads the base SDXL model from Hugging Face, swaps in the fixed VAE (madebyollin/sdxl-vae-fp16-fix) to avoid fp16 numerical issues, then calls pipeline.load_lora_weights() to merge your trained adapter into the UNet attention layers.

The cross_attention_kwargs={"scale": args.lora_scale} parameter lets you control LoRA strength at inference time without retraining. This is useful for dialing in the balance between your style and the base model's generality.

Each image is saved with a structured naming pattern (prompt_00_image_00.png, prompt_00_image_01.png) so you can easily identify which prompt produced which output.

• Save src/inference.py.

Seeing an import error?

Make sure your virtual environment is active. The diffusers and torch packages must be installed in the same environment you are running from.

If you see an error about AutoencoderKL, verify your diffusers version is 0.38.0 by running pip show diffusers.

Help me debug an import error in my inference script.

✔️ Awesome, I've got everything!

Great. Double check you have saved src/inference.py before moving on.

ⓧ I'd like to double check the full code

#!/usr/bin/env python3
"""Inference script for loading LoRA weights and generating images."""

import argparse
import os
from pathlib import Path

import torch
from diffusers import DiffusionPipeline, AutoencoderKL


def parse_args():
    parser = argparse.ArgumentParser(description="Generate images with trained LoRA")
    parser.add_argument("--lora_weights_path", type=str, required=True)
    parser.add_argument("--base_model", type=str, default="stabilityai/stable-diffusion-xl-base-1.0")
    parser.add_argument("--vae_model", type=str, default="madebyollin/sdxl-vae-fp16-fix")
    parser.add_argument("--prompt", type=str, nargs="+", default=["a photo of sks subject in a garden"])
    parser.add_argument("--negative_prompt", type=str, default="blurry, low quality, distorted")
    parser.add_argument("--num_images_per_prompt", type=int, default=2)
    parser.add_argument("--guidance_scale", type=float, default=7.5)
    parser.add_argument("--num_inference_steps", type=int, default=30)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output_dir", type=str, default="./outputs/generated_images")
    parser.add_argument("--lora_scale", type=float, default=1.0)
    return parser.parse_args()


def main():
    args = parse_args()
    vae = AutoencoderKL.from_pretrained(args.vae_model, torch_dtype=torch.float16)
    pipeline = DiffusionPipeline.from_pretrained(
        args.base_model,
        vae=vae,
        torch_dtype=torch.float16,
    ).to("cuda")
    pipeline.load_lora_weights(args.lora_weights_path)
    os.makedirs(args.output_dir, exist_ok=True)
    generator = torch.Generator(device="cuda").manual_seed(args.seed)
    for i, prompt in enumerate(args.prompt):
        images = pipeline(
            prompt=prompt,
            negative_prompt=args.negative_prompt,
            num_images_per_prompt=args.num_images_per_prompt,
            guidance_scale=args.guidance_scale,
            num_inference_steps=args.num_inference_steps,
            generator=generator,
            cross_attention_kwargs={"scale": args.lora_scale},
        ).images
        for j, image in enumerate(images):
            filename = f"prompt_{i:02d}_image_{j:02d}.png"
            filepath = Path(args.output_dir) / filename
            image.save(filepath)


if __name__ == "__main__":
    main()

Test inference with your trained LoRA weights

You have trained weights from your earlier 50-step test run sitting in outputs/lora_weights/. Time to point the inference script at them and generate some images.

• Run the inference script with multiple prompts by running this command:

python src/inference.py \
    --lora_weights_path ./outputs/lora_weights \
    --prompt "a photo of sks subject in a garden" "a photo of sks subject at the beach" \
    --num_images_per_prompt 2 \
    --guidance_scale 7.5 \
    --num_inference_steps 30

What does this command do?

This passes two prompts to the script, generating 2 images per prompt for a total of 4 evaluation images. The --guidance_scale 7.5 controls how closely the model follows the text prompt (higher values mean stricter adherence).

The --num_inference_steps 30 sets the number of denoising steps. More steps generally produce higher quality but take longer.

• Confirm the images were generated by running:

ls outputs/generated_images/

You should see files named prompt_00_image_00.png, prompt_00_image_01.png, prompt_01_image_00.png, and prompt_01_image_01.png in the output.

What should I see in the images?

With only 50 training steps, the LoRA adapter has barely started learning. Your images may look mostly like the base SDXL model with subtle hints of your training data. This is expected.

For a production fine-tune, you would train for the full 1000 steps configured in config/default.yaml. The infrastructure you have built supports this. Just update --max_train_steps in your training run.

Seeing a CUDA out of memory error?

SDXL inference at 1024x1024 requires approximately 10-12GB of VRAM. If your GPU runs out of memory, try reducing the resolution by adding --resolution 768 or generating one image at a time with --num_images_per_prompt 1.

If you see a device mismatch error, make sure you have a CUDA-capable GPU available and that PyTorch detects it. Run python -c "import torch; print(torch.cuda.is_available())" to check.

Help me fix a CUDA memory or device error during LoRA inference.

Your full pipeline is operational. You can train LoRA adapters locally or on SageMaker, then generate images with a single command. Next up, you will take this further with an automatic upload to the Hugging Face Hub.

Secret mission

Automatic Hub Upload with Model Card

Your trained LoRA weights are sitting locally. In a real team workflow, you would share them through Hugging Face Hub with a professional model card that documents the training configuration. Can you extend the training script to generate and upload everything automatically?

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

SageMaker training jobs are billed per second and stop automatically when training completes. The main ongoing cost is S3 storage for your training data and model artifacts.

Cost warning

If you used the SageMaker path, your S3 bucket still holds training data and model output. S3 storage costs are small (cents per GB per month), but they accumulate if forgotten.

Delete your S3 bucket contents if you do not plan to retrain soon.

Resources you used:

• S3 bucket with training data and model artifacts (SageMaker path only).
• Local diffusers-lora-dual-pipeline project directory with trained weights, generated images, and virtual environment.
• Hugging Face Hub repository with uploaded LoRA weights (Secret Mission only).

Keep everything running

No action needed. Choose this if you plan to continue experimenting with your LoRA weights or retraining with new images.

• Your trained LoRA weights (pytorch_lora_weights.safetensors) are only 3-8 MB. They cost nothing to keep on disk.
• Your S3 bucket will incur small ongoing storage charges if you used the SageMaker path.
• Your Hugging Face Hub repository (if created during the Secret Mission) remains publicly accessible.

Pause - I'll come back to this later

No persistent compute resources are running. SageMaker training jobs stop automatically after completion.

• Your local project files stay on disk until you delete them manually.
• Your S3 data stays in the bucket until you remove it. Storage costs are minimal for the small dataset and model files.
• No action is required to pause this project.

Delete - I don't want to use this again

Remove all project resources and start fresh. Follow each item below to fully clean up.

Delete S3 bucket contents (SageMaker path only)

• Delete all training data and model artifacts from your S3 bucket by running:

aws s3 rm s3://[[S3_BUCKET="your-sagemaker-bucket-name"]]/lora-training/ --recursive

Verify no running SageMaker endpoints

• Confirm you have no active endpoints by running:

aws sagemaker list-endpoints

You should see an empty list. This project does not create persistent endpoints, but it is good practice to verify.

Delete Hugging Face Hub repository (Secret Mission only)

• If you uploaded LoRA weights to the Hub during the Secret Mission, delete the repository by running:

from huggingface_hub import delete_repo

delete_repo(repo_id="your-username/your-model-id", repo_type="model")

Replace the repo_id with the Hub model ID you used during training.

Delete local project files

• Remove the entire project directory from your Desktop by running:

macOS

cd ~/Desktop
rm -rf diffusers-lora-dual-pipeline

Windows

cd ~\Desktop
rmdir /s /q diffusers-lora-dual-pipeline

This removes all source code, trained weights, generated images, and the virtual environment.

Nice Work!

Nice Work!

You've built a professional dual-pipeline MLOps repository for LoRA fine-tuning! Your diffusers-lora-dual-pipeline project proves you can design ML training infrastructure that works identically across local development and cloud production environments using Stable Diffusion XL.

You've learned how to:

• Built a unified training script using diffusers + PEFT that runs identically on a local GPU and on Amazon SageMaker, with configurable LoRA adapters targeting the UNet attention layers.
• Created a SageMaker submission pipeline using PyTorch Estimator that packages data to S3, provisions a GPU instance, and retrieves trained LoRA weights automatically.
• Implemented an inference pipeline that loads the base SDXL model, injects trained LoRA weights via pipeline.load_lora_weights(), and generates evaluation images with configurable prompts and guidance scale.
• Secret Mission: Added automatic model upload to Hugging Face Hub with a generated model card containing training metadata, usage code snippets, and the trained LoRA safetensors.

Ready to quiz yourself?