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

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

Browse source 
P1 #1: ``_preflight_full_diffusers_repo(effective_base, hf_token)``
now runs for every load mode, including the GGUF-with-auto-base
path. Round 19 only preflighted the full repo or an explicit
``base_repo``, so an auto-picked companion that turned out to be
gated / private / missing still unloaded the user's chat model
before ``from_pretrained`` failed. ``effective_base`` is the same
value that feeds every downstream allocation, so preflighting it
unconditionally catches all three modes.

P1 #2: ``diffusers.GGUFQuantizationConfig`` (which imports the
``gguf`` package at construction time) is now built up front,
inside the same try block that surfaces "Re-run Studio setup".
Previously the missing-dependency exception fired AFTER
``_release_other_gpu_owners_for_diffusion`` and
``_release_chat_backend_for_diffusion`` had already taken the
chat / export models down. The downstream from_single_file call
reuses the same ``quant_config`` reference.

P1 #4: ``studio/backend/requirements/studio.txt`` now lists
``diffusers>=0.37.0`` and ``gguf>=0.10.0``. These were only in
the extras files, so fresh standard Studio installs failed on
/images/load with the round 20 P1 #2 dependency error message.

P1 #5: ``LoadRequest``, ``UnloadRequest``, and
``ValidateModelRequest`` now apply the same control-character +
embedded-HF-token validators that ``DiffusionLoadRequest``
already had. /api/inference/load, /api/inference/validate, and
/api/inference/unload used to accept newline / tab / control
characters in ``model_path`` (log-line smuggling) and URL-form
``https://hf_xxxxx@huggingface.co/...`` (credential leak through
structured log sinks).

P2 #6: ``_collapse_local`` in the diffusion load-error scrubber
now resolves relative candidates and adds the absolute form to
the substring set. A relative ``exports/my-flux`` used to leak
``/mnt/disks/.../exports/my-flux/...`` via downstream library
errors because the scrubber only matched the original literal.
Replacement is longest-first so a leaf-only context survives.

All 85 diffusion-relevant + 35 related model-validation tests
pass locally.

(P1 #3 cross-workload GPU handoff lock is deferred: deserves a
focused design pass across /images/load, /chat/load (both
branches), /training/start, and /export/load to pick a lock
boundary that does not deadlock against the backend load locks
or stall the SSE log stream.)
This commit is contained in:
Daniel Han-Chen 2026-05-25 10:07:11 +00:00
commit ff3bad37fe
4 changed files with 124 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -893,20 +893,39 @@ class DiffusionBackend:
token = hf_token,
)
# Round 19 P1 #3: the GGUF branch above already
# proved repo + filename are accessible via
# ``hf_hub_download``. The full-diffusers path (no
# ``gguf_filename``) did NOT, so a typo / private /
# gated full repo only surfaced inside
# ``from_pretrained`` AFTER chat was unloaded. Probe
# ``effective_base`` for ``model_index.json`` here so
# the chat model is preserved on a bad full-repo
# request. Also probe when the GGUF caller supplied
# an explicit ``base_repo`` (the base companion is
# ALSO downloaded via from_pretrained further down
# and would OOM-then-fail past the unload).
if not gguf_filename or base_repo:
_preflight_full_diffusers_repo(effective_base, hf_token)
# Round 20 P1 #1: every load mode (full diffusers
# repo, GGUF + explicit base_repo, GGUF + auto-picked
# base_repo) feeds ``effective_base`` into
# ``from_pretrained`` further down. The round 19
# preflight only ran for the first two, so an
# auto-picked GGUF companion that turned out to be
# gated / private / missing still unloaded chat
# before the load failed. Always preflight
# ``effective_base`` so a bad companion repo is
# caught BEFORE chat / export are released.
_preflight_full_diffusers_repo(effective_base, hf_token)
# Round 20 P1 #2: ``diffusers.GGUFQuantizationConfig``
# imports the ``gguf`` package lazily at construction
# time. Partial Studio installs (``diffusers`` present,
# ``gguf`` not) used to discover that AFTER the chat /
# export release calls. Build the quant config up
# front so the missing-dependency surface raises
# while the user's chat model is still resident.
quant_config = None
if gguf_filename:
try:
quant_config = diffusers.GGUFQuantizationConfig(
compute_dtype = dtype
)
except ModuleNotFoundError as exc:
missing = exc.name or str(exc)
raise RuntimeError(
"Diffusion GGUF loading requires the gguf "
"runtime package. Missing dependency: "
f"{missing}. Re-run Studio setup before "
"loading an image GGUF."
) from exc
# All cheap failure points (bad gguf_filename, missing
# pipeline / transformer class, gated download token,
@ -965,7 +984,8 @@ class DiffusionBackend:
_drain_cuda_cache()
if gguf_filename:
quant_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype)
# ``quant_config`` was already constructed above
# (round 20 P1 #2 pre-release fail-fast).
# Diffusers-format GGUFs (FLUX.2 klein / Qwen-Image /
# SD3) need the matching base repo's component config
# at config=<base_repo>, subfolder="transformer".
@ -1098,20 +1118,34 @@ class DiffusionBackend:
except (OSError, ValueError):
return msg
leaf = p.name or candidate
abs_str = None
if p.is_absolute() or p.exists():
try:
abs_str = str(p)
except (OSError, ValueError):
abs_str = None
if abs_str and abs_str in msg:
msg = msg.replace(abs_str, leaf)
if (
candidate != leaf
and candidate in msg
and ("/" in candidate or "\\" in candidate)
needles: set[str] = set()
# Round 20 P2 #6: a relative candidate like
# ``exports/my-flux`` used to collapse only the
# exact ``exports/my-flux`` substring, but
# downstream libraries (diffusers / safetensors)
# resolve and emit ``/mnt/disks/.../exports/my-flux/...``
# absolute strings that leaked the operator's
# filesystem layout. Also scrub the resolved
# absolute form so the leaf is the only path
# fragment that survives.
try:
if p.exists():
needles.add(str(p.resolve()))
elif p.is_absolute():
needles.add(str(p))
except (OSError, ValueError):
pass
if "/" in candidate or "\\" in candidate:
needles.add(candidate)
# Replace longest first so a parent-directory
# substring does not blank out the leaf-only
# context the user needs.
for needle in sorted(
(n for n in needles if n and n != leaf),
key = len,
reverse = True,
):
msg = msg.replace(candidate, leaf)
msg = msg.replace(needle, leaf)
return msg
# ``effective_base`` and ``gguf_filename`` are local

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

@ -60,6 +60,24 @@ class LoadRequest(BaseModel):
return None
return value
# Round 20 P1 #5: extend the diffusion-side identifier hardening
# (round 5 P2 / round 15 P1 #5) to the chat LoadRequest. Newline
# / tab / control characters in ``model_path`` or ``gguf_variant``
# would otherwise be echoed verbatim into structured-log lines
# ("Loading model %s") and let a caller smuggle in fake log
# entries, and an embedded ``hf_...`` token in a URL-form path
# would leak the credential into the same log sinks the
# diffusion route already redacts.
@field_validator("model_path", "gguf_variant")
@classmethod
def _no_identifier_control_chars(cls, v, info):
return _no_control_chars(v, info.field_name)
@field_validator("model_path")
@classmethod
def _no_embedded_hf_tokens(cls, v, info):
return _reject_embedded_hf_token(v, info.field_name)
cache_type_kv: Optional[str] = Field(
None,
description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')",
@ -110,6 +128,20 @@ class UnloadRequest(BaseModel):
model_path: str = Field(..., description = "Model identifier to unload")
# Round 20 P1 #5: mirror the LoadRequest identifier hardening so
# /api/inference/unload also rejects control characters and
# URL-embedded HF tokens before the path reaches structured log
# sinks.
@field_validator("model_path")
@classmethod
def _no_identifier_control_chars(cls, v, info):
return _no_control_chars(v, info.field_name)
@field_validator("model_path")
@classmethod
def _no_embedded_hf_tokens(cls, v, info):
return _reject_embedded_hf_token(v, info.field_name)
class ValidateModelRequest(BaseModel):
"""
@ -130,6 +162,20 @@ class ValidateModelRequest(BaseModel):
None, description = "GGUF quantization variant (e.g. 'Q4_K_M')"
)
# Round 20 P1 #5: same identifier hardening as LoadRequest /
# UnloadRequest. /api/inference/validate flows directly into
# ``ModelConfig.from_identifier`` and the resulting log lines, so
# control characters and embedded HF tokens must not survive.
@field_validator("model_path", "gguf_variant")
@classmethod
def _no_identifier_control_chars(cls, v, info):
return _no_control_chars(v, info.field_name)
@field_validator("model_path")
@classmethod
def _no_embedded_hf_tokens(cls, v, info):
return _reject_embedded_hf_token(v, info.field_name)
class ValidateModelResponse(BaseModel):
"""

7
studio/backend/requirements/studio.txt
View file

@ -23,3 +23,10 @@ diceware
ddgs
cryptography>=42.0.0
httpx>=0.27.0
# Studio Images page runtime. Flux2KleinPipeline / Flux2Pipeline /
# QwenImagePipeline / StableDiffusion3Pipeline are available in
# diffusers>=0.37.0, and GGUFQuantizationConfig requires the gguf
# package (round 20 P1 #4: fresh standard Studio installs failed on
# /images/load because these were only listed in the extras files).
diffusers>=0.37.0
gguf>=0.10.0

10
studio/backend/tests/test_diffusion_backend.py
View file

@ -1305,7 +1305,15 @@ def test_load_model_accepts_relative_local_dir(monkeypatch, tmp_path):
),
)
def _boom(**_):
def _boom(**kwargs):
# Round 20 P1 #1 added a base-repo preflight that downloads
# the diffusers ``model_index.json`` of the auto-picked
# companion repo BEFORE the chat unload. Allow that call
# through (it would otherwise hit the network) but still
# reject any attempt to download the GGUF itself, which is
# what this test guards.
if kwargs.get("filename") == "model_index.json":
return "/tmp/model_index.json"
raise AssertionError("hf_hub_download must not run for a local dir")
fake_hub = SimpleNamespace(hf_hub_download = _boom)