unsloth/studio/backend/tests/test_diffusion_lora.py

"""Tests for diffusion LoRA support: the shared helpers, request-model validation, the
native prompt-tag/dir wiring, and the diffusers set_adapters manager."""

from __future__ import annotations

import os
import types
from pathlib import Path

import pytest

from core.inference import diffusion_lora as dl


# ── Pure helpers ────────────────────────────────────────────────────────────


def test_sanitize_alias_strips_path_ext_and_unsafe_chars():
    assert dl.sanitize_alias("My Cool/LoRA v2.safetensors") == "LoRA_v2"
    assert dl.sanitize_alias("owner/repo-name") == "repo-name"
    assert dl.sanitize_alias("weird:<>chars.gguf") == "weird_chars"
    assert dl.sanitize_alias("") == "lora"
    # Internal dots (version tags like "V1.0") must be replaced: the alias becomes a
    # diffusers PEFT adapter name and PEFT rejects "." in module/adapter names.
    assert (
        dl.sanitize_alias("Qwen-Image-2512-Lightning-8steps-V1.0-bf16")
        == "Qwen-Image-2512-Lightning-8steps-V1_0-bf16"
    )
    assert "." not in dl.sanitize_alias("model.v1.0.safetensors")


def test_inject_prompt_tags_appends_with_spacing():
    r = dl.ResolvedLora("id", "style", "/p.safetensors", "safetensors", 0.8)
    assert dl.inject_prompt_tags("a cat", [r]) == "a cat <lora:style:0.8>"
    # weight formatting: 1.0 -> "1", trailing zeros trimmed
    r1 = dl.ResolvedLora("id", "s", "/p", "safetensors", 1.0)
    assert dl.inject_prompt_tags("x", [r1]) == "x <lora:s:1>"


def test_inject_prompt_tags_validated_weight_overrides_user_typed():
    r = dl.ResolvedLora("id", "style", "/p", "safetensors", 0.8)
    # A user-typed tag for a SELECTED adapter is replaced by the backend-validated weight
    # (so the recorded/validated 0-2 weight wins over whatever was typed), not duplicated.
    assert dl.inject_prompt_tags("a cat <lora:style:1>", [r]) == "a cat <lora:style:0.8>"


def test_inject_prompt_tags_strips_unselected_user_tags():
    r = dl.ResolvedLora("id", "style", "/p", "safetensors", 0.8)
    # A user tag for an alias that is NOT selected is stripped: only selected adapters are
    # materialized in the managed --lora-model-dir, so sd-cli would drop the dead tag anyway;
    # removing it keeps the prompt clean and unambiguous.
    out = dl.inject_prompt_tags("a cat <lora:other:0.5>", [r])
    assert "<lora:other:0.5>" not in out
    assert out == "a cat <lora:style:0.8>"


def test_inject_prompt_tags_empty_returns_prompt():
    assert dl.inject_prompt_tags("hello", []) == "hello"


def test_supports_lora_matrix():
    # native: flux/z-image yes, qwen no
    assert dl.supports_lora(
        engine = "sd_cpp", family = "flux.1", model_kind = "gguf", transformer_quant = None
    )
    assert dl.supports_lora(
        engine = "sd_cpp", family = "z-image", model_kind = "gguf", transformer_quant = None
    )
    assert not dl.supports_lora(
        engine = "sd_cpp", family = "qwen-image", model_kind = "gguf", transformer_quant = None
    )
    # diffusers: bf16 yes, fp8/int8 dense no, gguf-diffusers no
    assert dl.supports_lora(
        engine = "diffusers", family = "flux.1", model_kind = "pipeline", transformer_quant = None
    )
    assert dl.supports_lora(
        engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = None
    )
    assert not dl.supports_lora(
        engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "fp8"
    )
    assert not dl.supports_lora(
        engine = "diffusers", family = "flux.1", model_kind = "single_file", transformer_quant = "int8"
    )
    assert not dl.supports_lora(
        engine = "diffusers", family = "flux.1", model_kind = "gguf", transformer_quant = None
    )
    # A torch.compile'd diffusers transformer (Speed=default/max) can't take a non-hotswap
    # adapter: diffusers needs the adapter loaded before compilation.
    assert not dl.supports_lora(
        engine = "diffusers",
        family = "flux.1",
        model_kind = "pipeline",
        transformer_quant = None,
        compiled = True,
    )
    # compiled is diffusers-only; the native path ignores it.
    assert dl.supports_lora(
        engine = "sd_cpp",
        family = "flux.1",
        model_kind = "gguf",
        transformer_quant = None,
        compiled = True,
    )


