Match edit keywords by id segment and sanitize generate-route errors

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):

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.")

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
 
 

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.")

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