14727cc744
First diffusion training path in Studio: train a LoRA on the SDXL U-Net from an image + caption dataset and export it as a diffusers .safetensors that the existing diffusion LoRA loader (and any diffusers pipeline) can load. core/training/diffusion_lora_trainer.py: - DiffusionLoraConfig with validation/defaults (rank, alpha, targets, lr, steps, grad accumulation, resolution, min-SNR gamma, gradient checkpointing, lr scheduler, seed, mixed precision). - discover_image_caption_pairs: captions from metadata.jsonl / captions.jsonl, per-image .txt/.caption sidecars, or a dreambooth instance_prompt fallback (pure, unit-tested). - run_diffusion_lora_training: the loop -- freeze base, PEFT-wrap the U-Net attention projections, VAE-encode (fp32 VAE to avoid the SDXL fp16 overflow), sample noise + timesteps, predict, MSE loss with optional min-SNR weighting (epsilon / v-prediction), AdamW + get_scheduler + grad accumulation + grad clipping, then export via save_lora_weights. Emits worker-protocol events (model_load_*, progress, complete) and polls should_stop for a clean stop with a partial save. - run_diffusion_training_process: mp.Queue subprocess adapter (event_queue / stop_queue), so the training worker can spawn it; plus a CLI entry point. Only SDXL (U-Net) is trained here; DiT families and the Studio UI form + route wiring are follow-ups. The trainer is decoupled and worker-ready. Tests: test_diffusion_lora_trainer.py covers caption discovery (metadata / sidecar / instance prompt / skip-uncaptioned / errors), config normalisation + validation, the SDXL add-time-ids, and the dict->config adapter. Verified live on GPU: a 60-step SDXL LoRA run lowers the loss, exports a ~45 MB adapter, and loading it back shifts generation from baseline (mean abs pixel diff ~55/255).
119 lines
4 KiB
Python
119 lines
4 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""CPU-only unit tests for the diffusion LoRA trainer's pure helpers.
|
|
|
|
The training loop needs a GPU + weights, but dataset discovery, config normalisation,
|
|
the SDXL add-time-ids, and the dict->config adapter are pure and tested here.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
from core.training.diffusion_lora_trainer import (
|
|
DEFAULT_LORA_TARGETS,
|
|
DiffusionLoraConfig,
|
|
_config_from_dict,
|
|
compute_sdxl_add_time_ids,
|
|
discover_image_caption_pairs,
|
|
)
|
|
|
|
|
|
def _touch(p):
|
|
p.write_bytes(b"")
|
|
|
|
|
|
def test_discover_prefers_metadata_then_sidecar_then_instance(tmp_path):
|
|
_touch(tmp_path / "a.png")
|
|
_touch(tmp_path / "b.jpg")
|
|
_touch(tmp_path / "c.webp")
|
|
# a.png captioned via metadata.jsonl
|
|
(tmp_path / "metadata.jsonl").write_text(
|
|
json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n", encoding="utf-8"
|
|
)
|
|
# b.jpg captioned via sidecar
|
|
(tmp_path / "b.txt").write_text("from sidecar", encoding="utf-8")
|
|
# c.webp falls back to the instance prompt
|
|
pairs = dict(discover_image_caption_pairs(tmp_path, instance_prompt="from instance"))
|
|
assert pairs[str(tmp_path / "a.png")] == "from metadata"
|
|
assert pairs[str(tmp_path / "b.jpg")] == "from sidecar"
|
|
assert pairs[str(tmp_path / "c.webp")] == "from instance"
|
|
|
|
|
|
def test_discover_skips_uncaptioned_without_instance_prompt(tmp_path):
|
|
_touch(tmp_path / "cap.png")
|
|
_touch(tmp_path / "nocap.png")
|
|
(tmp_path / "cap.caption").write_text("a caption", encoding="utf-8")
|
|
pairs = discover_image_caption_pairs(tmp_path)
|
|
assert pairs == [(str(tmp_path / "cap.png"), "a caption")]
|
|
|
|
|
|
def test_discover_captions_jsonl_and_image_key(tmp_path):
|
|
_touch(tmp_path / "x.png")
|
|
(tmp_path / "captions.jsonl").write_text(
|
|
json.dumps({"image": "x.png", "text": "hi"}) + "\n", encoding="utf-8"
|
|
)
|
|
assert discover_image_caption_pairs(tmp_path) == [(str(tmp_path / "x.png"), "hi")]
|
|
|
|
|
|
def test_discover_custom_caption_column(tmp_path):
|
|
_touch(tmp_path / "x.png")
|
|
(tmp_path / "metadata.jsonl").write_text(
|
|
json.dumps({"file_name": "x.png", "caption": "col"}) + "\n", encoding="utf-8"
|
|
)
|
|
assert discover_image_caption_pairs(tmp_path, caption_column="caption")[0][1] == "col"
|
|
|
|
|
|
def test_discover_empty_raises(tmp_path):
|
|
_touch(tmp_path / "x.png") # no captions anywhere, no instance prompt
|
|
with pytest.raises(ValueError, match="No captioned images"):
|
|
discover_image_caption_pairs(tmp_path)
|
|
|
|
|
|
def test_discover_missing_dir_raises(tmp_path):
|
|
with pytest.raises(FileNotFoundError):
|
|
discover_image_caption_pairs(tmp_path / "nope")
|
|
|
|
|
|
def test_config_normalized_defaults():
|
|
cfg = DiffusionLoraConfig(base_model="b", data_dir="d", output_dir="o").normalized()
|
|
assert cfg.lora_alpha == cfg.lora_rank # alpha defaults to rank
|
|
assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"kw",
|
|
[
|
|
{"train_steps": 0},
|
|
{"train_batch_size": 0},
|
|
{"gradient_accumulation_steps": 0},
|
|
{"lora_rank": 0},
|
|
{"resolution": 100}, # not a multiple of 8
|
|
{"resolution": 32}, # too small
|
|
{"mixed_precision": "int4"},
|
|
],
|
|
)
|
|
def test_config_normalized_validation(kw):
|
|
with pytest.raises(ValueError):
|
|
DiffusionLoraConfig(base_model="b", data_dir="d", output_dir="o", **kw).normalized()
|
|
|
|
|
|
def test_compute_sdxl_add_time_ids():
|
|
assert compute_sdxl_add_time_ids(1024) == (1024, 1024, 0, 0, 1024, 1024)
|
|
|
|
|
|
def test_config_from_dict_ignores_unknown_and_tuples_targets():
|
|
cfg = _config_from_dict(
|
|
{
|
|
"base_model": "b",
|
|
"data_dir": "d",
|
|
"output_dir": "o",
|
|
"lora_target_modules": ["to_q", "to_v"],
|
|
"unknown_field": 123, # must be ignored, not crash
|
|
}
|
|
)
|
|
assert cfg.lora_target_modules == ("to_q", "to_v")
|
|
assert not hasattr(cfg, "unknown_field")
|