diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index aebc12e983..c15f10a171 100644 --- a/scripts/build_prequant_checkpoint.py +++ b/scripts/build_prequant_checkpoint.py @@ -55,6 +55,7 @@ def main(argv = None) -> int: # Reuse the runtime quant factory + filter so offline == runtime (the LPIPS-0 invariant). from core.inference.diffusion_transformer_quant import ( + FP8_GRANULARITY, TQ_FP8, TQ_SCHEMES, _make_quant_config, @@ -100,25 +101,30 @@ def main(argv = None) -> int: k: (v.detach().to("cpu") if hasattr(v, "detach") else v) for k, v in transformer.state_dict().items() } + metadata = { + "base_model_id": args.base, + "family": fam.name, + "scheme": scheme, + "min_features": args.min_features, + # The layers skipped for this scheme (int8's M=1 modulation projections; () for + # the scaled_mm schemes) and, for fp8, the baked accumulate mode. Both let the + # loader reject a checkpoint that would not match the runtime path. + "exclude_name_tokens": list(exclude_name_tokens), + "fast_accum": fast_accum, + "torch_dtype": args.dtype, + "quant_backend": "torchao", + "transformer_class": fam.transformer_class, + "torch_version": torch.__version__, + "torchao_version": getattr(torchao, "__version__", "?"), + "diffusers_version": diffusers.__version__, + } + # Record the fp8 granularity so the loader can reject a stale per-tensor checkpoint + # (the runtime now requires per-row; see FP8_GRANULARITY). + if scheme == TQ_FP8: + metadata["fp8_granularity"] = FP8_GRANULARITY ckpt = { "format": PREQUANT_FORMAT, - "metadata": { - "base_model_id": args.base, - "family": fam.name, - "scheme": scheme, - "min_features": args.min_features, - # The layers skipped for this scheme (int8's M=1 modulation projections; () for - # the scaled_mm schemes) and, for fp8, the baked accumulate mode. Both let the - # loader reject a checkpoint that would not match the runtime path. - "exclude_name_tokens": list(exclude_name_tokens), - "fast_accum": fast_accum, - "torch_dtype": args.dtype, - "quant_backend": "torchao", - "transformer_class": fam.transformer_class, - "torch_version": torch.__version__, - "torchao_version": getattr(torchao, "__version__", "?"), - "diffusers_version": diffusers.__version__, - }, + "metadata": metadata, "state_dict": state_dict, } diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index ec4c2fc41c..957df10ef7 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, @@ -148,20 +153,37 @@ def select_and_activate_engine( 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 ) @@ -173,8 +195,8 @@ def select_and_activate_engine( 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/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 775be1c962..f49843ed3d 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -421,6 +421,38 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str: return base or fam.base_repo +# Default (steps, guidance) per model for callers that can't pass them — namely +# the OpenAI /v1/images/generations endpoint, whose spec has no step/guidance +# knobs. Distilled "turbo/schnell" models want few steps and no CFG; the full +# "dev" models want more steps and real CFG. Matched by substring, most specific +# first — the same scheme and values as the UI's MODEL_DEFAULTS table +# (studio/frontend/src/features/images/images-page.tsx); keep the two in sync. +_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = ( + ("z-image-turbo", 9, 0.0), + ("flux.1-schnell", 4, 0.0), + ("flux.1", 28, 3.5), + ("flux.2-klein", 4, 0.0), + ("qwen-image", 20, 4.0), + ("z-image", 20, 4.0), +) +# Unrecognised model: distilled few-step / no-CFG shape, matching the UI fallback. +_GENERATION_DEFAULT_FALLBACK = (9, 0.0) + + +def default_generation_params(*identifiers: Optional[str]) -> tuple[int, float]: + """Default ``(steps, guidance)`` for a loaded model. The first identifier that + names a known model wins (the repo id, then the resolved base repo), so a + local-path load — whose repo id is just a filesystem path that may not name + the model — still resolves via its base repo. Within an identifier, keys are + matched as substrings, most specific first (the same scheme as the UI).""" + for identifier in identifiers: + needle = (identifier or "").lower() + for key, steps, guidance in _GENERATION_DEFAULTS: + if key in needle: + return steps, guidance + return _GENERATION_DEFAULT_FALLBACK + + def family_prequant_repo(fam: DiffusionFamily, scheme: str) -> Optional[str]: """The hosted pre-quantized transformer repo for ``scheme`` in this family, or None.""" for entry_scheme, repo_id in fam.prequant_repos: diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 088ebc3575..485c4a1847 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -268,6 +268,22 @@ def _validate_checkpoint( if meta.get("scheme") != scheme: _warn(logger, scheme, ValueError(f"checkpoint scheme {meta.get('scheme')!r} != {scheme!r}")) return False + # fp8 REQUIRES per-row granularity (per-tensor collapses outlier-heavy DiTs to noise). A + # checkpoint built before that fix carries the old per-tensor layout and either omits + # ``fp8_granularity`` or records something other than per-row; reject it so the loader + # falls back to rebuilding/re-quantising instead of installing a broken fp8 transformer. + from .diffusion_transformer_quant import FP8_GRANULARITY, TQ_FP8 + + if scheme == TQ_FP8 and meta.get("fp8_granularity") != FP8_GRANULARITY: + _warn( + logger, + scheme, + ValueError( + f"fp8 checkpoint granularity {meta.get('fp8_granularity')!r} != " + f"{FP8_GRANULARITY!r} (stale per-tensor artifact); rebuild it" + ), + ) + return False ckpt_base = meta.get("base_model_id") if base: # A checkpoint whose keys happen to match a different base can load strict=True and diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 9e2e56a2a5..39aa751778 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -37,6 +37,12 @@ TQ_AUTO = "auto" TQ_SCHEMES = (TQ_INT8, TQ_FP8, TQ_NVFP4, TQ_MXFP8) TQ_MODES = (TQ_AUTO,) + TQ_SCHEMES +# The fp8 weight/activation granularity the runtime config uses (see _make_quant_config): +# per-ROW is REQUIRED for correctness on outlier-heavy DiTs. Stamped into a pre-quantized +# fp8 checkpoint's metadata at build time and required by the loader, so a stale checkpoint +# baked with the old per-TENSOR layout is rejected and rebuilt rather than reproducing noise. +FP8_GRANULARITY = "per_row" + # Skip linears whose in/out features are below this. A small share of the FLOPs, so leaving # them bf16 costs ~nothing and keeps quality a touch higher. DEFAULT_MIN_LINEAR_FEATURES = 512 @@ -298,19 +304,29 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any: if scheme == TQ_INT8: return Int8DynamicActivationInt8WeightConfig() if scheme == TQ_FP8: - # Choose fp8 accumulate by GPU class (unless forced). On consumer / workstation - # cards (GDDR) the fp8 tensor cores run ~2x faster with FP16 (fast) accumulate - # than FP32 (e.g. ~838 vs ~419 TFLOPS on RTX 50xx), so fast accumulate is a real - # win there. Data-center HBM parts default to the higher-precision accumulate. - # fast accumulate is a precision (not overflow) tradeoff and stays below the fp8 - # quant noise floor (measured 0 non-finite even on Z-Image's ~1e6 activations). + # Per-ROW granularity (per-token activation scale + per-output-channel weight scale) + # is REQUIRED for correctness, not just quality. torchao's default fp8 granularity is + # per-TENSOR: one scale for the whole activation tensor. A DiT with extreme activation + # outliers breaks under that -- z-image's MLP activations peak near 6.6e4, and a single + # such outlier forces a tensor-wide scale that pushes every normal value (~1-30) below + # the fp8 resolution, so they quantise to ~0 and the denoise collapses to pure noise + # (measured on B200: per-tensor fp8 = noise, per-row fp8 = matches bf16). Per-row + # confines each outlier to its own token/channel. This is also why int8 dynamic was + # always fine here: it is per-token by default. The per-row scaled_mm is probed by + # _smoke_probe, so an arch/build without it falls through the ladder to int8. + # + # fast accumulate (fp8 only) is chosen by GPU class unless forced: consumer / workstation + # cards (GDDR) run the fp8 tensor cores ~2x faster with FP16 (fast) accumulate than FP32 + # (e.g. ~838 vs ~419 TFLOPS on RTX 50xx); data-center HBM parts keep precise accumulate. + from torchao.quantization import PerRow try: from torchao.float8 import Float8MMConfig return Float8DynamicActivationFloat8WeightConfig( - mm_config = Float8MMConfig(use_fast_accum = _resolve_fast_accum(fast_accum)) + granularity = PerRow(), + mm_config = Float8MMConfig(use_fast_accum = _resolve_fast_accum(fast_accum)), ) - except Exception: # noqa: BLE001 — older torchao without the explicit knob - return Float8DynamicActivationFloat8WeightConfig() + except Exception: # noqa: BLE001 — older torchao without the explicit mm knob + return Float8DynamicActivationFloat8WeightConfig(granularity = PerRow()) if scheme == TQ_NVFP4: from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 7d70309ca1..2edd230f26 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -326,6 +326,142 @@ 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", + lora: Optional[list[dict]] = None, +) -> 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 + # Structured LoRA list: the sdcpp API resolves each ``path`` against the server's + # scanned ``--lora-model-dir`` (prompt-embedded ```` tags are intentionally + # unsupported server-side), so LoRAs must be staged there and named here. + if lora: + req["lora"] = lora + 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 43df6aa12b..7632551661 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -50,13 +50,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__) @@ -69,6 +78,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. @@ -106,9 +154,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 @@ -121,6 +211,8 @@ class _SdState: threads: Optional[int] = None sampling_method: Optional[str] = None flow_shift: Optional[float] = None + server: Optional[SdCppServer] = None + mode: str = "server" # Token kept so LoRA adapters selected at generate time can be fetched from the Hub. hf_token: Optional[str] = None @@ -194,11 +286,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 @@ -215,6 +315,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( @@ -308,10 +440,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) @@ -328,38 +485,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, - hf_token = hf_token, - ) - # 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 @@ -369,6 +509,79 @@ 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, + hf_token = hf_token, + ) + 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: @@ -496,11 +709,13 @@ class SdCppDiffusionBackend: upscale: Optional[float] = None, # Reference workflow is GPU/diffusers-only (FLUX.2); accepted for interface parity. reference_images: Optional[list[str]] = None, - # LoRA adapters as (id, weight) pairs; resolved + materialized into a managed dir - # and injected as prompt tags per sd-cli run. None/empty = no LoRA. + # LoRA adapters as (id, weight) pairs; resolved up front, then applied per engine + # path: prompt tags for one-shot sd-cli, structured `lora` entries + # for the resident sd-server. None/empty = no LoRA. loras: Optional[list[tuple[str, float]]] = None, - # ControlNet: accepted for interface parity with the diffusers backend. Native sd.cpp - # ControlNet (sd-cli --control-net) is a follow-up; a request is rejected clearly. + # Accepted for the uniform engine interface; the guard below rejects it on the native + # engine (ControlNet is diffusers-only) like img2img/inpaint, so a direct API call with + # ControlNet set fails clearly instead of TypeError'ing on an unexpected kwarg. controlnet: Optional[tuple[str, str, str, float, float, float]] = None, ) -> dict[str, Any]: import tempfile @@ -534,8 +749,17 @@ class SdCppDiffusionBackend: 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) @@ -543,11 +767,10 @@ class SdCppDiffusionBackend: seed = int(seed) cfg_scale, flux_guidance = _map_guidance(state.family, guidance) # Resolve any selected LoRA adapters up front (downloads land in the HF - # cache; a bad id fails here as a clear 400 before we spawn sd-cli). - # Drop weight-0 rows BEFORE the support gate: LoraSpec documents weight 0 - # as disabling the adapter (the diffusers path treats it as empty), so a - # request carrying only disabled rows must stay a no-op even on a family - # where native LoRA is unsupported, rather than 400 on a dead selection. + # cache; a bad id fails here as a clear 400 before we generate). Drop + # weight-0 rows BEFORE the support gate: weight 0 disables an adapter, so a + # request carrying only disabled rows stays a no-op even on a family where + # native LoRA is unsupported, rather than 400 on a dead selection. lora_resolved: list = [] active_loras = [(i, w) for (i, w) in (loras or []) if w != 0] if active_loras: @@ -564,71 +787,41 @@ class SdCppDiffusionBackend: lora_resolved = diffusion_lora.resolve_specs( active_loras, hf_token = state.hf_token, cancel_event = cancel ) - 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: - # Materialize selected LoRAs into a managed dir sd-cli can scan, and - # inject matching tags into the prompt (deduped against - # any the user typed). Empty -> prompt/dir unchanged. - eff_prompt = prompt - lora_dir: Optional[str] = None - if lora_resolved: - materialized = diffusion_lora.materialize_native_dir( - lora_resolved, Path(tmpdir) / "loras" - ) - eff_prompt = diffusion_lora.inject_prompt_tags(prompt, materialized) - lora_dir = str(Path(tmpdir) / "loras") - 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 = eff_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, - lora_dir = lora_dir, - lora_apply_mode = "auto" if lora_dir else None, - ) - 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, + lora_resolved = lora_resolved, + 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, + lora_resolved = lora_resolved, + 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), @@ -643,6 +836,195 @@ 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], + lora_resolved: list, + 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. + + LoRA on the server goes through the structured ``lora`` request field, NOT prompt + tags (the sdcpp API intentionally ignores ```` in the prompt). Selected + adapters are staged into the server's ``--lora-model-dir`` scratch dir, which the + server rescans per request, and referenced by their staged filename.""" + import io + import os + import shutil + + from PIL import Image + + from core.inference import diffusion_lora + + 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] = [] + # Stage selected LoRAs into a per-request subdir of the server's lora-model-dir so a + # previous request's adapters can't leak into this one; reference them by the path + # relative to that dir (what the server's recursive scan resolves against). The + # subdir is removed after the batch. supports_lora already gated the family upstream. + lora_payload: Optional[list[dict]] = None + lora_stage: Optional[Path] = None + if lora_resolved: + server_lora_dir = state.server.lora_dir + if server_lora_dir: + lora_stage = Path(server_lora_dir) / f"gen_{os.urandom(6).hex()}" + materialized = diffusion_lora.materialize_native_dir(lora_resolved, lora_stage) + lora_payload = [ + {"path": f"{lora_stage.name}/{Path(m.path).name}", "multiplier": float(m.weight)} + for m in materialized + ] + try: + 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, + lora = lora_payload, + ) + 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))) + finally: + if lora_stage is not None: + shutil.rmtree(lora_stage, ignore_errors = True) + 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], + lora_resolved: list, + cancel: threading.Event, + ) -> tuple[list, list[int]]: + """Fallback path: re-run one-shot sd-cli per image (reloads the model each time). + + LoRA on the one-shot path uses sd-cli's own mechanism: materialize the selected + adapters into a ``--lora-model-dir`` and inject matching ```` tags + into the prompt (sd-cli parses and strips them). supports_lora already gated the + family upstream, so a non-empty ``lora_resolved`` is safe to apply here.""" + import tempfile + + from PIL import Image + + from core.inference import diffusion_lora + + 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: + # Materialize selected LoRAs into a managed dir sd-cli can scan, and inject + # matching tags into the prompt (deduped against any the user + # typed). Empty -> prompt/dir unchanged. + eff_prompt = prompt + lora_dir: Optional[str] = None + if lora_resolved: + materialized = diffusion_lora.materialize_native_dir( + lora_resolved, Path(tmpdir) / "loras" + ) + eff_prompt = diffusion_lora.inject_prompt_tags(prompt, materialized) + lora_dir = str(Path(tmpdir) / "loras") + 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 = eff_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, + lora_dir = lora_dir, + lora_apply_mode = "auto" if lora_dir else None, + ) + 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: @@ -680,13 +1062,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, @@ -706,6 +1115,7 @@ class SdCppDiffusionBackend: "attention_backend": None, "transformer_cache": None, "engine": "sd_cpp", + "native_mode": None, "supports_lora": False, "supports_controlnet": False, "workflows": [], @@ -740,8 +1150,10 @@ class SdCppDiffusionBackend: model_kind = "gguf", transformer_quant = None, ), - # Native ControlNet (sd-cli --control-net) is a follow-up; off for now. + # ControlNet is diffusers-only; the native engine's generate() rejects it. "supports_controlnet": False, + # "server" = resident sd-server (load once); "oneshot" = legacy per-image sd-cli. + "native_mode": state.mode, # The native engine supports plain text-to-image only (generate() rejects # img2img / inpaint / reference / upscale), so advertise just txt2img. Without # this the status omits workflows, the UI reads [], and it disables the Create diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 18a64d9ea0..52e4d6d693 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,37 +136,39 @@ 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 default install root build layouts (the installer target); honors ``UNSLOTH_STUDIO_HOME`` / ``STUDIO_HOME``, 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 @@ -177,27 +182,55 @@ def find_sd_cpp_binary() -> Optional[str]: if studio_home else Path.home() / ".unsloth" / "stable-diffusion.cpp" ) - hit = _first_file(_layout_candidates(default_root)) + 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..d66f3cdded --- /dev/null +++ b/studio/backend/core/inference/sd_cpp_server.py @@ -0,0 +1,510 @@ +# 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}" + + @property + def lora_dir(self) -> Optional[str]: + """The server's ``--lora-model-dir`` scratch dir. Adapters staged here are picked + up per request (the sdcpp ``img_gen`` route refreshes its LoRA scan each call), so + server-mode LoRA works by materializing here and referencing the files via the + structured ``lora`` request field (the server ignores ```` prompt tags).""" + return self._scratch_dir + + 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 9348fa448d..58701b5828 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2089,6 +2089,11 @@ class DiffusionStatusResponse(BaseModel): "txt2img, img2img, inpaint. Empty when nothing is loaded or on the native engine.", ) 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)", @@ -2105,3 +2110,68 @@ class DiffusionStatusResponse(BaseModel): "picker's enabled state). Diffusers only, for families with a ControlNet pipeline; False " "for the native engine, GGUF-via-diffusers, and torchao fp8/int8 dense.", ) + + +# ── OpenAI-compatible images API (POST /v1/images/generations) ── +# +# Shapes mirror OpenAI's CreateImageRequest / ImagesResponse so off-the-shelf +# OpenAI clients work unchanged. The loaded image GGUF stands in for the model; +# GPT-image-only knobs (quality, style, background, output_format, ...) are +# accepted and ignored, exactly as dall-e-2 ignores them. The size string is +# parsed and `stream` is rejected in the route, where the diffusion backend is +# in reach; everything Pydantic can check declaratively lives here. + + +class ImageGenerationRequest(BaseModel): + """OpenAI ``CreateImageRequest`` for ``POST /v1/images/generations``. + + ``prompt`` is the only required field, per the spec. Unlisted OpenAI fields + are ignored (Pydantic's default), matching dall-e-2's treatment of the + GPT-image-only parameters.""" + + prompt: str = Field(..., min_length = 1, description = "Text description of the image(s).") + model: Optional[str] = Field( + None, description = "Model id (informational; the loaded image model is used)." + ) + n: int = Field(1, ge = 1, le = 10, description = "Number of images to generate (1-10).") + size: str = Field( + "auto", description = "'auto' or 'x' (256-2048, each a multiple of 16)." + ) + response_format: Literal["url", "b64_json"] = Field( + "url", description = "Return each image as a URL or a base64-encoded PNG." + ) + user: Optional[str] = Field(None, description = "End-user identifier (accepted, unused).") + # gpt-image-only; declared so we can reject it with a clear error instead of + # silently returning JSON to a client that asked for an SSE stream. + stream: Optional[bool] = Field( + None, description = "Streaming image generation is not supported; omit or set false." + ) + + @field_validator("n", "size", "response_format", mode = "before") + @classmethod + def _null_means_default(cls, value, info): + # OpenAI marks these nullable WITH a default, so an explicit null means + # "use the default" — coalesce it instead of 400-ing a spec-valid body. + if value is None: + return cls.model_fields[info.field_name].default + return value + + +class ImageGenerationData(BaseModel): + """One image in an ``ImagesResponse`` (OpenAI ``Image``). Exactly one of + ``url`` / ``b64_json`` is set, per the request's ``response_format``; the + route serializes with ``exclude_none`` so the unused key is omitted.""" + + b64_json: Optional[str] = Field( + None, description = "Base64-encoded PNG (response_format=b64_json)." + ) + url: Optional[str] = Field(None, description = "URL to the PNG bytes (response_format=url).") + + +class ImageGenerationResponse(BaseModel): + """OpenAI ``ImagesResponse``. dall-e-shaped: the GPT-image-only top-level + fields (background/output_format/size/quality/usage) are omitted, since our + sizes wouldn't satisfy their fixed enums and we report no token usage.""" + + created: int = Field(..., description = "Unix timestamp (seconds) the images were created.") + data: list[ImageGenerationData] = Field(..., description = "The generated images.") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0beffd9187..5f2169f86a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1062,6 +1062,9 @@ from models.inference import ( DiffusionLoadProgressResponse, GalleryImage, GalleryListResponse, + ImageGenerationRequest, + ImageGenerationData, + ImageGenerationResponse, LoadResponse, LoadProgressResponse, UnloadResponse, @@ -11375,3 +11378,167 @@ async def diffusion_load_progress(current_subject: str = Depends(get_current_sub async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)): from core.inference.diffusion_engine_router import get_active_diffusion_engine return DiffusionGenerateProgressResponse(**get_active_diffusion_engine().generate_progress()) + + +# ────────────────────────────────────────────────────────────────────────── +# OpenAI-compatible images API (POST /v1/images/generations) +# +# The inference router is mounted at both /api/inference and /v1, so this also +# answers /v1/images/generations for off-the-shelf OpenAI clients. It maps +# OpenAI's CreateImageRequest onto the in-process diffusion backend and returns +# an ImagesResponse. Studio's own Image tab uses the richer /images/generate +# route above; this is the spec-shaped surface, and the single error boundary +# mapping backend exceptions to OpenAI error envelopes (the global /v1 handler +# wraps HTTPException detail into the envelope). +# ────────────────────────────────────────────────────────────────────────── + + +# Diffusion dims must land in [256, 2048] on a multiple of 16 (8x VAE downsample +# x 2x patch); the named OpenAI sizes (1024x1024, 1536x1024, 256x256, ...) all +# satisfy this. Mirrors DiffusionGenerateRequest's width/height bounds so both +# generate paths accept the same geometry. +_IMAGE_SIZE_RE = _re.compile(r"^(\d{1,5})\s*x\s*(\d{1,5})$") +_IMAGE_DIM_MIN, _IMAGE_DIM_MAX = 256, 2048 +# Sanitized 503 detail shared by the pre-check and the unload-race branch, so both +# "no image model" responses stay identical. +_NO_IMAGE_MODEL_MSG = "No image model loaded. Load an image model first." + + +def _parse_openai_image_size(size: str) -> tuple[int, int]: + """OpenAI ``size`` -> (width, height). ``auto``/empty -> 1024x1024 (~1MP, what + these models target). Raises ValueError with a client-facing message.""" + text = (size or "").strip().lower() + if text in ("", "auto"): + return 1024, 1024 + match = _IMAGE_SIZE_RE.match(text) + if not match: + raise ValueError("size must be 'auto' or 'x', e.g. '1024x1024'.") + width, height = int(match.group(1)), int(match.group(2)) + for label, value in (("width", width), ("height", height)): + if not _IMAGE_DIM_MIN <= value <= _IMAGE_DIM_MAX: + raise ValueError(f"size {label} must be between {_IMAGE_DIM_MIN} and {_IMAGE_DIM_MAX}.") + if value % 16 != 0: + raise ValueError(f"size {label} must be a multiple of 16.") + return width, height + + +def _absolute_image_url(request: Request, relative: str) -> str: + """Join a relative gallery path onto the request's own scheme+host, for the + response_format=url links. Like every Studio route, the target needs the + bearer token; b64_json avoids that for clients that can't carry it.""" + return str(request.base_url).rstrip("/") + relative + + +@router.post( + "/images/generations", + response_model = ImageGenerationResponse, + response_model_exclude_none = True, +) +async def openai_image_generations( + body: ImageGenerationRequest, + request: Request, + current_subject: str = Depends(get_current_subject), +): + """OpenAI-compatible text-to-image (POST /v1/images/generations). + + Generates ``n`` images from ``prompt`` on the loaded diffusion model and + returns them as URLs (default) or base64 PNGs per ``response_format``. Steps + and guidance have no OpenAI knob, so they default per loaded model.""" + from core.inference import image_gallery + from core.inference.diffusion_engine_router import get_active_diffusion_engine + from core.inference.diffusion_families import default_generation_params + + if body.stream: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + "Streaming image generation is not supported.", status = 400, param = "stream" + ), + ) + try: + width, height = _parse_openai_image_size(body.size) + except ValueError as exc: + raise HTTPException( + status_code = 400, detail = openai_error_body(str(exc), status = 400, param = "size") + ) + + # Use the active engine (diffusers OR native sd.cpp on a no-GPU host), the same + # accessor /images/generate uses, so a model loaded on the native engine isn't + # wrongly reported unloaded here. + backend = get_active_diffusion_engine() + status = backend.status() + if not status.get("loaded"): + # Mirror /v1/completions and /v1/embeddings, which 503 when their backend + # isn't loaded; the global handler turns this into the OpenAI envelope. + raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG) + + # Fall back to the resolved base repo so a local-path load (whose repo_id is a + # filesystem path) still gets the right per-model steps/guidance. + steps, guidance = default_generation_params(status.get("repo_id"), status.get("base_repo")) + try: + result = await asyncio.to_thread( + backend.generate, + prompt = body.prompt, + width = width, + height = height, + steps = steps, + guidance = guidance, + batch_size = body.n, + ) + except Exception as exc: # noqa: BLE001 (single boundary, sanitized envelope) + # A RuntimeError with the model now unloaded means it was evicted/unloaded + # between the readiness check above and the call (a transient race): 503. + # Every other failure (CUDA OOM, a diffusers shape/device error, both also + # RuntimeError) is a real 500, and its raw message must not reach the client. + if isinstance(exc, RuntimeError) and not backend.is_loaded: + raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG) + logger.error("openai_images.generate_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Image generation failed.") + + created = int(time.time()) + want_b64 = body.response_format == "b64_json" + # Persist each image with its full recipe embedded, like /images/generate, so + # the response_format=url links resolve and the images show up in the gallery. + recipe = { + "prompt": body.prompt, + "negative_prompt": None, + "width": width, + "height": height, + "steps": steps, + "guidance": guidance, + # The batch shares one base seed, so restoring a batch_index>0 sibling needs the + # original batch_size to replay it (same as /images/generate); persist it. + "batch_size": body.n, + "model": result.get("repo_id"), + "created_at": float(created), + } + # The diffusers batch shares one seed; the native sd.cpp batch uses a distinct seed + # per image (returned in ``seeds``), so record each image's own seed, like + # /images/generate, or a native batch_index>0 image shows the wrong seed. + per_image_seeds = result.get("seeds") + + def _persist() -> list[ImageGenerationData]: + items: list[ImageGenerationData] = [] + for index, image in enumerate(result["images"]): + seed = ( + per_image_seeds[index] + if per_image_seeds and index < len(per_image_seeds) + else result["seed"] + ) + record = image_gallery.save(image, {**recipe, "batch_index": index, "seed": seed}) + if want_b64: + encoded = image_gallery.image_b64(record["id"]) + if encoded is None: # vanished between write and read — fail the call + raise RuntimeError("generated image could not be read back for encoding") + items.append(ImageGenerationData(b64_json = encoded)) + else: + items.append(ImageGenerationData(url = _absolute_image_url(request, record["url"]))) + return items + + try: + data = await asyncio.to_thread(_persist) + except Exception as exc: # noqa: BLE001 + logger.error("openai_images.persist_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Failed to save the generated image.") + + return ImageGenerationResponse(created = created, data = data) 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_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index aebbb40af6..a84ffcf24a 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -133,9 +133,13 @@ def _stub_torch_accelerate( def _good_ckpt(scheme = "fp8", base = "Tongyi-MAI/Z-Image-Turbo"): + meta = {"scheme": scheme, "base_model_id": base} + # fp8 checkpoints must record per-row granularity or the loader rejects them as stale. + if scheme == "fp8": + meta["fp8_granularity"] = "per_row" return { "format": PREQUANT_FORMAT, - "metadata": {"scheme": scheme, "base_model_id": base}, + "metadata": meta, "state_dict": {"weight": object()}, } @@ -220,6 +224,23 @@ def test_load_base_mismatch_is_none(monkeypatch, tmp_path): assert _load(monkeypatch, tmp_path, _good_ckpt(base = "other/model")) is None +def test_load_fp8_stale_per_tensor_is_rejected(monkeypatch, tmp_path): + # A pre-fix fp8 checkpoint has no fp8_granularity (old per-tensor layout); it must be + # rejected so the loader rebuilds instead of reproducing the noise failure. + stale = _good_ckpt(scheme = "fp8") + del stale["metadata"]["fp8_granularity"] + assert _load(monkeypatch, tmp_path, stale, scheme = "fp8") is None + # An explicit per-tensor granularity is likewise rejected. + per_tensor = _good_ckpt(scheme = "fp8") + per_tensor["metadata"]["fp8_granularity"] = "per_tensor" + assert _load(monkeypatch, tmp_path, per_tensor, scheme = "fp8") is None + + +def test_load_int8_ignores_fp8_granularity(monkeypatch, tmp_path): + # The granularity gate is fp8-only: an int8 checkpoint without it still loads. + assert _load(monkeypatch, tmp_path, _good_ckpt(scheme = "int8"), scheme = "int8") is not None + + def test_load_missing_base_metadata_is_none(monkeypatch, tmp_path): # A checkpoint whose keys happen to match a different base can load strict=True and then # render from the wrong weights, so a base was requested but none recorded must be refused. diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index 4953380412..60296bdfc5 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -415,6 +415,23 @@ def test_resolve_fast_accum(monkeypatch): assert tq._resolve_fast_accum(False) is False # forced off (e.g. on a consumer card) +def test_fp8_config_uses_per_row_granularity(): + """FP8 must use PerRow (per-token activation + per-channel weight) scaling. torchao's + default is per-TENSOR: on a DiT with extreme activation outliers (z-image's ~6.6e4) one + outlier forces a tensor-wide scale that pushes normal values below fp8 resolution and the + denoise collapses to noise. This is the regression guard for that fix (validated on B200: + per-tensor fp8 = noise, per-row fp8 = matches bf16).""" + torchao_quant = pytest.importorskip("torchao.quantization") + per_row = getattr(torchao_quant, "PerRow", None) + if per_row is None: + pytest.skip("torchao build without PerRow granularity") + cfg = tq._make_quant_config(TQ_FP8) + gran = getattr(cfg, "granularity", None) + assert gran is not None, "fp8 config must set an explicit granularity, not torchao's default" + grans = gran if isinstance(gran, (list, tuple)) else [gran] + assert grans and all(isinstance(g, per_row) for g in grans), f"expected all PerRow, got {gran}" + + def test_quantize_transformer_applies_and_marks(monkeypatch): monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_FP8) seen: dict = {} diff --git a/studio/backend/tests/test_openai_images_generations_route.py b/studio/backend/tests/test_openai_images_generations_route.py new file mode 100644 index 0000000000..f82586a173 --- /dev/null +++ b/studio/backend/tests/test_openai_images_generations_route.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""FastAPI round-trip tests for the OpenAI-compatible POST /v1/images/generations. + +The diffusion backend and image gallery are replaced with light fakes, so these +exercise the route wiring, OpenAI param mapping, validation, error envelopes, and +response shape without torch, diffusers, weights, or a GPU. The pure helpers +(`_parse_openai_image_size`, `default_generation_params`) are unit-tested directly. +""" + +from __future__ import annotations + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import core.inference.diffusion as diffusion_module +import core.inference.diffusion_engine_router as engine_router +import core.inference.image_gallery as gallery_module +from auth.authentication import get_current_subject +from core.inference.diffusion_families import default_generation_params +from routes.inference import router, _parse_openai_image_size +from utils.api_errors import install_api_error_handlers + + +# ── pure helpers ──────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "repo_id, expected", + [ + ("unsloth/Z-Image-Turbo-GGUF", (9, 0.0)), # turbo entry, before the z-image fallback + ("unsloth/Z-Image-GGUF", (20, 4.0)), + ("unsloth/FLUX.1-schnell-GGUF", (4, 0.0)), # schnell entry, before the flux.1 entry + ("black-forest-labs/FLUX.1-dev", (28, 3.5)), + ("unsloth/FLUX.2-klein-4B-GGUF", (4, 0.0)), + ("unsloth/Qwen-Image-2512-GGUF", (20, 4.0)), + ("some/unknown-model", (9, 0.0)), # fallback + ("", (9, 0.0)), + ], +) +def test_default_generation_params(repo_id, expected): + assert default_generation_params(repo_id) == expected + + +def test_default_generation_params_specificity_ordering(): + # The "-turbo" / "-schnell" entries must win over their broader siblings; a + # reorder that broke this would silently mis-default. + assert default_generation_params("x/Z-Image-Turbo") != default_generation_params("x/Z-Image") + assert default_generation_params("x/FLUX.1-schnell") != default_generation_params( + "x/FLUX.1-dev" + ) + + +def test_default_generation_params_falls_back_to_base_repo(): + # A local-path load: repo_id is a filesystem path that names no model, so the + # resolved base repo is what identifies it (and distinguishes dev from schnell). + assert default_generation_params("/models/my-ckpt", "black-forest-labs/FLUX.1-dev") == (28, 3.5) + assert default_generation_params("/models/my-ckpt", "black-forest-labs/FLUX.1-schnell") == ( + 4, + 0.0, + ) + assert default_generation_params("/models/my-ckpt", "Qwen/Qwen-Image") == (20, 4.0) + # repo_id wins when it already names the model; base repo is only a fallback. + assert default_generation_params("unsloth/Z-Image-Turbo-GGUF", "Tongyi-MAI/Z-Image") == (9, 0.0) + # Nothing identifiable -> fallback; None identifiers are skipped. + assert default_generation_params(None, None) == (9, 0.0) + assert default_generation_params("/models/x", None) == (9, 0.0) + + +@pytest.mark.parametrize( + "size, expected", + [ + ("auto", (1024, 1024)), + ("", (1024, 1024)), + ("AUTO", (1024, 1024)), + ("512x512", (512, 512)), + ("512x256", (512, 256)), + ("1792x1024", (1792, 1024)), # dall-e-3 named size: must pass the bounds + ("1024x1792", (1024, 1792)), + (" 256 x 256 ", (256, 256)), + ], +) +def test_parse_image_size_ok(size, expected): + assert _parse_openai_image_size(size) == expected + + +@pytest.mark.parametrize("size", ["abc", "100x100", "4096x4096", "300x300", "512", "x512", "0x0"]) +def test_parse_image_size_rejects(size): + with pytest.raises(ValueError): + _parse_openai_image_size(size) + + +# ── route round-trip ──────────────────────────────────────────────────── + + +class _FakeBackend: + def __init__( + self, + loaded = True, + repo_id = "unsloth/Z-Image-Turbo-GGUF", + base_repo = None, + generate_error = None, + unload_on_generate = False, + native_seeds = False, + ) -> None: + self._loaded = loaded + self._repo_id = repo_id + self._base_repo = base_repo + # Model the native sd.cpp engine, which returns a distinct seed per image. + self._native_seeds = native_seeds + # When set, generate() raises this; unload_on_generate flips is_loaded off + # first, to model the eviction/unload race vs an in-pipeline failure (OOM). + self._generate_error = generate_error + self._unload_on_generate = unload_on_generate + self.calls = [] + + @property + def is_loaded(self): + return self._loaded + + def status(self): + return { + "loaded": self._loaded, + "repo_id": self._repo_id if self._loaded else None, + "family": "z-image" if self._loaded else None, + "base_repo": self._base_repo if self._loaded else None, + "device": "cpu", + "dtype": "float32", + "cpu_offload": False, + } + + def generate( + self, + *, + prompt, + width, + height, + steps, + guidance, + batch_size = 1, + ): + if not self._loaded: + raise RuntimeError("No diffusion model is loaded.") + if self._generate_error is not None: + if self._unload_on_generate: + self._loaded = False + raise self._generate_error + self.calls.append( + dict( + prompt = prompt, + width = width, + height = height, + steps = steps, + guidance = guidance, + batch_size = batch_size, + ) + ) + out = { + "images": [object() for _ in range(batch_size)], + "seed": 4242, + "repo_id": self._repo_id, + } + if self._native_seeds: + out["seeds"] = [4242 + i for i in range(batch_size)] + return out + + +def _make_client(backend): + store = {} + + def _save(image, meta): + image_id = f"img{len(store)}" + record = {**meta, "id": image_id, "url": f"/api/inference/images/gallery/{image_id}/file"} + store[image_id] = record + return record + + app = FastAPI() + install_api_error_handlers(app) + app.include_router(router, prefix = "/v1") + app.dependency_overrides[get_current_subject] = lambda: "test-user" + return TestClient(app), store, _save + + +@pytest.fixture +def client(monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None) + cli.backend = backend # type: ignore[attr-defined] + return cli + + +def _post(client, body): + return client.post("/v1/images/generations", json = body) + + +def test_url_response_shape(client): + resp = _post(client, {"prompt": "a sloth", "size": "256x256"}) + assert resp.status_code == 200 + body = resp.json() + assert set(body.keys()) == {"created", "data"} + assert isinstance(body["created"], int) and body["created"] > 0 + assert len(body["data"]) == 1 + item = body["data"][0] + assert "url" in item and "b64_json" not in item # exclude_none drops the unused key + assert item["url"].endswith("/file") + # Z-Image-Turbo defaults (9 steps, 0 guidance) flow into the backend call. + assert client.backend.calls[0] == dict( + prompt = "a sloth", width = 256, height = 256, steps = 9, guidance = 0.0, batch_size = 1 + ) + + +def test_b64_response_shape(client): + resp = _post(client, {"prompt": "a sloth", "size": "256x256", "response_format": "b64_json"}) + assert resp.status_code == 200 + item = resp.json()["data"][0] + assert "b64_json" in item and "url" not in item + assert item["b64_json"] == "QUJD" + + +def test_local_load_uses_base_repo_for_defaults(monkeypatch): + # repo_id is a local path that names no model; base_repo identifies FLUX.1-dev, + # so the route must pick 28 steps / 3.5 guidance, not the 9/0 fallback. + backend = _FakeBackend(repo_id = "/models/my-flux", base_repo = "black-forest-labs/FLUX.1-dev") + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"}) + assert resp.status_code == 200 + assert backend.calls[0]["steps"] == 28 and backend.calls[0]["guidance"] == 3.5 + + +def test_pipeline_runtime_error_is_sanitized_500(monkeypatch): + # A RuntimeError raised inside the pipeline while the model stays loaded (e.g. + # CUDA OOM, a RuntimeError subclass) must be a sanitized 500, not a 503 that + # echoes the raw exception text. + oom = RuntimeError("CUDA out of memory. Tried to allocate 20.00 GiB (GPU 0; 47.5 GiB total)") + backend = _FakeBackend(generate_error = oom) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"}) + assert resp.status_code == 500 + assert resp.json()["error"]["message"] == "Image generation failed." + assert "CUDA" not in resp.text # raw exception text must not leak + + +def test_unload_race_returns_503(monkeypatch): + # The model is evicted/unloaded between the readiness check and the call: the + # RuntimeError with is_loaded now False is the one case that maps to 503. + backend = _FakeBackend( + generate_error = RuntimeError("No diffusion model is loaded."), + unload_on_generate = True, + ) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"}) + assert resp.status_code == 503 + err = resp.json()["error"] + assert err["type"] == "api_error" + # The 503 carries the fixed sanitized message, not the raw exception text. + assert err["message"] == "No image model loaded. Load an image model first." + + +def test_non_runtime_pipeline_error_is_500(monkeypatch): + # A non-RuntimeError from the pipeline (the model stays loaded) must not route + # to the 503 branch (which is gated on isinstance RuntimeError) -> sanitized 500. + backend = _FakeBackend(generate_error = ValueError("bad tensor shape")) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"}) + assert resp.status_code == 500 + assert "shape" not in resp.text + + +def test_n_maps_to_batch(client): + resp = _post(client, {"prompt": "p", "size": "256x256", "n": 3}) + assert resp.status_code == 200 + assert len(resp.json()["data"]) == 3 + assert client.backend.calls[0]["batch_size"] == 3 + + +def test_batch_persists_batch_size(monkeypatch): + # n>1 must persist batch_size in each gallery record so the Studio restore path + # can replay a batch_index>0 sibling (which shares the batch's single seed). + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256", "n": 3}) + assert resp.status_code == 200 + records = sorted(store.values(), key = lambda r: r["batch_index"]) + assert [r["batch_index"] for r in records] == [0, 1, 2] + assert all(r["batch_size"] == 3 for r in records) + + +def test_uses_active_engine_not_diffusers_singleton(monkeypatch): + # On a no-GPU host the loaded model lives behind the native sd_cpp engine, not the + # diffusers singleton. The route must query get_active_diffusion_engine (like + # /images/generate) or it 503s a model that is loaded and usable. + active = _FakeBackend(loaded = True) # the active (e.g. sd_cpp) engine, loaded + idle_diffusers = _FakeBackend(loaded = False) # diffusers singleton, empty + monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: active) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: idle_diffusers) + cli, store, _save = _make_client(active) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"}) + assert resp.status_code == 200 + assert len(active.calls) == 1 # the active engine did the work, not the idle singleton + + +def test_native_batch_persists_per_image_seed(monkeypatch): + # The native sd.cpp engine returns a distinct seed per image (base+index) in + # "seeds"; each gallery record must store its own seed (like /images/generate), + # not the shared base, or a restored batch_index>0 image shows the wrong seed. + backend = _FakeBackend(native_seeds = True) + monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256", "n": 3}) + assert resp.status_code == 200 + records = sorted(store.values(), key = lambda r: r["batch_index"]) + assert [r["seed"] for r in records] == [4242, 4243, 4244] + + +def test_null_fields_coalesce_to_defaults(client): + # OpenAI marks n/size/response_format nullable-with-default: null -> default. + resp = _post(client, {"prompt": "p", "n": None, "size": None, "response_format": None}) + assert resp.status_code == 200 + assert len(resp.json()["data"]) == 1 + assert "url" in resp.json()["data"][0] + assert client.backend.calls[0]["width"] == 1024 # size null -> auto -> 1024 + + +@pytest.mark.parametrize( + "body, param", + [ + ({"size": "256x256"}, "prompt"), # missing prompt + ({"prompt": "", "size": "256x256"}, "prompt"), # empty prompt + ({"prompt": "p", "size": "300x300"}, "size"), # not multiple of 16 + ({"prompt": "p", "size": "abc"}, "size"), # unparseable + ({"prompt": "p", "stream": True}, "stream"), # streaming unsupported + ], +) +def test_validation_400_with_param(client, body, param): + resp = _post(client, body) + assert resp.status_code == 400 + err = resp.json()["error"] + assert err["type"] == "invalid_request_error" + assert err["param"] == param + for k in ("message", "code"): + assert k in err + + +@pytest.mark.parametrize("n", [0, 11, -1]) +def test_n_out_of_range_400(client, n): + resp = _post(client, {"prompt": "p", "size": "256x256", "n": n}) + assert resp.status_code == 400 + assert resp.json()["error"]["type"] == "invalid_request_error" + + +def test_no_model_loaded_503(monkeypatch): + backend = _FakeBackend(loaded = False) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + cli, store, _save = _make_client(backend) + monkeypatch.setattr(gallery_module, "save", _save) + resp = cli.post("/v1/images/generations", json = {"prompt": "p"}) + assert resp.status_code == 503 + # 503 still wears the OpenAI envelope (api_error) on the /v1 surface. + assert resp.json()["error"]["type"] == "api_error" + + +def test_auth_required(): + backend = _FakeBackend() + app = FastAPI() + install_api_error_handlers(app) + app.include_router(router, prefix = "/v1") + # No dependency override: the real auth dependency runs and rejects. + resp = TestClient(app).post("/v1/images/generations", json = {"prompt": "p"}) + assert resp.status_code in (401, 403) 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..0b52ea6e14 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -77,10 +77,70 @@ 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 + self.lora_dir = None # set by a test to the server's --lora-model-dir scratch dir + + def is_alive(self): + return self.alive and not self.stopped + + def start( + self, + files, + *, + vae_format = None, + offload = None, + native_speed = None, + threads = None, + ): + self.started = dict( + files = files, + vae_format = vae_format, + offload = offload, + native_speed = native_speed, + threads = threads, + ) + + def img_gen( + self, + payload, + *, + on_step = None, + cancel_event = None, + total_timeout = None, + ): + import io as _io + + self.payloads.append(payload) + self.timeouts.append(total_timeout) + if on_step is not None: + steps = payload.get("sample_params", {}).get("sample_steps", 0) + on_step(f" {steps}/{steps}") + n = int(payload.get("batch_count", 1)) + blobs = [] + for i in range(n): + buf = _io.BytesIO() + Image.new("RGB", (1, 1), (i, i, i)).save(buf, format = "PNG") + blobs.append(buf.getvalue()) + return blobs + + def stop(self): + self.stopped = True + + # ── asset resolution ────────────────────────────────────────────────────────── @@ -321,6 +381,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. @@ -355,3 +669,108 @@ def test_run_load_redacts_paths_in_progress_error(monkeypatch): with npl._REDACTION_LOCK: if secret_root in npl._NATIVE_PATH_REDACTIONS: npl._NATIVE_PATH_REDACTIONS.remove(secret_root) + + +# ── LoRA (native engine) ──────────────────────────────────────────────────────── + + +def _fake_materialize(resolved, dest): + """Stand-in for diffusion_lora.materialize_native_dir: write a stub file per adapter + into ``dest`` and return the resolved list pointing at the written paths (mirroring the + real helper's contract without touching the Hub / real weights).""" + from pathlib import Path as _P + + from core.inference import diffusion_lora as dl + + dest.mkdir(parents = True, exist_ok = True) + out = [] + for r in resolved: + p = _P(dest) / f"{r.alias}.safetensors" + p.write_bytes(b"stub") + out.append(dl.ResolvedLora(r.id, r.alias, str(p), r.fmt, r.weight)) + return out + + +def _patch_lora(monkeypatch, resolved, supported = True): + from core.inference import diffusion_lora as dl + + monkeypatch.setattr(dl, "supports_lora", lambda **k: supported) + monkeypatch.setattr(dl, "resolve_specs", lambda specs, **k: list(resolved)) + monkeypatch.setattr(dl, "materialize_native_dir", _fake_materialize) + + +def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch): + # One-shot sd-cli LoRA: adapters materialized into a --lora-model-dir and selected with + # tags injected into the prompt (the real inject_prompt_tags runs here). + from core.inference import diffusion_lora as dl + + eng = _FakeEngine() + b = _loaded_backend(engine = eng) # mode = "oneshot" + _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)]) + b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)]) + _, params, _, _ = eng.calls[0] + assert params.lora_dir is not None and params.lora_apply_mode == "auto" + assert "" in params.prompt + + +def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tmp_path): + # Server-mode LoRA rides the structured `lora` request field (the sdcpp API ignores + # prompt tags): adapters staged into the server's --lora-model-dir, referenced + # by their path relative to it + the validated multiplier. + from pathlib import Path as _P + + from core.inference import diffusion_lora as dl + + b = SdCppDiffusionBackend() + servers: list = [] + _run_server_load(monkeypatch, b, servers) + servers[0].lora_dir = str(tmp_path) + _patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)]) + b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)]) + payload = servers[0].payloads[0] + assert "lora" in payload and len(payload["lora"]) == 1 + assert payload["lora"][0]["multiplier"] == 0.7 + assert payload["lora"][0]["path"].endswith("myalias.safetensors") + assert " 400) rather + # than TypeError on an unexpected kwarg (-> opaque 500). + b = _loaded_backend(engine = _FakeEngine()) + with pytest.raises(ValueError, match = "ControlNet is not yet supported on the native"): + b.generate(prompt = "x", steps = 4, seed = 1, controlnet = ("id", "img", "canny", 1.0, 0.0, 1.0)) + + +def test_generate_rejects_image_conditioned_on_native_engine(): + # img2img / inpaint / reference / upscale are likewise diffusers-only; a direct API call + # with an init image on the native engine gets a clean ValueError, not a silent txt2img. + b = _loaded_backend(engine = _FakeEngine()) + with pytest.raises(ValueError, match = "not yet supported on the native"): + b.generate(prompt = "x", steps = 4, seed = 1, init_image = "data:image/png;base64,AAAA") + + +def test_status_native_reports_supports_controlnet_false(): + b = _loaded_backend() + assert b.status()["supports_controlnet"] is False 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_install.py b/studio/backend/tests/test_sd_cpp_install.py index 4e06a3cdba..cc5051e784 100644 --- a/studio/backend/tests/test_sd_cpp_install.py +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -148,11 +148,12 @@ def test_pinned_tag_default_and_override(monkeypatch): assert _pinned_tag() is None -def test_repo_default_and_mirror_override(monkeypatch): +def test_repo_default_and_override(monkeypatch): monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False) - assert _repo() == DEFAULT_REPO == "leejet/stable-diffusion.cpp" - monkeypatch.setenv("UNSLOTH_SD_CPP_REPO", "unslothai/stable-diffusion.cpp") - assert _repo() == "unslothai/stable-diffusion.cpp" + # Default is the Unsloth mirror; the env override can point back to leejet upstream. + assert _repo() == DEFAULT_REPO == "unslothai/stable-diffusion.cpp" + monkeypatch.setenv("UNSLOTH_SD_CPP_REPO", "leejet/stable-diffusion.cpp") + assert _repo() == "leejet/stable-diffusion.cpp" # ── sha256 integrity check ────────────────────────────────────────────────── @@ -306,3 +307,257 @@ def test_find_sd_cpp_binary_honors_studio_home(tmp_path, monkeypatch): binary.parent.mkdir(parents = True) binary.write_bytes(b"x") assert eng.find_sd_cpp_binary() == str(binary) + + +# ── Unsloth mirror: default source + the CPU/Apple asset set it publishes ───── + +_TAG = "master-741-484baa4" +# Exactly what unslothai/stable-diffusion.cpp's CI publishes (CPU + Apple only; GPU +# hosts run diffusers and never reach the native installer). +_MIRROR_ASSETS = [ + f"sd-{_TAG}-bin-Darwin-macOS-arm64.zip", + f"sd-{_TAG}-bin-Darwin-macOS-x86_64.zip", + f"sd-{_TAG}-bin-Linux-Ubuntu-22.04-x86_64.zip", + f"sd-{_TAG}-bin-Linux-Ubuntu-24.04-aarch64.zip", + f"sd-{_TAG}-bin-win-cpu-x64.zip", +] + + +def _mresolve( + system, + machine, + accelerator = "auto", +): + return resolve_release_asset( + _MIRROR_ASSETS, system = system, machine = machine, accelerator = accelerator + ) + + +def test_default_repo_is_the_unsloth_mirror(): + assert sdmod.DEFAULT_REPO == "unslothai/stable-diffusion.cpp" + assert sdmod.UPSTREAM_FALLBACK_REPO == "leejet/stable-diffusion.cpp" + + +def test_mirror_matrix_resolves_every_cpu_apple_host(): + assert _mresolve("Darwin", "arm64") == f"sd-{_TAG}-bin-Darwin-macOS-arm64.zip" + assert _mresolve("Darwin", "x86_64") == f"sd-{_TAG}-bin-Darwin-macOS-x86_64.zip" + # WSL reports as Linux x86_64, so this also covers WSL. + assert _mresolve("Linux", "x86_64") == f"sd-{_TAG}-bin-Linux-Ubuntu-22.04-x86_64.zip" + assert _mresolve("Linux", "aarch64") == f"sd-{_TAG}-bin-Linux-Ubuntu-24.04-aarch64.zip" + assert _mresolve("Windows", "AMD64") == f"sd-{_TAG}-bin-win-cpu-x64.zip" + assert _mresolve("Windows", "AMD64", "cpu") == f"sd-{_TAG}-bin-win-cpu-x64.zip" + + +# ── mirror -> upstream fallback in install() ───────────────────────────────── + + +def _stub_two_repos(monkeypatch, *, mirror_serves, upstream_serves, zip_bytes, digest): + """Stub _fetch_release to serve (or 404) per repo, so install()'s mirror->upstream + fallback can be exercised without network. Host pinned to Linux x86_64.""" + monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False) + monkeypatch.delenv("UNSLOTH_SD_CPP_TAG", raising = False) + mirror_name = f"sd-{_TAG}-bin-Linux-Ubuntu-22.04-x86_64.zip" + upstream_name = "sd-master-741-484baa4-bin-Linux-Ubuntu-24.04-x86_64.zip" + + def _rel(name): + return { + "tag_name": _TAG, + "assets": [ + { + "name": name, + "browser_download_url": f"https://example.invalid/{name}", + "digest": digest, + } + ], + } + + def fake_fetch( + tag = None, + *, + repo = None, + token = None, + timeout = 30.0, + allow_latest = True, + ): + r = repo or sdmod.DEFAULT_REPO + if r == sdmod.DEFAULT_REPO: + if mirror_serves: + return _rel(mirror_name) + raise urllib.error.HTTPError(f"https://api/{r}", 404, "not found", None, None) + if upstream_serves: + return _rel(upstream_name) + raise urllib.error.HTTPError(f"https://api/{r}", 404, "not found", None, None) + + monkeypatch.setattr(sdmod, "_fetch_release", fake_fetch) + monkeypatch.setattr(sdmod, "_download", lambda url, dest, **k: dest.write_bytes(zip_bytes)) + monkeypatch.setattr(sdmod.platform, "system", lambda: "Linux") + monkeypatch.setattr(sdmod.platform, "machine", lambda: "x86_64") + return mirror_name, upstream_name + + +def test_install_uses_mirror_when_available(tmp_path, monkeypatch, capsys): + zb = _zip_with_sd_cli() + _stub_two_repos( + monkeypatch, + mirror_serves = True, + upstream_serves = True, + zip_bytes = zb, + digest = "sha256:" + hashlib.sha256(zb).hexdigest(), + ) + sd_cli = install(install_dir = tmp_path) + assert sd_cli.name == "sd-cli" and sd_cli.is_file() + assert "unslothai/stable-diffusion.cpp" in capsys.readouterr().out + + +def test_install_falls_back_to_upstream_when_mirror_missing(tmp_path, monkeypatch, capsys): + zb = _zip_with_sd_cli() + _stub_two_repos( + monkeypatch, + mirror_serves = False, + upstream_serves = True, + zip_bytes = zb, + digest = "sha256:" + hashlib.sha256(zb).hexdigest(), + ) + sd_cli = install(install_dir = tmp_path) + assert sd_cli.name == "sd-cli" and sd_cli.is_file() + captured = capsys.readouterr() + # The repo-fallback diagnostic goes to stderr so --print-asset's stdout stays a single + # asset line; the "source ..." install note still prints to stdout. + assert "falling back to leejet/stable-diffusion.cpp" in captured.err + assert "source leejet/stable-diffusion.cpp" in captured.out + + +def test_install_errors_when_neither_source_serves(tmp_path, monkeypatch): + _stub_two_repos( + monkeypatch, mirror_serves = False, upstream_serves = False, zip_bytes = b"", digest = "" + ) + with pytest.raises(RuntimeError, match = "No prebuilt sd-cli"): + install(install_dir = tmp_path) + + +# ── explicit GPU accelerator on a CPU-only mirror -> no CPU substitution ────── + + +def test_mirror_windows_gpu_accel_is_no_match_not_cpu(): + # The mirror ships only a CPU win zip. An explicit --accelerator cuda/vulkan/rocm must + # NOT silently resolve to that CPU build; it returns None so install() falls back to + # upstream, which does build the accelerated asset. + for accel in ("cuda", "vulkan", "rocm"): + assert _mresolve("Windows", "AMD64", accel) is None + # auto / cpu still resolve to the CPU build. + assert _mresolve("Windows", "AMD64", "cpu") == f"sd-{_TAG}-bin-win-cpu-x64.zip" + + +def test_mirror_linux_gpu_accel_is_no_match_not_cpu(): + for accel in ("cuda", "vulkan", "rocm"): + assert _mresolve("Linux", "x86_64", accel) is None + assert _mresolve("Linux", "x86_64", "cpu") == f"sd-{_TAG}-bin-Linux-Ubuntu-22.04-x86_64.zip" + + +def test_upstream_full_matrix_still_resolves_gpu_accel(): + # Regression guard: the "explicit GPU accel -> None on miss" change must not break + # the upstream matrix, which does publish the accelerated builds. + assert _resolve("Windows", "AMD64", "cuda") == "sd-master-8caa3f9-bin-win-cuda12-x64.zip" + assert _resolve("Linux", "x86_64", "vulkan").endswith("x86_64-vulkan.zip") + assert "rocm" in _resolve("Linux", "x86_64", "rocm") + + +# ── explicit repo override suppresses the upstream fallback ─────────────────── + + +def test_explicit_repo_override_equal_to_default_suppresses_fallback(tmp_path, monkeypatch): + # A user who pins UNSLOTH_SD_CPP_REPO (even to the default value) must get exactly that + # repo -- no surprise leejet substitution -- so a missing release errors instead. + _stub_two_repos( + monkeypatch, mirror_serves = False, upstream_serves = True, zip_bytes = b"", digest = "" + ) + monkeypatch.setenv("UNSLOTH_SD_CPP_REPO", sdmod.DEFAULT_REPO) + with pytest.raises(RuntimeError, match = "No prebuilt sd-cli"): + install(install_dir = tmp_path) + + +# ── pinned tag missing on mirror -> pinned upstream before mirror latest ────── + + +def test_pinned_tag_prefers_upstream_pin_over_mirror_latest(tmp_path, monkeypatch, capsys): + # The mirror lacks the pinned tag but has a newer latest; upstream has the pin. The + # pinned upstream build must win over the unpinned mirror-latest (reproducibility). + monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False) + monkeypatch.setenv("UNSLOTH_SD_CPP_TAG", "master-999-pinned") + zb = _zip_with_sd_cli() + digest = "sha256:" + hashlib.sha256(zb).hexdigest() + + def _rel(name, tag): + return { + "tag_name": tag, + "assets": [ + { + "name": name, + "browser_download_url": f"https://example.invalid/{name}", + "digest": digest, + } + ], + } + + mirror_latest = "sd-master-000-latest-bin-Linux-Ubuntu-22.04-x86_64.zip" + upstream_pinned = "sd-master-999-pinned-bin-Linux-Ubuntu-24.04-x86_64.zip" + + def fake_fetch( + tag = None, + *, + repo = None, + token = None, + timeout = 30.0, + allow_latest = True, + ): + r = repo or sdmod.DEFAULT_REPO + if r == sdmod.DEFAULT_REPO: + if tag == "master-999-pinned": # mirror lacks the pin + if not allow_latest: + return None + return _rel(mirror_latest, "master-000-latest") + return _rel(mirror_latest, "master-000-latest") + # upstream HAS the pin + if tag == "master-999-pinned": + return _rel(upstream_pinned, "master-999-pinned") + return _rel(upstream_pinned, "master-999-pinned") + + monkeypatch.setattr(sdmod, "_fetch_release", fake_fetch) + monkeypatch.setattr(sdmod, "_download", lambda url, dest, **k: dest.write_bytes(zb)) + monkeypatch.setattr(sdmod.platform, "system", lambda: "Linux") + monkeypatch.setattr(sdmod.platform, "machine", lambda: "x86_64") + + install(install_dir = tmp_path) + out = capsys.readouterr().out + assert "source leejet/stable-diffusion.cpp release master-999-pinned" in out + + +# ── --print-asset routes through the primary/upstream fallback ──────────────── + + +def test_print_asset_uses_upstream_fallback(monkeypatch, capsys): + # A host the mirror does not build (Linux Vulkan) must print the upstream asset that a + # real install would fetch, not "no matching prebuilt". + monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False) + monkeypatch.delenv("UNSLOTH_SD_CPP_TAG", raising = False) + + def fake_fetch( + tag = None, + *, + repo = None, + token = None, + timeout = 30.0, + allow_latest = True, + ): + r = repo or sdmod.DEFAULT_REPO + if r == sdmod.DEFAULT_REPO: + return {"tag_name": _TAG, "assets": [{"name": n} for n in _MIRROR_ASSETS]} + return {"tag_name": "master-8caa3f9", "assets": [{"name": n} for n in _ASSETS]} + + monkeypatch.setattr(sdmod, "_fetch_release", fake_fetch) + monkeypatch.setattr(sdmod.platform, "system", lambda: "Linux") + monkeypatch.setattr(sdmod.platform, "machine", lambda: "x86_64") + rc = sdmod.main(["--print-asset", "--accelerator", "vulkan"]) + out = capsys.readouterr().out + assert rc == 0 + assert "vulkan" in out and "no matching prebuilt" not in out 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/frontend/src/features/images/diffusion-train-dialog.tsx b/studio/frontend/src/features/images/diffusion-train-dialog.tsx new file mode 100644 index 0000000000..25b260a03e --- /dev/null +++ b/studio/frontend/src/features/images/diffusion-train-dialog.tsx @@ -0,0 +1,510 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { useCallback, useEffect, useRef, useState } from "react"; + +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store"; +import { toast } from "@/lib/toast"; + +import { + type DiffusionTrainingInfo, + type DiffusionTrainingStatus, + getDiffusionTrainingInfo, + getDiffusionTrainingStatus, + startDiffusionTraining, + stopDiffusionTraining, + uploadDiffusionDataset, +} from "./api"; + +// The two official SDXL bases the backend allowlists for non-GGUF loads. Everything the +// dropdown offers is trainable; "custom" is the escape hatch for local SDXL checkpoints. +const SDXL_BASES: Array<{ id: string; label: string }> = [ + { id: "stabilityai/stable-diffusion-xl-base-1.0", label: "SDXL Base 1.0 (best quality)" }, + { id: "stabilityai/sdxl-turbo", label: "SDXL Turbo (fast, good for quick tests)" }, +]; +const CUSTOM_BASE = "__custom__"; +const UPLOAD_DATASET = "__upload__"; +const DATASET_FILE_ACCEPT = ".png,.jpg,.jpeg,.webp,.bmp,.txt,.caption,.jsonl"; + +const selectClass = + "h-8 w-full rounded-md border border-input bg-background px-2 text-xs"; + +// A self-contained "Train an SDXL LoRA" dialog. It posts to /api/train/diffusion/start +// and polls /status while open, so it never blocks the page and works whether or not a +// model is loaded for generation. Only SDXL is trainable today; the backend refuses +// known non-SDXL picks instantly, and the base-model dropdown keeps users on safe picks. +export function DiffusionTrainDialog({ + open, + onOpenChange, + defaultBaseModel, + onTrainingComplete, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; + defaultBaseModel?: string; + // Called once when a run finishes so the page can rescan the LoRA picker. + onTrainingComplete?: () => void; +}) { + const [baseChoice, setBaseChoice] = useState(SDXL_BASES[0].id); + const [customBase, setCustomBase] = useState(""); + const [info, setInfo] = useState(null); + const [dataset, setDataset] = useState(UPLOAD_DATASET); + const [uploadName, setUploadName] = useState("my-images"); + const [uploading, setUploading] = useState(false); + const fileInputRef = useRef(null); + const [outputDir, setOutputDir] = useState(""); + const [instancePrompt, setInstancePrompt] = useState(""); + const [showAdvanced, setShowAdvanced] = useState(false); + const [steps, setSteps] = useState(500); + const [learningRate, setLearningRate] = useState(0.0001); + const [rank, setRank] = useState(16); + const [resolution, setResolution] = useState(1024); + const [batchSize, setBatchSize] = useState(1); + const [precision, setPrecision] = useState<"bf16" | "fp16" | "no">("bf16"); + const [starting, setStarting] = useState(false); + const [status, setStatus] = useState(null); + + // The dialog stays mounted (ImagesPage is keep-alive), so seed per-open state here: + // the base-model choice from the currently loaded SDXL pipeline (when there is one), + // and the dataset list from the backend. + const refreshInfo = useCallback(async (): Promise => { + try { + const i = await getDiffusionTrainingInfo(); + setInfo(i); + return i; + } catch { + return null; // older backends: keep the upload-only flow usable + } + }, []); + + useEffect(() => { + if (!open) return; + if (defaultBaseModel) { + setBaseChoice(defaultBaseModel); + } + void refreshInfo().then((i) => { + // Preselect the only dataset, or the freshest-looking state: with no datasets + // yet, the picker sits on "Upload new images". + setDataset((cur) => { + if (cur !== UPLOAD_DATASET && i?.datasets.some((d) => d.name === cur)) return cur; + return i && i.datasets.length > 0 ? i.datasets[0].name : UPLOAD_DATASET; + }); + }); + }, [open, defaultBaseModel, refreshInfo]); + + const poll = useCallback(async () => { + try { + setStatus(await getDiffusionTrainingStatus()); + } catch { + // Best-effort; a failed poll should not surface an error while the dialog is open. + } + }, []); + + // Poll status only while the dialog is open. + useEffect(() => { + if (!open) return; + void poll(); + const id = window.setInterval(() => void poll(), 1500); + return () => window.clearInterval(id); + }, [open, poll]); + + const active = Boolean(status?.active) || status?.status === "running"; + const completed = status?.status === "completed"; + const pct = + status && status.total_steps > 0 + ? Math.min(100, Math.round((status.step / status.total_steps) * 100)) + : 0; + + // Notify the parent exactly once when a run reaches "completed", so it can rescan the + // LoRA picker (a LoRA trained while a model is loaded is otherwise invisible until a + // model swap re-runs the discovery effect). + const [notifiedComplete, setNotifiedComplete] = useState(false); + useEffect(() => { + if (status?.status === "completed" && !notifiedComplete) { + setNotifiedComplete(true); + onTrainingComplete?.(); + } else if (status?.status === "running" && notifiedComplete) { + setNotifiedComplete(false); // arm again for the next run + } + }, [status?.status, notifiedComplete, onTrainingComplete]); + + const selectedDataset = + dataset !== UPLOAD_DATASET ? info?.datasets.find((d) => d.name === dataset) : undefined; + + const onUpload = useCallback(async () => { + const files = Array.from(fileInputRef.current?.files ?? []); + if (files.length === 0) { + toast.error("Choose the images to upload first."); + return; + } + const name = uploadName.trim(); + if (!name) { + toast.error("Give the dataset a folder name, e.g. my-style-photos."); + return; + } + setUploading(true); + try { + const res = await uploadDiffusionDataset(name, files); + toast.success( + `Uploaded ${res.uploaded} file${res.uploaded === 1 ? "" : "s"} - ` + + `"${res.name}" now has ${res.image_count} images`, + ); + if (fileInputRef.current) fileInputRef.current.value = ""; + await refreshInfo(); + setDataset(res.name); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Upload failed"); + } finally { + setUploading(false); + } + }, [uploadName, refreshInfo]); + + const onStart = useCallback(async () => { + const baseModel = (baseChoice === CUSTOM_BASE ? customBase : baseChoice).trim(); + if (!baseModel) { + toast.error("Pick a base model (or fill in the custom repo/path)."); + return; + } + if (dataset === UPLOAD_DATASET) { + toast.error("Upload your training images first (or pick an existing dataset)."); + return; + } + if (!outputDir.trim()) { + toast.error("Name the adapter (this becomes its folder under Studio outputs)."); + return; + } + if (selectedDataset && selectedDataset.caption_count === 0 && !instancePrompt.trim()) { + toast.error( + "These images have no captions - add a trigger prompt so the trainer knows " + + "what to learn (it becomes the caption for every image).", + ); + return; + } + // Mirror the backend's numeric validation so obvious mistakes are caught before the + // request (the backend returns 400 for these; catching here gives a clearer message). + if (steps < 1) return toast.error("Steps must be at least 1."); + if (rank < 1) return toast.error("LoRA rank must be at least 1."); + if (resolution < 64 || resolution % 8 !== 0) { + return toast.error("Resolution must be a multiple of 8 and at least 64."); + } + if (batchSize < 1) return toast.error("Batch size must be at least 1."); + if (learningRate <= 0) return toast.error("Learning rate must be greater than 0."); + setStarting(true); + try { + await startDiffusionTraining({ + base_model: baseModel, + data_dir: dataset, + output_dir: outputDir.trim(), + instance_prompt: instancePrompt.trim() || undefined, + resolution, + train_steps: steps, + learning_rate: learningRate, + train_batch_size: batchSize, + lora_rank: rank, + mixed_precision: precision, + // Forward the saved Hub token so a gated/private SDXL base can be trained (the + // image load flow already sends it, so a model you can load, you can also train). + hf_token: hfApiToken(getHfToken()) || undefined, + }); + toast.success("Training started"); + void poll(); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Failed to start training"); + } finally { + setStarting(false); + } + }, [ + baseChoice, + customBase, + dataset, + selectedDataset, + outputDir, + instancePrompt, + resolution, + steps, + learningRate, + batchSize, + rank, + precision, + poll, + ]); + + const onStop = useCallback(async () => { + try { + await stopDiffusionTraining(); + toast.success("Stop requested; finishing the current step."); + void poll(); + } catch (e) { + toast.error(e instanceof Error ? e.message : "Failed to stop training"); + } + }, [poll]); + + return ( + + + + Train an SDXL LoRA + + Teach SDXL a style, character, or subject from your own images. The finished + adapter shows up in this page's LoRA picker. Only SDXL can be trained for + now - FLUX, Qwen-Image and Z-Image load LoRAs but can't train them yet. + + + +
+ {/* 1. Base model: a constrained dropdown instead of free text, so the "SDXL + only" rule is embodied by the control. Custom stays available for local + SDXL checkpoints; a known non-SDXL pick is refused instantly by the API. */} +
+ + + {baseChoice === CUSTOM_BASE && ( + setCustomBase(e.target.value)} + className="h-8 text-xs" + /> + )} +
+ + {/* 2. Training images: pick an existing dataset folder or upload straight from + the browser - no shell access or Studio-home knowledge needed. */} +
+ + + {dataset === UPLOAD_DATASET && ( +
+ setUploadName(e.target.value)} + className="h-8 text-xs" + aria-label="New dataset name" + /> +
+ + +
+

+ 10-50 images work well. Optional captions: a .txt per image (same + filename) or a metadata.jsonl; without them the trigger prompt below + captions every image. You can upload more into the same name later. +

+
+ )} + {selectedDataset && selectedDataset.caption_count === 0 && ( +

+ No caption files in this dataset - the trigger prompt below will be used + as the caption for every image. +

+ )} +
+ + {/* 3. What to call the result + how to trigger it. */} +
+ + setOutputDir(e.target.value)} + className="h-8 text-xs" + /> +
+
+ + setInstancePrompt(e.target.value)} + className="h-8 text-xs" + /> +
+ + {/* 4. Hyperparameters, collapsed: the defaults suit a first run, and hiding + them keeps the primary flow at three decisions. */} + + {showAdvanced && ( + <> +
+
+ + setSteps(Number(e.target.value) || 1)} + className="h-8 text-xs" + /> +
+
+ + setRank(Number(e.target.value) || 1)} + className="h-8 text-xs" + /> +
+
+ + setResolution(Number(e.target.value) || 1024)} + className="h-8 text-xs" + /> +
+
+ + setBatchSize(Number(e.target.value) || 1)} + className="h-8 text-xs" + /> +
+
+
+
+ + setLearningRate(Number(e.target.value) || 0.0001)} + className="h-8 text-xs" + /> +
+
+ + +
+
+ + )} + + {status && status.status !== "idle" && ( +
+
+ {status.status} + + {status.total_steps > 0 ? `${status.step}/${status.total_steps}` : ""} + +
+
+
+
+
+ {completed + ? "Adapter ready - find it in the LoRA picker on this page." + : status.message} + {status.loss != null && !completed && <> · loss {status.loss.toFixed(4)}} + {status.lora_path && ( +
Saved: {status.lora_path}
+ )} +
+
+ )} +
+ + + {active ? ( + + ) : completed ? ( + <> + + + + ) : ( + + )} + + +
+ ); +} diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index bff5f93e9d..481675ab32 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -2006,7 +2006,6 @@ export function ImagesPage({ active = true }: { active?: boolean }) { )} - {/* Train mode: the full-page training workspace. Kept unmounted in Create mode so its polling stops; Create's own state (gallery, model, workflow) is untouched. */} {pageMode === "train" ? ( diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py index 639e3a894d..52229d5f60 100644 --- a/studio/install_sd_cpp_prebuilt.py +++ b/studio/install_sd_cpp_prebuilt.py @@ -38,14 +38,18 @@ import zipfile from pathlib import Path from typing import Optional, Sequence -# Default upstream source. Overridable with UNSLOTH_SD_CPP_REPO so a pinned unslothai -# mirror (built the same way as unslothai/llama.cpp's prebuilts) can be used without a -# code change once it exists; otherwise this falls back to leejet upstream. -DEFAULT_REPO = "leejet/stable-diffusion.cpp" +# Default source: the Unsloth-built mirror, whose CPU/Apple prebuilts are compiled and +# published by unslothai/stable-diffusion.cpp (the same way unslothai/llama.cpp ships its +# prebuilts). Override with UNSLOTH_SD_CPP_REPO to point elsewhere (e.g. back to leejet). +# GPU hosts never reach here -- they run diffusers -- so only CPU/Apple assets are needed. +DEFAULT_REPO = "unslothai/stable-diffusion.cpp" +# Upstream we fall back to if the mirror can't serve this host (mirror release missing, or +# a host we don't yet build): resolve against leejet so native install still works. +UPSTREAM_FALLBACK_REPO = "leejet/stable-diffusion.cpp" # Pinned release tag for REPRODUCIBILITY: "releases/latest" silently swaps the binary -# under users on every upstream push. Override with UNSLOTH_SD_CPP_TAG; set it empty to -# track latest. If the pinned tag is gone upstream, install falls back to latest. -DEFAULT_TAG = "master-737-3b6c9ca" +# under users on every push. Override with UNSLOTH_SD_CPP_TAG; set it empty to track +# latest. If the pinned tag is gone, install falls back to that repo's latest. +DEFAULT_TAG = "master-741-484baa4" # Back-compat alias (some callers/tests import REPO). REPO = DEFAULT_REPO @@ -118,14 +122,24 @@ def resolve_release_asset( pool = [a for a in zips if "bin-win" in a.lower()] token = _WINDOWS_ACCEL_TOKEN.get(accel, accel) sel = [a for a in pool if token in a.lower()] - if not sel: # fall back to a plain avx2 CPU build - sel = [a for a in pool if "avx2" in a.lower()] - return sel[0] if sel else (pool[0] if pool else None) + if sel: + return sel[0] + # An EXPLICIT GPU accelerator with no matching asset is a real miss, not a CPU + # request: return None so the caller can fall back to a repo that builds it, + # rather than silently installing a CPU build for a --accelerator cuda request. + if accel in ("cuda", "vulkan", "rocm"): + return None + # auto / cpu -> a plain avx2 CPU build, else any windows build. + cpu = [a for a in pool if "avx2" in a.lower()] + return cpu[0] if cpu else (pool[0] if pool else None) # linux (and anything else unix-like) pool = [a for a in zips if "linux" in a.lower() and any(t in a.lower() for t in arch)] - if accel in _LINUX_ACCEL_TOKEN: - sel = [a for a in pool if _LINUX_ACCEL_TOKEN[accel] in a.lower()] + if accel in ("cuda", "vulkan", "rocm"): + # Explicit GPU accelerator: require its marker, else no match (let the caller + # fall back) -- never hand back a plain CPU build for a GPU request. + marker = _LINUX_ACCEL_TOKEN.get(accel, accel) + sel = [a for a in pool if marker in a.lower()] else: # auto / cpu -> the plain build with no accelerator marker sel = [a for a in pool if not any(m in a.lower() for m in _LINUX_ACCEL_MARKERS)] return sel[0] if sel else None @@ -137,10 +151,14 @@ def _fetch_release( repo: Optional[str] = None, token: Optional[str] = None, timeout: float = 30.0, -) -> dict: - """GET a release JSON from GitHub. With ``tag`` set, fetch that exact release (and fall - back to latest if the tag is gone upstream); otherwise fetch latest. ``token`` is - optional and lifts the API rate limit.""" + allow_latest: bool = True, +) -> Optional[dict]: + """GET a release JSON from GitHub. With ``tag`` set, fetch that exact release; otherwise + fetch latest. ``token`` is optional and lifts the API rate limit. + + When the pinned ``tag`` is missing (404): if ``allow_latest`` fall back to that repo's + latest, else return ``None`` so the caller can try the SAME pin on another repo before + settling for any repo's unpinned latest.""" repo = repo or _repo() token = token or os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN") @@ -155,9 +173,11 @@ def _fetch_release( if tag: try: return _get(f"{base}/tags/{tag}") - except urllib.error.HTTPError as exc: # pinned tag removed upstream -> latest + except urllib.error.HTTPError as exc: # pinned tag removed -> maybe latest if exc.code != 404: raise + if not allow_latest: + return None print( f"sd-cli: pinned tag {tag} not found on {repo}; falling back to latest", flush = True ) @@ -213,6 +233,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, @@ -272,6 +303,90 @@ def _maybe_fetch_windows_cudart(release: dict, chosen: str, target: Path) -> Non dest.unlink(missing_ok = True) +def _resolve_repo_asset( + repo: str, + tag: Optional[str], + accelerator: str, + token: Optional[str], + *, + allow_latest: bool = True, +) -> tuple[Optional[dict], Optional[str]]: + """Fetch ``repo``'s release and pick the asset for this host. Returns + ``(release, asset_name)`` or ``(None, None)`` when the repo has no usable release + (fetch failed, or the pinned tag is missing and ``allow_latest`` is False) or no + asset for this host, so the caller can fall back.""" + try: + release = _fetch_release(tag, repo = repo, token = token, allow_latest = allow_latest) + except Exception as exc: # noqa: BLE001 - network / rate limit -> fall back + print(f"sd-cli: {repo} release fetch failed ({exc})", flush = True) + return None, None + if release is None: # pinned tag missing and the latest fallback was withheld + return None, None + names = [a["name"] for a in (release.get("assets") or [])] + chosen = resolve_release_asset( + names, + system = platform.system(), + machine = platform.machine(), + accelerator = accelerator, + ) + return release, chosen + + +def _resolve_with_fallback( + accelerator: str, token: Optional[str] +) -> tuple[str, Optional[dict], Optional[str]]: + """Resolve ``(used_repo, release, asset_name)`` for this host across the primary repo + and -- only when the built-in default is in use and the user did not pin a repo -- the + upstream fallback. + + Ordering guarantees reproducibility: a pinned tag is tried EXACTLY on every candidate + repo before any repo's unpinned latest, so a mirror that is missing the pinned release + prefers the pinned upstream build over an unpinned mirror-latest. Returns + ``(primary, None, None)`` when nothing serves this host. Shared by ``install`` and + ``--print-asset`` so both honour the same fallback.""" + tag = _pinned_tag() + primary = _repo() + # Fall back to upstream ONLY when the user did not pin a repo (env unset) and the + # built-in default is in use: an explicit UNSLOTH_SD_CPP_REPO (even one equal to the + # default) gets exactly that repo, with no surprise upstream substitution. + repo_pinned = bool((os.environ.get("UNSLOTH_SD_CPP_REPO") or "").strip()) + allow_upstream = ( + not repo_pinned and primary == DEFAULT_REPO and DEFAULT_REPO != UPSTREAM_FALLBACK_REPO + ) + + # (repo, tag_to_fetch, allow_latest). With a pin set, try the exact pin on every repo + # first (allow_latest = False), then each repo's latest (tag = None). + attempts: list[tuple[str, Optional[str], bool]] = [] + if tag: + attempts.append((primary, tag, False)) + if allow_upstream: + attempts.append((UPSTREAM_FALLBACK_REPO, tag, False)) + attempts.append((primary, None, True)) + if allow_upstream: + attempts.append((UPSTREAM_FALLBACK_REPO, None, True)) + else: + attempts.append((primary, None, True)) + if allow_upstream: + attempts.append((UPSTREAM_FALLBACK_REPO, None, True)) + + for repo, want_tag, allow_latest in attempts: + release, chosen = _resolve_repo_asset( + repo, want_tag, accelerator, token, allow_latest = allow_latest + ) + if release is not None and chosen: + if repo != primary: + # Diagnostic goes to stderr, not stdout: --print-asset documents its + # stdout as the asset name only, so a caller parsing it as a single line + # must not see this fallback log mixed in. + print( + f"falling back to {repo} for {platform.system()}/{platform.machine()}", + file = sys.stderr, + flush = True, + ) + return repo, release, chosen + return primary, None, None + + def install( *, install_dir: Optional[Path] = None, @@ -280,25 +395,22 @@ def install( ) -> Path: """Download + extract the prebuilt for this host. Returns the sd-cli path. - Raises ``RuntimeError`` if no asset matches the host (the caller should then - build from source) or the archive has no ``sd-cli``. + Resolves against the Unsloth mirror (``DEFAULT_REPO``) first; if the mirror can't + serve this host (release missing, or a host we don't build) AND the default repo is + in use, falls back to leejet upstream so native install still works. Raises + ``RuntimeError`` only when neither source has an asset for the host, or the archive + has no ``sd-cli``. """ target = install_dir or default_install_dir() - release = _fetch_release(_pinned_tag(), token = token) - print(f"sd-cli: source {_repo()} release {release.get('tag_name', '?')}", flush = True) - names = [a["name"] for a in release.get("assets", [])] - chosen = resolve_release_asset( - names, - system = platform.system(), - machine = platform.machine(), - accelerator = accelerator, - ) - if not chosen: + used_repo, release, chosen = _resolve_with_fallback(accelerator, token) + + if release is None or not chosen: raise RuntimeError( f"No prebuilt sd-cli for {platform.system()}/{platform.machine()} " - f"(accelerator={accelerator}). Build from source: " - f"https://github.com/{_repo()}" + f"(accelerator={accelerator}) from {used_repo}. Build from source: " + f"https://github.com/{used_repo}" ) + print(f"sd-cli: source {used_repo} release {release.get('tag_name', '?')}", flush = True) asset = next(a for a in release["assets"] if a["name"] == chosen) url = asset["browser_download_url"] target.mkdir(parents = True, exist_ok = True) @@ -323,6 +435,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 @@ -338,14 +457,10 @@ def main(argv: Optional[list[str]] = None) -> int: args = p.parse_args(argv) if args.print_asset: - release = _fetch_release(_pinned_tag()) - names = [a["name"] for a in release.get("assets", [])] - chosen = resolve_release_asset( - names, - system = platform.system(), - machine = platform.machine(), - accelerator = args.accelerator, - ) + # Route through the same primary/fallback resolution as install(), so a host the + # mirror does not build (e.g. a Linux Vulkan request) reports the upstream asset + # it would actually download instead of a false "no matching prebuilt". + _used, _release, chosen = _resolve_with_fallback(args.accelerator, None) print(chosen or "(no matching prebuilt; build from source)") return 0 if chosen else 2