diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index c26388985a..feb5b1cfee 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -834,6 +834,13 @@ class DiffusionBackend: device, dtype = self._pick_device_and_dtype() + # Round 32 P1 #3: track whether the backend-side + # helper-busy check published a "diffusion-backend" pending + # entry so the outer finally clears the matching publish + # exactly once. Set inside the try below right after the + # snapshot succeeds. + backend_pending_published = False + # _load_lock serialises the entire load so two concurrent calls # cannot both kick off a multi-GB download + GPU upload at once. # The second caller waits behind the first and then loads on top @@ -1062,7 +1069,19 @@ class DiffusionBackend: # _release_other_gpu_owners_for_diffusion raises # RuntimeError early when training/export is active # without touching the chat backend. - _raise_if_helper_advisor_busy_for_diffusion() + # Round 32 P1 #3: publish a backend-side pending + # entry under the helper-advisor start lock so a + # direct / test / future caller of this method is + # symmetric with the route layer's + # _raise_if_helper_advisor_busy("diffusion"). The + # route's "diffusion" tag and this "diffusion- + # backend" tag refcount independently; both + # contribute to public_load_pending(). + backend_pending_published = ( + _raise_if_helper_advisor_busy_for_diffusion( + publish_pending = True, + ) + ) _release_other_gpu_owners_for_diffusion() _release_chat_backend_for_diffusion(check_helper_advisor = False) @@ -1306,6 +1325,12 @@ class DiffusionBackend: self._pending_repo_id = None self._pending_base_repo = None self._pending_gguf_filename = None + # Round 32 P1 #3: clear the backend-side public-load + # pending publish if it was set. Skipped when the + # helper-busy snapshot raised (no publish to clear) + # so the counter stays in sync with publishes. + if backend_pending_published: + _clear_diffusion_backend_pending() def unload_model(self) -> dict[str, Any]: # Take the load lock and the generate lock so unload cannot: @@ -1528,23 +1553,64 @@ def encode_png_base64(pil_image: "Any") -> str: # ─── Helpers ────────────────────────────────────────────────────────── -def _raise_if_helper_advisor_busy_for_diffusion() -> None: +def _raise_if_helper_advisor_busy_for_diffusion( + *, + publish_pending: bool = False, +) -> bool: """Round 29 P1 #1: split the helper-busy check out of _release_chat_backend_for_diffusion so the diffusion load can check ALL conflicts (helper, training, export) BEFORE doing ANY destructive unloads. Otherwise a route-precheck race or a direct backend call would unload the user's chat while training was active, then 409 with the user holding no model at all. + + Round 32 P1 #3: when ``publish_pending=True`` also takes + ``_HELPER_ADVISOR_START_LOCK`` and publishes a + ``diffusion-backend`` public-load pending entry so a concurrent + AI Assist helper / advisor start that wins the start lock sees + the pending public owner and refuses VRAM. The route layer + publishes its own ``diffusion`` tag (refcount semantics, so the + two publishes coexist without erasing each other). Returns True + when a pending entry was actually published so the caller can + pair it with ``_clear_diffusion_backend_pending`` in finally. + Direct callers (tests, scripts) opt in with ``publish_pending= + True`` to get the same atomic check + publish the route gets. + The ``check_helper_advisor`` callback in + ``_release_chat_backend_for_diffusion`` keeps the default False + so legacy callers do not double-publish or leak pending entries. """ try: - from utils.datasets.llm_assist import helper_advisor_busy + from utils.datasets.llm_assist import ( + _HELPER_ADVISOR_START_LOCK, + _publish_public_load_pending, + helper_advisor_busy, + ) + except Exception: + return False + with _HELPER_ADVISOR_START_LOCK: + 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." + ) + if publish_pending: + _publish_public_load_pending("diffusion-backend") + return True + return False + + +def _clear_diffusion_backend_pending() -> None: + """Round 32 P1 #3: paired clear for + ``_raise_if_helper_advisor_busy_for_diffusion(publish_pending=True)``. + Safe to call when the helpers module is unavailable (no-op).""" + try: + from utils.datasets.llm_assist import _release_public_load_pending except Exception: return - 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." - ) + try: + _release_public_load_pending("diffusion-backend") + except Exception: + pass def _release_chat_backend_for_diffusion(*, check_helper_advisor: bool = True) -> None: diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index df5aed2201..5e8ee29bf9 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -27,7 +27,11 @@ def _validate_save_directory(value: str) -> str: raise ValueError("save_directory must not be empty") if "\x00" in raw: raise ValueError("save_directory may not contain null bytes") - if any(ch in raw for ch in ("\r", "\n")): + # Round 32 P1: reject ALL ASCII control characters (including + # TAB / VT / FF) so a caller cannot smuggle log-line breaks or + # subprocess argv splitters past the export worker. The earlier + # CR / LF check missed every other C0 byte. + if any(ord(ch) < 0x20 or ord(ch) == 0x7f for ch in raw): raise ValueError("save_directory may not contain control characters") if len(raw) > 255: raise ValueError("save_directory must be <= 255 characters") diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 6b99cc4ec1..749e5d8dc4 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1557,20 +1557,30 @@ class DiffusionLoadRequest(BaseModel): VAE / text encoders when loading a GGUF-only repo. """ - # repo_id and base_repo can be absolute local paths (Studio - # exports under deeply nested ``outputs/...`` directories, - # Windows paths with drive letter, etc.). 1024 chars matches - # POSIX PATH_MAX-class limits and Windows long-path support; - # the rounds-of-256 cap was rejecting realistic export paths. + # repo_id and base_repo are HF Hub identifiers in this release. + # Local-path support is gated behind a frontend / Tauri + # ``load-diffusion-model`` directory lease producer that has not + # shipped yet (round 32 P1 #3 in the PR reviewer trail). The + # 1024-char cap matches POSIX PATH_MAX so future local-path + # support can flip on without re-validating the field width. repo_id: str = Field( - ..., min_length = 1, max_length = 1024, description = "HF repo id or local path" + ..., + min_length = 1, + max_length = 1024, + description = ( + "HF repo id (owner/name). Local filesystem paths are reserved " + "for a future native-lease flow and currently rejected by the " + "route's _looks_like_local_diffusion_path guard." + ), ) # Round 30 P1 #4: chat /api/inference/load gates native local paths # through a signed native_path_lease grant before the backend # touches the filesystem. Mirror that here so /api/inference/images/ # load cannot be used as an authenticated probe for arbitrary - # local directories. Optional: Hub ids (no leading slash / tilde) - # skip the lease check entirely. + # local directories. Optional; Hub ids (no leading slash / tilde) + # skip the lease check entirely. The Images UI does not yet + # surface a local-path picker, so callers that omit this field + # always get the Hub-id code path. native_path_lease: Optional[str] = Field( None, description = "Frontend-visible signed native path grant for a local repo_id", @@ -1583,7 +1593,11 @@ class DiffusionLoadRequest(BaseModel): base_repo: Optional[str] = Field( None, max_length = 1024, - description = "Diffusers base repo (HF id or local path) for VAE + text encoders", + description = ( + "Diffusers base repo (HF id) for VAE + text encoders. Local " + "paths are gated on the same future native-lease flow as " + "repo_id." + ), ) base_repo_native_path_lease: Optional[str] = Field( None,