From 8074a2b67b7785be3d73c17b8b361712632e4152 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 00:04:12 +0000 Subject: [PATCH] Fix/adjust diffusion: smart base, safetensors, peak VRAM, GGUF guard - _smart_base_repo: pick 9B base for unsloth/FLUX.2-klein-9B-GGUF and -base- variants per the repo id, instead of always falling back to the 4B family default. - pipe_kwargs use_safetensors=True so diffusers refuses pickle .bin weights at load time (defends against compromised base_repo). - Release the previous pipeline BEFORE allocating the new one so peak VRAM stays at one model's worth instead of two on swap. - Reject empty gguf_filename when repo_id ends with -GGUF; the prior behavior tried from_pretrained on a GGUF-only repo and 500'd deep in diffusers with a confusing model-index error. - Status returns gguf_filename (basename) instead of gguf_path so the local cache path / username does not leak to authenticated Studio sessions. - requirements/no-torch-runtime.txt: pin diffusers>=0.37.0 so older installs cannot resolve a version without Flux2KleinPipeline. - Frontend curated distilled klein entries now point at the matching non-base diffusers repos (FLUX.2-klein-4B / -9B) per the published model cards. Update api.ts to mirror the renamed status field. --- studio/backend/core/inference/diffusion.py | 77 +++++++++++++++-- .../backend/requirements/no-torch-runtime.txt | 6 +- .../backend/tests/test_diffusion_backend.py | 85 +++++++++++++++++-- studio/backend/tests/test_diffusion_routes.py | 2 +- studio/frontend/src/features/images/api.ts | 2 +- .../src/features/images/images-page.tsx | 11 ++- 6 files changed, 163 insertions(+), 20 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index bc7d20ad0b..d87d4c7623 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -141,6 +141,31 @@ _FULL_REPO_FAMILIES: tuple[DiffusionFamily, ...] = ( ) +def _smart_base_repo(fam: DiffusionFamily, repo_id: str) -> str: + """Pick the best matching base diffusers repo for a given GGUF repo + when the caller did not pass an explicit base_repo. + + Currently only specialises the flux.2-klein family: a repo name + containing "9b" gets the 9B base, "base-4b" / "base-9b" map to the + Base variants, everything else falls back to the family default + (Apache 2.0 4B Base). + """ + if fam.name != "flux.2-klein": + return fam.base_repo + lower = (repo_id or "").lower() + is_9b = "9b" in lower + is_base = "base" in lower + if is_9b and is_base: + return "black-forest-labs/FLUX.2-klein-base-9B" + if is_9b: + return "black-forest-labs/FLUX.2-klein-9B" + if is_base: + return "black-forest-labs/FLUX.2-klein-base-4B" + # Distilled 4B is the default for any flux-2-klein GGUF that does + # not advertise 9B or "base". + return "black-forest-labs/FLUX.2-klein-4B" + + def detect_family( repo_id: str, *, override_family: Optional[str] = None ) -> Optional[DiffusionFamily]: @@ -225,6 +250,12 @@ class DiffusionBackend: return self._repo_id def status(self) -> dict[str, Any]: + # Only echo the GGUF basename; full absolute path leaks the + # local HF cache layout (and the system username on default + # POSIX layouts) to any authenticated Studio session. + gguf_basename = ( + Path(self._gguf_path).name if self._gguf_path else None + ) return { "is_loaded": self.is_loaded, "is_loading": self._loading, @@ -232,7 +263,7 @@ class DiffusionBackend: "family": self._family.name if self._family else None, "pipeline_class": self._family.pipeline_class if self._family else None, "base_repo": self._base_repo, - "gguf_path": self._gguf_path, + "gguf_filename": gguf_basename, "device": self._device, "dtype": self._dtype, "loaded_at": self._loaded_at, @@ -334,13 +365,26 @@ class DiffusionBackend: # 2. if no GGUF file was requested the user is loading a # full diffusers repo; use repo_id directly so we do # not silently substitute the family default - # 3. otherwise fall back to the family default + # 3. otherwise use the family + repo_id heuristic so a + # 9B GGUF picks the 9B base, not the 4B fallback if base_repo: effective_base = base_repo elif not gguf_filename: + # Guard: a repo that ends in "-GGUF" (the unsloth + # convention) is GGUF-only and will 500 on + # from_pretrained; surface a clear error instead of + # letting diffusers raise a confusing model-index + # failure deep in the loader. + if repo_id.lower().endswith("-gguf"): + raise RuntimeError( + f"'{repo_id}' looks like a GGUF-only repo. " + "Either provide gguf_filename to pick a quant, " + "or pass base_repo to override the full-repo " + "load target." + ) effective_base = repo_id else: - effective_base = fam.base_repo + effective_base = _smart_base_repo(fam, repo_id) logger.info( "Loading diffusion model %s (family=%s, device=%s, dtype=%s, base=%s)", repo_id, @@ -370,20 +414,38 @@ class DiffusionBackend: torch_dtype = dtype, ) - pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} + pipe_kwargs: dict[str, Any] = { + "torch_dtype": dtype, + # use_safetensors=True refuses pickle-backed .bin + # weights at load time. Diffusers will fall back to + # safetensors variants on repos that publish both, + # and hard-error on repos that only ship .bin (which + # is the threat model we want to block since pickle + # files can execute arbitrary code in this process). + "use_safetensors": True, + } if transformer is not None: pipe_kwargs["transformer"] = transformer if hf_token: pipe_kwargs["token"] = hf_token + # Release the previous pipeline BEFORE allocating the + # new one so peak VRAM stays at one model's worth, not + # two. This matters on 16-24 GB consumer GPUs where the + # combined footprint would OOM the from_pretrained call. + old = self._pipe + if old is not None: + with self._lock: + self._pipe = None + _release(old) + old = None + pipe = pipeline_cls.from_pretrained(effective_base, **pipe_kwargs) if enable_model_cpu_offload and device == "cuda": pipe.enable_model_cpu_offload() else: pipe.to(device) - # Drop the old pipeline only after the new one is in place. - old = self._pipe with self._lock: self._pipe = pipe self._family = fam @@ -393,7 +455,8 @@ class DiffusionBackend: self._device = device self._dtype = str(dtype).replace("torch.", "") self._loaded_at = time.time() - _release(old) + # ``old`` was released above before the new allocation; + # nothing left to free here. return self.status() except Exception as exc: diff --git a/studio/backend/requirements/no-torch-runtime.txt b/studio/backend/requirements/no-torch-runtime.txt index fa3f33757e..6b7c17a0be 100644 --- a/studio/backend/requirements/no-torch-runtime.txt +++ b/studio/backend/requirements/no-torch-runtime.txt @@ -45,9 +45,11 @@ accelerate>=0.34.1 peft>=0.18.0,!=0.11.0 huggingface_hub>=0.34.0 hf_transfer -diffusers +# Floor 0.37.0 introduces Flux2KleinPipeline + Flux2Pipeline which the +# Studio Images page imports for the default curated picker. +diffusers>=0.37.0 # Required by diffusers.GGUFQuantizationConfig (used by the Images page -# to load FLUX.2 / FLUX.1 / Qwen-Image / SDXL GGUFs from the Hub). +# to load FLUX.2 / FLUX.1 / Qwen-Image GGUFs from the Hub). gguf # Transitive deps required because this file is installed with --no-deps. diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 1a6a8de1c3..b3cf13f10d 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -146,7 +146,7 @@ def test_status_shape_unloaded(): "family", "pipeline_class", "base_repo", - "gguf_path", + "gguf_filename", "device", "dtype", "loaded_at", @@ -373,10 +373,11 @@ def test_load_model_gguf_path_happy(monkeypatch): assert status["is_loaded"] is True assert status["family"] == "flux.2-klein" assert status["pipeline_class"] == "Flux2KleinPipeline" - assert status["base_repo"] == "black-forest-labs/FLUX.2-klein-base-4B" - assert status["gguf_path"] == ( - "/fake/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q4_K_S.gguf" - ) + # _smart_base_repo picks the distilled 4B (not the Base) for the + # "FLUX.2-klein-4B-GGUF" repo name. The Base variant kicks in only + # when "base" is part of the repo id. + assert status["base_repo"] == "black-forest-labs/FLUX.2-klein-4B" + assert status["gguf_filename"] == "flux-2-klein-4b-Q4_K_S.gguf" def test_load_model_recovers_after_failure(monkeypatch): @@ -426,6 +427,80 @@ def test_load_model_base_repo_override(monkeypatch): assert status["base_repo"] == "black-forest-labs/FLUX.2-klein-base-9B" +def test_load_model_gguf_only_repo_without_filename_errors(monkeypatch): + """When the caller points at a -GGUF repo but forgets the filename, + surface a clear error instead of calling from_pretrained on the + GGUF-only repo (which 500s deep in diffusers).""" + _install_fake_diffusers(monkeypatch) + from core.inference.diffusion import get_diffusion_backend + + backend = get_diffusion_backend() + with pytest.raises(RuntimeError, match = "looks like a GGUF-only repo"): + backend.load_model("unsloth/FLUX.2-klein-4B-GGUF") + + +def test_smart_base_repo_picks_9b(monkeypatch): + """For unsloth/FLUX.2-klein-9B-GGUF without an explicit base_repo, + the backend must fall through to FLUX.2-klein-9B, not the 4B base.""" + _install_fake_diffusers(monkeypatch) + from core.inference.diffusion import get_diffusion_backend + + backend = get_diffusion_backend() + status = backend.load_model( + "unsloth/FLUX.2-klein-9B-GGUF", + gguf_filename = "flux-2-klein-9b-Q4_K_S.gguf", + ) + assert status["base_repo"] == "black-forest-labs/FLUX.2-klein-9B" + + +def test_smart_base_repo_picks_base_9b(monkeypatch): + _install_fake_diffusers(monkeypatch) + from core.inference.diffusion import get_diffusion_backend + + backend = get_diffusion_backend() + status = backend.load_model( + "unsloth/FLUX.2-klein-base-9B-GGUF", + gguf_filename = "flux-2-klein-base-9b-Q4_K_S.gguf", + ) + assert status["base_repo"] == "black-forest-labs/FLUX.2-klein-base-9B" + + +def test_smart_base_repo_picks_base_4b(monkeypatch): + _install_fake_diffusers(monkeypatch) + from core.inference.diffusion import get_diffusion_backend + + backend = get_diffusion_backend() + status = backend.load_model( + "unsloth/FLUX.2-klein-base-4B-GGUF", + gguf_filename = "flux-2-klein-base-4b-Q4_K_S.gguf", + ) + assert status["base_repo"] == "black-forest-labs/FLUX.2-klein-base-4B" + + +def test_load_model_uses_safetensors_flag(monkeypatch): + """The pipeline.from_pretrained call must pass use_safetensors=True + so pickle-backed .bin weights are refused at load time.""" + fake = _install_fake_diffusers(monkeypatch) + from core.inference.diffusion import get_diffusion_backend + + captured: dict = {} + + original = fake.Flux2KleinPipeline.from_pretrained.__func__ + + def _capture(cls, base_repo, **kw): + captured.update(kw) + return original(cls, base_repo, **kw) + + fake.Flux2KleinPipeline.from_pretrained = classmethod(_capture) + + backend = get_diffusion_backend() + backend.load_model( + "unsloth/FLUX.2-klein-base-4B-GGUF", + gguf_filename = "flux-2-klein-base-4b-Q4_K_S.gguf", + ) + assert captured.get("use_safetensors") is True + + def test_load_model_full_repo_does_not_substitute(monkeypatch): """A full diffusers repo (no gguf_filename) must call from_pretrained with the user-supplied repo, not the family default. This was the diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index af759ab234..ca420f72c5 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -43,7 +43,7 @@ class _FakeBackend: "family": "flux.2-klein" if self._loaded else None, "pipeline_class": "Flux2KleinPipeline" if self._loaded else None, "base_repo": "black-forest-labs/FLUX.2-klein" if self._loaded else None, - "gguf_path": None, + "gguf_filename": None, "device": "cpu", "dtype": "torch.bfloat16", "loaded_at": 0, diff --git a/studio/frontend/src/features/images/api.ts b/studio/frontend/src/features/images/api.ts index 017b856b5a..e576f3987e 100644 --- a/studio/frontend/src/features/images/api.ts +++ b/studio/frontend/src/features/images/api.ts @@ -23,7 +23,7 @@ export interface DiffusionStatus { family: string | null; pipeline_class: string | null; base_repo: string | null; - gguf_path: string | null; + gguf_filename: string | null; device: string | null; dtype: string | null; loaded_at: number | null; diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index abca68c0f2..1cc795345c 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -59,17 +59,20 @@ const CURATED_MODELS: Array<{ label: "FLUX.2 klein 4B (Q4_K_S, distilled)", repo_id: "unsloth/FLUX.2-klein-4B-GGUF", default_gguf: "flux-2-klein-4b-Q4_K_S.gguf", - base_repo: "black-forest-labs/FLUX.2-klein-base-4B", + // Distilled GGUF must pair with the distilled base, not the Base + // checkpoint. The Hub model card for the GGUF lists + // base_model: black-forest-labs/FLUX.2-klein-4B. + base_repo: "black-forest-labs/FLUX.2-klein-4B", family: "flux.2-klein", - notes: "13 GB VRAM. Distilled klein 4B with the Apache base.", + notes: "13 GB VRAM. Distilled klein 4B. Requires HF access to FLUX.2 klein 4B.", }, { label: "FLUX.2 klein 9B (Q4_K_S, gated)", repo_id: "unsloth/FLUX.2-klein-9B-GGUF", default_gguf: "flux-2-klein-9b-Q4_K_S.gguf", - base_repo: "black-forest-labs/FLUX.2-klein-base-9B", + base_repo: "black-forest-labs/FLUX.2-klein-9B", family: "flux.2-klein", - notes: "17 GB VRAM. Higher quality. Requires HF access to FLUX.2 klein base 9B.", + notes: "17 GB VRAM. Higher quality distilled. Requires HF access to FLUX.2 klein 9B.", }, { label: "FLUX.2 dev (Q4_K_S, gated)",