Fix/adjust diffusion: round 26 P1 batch for PR #5754

Twelve P1 findings from round 26 reviewer aggregate, plus the CI
revert of round 25 P1 #5 to a less invasive location.

1. requirements/studio.txt + requirements/single-env/constraints.txt:
   revert the round 25 huggingface-hub bump (broke Studio Update CI,
   Mac Studio Update CI, Mac Studio UI CI, Studio UI CI all with
   ResolutionImpossible against transformers==4.57.6 which requires
   hub<1.0). Standard install path stays on the well-tested 4.57.6 +
   0.36.2 + trl 0.23.1 trio.

2. requirements/no-torch-runtime.txt + pyproject.toml
   [huggingfacenotorch]: bump huggingface_hub floor from >=0.34.0 to
   >=1.3.0,<2.0 -- this is where the actual transformers 5.x +
   hub 0.36.2 broken combo can land because the file installs
   --no-deps. transformers 5.x calls hub.is_offline_mode which only
   exists in hub 1.x.

3. utils/datasets/llm_assist.py: revert round 25 P1 #4 (helper/advisor
   sharing the global llama backend) which introduced three
   regressions: a chat-evict load race after the busy precheck, a
   finally-block that could unload a user chat model, and an
   identifier mismatch the delete guard could not canonicalize. Go
   back to PRIVATE LlamaCppBackend instances and expose the active
   helper/advisor repos through a new thread-safe registry
   (helper_advisor_owns_repo / _register_helper_advisor_repo /
   _unregister_helper_advisor_repo) so DELETE /api/models/delete-cached
   can still block the rmtree.

4. routes/models.py delete_cached_model: check the new helper/advisor
   registry up front and 409 if a helper/advisor still owns the
   target repo. Closes round 26 P1 #13 and #14 (helper/advisor
   identifiers were prefixed and would never equal the raw repo id).

