mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-04 13:20:38 +00:00
cbb2eaa729
Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.)
336 lines
12 KiB
Python
336 lines
12 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""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 (
|
|
SdCppGenParams,
|
|
SdCppModelFiles,
|
|
SdCppUpscaleParams,
|
|
build_sd_cpp_command,
|
|
build_sd_cpp_upscale_command,
|
|
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") == ()
|
|
|
|
|
|
# ── offload policy -> sd-cli flags ──────────────────────────────────────────
|
|
|
|
|
|
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_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", batch_count = 3)
|
|
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
|
|
assert _pair(cmd, "--negative-prompt") == "blurry"
|
|
assert _pair(cmd, "--batch-count") == "3"
|
|
|
|
|
|
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_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 left unset must NOT force --width/--height,
|
|
# so sd.cpp derives the size from the input image instead of resizing it to 1024.
|
|
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",
|
|
)
|