diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 733a5d1bbe..be86e47282 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -802,8 +802,13 @@ class DiffusionBackend: fam = detect_family(repo_id, override_family = family_override) if fam is None: + # Round 22 P2 #4: route the repo label through + # ``_display_repo_id`` so a local absolute path that did + # not match any family does not leak the operator's + # filesystem layout via the error message / last_error + # / 400 response body. raise RuntimeError( - f"Could not infer a diffusion family for '{repo_id}'. " + f"Could not infer a diffusion family for '{_display_repo_id(repo_id)}'. " "Pass family_override = 'flux.2-klein' / 'flux.2' / " "'flux.1' / 'qwen-image' / 'stable-diffusion-3' / " "'stable-diffusion-xl' to disambiguate." diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 7c53b0fee5..e0eec81197 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -8,6 +8,13 @@ Pydantic schemas for Training API from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing import Any, Optional, List, Dict, Literal +# Round 22 P1 #1: reuse the chat / diffusion identifier validators +# so /api/training/start rejects newline / tab / control characters +# and URL-form ``hf_xxxxx`` tokens in ``model_name``. Without these +# a caller could log-line-smuggle through "Loading model %s" lines +# and leak the bearer token into structured-log sinks. +from models.inference import _no_control_chars, _reject_embedded_hf_token + _MAX_BATCH_SIZE = 4096 _MAX_GRAD_ACCUM = 4096 @@ -49,6 +56,20 @@ class TrainingStartRequest(BaseModel): model_name: str = Field( ..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')" ) + + # 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") + @classmethod + def _no_model_name_control_chars(cls, v, info): + return _no_control_chars(v, info.field_name) + + @field_validator("model_name") + @classmethod + def _no_model_name_embedded_hf_tokens(cls, v, info): + return _reject_embedded_hf_token(v, info.field_name) + training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = ( Field( ..., diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 4c66d2ebf6..8aad3230e0 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -109,6 +109,28 @@ def precache_helper_gguf(): pass +def _diffusion_image_model_busy() -> bool: + """Round 22 P1 #2 / #3: helper / advisor GGUFs share VRAM with + the Images page diffusion pipeline. Public chat / training / + export routes call the strict ``_release_diffusion_for`` helper + before allocating, but these dataset-side helpers used to load + llama-server directly with no diffusion guard at all. Skip the + helper GGUF when ``DiffusionBackend.status()`` reports loaded / + loading so we do not double-own VRAM. Fail closed (treat as + busy) on any status() error to preserve the resident image + model rather than racing it for memory. + """ + try: + from core.inference.diffusion import get_diffusion_backend + except Exception: + return False + try: + status = get_diffusion_backend().status() + except Exception: + return True + return bool(status.get("is_loaded") or status.get("is_loading")) + + def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]: """ Load helper model, run one chat completion, unload. @@ -118,6 +140,12 @@ 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" + ) + return None + repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) variant = os.environ.get( "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT @@ -508,6 +536,13 @@ 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" + ) + return None + repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO) variant = os.environ.get( "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT