feat: add caption_extension, GPU arch check, 8GB config, and tests - #8
Open
wilsontr wants to merge 1 commit into
Open
feat: add caption_extension, GPU arch check, 8GB config, and tests#8wilsontr wants to merge 1 commit into
wilsontr wants to merge 1 commit into
Conversation
datasets.py: - Add `caption_extension` parameter to DreamBoothDataset for loading per-image captions from sidecar .txt files. Applies dataset repeats to captions for correct index alignment. train.py: - Parse `caption_extension` from config and pass it through to the dataset. - Add early GPU compute capability check at startup. If the installed PyTorch build does not include kernels for the current GPU (e.g. Blackwell sm_120), raise a RuntimeError with a clear fix message instead of failing cryptically during model loading. tests: - test_datasets.py: 5 tests covering sidecar loading, missing sidecars, empty sidecars, caption_extension disabled, and repeat alignment. - test_train.py: 2 tests for the GPU arch check (unsupported raises, supported passes). config/examples/train_lora_sdxl_8gb_1.0.yaml: - New 8GB VRAM config using 4-bit NF4 QLoRA, 768px resolution, 8-bit AdamW, gradient checkpointing, and latent caching. Demonstrates caption_extension usage. README.md: - Document caption_extension feature and sidecar file format. - Add 8GB config to the config table. - Add Blackwell GPU troubleshooting entry.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
caption_extension: Adds per-image sidecar caption support to DreamBoothDataset. When
caption_extension: ".txt"is set in the config, the dataset reads same-name.txtfiles alongside each image for tag-based LoRA training (Illustrious-style), falling back toinstance_promptwhen no sidecar exists.GPU arch check: Early startup check that compares the GPU's compute capability against PyTorch's compiled architectures. Raises a clear
RuntimeErrorwith actionable fix instructions instead of failing cryptically during model loading (e.g., Blackwellsm_120on acu126PyTorch build).8GB example config: New
config/examples/train_lora_sdxl_8gb_1.0.yamlusing 4-bit NF4 QLoRA, 768px resolution, 8-bit AdamW, gradient checkpointing, and latent caching.Tests: 7 passing tests covering sidecar loading edge cases and GPU arch check behavior (
pytest>=8.0added to dev dependencies).README: Documented caption_extension, 8GB config, and Blackwell GPU troubleshooting.
Changes
src/lorakit/datasets.pysrc/lorakit/train.pytests/test_datasets.pytests/test_train.pyconfig/examples/train_lora_sdxl_8gb_1.0.yamlREADME.mdpyproject.toml