Fix/adjust diffusion: round 27 P1 + P2 batch for PR unslothai#5754

danielhanchen · danielhanchen · commit 6c528fb01386 · 2026-05-25T13:39:40.000Z
Round 27 findings (Opus parallel concurrency + frontend reviews). Backend P1 fixes: 1. utils/datasets/llm_assist.py: the round 26 helper/advisor active registry used a plain set, so two concurrent helper / advisor loads of the same DEFAULT_HELPER_MODEL_REPO would both set.add() (no-op the second time) and then the first finally set.discard() would underflow the registration while the second call was still mmap'ing the GGUF. Switch to a Counter with proper refcount increment/decrement so the repo stays registered until the last user releases it. 2. routes/inference.py _release_chat_for and core/inference/diffusion.py _release_chat_backend_for_diffusion: helper/advisor GGUF runs on a PRIVATE LlamaCppBackend (round 26 P1 #1), so the global llama checks below could not see them. A user-driven /training/start, /export/load-checkpoint, or /images/load would skip the unload and allocate FLUX VRAM on top of the helper's resident weights, OOMing on 16-24 GB consumer GPUs. Both release paths now consult helper_advisor_busy() and fail 503 (or RuntimeError for the in-backend path) so the user retries instead of double-owning VRAM. Frontend P2 fixes: 3. studio/frontend/src/features/images/images-page.tsx: handleUnload now calls refreshStatus() in the catch path so a partial unload (503 from the backend) does not leave the UI showing a stale "Loaded:" label. Matches the handleLoad pattern. 4. images-page.tsx: when status.is_loading is true, auto-poll refreshStatus every 2 s so the user sees real progress instead of a frozen "Loading..." label until they manually click Refresh. 5. images-page.tsx: aria-label="Inference steps" / "Guidance scale" on the two sliders so screen readers can announce them. 6. images-page.tsx: defensive (r.guidance_scale ?? 0).toFixed(1) in the results caption so a future backend that serialises NaN/None for guidance does not throw at render. Tests: 105 targeted (diffusion + cached_gguf + inference_validation) and 1768 broader backend tests pass locally. Frontend `npm run typecheck` passes.
diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
@@ -1516,6 +1516,20 @@ def _release_chat_backend_for_diffusion() -> None:
     diffusion ``load_model`` bails out instead of double-owning VRAM
     (round 17 P1 #2).
     """
+    # Round 27 P1 #2: helper / advisor GGUF loads run on a PRIVATE
+    # LlamaCppBackend so the global llama check below cannot see them.
+    # Refuse the diffusion handoff while a helper / advisor still owns
+    # its private backend so we do not allocate FLUX VRAM on top.
+    try:
+        from utils.datasets.llm_assist import helper_advisor_busy
+    except Exception:
+        pass
+    else:
+        if helper_advisor_busy():
+            raise RuntimeError(
+                "AI Assist (helper / advisor GGUF) is still using the GPU. "
+                "Wait for it to finish before loading a diffusion image model."
+            )
     # 1. GGUF chat backend (llama-server subprocess). We unload when
     #    EITHER is_loaded is True (resident model) OR is_active is
     #    True (mid-download / startup) OR loading_model_identifier is
diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
@@ -531,6 +531,24 @@ async def _release_chat_for(workload: str) -> None:
     start. Conversely, the standard chat-load path releases only
     the llama side.
     """
+    # Round 27 P1 #2: helper / advisor GGUF loads run on a PRIVATE
+    # LlamaCppBackend (round 26 P1 #1) so the global llama checks
+    # below do not see them. Refuse the handoff while a helper /
+    # advisor still owns its private backend so a new GPU workload
+    # does not allocate on top of helper VRAM and OOM.
+    try:
+        from utils.datasets.llm_assist import helper_advisor_busy
+    except Exception:
+        pass
+    else:
+        if helper_advisor_busy():
+            raise HTTPException(
+                status_code = 503,
+                detail = (
+                    f"AI Assist (helper / advisor GGUF) is still using the "
+                    f"GPU. Wait for it to finish before starting {workload}."
+                ),
+            )
     await _release_llama_for(workload)
     await _release_safetensors_chat_for(workload)
 
diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py
@@ -20,6 +20,7 @@
 import textwrap
 import threading
 import time
+from collections import Counter
 from itertools import islice
 from typing import Any, Optional
 
@@ -37,9 +38,15 @@
 # caused chat-evict races and finally-eviction bugs and still left
 # delete-cache blind because helper/advisor publish prefixed
 # identifiers the guard could not match). Expose loading repo ids
-# through a thread-safe set so DELETE /api/models/delete-cached can
-# block while a helper or advisor still owns the cache.
-_HELPER_ADVISOR_ACTIVE_REPOS: set[str] = set()
+# through a thread-safe Counter so DELETE /api/models/delete-cached
+# can block while a helper or advisor still owns the cache.
+#
+# Round 27 P1 #1: must refcount, not a plain set. A helper and an
+# advisor (or two concurrent helpers) often share the default repo
+# unsloth/gemma-4-E2B-it-GGUF. With a set, the first finally call
+# discarded the repo while the second invocation was still loading,
+# and the delete-cache guard then let rmtree race the live mmap.
+_HELPER_ADVISOR_REFCOUNT: Counter[str] = Counter()
 _HELPER_ADVISOR_LOCK = threading.Lock()
 
 
@@ -51,21 +58,33 @@ def helper_advisor_owns_repo(repo_id: str) -> bool:
         return False
     needle = repo_id.lower()
     with _HELPER_ADVISOR_LOCK:
-        return needle in _HELPER_ADVISOR_ACTIVE_REPOS
+        return _HELPER_ADVISOR_REFCOUNT.get(needle, 0) > 0
+
+
+def helper_advisor_busy() -> bool:
+    """Round 27 P1 #2: True if ANY helper/advisor load is in flight.
+    Used by diffusion / training / export release paths so they do
+    not allocate on top of the helper's VRAM while it owns its
+    private LlamaCppBackend instance."""
+    with _HELPER_ADVISOR_LOCK:
+        return sum(_HELPER_ADVISOR_REFCOUNT.values()) > 0
 
 
 def _register_helper_advisor_repo(repo_id: str) -> None:
     if not repo_id:
         return
     with _HELPER_ADVISOR_LOCK:
-        _HELPER_ADVISOR_ACTIVE_REPOS.add(repo_id.lower())
+        _HELPER_ADVISOR_REFCOUNT[repo_id.lower()] += 1
 
 
 def _unregister_helper_advisor_repo(repo_id: str) -> None:
     if not repo_id:
         return
+    needle = repo_id.lower()
     with _HELPER_ADVISOR_LOCK:
-        _HELPER_ADVISOR_ACTIVE_REPOS.discard(repo_id.lower())
+        _HELPER_ADVISOR_REFCOUNT[needle] -= 1
+        if _HELPER_ADVISOR_REFCOUNT[needle] <= 0:
+            _HELPER_ADVISOR_REFCOUNT.pop(needle, None)
 
 
 def _strip_think_tags(text: str) -> str:
diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx
@@ -148,6 +148,18 @@ export function ImagesPage() {
     void refreshStatus();
   }, [refreshStatus]);
 
