feat: add caption_extension, GPU arch check, 8GB config, and tests

README.md

@@ -57,7 +57,20 @@ D:/datasets/my_subject/
 └── ...
 ```
 
-Supported formats: `.jpg`, `.jpeg`, `.png`, `.webp`. Images are auto-resized and cropped to the training `resolution` (default 1024). No captions or subfolders needed.
+Supported formats: `.jpg`, `.jpeg`, `.png`, `.webp`. Images are auto-resized and cropped to the training `resolution` (default 1024).
+
+**Captions (optional).** To train with per-image captions instead of a single trigger token for all images (e.g. for tag-based character LoRAs), place a `.txt` file alongside each image with the same stem:
+
+```
+D:/datasets/my_subject/
+├── img01.jpg
+├── img01.txt    # "1girl, purple hair, blue eyes, smiling, outdoors"
+├── img02.png
+├── img02.txt    # "1girl, purple hair, blue eyes, sitting, glasses"
+└── ...
+```
+
+Then set `caption_extension: ".txt"` (next to `dataset_folder`) in your config. When sidecar files are present, each image is trained with its own caption — great for Illustrious-style LoRAs — instead of a single DreamBooth trigger token.
 
 **2. Copy and edit a config.** Pick one of the [example configs](#choosing-a-config), copy it out of `config/examples/`, and change at least:
 
@@ -97,6 +110,7 @@ All examples live in `config/examples/`. Copy one into `config/` and edit it.
 | `train_lora_sdxl_24gb_4090_4bit_1.0.yaml` | 24 GB | Fast | 4-bit QLoRA base. Lower VRAM, slightly slower than bf16. |
 | `train_lora_sdxl_24gb_4090_1.0.yaml` | 24 GB | — | Fully-commented reference of every option. |
 | `train_lora_sdxl_16gb_t4_1.0.yaml` | 16 GB | — | Low-VRAM (e.g. T4) using 4-bit quantization. |
+| `train_lora_sdxl_8gb_1.0.yaml` | 8 GB | — | 4-bit QLoRA + 768 px + 8-bit AdamW. For RTX 5060 Ti, 4060 Ti, etc. |
 
 Not sure? On a 24 GB card use the **bf16 metal** config. On 16 GB, use the **T4** config.
 
@@ -216,6 +230,7 @@ Ready-made profiling configs live in `config/examples/profile_4bit.yaml` and `co
 - **Out of memory** — set `gradient_checkpointing: true`, switch to a 4-bit config, or lower `resolution`.
 - **`Cannot find a working triton installation`** — run `uv sync` (Triton is a declared dependency); avoid running with a stale/hand-modified environment. `torch.compile` requires it.
 - **Quantization import errors** — run `uv sync` to (re)install `bitsandbytes`.
+- **`no kernel image is available for execution on the device` / unsupported compute capability** — your GPU (e.g. RTX 5060 Ti, Blackwell architecture) requires a PyTorch build with CUDA 13.0+ support. lorakit checks this at startup and prints the fix. In short: install a newer PyTorch wheel (`uv pip install torch --index-url https://download.pytorch.org/whl/cu130`), then reinstall `bitsandbytes` from the same session so its kernels are compiled for your GPU.
 
 ---
 

config/examples/train_lora_sdxl_8gb_1.0.yaml

