From a0472d2bffefcd9dce68d361a8a0b52292a7b610 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Fri, 26 Jun 2026 01:44:59 -0300 Subject: [PATCH] Match edit keywords by id segment and sanitize generate-route errors --- .../backend/core/inference/diffusion_families.py | 11 +++++++++-- studio/backend/routes/inference.py | 10 ++++++++-- studio/backend/tests/test_diffusion_backend.py | 6 +++++- studio/backend/tests/test_diffusion_routes.py | 16 ++++++++++++++++ .../assistant-ui/model-selector/pickers.tsx | 9 ++++++--- 5 files changed, 44 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 51987061a0..450fa6f63a 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -14,6 +14,7 @@ diffusers classes and base repo needed to assemble the full pipeline. from __future__ import annotations +import re from dataclasses import dataclass, field from pathlib import Path, PurePosixPath from typing import Optional @@ -75,7 +76,7 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( # Editing / inpaint checkpoints share an arch keyword but need a different # pipeline and an input image, which this text-to-image backend doesn't drive. -_EDIT_KEYWORDS = ("edit", "kontext", "inpaint") +_EDIT_KEYWORDS = ("edit", "kontext", "inpaint", "inpainting") def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[DiffusionFamily]: @@ -92,7 +93,13 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff return fam return None needle = repo_id.lower() - if any(kw in needle for kw in _EDIT_KEYWORDS): + # 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)) + if any(kw in segments for kw in _EDIT_KEYWORDS): return None for fam in _FAMILIES: if fam.name in needle or any(alias in needle for alias in fam.aliases): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 91da38109e..ca3c91cb98 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10104,8 +10104,14 @@ async def generate_diffusion_image( batch_size = request.batch_size, ) except RuntimeError as exc: - # No model loaded (or unloaded mid-flight) — a client-state problem. - raise HTTPException(status_code = 409, detail = str(exc)) + if not backend.is_loaded: + # The only genuine client-state 409: nothing is loaded to generate with. + raise HTTPException(status_code = 409, detail = "No diffusion model is loaded.") + # A pipeline RuntimeError (CUDA OOM, shape/device) is a server failure; fall + # through to the sanitized 500 instead of echoing raw exception text (which + # would 409 an OOM as retryable and leak VRAM totals / tensor shapes). + logger.error("diffusion.generate_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Image generation failed.") except Exception as exc: logger.error("diffusion.generate_failed: %s", exc) raise HTTPException(status_code = 500, detail = "Image generation failed.") diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index e38595bb01..5fc5d23bbe 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -47,9 +47,13 @@ def test_detect_family_from_repo_id(): # Qwen-Image guides via true_cfg_scale, not guidance_scale. assert detect_family("unsloth/Qwen-Image-2512-GGUF").cfg_kwarg == "true_cfg_scale" assert detect_family("unsloth/Z-Image-GGUF").cfg_kwarg == "guidance_scale" - # Image-editing checkpoints are rejected (text-to-image backend only). + # Image-editing checkpoints are rejected (text-to-image backend only): the + # edit keyword is matched as a whole id segment, so an "edit" that's only a + # substring of a normal word ("Edition") still loads. assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") is None assert detect_family("unsloth/FLUX.1-Kontext-dev-GGUF") is None + assert detect_family("unsloth/Qwen-Image-Inpainting-GGUF") is None + assert detect_family("unsloth/Z-Image-Edition-GGUF").name == "z-image" assert detect_family("meta-llama/Llama-3-8B") is None diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 5de10ede19..d22bdc57e7 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -223,6 +223,22 @@ def test_generate_without_load_returns_409(client): assert resp.status_code == 409 +def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch): + # A loaded model that fails mid-pipeline (CUDA OOM, a RuntimeError) is a server + # failure: 500 with a generic message, not a 409 echoing the raw exception. + backend = diffusion_module.get_diffusion_backend() + backend.loaded = True + + def _oom(**kwargs): + raise RuntimeError("CUDA out of memory. Tried to allocate 20.00 GiB (24.00 GiB total)") + + monkeypatch.setattr(backend, "generate", _oom) + resp = client.post("/api/inference/images/generate", json = {"prompt": "p"}) + assert resp.status_code == 500 + assert resp.json()["detail"] == "Image generation failed." + assert "CUDA" not in resp.json()["detail"] + + def test_load_unknown_family_returns_400(client, monkeypatch): def _raise(*a, **k): raise ValueError("'x/y' isn't a supported image-generation model. Supported: Z-Image.") diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 5dfc22be14..fab86270ff 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -990,11 +990,14 @@ export const IMAGE_GEN_TASKS = [ // id so they don't show in the Images picker only to 400 on load. Keeping the // image-to-image task itself is required: some supported models (FLUX.2-klein) // carry that tag too. -const IMAGE_EDIT_KEYWORDS = ["edit", "kontext", "inpaint"] as const; +const IMAGE_EDIT_KEYWORDS = ["edit", "kontext", "inpaint", "inpainting"] as const; function isImageEditModel(repoId: string | null | undefined): boolean { if (!repoId) return false; - const id = repoId.toLowerCase(); - return IMAGE_EDIT_KEYWORDS.some((kw) => id.includes(kw)); + // Whole-segment match (not substring) so a normal model like "...-edition" + // isn't hidden; mirrors the backend detect_family segment check. Split on + // both path separators so a Windows local path is segmented too. + const segments = new Set(repoId.toLowerCase().split(/[-_./\\]+/)); + return IMAGE_EDIT_KEYWORDS.some((kw) => segments.has(kw)); } // Gate an on-device model by the picker's task scope. With a filter (the Images