+  // Round 27 P2: when the backend is mid-load (is_loading=true) the
+  // status label froze at "Loading..." until the user clicked
+  // Refresh. Auto-poll every 2 s while a load is in flight so the
+  // UI tracks real backend progress.
+  useEffect(() => {
+    if (!status?.is_loading) return;
+    const id = window.setInterval(() => {
+      void refreshStatus();
+    }, 2000);
+    return () => window.clearInterval(id);
+  }, [status?.is_loading, refreshStatus]);
+
   const handleLoad = useCallback(async () => {
     setBusy("loading");
     try {
@@ -207,6 +219,12 @@ export function ImagesPage() {
       toast.error("Failed to unload image model", {
         description: err instanceof Error ? err.message : String(err),
       });
+      // Round 27 P2: a partial unload (subprocess refused to terminate,
+      // 503 from the backend) used to leave the UI showing the old
+      // "Loaded:" label even though the backend state was half torn
+      // down. Refresh so the button states match reality (mirrors
+      // handleLoad above which always re-fetches on catch).
+      await refreshStatus();
     } finally {
       setBusy("idle");
     }
@@ -491,6 +509,7 @@ export function ImagesPage() {
             <div className="flex flex-col gap-1">
               <Label>Steps: {steps}</Label>
               <Slider
+                aria-label="Inference steps"
                 min={1}
                 max={60}
                 step={1}
@@ -501,6 +520,7 @@ export function ImagesPage() {
             <div className="flex flex-col gap-1">
               <Label>Guidance: {guidance.toFixed(1)}</Label>
               <Slider
+                aria-label="Guidance scale"
                 min={0}
                 max={15}
                 step={0.1}
@@ -554,7 +574,7 @@ export function ImagesPage() {
                   data-testid="diffusion-result-image"
                 />
                 <figcaption className="text-xs text-muted-foreground">
-                  {r.width}x{r.height} - {r.num_inference_steps} steps - g={r.guidance_scale.toFixed(1)}
+                  {r.width}x{r.height} - {r.num_inference_steps} steps - g={(r.guidance_scale ?? 0).toFixed(1)}
                   {/* Prefer seed_str (full uint64 precision) since the
                        numeric seed gets rounded by JSON.parse above
                        Number.MAX_SAFE_INTEGER and would otherwise