@@ -0,0 +1,85 @@
+---
+# 4-bit QLoRA SDXL LoRA config for 8 GB GPUs (RTX 5060 Ti, RTX 4060 Ti, etc.)
+# Uses NF4 quantization, gradient checkpointing, 8-bit AdamW, and 768 px resolution
+# to fit training into 8 GB of VRAM.
+#
+# Copy this file, then set name, instant_prompt, class_prompt, and dataset_folder.
+job: train
+version: "1.0"
+name: "my_model_name"
+output_folder: "output"
+config:
+  logging_folder: "logs"
+  device: "cuda:0"
+  allow_tf32: true
+  matmul_precision: "high"
+  cudnn_benchmark: true
+  instant_prompt: "SKS"
+  class_prompt: "man"
+  train:
+    dataset_folder: "/path/to/images/folder"
+    #caption_extension: ".txt"  # uncomment to use per-image .txt sidecar captions
+    #resume_from_checkpoint: "latest"
+    batch_size: 1
+    max_train_steps: 2000
+    save_every: 250
+    checkpoints_total_limit: 3
+    max_grad_norm: 1.0
+    gradient_accumulation_steps: 1
+    seed: 42
+    dtype: bf16
+    gradient_checkpointing: true    # critical: saves ~30% VRAM
+    cache_latents: true             # encode images once, skip per-step VAE
+    train_text_encoder: false       # text encoders add ~2 GB
+    resolution: 768                 # 1024 OOMs on 8 GB; 768 still produces good quality
+    num_workers: 2
+    loss_decay: 0.995
+    disable_sampling: true          # sampling loads fp32 VAE, OOM risk on 8 GB
+    optimizer: "adamw8bit"          # 8-bit states save ~1 GB vs full AdamW
+    optimizer_params:
+      lr: 0.0001
+      betas: [0.9, 0.999]
+      weight_decay: 0.02
+      eps: 0.00000001
+    lora:
+      use_dora: false
+      init_lora_weights: "gaussian"
+      target_modules:
+        "to_k": [4, 4]
+        "to_q": [4, 4]
+        "to_v": [4, 4]
+        "to_out.0": [4, 4]
+    text_encoder_lora:
+      use_dora: false
+      init_lora_weights: "gaussian"
+      target_modules:
+        "q_proj": [4, 4]
+        "k_proj": [4, 4]
+        "v_proj": [4, 4]
+        "out_proj": [4, 4]
+    lr_scheduler: "constant"
+    lr_scheduler_params:
+      num_warmup_steps: 0
+      num_training_steps: 2000
+      num_cycles: 1
+      power: 1
+  model:
+    name_or_path: "stabilityai/stable-diffusion-xl-base-1.0"
+    quantization:
+      bits: 4                        # NF4 QLoRA: ~2 GB UNet vs ~7 GB fp16
+      quantize_text_encoder: false
+      bnb_4bit_quant_type: "nf4"
+      bnb_4bit_use_double_quant: true
+  # Sampling is disabled above; to enable preview images:
+  # 1. Set disable_sampling: false
+  # 2. Uncomment the block below and set your prompts
+  #sample:
+  #  sample_every: 250
+  #  prompts:
+  #    - "SKS man, portrait"
+  #  neg:
+  #    - "low quality, blurry"
+  #  seed: 42
+  #  walk_seed: true
+  #  guidance_scale: 7
+  #  sample_steps: 20

pyproject.toml

@@ -31,7 +31,7 @@ dependencies = [
 ]
 
 [dependency-groups]
-dev = ["ruff", "vulture"]
+dev = ["ruff", "vulture", "pytest>=8.0"]
 
 [project.scripts]
 lorakit = "lorakit.cli.lorakit:main"

src/lorakit/datasets.py

@@ -27,6 +27,7 @@ def __init__(
         repeats=1,
         center_crop=False,
         random_flip=False,
+        caption_extension=None,
     ):
         self.resolution = resolution
         self.center_crop = center_crop
@@ -71,12 +72,30 @@ def __init__(
                 )
             )
         instance_images = [Image.open(path) for path in image_paths]
+
         self.custom_instance_prompts = None
+        if caption_extension:
+            captions = []
+            for img_path in image_paths:
+                cap_path = img_path.with_suffix(caption_extension)
+                if cap_path.exists():
+                    content = cap_path.read_text(encoding="utf-8").strip()
+                    captions.append(content if content else None)
+                else:
+                    captions.append(None)
+            if any(captions):
+                self.custom_instance_prompts = captions
 
         self.instance_images = []
         for img in instance_images:
             self.instance_images.extend(itertools.repeat(img, repeats))
 
+        if self.custom_instance_prompts:
+            repeated = []
+            for cap in self.custom_instance_prompts:
+                repeated.extend([cap] * repeats)
+            self.custom_instance_prompts = repeated
+
         # image processing to prepare for using SD-XL micro-conditioning
         self.original_sizes = []
         self.crop_top_lefts = []

src/lorakit/train.py

@@ -256,6 +256,19 @@ def __init__(self, config, version, name, root_folder, *, config_path=None):
 
         print(f"Using device: {self._device}")
 
+        if torch.cuda.is_available() and self._device != "cpu":
+            device_idx = int(self._device.split(":")[-1]) if ":" in self._device else 0
+            cc = torch.cuda.get_device_capability(device_idx)
+            cc_str = f"sm_{cc[0]}{cc[1]}"
+            if cc_str not in torch.cuda.get_arch_list():
+                raise RuntimeError(
+                    f"\nGPU {torch.cuda.get_device_name(device_idx)} "
+                    f"(compute capability {cc_str}) is not supported by this PyTorch build.\n"
+                    f"Supported architectures: {' '.join(torch.cuda.get_arch_list())}\n"
+                    f"Install a PyTorch build that includes {cc_str} support, e.g.:\n"
+                    f"  pip install torch --index-url https://download.pytorch.org/whl/cu130\n"
+                )
+
         self._allow_tf32 = config.get("allow_tf32", False)
 
         if self._allow_tf32:
@@ -470,6 +483,8 @@ def __init__(self, config, version, name, root_folder, *, config_path=None):
             )
         )
 
+        self._caption_extension = self._train_config.get("caption_extension", None)
+
         self._num_workers = self._train_config.get("num_workers", 1)
         if self._num_workers > 1:
             print(f"Using {self._num_workers} workers for data loading")
@@ -1269,6 +1284,7 @@ def load_model_hook(models, input_dir):
             class_data_root=self._class_data_folder if self._with_prior_preservation else None,
             class_num=self._num_class_images,
             resolution=self._resolution,
+            caption_extension=self._caption_extension,
         )
         # `persistent_workers` avoids respawning the worker processes (and re-running
         # the dataset's expensive __init__ image preprocessing) on every epoch, which

tests/test_datasets.py

@@ -0,0 +1,116 @@
+import tempfile
+from pathlib import Path
+
+from PIL import Image
+
+from lorakit.datasets import DreamBoothDataset
+
+
+def test_loads_sidecar_captions():
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp = Path(tmp)
+        img = Image.new("RGB", (64, 64))
+        img.save(tmp / "test.png")
+        (tmp / "test.txt").write_text("cat, hat, mat", encoding="utf-8")
+
+        ds = DreamBoothDataset(
+            str(tmp),
+            instance_prompt="fallback",
+            class_prompt="thing",
+            resolution=64,
+            caption_extension=".txt",
+        )
+        assert ds.custom_instance_prompts == ["cat, hat, mat"], (
+            "custom_instance_prompts should load the sidecar caption"
+        )
+        assert ds[0]["instance_prompt"] == "cat, hat, mat", (
+            "__getitem__ should return the sidecar caption as instance_prompt"
+        )
+
+
+def test_falls_back_to_instance_prompt_when_no_sidecar():
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp = Path(tmp)
+        img = Image.new("RGB", (64, 64))
+        img.save(tmp / "test.png")
+
+        ds = DreamBoothDataset(
+            str(tmp),
+            instance_prompt="fallback",
+            class_prompt="thing",
+            resolution=64,
+            caption_extension=".txt",
+        )
+        assert ds.custom_instance_prompts is None, (
+            "custom_instance_prompts should be None when no sidecars exist"
+        )
+        assert ds[0]["instance_prompt"] == "fallback", (
+            "__getitem__ should fall back to instance_prompt"
+        )
+
+
+def test_empty_sidecar_falls_back_to_instance_prompt():
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp = Path(tmp)
+        img = Image.new("RGB", (64, 64))
+        img.save(tmp / "test.png")
+        (tmp / "test.txt").write_text("", encoding="utf-8")
+
+        ds = DreamBoothDataset(
+            str(tmp),
+            instance_prompt="fallback",
+            class_prompt="thing",
+            resolution=64,
+            caption_extension=".txt",
+        )
+        assert ds.custom_instance_prompts is None, (
+            "empty sidecar should not populate custom_instance_prompts"
+        )
+        assert ds[0]["instance_prompt"] == "fallback", (
+            "__getitem__ should fall back when caption is empty"
+        )
+
+
+def test_no_caption_extension_disables_feature():
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp = Path(tmp)
+        img = Image.new("RGB", (64, 64))
+        img.save(tmp / "test.png")
+        (tmp / "test.txt").write_text("cat, hat, mat", encoding="utf-8")
+
+        ds = DreamBoothDataset(
+            str(tmp),
+            instance_prompt="fallback",
+            class_prompt="thing",
+            resolution=64,
+        )
+        assert ds.custom_instance_prompts is None, (
+            "custom_instance_prompts should be None when caption_extension is not set"
+        )
+        assert ds[0]["instance_prompt"] == "fallback", (
+            "__getitem__ should use instance_prompt when caption_extension is not set"
+        )
+
+
+def test_repeats_apply_to_captions():
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp = Path(tmp)
+        img = Image.new("RGB", (64, 64))
+        img.save(tmp / "test.png")
+        (tmp / "test.txt").write_text("cat, hat, mat", encoding="utf-8")
+
+        ds = DreamBoothDataset(
+            str(tmp),
+            instance_prompt="fallback",
+            class_prompt="thing",
+            resolution=64,
+            repeats=3,
+            caption_extension=".txt",
+        )
+        assert ds.num_instance_images == 3, "1 image × 3 repeats = 3"
+        assert len(ds.custom_instance_prompts) == 3, (
+            "captions should be repeated to match images"
+        )
+        assert ds[0]["instance_prompt"] == "cat, hat, mat"
+        assert ds[1]["instance_prompt"] == "cat, hat, mat"
+        assert ds[2]["instance_prompt"] == "cat, hat, mat"