95f783ae1e
# Conflicts: # studio/backend/core/training/diffusion_train_common.py # studio/backend/tests/test_diffusion_dit_trainer.py # studio/backend/tests/test_diffusion_training.py
434 lines
16 KiB
Python
434 lines
16 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,
|
|
_coerce_gradient_checkpointing,
|
|
_config_from_dict,
|
|
compute_sdxl_add_time_ids,
|
|
discover_image_caption_pairs,
|
|
resolve_train_steps,
|
|
)
|
|
|
|
|
|
def _touch(p):
|
|
p.write_bytes(b"")
|
|
|
|
|
|
def test_discover_prefers_sidecar_then_metadata_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 only
|
|
(tmp_path / "metadata.jsonl").write_text(
|
|
json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n", encoding = "utf-8"
|
|
)
|
|
# b.jpg captioned via sidecar only
|
|
(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_sidecar_overrides_metadata_row(tmp_path):
|
|
# A per-image sidecar is the user's explicit edit and must win over a metadata row
|
|
# for the same image (the labeling grid writes sidecars).
|
|
_touch(tmp_path / "a.png")
|
|
(tmp_path / "metadata.jsonl").write_text(
|
|
json.dumps({"file_name": "a.png", "text": "from metadata"}) + "\n", encoding = "utf-8"
|
|
)
|
|
(tmp_path / "a.txt").write_text("edited sidecar", encoding = "utf-8")
|
|
pairs = dict(discover_image_caption_pairs(tmp_path))
|
|
assert pairs[str(tmp_path / "a.png")] == "edited sidecar"
|
|
|
|
|
|
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_config_normalized_accepts_mxfp8_dense_base():
|
|
# mxfp8 is a dense speed mode: a dense base + bf16 compute normalises through.
|
|
cfg = DiffusionLoraConfig(
|
|
base_model = "black-forest-labs/FLUX.1-dev",
|
|
data_dir = "d",
|
|
output_dir = "o",
|
|
base_precision = "mxfp8",
|
|
).normalized()
|
|
assert cfg.base_precision == "mxfp8"
|
|
|
|
|
|
def test_config_normalized_mxfp8_rejects_prequant_base():
|
|
# A prequant (bnb-4bit) base cannot serve the dense mxfp8 base precision.
|
|
with pytest.raises(ValueError, match = "mxfp8"):
|
|
DiffusionLoraConfig(
|
|
base_model = "unsloth/Qwen-Image-2512-unsloth-bnb-4bit",
|
|
data_dir = "d",
|
|
output_dir = "o",
|
|
base_precision = "mxfp8",
|
|
).normalized()
|
|
|
|
|
|
def test_config_normalized_mxfp8_requires_bf16_compute():
|
|
# Like the other dense modes, mxfp8 trains in bf16 compute; fp16 is refused.
|
|
with pytest.raises(ValueError, match = "mxfp8"):
|
|
DiffusionLoraConfig(
|
|
base_model = "black-forest-labs/FLUX.1-dev",
|
|
data_dir = "d",
|
|
output_dir = "o",
|
|
base_precision = "mxfp8",
|
|
mixed_precision = "fp16",
|
|
).normalized()
|
|
|
|
|
|
def test_config_normalized_lists_mxfp8_in_invalid_mode_error():
|
|
# The invalid-base_precision message enumerates the allowed modes, including mxfp8.
|
|
with pytest.raises(ValueError, match = "mxfp8"):
|
|
DiffusionLoraConfig(
|
|
base_model = "b", data_dir = "d", output_dir = "o", base_precision = "bogus"
|
|
).normalized()
|
|
|
|
|
|
def _cfg(**kw):
|
|
return DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw)
|
|
|
|
|
|
def test_resolve_train_steps_uses_train_steps_when_epochs_disabled():
|
|
# num_epochs == 0 leaves the explicit train_steps untouched, whatever the image count.
|
|
cfg = _cfg(train_steps = 300, num_epochs = 0)
|
|
assert resolve_train_steps(cfg, 20) == 300
|
|
assert resolve_train_steps(cfg, 1) == 300
|
|
|
|
|
|
def test_resolve_train_steps_epochs_ceil_over_batch_and_grad_accum():
|
|
# One epoch = ceil(N / (batch x grad_accum)) optimizer steps; num_epochs multiplies it.
|
|
# 10 images, batch 4, grad_accum 1 -> ceil(10/4)=3 steps/epoch.
|
|
assert resolve_train_steps(_cfg(num_epochs = 1, train_batch_size = 4), 10) == 3
|
|
assert resolve_train_steps(_cfg(num_epochs = 5, train_batch_size = 4), 10) == 15
|
|
# grad_accum widens the effective batch: 100 images, batch 2, grad_accum 3 -> per_step=6,
|
|
# ceil(100/6)=17 steps/epoch, 2 epochs -> 34.
|
|
cfg = _cfg(num_epochs = 2, train_batch_size = 2, gradient_accumulation_steps = 3)
|
|
assert resolve_train_steps(cfg, 100) == 34
|
|
# An exact multiple does not round up: 8 images / batch 4 -> 2 steps/epoch.
|
|
assert resolve_train_steps(_cfg(num_epochs = 3, train_batch_size = 4), 8) == 6
|
|
|
|
|
|
def test_resolve_train_steps_single_image_dataset():
|
|
# A one-image dataset is one optimizer step per epoch, so num_epochs == steps.
|
|
assert resolve_train_steps(_cfg(num_epochs = 7, train_batch_size = 4), 1) == 7
|
|
|
|
|
|
def test_resolve_train_steps_caps_at_100000():
|
|
# The run length is capped at 100000 even for absurd epoch counts (matches the request
|
|
# model's train_steps ceiling), so a huge epochs x dataset never overflows the loop.
|
|
cfg = _cfg(num_epochs = 1000, train_batch_size = 1)
|
|
assert resolve_train_steps(cfg, 10_000) == 100000
|
|
|
|
|
|
def test_config_normalized_num_epochs_bounds():
|
|
# 0 (disabled) and the 1..1000 range normalise; out-of-range is rejected.
|
|
assert _cfg(num_epochs = 0).normalized().num_epochs == 0
|
|
assert _cfg(num_epochs = 1000).normalized().num_epochs == 1000
|
|
with pytest.raises(ValueError, match = "num_epochs"):
|
|
_cfg(num_epochs = -1).normalized()
|
|
with pytest.raises(ValueError, match = "num_epochs"):
|
|
_cfg(num_epochs = 1001).normalized()
|
|
|
|
|
|
def test_config_from_dict_threads_num_epochs():
|
|
# num_epochs flows through the shared-payload adapter onto the diffusion field.
|
|
cfg = _config_from_dict(
|
|
{"base_model": "b", "data_dir": "d", "output_dir": "o", "num_epochs": 12}
|
|
)
|
|
assert cfg.num_epochs == 12
|
|
|
|
|
|
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")
|
|
|
|
|
|
def test_config_rejects_zero_lora_alpha():
|
|
# An explicit zero alpha would scale the adapter to nothing; reject it.
|
|
with pytest.raises(ValueError, match = "lora_alpha"):
|
|
DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", lora_alpha = 0).normalized()
|
|
|
|
|
|
def test_config_coerces_string_learning_rate():
|
|
# The Studio config path preserves learning_rate as a string; normalize to float.
|
|
cfg = DiffusionLoraConfig(
|
|
base_model = "b", data_dir = "d", output_dir = "o", learning_rate = "1e-4"
|
|
).normalized()
|
|
assert cfg.learning_rate == 1e-4
|
|
with pytest.raises(ValueError, match = "learning_rate"):
|
|
DiffusionLoraConfig(
|
|
base_model = "b", data_dir = "d", output_dir = "o", learning_rate = "abc"
|
|
).normalized()
|
|
|
|
|
|
def test_config_blank_hf_token_is_anonymous():
|
|
cfg = DiffusionLoraConfig(
|
|
base_model = "b", data_dir = "d", output_dir = "o", hf_token = " "
|
|
).normalized()
|
|
assert cfg.hf_token is None
|
|
|
|
|
|
def test_config_from_dict_aliases_generic_studio_keys():
|
|
# The generic Studio training payload uses different key names; alias them.
|
|
cfg = _config_from_dict(
|
|
{
|
|
"model_name": "b",
|
|
"data_dir": "d",
|
|
"output_dir": "o",
|
|
"max_steps": 25,
|
|
"batch_size": 3,
|
|
"lora_r": 8,
|
|
"lr_scheduler_type": "cosine",
|
|
"random_seed": 7,
|
|
}
|
|
)
|
|
assert cfg.base_model == "b"
|
|
assert cfg.train_steps == 25
|
|
assert cfg.train_batch_size == 3
|
|
assert cfg.lora_rank == 8
|
|
assert cfg.lr_scheduler == "cosine"
|
|
assert cfg.seed == 7
|
|
|
|
|
|
def test_config_from_dict_canonical_key_beats_alias():
|
|
cfg = _config_from_dict(
|
|
{"base_model": "canon", "model_name": "alias", "data_dir": "d", "output_dir": "o"}
|
|
)
|
|
assert cfg.base_model == "canon"
|
|
|
|
|
|
def test_gradient_checkpointing_string_coercion():
|
|
# Studio sends a string; the disable words are False, everything else truthy True.
|
|
for off in ("none", "None", "false", "0", "no", "off", ""):
|
|
assert _coerce_gradient_checkpointing(off) is False
|
|
for on in ("true", "unsloth", "yes"):
|
|
assert _coerce_gradient_checkpointing(on) is True
|
|
assert _coerce_gradient_checkpointing(True) is True
|
|
assert _coerce_gradient_checkpointing(False) is False
|
|
cfg = _config_from_dict(
|
|
{"base_model": "b", "data_dir": "d", "output_dir": "o", "gradient_checkpointing": "none"}
|
|
)
|
|
assert cfg.gradient_checkpointing is False
|
|
|
|
|
|
def test_config_rejects_nonpositive_learning_rate():
|
|
with pytest.raises(ValueError, match = "learning_rate"):
|
|
DiffusionLoraConfig(
|
|
base_model = "b", data_dir = "d", output_dir = "o", learning_rate = 0
|
|
).normalized()
|
|
|
|
|
|
def test_config_rejects_untrainable_base_models():
|
|
# GGUF checkpoints and families without a trainer (Kontext editing, SD3) must fail at
|
|
# normalise time (an instant 400 via the API), not minutes later inside from_pretrained.
|
|
for bad in (
|
|
"unsloth/FLUX.1-dev-GGUF",
|
|
"z-image-turbo-Q4_K_M.gguf",
|
|
"stabilityai/stable-diffusion-3-medium",
|
|
"unsloth/FLUX.1-Kontext-dev",
|
|
):
|
|
with pytest.raises(ValueError):
|
|
DiffusionLoraConfig(base_model = bad, data_dir = "d", output_dir = "o").normalized()
|
|
|
|
|
|
def test_config_resolves_dit_families():
|
|
# FLUX.1 / Qwen-Image / Z-Image bases now resolve to their DiT trainer families.
|
|
for base, fam in (
|
|
("black-forest-labs/FLUX.1-dev", "flux.1"),
|
|
("black-forest-labs/FLUX.1-schnell", "flux.1"),
|
|
("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", "qwen-image"),
|
|
("Tongyi-MAI/Z-Image-Turbo", "z-image"),
|
|
):
|
|
cfg = DiffusionLoraConfig(base_model = base, data_dir = "d", output_dir = "o").normalized()
|
|
assert cfg.resolved_family == fam
|
|
|
|
|
|
def test_config_accepts_sdxl_and_unknown_base_models():
|
|
# SDXL names and unclassifiable custom names/paths must pass the guard (a wrong
|
|
# custom pick still fails cleanly in from_pretrained).
|
|
for ok in (
|
|
"stabilityai/stable-diffusion-xl-base-1.0",
|
|
"stabilityai/sdxl-turbo",
|
|
"/data/checkpoints/my-custom-sdxl",
|
|
"my-finetune",
|
|
):
|
|
cfg = DiffusionLoraConfig(base_model = ok, data_dir = "d", output_dir = "o").normalized()
|
|
assert cfg.base_model == ok
|
|
|
|
|
|
# ── trainer registry + family resolution + metadata sidecar (PR A platform) ──
|
|
def test_get_trainer_resolves_sdxl():
|
|
from core.training.diffusion_lora_trainer import get_trainer, run_diffusion_lora_training
|
|
assert get_trainer("sdxl") is run_diffusion_lora_training
|
|
assert get_trainer("SDXL") is run_diffusion_lora_training # case-insensitive
|
|
|
|
|
|
def test_get_trainer_unknown_family_raises():
|
|
from core.training.diffusion_lora_trainer import get_trainer
|
|
with pytest.raises(ValueError, match = "No trainer"):
|
|
get_trainer("flux.2-dev") # a real family with no registered trainer
|
|
|
|
|
|
def test_get_trainer_resolves_dit_families():
|
|
from core.training.diffusion_dit_trainer import run_dit_lora_training
|
|
from core.training.diffusion_lora_trainer import get_trainer
|
|
for fam in ("flux.1", "qwen-image", "z-image"):
|
|
assert get_trainer(fam) is run_dit_lora_training
|
|
|
|
|
|
def test_normalized_sets_resolved_family():
|
|
cfg = DiffusionLoraConfig(
|
|
base_model = "stabilityai/stable-diffusion-xl-base-1.0", data_dir = "d", output_dir = "o"
|
|
).normalized()
|
|
assert cfg.resolved_family == "sdxl"
|
|
cfg2 = DiffusionLoraConfig(
|
|
base_model = "my-custom-thing", data_dir = "d", output_dir = "o"
|
|
).normalized()
|
|
assert cfg2.resolved_family == "sdxl" # unknown -> default SDXL trainer
|
|
|
|
|
|
def test_explicit_model_family_validated():
|
|
from core.training.diffusion_lora_trainer import DiffusionLoraConfig as C
|
|
|
|
# A bogus explicit family is rejected up front.
|
|
with pytest.raises(ValueError, match = "Unknown model_family"):
|
|
C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "not-a-family").normalized()
|
|
# A known-but-not-trainable family (Kontext editing) is rejected with a helpful hint.
|
|
with pytest.raises(ValueError):
|
|
C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "flux.1-kontext").normalized()
|
|
# A DiT family that IS trainable resolves to itself.
|
|
assert (
|
|
C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "flux.1")
|
|
.normalized()
|
|
.resolved_family
|
|
== "flux.1"
|
|
)
|
|
# SDXL explicit passes.
|
|
assert (
|
|
C(base_model = "b", data_dir = "d", output_dir = "o", model_family = "sdxl")
|
|
.normalized()
|
|
.resolved_family
|
|
== "sdxl"
|
|
)
|
|
|
|
|
|
def test_publish_writes_metadata_sidecar(tmp_path, monkeypatch):
|
|
import json as _json
|
|
from pathlib import Path
|
|
|
|
from core.inference import diffusion_lora
|
|
from core.training.diffusion_lora_trainer import _publish_to_lora_catalog
|
|
|
|
loras = tmp_path / "loras"
|
|
loras.mkdir()
|
|
monkeypatch.setattr(diffusion_lora, "loras_dir", lambda: loras)
|
|
|
|
src = tmp_path / "run" / "pytorch_lora_weights.safetensors"
|
|
src.parent.mkdir(parents = True)
|
|
src.write_bytes(b"fake-adapter")
|
|
cfg = DiffusionLoraConfig(
|
|
base_model = "stabilityai/sdxl-turbo",
|
|
data_dir = "d",
|
|
output_dir = str(tmp_path / "run"),
|
|
adapter_name = "my.style",
|
|
instance_prompt = "a photo in sks style",
|
|
lora_rank = 8,
|
|
).normalized()
|
|
|
|
dest = _publish_to_lora_catalog(str(src), cfg)
|
|
assert dest is not None
|
|
sidecar = Path(dest).with_suffix(".json")
|
|
assert sidecar.is_file()
|
|
meta = _json.loads(sidecar.read_text())
|
|
assert meta["family"] == "sdxl"
|
|
assert meta["families"] == ["sdxl"]
|
|
assert meta["base_model"] == "stabilityai/sdxl-turbo"
|
|
assert meta["lora_rank"] == 8
|
|
assert meta["trigger_prompt"] == "a photo in sks style"
|
|
assert meta["source"] == "studio-trained"
|