sdxl-training

Changes

This update introduces benchmarking and metadata-related features, along with improvements to the trainer engine and documentation.

Detailed changes:

• Add benchmark.py for training performance evaluation
• Add test_benchmark.py and test_metadata_features.py to improve test coverage
• Enhance the trainer engine with advanced evaluation capabilities
• Improve LoRA method exporting with detailed example configurations
• Update documentation with comprehensive usage and configuration examples

Status

With this update completed, the training framework is now ready for official use.

For detailed usage instructions and examples, please refer to:

• Demo Guide: DEMO.md
• Usage Guide: USAGE.md

Minimum Hardware Requirements

Please note that this project has relatively high resource requirements. The following configuration is recommended as the minimum supported environment:

OS: Ubuntu 24.04 LTS
RAM: 64 GB
GPU VRAM: At least 12 GB
NVMe SSD (recommended for latent/text-encoder caching)

Running below these requirements may result in insufficient memory, significantly reduced performance, or inability to complete training.

Thank you for following the development of this project.

Commit

Changes Made

Checkpoint System:

• Added thread-safe locking for checkpoint writes to prevent corruption
• Required optimizer_state_dict validation for checkpoint integrity
• Stopped background writers in test cleanups to prevent resource leaks
• Added cache_vae_dtype config field with fp32/bf16 tradeoff options

Dataset & Caching:

• Migrated dataset caching from pickle to safetensors format for faster/safer storage
• Added cache_vae_dtype VAE precision control (fp32 safe, bf16 faster/lower VRAM)
• Added seed parameter for reproducible dataset shuffling and deterministic sampling

Performance Optimizations:

• Added TF32 matmul and cuDNN benchmark for CUDA hardware
• Added torch.compile support for UNet with dynamic shapes
• Enabled pin_memory on CUDA for faster data transfer
• Added VRAM monitoring throughout training phases
• Reset peak memory stats per epoch for accurate telemetry

Hardware Guardrails:

• Added GPU-first guard verification test (skips on CUDA-enabled machines)
• Added mixed_precision=‘no’ VRAM usage warnings for <16GB GPUs
• Better handling of CPU RNG states after map_location=“cuda” loads
• GPU memory cleanup after model offloading

Testing:

• Fix test_checkpoint.py to shutdown background writers
• Add GPU-first guard test that requires non-CUDA environment

Commit

feat: add gradient accumulation and offload checkpoint writes to a background worker

• Accumulate gradients across micro-steps: zero grad only at window start, scale loss by acc_steps, and step optimizer/scheduler at window end.
• Move snapshot serialization to a background thread so disk I/O and the “latest” copy no longer stall training; recovery checkpoints stay synchronous for crash-durable resume.
• Clone state to CPU (_to_cpu) before enqueueing to avoid GPU OOM and in-place corruption by later training updates.
• Treat rolling recovery checkpoints as a minimal resume set (LoRA+RNG only); full snapshots still store optimizer/scheduler/scaler state. Resume from a rolling checkpoint rebuilds optimizer momentum and LR schedule from scratch.
• Make CaptionProcessor RNG thread-safe with a lock for concurrent prefetch.
• Add flush()/shutdown() for deterministic test state and clean process exit.

Commit

Two small polish passes today!

• 📜 Quieter training logs — our console used to shout a full stats line every single step (so noisy 😵). Now it whispers just once every 1,000 steps. Don’t worry though — the little tqdm progress bar still wiggles along in real time, so you always know we’re alive and training! 🐣
• 🏷️ LoRA name tags — when we export a LoRA, we now tuck proper kohya-standard metadata into the safetensors header. That means sd-webui, ComfyUI, and friends finally recognize our SDXL LoRA and scale it correctly instead of squinting confused. We even peek at a lora_down weight to auto-detect the rank~ 🔍
  Small changes, cozy results!

Commit

fixed a sneaky little bug: SDXL micro-conditioning was using the resized image size and crop coords, which was a bit off~ now it happily uses the TRUE original size and original-image-space crop offsets, matching the kohya / diffusers convention, so the model finally understands each image’s real backstory

