Address a further round of Codex review findings on the image PR
Backend: - validate_load_request rejects a non-.gguf single-file name before the GPU handoff, so a family-looking repo paired with README.md no longer evicts the chat model and only fails in the background load. - detect_family scopes the edit/kontext/inpaint keyword check to the model id or filename basename, not arbitrary parent directories, so a valid text-to-image file under a folder named edit is no longer rejected. - the images gallery listing skips records that fail schema validation, so one corrupt or hand-dropped PNG can no longer 500 the whole endpoint. - _terminate reaps the killed sd-cli child so cancellation and timeout paths do not leak zombie process-table entries. - the images load route gates the chat-eviction handoff on the resolved device being non-CPU, so a CPU-only diffusers fallback no longer evicts a resident chat model for a load that cannot use the GPU. Frontend: - treat Images as a chat-like full-height route (no outer padding or scroll) so its picker is not pushed down and the gallery is not clipped. - allow /images under the chat-only guard so the native CPU/MPS image path is reachable on the no-GPU hosts it was built for. - roll the optimistic quant label back when a same-repo swap fails after the load started, so the selector never advertises a quant that is not loaded.
This commit is contained in:
Daniel Han
parent
a4277a01e4
commit
18f9510d11
8 changed files with 136 additions and 22 deletions
|
|
@ -290,6 +290,14 @@ class DiffusionBackend:
|
|||
raise ValueError(
|
||||
"gguf_filename is required: this backend loads single-file GGUF checkpoints only."
|
||||
)
|
||||
# Reject a non-GGUF single-file name (e.g. README.md, config.json) here, before
|
||||
# the route hands the GPU over: without this a family-looking repo_id paired with
|
||||
# a non-GGUF filename passes preflight, evicts the chat model, and only fails in
|
||||
# the background from_single_file -- exactly the eviction this validation prevents.
|
||||
if not gguf_filename.lower().endswith(".gguf"):
|
||||
raise ValueError(
|
||||
f"gguf_filename must name a .gguf single-file checkpoint; got '{gguf_filename}'."
|
||||
)
|
||||
fam = detect_family_for_pick(repo_id, gguf_filename, family_override)
|
||||
if fam is None:
|
||||
raise ValueError(
|
||||
|
|
|
|||
|
|
@ -170,9 +170,13 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff
|
|||
# Match edit keywords as whole id segments, not raw substrings, so a normal
|
||||
# text-to-image repo like ".../some-image-edition" isn't misread as an editing
|
||||
# checkpoint. Qwen-Image-Edit / FLUX.1-Kontext still match (edit/kontext are
|
||||
# whole tokens there). Split on both path separators so a Windows local path
|
||||
# is segmented too.
|
||||
segments = set(re.split(r"[-_./\\]+", needle))
|
||||
# whole tokens there). Scope the check to the LAST path component (the model id
|
||||
# or filename), not arbitrary parent directories: a valid file selected as
|
||||
# repo_id `/models/edit` + filename `Z-Image-Turbo-Q4.gguf` must not be rejected
|
||||
# just because a parent folder happens to be named `edit`. The combined
|
||||
# `repo_id/gguf_filename` fallback passes the filename as that last segment.
|
||||
basename = re.split(r"[/\\]+", needle)[-1]
|
||||
segments = set(re.split(r"[-_.]+", basename))
|
||||
if any(kw in segments for kw in _EDIT_KEYWORDS):
|
||||
return None
|
||||
for fam in _FAMILIES:
|
||||
|
|
|
|||
|
|
@ -75,6 +75,14 @@ def _terminate(proc: "subprocess.Popen") -> None:
|
|||
proc.kill()
|
||||
except Exception: # noqa: BLE001 -- best-effort teardown
|
||||
pass
|
||||
# Reap the killed child so it does not linger as a zombie until the next Popen
|
||||
# cleanup / interpreter exit. Callers raise immediately after _terminate (the
|
||||
# cancellation and timeout paths), so without this a burst of image cancellations
|
||||
# leaks process-table entries. SIGKILL is prompt, so a short bounded wait suffices.
|
||||
try:
|
||||
proc.wait(timeout = 5)
|
||||
except Exception: # noqa: BLE001 -- best-effort reap; never block teardown
|
||||
pass
|
||||
|
||||
|
||||
def _binary_name(stem: str) -> str:
|
||||
|
|
|
|||
|
|
@ -11051,12 +11051,10 @@ async def load_diffusion_model(
|
|||
from core.inference.diffusion import get_diffusion_backend
|
||||
from core.inference.diffusion_device import resolve_diffusion_device_target
|
||||
from core.inference.diffusion_engine_router import (
|
||||
active_engine_name,
|
||||
annotate_status,
|
||||
select_and_activate_engine,
|
||||
)
|
||||
from core.inference.gpu_arbiter import acquire_for, release, DIFFUSION
|
||||
from core.inference.sd_cpp_engine import ENGINE_SD_CPP
|
||||
from utils.native_path_leases import redact_native_paths
|
||||
|
||||
backend = get_diffusion_backend()
|
||||
|
|
@ -11078,12 +11076,15 @@ async def load_diffusion_model(
|
|||
# installing the sd-cli binary if needed -- all BEFORE evicting chat, so a
|
||||
# native fallback never strands a half-loaded state.
|
||||
engine = await asyncio.to_thread(select_and_activate_engine, fam, hf_token = request.hf_token)
|
||||
# Take the GPU from the chat backend only when this load will actually use it.
|
||||
# diffusers always does; a *force-native* sd.cpp load on a CUDA/XPU/MPS box does
|
||||
# too. But a native sd.cpp load on a pure-CPU host never touches the GPU, so
|
||||
# acquiring would evict the resident chat model for nothing -- skip the handoff.
|
||||
# Take the GPU from the chat backend only when this load will actually use it,
|
||||
# which is exactly the resolved device being non-CPU. diffusers on an accelerator
|
||||
# and a force-native sd.cpp load on CUDA/XPU/MPS both resolve to that device; a
|
||||
# native sd.cpp load on a pure-CPU host does not. Crucially, a CPU-only host with
|
||||
# no usable sd-cli falls back to diffusers ON CPU -- that also never touches GPU
|
||||
# memory, so keying off the engine name (not the device) would wrongly evict a
|
||||
# resident chat model for a load that can't use the GPU. Gate on the device.
|
||||
device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device)
|
||||
needs_gpu = active_engine_name() != ENGINE_SD_CPP or device != "cpu"
|
||||
needs_gpu = device != "cpu"
|
||||
if needs_gpu:
|
||||
# Then kick the (slow) load onto a background thread and return at once --
|
||||
# the client polls images/load-progress.
|
||||
|
|
@ -11216,6 +11217,8 @@ async def list_gallery_images(
|
|||
offset: int = 0,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
from pydantic import ValidationError
|
||||
|
||||
from core.inference import image_gallery
|
||||
|
||||
limit = max(1, min(limit, 200))
|
||||
|
|
@ -11223,10 +11226,18 @@ async def list_gallery_images(
|
|||
# Fetch one extra to learn whether more remain, without a second scan.
|
||||
records = await asyncio.to_thread(image_gallery.list_images, limit + 1, offset)
|
||||
has_more = len(records) > limit
|
||||
return GalleryListResponse(
|
||||
images = [GalleryImage(**r) for r in records[:limit]],
|
||||
has_more = has_more,
|
||||
)
|
||||
# Build the response per record and drop any that fail schema validation: a PNG
|
||||
# whose recipe chunk has all required keys but a wrong value type (e.g. a
|
||||
# hand-dropped or corrupted file) passes the presence-only read but would raise
|
||||
# inside GalleryImage(**r). Skipping it keeps one bad file from 500-ing the whole
|
||||
# gallery listing.
|
||||
images = []
|
||||
for r in records[:limit]:
|
||||
try:
|
||||
images.append(GalleryImage(**r))
|
||||
except ValidationError:
|
||||
continue
|
||||
return GalleryListResponse(images = images, has_more = has_more)
|
||||
|
||||
|
||||
@studio_router.get("/images/gallery/{image_id}/file")
|
||||
|
|
|
|||
|
|
@ -56,6 +56,20 @@ def test_detect_family_from_repo_id():
|
|||
assert detect_family("meta-llama/Llama-3-8B") is None
|
||||
|
||||
|
||||
def test_detect_family_edit_keyword_scoped_to_basename():
|
||||
from core.inference.diffusion_families import detect_family_for_pick
|
||||
|
||||
# A parent directory named `edit`/`kontext`/`inpaint` must NOT reject a valid
|
||||
# text-to-image file: only the model id / filename basename is scanned for the
|
||||
# edit keyword. A direct local pick arrives as (parent_dir, filename).
|
||||
assert detect_family("/models/edit") is None # dir alone is ambiguous
|
||||
assert detect_family_for_pick("/models/edit", "Z-Image-Turbo-Q4.gguf").name == "z-image"
|
||||
assert detect_family_for_pick("/models/kontext", "qwen-image-2512-Q4.gguf").name == "qwen-image"
|
||||
# But a genuine editing checkpoint (keyword in the id/filename) is still rejected.
|
||||
assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") is None
|
||||
assert detect_family_for_pick("/models/misc", "Qwen-Image-Edit-2511-Q4.gguf") is None
|
||||
|
||||
|
||||
def test_detect_family_override():
|
||||
assert detect_family("local/path", override = "z-image").name == "z-image"
|
||||
assert detect_family("local/path", override = "zimage").name == "z-image"
|
||||
|
|
@ -760,6 +774,11 @@ def test_validate_load_request(tmp_path):
|
|||
backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF")
|
||||
with pytest.raises(ValueError, match = "family"):
|
||||
backend.validate_load_request("meta/Llama-3", gguf_filename = "q.gguf")
|
||||
# A family-looking repo paired with a non-GGUF single-file name is rejected here,
|
||||
# BEFORE the route evicts chat and hands over the GPU (the background load would
|
||||
# otherwise be the first to notice README.md is not a checkpoint).
|
||||
with pytest.raises(ValueError, match = r"\.gguf"):
|
||||
backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "README.md")
|
||||
assert (
|
||||
backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "q.gguf").name
|
||||
== "z-image"
|
||||
|
|
|
|||
|
|
@ -135,6 +135,26 @@ def test_runtime_env_handles_missing_lib_path():
|
|||
assert env[var] == "/opt/sdcpp/bin"
|
||||
|
||||
|
||||
def test_terminate_reaps_killed_child():
|
||||
# Cancellation/timeout paths call _terminate then immediately raise, so it must
|
||||
# reap the killed child itself -- otherwise a burst of image cancellations leaves
|
||||
# zombies until a later Popen cleanup. After _terminate the returncode is set
|
||||
# (the child has been waited on), so nothing lingers.
|
||||
import subprocess
|
||||
|
||||
proc = subprocess.Popen(
|
||||
[sys.executable, "-c", "import time; time.sleep(30)"],
|
||||
start_new_session = (os.name == "posix"),
|
||||
)
|
||||
try:
|
||||
eng._terminate(proc)
|
||||
assert proc.returncode is not None
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
|
||||
|
||||
# ── generate (fake subprocess) ──────────────────────────────────────────────
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -83,6 +83,11 @@ const CHAT_ONLY_ALLOWED = new Set([
|
|||
function isChatOnlyAllowed(pathname: string): boolean {
|
||||
if (CHAT_ONLY_ALLOWED.has(pathname)) return true;
|
||||
if (pathname === "/data-recipes" || pathname.startsWith("/data-recipes/")) return true;
|
||||
// Images runs on CPU/MPS via the native sd.cpp engine, which is exactly the
|
||||
// no-GPU (chat-only) setup it was added for. The generic chat-only flag is about
|
||||
// training/export needing a GPU, so it must not redirect /images away here or the
|
||||
// native image path is unreachable on the hosts that need it.
|
||||
if (pathname === "/images" || pathname.startsWith("/images/")) return true;
|
||||
return false;
|
||||
}
|
||||
|
||||
|
|
@ -159,6 +164,13 @@ function RootLayout() {
|
|||
setImagesMounted(true);
|
||||
}
|
||||
const shouldMountImages = isImagesRoute || imagesMounted;
|
||||
// Chat and Images both render their own full-height shell (a fixed top rail + an
|
||||
// internally-scrolling body), so both want the chat-style layout: no outer pt-14
|
||||
// inset and no outer scroll. Keying the layout off isChatRoute alone gave /images
|
||||
// the non-chat pt-14 + outer overflow, pushing its picker down and clipping the
|
||||
// bottom gallery. Treat them the same for the container padding/overflow only; the
|
||||
// keep-alive mounts below stay keyed to each specific route.
|
||||
const isChatLike = isChatRoute || isImagesRoute;
|
||||
|
||||
useTrainingUnloadGuard();
|
||||
// Global export driver: streams worker logs and tracks status from any route
|
||||
|
|
@ -251,10 +263,10 @@ function RootLayout() {
|
|||
className="!min-h-0 h-[calc(100dvh-var(--studio-titlebar-height,0px))] overflow-hidden"
|
||||
>
|
||||
<AppSidebar />
|
||||
<SidebarInset className={isChatRoute ? "overflow-hidden" : "overflow-y-auto"}>
|
||||
<SidebarInset className={isChatLike ? "overflow-hidden" : "overflow-y-auto"}>
|
||||
<Navbar />
|
||||
<div
|
||||
className={`relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"} ${isChatRoute ? "" : "pt-14 md:pt-[var(--studio-non-chat-content-top-inset,var(--studio-content-top-inset,0px))] md:[--studio-titlebar-height:var(--studio-non-chat-content-top-inset,var(--studio-content-top-inset,0px))]"}`}
|
||||
className={`relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatLike ? "overflow-hidden" : "overflow-visible"} ${isChatLike ? "" : "pt-14 md:pt-[var(--studio-non-chat-content-top-inset,var(--studio-content-top-inset,0px))] md:[--studio-titlebar-height:var(--studio-non-chat-content-top-inset,var(--studio-content-top-inset,0px))]"}`}
|
||||
>
|
||||
{/* Stays mounted across navigation so an in-flight generation is
|
||||
not cancelled when leaving /chat; hidden (not unmounted) off-route.
|
||||
|
|
@ -280,7 +292,7 @@ function RootLayout() {
|
|||
<div
|
||||
className={
|
||||
isImagesRoute
|
||||
? "flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-visible"
|
||||
? "flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden"
|
||||
: "hidden"
|
||||
}
|
||||
inert={!isImagesRoute || undefined}
|
||||
|
|
|
|||
|
|
@ -452,6 +452,12 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
const loadToastId = useRef<string | number | null>(null);
|
||||
// Last load-progress signature shown, so a tick that moved nothing skips the toast.
|
||||
const lastLoadSig = useRef<string | null>(null);
|
||||
// The quant to restore if the current optimistic swap fails. A same-repo quant
|
||||
// change sets `quant` immediately for picker feedback; if the load then fails
|
||||
// AFTER starting (an error/eviction during download), the old pipeline stays
|
||||
// loaded, so the poll must roll the label back rather than advertise the failed
|
||||
// quant. `{ prev }` distinguishes "revert to null" from "nothing pending".
|
||||
const quantRevert = useRef<{ prev: string | null } | null>(null);
|
||||
|
||||
const dismissLoadToast = useCallback(() => {
|
||||
if (loadToastId.current != null) toast.dismiss(loadToastId.current);
|
||||
|
|
@ -644,12 +650,22 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
setStatus(await getDiffusionStatus());
|
||||
toast.success("Model loaded");
|
||||
setBusy(null);
|
||||
// Load succeeded: the optimistic quant is now the real one, so drop the
|
||||
// pending revert.
|
||||
quantRevert.current = null;
|
||||
return;
|
||||
}
|
||||
if (p.phase === "error") {
|
||||
dismissLoadToast();
|
||||
toast.error(p.error || "Failed to load model");
|
||||
setBusy(null);
|
||||
// A load that failed AFTER starting leaves the previous pipeline loaded, so
|
||||
// roll the optimistic quant label back to what is actually loaded (status
|
||||
// does not carry the quant, so refreshStatus alone can't correct it).
|
||||
if (quantRevert.current) {
|
||||
setQuant(quantRevert.current.prev);
|
||||
quantRevert.current = null;
|
||||
}
|
||||
// A failed load may have freed a previously-loaded model, so resync to
|
||||
// the real backend state (the synchronous failure path does the same).
|
||||
void refreshStatus();
|
||||
|
|
@ -662,6 +678,11 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
// busy stuck on "loading", deadening the picker and Generate button.
|
||||
dismissLoadToast();
|
||||
setBusy(null);
|
||||
// Same optimistic-quant rollback as the error path: the swap did not take.
|
||||
if (quantRevert.current) {
|
||||
setQuant(quantRevert.current.prev);
|
||||
quantRevert.current = null;
|
||||
}
|
||||
void refreshStatus();
|
||||
return;
|
||||
}
|
||||
|
|
@ -751,15 +772,21 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
if (busy !== null) return;
|
||||
if (meta.ggufVariant && meta.ggufFilename) {
|
||||
// Optimistic for instant picker feedback, but revert if the load fails to
|
||||
// START (400/409/network): the selector must not advertise a quant that
|
||||
// is not the loaded one. Poll-phase failures re-sync via refreshStatus.
|
||||
// START (400/409/network) or LATER during the poll (download/preflight
|
||||
// error/eviction) -- in both cases the old pipeline stays loaded, so the
|
||||
// selector must not advertise the failed quant. The poll owns the after-start
|
||||
// revert via quantRevert; here we only handle the never-started case.
|
||||
const prevQuant = quant;
|
||||
quantRevert.current = { prev: prevQuant };
|
||||
setQuant(meta.ggufVariant);
|
||||
const d = defaultsFor(id);
|
||||
setSteps(d.steps);
|
||||
setGuidance(d.guidance);
|
||||
void handleLoad(id, meta.ggufFilename).then((started) => {
|
||||
if (!started) setQuant(prevQuant);
|
||||
if (!started) {
|
||||
setQuant(prevQuant);
|
||||
quantRevert.current = null;
|
||||
}
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
|
@ -774,14 +801,19 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
|
|||
if (!filename.toLowerCase().endsWith(".gguf")) return;
|
||||
// A direct pick carries no curated variant label; surface the filename so
|
||||
// the selector stops advertising the previously loaded quant. Optimistic,
|
||||
// reverted if the load fails to start (mirrors the curated branch above).
|
||||
// reverted if the load fails to start OR fails later in the poll (mirrors the
|
||||
// curated branch above; the poll owns the after-start revert via quantRevert).
|
||||
const prevQuant = quant;
|
||||
quantRevert.current = { prev: prevQuant };
|
||||
setQuant(filename);
|
||||
const d = defaultsFor(id);
|
||||
setSteps(d.steps);
|
||||
setGuidance(d.guidance);
|
||||
void handleLoad(dir, filename).then((started) => {
|
||||
if (!started) setQuant(prevQuant);
|
||||
if (!started) {
|
||||
setQuant(prevQuant);
|
||||
quantRevert.current = null;
|
||||
}
|
||||
});
|
||||
}
|
||||
},
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue