ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler

Browse source 
Address review feedback on #6724:
- engine router: unload the engine being deactivated on a switch, so the old
  model is not left resident-but-unreachable (the evictor only targets the active
  engine).
- generate route: sd.cpp execution errors (nonzero exit / timeout / missing
  output) now map to 500, not 409 (which only means not-loaded / cancelled).
- native batch: return per-image seeds and persist the actual seed for each image
  so every batch image is reproducible.
- Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the
  stable-diffusion.cpp docs; other families keep sd-cli defaults.
- honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload
  offload flags on the native load instead of hardcoding them off.
- fail the load when the sd-cli binary is present but not runnable (version()
  now returns None on exec error / nonzero exit).
- size estimate: only treat the transformer asset as a possible local path.
This commit is contained in:
Daniel Han 2026-06-28 05:43:45 +00:00
commit 9a2cc341c5
7 changed files with 122 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

8
studio/backend/core/inference/diffusion_engine_router.py
View file

@ -69,6 +69,14 @@ def active_engine_name() -> str:
def _activate(name: str, reason: Optional[str]) -> Any:
global _active_engine_name, _fallback_reason
with _lock:
# Switching engines: unload the one being deactivated first, or its model
# stays resident but unreachable (the arbiter evictor only targets the active
# engine), leaking 10+ GB and defeating the chat<->diffusion handoff.
if name != _active_engine_name:
try:
get_active_diffusion_engine().unload()
except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch
logger.warning("failed to unload previous engine %s: %s", _active_engine_name, exc)
_active_engine_name = name
_fallback_reason = reason if name == ENGINE_DIFFUSERS else None
if name == ENGINE_SD_CPP:

9
studio/backend/core/inference/diffusion_families.py
View file

@ -58,6 +58,12 @@ class DiffusionFamily:
# autoencoder, None (auto) otherwise.
sd_cpp_vae_format: Optional[str] = None
sd_cpp_text_encoders: tuple[tuple[str, str, str], ...] = field(default_factory = tuple)
# Family-specific sd-cli sampler settings, applied on the native path so the
# output matches the model's supported invocation (e.g. Qwen-Image needs
# --sampling-method euler --flow-shift 3 per stable-diffusion.cpp's docs). None
# leaves sd-cli's defaults (correct for the distilled flux/z-image families).
sd_cpp_sampling_method: Optional[str] = None
sd_cpp_flow_shift: Optional[float] = None
# Keyed by architecture, not per model variant: a checkpoint's specific base repo
@ -114,6 +120,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
"qwen2vl",
),
),
# Qwen-Image's supported sd.cpp invocation (docs/qwen_image.md).
sd_cpp_sampling_method = "euler",
sd_cpp_flow_shift = 3.0,
),
DiffusionFamily(
name = "z-image",

74
studio/backend/core/inference/sd_cpp_backend.py
View file

@ -41,7 +41,13 @@ from core.inference.diffusion_families import (
resolve_base_repo,
resolve_local_gguf_child,
)
from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles
from core.inference.diffusion_memory import (
OFFLOAD_GROUP,
OFFLOAD_MODEL,
OFFLOAD_NONE,
OFFLOAD_SEQUENTIAL,
)
from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles, offload_flags
from core.inference.sd_cpp_engine import (
SdCppCancelled,
SdCppEngine,
@ -110,6 +116,26 @@ class _SdState:
native_speed: str = "off"
offload_flags: tuple[str, ...] = ()
threads: Optional[int] = None
sampling_method: Optional[str] = None
flow_shift: Optional[float] = None
def _memory_policy(memory_mode: Optional[str], cpu_offload: bool) -> str:
"""Map the diffusers memory knobs onto an sd-cli offload policy. Only meaningful
off-CPU (forced sd_cpp / MPS); on CPU everything is resident in RAM anyway."""
mode = (memory_mode or "").strip().lower()
if mode == "low_vram":
return OFFLOAD_SEQUENTIAL
if mode == "balanced":
return OFFLOAD_GROUP
if cpu_offload and mode in ("", "auto"):
return OFFLOAD_MODEL
return OFFLOAD_NONE
def _native_speed_for(speed_mode: Optional[str]) -> str:
mode = (speed_mode or "off").strip().lower()
return mode if mode in ("default", "max") else "off"
@dataclass
@ -235,6 +261,9 @@ class SdCppDiffusionBackend:
base = base,
fam = fam,
hf_token = hf_token,
cpu_offload = cpu_offload,
memory_mode = memory_mode,
speed_mode = speed_mode,
_load_token = token,
),
daemon = True,
@ -249,6 +278,9 @@ class SdCppDiffusionBackend:
base: str,
fam: DiffusionFamily,
hf_token: Optional[str],
cpu_offload: bool = False,
memory_mode: Optional[str] = None,
speed_mode: Optional[str] = None,
_load_token: int,
) -> None:
try:
@ -271,6 +303,12 @@ class SdCppDiffusionBackend:
qwen2vl = paths.get("qwen2vl"),
)
device = resolve_diffusion_device_target().device
# Honor the requested speed everywhere; offload only off-CPU (forced
# sd_cpp / MPS), since on CPU sd-cli is resident in RAM and the offload
# flags are no-ops.
offload: tuple[str, ...] = ()
if device != "cpu":
offload = tuple(offload_flags(_memory_policy(memory_mode, cpu_offload)))
state = _SdState(
repo_id = repo_id,
base_repo = base,
@ -278,12 +316,17 @@ class SdCppDiffusionBackend:
device = device,
files = files,
vae_format = fam.sd_cpp_vae_format,
native_speed = "off",
offload_flags = (),
native_speed = _native_speed_for(speed_mode),
offload_flags = offload,
threads = None,
sampling_method = fam.sd_cpp_sampling_method,
flow_shift = fam.sd_cpp_flow_shift,
)
# Probe the binary version (cheap) so a broken install fails the load.
engine.version()
# Probe the binary: version() returns None when the present binary cannot
# run (bad permissions / missing shared libs), so fail the load now rather
# than commit a "ready" state that crashes on the first generation.
if engine.version() is None:
raise RuntimeError("sd-cli binary is present but not runnable.")
with self._lock:
if self._load_token != _load_token:
return # superseded / cancelled
@ -319,8 +362,10 @@ class SdCppDiffusionBackend:
try:
from huggingface_hub import HfApi
api = HfApi(token = hf_token)
for repo, fn, _ in assets:
if Path(repo).expanduser().exists():
for repo, fn, kind in assets:
# Only the transformer can be a local path; for the others ``repo`` is
# an HF id (a same-named local dir must not skip the size estimate).
if kind == "diffusion_model" and Path(repo).expanduser().exists():
continue
try:
info = api.get_paths_info(repo, paths = [fn], expand = False)
@ -405,10 +450,15 @@ class SdCppDiffusionBackend:
else:
seed = int(seed)
cfg_scale, flux_guidance = _map_guidance(state.family, guidance)
vae_extra = ["--vae-format", state.vae_format] if state.vae_format else None
extra_args: list[str] = []
if state.vae_format:
extra_args += ["--vae-format", state.vae_format]
if state.flow_shift is not None:
extra_args += ["--flow-shift", repr(float(state.flow_shift))]
self._gen = _SdGen(total_steps = int(steps))
images = []
seeds: list[int] = []
with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir:
for index in range(max(1, int(batch_size))):
if cancel.is_set():
@ -426,6 +476,7 @@ class SdCppDiffusionBackend:
cfg_scale = cfg_scale,
guidance = flux_guidance,
seed = seed_i,
sampling_method = state.sampling_method,
batch_count = 1,
)
engine.generate(
@ -435,15 +486,18 @@ class SdCppDiffusionBackend:
offload = list(state.offload_flags) or None,
native_speed = state.native_speed,
threads = state.threads,
extra_args = vae_extra,
extra_args = extra_args or None,
on_log = self._on_log,
cancel_event = cancel,
)
with Image.open(out_path) as im:
images.append(im.copy())
seeds.append(seed_i)
if cancel.is_set():
raise RuntimeError("Diffusion generation was cancelled.")
return {"images": images, "seed": int(seed), "repo_id": state.repo_id}
# ``seeds`` is the per-image seed (each sd-cli run used seed+index), so
# the route can persist the real seed for every image in the batch.
return {"images": images, "seed": int(seed), "seeds": seeds, "repo_id": state.repo_id}
except SdCppCancelled as exc:
raise RuntimeError("Diffusion generation was cancelled.") from exc
finally:

13
studio/backend/core/inference/sd_cpp_engine.py
View file

@ -189,7 +189,10 @@ class SdCppEngine:
return bool(self.binary) and Path(self.binary).is_file()
def version(self, *, timeout: float = 10.0) -> Optional[str]:
"""First line of ``sd-cli --version``, cached. None if it can't run."""
"""First line of ``sd-cli --version``, cached on success. ``None`` when the
binary is absent OR present-but-unrunnable (exec error / nonzero exit, e.g.
missing shared libraries / bad permissions), so callers can fail a load early
instead of committing a "ready" state that crashes on first generation."""
if not self.is_available():
return None
if self._version is not None:
@ -204,10 +207,12 @@ class SdCppEngine:
check = False,
env = runtime_env(self.binary),
)
text = ((res.stdout or "") + "\n" + (res.stderr or "")).strip()
self._version = text.splitlines()[0] if text else ""
except (OSError, subprocess.SubprocessError):
self._version = ""
return None
if res.returncode != 0:
return None
text = ((res.stdout or "") + "\n" + (res.stderr or "")).strip()
self._version = text.splitlines()[0] if text else ""
return self._version
def generate(

24
studio/backend/routes/inference.py
View file

@ -10127,20 +10127,28 @@ async def generate_diffusion_image(
batch_size = request.batch_size,
)
except RuntimeError as exc:
# No model loaded (or unloaded mid-flight) — a client-state problem.
raise HTTPException(status_code = 409, detail = str(exc))
# Only "no model loaded" / cancelled are client-state (409). The native
# sd.cpp engine also raises RuntimeError for execution failures (nonzero
# exit, timeout, missing output), which are server errors (500).
msg = str(exc)
if "No diffusion model is loaded" in msg or "cancelled" in msg.lower():
raise HTTPException(status_code = 409, detail = msg)
logger.error("diffusion.generate_failed: %s", exc)
raise HTTPException(status_code = 500, detail = "Image generation failed.")
except Exception as exc:
logger.error("diffusion.generate_failed: %s", exc)
raise HTTPException(status_code = 500, detail = "Image generation failed.")
# Persist each image with its full recipe embedded. The whole batch shares
# one seed (drawn sequentially from one generator) and one timestamp (the
# images are generated together), so their gallery order is stable on reload.
# Persist each image with its full recipe embedded. The diffusers batch shares
# one seed (drawn sequentially from one generator); the native sd.cpp batch uses a
# distinct seed per image and returns them in ``seeds`` so each is reproducible.
created_at = time.time()
per_image_seeds = result.get("seeds")
def _persist() -> list[dict]:
records = []
for index, image in enumerate(result["images"]):
seed = per_image_seeds[index] if per_image_seeds and index < len(per_image_seeds) else result["seed"]
records.append(
image_gallery.save(
image,
@ -10151,9 +10159,9 @@ async def generate_diffusion_image(
"height": request.height,
"steps": request.steps,
"guidance": request.guidance,
"seed": result["seed"],
# Position within the batch: images here share a seed + timestamp,
# so the export filename needs this to stay unique.
"seed": seed,
# Position within the batch: shared timestamp, so the export
# filename needs this to stay unique.
"batch_index": index,
"model": result.get("repo_id"),
"created_at": created_at,

13
studio/backend/tests/test_sd_cpp_backend.py
View file

@ -74,6 +74,8 @@ def _loaded_backend(fam_name = "z-image", engine = None):
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,
)
return b
@ -153,6 +155,17 @@ def test_generate_returns_images_and_seed():
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():

4
studio/backend/tests/test_sd_cpp_engine.py
View file

@ -95,7 +95,9 @@ def test_engine_version_parsed_and_cached(tmp_path, monkeypatch):
def _fake_run(*_a, **_k):
calls["n"] += 1
return types.SimpleNamespace(stdout = "stable-diffusion.cpp version master-721\n", stderr = "")
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"