# 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, ) 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_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_rejects_nonpositive_snr_gamma(): # gamma <= 0 zeroes/inverts the min-SNR weight; None is the documented disable. with pytest.raises(ValueError, match = "snr_gamma"): DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", snr_gamma = 0).normalized() cfg = DiffusionLoraConfig( base_model = "b", data_dir = "d", output_dir = "o", snr_gamma = None ).normalized() assert cfg.snr_gamma is None 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"