116 lines
4.7 KiB
Python
116 lines
4.7 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
|
|
|
|
"""Unit tests for the flow-matching DiT LoRA trainer (FLUX.1 / Qwen-Image / Z-Image).
|
|
|
|
CPU-only: cover family resolution, the per-family spec table, the QLoRA prequant
|
|
heuristic, the bf16-only guard, and the gated-repo name check. The full training loop is
|
|
exercised by the live GPU smokes, not here."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from core.training.diffusion_dit_trainer import (
|
|
_FLUX_TARGETS,
|
|
_GATED_TRAIN_REPOS,
|
|
_QWEN_TARGETS,
|
|
_SPECS,
|
|
_ZIMAGE_TARGETS,
|
|
_assert_gated_access,
|
|
_repo_is_prequantized,
|
|
_select_lora_targets,
|
|
run_dit_lora_training,
|
|
)
|
|
from core.training.diffusion_train_common import (
|
|
DEFAULT_LORA_TARGETS,
|
|
DiffusionLoraConfig,
|
|
family_train_infos,
|
|
)
|
|
|
|
|
|
def test_specs_cover_the_three_dit_families():
|
|
assert set(_SPECS) == {"flux.1", "qwen-image", "z-image"}
|
|
# FLUX / Qwen share the added-kv attention target set; Z-Image is single-stream.
|
|
assert "add_q_proj" in _SPECS["flux.1"].lora_targets
|
|
assert "add_q_proj" in _SPECS["qwen-image"].lora_targets
|
|
assert "add_q_proj" not in _SPECS["z-image"].lora_targets
|
|
# Z-Image and Qwen are bf16-only.
|
|
assert _SPECS["z-image"].force_bf16 is True
|
|
assert _SPECS["qwen-image"].force_bf16 is True
|
|
|
|
|
|
def test_select_lora_targets_uses_family_default_for_generic_config():
|
|
# normalized() fills lora_target_modules with the generic DEFAULT_LORA_TARGETS when a
|
|
# caller doesn't set it, so that value must resolve to the family's targets (which add
|
|
# the DiT-specific projections), not stay stuck on the generic SDXL list.
|
|
assert _select_lora_targets(DEFAULT_LORA_TARGETS, _FLUX_TARGETS) == _FLUX_TARGETS
|
|
assert _select_lora_targets(DEFAULT_LORA_TARGETS, _QWEN_TARGETS) == _QWEN_TARGETS
|
|
assert _select_lora_targets(DEFAULT_LORA_TARGETS, _ZIMAGE_TARGETS) == _ZIMAGE_TARGETS
|
|
|
|
|
|
def test_select_lora_targets_explicit_override_wins():
|
|
# Any OTHER explicit tuple is a deliberate override and must win over the family spec.
|
|
override = ("to_q", "to_k")
|
|
assert _select_lora_targets(override, _FLUX_TARGETS) == override
|
|
# The default request path (config carrying the generic default) reaches the spec.
|
|
cfg = DiffusionLoraConfig(
|
|
base_model = "black-forest-labs/FLUX.1-dev", data_dir = "d", output_dir = "o"
|
|
).normalized()
|
|
assert cfg.lora_target_modules == DEFAULT_LORA_TARGETS
|
|
assert (
|
|
_select_lora_targets(cfg.lora_target_modules, _SPECS["flux.1"].lora_targets)
|
|
== _FLUX_TARGETS
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"repo, expected",
|
|
[
|
|
("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", True),
|
|
("unsloth/Z-Image-Turbo-unsloth-bnb-4bit", True),
|
|
("some/model-int4", True),
|
|
("black-forest-labs/FLUX.1-dev", False),
|
|
("Tongyi-MAI/Z-Image-Turbo", False),
|
|
],
|
|
)
|
|
def test_prequant_heuristic(repo, expected):
|
|
assert _repo_is_prequantized(repo) is expected
|
|
|
|
|
|
def test_zimage_rejects_fp16_before_loading():
|
|
# bf16-only families must refuse an explicit fp16 request up front (no model load).
|
|
cfg = DiffusionLoraConfig(
|
|
base_model = "Tongyi-MAI/Z-Image-Turbo",
|
|
data_dir = "does-not-exist",
|
|
output_dir = "o",
|
|
mixed_precision = "fp16",
|
|
)
|
|
with pytest.raises(ValueError, match = "bf16"):
|
|
run_dit_lora_training(cfg)
|
|
|
|
|
|
def test_gated_access_requires_token():
|
|
assert "black-forest-labs/flux.1-dev" in _GATED_TRAIN_REPOS
|
|
# No token -> clear, actionable error before any download.
|
|
with pytest.raises(ValueError, match = "gated"):
|
|
_assert_gated_access("black-forest-labs/FLUX.1-dev", None)
|
|
with pytest.raises(ValueError, match = "gated"):
|
|
_assert_gated_access("black-forest-labs/FLUX.1-dev", " ")
|
|
# With a token, or for a non-gated repo, it is a no-op.
|
|
_assert_gated_access("black-forest-labs/FLUX.1-dev", "hf_realtoken")
|
|
_assert_gated_access("Tongyi-MAI/Z-Image-Turbo", None)
|
|
|
|
|
|
def test_family_train_infos_lists_dit_families():
|
|
infos = {i["name"]: i for i in family_train_infos()}
|
|
for fam in ("sdxl", "flux.1", "qwen-image", "z-image"):
|
|
assert fam in infos, f"{fam} missing from family_train_infos"
|
|
assert infos[fam]["default_base"]
|
|
assert infos[fam]["base_repos"]
|
|
assert "resolution" in infos[fam]["defaults"]
|
|
# FLUX default base is the gated dev repo; its note flags the license requirement.
|
|
assert infos["flux.1"]["default_base"] == "black-forest-labs/FLUX.1-dev"
|
|
assert "gated" in infos["flux.1"]["vram_note"].lower()
|
|
# Z-Image defaults to the prequant nf4 repo for QLoRA.
|
|
assert "4bit" in infos["z-image"]["default_base"].lower()
|