From 09ca2b27d319eb24f357adfe121538161a5e2c19 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 17:41:14 +0000 Subject: [PATCH] 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. --- studio/backend/core/inference/diffusion.py | 31 +++++++++++----------- studio/backend/routes/datasets.py | 12 +++++++++ 2 files changed, 28 insertions(+), 15 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index efc8799ba0..5c4450fd15 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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: diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index b0e11fd248..9ac2edb44d 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -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))