5. routes/models.py get_lora_base_model: validate lora_path with
   _validate_logged_identifier before it is reflected in 404 detail
   and error logs (round 26 P1 #12).

6. routes/inference.py /unload: round 21 P1 #3 added a "or not
   is_loaded" fallback that let an unload of owner/B cancel a pending
   llama load of owner/A. Replace it with a narrow
   llama_is_starting_without_identifier branch that only fires when
   llama-server is mid-startup with neither identifier set (round 26
   P1 #5).

7. routes/inference.py /unload: poll loading_model_identifier for up
   to 5 s after asyncio.to_thread(unload_model) so a legitimate
   pending-load cancel does not 503 because the load thread has not
   yet observed _cancel_event in its finally (round 26 P2 #15).

8. models/training.py TrainingStartRequest: extend identifier
   hardening to hf_dataset, subset, train_split, eval_split. Round 22
   only guarded model_name (round 26 P1 #10).

9. models/data_recipe.py SeedInspectRequest: add _no_control_chars +
   _reject_embedded_hf_token field_validators on dataset_name (round
   26 P1 #11).

Tests: 105 targeted (diffusion + cached_gguf + llama_cpp_cache +
inference_model_validation + models_get_model_config) and 1768
broader backend tests pass locally. Pre-existing
test_desktop_auth.py, test_studio_api.py, and
test_training_worker_flash_attn.py failures reproduce on HEAD
without these changes.

pyproject.toml

@@ -81,7 +81,10 @@ huggingfacenotorch = [
     "datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0",
     "accelerate>=0.34.1",
     "peft>=0.18.0,!=0.11.0",
-    "huggingface_hub>=0.34.0",
+    # Round 26 P1 #9: floor at 1.3.0 because the diffusion stack below
+    # pulls transformers 5.x which calls hub.is_offline_mode (hub 1.x).
+    # Keep <2.0 to avoid any future hub ABI break.
+    "huggingface_hub>=1.3.0,<2.0",
     "hf_transfer",
     # Studio Images page depends on Flux2KleinPipeline /
     # Flux2Pipeline, both shipped in diffusers>=0.37.0. Floor was

studio/backend/models/data_recipe.py

@@ -90,6 +90,18 @@ class SeedInspectRequest(BaseModel):
     split: str | None = "train"
     preview_size: int = Field(default = 10, ge = 1, le = 50)
 
+    # Round 26 P1 #11: dataset_name reaches HF + log/echo paths, so
+    # mirror the hardening other dataset request models already do.
+    @field_validator("dataset_name")
+    @classmethod
+    def _no_dataset_name_control_chars(cls, v, info):
+        return _no_control_chars(v, info.field_name)
+
+    @field_validator("dataset_name")
+    @classmethod
+    def _no_dataset_name_embedded_hf_tokens(cls, v, info):
+        return _reject_embedded_hf_token(v, info.field_name)
+
 
 class SeedInspectUploadRequest(BaseModel):
     # Legacy single-file flow (mutually exclusive with file_ids)

studio/backend/models/training.py

@@ -60,12 +60,14 @@ class TrainingStartRequest(BaseModel):
     # Round 22 P1 #1: identifier hardening (round 5 / 15 / 20 / 21
     # extended these to chat + diffusion request models; training
     # was the last unguarded entry point).
-    @field_validator("model_name")
+    # Round 26 P1 #10: hf_dataset / subset / train_split / eval_split
+    # are reflected in status + error messages, harden them too.
+    @field_validator("model_name", "hf_dataset", "subset", "train_split", "eval_split")
     @classmethod
     def _no_model_name_control_chars(cls, v, info):
         return _no_control_chars(v, info.field_name)
 
-    @field_validator("model_name")
+    @field_validator("model_name", "hf_dataset")
     @classmethod
     def _no_model_name_embedded_hf_tokens(cls, v, info):
         return _reject_embedded_hf_token(v, info.field_name)

studio/backend/requirements/no-torch-runtime.txt

@@ -43,7 +43,13 @@ safetensors>=0.4.3
 datasets>=3.4.1,!=4.0.*,!=4.1.0,<4.4.0
 accelerate>=0.34.1
 peft>=0.18.0,!=0.11.0
-huggingface_hub>=0.34.0
+# Round 26 P1 #8: floor at 1.3.0 because transformers 5.x (allowed by
+# the range below) calls huggingface_hub.is_offline_mode, only present
+# in hub 1.x. Under --no-deps the resolver does not enforce this
+# transitively, so a pre-existing 0.36.2 used to be kept and the next
+# `from transformers import AutoConfig` raised ImportError. Upper bound
+# <2.0 keeps us off any future ABI break.
+huggingface_hub>=1.3.0,<2.0
 hf_transfer
 # Floor 0.37.0 introduces Flux2KleinPipeline + Flux2Pipeline which the
 # Studio Images page imports for the default curated picker.

studio/backend/requirements/single-env/constraints.txt

@@ -2,12 +2,7 @@
 # Keep compatible with unsloth transformers bounds.
 transformers==4.57.6
 trl==0.23.1
-# Round 25 P1 #5 follow-up: bumped from 0.36.2 because studio.txt now
+huggingface-hub==0.36.2
-# requires >=1.3.0,<2.0 (Flux2KleinPipeline transitively imports
-# transformers 5.x which needs huggingface_hub.is_offline_mode,
-# introduced in hub 1.x). 1.8.0 matches the explicit pin used by
-# studio/setup.sh / setup.ps1 for the t5 sub-envs.
-huggingface-hub==1.8.0
 
 # Studio stack
 datasets==4.3.0

studio/backend/requirements/studio.txt

@@ -17,25 +17,7 @@ pyjwt
 easydict
 addict
 # gradio>=4.0.0                  # 148 MB - Studio uses React + FastAPI, not Gradio
-# Round 25 P1 #5: keep the Studio Images dependency set internally
+huggingface-hub==0.36.2
-# compatible. ``diffusers>=0.37.0`` ships Flux2KleinPipeline /
-# Flux2Pipeline, which transitively import the newer ``transformers``
-# (>=4.56) that requires ``huggingface_hub.is_offline_mode`` -- only
-# available in ``huggingface_hub>=1.0``. The previous ``==0.36.2``
-# pin let fresh installs end up with ``transformers 5.x`` +
-# ``huggingface_hub 0.36.2``, which crashed on the first
-# ``/api/inference/images/load`` with
-# ``Flux2KleinPipeline ... no attribute 'is_offline_mode'``. Bump
-# the floor so ``diffusers`` and ``transformers`` resolve into a
-# runtime they can actually import.
-huggingface-hub>=1.3.0,<2.0
-# Mirror the ``transformers`` constraint from
-# ``no-torch-runtime.txt``. Without it, the standard install can
-# resolve ``transformers 5.4.0+`` which drops Studio-supported
-# trainers. ``tokenizers<=0.23.0`` is required because
-# ``transformers 4.56..5.3`` declares it explicitly.
-tokenizers<=0.23.0
-transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.3.0
 structlog>=24.1.0
 diceware
 ddgs

studio/backend/routes/inference.py

@@ -1683,9 +1683,27 @@ async def unload_model(
             or is_registered_native_path_label(loaded_identifier, request.model_path)
             or is_registered_native_path_label(loading_identifier, request.model_path)
         )
-        if (getattr(llama_backend, "is_active", False) or loading_identifier) and (
+        # Round 26 P1 #5: the previous ``or not is_loaded`` fallback
-            llama_matches_request or not getattr(llama_backend, "is_loaded", False)
+        # let an unload of ``owner/B`` cancel a pending llama download
-        ):
+        # of ``owner/A`` and silently leave safetensors ``owner/B``
+        # alive. Only enter the llama branch when the request actually
+        # matches the loaded/loading identifier, OR when llama-server
+        # is starting up without any identifier yet (the original
+        # narrow case we wanted to catch).
+        llama_is_starting_without_identifier = (
+            getattr(llama_backend, "is_active", False)
+            and not getattr(llama_backend, "is_loaded", False)
+            and not loaded_identifier
+            and not loading_identifier
+        )
+        should_unload_llama = (
+            llama_matches_request
+            and (
+                getattr(llama_backend, "is_active", False)
+                or loading_identifier
+            )
+        ) or llama_is_starting_without_identifier
+        if should_unload_llama:
             # Round 19 P1 #6: previously this called
             # ``llama_backend.unload_model()`` and unconditionally
             # returned ``status="unloaded"`` even when the subprocess
@@ -1694,6 +1712,16 @@ async def unload_model(
             # still resident. Treat ``False`` / leftover state as a
             # 503 so the user retries.
             ok = await asyncio.to_thread(llama_backend.unload_model)
+            # Round 26 P2 #15: explicit cancel of a pending GGUF load
+            # leaves loading_model_identifier set briefly until the
+            # load thread observes _cancel_event in its finally. Wait
+            # up to 5s so a legitimate cancel does not 503.
+            deadline = time.monotonic() + 5.0
+            while (
+                getattr(llama_backend, "loading_model_identifier", None)
+                and time.monotonic() < deadline
+            ):
+                await asyncio.sleep(0.1)
             if (
                 ok is False
                 or getattr(llama_backend, "is_loaded", False)

studio/backend/routes/models.py

@@ -2219,6 +2219,9 @@ async def get_lora_base_model(
 
     This endpoint wraps the backend get_base_model_from_lora function.
     """
+    # Round 26 P1 #12: lora_path is echoed back in 404 detail and logs;
+    # harden it the same way other reflected identifiers are.
+    lora_path = _validate_logged_identifier(lora_path, "lora_path")
     try:
         base_model = get_base_model_from_lora(lora_path)
 
@@ -2853,6 +2856,31 @@ async def delete_cached_model(
                 continue
         return False
 
+    # Round 26 P1 #13 / #14: helper/advisor GGUF loads run on a
+    # PRIVATE LlamaCppBackend, so the global backend below cannot see
+    # them. utils/datasets/llm_assist.py publishes the active repo
+    # via helper_advisor_owns_repo() for exactly this guard. Fail
+    # closed on the variant question (block any variant of the repo)
+    # because helper/advisor flows do not pass a variant through.
+    try:
+        from utils.datasets.llm_assist import helper_advisor_owns_repo
+
+        if helper_advisor_owns_repo(repo_id):
+            raise HTTPException(
+                status_code = 409,
+                detail = "Cannot delete a model while AI Assist is using it",
+            )
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.warning(
+            "Could not check helper/advisor backend status before cache delete: %s", e
+        )
+        raise HTTPException(
+            status_code = 503,
+            detail = "Could not verify AI Assist load status before deleting cache",
+        ) from e
+
     # Check if model is currently loaded OR loading. is_active and
     # not is_loaded means an llama-server download / startup is in
     # flight; the cache delete would race the hf_hub_download / mmap.

studio/backend/utils/datasets/llm_assist.py

@@ -18,6 +18,7 @@ import logging
 import os
 import re
 import textwrap
+import threading
 import time
 from itertools import islice
 from typing import Any, Optional
@@ -31,6 +32,41 @@ DEFAULT_HELPER_MODEL_VARIANT = "UD-Q4_K_XL"
 
 README_MAX_CHARS = 1500
 
+# Round 26 P1 #13 / #14: helper/advisor run on PRIVATE LlamaCppBackend
+# instances (round 25 P1 #4 briefly used the global singleton, which
+# 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()
+_HELPER_ADVISOR_LOCK = threading.Lock()
+
+
+def helper_advisor_owns_repo(repo_id: str) -> bool:
+    """Return True if any helper/advisor load currently owns this
+    HF repo id. Comparison is case-insensitive to match the chat
+    backend's lowercased needle."""
+    if not repo_id:
+        return False
+    needle = repo_id.lower()
+    with _HELPER_ADVISOR_LOCK:
+        return needle in _HELPER_ADVISOR_ACTIVE_REPOS
+
+
+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())
+
+
+def _unregister_helper_advisor_repo(repo_id: str) -> None:
+    if not repo_id:
+        return
+    with _HELPER_ADVISOR_LOCK:
+        _HELPER_ADVISOR_ACTIVE_REPOS.discard(repo_id.lower())
+
 
 def _strip_think_tags(text: str) -> str:
     """Strip <think>...</think> reasoning blocks emitted by some models.
@@ -244,19 +280,17 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
     )
 
     backend = None
+    _register_helper_advisor_repo(repo)
     try:
-        # Round 25 P1 #4: use the GLOBAL llama backend instead of a
+        # Round 26 P1 #1 / #3 / #13 / #14: use a PRIVATE backend so the
-        # private ``LlamaCppBackend()`` instance. The private instance
+        # helper can never preempt or be preempted by the user's
-        # was invisible to ``DELETE /api/models/delete-cached`` and the
+        # chat backend and cannot accidentally unload it in finally.
-        # other global delete guards because they inspect the singleton
+        # The active repo is published via _register_helper_advisor_repo
-        # returned by ``get_llama_cpp_backend()``. A concurrent cache
+        # above so DELETE /api/models/delete-cached can still block the
-        # delete could rmtree the helper's mid-flight download or
+        # cache rmtree while the helper is downloading or mmap'ing.
-        # mmap'd snapshot. ``_gpu_workload_busy_for_helper`` above
+        from core.inference.llama_cpp import LlamaCppBackend
-        # already ensures the global backend is idle before we reach
-        # here, so taking it over is safe.
-        from routes.inference import get_llama_cpp_backend
 
-        backend = get_llama_cpp_backend()
+        backend = LlamaCppBackend()
         logger.info(f"Loading helper model: {repo} ({variant})")
 
         ok = backend.load_model(
@@ -305,6 +339,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
                 logger.info("Helper model unloaded")
             except Exception:
                 pass
+        _unregister_helper_advisor_repo(repo)
 
 
 # ─── Public API ───────────────────────────────────────────────────────
@@ -649,16 +684,15 @@ def _run_multi_pass_advisor(
     )
 
     backend = None
+    _register_helper_advisor_repo(repo)
     try:
-        # Round 25 P1 #4: mirror ``_run_with_helper`` and acquire the
+        # Round 26 P1 #2 / #4 / #13 / #14: mirror ``_run_with_helper``
-        # GLOBAL llama backend so cache-delete and unload guards see
+        # and use a PRIVATE backend. Round 25's global-backend swap
-        # this advisor load via the singleton's
+        # introduced chat-evict races and finally-eviction bugs.
-        # ``loading_model_identifier`` / ``model_identifier``. The
+        # The registry above keeps delete-cache safe.
-        # round 23/24 ``_gpu_workload_busy_for_helper`` already
+        from core.inference.llama_cpp import LlamaCppBackend
-        # blocks reach here unless the global llama backend is idle.
-        from routes.inference import get_llama_cpp_backend
 
-        backend = get_llama_cpp_backend()
+        backend = LlamaCppBackend()
         logger.info(f"Loading advisor model: {repo} ({variant})")
         t0 = time.monotonic()
 
@@ -990,6 +1024,7 @@ def _run_multi_pass_advisor(
                 logger.info("Advisor model unloaded")
             except Exception:
                 pass
+        _unregister_helper_advisor_repo(repo)
 
 
 def llm_conversion_advisor(