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

P1 #1: ``_gpu_workload_busy_for_helper`` in
``utils/datasets/llm_assist.py`` now also gates on the GGUF chat
backend (llama-server) AND the safetensors chat backend. Round 23
extended it to training + export but missed Chat, so a helper /
advisor GGUF could still race a loaded chat model for VRAM.
Both checks fail closed when status is unverifiable.

P1 #2 / #3 / #4 / #5: re-ordered the route-level GPU-handoff
unloads so the diffusion release runs BEFORE the chat releases.
A wedged diffusion unload used to fire AFTER chat was already
gone, so the user lost both on a single failure. Drop chat last
so an earlier failure preserves it. Applied to
``/training/start`` (training.py), ``/export/load`` (export.py),
``/chat/load`` GGUF branch and ``/chat/load`` safetensors branch
(routes/inference.py).

P1 #7 + P2 #13: ``/delete-finetuned`` body now hardens
``model_path`` and ``gguf_variant`` via the shared
``_validate_logged_identifier`` helper, so control characters
and URL-form HF tokens can no longer log-line-smuggle.

P1 #8 + #10: ``/delete-cached`` body hardens ``repo_id`` and
``variant`` the same way.

P1 #9: ``/download-progress`` ``repo_id`` query parameter is
also hardened; the value flows into log lines deep inside
``_get_repo_size_cached`` on lookup failure.

P1 #11: ``CheckFormatRequest.dataset_name`` and
``AiAssistMappingRequest.{dataset_name, model_name}`` in
``models/datasets.py`` now apply the same control-char +
embedded-HF-token validators, matching every other public
request-body model.

All 115 diffusion + training-validation + cached_gguf + export
+ inference model-validation tests pass locally.

(P1 #6 native-path-lease enforcement for diffusion local paths
and P1 #12 React Compiler frontend lint deferred -- both need
focused design / frontend touchups separate from this batch.)

studio/backend/models/datasets.py

@@ -7,7 +7,12 @@ Dataset-related Pydantic models for API requests and responses.
 
 from typing import Any, Dict, List, Optional
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+# Round 24 P1 #11: reuse the chat / diffusion / export identifier
+# hardening so dataset routes also reject control characters and
+# URL-embedded HF tokens in user-controlled identifiers.
+from models.inference import _no_control_chars, _reject_embedded_hf_token
 
 
 class CheckFormatRequest(BaseModel):
@@ -27,6 +32,16 @@ class CheckFormatRequest(BaseModel):
             values.setdefault("train_split", values.pop("split"))
         return values
 
+    @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 CheckFormatResponse(BaseModel):
     """Response for dataset format check"""
@@ -57,6 +72,16 @@ class AiAssistMappingRequest(BaseModel):
     model_name: Optional[str] = None
     model_type: Optional[str] = None
 
+    @field_validator("dataset_name", "model_name")
+    @classmethod
+    def _no_identifier_control_chars(cls, v, info):
+        return _no_control_chars(v, info.field_name)
+
+    @field_validator("dataset_name", "model_name")
+    @classmethod
+    def _no_identifier_embedded_hf_tokens(cls, v, info):
+        return _reject_embedded_hf_token(v, info.field_name)
+
 
 class AiAssistMappingResponse(BaseModel):
     """Response from LLM-assisted column classification and conversion advice."""

studio/backend/routes/export.py

