mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-31 02:23:59 +00:00
784 lines
30 KiB
Python
784 lines
30 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
|
|
|
|
"""Tests for the native sd.cpp diffusion backend (the no-GPU engine)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import types
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from core.inference import sd_cpp_backend as bk
|
|
from core.inference.diffusion_families import detect_family
|
|
from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles
|
|
from core.inference.sd_cpp_backend import (
|
|
SdCppDiffusionBackend,
|
|
_map_guidance,
|
|
ensure_sd_cpp_binary,
|
|
)
|
|
from core.inference.sd_cpp_engine import SdCppCancelled
|
|
|
|
|
|
class _FakeEngine:
|
|
"""Stands in for SdCppEngine: writes a 1x1 PNG and records the args."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
fail = None,
|
|
cancel_on_call = False,
|
|
):
|
|
self.calls = []
|
|
self.fail = fail
|
|
self.cancel_on_call = cancel_on_call
|
|
|
|
def is_available(self):
|
|
return True
|
|
|
|
def version(self, **_):
|
|
return "fake sd-cli"
|
|
|
|
def generate(
|
|
self,
|
|
files,
|
|
params,
|
|
*,
|
|
output_path,
|
|
cancel_event = None,
|
|
**kw,
|
|
):
|
|
self.calls.append((files, params, output_path, kw))
|
|
if self.cancel_on_call and cancel_event is not None:
|
|
cancel_event.set()
|
|
if self.fail is not None:
|
|
raise self.fail
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
raise SdCppCancelled("cancelled")
|
|
Image.new("RGB", (1, 1), (10, 20, 30)).save(output_path)
|
|
from pathlib import Path
|
|
|
|
return Path(output_path)
|
|
|
|
|
|
def _loaded_backend(fam_name = "z-image", engine = None):
|
|
b = SdCppDiffusionBackend(engine = engine or _FakeEngine())
|
|
fam = detect_family(fam_name)
|
|
b._state = bk._SdState(
|
|
repo_id = "unsloth/Z-Image-Turbo-GGUF",
|
|
base_repo = fam.base_repo,
|
|
family = fam,
|
|
device = "cpu",
|
|
files = SdCppModelFiles(
|
|
diffusion_model = "/m/z.gguf", vae = "/m/vae.safetensors", llm = "/m/llm.safetensors"
|
|
),
|
|
vae_format = fam.sd_cpp_vae_format,
|
|
sampling_method = fam.sd_cpp_sampling_method,
|
|
flow_shift = fam.sd_cpp_flow_shift,
|
|
mode = "oneshot", # this fixture injects an engine, so it exercises the one-shot path
|
|
)
|
|
return b
|
|
|
|
|
|
class _FakeServer:
|
|
"""Stands in for SdCppServer: records the spawn + one img_gen per whole batch."""
|
|
|
|
def __init__(self, binary):
|
|
self.binary = binary
|
|
self.started = None
|
|
self.stopped = False
|
|
self.payloads = []
|
|
self.timeouts = []
|
|
self.alive = True
|
|
self.lora_dir = None # set by a test to the server's --lora-model-dir scratch dir
|
|
|
|
def is_alive(self):
|
|
return self.alive and not self.stopped
|
|
|
|
def start(
|
|
self,
|
|
files,
|
|
*,
|
|
vae_format = None,
|
|
offload = None,
|
|
native_speed = None,
|
|
threads = None,
|
|
):
|
|
self.started = dict(
|
|
files = files,
|
|
vae_format = vae_format,
|
|
offload = offload,
|
|
native_speed = native_speed,
|
|
threads = threads,
|
|
)
|
|
|
|
def img_gen(
|
|
self,
|
|
payload,
|
|
*,
|
|
on_step = None,
|
|
cancel_event = None,
|
|
total_timeout = None,
|
|
):
|
|
import io as _io
|
|
|
|
self.payloads.append(payload)
|
|
self.timeouts.append(total_timeout)
|
|
if on_step is not None:
|
|
steps = payload.get("sample_params", {}).get("sample_steps", 0)
|
|
on_step(f" {steps}/{steps}")
|
|
n = int(payload.get("batch_count", 1))
|
|
blobs = []
|
|
for i in range(n):
|
|
buf = _io.BytesIO()
|
|
Image.new("RGB", (1, 1), (i, i, i)).save(buf, format = "PNG")
|
|
blobs.append(buf.getvalue())
|
|
return blobs
|
|
|
|
def stop(self):
|
|
self.stopped = True
|
|
|
|
|
|
# ── asset resolution ──────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"fam_name,expect_kinds",
|
|
[
|
|
("flux.1", {"diffusion_model", "vae", "clip_l", "t5xxl"}),
|
|
("z-image", {"diffusion_model", "vae", "llm"}),
|
|
("qwen-image", {"diffusion_model", "vae", "qwen2vl"}),
|
|
("flux.2-klein", {"diffusion_model", "vae", "llm"}),
|
|
],
|
|
)
|
|
def test_asset_specs_cover_required_files(fam_name, expect_kinds):
|
|
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
|
fam = detect_family(fam_name)
|
|
specs = b._asset_specs("unsloth/x-GGUF", "x-Q4_K_M.gguf", fam)
|
|
kinds = {kind for _, _, kind in specs}
|
|
assert kinds == expect_kinds
|
|
# Every spec has a non-empty repo + filename.
|
|
assert all(repo and fn for repo, fn, _ in specs)
|
|
# The transformer reuses the requested GGUF, not a registry file.
|
|
tr = [s for s in specs if s[2] == "diffusion_model"][0]
|
|
assert tr[0] == "unsloth/x-GGUF" and tr[1] == "x-Q4_K_M.gguf"
|
|
|
|
|
|
# ── guidance mapping ──────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_map_guidance_flux_uses_distilled_guidance():
|
|
cfg, g = _map_guidance(detect_family("flux.1"), 3.5)
|
|
assert cfg is None and g == 3.5
|
|
|
|
|
|
def test_map_guidance_cfg_family_off_when_distilled():
|
|
# qwen-image uses real CFG; a distilled 0 -> CFG off (1.0), a >1 value passes through.
|
|
assert _map_guidance(detect_family("qwen-image"), 0.0) == (1.0, None)
|
|
assert _map_guidance(detect_family("qwen-image"), 4.0) == (4.0, None)
|
|
|
|
|
|
# ── status ────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_status_unloaded_reports_sd_cpp_engine():
|
|
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
|
st = b.status()
|
|
assert st["loaded"] is False and st["engine"] == "sd_cpp"
|
|
|
|
|
|
def test_status_loaded_shape():
|
|
b = _loaded_backend()
|
|
st = b.status()
|
|
assert st["loaded"] is True
|
|
assert st["engine"] == "sd_cpp"
|
|
assert st["family"] == "z-image"
|
|
assert st["device"] == "cpu"
|
|
# diffusers-only fields are present (route response parity) but null.
|
|
for k in ("transformer_quant", "attention_backend", "transformer_cache", "text_encoder_quant"):
|
|
assert st[k] is None
|
|
|
|
|
|
# ── generate ──────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_generate_returns_images_and_seed():
|
|
eng = _FakeEngine()
|
|
b = _loaded_backend(engine = eng)
|
|
out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, seed = 123, batch_size = 2)
|
|
assert out["seed"] == 123
|
|
assert out["repo_id"] == "unsloth/Z-Image-Turbo-GGUF"
|
|
assert len(out["images"]) == 2
|
|
assert all(isinstance(im, Image.Image) for im in out["images"])
|
|
# One sd-cli run per batch image, each a distinct seed from the base.
|
|
assert len(eng.calls) == 2
|
|
seeds = [params.seed for _, params, _, _ in eng.calls]
|
|
assert seeds == [123, 124]
|
|
# The per-image seeds are returned so the route can persist each one.
|
|
assert out["seeds"] == [123, 124]
|
|
|
|
|
|
def test_generate_qwen_passes_sampling_args():
|
|
eng = _FakeEngine()
|
|
b = _loaded_backend(fam_name = "qwen-image", engine = eng)
|
|
b.generate(prompt = "x", steps = 20, guidance = 4.0, seed = 1)
|
|
_, params, _, kw = eng.calls[0]
|
|
assert params.sampling_method == "euler" # Qwen's supported sd.cpp sampler
|
|
assert "--flow-shift" in (kw.get("extra_args") or [])
|
|
|
|
|
|
def test_generate_raises_when_not_loaded():
|
|
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
|
with pytest.raises(RuntimeError, match = "No diffusion model is loaded"):
|
|
b.generate(prompt = "x")
|
|
|
|
|
|
def test_generate_passes_vae_format_for_flux2():
|
|
eng = _FakeEngine()
|
|
b = _loaded_backend(fam_name = "flux.2-klein", engine = eng)
|
|
b.generate(prompt = "x", steps = 4, seed = 1)
|
|
_, _, _, kw = eng.calls[0]
|
|
assert kw.get("extra_args") == ["--vae-format", "flux2"]
|
|
|
|
|
|
def test_generate_cancellation_raises_cancelled_not_failure():
|
|
# The engine cancels mid-run; the backend surfaces a cancellation, not a crash.
|
|
eng = _FakeEngine(cancel_on_call = True)
|
|
b = _loaded_backend(engine = eng)
|
|
with pytest.raises(RuntimeError, match = "cancelled"):
|
|
b.generate(prompt = "x", steps = 8, seed = 5)
|
|
|
|
|
|
def test_generate_progress_tracks_parsed_steps():
|
|
b = _loaded_backend()
|
|
b._gen = bk._SdGen(total_steps = 8)
|
|
b._on_log(" sampling 4/8 done")
|
|
p = b.generate_progress()
|
|
assert p["active"] is True and p["step"] == 4 and p["total_steps"] == 8
|
|
# A fraction with a different denominator must not move the bar.
|
|
b._on_log("loaded 1/3 tensors")
|
|
assert b.generate_progress()["step"] == 4
|
|
|
|
|
|
# ── load validation + binary install ──────────────────────────────────────────
|
|
|
|
|
|
def test_begin_load_rejects_unsupported_family(monkeypatch):
|
|
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
|
# A family with no native asset mapping must be rejected (router falls back).
|
|
monkeypatch.setattr(bk, "family_sd_cpp_supported", lambda fam: False)
|
|
with pytest.raises(ValueError, match = "no native sd.cpp asset mapping"):
|
|
b.begin_load("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "z.gguf")
|
|
|
|
|
|
def test_begin_load_requires_gguf_filename():
|
|
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
|
with pytest.raises(ValueError, match = "gguf_filename is required"):
|
|
b.begin_load("unsloth/Z-Image-Turbo-GGUF")
|
|
|
|
|
|
def test_begin_load_resolves_family_from_filename_only(monkeypatch):
|
|
# A local .gguf pick whose family keyword lives only in the basename (parent dir
|
|
# carries none) must resolve via the same filename fallback the route validated
|
|
# with -- not dead-end with "Could not infer" on a native (no-GPU) host.
|
|
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
|
monkeypatch.setattr(b, "_run_load", lambda **kwargs: None) # skip the download thread
|
|
b.begin_load("/models/gguf-store", gguf_filename = "Z-Image-Turbo-Q4_K_M.gguf")
|
|
# Validation passed (no ValueError) and the family was inferred from the filename.
|
|
assert b._loading is not None and b._loading.repo_id == "/models/gguf-store"
|
|
|
|
|
|
def test_ensure_binary_returns_found(monkeypatch):
|
|
monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli")
|
|
assert ensure_sd_cpp_binary() == "/usr/bin/sd-cli"
|
|
|
|
|
|
def test_ensure_binary_install_disabled_returns_none(monkeypatch):
|
|
monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: None)
|
|
assert ensure_sd_cpp_binary(allow_install = False) is None
|
|
|
|
|
|
def test_unload_clears_state_and_signals_cancel():
|
|
cancel = threading.Event()
|
|
b = _loaded_backend()
|
|
b._active_generate_cancel = cancel
|
|
st = b.unload()
|
|
assert st["loaded"] is False
|
|
assert cancel.is_set()
|
|
assert b._cancel_event.is_set()
|
|
|
|
|
|
def test_status_reports_offload_when_flags_active():
|
|
# status must reflect the offload flags actually passed to sd-cli, not always "none",
|
|
# so a balanced/low_vram (or cpu_offload) load is verifiable.
|
|
b = _loaded_backend()
|
|
# No flags (CPU default) -> none.
|
|
assert b.status()["offload_policy"] == "none" and b.status()["cpu_offload"] is False
|
|
# Flags present (off-CPU offload) -> reported active.
|
|
s = b._state
|
|
b._state = bk._SdState(
|
|
repo_id = s.repo_id,
|
|
base_repo = s.base_repo,
|
|
family = s.family,
|
|
device = "cuda",
|
|
files = s.files,
|
|
offload_flags = ("--vae-on-cpu", "--clip-on-cpu"),
|
|
)
|
|
st = b.status()
|
|
assert st["cpu_offload"] is True and st["offload_policy"] == "active"
|
|
|
|
|
|
def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch):
|
|
# A generation that started during the asset download is still running against the OLD
|
|
# model. _run_load must cancel it AND wait on _generate_lock before committing the new
|
|
# state, or a stale sd-cli run finishes afterward and persists an image from the previous
|
|
# model once the new load reports ready.
|
|
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
|
fam = detect_family("z-image")
|
|
monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: [])
|
|
monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None)
|
|
monkeypatch.setattr(
|
|
b,
|
|
"_fetch_assets",
|
|
lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"},
|
|
)
|
|
# Avoid importing torch from the worker thread (its first import deadlocks off the main
|
|
# thread -- a test artifact, not a production path); the device only needs to be CPU here.
|
|
monkeypatch.setattr(
|
|
bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu")
|
|
)
|
|
|
|
b._load_token = 5
|
|
cancel = threading.Event()
|
|
b._active_generate_cancel = cancel # a generation is "in flight"
|
|
|
|
committed = threading.Event()
|
|
|
|
def _load():
|
|
b._run_load(
|
|
repo_id = "unsloth/Z-Image-Turbo-GGUF",
|
|
gguf_filename = "z.gguf",
|
|
base = fam.base_repo,
|
|
fam = fam,
|
|
hf_token = None,
|
|
_load_token = 5,
|
|
)
|
|
committed.set()
|
|
|
|
b._generate_lock.acquire() # simulate the live denoise holding _generate_lock
|
|
try:
|
|
threading.Thread(target = _load, daemon = True).start()
|
|
# The commit must block behind the live generation and not publish the new state,
|
|
# but must already have signalled the in-flight cancel.
|
|
assert not committed.wait(0.5)
|
|
assert b._state is None
|
|
assert cancel.is_set()
|
|
finally:
|
|
b._generate_lock.release()
|
|
assert committed.wait(5) # only now does the commit run
|
|
assert b._state is not None and b._state.repo_id == "unsloth/Z-Image-Turbo-GGUF"
|
|
|
|
|
|
# ── persistent sd-server mode ──────────────────────────────────────────────────
|
|
|
|
|
|
def test_resolve_backend_prefers_server(monkeypatch):
|
|
b = SdCppDiffusionBackend() # no injected engine
|
|
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
|
|
mode, binary, engine = b._resolve_backend()
|
|
assert mode == "server" and binary == "/x/sd-server" and engine is None
|
|
|
|
|
|
def test_resolve_backend_injected_engine_forces_oneshot():
|
|
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
|
mode, binary, engine = b._resolve_backend()
|
|
assert mode == "oneshot" and binary is None and engine is not None
|
|
|
|
|
|
def test_resolve_backend_falls_back_to_oneshot_without_server(monkeypatch):
|
|
b = SdCppDiffusionBackend()
|
|
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: None)
|
|
monkeypatch.setattr(bk, "_install_allowed", lambda: False) # don't attempt a real install
|
|
monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli")
|
|
mode, binary, engine = b._resolve_backend()
|
|
assert mode == "oneshot" and engine is not None
|
|
|
|
|
|
def test_resolve_backend_cached_fallback_engine_does_not_pin_oneshot(monkeypatch):
|
|
# A lazily cached fallback engine (NOT an explicit injection) must not force one-shot:
|
|
# once a server is available again, the next load can use it.
|
|
b = SdCppDiffusionBackend() # no injected engine
|
|
b._engine = _FakeEngine() # simulate a prior lazy one-shot fallback caching the engine
|
|
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
|
|
mode, binary, engine = b._resolve_backend()
|
|
assert mode == "server" and binary == "/x/sd-server" and engine is None
|
|
|
|
|
|
def _run_server_load(
|
|
monkeypatch,
|
|
b,
|
|
servers,
|
|
fam_name = "z-image",
|
|
):
|
|
fam = detect_family(fam_name)
|
|
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
|
|
# The fake binary path is not a real executable; skip the up-front runnability probe.
|
|
monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True)
|
|
|
|
def _factory(binary):
|
|
s = _FakeServer(binary)
|
|
servers.append(s)
|
|
return s
|
|
|
|
monkeypatch.setattr(bk, "SdCppServer", _factory)
|
|
monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: [])
|
|
monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None)
|
|
monkeypatch.setattr(
|
|
b,
|
|
"_fetch_assets",
|
|
lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"},
|
|
)
|
|
monkeypatch.setattr(
|
|
bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu")
|
|
)
|
|
b._load_token = 1
|
|
b._run_load(
|
|
repo_id = "unsloth/Z-Image-Turbo-GGUF",
|
|
gguf_filename = "z.gguf",
|
|
base = fam.base_repo,
|
|
fam = fam,
|
|
hf_token = None,
|
|
_load_token = 1,
|
|
)
|
|
|
|
|
|
def test_server_load_spawns_once_and_status_reports_mode(monkeypatch):
|
|
b = SdCppDiffusionBackend()
|
|
servers: list = []
|
|
_run_server_load(monkeypatch, b, servers)
|
|
assert len(servers) == 1
|
|
assert servers[0].started is not None # the model is loaded once, at spawn
|
|
assert b._state is not None and b._state.mode == "server" and b._state.server is servers[0]
|
|
assert b.status()["native_mode"] == "server"
|
|
|
|
|
|
def test_server_generate_uses_one_request_for_whole_batch(monkeypatch):
|
|
b = SdCppDiffusionBackend()
|
|
servers: list = []
|
|
_run_server_load(monkeypatch, b, servers)
|
|
out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, seed = 7, batch_size = 3)
|
|
assert len(out["images"]) == 3
|
|
assert all(isinstance(im, Image.Image) for im in out["images"])
|
|
# ONE job for the whole batch (no per-image model reload), unlike the one-shot path.
|
|
assert len(servers[0].payloads) == 1
|
|
assert servers[0].payloads[0]["batch_count"] == 3
|
|
assert out["seed"] == 7 and out["seeds"] == [7, 8, 9]
|
|
# step progress was driven from the server's stdout line.
|
|
assert b._gen is None # cleared after generate
|
|
|
|
|
|
def test_server_generate_splits_batches_above_server_limit(monkeypatch):
|
|
# A batch above the server's per-job limit is chunked (the one-shot path did these
|
|
# image-by-image); each chunk gets a timeout proportional to its image count.
|
|
b = SdCppDiffusionBackend()
|
|
servers: list = []
|
|
_run_server_load(monkeypatch, b, servers)
|
|
out = b.generate(prompt = "x", width = 64, height = 64, steps = 4, seed = 100, batch_size = 10)
|
|
assert len(out["images"]) == 10
|
|
counts = [p["batch_count"] for p in servers[0].payloads]
|
|
assert counts == [bk._MAX_SERVER_BATCH, 10 - bk._MAX_SERVER_BATCH] # [8, 2]
|
|
# Each chunk's timeout scales with its image count, not one fixed batch deadline.
|
|
assert servers[0].timeouts == [
|
|
bk._SERVER_PER_IMAGE_TIMEOUT_S * 8,
|
|
bk._SERVER_PER_IMAGE_TIMEOUT_S * 2,
|
|
]
|
|
# Seeds run contiguously across chunks (chunk 2 submitted at base + 8).
|
|
assert out["seeds"] == list(range(100, 110))
|
|
assert servers[0].payloads[1]["seed"] == 108
|
|
|
|
|
|
def test_server_generate_masks_large_seed(monkeypatch):
|
|
# sd.cpp's image seed is signed int64; a larger explicit seed must be masked before it
|
|
# reaches the server (the request model / diffusers accept up to 2**64 - 1).
|
|
b = SdCppDiffusionBackend()
|
|
servers: list = []
|
|
_run_server_load(monkeypatch, b, servers)
|
|
out = b.generate(prompt = "x", width = 64, height = 64, steps = 4, seed = 2**64 - 1, batch_size = 1)
|
|
assert servers[0].payloads[0]["seed"] <= (1 << 63) - 1
|
|
assert all(s <= (1 << 63) - 1 for s in out["seeds"])
|
|
|
|
|
|
def test_status_clears_when_server_died(monkeypatch):
|
|
b = SdCppDiffusionBackend()
|
|
servers: list = []
|
|
_run_server_load(monkeypatch, b, servers)
|
|
assert b.status()["loaded"] is True
|
|
servers[0].alive = False # the resident server crashed / was OOM-killed
|
|
st = b.status()
|
|
assert st["loaded"] is False
|
|
assert b._state is None # stale state was dropped so clients reload
|
|
|
|
|
|
def test_server_generate_progress_from_stdout(monkeypatch):
|
|
b = SdCppDiffusionBackend()
|
|
servers: list = []
|
|
_run_server_load(monkeypatch, b, servers)
|
|
|
|
seen = {}
|
|
|
|
class _WatchServer(_FakeServer):
|
|
def img_gen(
|
|
self,
|
|
payload,
|
|
*,
|
|
on_step = None,
|
|
cancel_event = None,
|
|
total_timeout = None,
|
|
):
|
|
on_step(" 4/8")
|
|
seen["mid"] = b.generate_progress()
|
|
return super().img_gen(
|
|
payload, on_step = on_step, cancel_event = cancel_event, total_timeout = total_timeout
|
|
)
|
|
|
|
b._state = bk._SdState(
|
|
repo_id = b._state.repo_id,
|
|
base_repo = b._state.base_repo,
|
|
family = b._state.family,
|
|
device = b._state.device,
|
|
files = b._state.files,
|
|
vae_format = b._state.vae_format,
|
|
sampling_method = b._state.sampling_method,
|
|
flow_shift = b._state.flow_shift,
|
|
server = _WatchServer("/x/sd-server"),
|
|
mode = "server",
|
|
)
|
|
b.generate(prompt = "x", steps = 8, seed = 1)
|
|
assert seen["mid"]["step"] == 4 and seen["mid"]["total_steps"] == 8
|
|
|
|
|
|
def test_server_unload_stops_server(monkeypatch):
|
|
b = SdCppDiffusionBackend()
|
|
servers: list = []
|
|
_run_server_load(monkeypatch, b, servers)
|
|
st = b.unload()
|
|
assert st["loaded"] is False
|
|
assert servers[0].stopped is True
|
|
assert b._state is None
|
|
|
|
|
|
def test_server_reload_stops_old_server_before_new(monkeypatch):
|
|
b = SdCppDiffusionBackend()
|
|
servers: list = []
|
|
_run_server_load(monkeypatch, b, servers)
|
|
# A second load must tear down the first server and start a fresh one.
|
|
b._load_token = 2
|
|
fam = detect_family("z-image")
|
|
b._run_load(
|
|
repo_id = "unsloth/Z-Image-Turbo-GGUF",
|
|
gguf_filename = "z.gguf",
|
|
base = fam.base_repo,
|
|
fam = fam,
|
|
hf_token = None,
|
|
_load_token = 2,
|
|
)
|
|
assert len(servers) == 2
|
|
assert servers[0].stopped is True # old server stopped
|
|
assert b._state.server is servers[1] and servers[1].stopped is False
|
|
|
|
|
|
def test_server_start_failure_falls_back_to_oneshot(monkeypatch):
|
|
# A present-but-broken sd-server must not fail the load when sd-cli works.
|
|
b = SdCppDiffusionBackend()
|
|
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
|
|
# Probe passes; the failure we exercise here is in start(), not the up-front probe.
|
|
monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True)
|
|
|
|
class _BadServer:
|
|
def __init__(self, binary):
|
|
self.stopped = False
|
|
|
|
def start(self, *a, **k):
|
|
raise RuntimeError("sd-server broken")
|
|
|
|
def stop(self):
|
|
self.stopped = True
|
|
|
|
monkeypatch.setattr(bk, "SdCppServer", _BadServer)
|
|
fake = _FakeEngine()
|
|
monkeypatch.setattr(b, "_resolve_engine", lambda: fake)
|
|
monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: [])
|
|
monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None)
|
|
monkeypatch.setattr(
|
|
b,
|
|
"_fetch_assets",
|
|
lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"},
|
|
)
|
|
monkeypatch.setattr(
|
|
bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu")
|
|
)
|
|
fam = detect_family("z-image")
|
|
b._load_token = 1
|
|
b._run_load(
|
|
repo_id = "unsloth/Z-Image-Turbo-GGUF",
|
|
gguf_filename = "z.gguf",
|
|
base = fam.base_repo,
|
|
fam = fam,
|
|
hf_token = None,
|
|
_load_token = 1,
|
|
)
|
|
assert b._state is not None and b._state.mode == "oneshot" and b._state.server is None
|
|
# and it can still generate via the one-shot engine
|
|
out = b.generate(prompt = "x", steps = 4, seed = 1)
|
|
assert len(out["images"]) == 1 and len(fake.calls) == 1
|
|
|
|
|
|
def test_run_load_redacts_paths_in_progress_error(monkeypatch):
|
|
# A load failure surfaced via load_progress() must run through redact_native_paths, the
|
|
# same scrub the diffusers load path applies, so a registered native path can't leak.
|
|
from utils import native_path_leases as npl
|
|
|
|
secret_root = "/managed/native/root"
|
|
npl._remember_native_path_for_redaction(secret_root, "model dir")
|
|
try:
|
|
b = SdCppDiffusionBackend(engine = _FakeEngine())
|
|
fam = detect_family("z-image")
|
|
monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: [])
|
|
monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None)
|
|
|
|
def _boom(*a, **k):
|
|
raise RuntimeError(f"failed to read {secret_root}/z.gguf")
|
|
|
|
monkeypatch.setattr(b, "_fetch_assets", _boom)
|
|
|
|
b._load_token = 1
|
|
b._loading = bk._SdLoading(repo_id = "unsloth/Z-Image-Turbo-GGUF", base_repo = fam.base_repo)
|
|
b._run_load(
|
|
repo_id = "unsloth/Z-Image-Turbo-GGUF",
|
|
gguf_filename = "z.gguf",
|
|
base = fam.base_repo,
|
|
fam = fam,
|
|
hf_token = None,
|
|
_load_token = 1,
|
|
)
|
|
err = b.load_progress()["error"]
|
|
assert err and secret_root not in err and "<native_path>" in err
|
|
finally:
|
|
with npl._REDACTION_LOCK:
|
|
if secret_root in npl._NATIVE_PATH_REDACTIONS:
|
|
npl._NATIVE_PATH_REDACTIONS.remove(secret_root)
|
|
|
|
|
|
# ── LoRA (native engine) ────────────────────────────────────────────────────────
|
|
|
|
|
|
def _fake_materialize(resolved, dest):
|
|
"""Stand-in for diffusion_lora.materialize_native_dir: write a stub file per adapter
|
|
into ``dest`` and return the resolved list pointing at the written paths (mirroring the
|
|
real helper's contract without touching the Hub / real weights)."""
|
|
from pathlib import Path as _P
|
|
|
|
from core.inference import diffusion_lora as dl
|
|
|
|
dest.mkdir(parents = True, exist_ok = True)
|
|
out = []
|
|
for r in resolved:
|
|
p = _P(dest) / f"{r.alias}.safetensors"
|
|
p.write_bytes(b"stub")
|
|
out.append(dl.ResolvedLora(r.id, r.alias, str(p), r.fmt, r.weight))
|
|
return out
|
|
|
|
|
|
def _patch_lora(
|
|
monkeypatch,
|
|
resolved,
|
|
supported = True,
|
|
):
|
|
from core.inference import diffusion_lora as dl
|
|
|
|
monkeypatch.setattr(dl, "supports_lora", lambda **k: supported)
|
|
monkeypatch.setattr(dl, "resolve_specs", lambda specs, **k: list(resolved))
|
|
monkeypatch.setattr(dl, "materialize_native_dir", _fake_materialize)
|
|
|
|
|
|
def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch):
|
|
# One-shot sd-cli LoRA: adapters materialized into a --lora-model-dir and selected with
|
|
# <lora:ALIAS:w> tags injected into the prompt (the real inject_prompt_tags runs here).
|
|
from core.inference import diffusion_lora as dl
|
|
|
|
eng = _FakeEngine()
|
|
b = _loaded_backend(engine = eng) # mode = "oneshot"
|
|
_patch_lora(
|
|
monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]
|
|
)
|
|
b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)])
|
|
_, params, _, _ = eng.calls[0]
|
|
assert params.lora_dir is not None and params.lora_apply_mode == "auto"
|
|
assert "<lora:myalias:0.8>" in params.prompt
|
|
|
|
|
|
def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tmp_path):
|
|
# Server-mode LoRA rides the structured `lora` request field (the sdcpp API ignores
|
|
# <lora:> prompt tags): adapters staged into the server's --lora-model-dir, referenced
|
|
# by their path relative to it + the validated multiplier.
|
|
from pathlib import Path as _P
|
|
|
|
from core.inference import diffusion_lora as dl
|
|
|
|
b = SdCppDiffusionBackend()
|
|
servers: list = []
|
|
_run_server_load(monkeypatch, b, servers)
|
|
servers[0].lora_dir = str(tmp_path)
|
|
_patch_lora(
|
|
monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]
|
|
)
|
|
b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)])
|
|
payload = servers[0].payloads[0]
|
|
assert "lora" in payload and len(payload["lora"]) == 1
|
|
assert payload["lora"][0]["multiplier"] == 0.7
|
|
assert payload["lora"][0]["path"].endswith("myalias.safetensors")
|
|
assert "<lora:" not in payload["prompt"] # no prompt-tag mechanism on the server
|
|
# The per-request stage subdir under the server's lora dir is removed after the batch.
|
|
assert not list(_P(tmp_path).glob("gen_*"))
|
|
|
|
|
|
def test_generate_rejects_loras_on_unsupported_family(monkeypatch):
|
|
b = _loaded_backend(engine = _FakeEngine())
|
|
_patch_lora(monkeypatch, [], supported = False)
|
|
with pytest.raises(ValueError, match = "LoRA is not supported"):
|
|
b.generate(prompt = "x", steps = 4, seed = 1, loras = [("id1", 1.0)])
|
|
|
|
|
|
def test_generate_zero_weight_loras_are_noop(monkeypatch):
|
|
# weight-0 rows are dropped BEFORE the support gate, so a request carrying only disabled
|
|
# adapters stays a no-op even on a family where native LoRA is unsupported.
|
|
eng = _FakeEngine()
|
|
b = _loaded_backend(engine = eng)
|
|
_patch_lora(monkeypatch, [], supported = False) # would raise if the gate were reached
|
|
b.generate(prompt = "x", steps = 4, seed = 1, loras = [("id1", 0.0)])
|
|
_, params, _, _ = eng.calls[0]
|
|
assert params.lora_dir is None # nothing applied
|
|
|
|
|
|
def test_generate_rejects_controlnet_on_native_engine():
|
|
# ControlNet is diffusers-only. The route passes `controlnet` to whichever engine is
|
|
# active, so the native backend must reject it with a clean ValueError (-> 400) rather
|
|
# than TypeError on an unexpected kwarg (-> opaque 500).
|
|
b = _loaded_backend(engine = _FakeEngine())
|
|
with pytest.raises(ValueError, match = "ControlNet is not yet supported on the native"):
|
|
b.generate(prompt = "x", steps = 4, seed = 1, controlnet = ("id", "img", "canny", 1.0, 0.0, 1.0))
|
|
|
|
|
|
def test_generate_rejects_image_conditioned_on_native_engine():
|
|
# img2img / inpaint / reference / upscale are likewise diffusers-only; a direct API call
|
|
# with an init image on the native engine gets a clean ValueError, not a silent txt2img.
|
|
b = _loaded_backend(engine = _FakeEngine())
|
|
with pytest.raises(ValueError, match = "not yet supported on the native"):
|
|
b.generate(prompt = "x", steps = 4, seed = 1, init_image = "data:image/png;base64,AAAA")
|
|
|
|
|
|
def test_status_native_reports_supports_controlnet_false():
|
|
b = _loaded_backend()
|
|
assert b.status()["supports_controlnet"] is False
|