def test_resolve_specs_maps_cancelled_to_diffusion_sentinel(tmp_path, monkeypatch):
    # A Hub download cancelled mid-flight raises RuntimeError("Cancelled"); resolve_specs
    # must convert it to the diffusion cancellation sentinel so the route maps it to 409,
    # not a generic 500 server-error toast.
    def _boom(spec_id, weight, **kw):
        raise RuntimeError("Cancelled")

    monkeypatch.setattr(dl, "resolve_one", _boom)
    with pytest.raises(RuntimeError) as ei:
        dl.resolve_specs([("a", 1.0)])
    assert str(ei.value) == dl.DIFFUSION_CANCELLED_MSG

    # A non-cancellation RuntimeError is left untouched.
    def _other(spec_id, weight, **kw):
        raise RuntimeError("disk full")

    monkeypatch.setattr(dl, "resolve_one", _other)
    with pytest.raises(RuntimeError) as ei2:
        dl.resolve_specs([("a", 1.0)])
    assert str(ei2.value) == "disk full"


def test_materialize_native_dir_symlinks_and_breaks_collisions(tmp_path):
    a = tmp_path / "a.safetensors"
    a.write_bytes(b"x")
    b = tmp_path / "sub"
    b.mkdir()
    b2 = b / "a.safetensors"  # same stem as `a` -> alias collision
    b2.write_bytes(b"y")
    resolved = [
        dl.ResolvedLora("a", "a", str(a), "safetensors", 1.0),
        dl.ResolvedLora("a2", "a", str(b2), "safetensors", 0.5),
    ]
    dest = tmp_path / "managed"
    out = dl.materialize_native_dir(resolved, dest)
    aliases = [r.alias for r in out]
    assert aliases == ["a", "a_2"]  # collision broken
    for r in out:
        assert os.path.exists(r.path)
        assert Path(r.path).parent == dest


def test_list_loras_scans_local(tmp_path, monkeypatch):
    d = tmp_path / "loras"
    d.mkdir()
    (d / "mystyle.safetensors").write_bytes(b"x")
    (d / "other.gguf").write_bytes(b"y")
    (d / "ignore.txt").write_bytes(b"z")
    monkeypatch.setattr(dl, "loras_dir", lambda: d)
    ids = {e.id for e in dl.list_loras()}
    assert ids == {"mystyle", "other"}
    fmts = {e.id: e.fmt for e in dl.list_loras()}
    assert fmts["other"] == "gguf" and fmts["mystyle"] == "safetensors"


def test_resolve_one_local_and_unknown(tmp_path, monkeypatch):
    d = tmp_path / "loras"
    d.mkdir()
    (d / "mystyle.safetensors").write_bytes(b"x")
    monkeypatch.setattr(dl, "loras_dir", lambda: d)
    r = dl.resolve_one("mystyle", 0.7)
    assert r.path.endswith("mystyle.safetensors") and r.weight == 0.7
    with pytest.raises(FileNotFoundError):
        dl.resolve_one("does-not-exist", 1.0)


def test_resolve_specs_drops_zero_weight(tmp_path, monkeypatch):
    d = tmp_path / "loras"
    d.mkdir()
    (d / "a.safetensors").write_bytes(b"x")
    monkeypatch.setattr(dl, "loras_dir", lambda: d)
    out = dl.resolve_specs([("a", 0.0), ("a", 1.0)])
    assert len(out) == 1 and out[0].weight == 1.0


def test_resolve_specs_maps_unknown_id_to_valueerror(tmp_path, monkeypatch):
    # An unknown / stale id raises FileNotFoundError in resolve_one; resolve_specs must
    # surface it as ValueError so the route returns 400, not a generic 500.
    d = tmp_path / "loras"
    d.mkdir()
    monkeypatch.setattr(dl, "loras_dir", lambda: d)
    with pytest.raises(ValueError):
        dl.resolve_specs([("nope", 1.0)])


def test_scan_local_disambiguates_identical_stems(tmp_path, monkeypatch):
    # foo.safetensors and foo.gguf must get distinct ids so each is addressable; a
    # unique stem keeps its clean stem id.
    d = tmp_path / "loras"
    d.mkdir()
    (d / "foo.safetensors").write_bytes(b"x")
    (d / "foo.gguf").write_bytes(b"y")
    (d / "solo.safetensors").write_bytes(b"z")
    monkeypatch.setattr(dl, "loras_dir", lambda: d)
    by_id = {e.id: e for e in dl.list_loras()}
    assert "foo.safetensors" in by_id and "foo.gguf" in by_id
    assert by_id["foo.safetensors"].fmt == "safetensors"
    assert by_id["foo.gguf"].fmt == "gguf"
    assert "solo" in by_id  # unique stem is untouched


def test_resolve_one_rejects_traversal_weight_name(tmp_path, monkeypatch):
    # A client-supplied weight file with traversal / absolute path is rejected before it
    # can reach the downloader (it must stay a plain filename inside the repo).
    monkeypatch.setattr(dl, "loras_dir", lambda: tmp_path)
    for bad in ("owner/name:../secret.safetensors", "owner/name:/etc/x.safetensors"):
        with pytest.raises(ValueError):
            dl.resolve_one(bad, 1.0)


# ── Request-model validation ────────────────────────────────────────────────


def test_lora_spec_and_request_validation():
    from models.inference import DiffusionGenerateRequest, LoraSpec

    # empty / missing loras -> unchanged behaviour
    assert DiffusionGenerateRequest(prompt = "x").loras is None
    req = DiffusionGenerateRequest(
        prompt = "x", loras = [{"id": "a", "weight": 0.5}, {"id": "b", "weight": 1.0}]
    )
    assert [l.id for l in req.loras] == ["a", "b"]
    # weight bounds enforced
    with pytest.raises(Exception):
        LoraSpec(id = "a", weight = 3.0)
    with pytest.raises(Exception):
        LoraSpec(id = "a", weight = -0.1)
    # default weight
    assert LoraSpec(id = "a").weight == 1.0
    # duplicate ids are rejected: repeating an id would load the same adapter as several
    # distinct suffixed adapters and stack its effect past the per-adapter weight bound.
    with pytest.raises(Exception):
        DiffusionGenerateRequest(
            prompt = "x", loras = [{"id": "a", "weight": 0.5}, {"id": "a", "weight": 1.0}]
        )


# ── Diffusers apply manager ─────────────────────────────────────────────────


class _FakePipe:
    def __init__(self):
        self.loaded: list[tuple[str, str]] = []
        self.active = None
        self.unloaded = 0

    def load_lora_weights(
        self,
        path,
        adapter_name = None,
    ):
        self.loaded.append((path, adapter_name))

    def set_adapters(
        self,
        names,
        adapter_weights = None,
    ):
        self.active = (list(names), list(adapter_weights) if adapter_weights else None)

    def unload_lora_weights(self):
        self.unloaded += 1
        self.loaded = []
        self.active = None


def _fake_state(
    pipe,
    *,
    kind = "pipeline",
    quant = None,
):
    fam = types.SimpleNamespace(name = "flux.1")
    return types.SimpleNamespace(
        pipe = pipe, family = fam, kind = kind, transformer_quant = quant, hf_token = None
    )


def _backend():
    from core.inference.diffusion import DiffusionBackend
    return DiffusionBackend()


def test_diffusers_apply_loads_and_sets_adapters(monkeypatch):
    import threading

    monkeypatch.setattr(
        dl,
        "resolve_specs",
        lambda specs, **_: [
            dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.safetensors", "safetensors", w)
            for i, w in specs
        ],
    )
    pipe = _FakePipe()
    _backend()._apply_loras(
        _fake_state(pipe), [("styleA", 0.8), ("styleB", 1.0)], threading.Event()
    )
    assert [n for _p, n in pipe.loaded] == ["styleA", "styleB"]
    assert pipe.active == (["styleA", "styleB"], [0.8, 1.0])
    assert getattr(pipe, "_unsloth_loras")  # marker recorded


def test_diffusers_apply_noop_when_unchanged(monkeypatch):
    import threading

    monkeypatch.setattr(
        dl,
        "resolve_specs",
        lambda specs, **_: [
            dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.safetensors", "safetensors", w)
            for i, w in specs
        ],
    )
    pipe = _FakePipe()
    b = _backend()
    b._apply_loras(_fake_state(pipe), [("styleA", 0.8)], threading.Event())
    first_loaded = list(pipe.loaded)
    b._apply_loras(_fake_state(pipe), [("styleA", 0.8)], threading.Event())
    assert pipe.loaded == first_loaded  # not reloaded
    assert pipe.unloaded == 0


def test_diffusers_apply_clears_when_empty(monkeypatch):
    import threading

    monkeypatch.setattr(
        dl,
        "resolve_specs",
        lambda specs, **_: [
            dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.safetensors", "safetensors", w)
            for i, w in specs
        ],
    )
    pipe = _FakePipe()
    b = _backend()
    b._apply_loras(_fake_state(pipe), [("styleA", 0.8)], threading.Event())
    b._apply_loras(_fake_state(pipe), [], threading.Event())
    assert pipe.unloaded == 1
    assert pipe._unsloth_loras == ()


def test_diffusers_apply_rejects_unsupported_quant():
    import threading
    pipe = _FakePipe()
    with pytest.raises(ValueError, match = "not supported"):
        _backend()._apply_loras(
            _fake_state(pipe, kind = "single_file", quant = "fp8"),
            [("styleA", 1.0)],
            threading.Event(),
        )


def test_diffusers_apply_rejects_gguf_adapter(monkeypatch):
    # A .gguf adapter (discoverable in the shared catalog) cannot load on the diffusers
    # engine; it must be rejected as a clean 400 before touching the pipe.
    import threading

    monkeypatch.setattr(
        dl,
        "resolve_specs",
        lambda specs, **_: [
            dl.ResolvedLora(i, dl.sanitize_alias(i), f"/{i}.gguf", "gguf", w) for i, w in specs
        ],
    )
    pipe = _FakePipe()
    with pytest.raises(ValueError, match = "GGUF LoRA"):
        _backend()._apply_loras(_fake_state(pipe), [("styleA", 1.0)], threading.Event())
    assert pipe.loaded == []  # never touched the pipe


def test_scan_local_reads_family_sidecar(tmp_path, monkeypatch):
    import json

    d = tmp_path / "loras"
    d.mkdir()
    (d / "trained.safetensors").write_bytes(b"x")
    (d / "trained.json").write_text(
        json.dumps({"family": "sdxl", "base_model": "b", "weight_default": 0.8})
    )
    (d / "plain.safetensors").write_bytes(b"y")  # no sidecar -> unknown family
    monkeypatch.setattr(dl, "loras_dir", lambda: d)

    by_id = {e.id: e for e in dl.list_loras()}
    assert by_id["trained"].families == ("sdxl",)
    assert by_id["trained"].weight_default == 0.8
    assert by_id["plain"].families == ()
    assert by_id["plain"].weight_default == 1.0

    # Family filter: the sdxl-tagged adapter is kept for sdxl and hidden for flux.1;
    # the untagged one is always shown (unknown compatibility).
    sdxl_ids = {e.id for e in dl.list_loras(family = "sdxl")}
    flux_ids = {e.id for e in dl.list_loras(family = "flux.1")}
    assert "trained" in sdxl_ids and "plain" in sdxl_ids
    assert "trained" not in flux_ids and "plain" in flux_ids


def test_scan_local_tolerates_bad_sidecar(tmp_path, monkeypatch):
    d = tmp_path / "loras"
    d.mkdir()
    (d / "a.safetensors").write_bytes(b"x")
    (d / "a.json").write_text("{ not valid json")
    monkeypatch.setattr(dl, "loras_dir", lambda: d)
    entry = next(e for e in dl.list_loras() if e.id == "a")
    assert entry.families == () and entry.weight_default == 1.0