Skip to content

Conversation

@bghira
Copy link
Owner

@bghira bghira commented Sep 11, 2025

The old VAE cache was doing a lot of Python threading gymnastics with queues and ThreadPoolExecutors. Replaced most of it with trainingsample's Rust batch operations since they're faster.

Changes:

  • VAE cache now groups images by aspect ratio and processes them in actual batches
  • Image resize/crop operations use ts.batch_resize_images() and ts.batch_*_crop_images()
  • Added batch luminance calculation for metadata
  • TrainingSample class uses trainingsample for resizing when possible
  • Simplified concurrent image reading logic
  • Kept all the fallbacks so it doesn't break if Rust explodes

Performance:

  • ~6.8x speedup on image resize operations
  • Less Python GIL contention
  • Actually batches things instead of pretending to

Tested with SD3 PEFT LoRA training.

@bghira bghira added the enhancement New feature or request label Sep 11, 2025
@bghira bghira self-assigned this Sep 11, 2025
@bghira bghira requested a review from Copilot September 11, 2025 01:07
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR replaces the Python-based VAE cache image processing with Rust-based batch operations from the new trainingsample library, delivering significant performance improvements. The implementation maintains backward compatibility through fallback mechanisms.

Key changes:

  • Integrated trainingsample library for ~6.8x faster batch image operations
  • Replaced complex Python threading with efficient batch processing grouped by aspect ratio
  • Added new BatchedTrainingSamples class as interface to Rust operations

Reviewed Changes

Copilot reviewed 6 out of 9 changed files in this pull request and generated 4 comments.

File Description
pyproject.toml, install/rocm/pyproject.toml, install/apple/pyproject.toml Added trainingsample ^0.1.0 dependency across all platform configurations
helpers/image_manipulation/batched_training_samples.py New wrapper class providing batch operations for resize, crop, and luminance calculations
helpers/image_manipulation/training_sample.py Enhanced TrainingSample to use trainingsample for resize/crop operations with PIL fallbacks
helpers/caching/vae.py Refactored VAE cache to use aspect-grouped batch processing instead of complex threading

Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.

@bghira bghira merged commit d7f6a53 into main Sep 11, 2025
1 check passed
@bghira bghira deleted the feature/trainingsample-rust branch September 11, 2025 01:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants