diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index 0c758501f8..ffe31d93a5 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -29,7 +29,12 @@ from typing import Any, Optional from core.inference.diffusion_device import resolve_diffusion_device_target from core.inference.diffusion_families import DiffusionFamily, family_sd_cpp_supported -from core.inference.sd_cpp_backend import _install_allowed, ensure_sd_cpp_binary +from core.inference.sd_cpp_backend import ( + _install_allowed, + _server_binary_runnable, + ensure_sd_cpp_binary, + ensure_sd_server_binary, +) from core.inference.sd_cpp_engine import ( ENGINE_DIFFUSERS, ENGINE_SD_CPP, @@ -137,20 +142,37 @@ def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] fam_ok = family_sd_cpp_supported(fam) binary = None + server_binary = None if policy_eligible and fam_ok: - binary = ensure_sd_cpp_binary( + # Probe the resident sd-server FIRST: the backend PREFERS it, and an sd-server-only + # install (no sd-cli) must still route to native rather than silently falling back to + # diffusers. Checking it before the sd-cli install also means a server-only host does + # not pay an avoidable sd-cli download. Install the accelerator-matched build (ROCm / + # Vulkan / CUDA) so a forced-native GPU load gets the GPU server, not the CPU one. + server_binary = ensure_sd_server_binary( allow_install = _install_allowed(), accelerator = _install_accelerator_for(backend), ) - # Probe runnability here, before committing the route to native: a present but - # non-runnable binary (wrong arch, missing shared libs, no execute bit) would - # otherwise pass as available and only fail inside the background load, instead - # of falling back to diffusers now. + if server_binary and not _server_binary_runnable(server_binary): + logger.warning( + "sd-server at %s is present but not runnable; not using it", server_binary + ) + server_binary = None + # sd-cli is the one-shot fallback. Always LOCATE an existing binary, but only + # auto-INSTALL it when there is no usable server, so a server-only install is not + # forced to also download a CLI it will never use. Probe runnability before + # committing native: a present but non-runnable binary (wrong arch, missing shared + # libs, no execute bit) would otherwise pass as available and only fail inside the + # background load, instead of falling back to diffusers now. + binary = ensure_sd_cpp_binary( + allow_install = _install_allowed() and server_binary is None, + accelerator = _install_accelerator_for(backend), + ) if binary and SdCppEngine(binary = binary).version() is None: - logger.warning("sd-cli at %s is present but not runnable; using diffusers", binary) + logger.warning("sd-cli at %s is present but not runnable; not using it", binary) binary = None - native_available = bool(binary) and policy_eligible and fam_ok + native_available = bool(binary or server_binary) and policy_eligible and fam_ok choice = select_diffusion_engine( backend, native_available = native_available, prefer_native = prefer_native ) @@ -162,8 +184,8 @@ def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] reason = f"GPU backend '{backend}' uses diffusers" elif not fam_ok: reason = f"family '{fam.name}' has no native sd.cpp asset mapping" - elif not binary: - reason = "sd-cli binary unavailable" + elif not (binary or server_binary): + reason = "native sd.cpp binary unavailable" else: reason = "diffusers selected" return _activate(ENGINE_DIFFUSERS, reason) diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 5e0db4f0ce..40a5613273 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -325,6 +325,136 @@ def build_sd_cpp_upscale_command( return cmd +def build_sd_cpp_server_command( + binary: str, + files: SdCppModelFiles, + *, + host: str, + port: int, + vae_format: Optional[str] = None, + offload: Optional[list[str]] = None, + native_speed: Optional[str] = None, + threads: Optional[int] = None, + scratch_dir: Optional[str] = None, + verbose: bool = False, + extra_args: Optional[list[str]] = None, +) -> list[str]: + """Build the ``sd-server`` argv: model + hardware/server flags only. + + ``sd-server`` (stable-diffusion.cpp ``examples/server``) loads the model once at + spawn from the SAME flags ``sd-cli`` takes (``--diffusion-model`` / ``--vae`` / + the text encoders / ``--vae-format`` / offload + speed), and adds ``--listen-ip`` + / ``--listen-port``. Per-generation parameters (prompt, size, steps, seed, cfg, + sampler, batch) are NOT here -- they go in each ``/sdcpp/v1/img_gen`` request, so + one resident process serves many generations without reloading the weights. + + ``offload`` / ``native_speed`` map to the exact same sd.cpp flags as the one-shot + engine (``--offload-to-cpu`` / ``--diffusion-fa`` / ...), verified to be accepted + by ``sd-server --help``. ``scratch_dir`` (if given) is pointed at by the LoRA / + hires-upscaler / embeddings directory flags: sd-server's img_gen handler recursively + iterates those dirs, and an unset / missing dir makes it fail the request, so we give + it a real (empty) directory. ``extra_args`` is appended last so a power user can + override anything (sd.cpp's parser is last-wins). + """ + if not files.diffusion_model: + raise ValueError("diffusion_model path is required") + + cmd: list[str] = [binary, "--diffusion-model", files.diffusion_model] + for flag, value in ( + ("--vae", files.vae), + ("--clip_l", files.clip_l), + ("--clip_g", files.clip_g), + ("--t5xxl", files.t5xxl), + ("--llm", files.llm), + ("--qwen2vl", files.qwen2vl), + ): + if value: + cmd += [flag, value] + if vae_format: + cmd += ["--vae-format", vae_format] + cmd += ["--listen-ip", str(host), "--listen-port", str(int(port))] + if scratch_dir: + cmd += [ + "--lora-model-dir", + scratch_dir, + "--hires-upscalers-dir", + scratch_dir, + "--embd-dir", + scratch_dir, + ] + if threads is not None: + cmd += ["--threads", str(int(threads))] + + offload = list(offload or []) + if offload: + cmd += offload + # De-dup speed flags against offload (offload may already include --diffusion-fa). + cmd += [f for f in native_speed_flags(native_speed) if f not in offload] + if verbose: + cmd += ["-v"] + if extra_args: + cmd += list(extra_args) + return cmd + + +def build_img_gen_request( + *, + prompt: str, + negative_prompt: Optional[str] = None, + width: int = 1024, + height: int = 1024, + steps: Optional[int] = None, + seed: Optional[int] = None, + batch_count: int = 1, + sample_method: Optional[str] = None, + flow_shift: Optional[float] = None, + cfg_scale: Optional[float] = None, + distilled_guidance: Optional[float] = None, + output_format: str = "png", +) -> dict: + """Build the ``POST /sdcpp/v1/img_gen`` JSON body for one text-to-image request. + + The native ``sdcpp`` API takes the whole batch in one request (``batch_count``), + so a batch reuses the resident model with no reload. Sampling lives under + ``sample_params``; guidance is split exactly like the one-shot engine's + ``_map_guidance``: a FLUX distilled value goes to ``guidance.distilled_guidance``, + a real classifier-free scale goes to ``guidance.txt_cfg``. Only set keys are + emitted so the server applies its own defaults for the rest. + """ + if not str(prompt).strip(): + raise ValueError("prompt is required") + + guidance: dict = {} + if cfg_scale is not None: + guidance["txt_cfg"] = float(cfg_scale) + if distilled_guidance is not None: + guidance["distilled_guidance"] = float(distilled_guidance) + + sample_params: dict = {} + if steps is not None: + sample_params["sample_steps"] = int(steps) + if sample_method: + sample_params["sample_method"] = str(sample_method) + if flow_shift is not None: + sample_params["flow_shift"] = float(flow_shift) + if guidance: + sample_params["guidance"] = guidance + + req: dict = { + "prompt": prompt, + "negative_prompt": negative_prompt or "", + "width": int(width), + "height": int(height), + "batch_count": max(1, int(batch_count)), + "output_format": output_format, + } + if seed is not None: + req["seed"] = int(seed) + if sample_params: + req["sample_params"] = sample_params + return req + + def _fmt_float(value: float) -> str: """Compact float -> str: drop a trailing ``.0`` so ``1.0`` -> ``1`` (sd-cli accepts both, but the tidy form keeps logged commands readable).""" diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 3c79b14f2d..53f7131e16 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -49,13 +49,22 @@ from core.inference.diffusion_memory import ( OFFLOAD_NONE, OFFLOAD_SEQUENTIAL, ) -from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles, offload_flags +from core.inference.sd_cpp_args import ( + SdCppGenParams, + SdCppModelFiles, + build_img_gen_request, + offload_flags, +) from core.inference.sd_cpp_engine import ( SdCppCancelled, SdCppEngine, find_sd_cpp_binary, + find_sd_server_binary, + runtime_env, ) +from core.inference.sd_cpp_server import SdCppServer from loggers import get_logger +from utils.subprocess_compat import windows_hidden_subprocess_kwargs logger = get_logger(__name__) @@ -68,6 +77,45 @@ _STEP_RE = re.compile(r"(\d+)\s*/\s*(\d+)") # download / extract / chmod. _install_lock = threading.Lock() +# sd-server accepts at most this many images per img_gen job; larger Studio batches +# (the request model allows up to 32) are split into chunks of this size, the way the +# one-shot path did them one image at a time. +_MAX_SERVER_BATCH = 8 + +# Per-image wall-clock budget for a server job, so a batch gets a timeout proportional to +# its image count (matching the one-shot path, where each image had its own budget) rather +# than one fixed deadline the whole batch has to finish within. +_SERVER_PER_IMAGE_TIMEOUT_S = 1800.0 + + +def _server_binary_runnable(binary: str) -> bool: + """Best-effort probe that ``binary`` can actually execute (not just exist). + + Runs `` --help`` with the same runtime env the server will use, so a present + but unrunnable build (wrong arch, missing shared libs, no execute bit) is caught before + a multi-GB asset download. Conservative: only a clear "cannot launch" signal (OSError, + or the dynamic-loader exit codes 126/127) returns False; anything else is treated as + runnable so a quirky ``--help`` exit code never blocks a working binary.""" + import subprocess + + try: + proc = subprocess.run( + [binary, "--help"], + capture_output = True, + timeout = 20, + env = runtime_env(binary), + **windows_hidden_subprocess_kwargs(), + ) + except OSError: + return False # cannot exec at all (wrong arch / no execute bit / missing loader) + except Exception: # noqa: BLE001 -- timeout or anything odd: don't block on a flaky probe + return True + # A negative return code is a signal death (e.g. -4 SIGILL from an incompatible + # prebuilt on an older CPU): the binary launches but immediately crashes, so treat it + # as unavailable and let the load fall back to diffusers instead of routing to a + # server that will die on startup. + return proc.returncode >= 0 and proc.returncode not in (126, 127) + def ensure_sd_cpp_binary(*, allow_install: bool = True, accelerator: str = "cpu") -> Optional[str]: """Path to a usable ``sd-cli`` binary, installing the prebuilt once if needed. @@ -105,9 +153,51 @@ def ensure_sd_cpp_binary(*, allow_install: bool = True, accelerator: str = "cpu" return None +def ensure_sd_server_binary( + *, allow_install: bool = True, accelerator: str = "cpu" +) -> Optional[str]: + """Path to a usable ``sd-server`` binary, installing the prebuilt once if needed. + + Unlike ``ensure_sd_cpp_binary``, this installs when *sd-server specifically* is + missing -- even if an ``sd-cli`` from an older install is already present -- so an + existing one-shot install is upgraded to the persistent server (the prebuilt archive + ships both). Returns None when it is absent and cannot be installed; the backend then + uses the one-shot fallback. Never raises. + """ + found = find_sd_server_binary() + if found: + return found + if not allow_install: + return None + with _install_lock: + found = find_sd_server_binary() + if found: + return found + try: + import sys + + studio_dir = Path(__file__).resolve().parents[3] # .../studio + if str(studio_dir) not in sys.path: + sys.path.insert(0, str(studio_dir)) + from install_sd_cpp_prebuilt import install as _install + except Exception as exc: # noqa: BLE001 -- import path / module issues are non-fatal + logger.warning("sd-server installer import failed: %s", exc) + return None + try: + _install(accelerator = accelerator) # extracts sd-cli AND sd-server + except Exception as exc: # noqa: BLE001 -- download/extract failure -> fall back + logger.warning("sd-server auto-install failed: %s", exc) + return None + return find_sd_server_binary() + + @dataclass(frozen = True) class _SdState: - """The loaded native checkpoint: resolved asset paths + run settings.""" + """The loaded native checkpoint: resolved asset paths + run settings. + + ``server`` is the resident ``sd-server`` process (the model is loaded once, inside + it) when ``mode == "server"``; in the ``"oneshot"`` fallback it is ``None`` and each + generation re-runs ``sd-cli``.""" repo_id: str base_repo: str @@ -120,6 +210,8 @@ class _SdState: threads: Optional[int] = None sampling_method: Optional[str] = None flow_shift: Optional[float] = None + server: Optional[SdCppServer] = None + mode: str = "server" def _memory_policy(memory_mode: Optional[str], cpu_offload: bool) -> str: @@ -191,11 +283,19 @@ class SdCppDiffusionBackend: self._lock = threading.Lock() self._generate_lock = threading.Lock() self._engine = engine # resolved lazily on first load so import stays cheap + # An engine passed in is an EXPLICIT injection (the test seam / escape hatch) and + # pins one-shot mode; an engine cached later by a runtime fallback must NOT, so a + # now-available server can still be used on the next load. + self._engine_injected = engine is not None self._state: Optional[_SdState] = None self._loading: Optional[_SdLoading] = None self._load_token = 0 self._cancel_event = threading.Event() self._active_generate_cancel: Optional[threading.Event] = None + # The sd-server being started for an in-flight load, before it is committed to + # _state. Tracked so an unload / superseding load can stop it mid-startup instead + # of leaving it loading (and holding the generate lock) for the whole timeout. + self._pending_server: Optional[SdCppServer] = None self._gen: Optional[_SdGen] = None @property @@ -212,6 +312,38 @@ class SdCppDiffusionBackend: self._engine = SdCppEngine(binary = binary) return self._engine + def _resolve_backend(self) -> tuple[str, Optional[str], Optional[SdCppEngine]]: + """Pick the native execution mode: ("server", binary, None) or ("oneshot", None, engine). + + The persistent ``sd-server`` is preferred (load once, serve many). The one-shot + ``sd-cli`` is the fallback for older / custom builds that lack the server target. + An explicitly injected engine forces one-shot (the unit-test seam and an escape + hatch), so a test never spawns a real server or triggers an install. A lazily + cached fallback engine does NOT force one-shot: once a resident server becomes + available (installed, or a per-model start that previously failed now works), the + next load can use it, instead of being pinned to one-shot for the whole session. + """ + if self._engine_injected and self._engine is not None: + return "oneshot", None, self._resolve_engine() + # Install the sd-server build matching the resolved device backend (ROCm / Vulkan / + # CUDA), not the default CPU build: a forced/enabled native load on a GPU host must + # not silently fetch the plain-CPU server. Lazy import avoids an import cycle with + # the router, which imports this backend during engine selection. + from core.inference.diffusion_engine_router import _install_accelerator_for + + accelerator = _install_accelerator_for( + getattr(resolve_diffusion_device_target(), "backend", "cpu") + ) + server_binary = ensure_sd_server_binary( + allow_install = _install_allowed(), accelerator = accelerator + ) + if server_binary is not None: + return "server", server_binary, None + logger.warning( + "sd-server not found; falling back to one-shot sd-cli (reloads the model per image)." + ) + return "oneshot", None, self._resolve_engine() + # ── Background load + progress ───────────────────────────────────────── def begin_load( @@ -298,10 +430,35 @@ class SdCppDiffusionBackend: _load_token: int, ) -> None: try: - # Ensure the binary up front so an install failure surfaces before the - # multi-GB asset pull (the router also pre-checks, but a forced reload here - # must not silently download then fail at generate). - engine = self._resolve_engine() + # Resolve the backend mode (persistent sd-server preferred, one-shot sd-cli + # fallback) and binary up front so an install / missing-binary failure + # surfaces before the multi-GB asset pull. + mode, server_binary, engine = self._resolve_backend() + if mode == "server": + # Probe the server binary before the multi-GB asset pull: a present but + # unrunnable build (wrong arch / missing libs) would otherwise download + # everything and only then fail to start. If it cannot run, fall back to + # the one-shot engine now (when it is usable), else surface the failure. + assert server_binary is not None + if not _server_binary_runnable(server_binary): + logger.warning( + "sd-server at %s is present but not runnable; trying one-shot sd-cli.", + server_binary, + ) + try: + usable = self._resolve_engine().version() is not None + except Exception: # noqa: BLE001 + usable = False + if not usable: + raise RuntimeError("sd-server binary is present but not runnable.") + mode, server_binary, engine = "oneshot", None, self._resolve_engine() + if mode == "oneshot": + # Probe the binary: version() returns None when the present binary cannot + # run (bad perms / missing libs), so fail now rather than commit a "ready" + # state that crashes on the first generation. + assert engine is not None + if engine.version() is None: + raise RuntimeError("sd-cli binary is present but not runnable.") assets = self._asset_specs(repo_id, gguf_filename, fam) self._set_expected_bytes(assets, hf_token) @@ -318,37 +475,21 @@ class SdCppDiffusionBackend: ) 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. + # sd_cpp / MPS), since on CPU the weights are 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, - family = fam, - device = device, - files = files, - vae_format = fam.sd_cpp_vae_format, - 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() 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.") - # A generation that started during the (slow) asset download is still running - # against the OLD model. Abort it, then WAIT on _generate_lock for it to exit - # before publishing the new state -- otherwise that stale sd-cli run can finish - # afterward and persist an image from the previous model once this load reports - # ready (mirrors the diffusers load_pipeline commit). _generate_lock is taken - # only here, not during the download, so the long fetch never serialises against - # generation; the inner token re-check guards an unload/newer load arriving while - # we waited. + native_speed = _native_speed_for(speed_mode) + + # Tear down any previously-loaded model, then commit the new one. A generation + # that started during the (slow) asset download is still running against the OLD + # model: abort it and WAIT on _generate_lock for it to exit before swapping, or + # a stale run could finish afterward and persist an image from the previous + # model. For server mode we stop the old server and start (load) the new one + # HERE, under _generate_lock, so generation never races a half-loaded server and + # two resident models never coexist. _generate_lock is taken only now, not during + # the download, so the long fetch never serialises against generation. with self._lock: if self._load_token != _load_token: return # superseded / cancelled @@ -358,6 +499,78 @@ class SdCppDiffusionBackend: with self._lock: if self._load_token != _load_token: return # superseded / cancelled while waiting + old_state = self._state + self._state = None # the old model is being torn down + if old_state is not None and old_state.server is not None: + old_state.server.stop() + server: Optional[SdCppServer] = None + if mode == "server": + assert server_binary is not None + server = SdCppServer(server_binary) + # Publish the not-yet-committed server so unload() / a superseding load + # can stop it mid-startup (SdCppServer.stop aborts the readiness wait + # without waiting on the lifecycle lock), instead of it loading for the + # full startup timeout while holding the generate lock. + with self._lock: + self._pending_server = server + try: + # Blocks until the server has loaded the model and is answering + # (its readiness check); raises with the log tail on a failed load. + server.start( + files, + vae_format = fam.sd_cpp_vae_format, + offload = list(offload), + native_speed = native_speed, + threads = None, + ) + except SdCppCancelled: + # Startup was aborted by an unload / superseding load: stop the + # half-started server and bail (the outer handler returns cleanly). + server.stop() + raise + except Exception as start_exc: # noqa: BLE001 + # A present-but-unusable sd-server must be no worse than the + # one-shot engine: fall back to sd-cli when it is usable, else + # surface the server error. + logger.warning( + "sd-server failed to start (%s); falling back to one-shot sd-cli.", + start_exc, + ) + server.stop() + server = None + try: + usable = self._resolve_engine().version() is not None + except Exception: # noqa: BLE001 + usable = False + if not usable: + raise start_exc + mode = "oneshot" + finally: + with self._lock: + if self._pending_server is server: + self._pending_server = None + state = _SdState( + repo_id = repo_id, + base_repo = base, + family = fam, + device = device, + files = files, + vae_format = fam.sd_cpp_vae_format, + native_speed = native_speed, + offload_flags = offload, + threads = None, + sampling_method = fam.sd_cpp_sampling_method, + flow_shift = fam.sd_cpp_flow_shift, + server = server, + mode = mode, + ) + with self._lock: + if self._load_token != _load_token: + # Superseded / unloaded while we were loading: discard the server + # we just started so it doesn't leak (and keep _state unloaded). + if server is not None: + server.stop() + return self._state = state self._loading = None except SdCppCancelled: @@ -475,76 +688,62 @@ class SdCppDiffusionBackend: seed: Optional[int] = None, batch_size: int = 1, ) -> dict[str, Any]: - import tempfile - - from PIL import Image - cancel = threading.Event() with self._generate_lock: with self._lock: state = self._state if state is None: raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) + # A resident server can exit while idle; if a client generates without first + # polling status, drop the stale loaded state and report not-loaded so it gets + # the recoverable reload path instead of a 500 from img_gen (not running). + if ( + state.mode == "server" + and state.server is not None + and not state.server.is_alive() + ): + self._state = None + raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) self._active_generate_cancel = cancel - engine = self._resolve_engine() try: if seed is None: seed = int.from_bytes(os.urandom(6), "big") & ((1 << 53) - 1) else: seed = int(seed) cfg_scale, flux_guidance = _map_guidance(state.family, guidance) - 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(): - raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # Distinct seed per batch image (sd-cli is one image/run here), - # so a batch is reproducible image-by-image from the base seed. - # Mask to sd-cli's int64 range, NOT 53 bits: the request model and - # the diffusers backend both accept large explicit seeds, so a tight - # 2**53 mask would silently truncate them (2**53 -> 0) and collide - # distinct requested seeds onto the same image. Randomly-drawn seeds - # above are already 53-bit (JS-safe); explicit seeds pass through. - seed_i = (seed + index) & ((1 << 63) - 1) - out_path = str(Path(tmpdir) / f"img_{index}.png") - params = SdCppGenParams( - prompt = prompt, - negative_prompt = negative_prompt or None, - width = int(width), - height = int(height), - steps = int(steps), - cfg_scale = cfg_scale, - guidance = flux_guidance, - seed = seed_i, - sampling_method = state.sampling_method, - batch_count = 1, - ) - engine.generate( - state.files, - params, - output_path = out_path, - offload = list(state.offload_flags) or None, - native_speed = state.native_speed, - threads = state.threads, - 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 state.mode == "server" and state.server is not None: + images, seeds = self._generate_server( + state, + prompt = prompt, + negative_prompt = negative_prompt, + width = width, + height = height, + steps = steps, + seed = seed, + batch_size = batch_size, + cfg_scale = cfg_scale, + flux_guidance = flux_guidance, + cancel = cancel, + ) + else: + images, seeds = self._generate_oneshot( + state, + prompt = prompt, + negative_prompt = negative_prompt, + width = width, + height = height, + steps = steps, + seed = seed, + batch_size = batch_size, + cfg_scale = cfg_scale, + flux_guidance = flux_guidance, + cancel = cancel, + ) if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # ``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. + # ``seeds`` is the per-image seed (image i used seed+i), so the route can + # persist the real seed for every image in the batch. return { "images": images, "seed": int(seed), @@ -559,6 +758,144 @@ class SdCppDiffusionBackend: if self._active_generate_cancel is cancel: self._active_generate_cancel = None + def _generate_server( + self, + state: _SdState, + *, + prompt: str, + negative_prompt: Optional[str], + width: int, + height: int, + steps: int, + seed: int, + batch_size: int, + cfg_scale: Optional[float], + flux_guidance: Optional[float], + cancel: threading.Event, + ) -> tuple[list, list[int]]: + """Generate via the resident sd-server (no model reload). + + A batch larger than the server's per-job limit is split into chunks: the server + rejects a batch_count above _MAX_SERVER_BATCH, and the one-shot path served large + batches image-by-image, so preserve that. The base seed is masked to sd.cpp's + signed-int64 range (the request model / diffusers accept larger seeds), and each + chunk is submitted at base+offset so the per-image seeds stay reproducible. Each + chunk gets a timeout proportional to its image count so a slow CPU batch is not + cancelled partway through on one fixed deadline.""" + import io + + from PIL import Image + + assert state.server is not None + total = max(1, int(batch_size)) + # sd.cpp's image seed is signed int64; mask the base (and every derived seed) so a + # large explicit seed is not rejected / wrapped inconsistently by the server. + base_seed = int(seed) & ((1 << 63) - 1) + images: list = [] + seeds: list[int] = [] + for offset in range(0, total, _MAX_SERVER_BATCH): + if cancel.is_set(): + raise SdCppCancelled("sd-server generation was cancelled.") + count = min(_MAX_SERVER_BATCH, total - offset) + chunk_seed = (base_seed + offset) & ((1 << 63) - 1) + payload = build_img_gen_request( + prompt = prompt, + negative_prompt = negative_prompt or None, + width = int(width), + height = int(height), + steps = int(steps), + seed = chunk_seed, + batch_count = count, + sample_method = state.sampling_method, + flow_shift = state.flow_shift, + cfg_scale = cfg_scale, + distilled_guidance = flux_guidance, + ) + blobs = state.server.img_gen( + payload, + on_step = self._on_log, + cancel_event = cancel, + total_timeout = _SERVER_PER_IMAGE_TIMEOUT_S * count, + ) + # All-or-nothing per chunk, like the one-shot path: if the server returns fewer + # blobs than requested (e.g. one image in the batch failed to encode), fail + # rather than silently dropping images from the user's requested batch. + if not cancel.is_set() and len(blobs) != count: + raise RuntimeError( + f"sd-server returned {len(blobs)} of {count} requested images in the batch." + ) + images.extend(Image.open(io.BytesIO(b)).convert("RGB") for b in blobs) + # sd.cpp advances the seed per image within a job, so report chunk_seed+i. + seeds.extend((chunk_seed + i) & ((1 << 63) - 1) for i in range(len(blobs))) + return images, seeds + + def _generate_oneshot( + self, + state: _SdState, + *, + prompt: str, + negative_prompt: Optional[str], + width: int, + height: int, + steps: int, + seed: int, + batch_size: int, + cfg_scale: Optional[float], + flux_guidance: Optional[float], + cancel: threading.Event, + ) -> tuple[list, list[int]]: + """Fallback path: re-run one-shot sd-cli per image (reloads the model each time).""" + import tempfile + + from PIL import Image + + engine = self._resolve_engine() + 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))] + + 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(): + raise RuntimeError(DIFFUSION_CANCELLED_MSG) + # Distinct seed per batch image, reproducible image-by-image from the base + # seed. Mask to int64, NOT 53 bits: the request model and the diffusers + # backend both accept large explicit seeds, so a tight 2**53 mask would + # truncate them and collide distinct requested seeds onto the same image. + seed_i = (seed + index) & ((1 << 63) - 1) + out_path = str(Path(tmpdir) / f"img_{index}.png") + params = SdCppGenParams( + prompt = prompt, + negative_prompt = negative_prompt or None, + width = int(width), + height = int(height), + steps = int(steps), + cfg_scale = cfg_scale, + guidance = flux_guidance, + seed = seed_i, + sampling_method = state.sampling_method, + batch_count = 1, + ) + engine.generate( + state.files, + params, + output_path = out_path, + offload = list(state.offload_flags) or None, + native_speed = state.native_speed, + threads = state.threads, + 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) + return images, seeds + def _on_log(self, line: str) -> None: gen = self._gen if gen is None or gen.total_steps <= 0: @@ -596,13 +933,40 @@ class SdCppDiffusionBackend: with self._lock: if self._active_generate_cancel is not None: self._active_generate_cancel.set() + state = self._state self._state = None self._load_token += 1 self._loading = None + # A load may be mid server.start() with the server not yet committed to _state; + # grab it too so we can stop it (its startup is abortable) instead of leaving it + # loading for the full startup timeout. + pending = self._pending_server + self._pending_server = None + # Stop the resident server outside the lock (terminate can take a few seconds). A + # mid-flight generation had its cancel event set above, so its poll loop unwinds + # as the process goes away. + if state is not None and state.server is not None: + state.server.stop() + if pending is not None and pending is not (state.server if state else None): + pending.stop() return self.status() def status(self) -> dict[str, Any]: state = self._state + # A resident sd-server can exit after load (OOM-killed / crashed while idle). If so, + # drop the stale loaded state so status reports not-loaded and clients reload, + # instead of every generation failing with a 500 against a dead process. + if ( + state is not None + and state.mode == "server" + and state.server is not None + and not state.server.is_alive() + ): + logger.warning("sd-server exited after load; clearing loaded state") + with self._lock: + if self._state is state: + self._state = None + state = None if state is None: return { "loaded": False, @@ -622,6 +986,7 @@ class SdCppDiffusionBackend: "attention_backend": None, "transformer_cache": None, "engine": "sd_cpp", + "native_mode": None, } return { "loaded": True, @@ -645,6 +1010,8 @@ class SdCppDiffusionBackend: "attention_backend": None, "transformer_cache": None, "engine": "sd_cpp", + # "server" = resident sd-server (load once); "oneshot" = legacy per-image sd-cli. + "native_mode": state.mode, } diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 182b25a1b7..1c8d8c7033 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -51,6 +51,9 @@ logger = logging.getLogger(__name__) # target is ``sd-cli``; older builds shipped ``sd`` -- both are probed on PATH. _BINARY_STEM = "sd-cli" _LEGACY_STEM = "sd" +# The persistent HTTP server target (stable-diffusion.cpp ``examples/server``). It +# ships next to ``sd-cli`` in both the prebuilt archives and the cmake build tree. +_SERVER_STEM = "sd-server" class SdCppCancelled(RuntimeError): @@ -119,11 +122,11 @@ def runtime_env(binary: str, base_env: Optional[dict[str, str]] = None) -> dict[ return env -def _layout_candidates(root: Path) -> list[Path]: - """sd-cli locations under a stable-diffusion.cpp checkout/install ``root``, +def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]: + """``stem`` locations under a stable-diffusion.cpp checkout/install ``root``, highest priority first: the cmake ``build/bin`` tree, then a Windows Release subdir, then the root itself.""" - name = _binary_name(_BINARY_STEM) + name = _binary_name(stem) cands = [ root / "build" / "bin" / name, root / "build" / "bin" / "Release" / name, @@ -133,66 +136,101 @@ def _layout_candidates(root: Path) -> list[Path]: return cands -def find_sd_cpp_binary() -> Optional[str]: - """Locate the ``sd-cli`` binary, or None. +def _first_file(paths: list[Path]) -> Optional[str]: + for p in paths: + try: + if p.is_file(): + return str(p) + except OSError: + continue + return None - Search order (mirrors the llama.cpp finder so a Studio install lands where - both engines look): - 1. ``SD_CLI_PATH`` env -- a direct path to the binary. + +def _find_binary( + *, direct_env: str, path_stems: tuple[str, ...], layout_stem: str +) -> Optional[str]: + """Shared finder for the stable-diffusion.cpp binaries. + + Search order (mirrors the llama.cpp finder so a Studio install lands where every + binary is looked for): + 1. ``direct_env`` -- a direct path to the binary. 2. ``UNSLOTH_SD_CPP_PATH`` env -- a stable-diffusion.cpp install dir. 3. the installer target: ``/../stable-diffusion.cpp`` when that env (or ``STUDIO_HOME``) is set, else ``~/.unsloth/stable-diffusion.cpp``. 4. ``./stable-diffusion.cpp`` in-tree build (developer checkout). - 5. ``sd-cli`` (then legacy ``sd``) on PATH. + 5. ``path_stems`` on PATH (in order). """ - - def _first_file(paths: list[Path]) -> Optional[str]: - for p in paths: - try: - if p.is_file(): - return str(p) - except OSError: - continue - return None - # 1. Direct binary path. - env_bin = os.environ.get("SD_CLI_PATH") + env_bin = os.environ.get(direct_env) if env_bin and Path(env_bin).is_file(): return env_bin # 2. Custom install dir. custom = os.environ.get("UNSLOTH_SD_CPP_PATH") if custom: - hit = _first_file(_layout_candidates(Path(custom))) + hit = _first_file(_layout_candidates(Path(custom), layout_stem)) if hit: return hit - # 3. Default install root: the installer's default_install_dir() -- a sibling of - # the llama.cpp install under UNSLOTH_STUDIO_HOME / STUDIO_HOME when set, else - # ~/.unsloth. Mirror that env resolution or a custom Studio home never resolves. + # 3. Default install root. Honors UNSLOTH_STUDIO_HOME / STUDIO_HOME the same way + # the installer's default_install_dir does (base = the Studio home's parent), so + # a binary installed under a custom Studio root is discovered and side-by-side + # Studios stay isolated; falls back to the sibling of ~/.unsloth/llama.cpp. studio_home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") - default_base = Path(studio_home).parent if studio_home else Path.home() / ".unsloth" - hit = _first_file(_layout_candidates(default_base / "stable-diffusion.cpp")) + default_root = ( + Path(studio_home).parent / "stable-diffusion.cpp" + if studio_home + else Path.home() / ".unsloth" / "stable-diffusion.cpp" + ) + hit = _first_file(_layout_candidates(default_root, layout_stem)) if hit: return hit # 4. In-tree developer build: /stable-diffusion.cpp. try: project_root = Path(__file__).resolve().parents[4] - hit = _first_file(_layout_candidates(project_root / "stable-diffusion.cpp")) + hit = _first_file(_layout_candidates(project_root / "stable-diffusion.cpp", layout_stem)) if hit: return hit except (OSError, IndexError): pass # 5. PATH. - for stem in (_BINARY_STEM, _LEGACY_STEM): + for stem in path_stems: on_path = shutil.which(stem) if on_path: return on_path return None +def find_sd_cpp_binary() -> Optional[str]: + """Locate the one-shot ``sd-cli`` binary (env ``SD_CLI_PATH``), or None. + + Probes ``sd-cli`` then legacy ``sd`` on PATH. This is the fallback engine once the + persistent ``sd-server`` exists; it also still backs the ESRGAN upscale mode. + """ + return _find_binary( + direct_env = "SD_CLI_PATH", + path_stems = (_BINARY_STEM, _LEGACY_STEM), + layout_stem = _BINARY_STEM, + ) + + +def find_sd_server_binary() -> Optional[str]: + """Locate the persistent ``sd-server`` binary (env ``SD_SERVER_PATH``), or None. + + Same precedence as ``find_sd_cpp_binary`` but keyed to the ``sd-server`` stem, so + a Studio install (prebuilt archive or cmake build, both of which ship ``sd-server`` + next to ``sd-cli``) is found in the same places. Preferred over the one-shot CLI: + it loads the model once and serves many generations without reloading from disk. + """ + return _find_binary( + direct_env = "SD_SERVER_PATH", + path_stems = (_SERVER_STEM,), + layout_stem = _SERVER_STEM, + ) + + class SdCppEngine: """A thin handle over a located ``sd-cli`` binary. diff --git a/studio/backend/core/inference/sd_cpp_server.py b/studio/backend/core/inference/sd_cpp_server.py new file mode 100644 index 0000000000..17b269e71d --- /dev/null +++ b/studio/backend/core/inference/sd_cpp_server.py @@ -0,0 +1,502 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persistent ``sd-server`` (stable-diffusion.cpp) process manager. + +The native diffusion tier used to shell out to one-shot ``sd-cli`` per image, which +reloaded the multi-GB GGUF from disk every generation. ``sd-server`` (the upstream +``examples/server`` target) loads the model once at spawn and serves many generations +over HTTP, exactly like the chat backend's persistent ``llama-server``. This manager +owns ONLY the process + HTTP lifecycle; the backend (``sd_cpp_backend.py``) still owns +asset resolution, request validation, and the public Studio surface. + +Shape mirrors ``core/rag/embed_llama_server.py``: + * ``start`` -- pick a free loopback port, spawn the server (model loads here), + drain stdout on a daemon thread, poll until ready. + * ``img_gen`` -- POST ``/sdcpp/v1/img_gen`` (the whole batch in one request), + poll the async job to a terminal state, return image bytes. + * ``stop`` -- SIGTERM -> wait -> SIGKILL, join the drain thread (idempotent). + +Readiness is real: upstream ``main.cpp`` loads the model BEFORE it binds the port and +prints ``listening on:``, so a 200 from ``GET /v1/models`` (a trivial handler) means the +model is loaded; a load failure exits the process before listening and is surfaced with +the captured log tail. (The richer ``/sdcpp/v1/capabilities`` handler can block in some +builds, so it is not used for readiness.) The job JSON has no per-step field, so step progress is recovered by +parsing the server's stdout (the same ``N/M`` lines ``sd-cli`` emits), routed to the +active generation's callback. + +Import-light on purpose (no torch / diffusers / PIL), so selecting the native tier on +a CPU box never drags the GPU stack into the process. +""" + +from __future__ import annotations + +import atexit +import base64 +import logging +import shutil +import socket +import subprocess +import tempfile +import threading +import time +from collections import deque +from typing import Any, Callable, Optional + +import httpx + +from core.inference.sd_cpp_args import SdCppModelFiles, build_sd_cpp_server_command +from core.inference.sd_cpp_engine import SdCppCancelled, runtime_env +from utils.native_path_leases import child_env_without_native_path_secret +from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid +from utils.subprocess_compat import windows_hidden_subprocess_kwargs + +logger = logging.getLogger(__name__) + +# httpx transport errors meaning "the server is gone / connection refused" -- treated +# as "not ready yet" while polling readiness, and as a fatal "server died" mid-request. +_TRANSPORT_ERRORS = ( + httpx.ConnectError, + httpx.ReadError, + httpx.RemoteProtocolError, + httpx.WriteError, +) + +# Readiness probe. Upstream binds the port only AFTER the model is loaded, so any 200 +# means ready. We use /v1/models (a trivial, always-fast handler) rather than +# /sdcpp/v1/capabilities: the capabilities handler can block in some builds (it enumerates +# model metadata), which would stall readiness even though the server is up. +_READY_PATH = "/v1/models" +# Native async sdcpp API. +_IMG_GEN_PATH = "/sdcpp/v1/img_gen" +_JOBS_PATH = "/sdcpp/v1/jobs" + +_TERMINAL_OK = "completed" +_TERMINAL_FAIL = "failed" +_TERMINAL_CANCELLED = "cancelled" + +# After a cancel is requested, how long to let the server reflect it in job status before +# abandoning the poll. The native cancel is best-effort, so without this cap a server that +# ignores/loses the cancel would keep this call (and the backend's generate lock) alive +# until the job finishes naturally, blocking a superseding load from swapping the model. +_CANCEL_GRACE_S = 5.0 + + +class SdCppServer: + """A resident ``sd-server`` subprocess plus the HTTP client that drives it.""" + + def __init__( + self, + binary: str, + *, + host: str = "127.0.0.1", + ) -> None: + self.binary = binary + self.host = host + self.port: Optional[int] = None + self._process: Optional[subprocess.Popen] = None + # Fixed-size, thread-safe tail buffer: the drain thread appends while lifecycle / + # request threads read it for diagnostics, so a deque(maxlen) is safer and cheaper + # than a list with manual slicing. + self._tail: deque[str] = deque(maxlen = 200) + self._stdout_thread: Optional[threading.Thread] = None + self._lifecycle_lock = threading.Lock() + # Set (lock-free) by stop() so a blocking start()/readiness wait can be aborted + # promptly without waiting on the lifecycle lock start() holds. + self._abort = threading.Event() + # Set for the duration of a generation so the continuous stdout drain can feed + # the active request's step-progress callback; cleared in img_gen's finally. + self._step_listener: Optional[Callable[[str], None]] = None + # trust_env=False: this client only ever talks to the loopback sd-server, so it must + # not route through HTTP_PROXY/HTTPS_PROXY (a proxy without 127.0.0.1 in NO_PROXY + # would break readiness/generation). Matches the local llama-server clients. + self._client = httpx.Client(timeout = 30.0, trust_env = False) + self._scratch_dir: Optional[str] = None + self._stopped = False + atexit.register(self.stop) + + # ── lifecycle ──────────────────────────────────────────────────────────── + + @property + def base_url(self) -> str: + return f"http://{self.host}:{self.port}" + + def is_alive(self) -> bool: + return self._process is not None and self._process.poll() is None + + @staticmethod + def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return int(s.getsockname()[1]) + + def start( + self, + files: SdCppModelFiles, + *, + vae_format: Optional[str] = None, + offload: Optional[list[str]] = None, + native_speed: Optional[str] = None, + threads: Optional[int] = None, + env: Optional[dict[str, str]] = None, + startup_timeout: float = 600.0, + ) -> None: + """Spawn the server (which loads the model) and block until it is ready. + + Raises ``RuntimeError`` (with the captured log tail) if the process exits during + startup or never answers within ``startup_timeout``. Holds the lifecycle lock so + a concurrent start/stop can't interleave. + """ + with self._lifecycle_lock: + # A stop()/unload that raced in AFTER the backend published this server as + # _pending_server but BEFORE start() took the lock has already set _abort and + # closed the httpx client. Honor that delivered stop instead of clearing the + # abort and spawning a model process the cancelled load would then leak. + if self._stopped or self._abort.is_set(): + raise SdCppCancelled("sd-server start was cancelled before launch.") + self._abort.clear() + port = self._find_free_port() + # An empty scratch dir for sd-server's LoRA / upscaler / embeddings scans + # (it recursively iterates them per request and errors on a missing dir). + self._scratch_dir = tempfile.mkdtemp(prefix = "sdcpp_dirs_") + cmd = build_sd_cpp_server_command( + self.binary, + files, + host = self.host, + port = port, + vae_format = vae_format, + offload = list(offload or []), + native_speed = native_speed, + threads = threads, + scratch_dir = self._scratch_dir, + verbose = True, # sd-server prints the per-step sampling lines we parse + ) + run_env = runtime_env(self.binary, child_env_without_native_path_secret()) + if env: + run_env.update(env) + logger.info("starting sd-server: %s", " ".join(cmd)) + # Clear in place: reassigning to [] drops the deque(maxlen=200) bound, so the + # continuous stdout drain would then grow the tail without limit for the whole + # resident-server lifetime. + self._tail.clear() + self._spawn_error: Optional[Exception] = None + spawned = threading.Event() + + # Spawn INSIDE the drain thread, which then reads stdout for the process's whole + # lifetime. child_popen_kwargs() sets PR_SET_PDEATHSIG, which on Linux is bound to + # the CREATING THREAD -- so the child must be created by a thread that outlives it, + # or a transient spawner thread ending would kill the server. The drain thread is + # exactly that long-lived owner; it dies only when the process exits or the + # interpreter goes away (the case we DO want to reap the GPU-resident server). + def _own_process() -> None: + try: + proc = subprocess.Popen( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + errors = "replace", + env = run_env, + **windows_hidden_subprocess_kwargs(), + **child_popen_kwargs(), + ) + except Exception as exc: # noqa: BLE001 -- surface the spawn failure to start() + self._spawn_error = exc + spawned.set() + return + self._process = proc + self.port = port + adopt_pid(proc.pid) # so a global shutdown sweep also reaps it + spawned.set() + self._drain_stdout(proc) + # stdout closed == the process exited; reap it so it is not left a zombie + # until the next stop()/reload. + try: + proc.wait(timeout = 5) + except Exception: # noqa: BLE001 + pass + + self._stdout_thread = threading.Thread( + target = _own_process, daemon = True, name = "sd-server-owner" + ) + self._stdout_thread.start() + spawned.wait() + if self._spawn_error is not None: + self._dispose() + raise RuntimeError(f"failed to spawn sd-server: {self._spawn_error}") + if not self._wait_ready(startup_timeout): + tail = "\n".join(list(self._tail)[-30:]) + aborted = self._abort.is_set() + self._kill_locked() + self._dispose() + if aborted: + raise SdCppCancelled("sd-server startup was cancelled.") + raise RuntimeError("sd-server failed to become ready. Last output:\n" + tail[:2000]) + + def _wait_ready( + self, + timeout: float, + interval: float = 0.5, + ) -> bool: + """Poll ``/v1/models`` until 200; bail early if the process exits. + + Upstream binds the port only AFTER the model is loaded, so a 200 here is a true + ready signal (no half-loaded race).""" + deadline = time.monotonic() + timeout + url = f"{self.base_url}{_READY_PATH}" + while time.monotonic() < deadline: + # A concurrent stop() (unload / superseding load) sets _abort so this wait can + # bail without holding the model-load hostage for the full startup_timeout. + if self._abort.is_set(): + logger.info("sd-server startup aborted before ready") + return False + if not self.is_alive(): + code = None if self._process is None else self._process.returncode + logger.error("sd-server exited early during load (code %s)", code) + return False + try: + if self._client.get(url, timeout = 2.0).status_code == 200: + return True + except (*_TRANSPORT_ERRORS, httpx.TimeoutException): + pass + time.sleep(interval) + logger.error("sd-server readiness timed out after %ss", timeout) + return False + + def _drain_stdout(self, proc: subprocess.Popen) -> None: + """Drain stdout so the pipe never deadlocks; keep a tail for diagnostics and + feed each line to the active generation's step callback.""" + try: + assert proc.stdout is not None + for raw in proc.stdout: + line = raw.rstrip() + if not line: + continue + self._tail.append(line) # deque(maxlen) discards the oldest automatically + logger.debug("[sd-server] %s", line) + cb = self._step_listener + if cb is not None: + try: + cb(line) + except Exception: # noqa: BLE001 -- a progress callback must never break drain + pass + except Exception: # noqa: BLE001 -- drain thread must never raise (pipe closed at teardown) + pass + + def stop(self) -> None: + """Terminate the server (SIGTERM -> SIGKILL), join the drain, and release the HTTP + client + atexit handler. Idempotent.""" + # Signal abort BEFORE contending for the lifecycle lock: a concurrent start() holds + # that lock for the whole (up to startup_timeout) readiness wait, so setting the + # event lets that wait bail immediately instead of stop() blocking behind it. + self._abort.set() + self._stopped = True + with self._lifecycle_lock: + self._kill_locked() + self._dispose() + + def _dispose(self) -> None: + """Release per-instance resources (on stop / failed start). The backend never + reuses a disposed server, so this closes the pooled httpx client and drops the + atexit handler that would otherwise pin every reloaded instance for the session.""" + try: + atexit.unregister(self.stop) + except Exception: # noqa: BLE001 + pass + try: + self._client.close() + except Exception: # noqa: BLE001 + pass + if self._scratch_dir: + shutil.rmtree(self._scratch_dir, ignore_errors = True) + self._scratch_dir = None + + def _kill_locked(self) -> None: + proc = self._process + if proc is None: + return + pid = proc.pid + try: + proc.terminate() + proc.wait(timeout = 5) + except subprocess.TimeoutExpired: + logger.warning("sd-server did not exit on SIGTERM; killing") + try: + proc.kill() + proc.wait(timeout = 5) + except Exception: # noqa: BLE001 -- best-effort teardown + pass + except Exception as exc: # noqa: BLE001 + logger.warning("error terminating sd-server: %s", exc) + finally: + forget_pid(pid) + self._process = None + self.port = None + if self._stdout_thread is not None: + self._stdout_thread.join(timeout = 2) + self._stdout_thread = None + + # ── generation ─────────────────────────────────────────────────────────── + + def img_gen( + self, + payload: dict[str, Any], + *, + on_step: Optional[Callable[[str], None]] = None, + cancel_event: Optional[threading.Event] = None, + poll_interval: float = 0.4, + submit_timeout: float = 60.0, + total_timeout: float = 1800.0, + ) -> list[bytes]: + """Submit one async ``img_gen`` job, poll it to completion, return image bytes. + + ``on_step`` receives each server stdout line (for the step bar). ``cancel_event``, + when set, cancels the job via the native endpoint and raises ``SdCppCancelled``. + Raises ``RuntimeError`` on submit/poll failures (including the server dying), with + the log tail attached. + """ + # If the server was already stopped for a cancel/unload/superseding load that set + # the cancel event before this submit began, report it as a cancellation (which the + # route maps to a client-state 409) rather than a generic "server died" 500. + if self._stopped or not self.is_alive(): + if cancel_event is not None and cancel_event.is_set(): + raise SdCppCancelled("sd-server generation was cancelled.") + raise RuntimeError("sd-server is not running.") + + self._step_listener = on_step + job_id: Optional[str] = None + try: + # Submit -> 202 Accepted + job id. + try: + resp = self._client.post( + f"{self.base_url}{_IMG_GEN_PATH}", json = payload, timeout = submit_timeout + ) + except (*_TRANSPORT_ERRORS, httpx.TimeoutException) as exc: + raise RuntimeError(self._died_message("img_gen submit", exc)) from exc + if resp.status_code == 429: + raise RuntimeError("sd-server job queue is full (HTTP 429).") + if resp.status_code not in (200, 202): + raise RuntimeError( + f"sd-server img_gen submit -> {resp.status_code}: {resp.text[:500]}" + ) + try: + job = resp.json() + except ValueError as exc: + raise RuntimeError( + f"sd-server img_gen returned a non-JSON submit response: {exc}" + ) from exc + if not isinstance(job, dict): + raise RuntimeError( + f"sd-server img_gen returned an unexpected submit response type: {type(job)}" + ) + job_id = job.get("id") + if not job_id: + raise RuntimeError(f"sd-server img_gen returned no job id: {job}") + + # Poll the job to a terminal state. + deadline = time.monotonic() + total_timeout + cancel_sent_at: Optional[float] = None + while True: + if cancel_event is not None and cancel_event.is_set(): + if cancel_sent_at is None: + self.cancel(job_id) + cancel_sent_at = time.monotonic() + elif time.monotonic() - cancel_sent_at > _CANCEL_GRACE_S: + # The best-effort cancel was not reflected in job status within the + # grace window; abandon the poll so the caller can stop the server + # instead of holding the generate lock until the job finishes. + raise SdCppCancelled("sd-server generation was cancelled.") + if not self.is_alive(): + # If we're unwinding a cancel (e.g. unload killed the server), surface a + # clean cancellation rather than a generic "server died" error. + if cancel_event is not None and cancel_event.is_set(): + raise SdCppCancelled("sd-server generation was cancelled.") + raise RuntimeError(self._died_message("img_gen poll", None)) + if time.monotonic() > deadline: + # Best-effort cancel, then tear the server down: current sd-server does + # not interrupt an already-generating job (cancel_generating=false / 409), + # so leaving it up would keep denoising the abandoned job and block later + # generations/reloads behind it. Stopping frees the slot; the backend sees + # the dead server on the next generate and takes the recoverable reload path. + self.cancel(job_id) + self.stop() + raise RuntimeError(f"sd-server generation timed out after {total_timeout}s") + try: + jr = self._client.get(f"{self.base_url}{_JOBS_PATH}/{job_id}", timeout = 10.0) + except (*_TRANSPORT_ERRORS, httpx.TimeoutException): + time.sleep(poll_interval) + continue + except RuntimeError as exc: + # A concurrent stop()/unload closes the shared httpx client; httpx then + # raises a plain RuntimeError ("client has been closed") that is NOT a + # transport error. When we are being cancelled, report it as a clean + # cancellation (route -> 409) instead of a generic 500 generation failure. + if cancel_event is not None and cancel_event.is_set(): + raise SdCppCancelled("sd-server generation was cancelled.") from exc + raise + if jr.status_code in (404, 410): + raise RuntimeError(f"sd-server job {job_id} is gone (HTTP {jr.status_code}).") + if jr.status_code != 200: + time.sleep(poll_interval) + continue + try: + jd = jr.json() + except ValueError as exc: + raise RuntimeError(f"sd-server job status was not JSON: {exc}") from exc + if not isinstance(jd, dict): + raise RuntimeError( + f"sd-server job status returned an unexpected response type: {type(jd)}" + ) + status = jd.get("status") + if status == _TERMINAL_OK: + return self._decode_images(jd) + if status == _TERMINAL_FAIL: + err = jd.get("error") or {} + raise RuntimeError( + "sd-server generation failed: " + f"{err.get('code', 'error')}: {err.get('message', '')}".strip() + ) + if status == _TERMINAL_CANCELLED: + raise SdCppCancelled("sd-server generation was cancelled.") + time.sleep(poll_interval) + finally: + self._step_listener = None + + def cancel(self, job_id: str) -> None: + """Best-effort native cancel of an in-flight job.""" + try: + self._client.post(f"{self.base_url}{_JOBS_PATH}/{job_id}/cancel", timeout = 5.0) + except Exception: # noqa: BLE001 -- cancel is best-effort + pass + + @staticmethod + def _decode_images(job: dict[str, Any]) -> list[bytes]: + # Defensive against an unexpected response shape (a misbehaving/older server): + # verify each level is the type we index before calling dict/list methods. + result = job.get("result") if isinstance(job, dict) else None + images = result.get("images") if isinstance(result, dict) else None + items = [it for it in images if isinstance(it, dict)] if isinstance(images, list) else [] + out: list[bytes] = [] + for item in sorted(items, key = lambda d: d.get("index", 0)): + b64 = item.get("b64_json") + if not b64: + continue + try: + out.append(base64.b64decode(b64)) + except Exception as exc: # noqa: BLE001 + raise RuntimeError(f"sd-server returned an undecodable image: {exc}") from exc + if not out: + raise RuntimeError("sd-server completed the job but returned no images.") + return out + + def _died_message(self, where: str, exc: Optional[Exception]) -> str: + tail = "\n".join(list(self._tail)[-20:]) + base = f"sd-server connection lost during {where}" + if not self.is_alive(): + code = None if self._process is None else self._process.returncode + base += f" (process exited, code {code})" + if exc is not None: + base += f": {exc}" + if tail: + base += "\nLast output:\n" + tail[:1500] + return base diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 04b70e7e11..57ce68b1cb 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1924,6 +1924,11 @@ class DiffusionStatusResponse(BaseModel): ) transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null") engine: Optional[str] = Field(None, description = "Active diffusion engine: diffusers | sd_cpp") + native_mode: Optional[str] = Field( + None, + description = "Native sd.cpp execution mode: server (resident sd-server) | oneshot " + "(per-image sd-cli) | null (diffusers engine)", + ) fallback_reason: Optional[str] = Field( None, description = "Why diffusers was chosen over the native sd.cpp engine (null when none)", diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index 898711e8bd..f4cebc893d 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -32,6 +32,10 @@ def _clean_env_and_state(monkeypatch): "get_active_diffusion_engine", lambda: SimpleNamespace(status = lambda: {"loaded": False, "repo_id": None}), ) + # Default: no resident sd-server (so existing tests exercise the sd-cli path only) and + # a stubbed runnability probe, so neither reaches the real install/exec path. + monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: None) + monkeypatch.setattr(r, "_server_binary_runnable", lambda *_a, **_k: True) yield @@ -70,6 +74,16 @@ def test_cpu_with_binary_and_supported_family_picks_sd_cpp(monkeypatch): assert r.active_engine_name() == ENGINE_SD_CPP +def test_cpu_with_only_sd_server_picks_sd_cpp(monkeypatch): + # An sd-server-only install (no runnable sd-cli) must still route to native: the + # backend prefers the resident server, so a runnable sd-server is native availability. + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, None) # no sd-cli + monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: None)) + monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: "/usr/bin/sd-server") + assert _select() == ENGINE_SD_CPP + + def test_present_but_not_runnable_binary_falls_back(monkeypatch): # A binary that exists but cannot run (version() -> None) must fall back to # diffusers at selection, not commit native and fail inside the load. diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index d3133b08d2..126eed4855 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -21,7 +21,9 @@ from core.inference.sd_cpp_args import ( SdCppGenParams, SdCppModelFiles, SdCppUpscaleParams, + build_img_gen_request, build_sd_cpp_command, + build_sd_cpp_server_command, build_sd_cpp_upscale_command, native_speed_flags, offload_flags, @@ -364,3 +366,104 @@ def test_build_upscale_requires_input_and_model(): SdCppUpscaleParams(input_image = "/i.png", upscale_model = ""), output_path = "/o.png", ) + + +# ── sd-server spawn command ────────────────────────────────────────────────── + + +def test_server_command_has_model_and_listen_but_no_request_params(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/q.gguf") + cmd = build_sd_cpp_server_command( + "/bin/sd-server", files, host = "127.0.0.1", port = 5678, vae_format = "flux2" + ) + assert _pair(cmd, "--diffusion-model") == "/m/z.gguf" + assert _pair(cmd, "--vae") == "/m/ae.sft" + assert _pair(cmd, "--llm") == "/m/q.gguf" + assert _pair(cmd, "--vae-format") == "flux2" + assert _pair(cmd, "--listen-ip") == "127.0.0.1" + assert _pair(cmd, "--listen-port") == "5678" + # Per-request parameters must NOT be baked into the spawn command. + for flag in ( + "--prompt", + "--seed", + "--steps", + "--cfg-scale", + "--guidance", + "--width", + "--height", + "--batch-count", + ): + assert flag not in cmd + + +def test_server_command_maps_offload_and_speed_and_dedupes(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + cmd = build_sd_cpp_server_command( + "/bin/sd-server", + files, + host = "127.0.0.1", + port = 1, + offload = ["--offload-to-cpu", "--diffusion-fa"], + native_speed = "default", # would add --diffusion-fa again + threads = 8, + ) + assert _pair(cmd, "--threads") == "8" + assert cmd.count("--diffusion-fa") == 1 # de-duped against offload + assert "--offload-to-cpu" in cmd + + +def test_server_command_scratch_dir_expands_to_lora_upscaler_embd(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + cmd = build_sd_cpp_server_command( + "/bin/sd-server", files, host = "127.0.0.1", port = 1, scratch_dir = "/tmp/scratch" + ) + assert _pair(cmd, "--lora-model-dir") == "/tmp/scratch" + assert _pair(cmd, "--hires-upscalers-dir") == "/tmp/scratch" + assert _pair(cmd, "--embd-dir") == "/tmp/scratch" + # Absent when not requested. + bare = build_sd_cpp_server_command("/bin/sd-server", files, host = "127.0.0.1", port = 1) + assert "--lora-model-dir" not in bare and "--hires-upscalers-dir" not in bare + + +def test_server_command_requires_diffusion_model(): + with pytest.raises(ValueError): + build_sd_cpp_server_command( + "/bin/sd-server", SdCppModelFiles(diffusion_model = ""), host = "127.0.0.1", port = 1 + ) + + +# ── img_gen request body ───────────────────────────────────────────────────── + + +def test_img_gen_request_maps_core_fields(): + req = build_img_gen_request( + prompt = "a fox", + negative_prompt = "blurry", + width = 512, + height = 768, + steps = 8, + seed = 42, + batch_count = 3, + sample_method = "euler", + cfg_scale = 4.0, + ) + assert req["prompt"] == "a fox" and req["negative_prompt"] == "blurry" + assert req["width"] == 512 and req["height"] == 768 + assert req["seed"] == 42 and req["batch_count"] == 3 + assert req["sample_params"]["sample_steps"] == 8 + assert req["sample_params"]["sample_method"] == "euler" + assert req["sample_params"]["guidance"]["txt_cfg"] == 4.0 + assert req["output_format"] == "png" + + +def test_img_gen_request_flux_uses_distilled_guidance(): + req = build_img_gen_request(prompt = "x", steps = 4, distilled_guidance = 3.5, flow_shift = 3.0) + g = req["sample_params"]["guidance"] + assert g["distilled_guidance"] == 3.5 + assert "txt_cfg" not in g + assert req["sample_params"]["flow_shift"] == 3.0 + + +def test_img_gen_request_requires_prompt(): + with pytest.raises(ValueError): + build_img_gen_request(prompt = " ", steps = 4) diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 46a0514332..fb64f1324b 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -77,10 +77,69 @@ def _loaded_backend(fam_name = "z-image", engine = None): 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 + + 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 ────────────────────────────────────────────────────────── @@ -321,6 +380,260 @@ def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch): 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. diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 5c9d3256af..daaf5223a8 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -24,6 +24,7 @@ from core.inference.sd_cpp_engine import ( ENGINE_SD_CPP, SdCppEngine, find_sd_cpp_binary, + find_sd_server_binary, runtime_env, select_diffusion_engine, ) @@ -75,6 +76,58 @@ def test_find_returns_none_when_absent(tmp_path, monkeypatch): 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 ────────────────────────────────────────────────── diff --git a/studio/backend/tests/test_sd_cpp_server.py b/studio/backend/tests/test_sd_cpp_server.py new file mode 100644 index 0000000000..9e040814e9 --- /dev/null +++ b/studio/backend/tests/test_sd_cpp_server.py @@ -0,0 +1,417 @@ +# 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 persistent sd-server process manager (SdCppServer). + +Hermetic: subprocess.Popen and the httpx client are faked, so nothing spawns a real +binary or opens a socket beyond the free-port probe.""" + +from __future__ import annotations + +import base64 +import io +import threading + +import pytest +from PIL import Image + +from core.inference import sd_cpp_server as srv +from core.inference.sd_cpp_args import SdCppModelFiles +from core.inference.sd_cpp_engine import SdCppCancelled +from core.inference.sd_cpp_server import SdCppServer + +_FILES = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/vae.sft", llm = "/m/llm.sft") + + +def _png_b64(shade: int) -> str: + buf = io.BytesIO() + Image.new("RGB", (1, 1), (shade, shade, shade)).save(buf, format = "PNG") + return base64.b64encode(buf.getvalue()).decode() + + +class _FakePopen: + """Minimal Popen stand-in. stdout yields the scripted lines then BLOCKS until the + process is terminated/killed/exited -- mirroring a real child that holds its pipe + open for its lifetime (so the owner/drain thread stays alive, as in production).""" + + def __init__( + self, + lines = (), + exit_code = None, + ): + self.pid = 4242 + self._lines = list(lines) + self._exit = exit_code # None == alive + self.returncode = exit_code + self.terminated = False + self.killed = False + self._done = threading.Event() + if exit_code is not None: + self._done.set() + + @property + def stdout(self): + def _gen(): + for ln in self._lines: + yield ln + self._done.wait() # hold the pipe open until the process ends + + return _gen() + + def poll(self): + return self._exit + + def terminate(self): + self.terminated = True + self._exit = 0 + self.returncode = 0 + self._done.set() + + def wait(self, timeout = None): + self._done.wait(timeout) + if self._exit is None: + self._exit = 0 + self.returncode = 0 + return self.returncode + + def kill(self): + self.killed = True + self._exit = -9 + self.returncode = -9 + self._done.set() + + +class _Resp: + def __init__( + self, + status_code, + payload = None, + text = "", + bad_json = False, + ): + self.status_code = status_code + self._payload = payload if payload is not None else {} + self.text = text + self._bad_json = bad_json + + def json(self): + if self._bad_json: + raise ValueError("not json") + return self._payload + + +class _FakeClient: + def __init__( + self, + *, + get = None, + post = None, + ): + self._get = get or (lambda url: _Resp(200, {})) + self._post = post or (lambda url, json: _Resp(202, {"id": "job1"})) + self.get_urls = [] + self.post_calls = [] + self.closed = False + + def get( + self, + url, + timeout = None, + ): + self.get_urls.append(url) + return self._get(url) + + def post( + self, + url, + json = None, + timeout = None, + ): + self.post_calls.append((url, json)) + return self._post(url, json) + + def close(self): + self.closed = True + + +@pytest.fixture +def patched(monkeypatch): + """Neutralise process-lifetime side effects for the manager under test.""" + monkeypatch.setattr(srv, "adopt_pid", lambda pid: None) + monkeypatch.setattr(srv, "forget_pid", lambda pid: None) + monkeypatch.setattr(srv, "child_popen_kwargs", lambda: {}) + monkeypatch.setattr(srv, "windows_hidden_subprocess_kwargs", lambda: {}) + return monkeypatch + + +def _server_with(popen, client): + s = SdCppServer("/x/sd-server") + s._client = client + # Attach the fake process + port so generation tests can run without start(). + s._process = popen + s.port = 1234 + return s + + +# ── start / readiness ────────────────────────────────────────────────────────── + + +def test_start_becomes_ready_when_capabilities_200(patched): + popen = _FakePopen(lines = ["loading model", "listening on: http://127.0.0.1:1"]) + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + s = _server_with( + popen, _FakeClient(get = lambda url: _Resp(200, {"model": {"path": "/m/z.gguf"}})) + ) + s.start(_FILES, startup_timeout = 5.0) + assert s.is_alive() is True + assert s.port is not None + + +def test_start_fails_fast_when_process_exits(patched): + # Model load failed -> process exits before listening; start must raise with the tail. + popen = _FakePopen(lines = ["error: bad model"], exit_code = 1) + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + # Capabilities never answers (connection refused) -> readiness relies on exit detection. + s = _server_with( + popen, _FakeClient(get = lambda url: (_ for _ in ()).throw(srv.httpx.ConnectError("refused"))) + ) + with pytest.raises(RuntimeError, match = "failed to become ready"): + s.start(_FILES, startup_timeout = 2.0) + + +# ── generation ─────────────────────────────────────────────────────────────── + + +def _completed_job(images_b64): + return _Resp( + 200, + { + "status": "completed", + "result": {"images": [{"index": i, "b64_json": b} for i, b in enumerate(images_b64)]}, + }, + ) + + +def test_img_gen_returns_image_bytes_in_index_order(patched): + popen = _FakePopen() + s = _server_with( + popen, + _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobA"}), + # result images deliberately out of order -> manager must sort by index. + get = lambda url: _Resp( + 200, + { + "status": "completed", + "result": { + "images": [ + {"index": 1, "b64_json": _png_b64(200)}, + {"index": 0, "b64_json": _png_b64(50)}, + ] + }, + }, + ), + ), + ) + blobs = s.img_gen({"prompt": "x", "batch_count": 2, "sample_params": {"sample_steps": 4}}) + assert len(blobs) == 2 + first = Image.open(io.BytesIO(blobs[0])).convert("RGB").getpixel((0, 0)) + assert first == (50, 50, 50) # index 0 first + + +def test_img_gen_failed_job_raises(patched): + popen = _FakePopen() + s = _server_with( + popen, + _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobF"}), + get = lambda url: _Resp( + 200, {"status": "failed", "error": {"code": "x", "message": "boom"}} + ), + ), + ) + with pytest.raises(RuntimeError, match = "generation failed.*boom"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_queue_full_raises(patched): + popen = _FakePopen() + s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(429, text = "busy"))) + with pytest.raises(RuntimeError, match = "queue is full"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_cancel_posts_cancel_and_raises(patched): + popen = _FakePopen() + cancel = threading.Event() + cancel.set() # already cancelled before the first poll + client = _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobC"}), + get = lambda url: _Resp( + 200, {"status": "cancelled", "error": {"code": "cancelled", "message": "c"}} + ), + ) + s = _server_with(popen, client) + with pytest.raises(SdCppCancelled): + s.img_gen({"prompt": "x"}, cancel_event = cancel) + assert any(url.endswith("/cancel") for url, _ in client.post_calls) + + +def test_img_gen_detects_server_death(patched): + popen = _FakePopen() + + def _die_get(url): + popen._exit = 137 # the process died between submit and poll + return _Resp(200, {"status": "generating"}) + + s = _server_with( + popen, _FakeClient(post = lambda url, json: _Resp(202, {"id": "jobD"}), get = _die_get) + ) + with pytest.raises(RuntimeError, match = "connection lost|process exited"): + s.img_gen({"prompt": "x"}) + + +# ── stdout routing + stop ────────────────────────────────────────────────────── + + +def test_drain_routes_lines_to_step_listener_and_tail(patched): + s = SdCppServer("/x/sd-server") + seen = [] + s._step_listener = seen.append + # exit_code set so stdout ends after the scripted lines (a live fake would block). + s._drain_stdout(_FakePopen(lines = ["sampling 1/8", "", "sampling 8/8", "done"], exit_code = 0)) + assert "sampling 1/8" in seen and "sampling 8/8" in seen + assert "" not in seen # blank lines skipped + assert s._tail[-1] == "done" + + +def test_stop_is_idempotent_and_terminates(patched): + popen = _FakePopen() + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + client = _FakeClient(get = lambda url: _Resp(200, {})) + s = _server_with(popen, client) + s.start(_FILES, startup_timeout = 5.0) + s.stop() + assert popen.terminated is True + assert s.is_alive() is False + assert client.closed is True # stop() releases the pooled HTTP client + s.stop() # second call must not raise + + +def test_img_gen_submit_error_raises(patched): + popen = _FakePopen() + s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(400, text = "bad params"))) + with pytest.raises(RuntimeError, match = "submit -> 400"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_malformed_submit_json_raises(patched): + popen = _FakePopen() + s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(202, bad_json = True))) + with pytest.raises(RuntimeError, match = "non-JSON submit"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_empty_result_raises(patched): + popen = _FakePopen() + s = _server_with( + popen, + _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobE"}), + get = lambda url: _Resp(200, {"status": "completed", "result": {"images": []}}), + ), + ) + with pytest.raises(RuntimeError, match = "no images"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_rejected_after_stop(patched): + popen = _FakePopen() + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + s = _server_with(popen, _FakeClient(get = lambda url: _Resp(200, {}))) + s.start(_FILES, startup_timeout = 5.0) + s.stop() + with pytest.raises(RuntimeError, match = "not running"): + s.img_gen({"prompt": "x"}) + + +# ── cancellation + defensive parsing (review follow-ups) ─────────────────────── + + +def test_img_gen_cancelled_before_submit_reports_cancellation(patched): + # The server was stopped for a cancel/unload before submit; with the cancel event set + # this must surface as a cancellation (route -> 409), not a generic "not running" 500. + popen = _FakePopen() + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + s = _server_with(popen, _FakeClient(get = lambda url: _Resp(200, {}))) + s.start(_FILES, startup_timeout = 5.0) + s.stop() + cancel = threading.Event() + cancel.set() + with pytest.raises(SdCppCancelled): + s.img_gen({"prompt": "x"}, cancel_event = cancel) + + +def test_img_gen_abandons_when_cancel_not_honored(patched): + # A best-effort cancel the server ignores must not pin this call (and the generate + # lock) until natural completion: after the grace window it raises cancellation. + patched.setattr(srv, "_CANCEL_GRACE_S", 0.0) + popen = _FakePopen() + cancel = threading.Event() + cancel.set() + client = _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobG"}), + get = lambda url: _Resp(200, {"status": "generating"}), # never terminal + ) + s = _server_with(popen, client) + with pytest.raises(SdCppCancelled): + s.img_gen({"prompt": "x"}, cancel_event = cancel, poll_interval = 0.01) + + +def test_img_gen_non_dict_submit_json_raises(patched): + popen = _FakePopen() + s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(202, ["not", "a", "dict"]))) + with pytest.raises(RuntimeError, match = "unexpected submit response"): + s.img_gen({"prompt": "x"}) + + +def test_img_gen_non_dict_status_json_raises(patched): + popen = _FakePopen() + s = _server_with( + popen, + _FakeClient( + post = lambda url, json: _Resp(202, {"id": "jobH"}), + get = lambda url: _Resp(200, ["unexpected"]), + ), + ) + with pytest.raises(RuntimeError, match = "unexpected response type"): + s.img_gen({"prompt": "x"}, poll_interval = 0.01) + + +def test_decode_images_tolerates_unexpected_shapes(): + # A misbehaving/older server can return non-dict result/images/items; _decode_images + # must raise a clean "no images" rather than an AttributeError on .get(). + for job in ({"result": ["x"]}, {"result": {"images": "nope"}}, {"result": {"images": [1, 2]}}): + with pytest.raises(RuntimeError, match = "no images"): + SdCppServer._decode_images(job) + + +def test_start_aborted_by_concurrent_stop(patched): + # A stop() during the readiness wait must abort start() promptly (without waiting out + # the startup timeout) and surface as a cancellation. + popen = _FakePopen(lines = ["loading model"]) + patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen) + + def _never_ready(url): + raise srv.httpx.ConnectError("refused") + + s = _server_with(popen, _FakeClient(get = _never_ready)) + + def _stop_soon(): + import time as _t + _t.sleep(0.2) + s.stop() + + threading.Thread(target = _stop_soon, daemon = True).start() + with pytest.raises(SdCppCancelled): + s.start(_FILES, startup_timeout = 30.0) diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py index 4b2b36f95a..c76aee11f8 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/studio/install_sd_cpp_prebuilt.py @@ -140,6 +140,17 @@ def _locate_sd_cli(root: Path) -> Optional[Path]: return None +def _locate_sd_server(root: Path) -> Optional[Path]: + """The persistent ``sd-server`` binary in the extracted tree, if the archive ships + one (modern stable-diffusion.cpp releases do). Best-effort: the native backend + falls back to one-shot ``sd-cli`` when it is absent.""" + name = "sd-server.exe" if sys.platform == "win32" else "sd-server" + for p in root.rglob(name): + if p.is_file(): + return p + return None + + def _download( url: str, dest: Path, @@ -237,6 +248,13 @@ def install( if sys.platform != "win32": _make_executable(sd_cli) print(f"installed sd-cli -> {sd_cli}", flush = True) + # The same archive ships the persistent sd-server; make it runnable too so the + # native backend can prefer it (load once, serve many) over one-shot sd-cli. + sd_server = _locate_sd_server(target) + if sd_server is not None and sys.platform != "win32": + _make_executable(sd_server) + if sd_server is not None: + print(f"installed sd-server -> {sd_server}", flush = True) return sd_cli