mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-03 20:14:03 +00:00
494 lines
18 KiB
Python
494 lines
18 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 sd-cli engine + routing (``sd_cpp_engine.py``).
|
|
|
|
Hermetic: the binary finder is driven against a tmp filesystem, and ``generate``
|
|
runs a fake ``subprocess.Popen`` that emits canned lines and writes the output
|
|
PNG -- no real ``sd-cli``, no GPU.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from core.inference import sd_cpp_engine as eng
|
|
from core.inference.sd_cpp_engine import (
|
|
ENGINE_DIFFUSERS,
|
|
ENGINE_SD_CPP,
|
|
SdCppEngine,
|
|
find_sd_cpp_binary,
|
|
find_sd_server_binary,
|
|
runtime_env,
|
|
select_diffusion_engine,
|
|
)
|
|
from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles, SdCppUpscaleParams
|
|
|
|
|
|
# ── binary discovery ────────────────────────────────────────────────────────
|
|
|
|
|
|
def _clear_env(monkeypatch):
|
|
monkeypatch.delenv("SD_CLI_PATH", raising = False)
|
|
monkeypatch.delenv("UNSLOTH_SD_CPP_PATH", raising = False)
|
|
|
|
|
|
def test_find_prefers_sd_cli_path_env(tmp_path, monkeypatch):
|
|
_clear_env(monkeypatch)
|
|
binary = tmp_path / "sd-cli"
|
|
binary.write_text("#!/bin/sh\n")
|
|
monkeypatch.setenv("SD_CLI_PATH", str(binary))
|
|
# even with PATH empty, the direct env wins
|
|
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
|
|
assert find_sd_cpp_binary() == str(binary)
|
|
|
|
|
|
def test_find_custom_install_dir_build_layout(tmp_path, monkeypatch):
|
|
_clear_env(monkeypatch)
|
|
root = tmp_path / "sdcpp"
|
|
built = root / "build" / "bin" / "sd-cli"
|
|
built.parent.mkdir(parents = True)
|
|
built.write_text("x")
|
|
monkeypatch.setenv("UNSLOTH_SD_CPP_PATH", str(root))
|
|
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
|
|
assert find_sd_cpp_binary() == str(built)
|
|
|
|
|
|
def test_find_falls_back_to_path(tmp_path, monkeypatch):
|
|
_clear_env(monkeypatch)
|
|
monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome"))
|
|
monkeypatch.setattr(
|
|
eng.shutil, "which", lambda stem: "/usr/bin/sd-cli" if stem == "sd-cli" else None
|
|
)
|
|
assert find_sd_cpp_binary() == "/usr/bin/sd-cli"
|
|
|
|
|
|
def test_find_returns_none_when_absent(tmp_path, monkeypatch):
|
|
_clear_env(monkeypatch)
|
|
monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome"))
|
|
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
|
|
assert find_sd_cpp_binary() is None
|
|
|
|
|
|
# ── sd-server discovery ──────────────────────────────────────────────────────
|
|
|
|
|
|
def _clear_server_env(monkeypatch):
|
|
monkeypatch.delenv("SD_SERVER_PATH", raising = False)
|
|
monkeypatch.delenv("SD_CLI_PATH", raising = False)
|
|
monkeypatch.delenv("UNSLOTH_SD_CPP_PATH", raising = False)
|
|
|
|
|
|
def test_find_server_prefers_sd_server_path_env(tmp_path, monkeypatch):
|
|
_clear_server_env(monkeypatch)
|
|
binary = tmp_path / "sd-server"
|
|
binary.write_text("x")
|
|
monkeypatch.setenv("SD_SERVER_PATH", str(binary))
|
|
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
|
|
assert find_sd_server_binary() == str(binary)
|
|
|
|
|
|
def test_find_server_build_layout(tmp_path, monkeypatch):
|
|
_clear_server_env(monkeypatch)
|
|
root = tmp_path / "sdcpp"
|
|
built = root / "build" / "bin" / "sd-server"
|
|
built.parent.mkdir(parents = True)
|
|
built.write_text("x")
|
|
monkeypatch.setenv("UNSLOTH_SD_CPP_PATH", str(root))
|
|
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
|
|
assert find_sd_server_binary() == str(built)
|
|
|
|
|
|
def test_find_server_path_fallback(tmp_path, monkeypatch):
|
|
_clear_server_env(monkeypatch)
|
|
monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome"))
|
|
monkeypatch.setattr(
|
|
eng.shutil, "which", lambda stem: "/usr/bin/sd-server" if stem == "sd-server" else None
|
|
)
|
|
assert find_sd_server_binary() == "/usr/bin/sd-server"
|
|
|
|
|
|
def test_find_server_not_confused_with_sd_cli(tmp_path, monkeypatch):
|
|
# A tree that has only sd-cli must NOT be reported as an sd-server (and vice versa),
|
|
# so the backend correctly falls back to one-shot when only the CLI is present.
|
|
_clear_server_env(monkeypatch)
|
|
root = tmp_path / "sdcpp"
|
|
(root / "build" / "bin").mkdir(parents = True)
|
|
(root / "build" / "bin" / "sd-cli").write_text("x")
|
|
monkeypatch.setenv("UNSLOTH_SD_CPP_PATH", str(root))
|
|
monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome"))
|
|
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
|
|
assert find_sd_server_binary() is None
|
|
assert find_sd_cpp_binary() == str(root / "build" / "bin" / "sd-cli")
|
|
|
|
|
|
# ── availability / version ──────────────────────────────────────────────────
|
|
|
|
|
|
def test_engine_unavailable_when_no_binary(monkeypatch):
|
|
# Force the "no binary anywhere" condition so the test is hermetic even on a host
|
|
# that happens to have sd-cli installed.
|
|
monkeypatch.setattr(eng, "find_sd_cpp_binary", lambda: None)
|
|
e = SdCppEngine(binary = None)
|
|
assert e.is_available() is False
|
|
assert e.version() is None
|
|
|
|
|
|
def test_engine_version_parsed_and_cached(tmp_path, monkeypatch):
|
|
binary = tmp_path / "sd-cli"
|
|
binary.write_text("x")
|
|
e = SdCppEngine(binary = str(binary))
|
|
calls = {"n": 0}
|
|
|
|
def _fake_run(*_a, **_k):
|
|
calls["n"] += 1
|
|
return types.SimpleNamespace(
|
|
stdout = "stable-diffusion.cpp version master-721\n", stderr = "", returncode = 0
|
|
)
|
|
|
|
monkeypatch.setattr(eng.subprocess, "run", _fake_run)
|
|
assert e.version() == "stable-diffusion.cpp version master-721"
|
|
assert e.version() == "stable-diffusion.cpp version master-721"
|
|
assert calls["n"] == 1 # cached after the first probe
|
|
|
|
|
|
# ── runtime env (bundled shared libs) ───────────────────────────────────────
|
|
|
|
|
|
def test_runtime_env_prepends_binary_dir_to_lib_path():
|
|
var = eng._lib_path_var()
|
|
env = runtime_env("/opt/sdcpp/bin/sd-cli", {var: "/existing"})
|
|
first = env[var].split(os.pathsep)[0]
|
|
assert first == "/opt/sdcpp/bin"
|
|
assert "/existing" in env[var]
|
|
|
|
|
|
def test_runtime_env_scrubs_native_path_lease_secret(monkeypatch):
|
|
# The sd-cli child is an external process and must never receive the native-path
|
|
# lease secret; every launch (version + generate/upscale) funnels through runtime_env.
|
|
monkeypatch.setenv("UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET", "top-secret")
|
|
from_os = runtime_env("/opt/sdcpp/bin/sd-cli")
|
|
assert "UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET" not in from_os
|
|
from_base = runtime_env(
|
|
"/opt/sdcpp/bin/sd-cli",
|
|
{"UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET": "top-secret"},
|
|
)
|
|
assert "UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET" not in from_base
|
|
|
|
|
|
def test_runtime_env_handles_missing_lib_path():
|
|
var = eng._lib_path_var()
|
|
env = runtime_env("/opt/sdcpp/bin/sd-cli", {})
|
|
assert env[var] == "/opt/sdcpp/bin"
|
|
|
|
|
|
def test_terminate_reaps_killed_child():
|
|
# Cancellation/timeout paths call _terminate then immediately raise, so it must
|
|
# reap the killed child itself -- otherwise a burst of image cancellations leaves
|
|
# zombies until a later Popen cleanup. After _terminate the returncode is set
|
|
# (the child has been waited on), so nothing lingers.
|
|
import subprocess
|
|
proc = subprocess.Popen(
|
|
[sys.executable, "-c", "import time; time.sleep(30)"],
|
|
start_new_session = (os.name == "posix"),
|
|
)
|
|
try:
|
|
eng._terminate(proc)
|
|
assert proc.returncode is not None
|
|
finally:
|
|
if proc.poll() is None:
|
|
proc.kill()
|
|
proc.wait()
|
|
|
|
|
|
# ── generate (fake subprocess) ──────────────────────────────────────────────
|
|
|
|
|
|
class _FakePopen:
|
|
"""Stand-in for subprocess.Popen: streams ``lines`` then writes ``out_file``
|
|
(unless ``write`` is False) and exits with ``returncode``."""
|
|
|
|
captured_cmd: list[str] = []
|
|
captured_env: dict = {}
|
|
|
|
def __init__(
|
|
self,
|
|
cmd,
|
|
*,
|
|
lines,
|
|
returncode,
|
|
out_file,
|
|
write,
|
|
env = None,
|
|
):
|
|
type(self).captured_cmd = list(cmd)
|
|
type(self).captured_env = dict(env or {})
|
|
self._lines = list(lines)
|
|
self.returncode = returncode
|
|
self._out_file = out_file
|
|
self._write = write
|
|
|
|
@property
|
|
def stdout(self):
|
|
return iter(self._lines)
|
|
|
|
def wait(self, timeout = None):
|
|
if self._write:
|
|
Path(self._out_file).write_bytes(b"\x89PNG\r\n")
|
|
return self.returncode
|
|
|
|
def poll(self):
|
|
return self.returncode
|
|
|
|
def kill(self):
|
|
pass
|
|
|
|
|
|
def _patch_popen(
|
|
monkeypatch,
|
|
*,
|
|
lines,
|
|
returncode,
|
|
out_file,
|
|
write = True,
|
|
):
|
|
def _factory(cmd, **kw):
|
|
return _FakePopen(
|
|
cmd,
|
|
lines = lines,
|
|
returncode = returncode,
|
|
out_file = out_file,
|
|
write = write,
|
|
env = kw.get("env"),
|
|
)
|
|
|
|
monkeypatch.setattr(eng.subprocess, "Popen", _factory)
|
|
|
|
|
|
def _engine(tmp_path):
|
|
binary = tmp_path / "sd-cli"
|
|
binary.write_text("x")
|
|
return SdCppEngine(binary = str(binary))
|
|
|
|
|
|
def test_generate_success_returns_path_and_collects_logs(tmp_path, monkeypatch):
|
|
e = _engine(tmp_path)
|
|
out = tmp_path / "img.png"
|
|
_patch_popen(
|
|
monkeypatch, lines = ["loading model", "step 1/8", "done"], returncode = 0, out_file = out
|
|
)
|
|
seen: list[str] = []
|
|
files = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/q.gguf")
|
|
params = SdCppGenParams(prompt = "a cat", steps = 8, seed = 1)
|
|
|
|
result = e.generate(files, params, output_path = str(out), on_log = seen.append)
|
|
|
|
assert result == out and out.is_file()
|
|
assert seen == ["loading model", "step 1/8", "done"]
|
|
# the real argv was built and handed to Popen
|
|
assert "--diffusion-model" in _FakePopen.captured_cmd
|
|
assert str(out) == _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--output") + 1]
|
|
# the subprocess env carries the binary's dir on the library path
|
|
var = eng._lib_path_var()
|
|
assert str(Path(e.binary).resolve().parent) in _FakePopen.captured_env.get(var, "")
|
|
|
|
|
|
def test_generate_raises_on_nonzero_exit(tmp_path, monkeypatch):
|
|
e = _engine(tmp_path)
|
|
out = tmp_path / "img.png"
|
|
_patch_popen(monkeypatch, lines = ["boom: bad gguf"], returncode = 1, out_file = out, write = False)
|
|
with pytest.raises(RuntimeError, match = "exited 1"):
|
|
e.generate(
|
|
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
|
SdCppGenParams(prompt = "x"),
|
|
output_path = str(out),
|
|
)
|
|
|
|
|
|
def test_generate_raises_when_no_output_despite_success(tmp_path, monkeypatch):
|
|
e = _engine(tmp_path)
|
|
out = tmp_path / "img.png"
|
|
_patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out, write = False)
|
|
with pytest.raises(RuntimeError, match = "no image"):
|
|
e.generate(
|
|
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
|
SdCppGenParams(prompt = "x"),
|
|
output_path = str(out),
|
|
)
|
|
|
|
|
|
def test_generate_does_not_return_stale_preexisting_output(tmp_path, monkeypatch):
|
|
# A leftover file at the target path must not satisfy the post-run output check
|
|
# when the run itself produced nothing: the target is cleared before the run.
|
|
e = _engine(tmp_path)
|
|
out = tmp_path / "img.png"
|
|
out.write_bytes(b"stale")
|
|
_patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out, write = False)
|
|
with pytest.raises(RuntimeError, match = "no image"):
|
|
e.generate(
|
|
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
|
SdCppGenParams(prompt = "x"),
|
|
output_path = str(out),
|
|
)
|
|
assert not out.exists()
|
|
|
|
|
|
def test_generate_raises_when_binary_missing():
|
|
e = SdCppEngine(binary = None)
|
|
with pytest.raises(RuntimeError, match = "not found"):
|
|
e.generate(
|
|
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
|
SdCppGenParams(prompt = "x"),
|
|
output_path = "/tmp/x.png",
|
|
)
|
|
|
|
|
|
class _HangingPopen:
|
|
"""A child that runs but never prints and never exits -- the case a plain
|
|
`for line in stdout` would block on forever, ignoring the timeout."""
|
|
|
|
def __init__(self, cmd, **_kw):
|
|
self._alive = True
|
|
|
|
class _Blocking:
|
|
def __init__(self, owner):
|
|
self.owner = owner
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
while self.owner._alive:
|
|
time.sleep(0.01)
|
|
raise StopIteration
|
|
|
|
@property
|
|
def stdout(self):
|
|
return self._Blocking(self)
|
|
|
|
def poll(self):
|
|
return None if self._alive else -9
|
|
|
|
def wait(self, timeout = None):
|
|
self._alive = False
|
|
return -9
|
|
|
|
def kill(self):
|
|
self._alive = False
|
|
|
|
|
|
def test_generate_times_out_on_silent_hang(tmp_path, monkeypatch):
|
|
e = _engine(tmp_path)
|
|
monkeypatch.setattr(eng.subprocess, "Popen", lambda cmd, **kw: _HangingPopen(cmd, **kw))
|
|
t0 = time.time()
|
|
with pytest.raises(RuntimeError, match = "timed out"):
|
|
e.generate(
|
|
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
|
SdCppGenParams(prompt = "x"),
|
|
output_path = str(tmp_path / "x.png"),
|
|
timeout = 0.3,
|
|
)
|
|
# The timeout is enforced promptly (not blocked until stdout EOF).
|
|
assert time.time() - t0 < 5.0
|
|
|
|
|
|
def test_img2img_generate_passes_init_image(tmp_path, monkeypatch):
|
|
e = _engine(tmp_path)
|
|
out = tmp_path / "img.png"
|
|
src = tmp_path / "src.png"
|
|
src.write_bytes(b"\x89PNG\r\n")
|
|
_patch_popen(monkeypatch, lines = ["img2img"], returncode = 0, out_file = out)
|
|
e.generate(
|
|
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
|
SdCppGenParams(prompt = "x", init_img = str(src), strength = 0.5),
|
|
output_path = str(out),
|
|
)
|
|
assert "--init-img" in _FakePopen.captured_cmd
|
|
assert str(src) == _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--init-img") + 1]
|
|
|
|
|
|
def test_generate_native_speed_dedupes_against_offload(tmp_path, monkeypatch):
|
|
e = _engine(tmp_path)
|
|
out = tmp_path / "img.png"
|
|
_patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out)
|
|
# offload already adds --diffusion-fa; native_speed="default" would add it again.
|
|
e.generate(
|
|
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
|
SdCppGenParams(prompt = "x"),
|
|
output_path = str(out),
|
|
offload = ["--offload-to-cpu", "--diffusion-fa"],
|
|
native_speed = "default",
|
|
)
|
|
# --diffusion-fa appears exactly once (de-duped), not twice.
|
|
assert _FakePopen.captured_cmd.count("--diffusion-fa") == 1
|
|
|
|
|
|
def test_generate_native_speed_adds_flag_when_not_offloaded(tmp_path, monkeypatch):
|
|
e = _engine(tmp_path)
|
|
out = tmp_path / "img.png"
|
|
_patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out)
|
|
e.generate(
|
|
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
|
SdCppGenParams(prompt = "x"),
|
|
output_path = str(out),
|
|
offload = [], # fast/resident tier: no offload, but speed flag still applies
|
|
native_speed = "default",
|
|
)
|
|
assert _FakePopen.captured_cmd.count("--diffusion-fa") == 1
|
|
|
|
|
|
def test_upscale_runs_and_returns_path(tmp_path, monkeypatch):
|
|
e = _engine(tmp_path)
|
|
out = tmp_path / "big.png"
|
|
_patch_popen(monkeypatch, lines = ["upscaling", "done"], returncode = 0, out_file = out)
|
|
result = e.upscale(
|
|
SdCppUpscaleParams(input_image = "/in/small.png", upscale_model = "/m/esrgan.pth", repeats = 2),
|
|
output_path = str(out),
|
|
)
|
|
assert result == out and out.is_file()
|
|
assert _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--mode") + 1] == "upscale"
|
|
assert "--upscale-model" in _FakePopen.captured_cmd
|
|
|
|
|
|
def test_upscale_raises_when_binary_missing(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(eng, "find_sd_cpp_binary", lambda: None)
|
|
e = SdCppEngine(binary = None)
|
|
with pytest.raises(RuntimeError, match = "not found"):
|
|
e.upscale(
|
|
SdCppUpscaleParams(input_image = "/i.png", upscale_model = "/m/e.pth"),
|
|
output_path = str(tmp_path / "x.png"),
|
|
)
|
|
|
|
|
|
# ── engine routing ──────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_routing_gpu_backends_use_diffusers():
|
|
for backend in ("cuda", "rocm", "xpu"):
|
|
assert select_diffusion_engine(backend, native_available = True) == ENGINE_DIFFUSERS
|
|
|
|
|
|
def test_routing_cpu_and_mps_use_native_when_available():
|
|
assert select_diffusion_engine("cpu", native_available = True) == ENGINE_SD_CPP
|
|
assert select_diffusion_engine("mps", native_available = True) == ENGINE_SD_CPP
|
|
|
|
|
|
def test_routing_cpu_falls_back_to_diffusers_without_binary():
|
|
assert select_diffusion_engine("cpu", native_available = False) == ENGINE_DIFFUSERS
|
|
|
|
|
|
def test_routing_prefer_native_overrides_gpu():
|
|
assert (
|
|
select_diffusion_engine("cuda", native_available = True, prefer_native = True) == ENGINE_SD_CPP
|
|
)
|
|
# but only if a binary is actually available
|
|
assert (
|
|
select_diffusion_engine("cuda", native_available = False, prefer_native = True)
|
|
== ENGINE_DIFFUSERS
|
|
)
|