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

P1 #1 + #2 + #6: extended the chat / diffusion / training
identifier hardening to every export-side request model.
ExportCommonOptions (parent of ExportMergedModelRequest /
ExportBaseModelRequest / ExportLoRAAdapterRequest) now applies
_no_control_chars and _reject_embedded_hf_token to repo_id and
base_model_id; ExportGGUFRequest gets the same on its repo_id
plus a control-char check on quantization_method; and
LoadCheckpointRequest validates checkpoint_path. Previously
"/api/export/*" accepted newline-smuggled identifiers and
URL-form ``hf_xxxxx`` tokens that flowed into log lines.

P1 #3 + #4: ``_run_with_helper`` and ``_run_multi_pass_advisor``
now use a shared ``_gpu_workload_busy_for_helper`` that gates on
diffusion (round 22 already), training, AND export. The round 22
guard only checked diffusion, so the dataset helper / advisor
could still load llama-server on top of an active training run
or a resident export checkpoint. Each step fails closed
(unverifiable status counts as busy) so the user's primary
workload is preserved.

P1 #5: PublishDatasetRequest in models/data_recipe.py also
applies the identifier hardening to repo_id; the publish path
previously accepted control characters and URL-form tokens.

P1 #7-10: added _validate_logged_identifier helper to
routes/models.py and applied it to the path / query parameter
endpoints that flow into logger.info(...) calls --
``/config/{model_name}``, ``/check-vision/{model_name}``,
``/check-embedding/{model_name}``, ``/gguf-variants``. Mapped
the validator's ValueError to HTTP 422 so the client sees the
same shape as a Pydantic validation failure.

P2 #11 + #12: ``Loading diffusion model %s`` and
``Diffusion load failed for %s`` log lines route ``repo_id`` /
``effective_base`` through ``_display_repo_id`` (collapses
absolute local paths to the leaf, still scrubs HF tokens)
instead of plain ``_redact_hf_tokens``. The error path was
already collapsed in the user-facing 400 / RuntimeError, but
the structured-log lines kept the full path.

All 97 diffusion + training-validation + related tests pass
locally.

studio/backend/core/inference/diffusion.py

@@ -907,13 +907,21 @@ class DiffusionBackend:
                 # Scrub them BEFORE the logger formats the line so the
                 # token never reaches structured-log sinks (round 14
                 # P2 #9).
+                # Round 23 P2 #11: ``_redact_hf_tokens`` only scrubs
+                # ``hf_xxxxx`` substrings, so an absolute local
+                # path like ``/home/alice/private/FLUX.2-klein-GGUF``
+                # used to land in this log line verbatim. Route
+                # through ``_display_repo_id`` so the leaf is
+                # logged when the value is a filesystem path, with
+                # the token-redaction step inside that helper as a
+                # belt-and-braces defence.
                 logger.info(
                     "Loading diffusion model %s (family=%s, device=%s, dtype=%s, base=%s)",
-                    _redact_hf_tokens(repo_id),
+                    _display_repo_id(repo_id),
                     fam.name,
                     device,
                     dtype,
-                    _redact_hf_tokens(effective_base),
+                    _display_repo_id(effective_base),
                 )
 
                 transformer = None
@@ -1249,9 +1257,12 @@ class DiffusionBackend:
                 # Use ``logger.error`` with the already-scrubbed
                 # message and exc_info=False so the bearer token
                 # cannot leak through structured logging sinks.
+                # Round 23 P2 #12: same fix as the start-of-load
+                # log above. ``_redact_hf_tokens`` alone left
+                # absolute local repo paths in this failure line.
                 logger.error(
                     "Diffusion load failed for %s: %s",
-                    _redact_hf_tokens(repo_id),
+                    _display_repo_id(repo_id),
                     exc_msg,
                 )
                 raise RuntimeError(

studio/backend/models/data_recipe.py

@@ -9,7 +9,13 @@ from __future__ import annotations
 
 from typing import Any
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+# Round 23 P1 #5: identifier hardening reused from the chat models
+# so /api/data_recipe/publish rejects control characters and
+# URL-form ``hf_xxxxx`` tokens in ``repo_id`` before they reach
+# log lines or the HF API.
+from models.inference import _no_control_chars, _reject_embedded_hf_token
 
 
 class RecipePayload(BaseModel):
@@ -60,6 +66,16 @@ class PublishDatasetRequest(BaseModel):
         description = "Execution artifact path captured by the UI for completed runs",
     )
 
+    @field_validator("repo_id")
+    @classmethod
+    def _no_repo_id_control_chars(cls, v, info):
+        return _no_control_chars(v, info.field_name)
+
+    @field_validator("repo_id")
+    @classmethod
+    def _no_repo_id_embedded_hf_tokens(cls, v, info):
+        return _reject_embedded_hf_token(v, info.field_name)
+
 
 class PublishDatasetResponse(BaseModel):
     success: bool = True

