vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-08-25 08:42:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29
unsloth/studio/backend/tests/test_sd_cpp_args.py
Maheswar Kumar e7ea3cb979
studio: honour the gpu selection for image and video loads (#8645) 
* studio: honour the gpu selection for image and video loads

The GPU picker was wired into the chat runtime and training only. Neither /images/load nor
/video/load carried gpu_ids at all, so both engines placed every module on ordinal 0 whatever
the user selected. On a mixed box that is frequently the smaller card, and for a checkpoint
that only fits on the larger one it decides whether the load runs at all.

Both request models, both routes and both begin_load chains now take gpu_ids, matching the
LLM path. Neither engine shards a diffusion checkpoint -- diffusers places whole components
and sd.cpp assigns whole modules per backend device -- so a selection of several cards
resolves to one, and the most free VRAM wins, the rule auto_select_gpu_ids already applies to
training. Picking the first id would land back on ordinal 0 whenever a user selects
everything, which is the case this exists to serve. An index this host does not have is
refused at the route, before the arbiter evicts chat.

For diffusers the index is a new DiffusionDeviceTarget field rather than part of device.
The memory, speed, attention and engine-routing policies all compare that string by value
against "cuda", so folding the index into it would make every one of them take its non-CUDA
branch and load the model with those optimisations silently disabled. Placement reads the new
torch_device; the policies keep reading device. The pin itself is torch.cuda.set_device,
applied where the target is built: it is thread-local, and the load, generate, ControlNet and
re-activation paths each resolve on their own thread. It is also the lever that moves the
offload budget, since diffusion_memory reads mem_get_info() for the current device, so the
weights and the budget they are sized against stay on the same card. The capability probe is
asked of the selected card too, or a mixed-generation box picks its dtype from the wrong GPU.

For sd.cpp the selection becomes --backend diffusion=CUDA<n>,te=<n>,vae=<n>, built after the
offload policy so --clip-on-cpu / --vae-on-cpu keep those modules on the CPU: they are the
deprecated spellings of te=cpu / vae=cpu, and the parser is last-wins, so pinning over them
would undo low_vram. Device names come from the load's own binary via --list-devices, and only
CUDA / ROCm names are matched, since Vulkan ordinals are a separate namespace with no defined
mapping to a physical index.

The image and video Advanced panels gain a GPU control, shown only where there is more than
one pinnable CUDA / ROCm device. A pick whose card has since disappeared falls back to
automatic rather than sending an index the backend would refuse.

Part of #8636.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* import pinnableGpuContext instead of relying on the re-export

The name was only re-exported from gpu-selection, not imported, so useDiffusionGpuChoices
failed the frontend build with TS2552.

* resolve the gpu selection once per load and pin every worker to it

The first cut resolved lazily from a candidate list, in the wrong id space, and pinned too
late. Each of those is a separate way to end up on the wrong card.

Physical ids now go through the hardware layer that owns the visibility mask.
resolve_requested_gpu_ids validates against the parent-visible set and get_parent_visible_gpu_ids
gives the order torch enumerates, so a physical id maps to its position in that list. Comparing
against torch.cuda.device_count() was wrong twice over: under CUDA_VISIBLE_DEVICES=4,5 it rejects
both valid picks, and under a reordered mask like 1,0 it accepts a pick that targets the other
card. Free VRAM is read on the translated ordinal for the same reason.

The winner is decided once, in begin_load, and carried to the worker as gpu_ordinal. Re-deriving
it per target meant the ranking changed the moment the checkpoint landed: the load picked the
emptiest card, filled it, and a generate-time resolution then moved the current device to a
different GPU while the pipeline stayed put. The resolved ordinal is committed onto _LoadState
with the pipeline, so a load in flight no longer moves the resident model's card, and a load that
fails before teardown leaves it alone.

Every worker that touches a loaded pipeline now pins through _state_device_target before it
builds any device object. torch.cuda.set_device is thread-local, so the load thread's pin did
nothing for the generate thread: image generation placed a freshly downloaded ControlNet with
state.device ("cuda", un-indexed) before the late pin, landing it on a different card than the
base pipeline, and the video generate worker never resolved a target at all, so its H3 memory
probe and every torch.Generator resolved against its own default device.

The precision gate runs against the selected card. assert_precision_available was reached before
the selection was known, so on heterogeneous GPUs it judged an explicit quantization scheme
against the default card: a valid Ampere-or-newer pick could be refused because GPU 0 is older,
or an unsupported one could pass preflight and fail only after eviction.

The sd.cpp --backend pin is built where the flags are handed to a binary, not once up front. A
deferred accelerator install, the post-download re-resolve, and a server start that falls back to
one-shot can all replace the build after the offload policy is computed, and the ggml device
names come from whichever one actually runs.

Both routes only validate a selection on a CUDA or ROCm target. Physical ids have no applicator
on XPU, MPS or CPU, and the request contract says to ignore them there, so validating turned a
documented no-op into a 400.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* judge every precision preflight and the h3 plan against the selected gpu

The rework carried the resolved ordinal inside begin_load, but three decisions upstream of it
still read the default card.

Both routes now resolve the selection once, through a shared _selected_gpu_ordinal that only runs
the CUDA resolver when the target reports a CUDA device, and pass the ordinal into every
assert_precision_available and assert_video_precision_available call, including the download-plan
preflights and the engine-switch re-check. Judging an explicit quantization scheme against the
default card refuses a pick the selected card supports, or passes one it cannot and fails after
eviction. The video load route also resolves the ordinal before its precision gate rather than
after it.

The MiniMax-H3 auto denoiser planner takes the ordinal too. It picks the file set and the memory
policy from device capacity, so against the wrong card it either chooses the pinned hosted INT8
denoiser without the dense fallback the destination cannot do without, or stages roughly 66 GB of
dense shards nothing opens.

The Advanced GPU control reads the torch inventory rather than the inference-backend-selected
list. On a host running a Vulkan llama.cpp build for chat with CUDA torch for diffusers, the
Vulkan branch returned entries marked diffusionPinnable false, so the control never appeared even
though the image and video routes honour those physical ids. The two runtimes are independent, so
the diffusion picker asks torch.

* scope the device pin and carry the ordinal into every remaining decision

Resolving a target with an ordinal is not the same as making that card current, and several
probes read the current device rather than the target.

The pin is now scoped where it has to be. Route preflights run on the asyncio.to_thread executor,
whose threads are reused, so an unrestored torch.cuda.set_device left this request's card current
for the next one, including an automatic load with no selection at all. diffusion_device_scope
restores the previous device, and both precision preflights use it around their probes, which
reach argument-less CUDA calls and bare "cuda" allocations. _resolve_device_target no longer pins
as a side effect: the dedicated load and generate workers take the permanent pin explicitly.

The ROCm dtype probe asks about the selected card. is_bf16_supported() takes no device argument,
so the only way to ask is to make that card current for the call; without it a bf16-capable pick
behind an older default was promoted to fp32, and the reverse attempted bf16 on a card that
cannot run it.

The route hands its ranked ordinal to begin_load instead of the raw list. Both were ranking from
live free VRAM, with the network preflight, engine activation and arbiter eviction in between, so
a scheme could be approved against one card and the checkpoint placed on another.

Prefetch and plan decisions read the selection too: the dense-quant probe and the pre-cast
text-encoder resolver take it from the load kwargs, the H3 auto denoiser planner takes it on the
download-plan path as well as the load path, and both plan routes pass it through. These size the
file set from device capacity, so against the wrong card they stage tens of GB the selected card
would replace, or omit what it needs and force an inline fallback after eviction.

The image and video plan requests send gpu_ids. Staging precedes loading, so a plan asked without
the pick was preflighted against the default card and could refuse a precision the selected card
supports before the correctly pinned load was ever reached.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* accept the ordinal in the native engine and stop the scope swallowing refusals

Two of these were regressions from the previous commit.

The route passes gpu_ordinal to whichever engine it activated, but the native sd.cpp backend
still accepted only gpu_ids, so every native GGUF image load raised an unexpected-keyword
TypeError before the background load started and the route answered 500. It takes the resolved
ordinal now, and stops re-ranking the raw list for itself.

diffusion_device_scope caught around its own yield, so an exception raised by the body was
swallowed and the generator yielded a second time, which contextlib turns into "generator did not
stop after throw()". An explicit-precision refusal reached the route as that instead of its real
message. Only entering the context is guarded now, and the body's exceptions travel untouched
while the previous device is still restored in a finally.

The rest are the same ordinal reaching further: the video pre-cast text-encoder resolver takes it
on both the plan and prefetch paths, the image DiT artifact planner reads it from the load kwargs
it already receives, and the H3 plan sizing uses the restoring scope rather than the permanently
pinning helper, since the plan route reaches it on a pooled executor thread.

The transformer-quant smoke cache is keyed by the card the probe ran on rather than by "cuda". A
pass on one selected GPU was standing in for a card that was never tested, and a failure on an
older card rejected the same scheme on a capable one.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* rank the download-plan gpu once per request

Both plan routes resolved the ordinal twice, once for the precision preflight and again for
download_plan. The preflight's smoke probe allocates on the card it tests and other workloads
move free VRAM too, so the second ranking could pick a different card and the plan then sized
its file set for a GPU the requested precision was never validated on.

* scope planning probes to the selected card and rank gpus only when training is idle

Removing the pin from _resolve_device_target left the planning helpers building an indexed
target without making that card current, and their selectors read the current device.

_dense_quant_prefetch_needed and _dit_prequant_plan_source now run their whole decision inside
diffusion_device_scope, so select_transformer_quant_scheme, resolve_dense_quant_candidate and
the memory snapshot all measure the card the load will use. The H3 plan scope had the same shape
of mistake: it closed right after building the target, leaving the capacity reader and the
encoder-scheme probe on the pooled thread's default card, so it now encloses the whole sizing
call.

The compile-cache fingerprint reads the current device instead of hardcoding index 0. A load
pinned to another card compiles for that architecture, and keying the bundle by GPU 0 let two
cards share or overwrite each other's artifacts.

Both plan routes resolve the GPU only after the training guard has answered. Ranking reads free
VRAM per candidate and opens a CUDA context on each, which is exactly what those routes refuse
to do while a training subprocess owns the cards.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: keep the diffusion GPU pin out of the offload policy

Two problems the device pin introduced into the native (sd.cpp) path, both
only visible once a card is actually selected.

The CPU-backend restart stopped working. `_restart_server_on_cpu_backend`
relaunches with `state.offload_flags` plus `--backend cpu`, and sd.cpp
declares `--backend` with `concat = ','`, so repeated values are JOINED into
one spec rather than replacing. An explicit per-module entry then outranks the
bare `cpu` default, so `diffusion=CUDA0,te=CUDA0,vae=CUDA0,cpu` still runs the
whole graph on the card that just aborted. The recovery that exists to survive
a ggml op the device cannot run became a silent no-op, and the abort recurs
immediately. Confirmed against the pinned prebuilt
(master-813-bfbef5b-u13b9d92), which logs the joined spec verbatim.

And the status lied about the offload. `status()` and the saved recipe derive
"was anything offloaded?" from whether the flag tuple is empty, which is how a
`fast` load reports `none`. The pin lives in the same tuple, so picking a GPU
reported CPU offload as on for a load that offloads nothing.

Both go through one helper, `without_device_backend_flags`, so the pin is
dropped where the flags are read as a policy and kept where they are argv.

Also: the native backend now resolves a bare `gpu_ids` itself, the way the
diffusers and video backends already do. The routes rank the selection and
pass the winner, so this only affects a direct caller, but it was the one
engine that dropped such a pick silently, and it leaves
`resolve_selected_cuda_ordinal` imported and unused.

Tests: the CPU restart argv, the `fast` and offloaded status pair, the native
self-resolve, both download-plan routes forwarding the selection and refusing a
bad index with a 400 (they had no coverage for this field, and they are the
routes that size a file set against a card), and the frontend rule behind the
Advanced control (single card offers nothing, Vulkan and XPU cards are not
offered, mixed namespaces are never one pool, a stale pick falls back to auto).

* Studio: update the Vulkan picker contract for the diffusion inventory

The contract pins the exact source line that gates the Vulkan branch, and this
branch deliberately scoped that gate: an image or video load runs on torch
rather than llama-server, so it reads the torch inventory even on a Vulkan chat
build. The assertion still holds the property it was written for, now against
`!forDiffusion && inference?.backend === "vulkan"`.

Fixes the Repo tests (CPU) failure on this PR.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: put a shared generate worker back on its own card

Three items from this round, two taken.

A pooled worker kept the previous model's GPU. `/images/generate` runs on an
`asyncio.to_thread` worker, and those threads are reused, so a load pinned to a
non-default card leaves that thread set to it permanently. The next model
loaded automatically has no ordinal, so nothing re-pinned the worker, and its
bare `cuda` Generators and allocations then targeted the previous card while
the weights sat on the default one: a cross-device error, or an allocation on a
GPU the user did not choose. The card the weights actually landed on is now
recorded with the pipeline as `placed_ordinal` and applied by every worker that
touches it. Kept separate from `gpu_ordinal`, so an automatic load still
resolves a bare, un-indexed device and reports exactly what it did before.

The download-plan routes dropped the selection entirely while training ran. The
training guard is there to keep a CUDA context from being opened, which only
the free-VRAM ranking does; validating and translating physical ids reads the
environment mask and nvidia-smi and opens nothing. Skipping all of it let the
plan answer 200 for a GPU the load would then refuse, and size its file set for
the default card, after tens of gigabytes had been staged. The ranking still
waits for the trainer; the validation no longer does, so the single-card
selection the UI sends resolves either way and a bad one is refused at the plan.

Not taken: refusing the load when the ggml device name cannot be resolved. The
fallback is deliberate. sd.cpp treats an unknown argument as fatal, and
`--list-devices` only arrived in upstream #1734, so a build older than that (a
user's own SD_CLI_PATH copy included) would go from "the selection is not
honoured here" to "this model cannot load at all" the moment a card is picked.
It runs on the build's own device, which is what every native load does today.
It now says so in the log rather than dropping the pick in silence.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: give the CPU-offload hooks the card the load was pinned to

Three of the four items from this round.

The offload hooks went to GPU 0 whatever was selected. apply_memory_plan hands
diffusers a bare "cuda", and enable_model_cpu_offload reads the index off that
device: with none, `_offload_gpu_id = gpu_id or torch_device.index or 0` settles
on 0 and the hooks onload every module to cuda:0 (pipeline_utils.py, diffusers
0.39, confirmed against the installed copy). Generation runs on the selected
card, so a low_vram or group-fallback load paged its modules onto the wrong GPU
and either failed across devices or filled the card the selection existed to
avoid. Placement now takes an indexed device string; the bare one stays for
anything reading it as a policy.

The GPU choices array changed identity on every render. pinnableGpuContext
builds a fresh filtered array per call, so the hook returned a new array each
time, which fed the load-advanced snapshot, which fed the download-footprint
resolver the GGUF picker's effect depends on: every status poll cleared the
companion sizes and re-POSTed a download plan per variant, discarding whatever
was in flight. Memoized on the device list.

And the pick did not survive a reload. Every other Advanced select is reseeded
from the loaded build; this one cannot be, because the status reports the device
a pipeline is on and not which physical card, so a refresh reset it to Auto
while the model stayed put and the next Reapply moved it to the default GPU. It
persists now, and the existing staleness guard still drops an id whose card has
gone rather than sending one the backend would refuse.

Not taken: sizing the picker's fit recommendations for the selected card. That
budget comes from the shared model-picker heuristic that chat and training also
use, and it is advisory: the decisions that actually gate a load (the precision
preflight, the memory plan, the capacity gate) all read the selected card
already. Rewiring the shared picker is a feature, not a fix for this branch.

* Tighten the device-selection comments

* Tighten the device-selection test comments

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
2026-08-13 07:45:04 -07:00

776 lines
30 KiB
Python
Raw Permalink Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Unit tests for the sd-cli command builder (``sd_cpp_args.py``).
Pure: no torch, no subprocess, no files. Just argv construction and the
policy -> offload-flag / family -> text-encoder-flag mappings.
"""
from __future__ import annotations
import pytest
from core.inference.diffusion_memory import (
OFFLOAD_GROUP,
OFFLOAD_MODEL,
OFFLOAD_NONE,
OFFLOAD_SEQUENTIAL,
)
from core.inference.sd_cpp_args import (
CPU_BACKEND_FLAGS,
SdCppGenParams,
SdCppModelFiles,
SdCppUpscaleParams,
SdCppVideoGenParams,
build_img_gen_request,
build_sd_cpp_command,
build_sd_cpp_server_command,
build_sd_cpp_upscale_command,
build_sd_cpp_video_command,
device_backend_flags,
is_ggml_unsupported_op_abort,
metal_text_encoder_flags,
native_speed_flags,
offload_flags,
text_encoder_flags_for_family,
)
def _pair(cmd: list[str], flag: str):
"""Value following ``flag`` in ``cmd``, or None if the flag is absent."""
return cmd[cmd.index(flag) + 1] if flag in cmd else None
# ── family text-encoder wiring ──────────────────────────────────────────────
def test_te_flags_by_family():
assert text_encoder_flags_for_family("z-image") == ("--llm",)
assert text_encoder_flags_for_family("qwen-image") == ("--qwen2vl",)
assert text_encoder_flags_for_family("flux.1") == ("--clip_l", "--t5xxl")
assert text_encoder_flags_for_family("flux.2-klein") == ("--llm",)
assert text_encoder_flags_for_family("unknown") == ()
# ── Metal text-encoder placement ────────────────────────────────────────────
def test_metal_keeps_the_text_encoder_off_the_gpu(monkeypatch):
# ggml's Metal backend aborts the process on RMS_NORM with non-contiguous rows and has no per-op CPU fallback, so an LLM
# text encoder killed sd-server mid-generation on macOS (macos-14, FLUX.2-klein-4B Q2_K: loads on mps, first generation exits -6).
monkeypatch.delenv("UNSLOTH_DIFFUSION_SD_CPP_METAL_TE_GPU", raising = False)
monkeypatch.setattr("sys.platform", "darwin")
assert metal_text_encoder_flags() == ["--clip-on-cpu"]
for other in ("linux", "win32"):
monkeypatch.setattr("sys.platform", other)
assert metal_text_encoder_flags() == []
# Opt back in once ggml grows the kernel.
monkeypatch.setattr("sys.platform", "darwin")
monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP_METAL_TE_GPU", "1")
assert metal_text_encoder_flags() == []
def test_metal_text_encoder_flag_reaches_both_command_builders(monkeypatch):
monkeypatch.delenv("UNSLOTH_DIFFUSION_SD_CPP_METAL_TE_GPU", raising = False)
monkeypatch.setattr("sys.platform", "darwin")
files = SdCppModelFiles(diffusion_model = "/m/x.gguf")
server = build_sd_cpp_server_command(
binary = "sd-server", files = files, host = "127.0.0.1", port = 1234
)
cli = build_sd_cpp_command(
binary = "sd-cli",
files = files,
params = SdCppGenParams(prompt = "x"),
output_path = "/o/x.png",
)
video = build_sd_cpp_video_command(
binary = "sd-cli",
files = SdCppModelFiles(
diffusion_model = "/m/h3.gguf",
vae = "/m/video.safetensors",
llm = "/m/qwen.gguf",
),
params = SdCppVideoGenParams(prompt = "x", width = 960, height = 544, num_frames = 124),
output_path = "/o/x.webm",
)
assert server.count("--clip-on-cpu") == 1
assert cli.count("--clip-on-cpu") == 1
assert video.count("--clip-on-cpu") == 1
# An offload policy that already pins the encoder must not emit it twice.
dual = build_sd_cpp_server_command(
binary = "sd-server",
files = files,
host = "127.0.0.1",
port = 1234,
offload = offload_flags("model"),
)
assert dual.count("--clip-on-cpu") == 1
monkeypatch.setattr("sys.platform", "linux")
assert "--clip-on-cpu" not in build_sd_cpp_server_command(
binary = "sd-server", files = files, host = "127.0.0.1", port = 1234
)
# ── offload policy -> sd-cli flags ──────────────────────────────────────────
def test_native_speed_flags():
assert native_speed_flags(None) == []
assert native_speed_flags("off") == []
assert native_speed_flags("") == []
# default now includes conv-direct: ~9% faster sampling on CPU (z-image Q8_0, 192 threads) at identical RSS and unchanged decode.
assert native_speed_flags("default") == ["--diffusion-fa", "--diffusion-conv-direct"]
assert native_speed_flags("max") == ["--diffusion-fa", "--diffusion-conv-direct"]
with pytest.raises(ValueError):
native_speed_flags("ludicrous")
def test_offload_none_is_empty():
assert offload_flags(OFFLOAD_NONE) == []
def test_offload_group_streams_with_flash_attention():
flags = offload_flags(OFFLOAD_GROUP)
assert "--offload-to-cpu" in flags
assert "--diffusion-fa" in flags
# group keeps CLIP/VAE resident -> no per-component cpu flags
assert "--clip-on-cpu" not in flags
assert "--vae-on-cpu" not in flags
def test_offload_model_pushes_everything_to_cpu_and_tiles():
flags = offload_flags(OFFLOAD_MODEL)
for expected in (
"--offload-to-cpu",
"--clip-on-cpu",
"--vae-on-cpu",
"--vae-tiling",
"--diffusion-fa",
):
assert expected in flags
# sequential maps the same as model
assert offload_flags(OFFLOAD_SEQUENTIAL) == flags
def test_offload_can_keep_the_vae_off_the_cpu_path():
"""H3's audio VAE aborts on the CPU path, so `low_vram` has to drop just that flag.
`ggml_conv_1d` hardcodes an F16 im2col destination and
`ggml_compute_forward_im2col_f16` then asserts the kernel is F16, while sd.cpp's
`audio_conv_weight_type` maps only BF16 to F16 and lets F32 through:
`GGML_ASSERT(src0->type == GGML_TYPE_F16) failed`, SIGABRT, exit 134. Converting the
checkpoint to fp16 does not help, since the type is imposed inside sd.cpp.
Everything else the policy asks for still applies. The denoiser dominates, so
`--offload-to-cpu` is where the saving is; dropping the mode entirely would cost far
more than dropping this one flag.
"""
for policy in (OFFLOAD_MODEL, OFFLOAD_SEQUENTIAL):
flags = offload_flags(policy, vae_on_cpu = False)
assert "--vae-on-cpu" not in flags
for expected in ("--offload-to-cpu", "--clip-on-cpu", "--vae-tiling", "--diffusion-fa"):
assert expected in flags, f"{policy}: {expected} should survive"
# The default is unchanged for every other family.
assert "--vae-on-cpu" in offload_flags(OFFLOAD_MODEL)
# And it is a no-op where the policy never emitted it.
assert offload_flags(OFFLOAD_GROUP, vae_on_cpu = False) == offload_flags(OFFLOAD_GROUP)
def test_offload_forced_flags_dedup():
# vae_tiling/diffusion_fa forced on with a policy that already sets them
flags = offload_flags(OFFLOAD_MODEL, vae_tiling = True, diffusion_fa = True)
assert flags.count("--vae-tiling") == 1
assert flags.count("--diffusion-fa") == 1
# forced on top of a no-offload policy
none_forced = offload_flags(OFFLOAD_NONE, vae_tiling = True, diffusion_fa = True)
assert none_forced == ["--diffusion-fa", "--vae-tiling"]
# ── full command construction ───────────────────────────────────────────────
def test_build_zimage_command_minimal():
files = SdCppModelFiles(
diffusion_model = "/m/z.gguf",
vae = "/m/ae.sft",
llm = "/m/qwen3.gguf",
)
params = SdCppGenParams(prompt = "a cat", width = 512, height = 768, steps = 8, cfg_scale = 1.0, seed = 42)
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/out/x.png")
assert cmd[0] == "/bin/sd-cli"
assert _pair(cmd, "--mode") == "img_gen"
assert _pair(cmd, "--diffusion-model") == "/m/z.gguf"
assert _pair(cmd, "--vae") == "/m/ae.sft"
assert _pair(cmd, "--llm") == "/m/qwen3.gguf"
assert _pair(cmd, "--prompt") == "a cat"
assert _pair(cmd, "--width") == "512"
assert _pair(cmd, "--height") == "768"
assert _pair(cmd, "--steps") == "8"
assert _pair(cmd, "--cfg-scale") == "1" # 1.0 -> "1"
assert _pair(cmd, "--seed") == "42"
assert _pair(cmd, "--output") == "/out/x.png"
# unset encoders are omitted
assert "--t5xxl" not in cmd
assert "--qwen2vl" not in cmd
def test_build_flux1_dual_text_encoders():
files = SdCppModelFiles(
diffusion_model = "/m/flux.gguf",
vae = "/m/ae.sft",
clip_l = "/m/clip_l.sft",
t5xxl = "/m/t5.gguf",
)
params = SdCppGenParams(prompt = "x", guidance = 3.5)
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
assert _pair(cmd, "--clip_l") == "/m/clip_l.sft"
assert _pair(cmd, "--t5xxl") == "/m/t5.gguf"
assert _pair(cmd, "--guidance") == "3.5"
assert "--llm" not in cmd
def test_build_appends_offload_and_extra_args_last():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
params = SdCppGenParams(prompt = "x")
off = offload_flags(OFFLOAD_GROUP)
cmd = build_sd_cpp_command(
"/bin/sd-cli",
files,
params,
output_path = "/o.png",
offload = off,
threads = 8,
verbose = True,
extra_args = ["--rng", "cuda"],
)
assert "--offload-to-cpu" in cmd
assert _pair(cmd, "--threads") == "8"
assert "-v" in cmd
# extra args come after everything Studio set (last-wins for power users)
assert cmd[-2:] == ["--rng", "cuda"]
def test_build_negative_prompt_and_batch():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
params = SdCppGenParams(prompt = "x", negative_prompt = "blurry")
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
assert _pair(cmd, "--negative-prompt") == "blurry"
# A CLI batch would silently drop every image after the first (the runner collects only the literal --output path), so the builder rejects it.
with pytest.raises(ValueError, match = "single-image"):
build_sd_cpp_command(
"/bin/sd-cli",
files,
SdCppGenParams(prompt = "x", batch_count = 3),
output_path = "/o.png",
)
def test_build_omits_unset_optional_params():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
params = SdCppGenParams(prompt = "x") # no steps/cfg/seed/sampler
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
for flag in (
"--steps",
"--cfg-scale",
"--guidance",
"--seed",
"--sampling-method",
"--batch-count",
"--threads",
"-v",
):
assert flag not in cmd
def test_build_requires_diffusion_model_and_prompt():
with pytest.raises(ValueError):
build_sd_cpp_command(
"/bin/sd-cli",
SdCppModelFiles(diffusion_model = ""),
SdCppGenParams(prompt = "x"),
output_path = "/o.png",
)
with pytest.raises(ValueError):
build_sd_cpp_command(
"/bin/sd-cli",
SdCppModelFiles(diffusion_model = "/m/z.gguf"),
SdCppGenParams(prompt = " "),
output_path = "/o.png",
)
# ── img2img / inpaint / edit / LoRA (Phase 6) ───────────────────────────────
def test_build_img2img_adds_init_and_strength():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/q.gguf")
params = SdCppGenParams(prompt = "make it autumn", init_img = "/in/src.png", strength = 0.6)
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
assert _pair(cmd, "--init-img") == "/in/src.png"
assert _pair(cmd, "--strength") == "0.6"
assert _pair(cmd, "--mode") == "img_gen" # img2img is still img_gen mode
def test_build_inpaint_adds_mask():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
params = SdCppGenParams(prompt = "x", init_img = "/in/src.png", mask = "/in/mask.png", strength = 0.8)
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
assert _pair(cmd, "--mask") == "/in/mask.png"
assert _pair(cmd, "--init-img") == "/in/src.png"
def test_build_inpaint_mask_without_init_img_rejected():
# sd-cli inpaint needs a source image, so a --mask with no --init-img is rejected up front instead of emitting doomed argv.
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
params = SdCppGenParams(prompt = "x", mask = "/in/mask.png")
with pytest.raises(ValueError, match = "init_img is required"):
build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
def test_build_rejects_none_prompt():
# A None prompt must be rejected, not coerced to the literal string "None" and forwarded into argv.
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
with pytest.raises(ValueError, match = "prompt is required"):
build_sd_cpp_command(
"/bin/sd-cli", files, SdCppGenParams(prompt = None), output_path = "/o.png"
)
def test_build_edit_repeats_ref_image():
files = SdCppModelFiles(diffusion_model = "/m/flux.gguf")
params = SdCppGenParams(prompt = "add a hat", ref_images = ("/r/a.png", "/r/b.png"))
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
# each ref image gets its own --ref-image flag
idxs = [i for i, t in enumerate(cmd) if t == "--ref-image"]
assert len(idxs) == 2
assert [cmd[i + 1] for i in idxs] == ["/r/a.png", "/r/b.png"]
def test_img2img_unset_dims_lets_sdcpp_derive_from_source():
# img2img/inpaint/edit with dims unset must NOT force --width/--height, so sd.cpp derives the size from the input image.
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
cmd = build_sd_cpp_command(
"/bin/sd-cli",
files,
SdCppGenParams(prompt = "x", init_img = "/in/src.png"),
output_path = "/o.png",
)
assert "--width" not in cmd and "--height" not in cmd
# an edit (ref-image) run derives its size too
cmd2 = build_sd_cpp_command(
"/bin/sd-cli",
files,
SdCppGenParams(prompt = "x", ref_images = ("/r/a.png",)),
output_path = "/o.png",
)
assert "--width" not in cmd2 and "--height" not in cmd2
def test_img2img_explicit_dims_are_emitted():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
cmd = build_sd_cpp_command(
"/bin/sd-cli",
files,
SdCppGenParams(prompt = "x", init_img = "/in/src.png", width = 768, height = 512),
output_path = "/o.png",
)
assert _pair(cmd, "--width") == "768" and _pair(cmd, "--height") == "512"
def test_txt2img_unset_dims_keep_1024_default():
# A plain txt2img run with no dims keeps the prior 1024x1024 default.
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
cmd = build_sd_cpp_command(
"/bin/sd-cli", files, SdCppGenParams(prompt = "x"), output_path = "/o.png"
)
assert _pair(cmd, "--width") == "1024" and _pair(cmd, "--height") == "1024"
def test_build_lora_dir_and_apply_mode():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
params = SdCppGenParams(
prompt = "a portrait <lora:mystyle:0.8>",
lora_dir = "/loras",
lora_apply_mode = "at_runtime",
)
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
assert _pair(cmd, "--lora-model-dir") == "/loras"
assert _pair(cmd, "--lora-apply-mode") == "at_runtime"
# the <lora:...> tag rides in the prompt unchanged
assert _pair(cmd, "--prompt") == "a portrait <lora:mystyle:0.8>"
def test_txt2img_omits_image_conditioning_flags():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
cmd = build_sd_cpp_command(
"/bin/sd-cli", files, SdCppGenParams(prompt = "x"), output_path = "/o.png"
)
for flag in ("--init-img", "--strength", "--mask", "--ref-image", "--lora-model-dir"):
assert flag not in cmd
# ── upscale mode ────────────────────────────────────────────────────────────
def test_build_upscale_command():
params = SdCppUpscaleParams(
input_image = "/in/small.png", upscale_model = "/m/esrgan.pth", repeats = 2
)
cmd = build_sd_cpp_upscale_command("/bin/sd-cli", params, output_path = "/out/big.png")
assert _pair(cmd, "--mode") == "upscale"
assert _pair(cmd, "--init-img") == "/in/small.png"
assert _pair(cmd, "--upscale-model") == "/m/esrgan.pth"
assert _pair(cmd, "--upscale-repeats") == "2"
assert _pair(cmd, "--output") == "/out/big.png"
# no prompt / text-encoder flags in upscale mode
assert "--prompt" not in cmd and "--llm" not in cmd
def test_build_upscale_rejects_non_positive_repeats():
# repeats=0 must not be silently swallowed into sd-cli's default of one pass.
with pytest.raises(ValueError, match = "repeats"):
build_sd_cpp_upscale_command(
"/bin/sd-cli",
SdCppUpscaleParams(input_image = "/i.png", upscale_model = "/m/e.pth", repeats = 0),
output_path = "/o.png",
)
def test_build_upscale_default_repeats_omits_flag():
cmd = build_sd_cpp_upscale_command(
"/bin/sd-cli",
SdCppUpscaleParams(input_image = "/i.png", upscale_model = "/m/e.pth"), # repeats=1
output_path = "/o.png",
)
assert "--upscale-repeats" not in cmd
def test_build_upscale_requires_input_and_model():
with pytest.raises(ValueError):
build_sd_cpp_upscale_command(
"/bin/sd-cli",
SdCppUpscaleParams(input_image = "", upscale_model = "/m/e.pth"),
output_path = "/o.png",
)
with pytest.raises(ValueError):
build_sd_cpp_upscale_command(
"/bin/sd-cli",
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)
def test_ggml_unsupported_op_abort_is_recognised_only_with_both_markers():
# The CPU-backend rescue fires only for the deterministic "this backend cannot run this graph" abort: an OOM kill or a plain crash surfaces as itself.
abort = (
"sd-server connection lost during img_gen poll (process exited, code -6)\n"
"Last output:\n"
"[ERROR] ggml_extend.hpp:70 - ggml_metal_op_encode_impl: error: unsupported op 'MUL_MAT'\n"
"1 sd-server 0x00000001044f8df4 ggml_abort + 156"
)
assert is_ggml_unsupported_op_abort(abort) is True
# The RMS_NORM shape of the same abort (text encoder) counts too.
assert (
is_ggml_unsupported_op_abort("error: unsupported op 'RMS_NORM'\nggml_abort + 156") is True
)
# Neither marker alone is enough, and an unrelated death is never a match.
assert is_ggml_unsupported_op_abort("unsupported op 'MUL_MAT'") is False
assert is_ggml_unsupported_op_abort("ggml_abort + 156") is False
assert is_ggml_unsupported_op_abort("process exited, code -9") is False
assert is_ggml_unsupported_op_abort("") is False
def test_server_command_appends_extra_args_last():
# --backend cpu is passed as extra_args by the abort rescue and sd.cpp is last-wins, so it must land after every normal flag.
cmd = build_sd_cpp_server_command(
"/x/sd-server",
SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/vae.sft"),
host = "127.0.0.1",
port = 1234,
native_speed = "default",
extra_args = list(CPU_BACKEND_FLAGS),
)
assert cmd[-2:] == ["--backend", "cpu"]
assert cmd[0] == "/x/sd-server"
def test_minimax_h3_video_command_has_all_joint_av_components():
files = SdCppModelFiles(
diffusion_model = "/m/minimax_h3_fl2va-Q4_K_M.gguf",
vae = "/m/video.safetensors",
audio_vae = "/m/audio.safetensors",
llm = "/m/qwen.gguf",
)
cmd = build_sd_cpp_video_command(
"/bin/sd-cli",
files,
SdCppVideoGenParams(
prompt = "a fox runs through snow",
width = 960,
height = 544,
num_frames = 124,
fps = 24,
steps = 30,
seed = 42,
),
output_path = "/out/result.webm",
offload = ["--diffusion-fa", "--offload-to-cpu"],
)
assert _pair(cmd, "--mode") == "vid_gen"
assert _pair(cmd, "--audio-vae") == "/m/audio.safetensors"
assert _pair(cmd, "--llm") == "/m/qwen.gguf"
assert _pair(cmd, "--video-frames") == "124"
assert _pair(cmd, "--fps") == "24"
assert _pair(cmd, "--cfg-scale") == "1"
assert "--rng" in cmd and _pair(cmd, "--rng") == "cpu"
assert cmd[-2:] == ["--diffusion-fa", "--offload-to-cpu"]
def test_video_build_appends_extra_args_verbatim_and_last():
"""The video mirror of the image builder's last-wins contract.
Token-wise de-duplication cannot express an override: the builder already sets --rng cpu, so
extra_args ["--rng", "cuda"] dropped the --rng it matched and appended a bare "cuda" for the
parser to choke on. Every sibling builder in this module appends the list verbatim.
"""
files = SdCppModelFiles(
diffusion_model = "/m/minimax_h3_fl2va-Q4_K_M.gguf",
vae = "/m/video.safetensors",
audio_vae = "/m/audio.safetensors",
llm = "/m/qwen.gguf",
)
cmd = build_sd_cpp_video_command(
"/bin/sd-cli",
files,
SdCppVideoGenParams(
prompt = "a fox runs through snow",
width = 960,
height = 544,
num_frames = 124,
fps = 24,
seed = 1,
),
output_path = "/o.webm",
extra_args = ["--rng", "cuda"],
)
assert cmd[-2:] == ["--rng", "cuda"]
assert cmd[-1] != "cuda" or cmd[-2] == "--rng"
# Studio's own value is still there, earlier, so sd.cpp's last-wins parser takes the override.
assert cmd.count("--rng") == 2
def _h3_video_cmd(**params):
return build_sd_cpp_video_command(
"/bin/sd-cli",
SdCppModelFiles(
diffusion_model = "/m/minimax_h3_fl2va-Q4_K_M.gguf",
vae = "/m/video.safetensors",
audio_vae = "/m/audio.safetensors",
llm = "/m/qwen.gguf",
),
SdCppVideoGenParams(
prompt = "a fox runs through snow",
width = 960,
height = 544,
num_frames = 124,
**params,
),
output_path = "/out/result.webm",
)
def test_minimax_h3_video_command_omits_keyframe_flags_for_text_only():
cmd = _h3_video_cmd()
assert "--init-img" not in cmd
assert "--end-img" not in cmd
def test_minimax_h3_video_command_carries_each_keyframe():
# Each keyframe combination maps to its sd.cpp flags.
assert _pair(_h3_video_cmd(init_img = "/k/first.png"), "--init-img") == "/k/first.png"
assert "--end-img" not in _h3_video_cmd(init_img = "/k/first.png")
end_only = _h3_video_cmd(end_img = "/k/last.png")
assert _pair(end_only, "--end-img") == "/k/last.png"
assert "--init-img" not in end_only
both = _h3_video_cmd(init_img = "/k/first.png", end_img = "/k/last.png")
assert _pair(both, "--init-img") == "/k/first.png"
assert _pair(both, "--end-img") == "/k/last.png"
def test_minimax_h3_video_command_packs_references_in_reading_order():
# Preserve the model's reference order.
cmd = _h3_video_cmd(
ref_images = ("/r/cat.png", "/r/style.png"),
ref_videos = ("/r/motion", "/r/orbit"),
ref_video_audios = ("/r/motion.wav",),
ref_audios = ("/r/voice.wav",),
)
assert [c for c in cmd if c.startswith("--ref")] == [
"--ref-image",
"--ref-image",
"--ref-video",
"--ref-video",
"--ref-video-audio",
"--ref-audio",
]
assert cmd[cmd.index("--ref-image") + 1] == "/r/cat.png"
assert cmd[cmd.index("--ref-video") + 1] == "/r/motion"
assert cmd[cmd.index("--ref-audio") + 1] == "/r/voice.wav"
def test_minimax_h3_video_command_refuses_keyframes_with_references():
# sd.cpp refuses the pair itself, but only once the model is resident.
with pytest.raises(ValueError, match = "different denoiser partitions"):
_h3_video_cmd(init_img = "/k/first.png", ref_images = ("/r/cat.png",))
with pytest.raises(ValueError, match = "different denoiser partitions"):
_h3_video_cmd(end_img = "/k/last.png", ref_audios = ("/r/voice.wav",))
def test_minimax_h3_video_command_refuses_an_unpairable_soundtrack():
# Reject soundtrack flags without a video at the same position.
with pytest.raises(ValueError, match = "reference video to pair with"):
_h3_video_cmd(ref_videos = ("/r/motion",), ref_video_audios = ("/r/a.wav", "/r/b.wav"))
def test_minimax_h3_video_command_carries_the_video_flow_shift():
# sd.cpp derives the audio schedule against a hardcoded 3.0, so only the video shift is a flag.
assert "--flow-shift" not in _h3_video_cmd()
assert _pair(_h3_video_cmd(flow_shift = 8.5), "--flow-shift") == "8.5"
assert _pair(_h3_video_cmd(flow_shift = 12.0), "--flow-shift") == "12"
def test_device_backend_flags_pin_all_three_graphs():
# sd.cpp defaults to ordinal 0 unless told otherwise, so a selection only takes effect through --backend.
assert device_backend_flags("CUDA1") == ["--backend", "diffusion=CUDA1,te=CUDA1,vae=CUDA1"]
# No selection keeps sd.cpp's own choice rather than pinning ordinal 0 explicitly.
assert device_backend_flags(None) == []
assert device_backend_flags("") == []
def test_device_backend_flags_leave_the_offloaded_modules_on_the_cpu():
# --clip-on-cpu / --vae-on-cpu ARE te=cpu / vae=cpu and the parser is last-wins, so pinning them would undo low_vram.
low_vram = offload_flags(OFFLOAD_MODEL)
assert "--clip-on-cpu" in low_vram and "--vae-on-cpu" in low_vram
assert device_backend_flags("CUDA1", low_vram) == [
"--backend",
"diffusion=CUDA1,te=cpu,vae=cpu",
]
# The group policy offloads parameters but runs both on the device, so both are pinned.
assert device_backend_flags("CUDA1", offload_flags(OFFLOAD_GROUP)) == [
"--backend",
"diffusion=CUDA1,te=CUDA1,vae=CUDA1",
]
def test_device_backend_flags_are_appended_not_folded_into_the_policy():
# The pin is built per binary at the point the flags are handed over, so the policy list it
# reads must stay unchanged: a deferred install can replace the build after this is computed.
policy = offload_flags(OFFLOAD_GROUP)
before = list(policy)
device_backend_flags("CUDA1", policy)
assert policy == before
Reference in a new issue View git blame Copy permalink