also, caption augmentation used the global random, so it got all mixed up on resume (´；ω；`) gave it its own little seeded RNG baby — its state is saved with the checkpoint and restored exactly, so training picks up right where it left off, no getting lost.

collate got a tiny shield too: if a batch mixes cached/uncached samples, it now complains LOUDLY instead of throwing a cryptic KeyError

Commit

Hi~ today the trainer went on a little diet!

• The VAE and Text Encoders were freeloaders. Once everything is cached, they just sit on the GPU sunbathing and eating ~2GB of VRAM for no reason. So now we politely walk them over to the CPU couch before training starts, and only invite them back if we actually need them.
• xformers joins the party. The UNet now uses memory-efficient attention when it’s installed, and quietly falls back to plain SDPA if it isn’t — no drama, no crashes.
• The caching progress bar finally tells the truth. Before, it counted already-cached samples as “work” and lied about the ETA. Now we pre-scan everything first, so the bar only moves for real encoding. Honest bar is best bar.
  Result: less VRAM, faster start, happier GPU. Nap successful.

Commit

Gave precaching a little diet today! Added a VAE-slicing toggle (off by default = faster), and now we hand VRAM back to the GPU after every batch (no more sneaky CUDA hoarding~). Single-file checkpoints now only carry the UNet home instead of lugging the whole pipeline. Shoved caption augmentation into the prefetch thread, and the progress bar got smarter — it counts only real encodes, not the lazy “already-cached” ones. The GPU is having an easy day, meow~

Commit

Two little wheels rolled into the trainer today:
🎯 Aspect-ratio bucketing: no more squishing every image into a square! Each pic now hops into the bucket whose ratio fits it best, gets a gentle cover-resize + center-crop, and BucketBatchSampler keeps same-bucket pals together so batches stack nice and tidy.
🏷️ Cute caption magic: Kohya-style tag shuffling, keep_tokens to pin the first few tags, and playful tag / whole-caption dropout — data augmentation now feels like sprinkling sugar 🍬.

Also fixed a sneaky SDXL detail: we now record each image’s real original_size and crop_ltrb and feed them into add_time_ids, instead of pretending everything is a 1024² square.
The cache grew up too — one combined cache_{hash}.pt carrying latents + text-encoder outputs + metadata, always saved to disk so reruns load in a blink, with shape checks that auto-recompute if you resize.
And the UNet now takes a nap 💤 until caching is done, freeing VRAM for the little GPUs that could.

Commit

✨ Precaching got a turbo boost! 🚀

Hi hi~ today the caching step stopped being such a lazy little snail! 🐌➡️🐇

• Before, we loaded one image, encoded it, saved it… one… by… one… zzz 😴 The GPU was basically napping while the disk did all the reading.

• Now we have a squad of worker threads (cache_workers) fetching images in parallel while the GPU chomps through them in batches (cache_batch_size). I/O and compute finally hold hands and work together! 🤝

• Bonus: a cute little progress log tells you img/s and ETA, so you know exactly when snack time ends. 🍪⏱️

• Also added a safety hug for resumes!
  If you try to resume a checkpoint with a different config, the trainer now gently stops and asks “are you suuure? type ‘yes’~” 🥺 No more silently mixing old weights with new settings and crying later. In non-interactive mode it just aborts to keep you safe. 💖

Commit

Added cute little tqdm progress bars to the caching & training loops so we can watch them go brrr~

Commit

Two handy little things landed in the trainer this time~

• 🧊 Caching: VAE latents and text-encoder outputs can now be pre-computed and stored ahead of time! Keep them in RAM or on disk. No more recomputing every step — less VRAM, faster training, yay~
• 🎯 UNet-only training: A new train_text_encoder switch (off by
  default). When off, only the UNet gets LoRA and the text encoders get to rest, so training stays light and breezy!
• 📝 Also refreshed the config example and added plenty of tests so nothing breaks ✅

Commit

Single-file SDXL checkpoints are now supported!

Commit

Initial implementation of the SDXL Training Toolkit! 
This is a very early and minimal version. The core training pipeline is now in place, but features like aspect ratio bucketing, VAE caching, and other optimizations are still on the way.

Commit

I just realized how complex it is to build an SDXL training pipeline from scratch… It might be a long while before the first version is ready to ship qwq