@@ -150,22 +150,15 @@ async def load_checkpoint(
         # reviews #1, #8, #9 flagged.
         from routes.inference import _release_chat_for, _release_diffusion_for
 
-        await _release_chat_for("export")
-
-        # Also unload any active diffusion pipeline (Images page); it
-        # competes for the same GPU and would survive the inference
-        # shutdown above. is_loading is treated like is_loaded so an
-        # in-flight load is also waited out (the diffusion unload
-        # acquires _load_lock + _generate_lock and blocks until the
-        # current load completes, then unloads).
-        # Round 17: previously this was a best-effort try/except that
-        # swallowed every failure with logger.debug, so a wedged
-        # diffusion backend let the export checkpoint load anyway and
-        # OOM at first allocation. ``_release_diffusion_for`` is
-        # strict: it raises HTTPException 503 if status() or
-        # unload_model() fails, or if the backend remains loaded or
-        # loading after the unload call.
+        # Round 24 P1 #3: release diffusion BEFORE chat so a failing
+        # diffusion unload does not leave the user with no chat
+        # model loaded. Same reasoning as the training-start flow
+        # (round 18 P1 #8 / round 24 P1 #2). Earlier rounds kept the
+        # chat release first because the helper was best-effort;
+        # now that ``_release_diffusion_for`` is strict it must run
+        # while chat is still resident so a failure preserves it.
         await _release_diffusion_for("export load")
+        await _release_chat_for("export")
 
         # load_checkpoint spawns and waits on a subprocess and can take
         # minutes. Run it in a worker thread so the event loop stays

studio/backend/routes/inference.py

@@ -1192,9 +1192,14 @@ async def load_model(
             # corrupt the user's exported artifact).
             _raise_if_training_active("chat")
             _raise_if_export_active("chat")
-            # Drop a settled export checkpoint that is just holding
-            # GPU memory but is not actively producing output.
+            # Round 24 P1 #4: release order is now
+            # export -> diffusion -> safetensors chat (was
+            # export -> safetensors chat -> diffusion). A wedged
+            # diffusion unload used to fire AFTER the safetensors
+            # chat was already gone, so the user lost both. Drop
+            # the chat last so an earlier failure preserves it.
             await _release_export_for("GGUF chat")
+            await _release_diffusion_for("GGUF chat load")
 
             llama_backend = get_llama_cpp_backend()
             # Round 19 P2 #8: previously also called
@@ -1206,19 +1211,13 @@ async def load_model(
             # ``_release_safetensors_chat_for`` below already
             # handles missing-backend cases as a no-op.
 
-            # Unload any safetensors / Unsloth model first to free
-            # VRAM. Uses the shared helper so we also drain
-            # ``loading_models`` (round 10 review #4); the inline
-            # version only checked ``active_model_name`` and let an
-            # in-flight safetensors load race the new GGUF allocation.
+            # Unload any safetensors / Unsloth model. Uses the shared
+            # helper so we also drain ``loading_models`` (round 10
+            # review #4); the inline version only checked
+            # ``active_model_name`` and let an in-flight safetensors
+            # load race the new GGUF allocation.
             await _release_safetensors_chat_for("GGUF chat")
 
-            # Round 17 P1 #4: route the diffusion unload through the
-            # strict ``_release_diffusion_for`` helper so a wedged
-            # diffusion pipeline blocks the GGUF chat load with 503
-            # instead of silently double-owning VRAM.
-            await _release_diffusion_for("GGUF chat load")
-
             # Inherit llama_extra_args from the previous load when the
             # request omits the field (the chat-settings Apply path
             # does not round-trip them; explicit [] still clears).
@@ -1392,22 +1391,22 @@ async def load_model(
         # and so we do not silently corrupt an in-flight export.
         _raise_if_training_active("chat")
         _raise_if_export_active("chat")
-        # Drop a settled export checkpoint that is just holding GPU
-        # memory but is not actively producing output.
+        # Round 24 P1 #5: release order is now
+        # export -> diffusion -> llama-chat (was
+        # export -> llama-chat -> diffusion). A wedged diffusion
+        # unload used to fire AFTER the GGUF chat was already gone,
+        # so the user lost both. Drop llama-chat last so an earlier
+        # failure preserves it.
         await _release_export_for("safetensors chat")
+        await _release_diffusion_for("safetensors chat load")
 
         backend = get_inference_backend()
 
-        # Unload any active or mid-download llama-server first.
-        # Shared helper so this stays in sync with the GGUF path's
+        # Unload any active or mid-download llama-server. Shared
+        # helper so this stays in sync with the GGUF path's
         # symmetric ``_release_safetensors_chat_for``.
         await _release_llama_for("safetensors chat")
 
-        # Round 17 P1 #5: strict diffusion unload via the shared
-        # helper so a wedged pipeline blocks the safetensors chat
-        # load with 503 instead of silently double-owning VRAM.
-        await _release_diffusion_for("safetensors chat load")
-
         # Export was already dropped above via the shared
         # ``await _release_export_for("safetensors chat")`` call
         # (which checks is_export_active() before the destructive

studio/backend/routes/models.py

@@ -1876,6 +1876,14 @@ async def delete_finetuned_model(
     Only paths under Studio's outputs/exports roots are accepted.  Exported
     GGUF entries can delete one quantization variant at a time.
     """
+    # Round 24 P1 #7 + P2 #13: harden both ``model_path`` and
+    # ``gguf_variant`` for control characters and embedded HF
+    # tokens, mirroring the chat / diffusion / training request
+    # validators. Both fields end up in logger.info(...) lines.
+    model_path = _validate_logged_identifier(model_path, "model_path")
+    if gguf_variant is not None:
+        gguf_variant = _validate_logged_identifier(gguf_variant, "gguf_variant")
+
     if source not in {"training", "exported"}:
         raise HTTPException(
             status_code = 400,
@@ -2506,6 +2514,10 @@ async def get_download_progress(
         "progress": 0,
         "cache_path": None,
     }
+    # Round 24 P1 #9: ``repo_id`` flows into log lines deep in
+    # ``_get_repo_size_cached`` on lookup failure, so the same
+    # hardening the request-body models use applies here too.
+    repo_id = _validate_logged_identifier(repo_id, "repo_id")
     try:
         if not _is_valid_repo_id(repo_id):
             return _empty
@@ -2769,6 +2781,12 @@ async def delete_cached_model(
     are removed (e.g. ``UD-Q4_K_XL``).  Otherwise the entire repo is deleted.
     Refuses if the model is currently loaded for inference.
     """
+    # Round 24 P1 #8 + #10: harden both ``repo_id`` and ``variant``
+    # against control characters / embedded HF tokens before they
+    # reach logger.info(...) lines or the HF cache scan.
+    repo_id = _validate_logged_identifier(repo_id, "repo_id")
+    if variant is not None:
+        variant = _validate_logged_identifier(variant, "variant")
     if not _is_valid_repo_id(repo_id):
         raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
 

studio/backend/routes/training.py

@@ -284,20 +284,18 @@ async def start_training(
         # model. The previous order (chat -> export) would drop chat
         # and then refuse training when a wedged idle export raised,
         # leaving the user with nothing loaded.
+        # Round 24 P1 #2: same reasoning extended to diffusion ->
+        # chat. A wedged diffusion unload used to fire AFTER the chat
+        # backend was already gone, so the user lost both chat and
+        # diffusion on a single failure mode. Order is now
+        # export -> diffusion -> chat, with chat as the last drop so
+        # earlier failures preserve it.
         await _release_export_for("training")
+        await _release_diffusion_for("training")
         await _release_chat_for("training")
 
-        # Also unload any loaded diffusion pipeline (Images page); it
-        # holds the same GPU and would survive the inference shutdown.
-        # is_loading=True is also handled (unload_model takes
-        # _load_lock + _generate_lock and waits the in-flight load out).
-        # Round 17: previously the diffusion unload was best-effort
-        # (try/except + logger.warning), so a stuck diffusion backend
-        # would let training start anyway and immediately OOM the
-        # subprocess. ``_release_diffusion_for`` is strict: it raises
-        # HTTPException 503 if status() or unload_model() fails, or if
-        # the backend remains loaded / loading after the unload call.
-        await _release_diffusion_for("training")
+        # (Diffusion release moved above chat in round 24 P1 #2;
+        # the old trailing call was removed to avoid double-unload.)
 
         # start_training now spawns a subprocess (non-blocking)
         success = backend.start_training(job_id = job_id, **training_kwargs)

studio/backend/utils/datasets/llm_assist.py

@@ -139,10 +139,59 @@ def _gpu_workload_busy_for_helper() -> bool:
     blocks the helper instead of double-owning VRAM. Each step
     fails closed: an unverifiable status counts as busy so the
     user's primary workload is preserved over the optional helper.
+
+    Round 24 P1 #1: extended to also catch a Chat-backend GPU owner.
+    The helper GGUF used to run on top of a loaded GGUF chat model
+    (llama-server) or safetensors chat model and OOM their shared
+    GPU; mirror the diffusion check by inspecting llama
+    ``is_loaded`` / ``is_active`` / ``loading_model_identifier`` and
+    safetensors ``active_model_name`` / ``loading_models``.
     """
     if _diffusion_image_model_busy():
         return True
 
+    try:
+        from routes.inference import get_llama_cpp_backend
+    except Exception:
+        pass
+    else:
+        try:
+            llama = get_llama_cpp_backend()
+            if (
+                getattr(llama, "is_loaded", False)
+                or getattr(llama, "is_active", False)
+                or getattr(llama, "loading_model_identifier", None)
+            ):
+                logger.info(
+                    "Skipping helper GGUF while a GGUF chat model is loaded/loading"
+                )
+                return True
+        except Exception:
+            logger.info(
+                "Skipping helper GGUF because llama-server status is unavailable"
+            )
+            return True
+
+    try:
+        from core.inference import get_inference_backend
+    except Exception:
+        pass
+    else:
+        try:
+            inf = get_inference_backend()
+            active = getattr(inf, "active_model_name", None)
+            loading = set(getattr(inf, "loading_models", set()) or set())
+            if active or loading:
+                logger.info(
+                    "Skipping helper GGUF while a safetensors chat model is loaded/loading"
+                )
+                return True
+        except Exception:
+            logger.info(
+                "Skipping helper GGUF because safetensors chat status is unavailable"
+            )
+            return True
+
     try:
         from core.training import get_training_backend
     except Exception: