diff --git a/README.md b/README.md index d5683e8..45be9d2 100644 --- a/README.md +++ b/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. --- diff --git a/config/examples/train_lora_sdxl_8gb_1.0.yaml b/config/examples/train_lora_sdxl_8gb_1.0.yaml new file mode 100644 index 0000000..d63b0b7 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index b080810..755e09d 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/lorakit/datasets.py b/src/lorakit/datasets.py index 1cd8b78..3961d4a 100644 --- a/src/lorakit/datasets.py +++ b/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 = [] diff --git a/src/lorakit/train.py b/src/lorakit/train.py index 83f67f5..84570f0 100644 --- a/src/lorakit/train.py +++ b/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 diff --git a/tests/test_datasets.py b/tests/test_datasets.py new file mode 100644 index 0000000..044bb8c --- /dev/null +++ b/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" diff --git a/tests/test_train.py b/tests/test_train.py new file mode 100644 index 0000000..7e187d4 --- /dev/null +++ b/tests/test_train.py @@ -0,0 +1,77 @@ +from pathlib import Path + +import pytest + + +def _make_minimal_config(): + """Return a minimal config dict that passes TrainJob.__init__ validation.""" + return { + "device": "cuda:0", + "allow_tf32": False, + "train": { + "dtype": "fp32", + "batch_size": 1, + "resolution": 64, + "gradient_accumulation_steps": 1, + "dataset_folder": "/fake", + "optimizer": "adamw", + "optimizer_params": {"lr": 1e-4}, + "lr_scheduler": "constant", + "lr_scheduler_params": {"num_warmup_steps": 0, "num_training_steps": 1}, + "lora": { + "target_modules": {"to_k": [4, 4], "to_q": [4, 4]}, + }, + }, + "model": {"name_or_path": "stabilityai/stable-diffusion-xl-base-1.0"}, + "class_prompt": "man", + "instant_prompt": "SKS", + } + + +def test_arch_check_raises_on_unsupported_gpu(monkeypatch): + """The arch check should raise RuntimeError when GPU CC is not in PyTorch's arch list.""" + + monkeypatch.setattr("torch.cuda.is_available", lambda: True) + monkeypatch.setattr("torch.cuda.get_device_capability", lambda idx=0: (12, 0)) + monkeypatch.setattr( + "torch.cuda.get_arch_list", + lambda: ["sm_75", "sm_80", "sm_86", "sm_90"], + ) + monkeypatch.setattr( + "torch.cuda.get_device_name", lambda idx=0: "NVIDIA GeForce RTX 5060 Ti" + ) + + # Prevent TrainJob.__init__ from checking that the dataset folder exists. + from lorakit import config as _config + + monkeypatch.setattr( + _config, "resolve_user_path", lambda path, **kw: Path(path) + ) + + from lorakit.train import TrainJob + + with pytest.raises(RuntimeError, match="is not supported by this PyTorch build"): + TrainJob(_make_minimal_config(), "1.0", "test_arch", "/tmp") + + +def test_arch_check_passes_on_supported_gpu(monkeypatch): + """The arch check should not raise when GPU CC is in PyTorch's arch list.""" + + monkeypatch.setattr("torch.cuda.is_available", lambda: True) + monkeypatch.setattr("torch.cuda.get_device_capability", lambda idx=0: (9, 0)) + monkeypatch.setattr( + "torch.cuda.get_arch_list", + lambda: ["sm_75", "sm_80", "sm_86", "sm_90"], + ) + monkeypatch.setattr("torch.cuda.get_device_name", lambda idx=0: "NVIDIA H100") + + from lorakit import config as _config + + monkeypatch.setattr( + _config, "resolve_user_path", lambda path, **kw: Path(path) + ) + + from lorakit.train import TrainJob + + # Should not raise + TrainJob(_make_minimal_config(), "1.0", "test_arch", "/tmp")