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

Fix/adjust diffusion: smart base, safetensors, peak VRAM, GGUF guard

Browse source 
- _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.
This commit is contained in:
Daniel Han-Chen 2026-05-25 00:04:12 +00:00
commit 8074a2b67b
6 changed files with 163 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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:

6
studio/backend/requirements/no-torch-runtime.txt
View file

@ -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.

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

@ -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

2
studio/backend/tests/test_diffusion_routes.py
View file

@ -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,

2
studio/frontend/src/features/images/api.ts
View file

@ -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;

11
studio/frontend/src/features/images/images-page.tsx
View file

@ -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)",