studio/backend/models/export.py

@@ -10,6 +10,13 @@ from pathlib import Path
 from pydantic import BaseModel, Field, field_validator
 from typing import List, Optional, Literal, Dict, Any
 
+# Round 23 P1 #1 / #2 / #6: reuse the chat identifier validators
+# so export requests reject newline / tab / control characters and
+# URL-form ``hf_xxxxx`` tokens in any user-supplied identifier
+# (Hub ``repo_id``, ``base_model_id``, the local
+# ``checkpoint_path``) that flows into log lines or HF API calls.
+from models.inference import _no_control_chars, _reject_embedded_hf_token
+
 
 def _validate_save_directory(value: str) -> str:
     """Reject save_directory values that escape the export root."""
@@ -54,6 +61,19 @@ class LoadCheckpointRequest(BaseModel):
         description = "Allow loading models with custom code. Only enable for checkpoints/base models you trust.",
     )
 
+    # Round 23 P1 #6: ``checkpoint_path`` is logged verbatim by the
+    # export route. Apply the same control-char + embedded-token
+    # rejection the chat / diffusion / training request models use.
+    @field_validator("checkpoint_path")
+    @classmethod
+    def _no_checkpoint_control_chars(cls, v, info):
+        return _no_control_chars(v, info.field_name)
+
+    @field_validator("checkpoint_path")
+    @classmethod
+    def _no_checkpoint_embedded_hf_tokens(cls, v, info):
+        return _reject_embedded_hf_token(v, info.field_name)
+
 
 class ExportStatusResponse(BaseModel):
     """Current export backend status."""
@@ -117,6 +137,20 @@ class ExportCommonOptions(BaseModel):
         description = "HuggingFace model ID of the base model (for model card metadata)",
     )
 
+    # Round 23 P1 #1: ``repo_id`` (Hub destination) and
+    # ``base_model_id`` (model card metadata) both feed log lines
+    # and the HF API. Reject control characters and URL-form
+    # ``hf_xxxxx`` tokens before they reach those sinks.
+    @field_validator("repo_id", "base_model_id")
+    @classmethod
+    def _no_identifier_control_chars(cls, v, info):
+        return _no_control_chars(v, info.field_name)
+
+    @field_validator("repo_id", "base_model_id")
+    @classmethod
+    def _no_identifier_embedded_hf_tokens(cls, v, info):
+        return _reject_embedded_hf_token(v, info.field_name)
+
 
 class ExportMergedModelRequest(ExportCommonOptions):
     """Request for exporting a merged PEFT model."""
@@ -163,6 +197,27 @@ class ExportGGUFRequest(BaseModel):
         description = "Hugging Face token for GGUF upload",
     )
 
+    # Round 23 P1 #2: GGUF export endpoint defines its own
+    # ``repo_id`` (does not inherit from ExportCommonOptions), so
+    # the chat-style hardening needs to be applied here separately.
+    # ``quantization_method`` is forwarded to the export worker
+    # command line, so it gets the control-char check too even
+    # though it does not normally carry tokens.
+    @field_validator("repo_id")
+    @classmethod
+    def _no_repo_id_control_chars(cls, v, info):
+        return _no_control_chars(v, info.field_name)
+
+    @field_validator("repo_id")
+    @classmethod
+    def _no_repo_id_embedded_hf_tokens(cls, v, info):
+        return _reject_embedded_hf_token(v, info.field_name)
+
+    @field_validator("quantization_method")
+    @classmethod
+    def _no_quantization_control_chars(cls, v, info):
+        return _no_control_chars(v, info.field_name)
+
 
 class ExportLoRAAdapterRequest(ExportCommonOptions):
     """Request for exporting only the LoRA adapter (not merged)."""

studio/backend/routes/models.py

@@ -134,11 +134,30 @@ from models.responses import (
     VisionCheckResponse,
     EmbeddingCheckResponse,
 )
+from models.inference import _no_control_chars, _reject_embedded_hf_token
 
 router = APIRouter()
 logger = get_logger(__name__)
 
 
+def _validate_logged_identifier(value: str, field_name: str) -> str:
+    """Round 23 P1 #7 / #8 / #9 / #10: path / query parameters that
+    flow into ``logger.info("... %s", value)`` lines were the last
+    unguarded entry points. Newline / tab / control characters let
+    a caller smuggle forged log entries; URL-form ``hf_xxxxx``
+    tokens would leak into structured-log sinks. Mirror the
+    request-body validators by running both checks here and
+    mapping the validator's ``ValueError`` to HTTP 422 so the
+    client sees the same shape as a Pydantic validation failure.
+    """
+    try:
+        value = _no_control_chars(value, field_name)
+        value = _reject_embedded_hf_token(value, field_name)
+    except ValueError as exc:
+        raise HTTPException(status_code = 422, detail = str(exc)) from exc
+    return value
+
+
 def derive_model_type(
     is_vision: bool, audio_type: Optional[str], is_embedding: bool = False
 ) -> ModelType:
@@ -1571,6 +1590,7 @@ async def get_model_config(
 
     This endpoint wraps the backend load_model_defaults function.
     """
+    model_name = _validate_logged_identifier(model_name, "model_name")
     try:
         if not is_local_path(model_name):
             resolved = resolve_cached_repo_id_case(model_name)
@@ -1580,7 +1600,11 @@ async def get_model_config(
                     resolved,
                     model_name,
                 )
-            model_name = resolved
+            # Round 23 P1 #7: re-validate the cache-resolved value
+            # (case-only resolver should be a no-op for these
+            # checks, but defend in depth in case the resolver
+            # ever broadens its match heuristic).
+            model_name = _validate_logged_identifier(resolved, "model_name")
 
         logger.info(f"Getting model config for: {model_name}")
         from utils.models.model_config import detect_audio_type
@@ -2220,6 +2244,7 @@ async def check_vision_model(
 
     This endpoint wraps the backend is_vision_model function.
     """
+    model_name = _validate_logged_identifier(model_name, "model_name")
     try:
         logger.info(f"Checking if vision model: {model_name}")
         is_vision = is_vision_model(model_name)
@@ -2248,6 +2273,7 @@ async def check_embedding_model(
 
     This endpoint wraps the backend is_embedding_model function.
     """
+    model_name = _validate_logged_identifier(model_name, "model_name")
     try:
         logger.info(f"Checking if embedding model: {model_name}")
         is_embedding = is_embedding_model(model_name, hf_token = hf_token)
@@ -2285,6 +2311,7 @@ async def get_gguf_variants(
     with file sizes, whether the model supports vision, and the recommended
     default variant.
     """
+    repo_id = _validate_logged_identifier(repo_id, "repo_id")
     try:
         from utils.models.model_config import is_local_path, list_local_gguf_variants
 

studio/backend/utils/datasets/llm_assist.py

@@ -131,6 +131,57 @@ def _diffusion_image_model_busy() -> bool:
     return bool(status.get("is_loaded") or status.get("is_loading"))
 
 
+def _gpu_workload_busy_for_helper() -> bool:
+    """Round 23 P1 #3 / #4: the diffusion-only guard from round 22
+    let the helper / advisor GGUF run on top of a live training run
+    or a resident export checkpoint. Extend the busy check to those
+    workloads too so any GPU owner (Images, Training, Export)
+    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.
+    """
+    if _diffusion_image_model_busy():
+        return True
+
+    try:
+        from core.training import get_training_backend
+    except Exception:
+        pass
+    else:
+        try:
+            if get_training_backend().is_training_active():
+                logger.info(
+                    "Skipping helper GGUF while training is active"
+                )
+                return True
+        except Exception:
+            logger.info(
+                "Skipping helper GGUF because training status is unavailable"
+            )
+            return True
+
+    try:
+        from core.export import get_export_backend
+    except Exception:
+        return False
+
+    try:
+        exp = get_export_backend()
+        is_active = getattr(exp, "is_export_active", None)
+        if (is_active and is_active()) or getattr(
+            exp, "current_checkpoint", None
+        ):
+            logger.info("Skipping helper GGUF while export owns the GPU")
+            return True
+    except Exception:
+        logger.info(
+            "Skipping helper GGUF because export status is unavailable"
+        )
+        return True
+
+    return False
+
+
 def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
     """
     Load helper model, run one chat completion, unload.
@@ -140,10 +191,10 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
     if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
         return None
 
-    if _diffusion_image_model_busy():
-        logger.info(
-            "Skipping helper GGUF while a diffusion image model is loaded/loading"
-        )
+    # Round 23 P1 #3: round 22 only guarded against a busy
+    # diffusion pipeline. Training / export own the same GPU too,
+    # so use the broader helper that gates on all three workloads.
+    if _gpu_workload_busy_for_helper():
         return None
 
     repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
@@ -536,11 +587,10 @@ def _run_multi_pass_advisor(
     if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
         return None
 
-    # Round 22 P1 #3: same diffusion-busy guard as ``_run_with_helper``.
-    if _diffusion_image_model_busy():
-        logger.info(
-            "Skipping advisor GGUF while a diffusion image model is loaded/loading"
-        )
+    # Round 23 P1 #4: extend the round 22 diffusion-only check to
+    # training + export so the advisor cannot race the user's
+    # active workload for GPU memory.
+    if _gpu_workload_busy_for_helper():
         return None
 
     repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)