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

SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers

Browse source 
Addresses review findings on the SDXL family:
- Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request,
  before the route evicts the current model; SDXL has no transformer-only GGUF variant.
- Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file
  (config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no
  longer triggers a multi-GB base download.
- Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline
  but this backend loads every sdxl repo as the base txt2img pipeline.
- Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch
  degrades to anonymous instead of erroring on a malformed token.
- Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and
  access state.family.denoiser_attr directly.

Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter.
This commit is contained in:
Daniel Han 2026-07-02 01:13:18 +00:00
commit a4fb348f9f
2 changed files with 113 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -191,10 +191,13 @@ def _snap_to_multiple(img: Any, multiple: int = 16) -> Any:
# no unsloth-hosted GGUF, so without this its curated catalog entry could not load.
# Exact-match, lowercased, so it cannot be widened by a typo-squat. Extend
# deliberately, and never add a repo that carries pickled weights or remote code.
# The SDXL refiner is intentionally NOT here: it is an img2img-only refiner pipeline
# (StableDiffusionXLImg2ImgPipeline), but this backend loads every ``sdxl`` repo as the
# base txt2img StableDiffusionXLPipeline and advertises txt2img, so allowlisting the
# refiner would surface the wrong workflow and call it without its required input image.
_TRUSTED_NON_GGUF_REPOS = frozenset(
{
"stabilityai/stable-diffusion-xl-base-1.0",
"stabilityai/stable-diffusion-xl-refiner-1.0",
"stabilityai/sdxl-turbo",
}
)
@ -434,6 +437,16 @@ class DiffusionBackend:
f"pass family_override with that family name. (Video models and image models "
f"whose diffusers transformer has no single-file loader are not supported.)"
)
# A GGUF load builds a transformer-only file via the generic GGUF branch
# (UNet2DConditionModel.from_single_file(subfolder="transformer", GGUFQuantizationConfig)).
# Families whose single file IS the whole pipeline (SDXL) have no transformer-only
# GGUF path, so reject GGUF here -- before the route evicts the current model and
# the background load fails deep in from_single_file.
if kind == "gguf" and fam.single_file_is_pipeline:
raise ValueError(
f"'{fam.name}' checkpoints are whole-pipeline single files and have no GGUF "
f"transformer variant; load the .safetensors pipeline instead of a GGUF."
)
# Non-GGUF loads (a single-file safetensors transformer, or a full pipeline)
# are gated to the unsloth org or a local path -- they fetch + deserialise
# weights, so an arbitrary remote repo is rejected here, before any work.
@ -568,6 +581,7 @@ class DiffusionBackend:
base,
kwargs.get("hf_token"),
kind = kind,
single_file_is_pipeline = bool(fam and fam.single_file_is_pipeline),
)
with self._lock:
# Stamp progress only if this load is still current; a superseding
@ -634,6 +648,7 @@ class DiffusionBackend:
hf_token: Optional[str],
*,
kind: str = "gguf",
single_file_is_pipeline: bool = False,
) -> tuple[int, list[str]]:
"""Total download size for the progress bar, plus the base-repo files to
fetch (the prefetch reuses this list, so the base is listed only once).
@ -641,7 +656,9 @@ class DiffusionBackend:
For a ``pipeline`` load the whole repo IS the pipeline (``base_repo`` is the
repo itself), so the transformer/ subfolder is INCLUDED -- unlike the GGUF /
single-file paths, where the transformer is the single file and the base repo
supplies only the companions."""
supplies only the companions. For a ``single_file_is_pipeline`` family (SDXL) the
single file is the WHOLE pipeline, so the base repo supplies only config/tokenizer
(no weights) and its weight files are skipped."""
from huggingface_hub import HfApi
api = HfApi()
@ -671,9 +688,16 @@ class DiffusionBackend:
if gguf_filename and not Path(repo_id).expanduser().exists():
info = api.model_info(repo_id, files_metadata = True, token = hf_token)
total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename)
# A whole-pipeline single file (SDXL) needs only the base repo's config/tokenizer,
# not its (unused, multi-GB) weight files.
base_filter = (
_base_config_file_downloaded
if (kind == "single_file" and single_file_is_pipeline)
else _base_file_downloaded
)
base_info = api.model_info(base_repo, files_metadata = True, token = hf_token)
for s in base_info.siblings:
if _base_file_downloaded(s.rfilename):
if base_filter(s.rfilename):
base_files.append(s.rfilename)
total += s.size or 0
except Exception as exc: # noqa: BLE001 — estimate is best-effort
@ -719,6 +743,13 @@ class DiffusionBackend:
model_kind: Optional[str] = None,
_load_token: Optional[int] = None,
) -> dict[str, Any]:
# A blank / whitespace-only token must degrade to anonymous access, not be passed
# as an explicit credential (from_single_file / from_pretrained / the Hub client
# can error on a malformed token instead of falling back). Normalize once here so
# every load branch and the size estimate below use a real token or None.
hf_token = hf_token.strip() if isinstance(hf_token, str) else hf_token
hf_token = hf_token or None
# Validate first (cheap, no torch/diffusers) so a direct call with a bad
# family fails with ValueError even in a no-diffusers runtime.
fam = self.validate_load_request(
@ -1288,7 +1319,10 @@ class DiffusionBackend:
if denoiser is None or vae is None:
return
try:
target_dtype = denoiser.dtype
# Read the dtype from a parameter (not denoiser.dtype): a plain nn.Module has no
# .dtype, and a torch.compile'd/ wrapped denoiser can obscure it; this also
# matches how the VAE dtype is read on the next line.
target_dtype = next(denoiser.parameters()).dtype
if next(vae.parameters()).dtype != target_dtype:
vae.to(dtype = target_dtype)
except (StopIteration, AttributeError, RuntimeError):
@ -1556,9 +1590,8 @@ class DiffusionBackend:
mask_pil = mask_pil.resize(init_pil.size, _PILImage.NEAREST)
if init_pil is not None:
# Keep the VAE encode dtype consistent with the input image.
self._align_vae_dtype(
pipe, getattr(state.family, "denoiser_attr", "transformer")
)
# state.family is always a DiffusionFamily, which defines denoiser_attr.
self._align_vae_dtype(pipe, state.family.denoiser_attr)
# Pipelines vary in which kwargs they accept (img2img derives size from the
# input image and may reject width/height; a distilled pipe may take no
@ -1876,6 +1909,33 @@ def _base_file_downloaded(rfilename: str) -> bool:
return not rfilename.startswith("assets/")
# Weight file extensions the base repo need NOT supply when the single file is the whole
# pipeline (SDXL): from_single_file(config=base) reads only the base repo's structure
# (config/tokenizer/scheduler) and takes the weights from the single file.
_BASE_WEIGHT_EXTS = (
".safetensors",
".bin",
".ckpt",
".pt",
".pth",
".gguf",
".onnx",
".onnx_data",
".msgpack",
".h5",
".pb",
)
def _base_config_file_downloaded(rfilename: str) -> bool:
"""True for base-repo files needed to BUILD a pipeline structure around a whole-pipeline
single file WITHOUT its weights: config / tokenizer / scheduler JSON, but no weight
tensors (the single file supplies those). Used for ``single_file_is_pipeline`` families."""
if not _base_file_downloaded(rfilename):
return False
return not rfilename.lower().endswith(_BASE_WEIGHT_EXTS)
def _pipeline_file_downloaded(rfilename: str) -> bool:
"""True for files a full-pipeline ``from_pretrained`` fetches.

50
studio/backend/tests/test_diffusion_sdxl.py
View file

@ -15,6 +15,8 @@ from __future__ import annotations
import types
import pytest
from core.inference import diffusion_lora
from core.inference.diffusion import (
DiffusionBackend,
@ -70,7 +72,8 @@ def test_sdxl_base_repos_are_trusted_non_gguf():
# Official safetensors-only base repos are allowlisted so their catalog entries load.
assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0")
assert _is_trusted_diffusion_repo("stabilityai/sdxl-turbo")
assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-refiner-1.0")
# The refiner is img2img-only and is intentionally NOT allowlisted (see
# test_sdxl_refiner_not_trusted).
# Case-insensitive match.
assert _is_trusted_diffusion_repo("StabilityAI/SDXL-Turbo")
# A random repo (even one that detects as SDXL) is NOT trusted for a non-GGUF load.
@ -100,9 +103,11 @@ class _FakeVae:
def test_align_vae_dtype_uses_unet_denoiser():
# For SDXL the denoiser lives at pipe.unet; _align_vae_dtype must read it (a pipe
# with only .unet and no .transformer) and cast the VAE to the U-Net's dtype.
# with only .unet and no .transformer) and cast the VAE to the U-Net's dtype. The
# dtype is read from a parameter (denoiser has no .dtype), so use a _FakeVae denoiser.
vae = _FakeVae(dtype = "float32")
pipe = types.SimpleNamespace(unet = types.SimpleNamespace(dtype = "bfloat16"), vae = vae)
unet = _FakeVae(dtype = "bfloat16")
pipe = types.SimpleNamespace(unet = unet, vae = vae)
DiffusionBackend._align_vae_dtype(pipe, "unet")
assert vae.moved_to == "bfloat16"
@ -110,7 +115,8 @@ def test_align_vae_dtype_uses_unet_denoiser():
def test_align_vae_dtype_transformer_default_unchanged():
# DiT default: reads pipe.transformer; a pipe with no transformer is a safe no-op.
vae = _FakeVae(dtype = "float32")
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace(dtype = "bfloat16"), vae = vae)
transformer = _FakeVae(dtype = "bfloat16")
pipe = types.SimpleNamespace(transformer = transformer, vae = vae)
DiffusionBackend._align_vae_dtype(pipe)
assert vae.moved_to == "bfloat16"
# No denoiser attribute -> no-op (does not raise, does not move the VAE).
@ -147,3 +153,39 @@ def test_pipeline_prefetch_skips_non_torch_artifacts():
assert not keep("unet/flax_model.msgpack")
assert not keep("vae_decoder/model.onnx_data")
assert not keep("assets/preview.png")
def test_sdxl_refiner_not_trusted():
# The refiner is an img2img-only pipeline; the sdxl family loads every repo as the
# base txt2img pipeline, so the refiner must NOT be allowlisted for a non-GGUF load.
assert not _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-refiner-1.0")
# The base and turbo remain trusted.
assert _is_trusted_diffusion_repo("stabilityai/stable-diffusion-xl-base-1.0")
assert _is_trusted_diffusion_repo("stabilityai/sdxl-turbo")
def test_sdxl_gguf_load_rejected_up_front():
# SDXL has no transformer-only GGUF variant (its single file is the whole pipeline),
# so a GGUF request must fail cheap validation before the GPU handoff.
backend = DiffusionBackend()
with pytest.raises(ValueError, match = "no GGUF"):
backend.validate_load_request(
"some-org/my-sdxl.gguf", gguf_filename = "my-sdxl.gguf", family_override = "sdxl"
)
def test_base_config_filter_skips_weights():
# For a whole-pipeline single file, the base repo supplies only config/tokenizer, not
# its (unused) weight tensors.
from core.inference.diffusion import _base_config_file_downloaded as keep
assert keep("model_index.json")
assert keep("text_encoder/config.json")
assert keep("tokenizer/vocab.json")
assert keep("scheduler/scheduler_config.json")
assert not keep("unet/diffusion_pytorch_model.safetensors")
assert not keep("vae/diffusion_pytorch_model.bin")
assert not keep("text_encoder/model.onnx")
# transformer/ and assets/ stay excluded (inherited from _base_file_downloaded).
assert not keep("transformer/config.json")
assert not keep("assets/x.png")