ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
Three round-32 reviewer findings, plus documentation cleanup for
the local-path Tauri/FE plumbing gap.

Concurrency: direct DiffusionBackend.load_model callers now publish
the helper/advisor pending marker symmetrically (round 32 P1 #3).
_raise_if_helper_advisor_busy_for_diffusion gains an optional
publish_pending flag; load_model passes True so the destructive
unload window is gated by a "diffusion-backend" tag published
under _HELPER_ADVISOR_START_LOCK. The route layer's "diffusion"
tag and the backend's "diffusion-backend" tag refcount
independently (sum > 0 still blocks helper starts), so neither
side's clear can erase the other's still-active marker. The
existing _release_chat_backend_for_diffusion(check_helper_advisor=
True) path stays snapshot-only (publish_pending defaults False) so
test / direct callers of that helper do not leak a counter.

Validation: export save_directory now rejects ALL ASCII control
characters (round 32 P1, save_directory tab finding). The earlier
CR / LF only guard missed TAB / VT / FF / DEL, which a caller
could smuggle past the export worker's logged subprocess argv.

Documentation: DiffusionLoadRequest.repo_id and base_repo updated
to reflect that local-path support is gated on a Tauri /
frontend load-diffusion-model directory lease producer that has
not shipped yet (round 32 P1 #1 from multiple reviewers). The
backend lease boundary is correct; what is missing is the FE /
native side that mints the matching grant. Until that lands,
local paths through the Images route always 400 with "Native
path grant is required", which the docstring now spells out.

Skipped (consistent with prior rounds):
  * Hub-pin findings (R32 P1 #4-#6): live B200 install with
    huggingface_hub==0.36.2 + transformers==4.57.6 + diffusers==
    0.37.1 verifiably imports Flux2KleinPipeline. Empirical
    justification documented in R30 / R30 follow-up commit msgs.
  * Tauri / native enum surgery (R32 P1 #1, 6 votes): real
    architectural work but out of scope for this PR's Python
    surface. Documented now; FE / Rust ticket to follow.

98 targeted backend tests pass (test_diffusion_routes,
test_diffusion_backend, test_inference_model_validation,
test_models_get_model_config_case_resolution, test_data_recipe_seed,
test_training_raw_support, test_export_log_cursor).
This commit is contained in:
Daniel Han-Chen 2026-05-25 16:19:06 +00:00
commit 90b51cc5c5
3 changed files with 102 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

82
studio/backend/core/inference/diffusion.py
View file

@ -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:

6
studio/backend/models/export.py
View file

@ -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")

32
studio/backend/models/inference.py
View file

@ -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,