diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 20d67a9cf3..abb541779e 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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=, 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 diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index d379989241..14c8414d74 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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): """ diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 6628eef7f7..b360261e44 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -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 diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 5d739043e0..85a45caf15 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -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)