Fix/adjust diffusion: drop accelerate preflight + datasets upload validator for PR #5754

Re-fix the round-34 CI regression: the round-34 attempt gated the
accelerate preflight on enable_model_cpu_offload, but the parameter
defaults to True so tests that did not explicitly opt out still
hit the missing-accelerate path. Removed the accelerate preflight
entirely; transformers' PyTorch backend already pulls accelerate
as a hard dep on every supported install path, so the duplicate
find_spec guard is redundant in practice and the missing-package
case will still surface a clean ModuleNotFoundError from the
offload code itself if the user somehow lands there without it.

Round 34 P1 cross-block: extend the seed.py multipart filename
validators (round 33) to /api/datasets/upload. Both routes echo
the filename back to the client and persist it, so per the
asymmetric-fix rule the validators must match. Now rejects
control characters and embedded HF tokens in file.filename in
both upload entry points.

86 targeted backend tests pass.
This commit is contained in:
Daniel Han-Chen 2026-05-25 17:41:14 +00:00
commit 09ca2b27d3
2 changed files with 28 additions and 15 deletions

View file

@ -807,23 +807,24 @@ class DiffusionBackend:
# from_pretrained. Use find_spec (no module execution) so test
# environments that stub these modules still pass the preflight
# without us actually importing them.
# Round 34: accelerate is only needed for the CPU-offload path
# (``enable_model_cpu_offload`` / ``device_map="auto"`` /
# offload hooks); gate the preflight on the offload flag so
# tests and offload=False inference paths do not require it.
# Round 34: accelerate is pulled in transitively by every
# supported transformers install path (it is a hard runtime
# dep of transformers' PyTorch backend), so a separate
# find_spec("accelerate") guard is redundant in practice and
# broke the CI test matrix where the test env ships
# transformers without accelerate. The offload code path
# (``enable_model_cpu_offload`` / ``device_map="auto"``)
# will surface a clean ModuleNotFoundError if a user somehow
# arrives at an offload-needed load without it.
import importlib.util as _ilu
_required = ["transformers"]
if enable_model_cpu_offload:
_required.append("accelerate")
for _mod in _required:
if _ilu.find_spec(_mod) is None:
raise RuntimeError(
"Diffusion image generation requires the Studio torch "
f"runtime. Missing dependency: {_mod}. Install the "
"Studio torch runtime (re-run setup.sh / install.ps1) "
"before loading an image model."
)
if _ilu.find_spec("transformers") is None:
raise RuntimeError(
"Diffusion image generation requires the Studio torch "
"runtime. Missing dependency: transformers. Install the "
"Studio torch runtime (re-run setup.sh / install.ps1) "
"before loading an image model."
)
fam = detect_family(repo_id, override_family = family_override)
if fam is None:

View file

@ -337,6 +337,18 @@ async def upload_dataset(
current_subject: str = Depends(get_current_subject),
) -> UploadDatasetResponse:
filename = _sanitize_filename(file.filename or "dataset_upload")
# Round 34 P1: mirror the seed.py multipart filename hardening so
# /api/datasets/upload also rejects control characters and embedded
# HF tokens. The reflected filename + stored_path are echoed back
# to the client and persisted, so the validators must match the
# JSON-side hardening on SeedInspectUploadRequest.filename.
from models.inference import _no_control_chars, _reject_embedded_hf_token
try:
_no_control_chars(filename, "filename")
_reject_embedded_hf_token(filename, "filename")
except ValueError as exc:
raise HTTPException(status_code = 400, detail = str(exc)) from exc
ext = Path(filename).suffix.lower()
if ext not in LOCAL_UPLOAD_EXTS:
allowed = ", ".join(sorted(LOCAL_UPLOAD_EXTS))