Add SDXL diffusion family (U-Net pipeline support)

SDXL is the first U-Net family in the diffusion backend: its denoiser is
pipe.unet (UNet2DConditionModel), not a DiT pipe.transformer, and a single-file
.safetensors is the whole pipeline rather than a transformer-only file. The
backend previously assumed a DiT transformer everywhere, so add the two hooks a
U-Net family needs and register SDXL.

DiffusionFamily gains denoiser_attr ("transformer" for DiT, "unet" for SDXL) and
single_file_is_pipeline (SDXL loads a single file via pipeline_class.from_single_file
with the base repo as config, instead of transformer_class.from_single_file plus a
companion assembly). _align_vae_dtype now reads the denoiser generically so img2img
and inpaint keep the VAE and U-Net dtypes aligned.

The non-GGUF trust gate is extended with a short, exact-match, safetensors-only
allowlist of official base repos (the SDXL base/refiner and sdxl-turbo), because
SDXL ships only as a full pipeline and has no unsloth-hosted GGUF. Local paths stay
trusted as before; a random repo, even one that detects as SDXL, is still rejected.

The image-conditioned and ControlNet workflows are the standard SDXL pipelines,
built around the resident modules via from_pipe like every other family, so SDXL
gets txt2img, img2img, inpaint, outpaint, upscale, LoRA and ControlNet. There is no
native sd.cpp mapping yet, so the no-GPU route falls back to diffusers.

Frontend catalog gains SDXL Base 1.0 and SDXL Turbo entries with SDXL step/guidance
defaults (Turbo: few steps, no CFG; base: ~30 steps, real CFG).

Tests: new test_diffusion_sdxl.py (family shape, detection, trust allowlist, model
kind, U-Net VAE-dtype alignment, LoRA gate) plus loader-branch tests in
test_diffusion_backend.py (pipeline-kind from_pretrained, single-file whole-pipeline
from_single_file, allowlist accept/reject). Verified live on GPU: sdxl-turbo loads
both as a pipeline and as a single file and generates coherent txt2img + img2img.
This commit is contained in:
Daniel Han 2026-07-01 13:57:07 +00:00
commit 977c9b7dd0
5 changed files with 290 additions and 13 deletions

View file

@ -183,17 +183,35 @@ def _snap_to_multiple(img: Any, multiple: int = 16) -> Any:
return img
# A small allowlist of well-known official base repos that may load as a full
# (non-GGUF) pipeline even though they are not under ``unsloth/``. These are
# safetensors-only checkpoints from their original publisher (no pickle, no remote
# code) that some architectures require: SDXL ships only as a full pipeline and has
# 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.
_TRUSTED_NON_GGUF_REPOS = frozenset(
{
"stabilityai/stable-diffusion-xl-base-1.0",
"stabilityai/stable-diffusion-xl-refiner-1.0",
"stabilityai/sdxl-turbo",
}
)
def _is_trusted_diffusion_repo(repo_id: str) -> bool:
"""Whether a NON-GGUF load is allowed for ``repo_id``.
Making ``gguf_filename`` optional opens a ``from_pretrained`` / ``from_single_file``
on an arbitrary repo, which fetches and deserialises third-party weights. So the
non-GGUF paths are gated to the ``unsloth/*`` org (the curated safetensors models) and
to local paths the user explicitly pointed at (already on their disk). The GGUF path
is unchanged and stays open to any repo, as before."""
non-GGUF paths are gated to the ``unsloth/*`` org (the curated safetensors models),
a short allowlist of official safetensors-only base repos (``_TRUSTED_NON_GGUF_REPOS``,
e.g. the SDXL base), and local paths the user explicitly pointed at (already on their
disk). The GGUF path is unchanged and stays open to any repo, as before."""
if Path(repo_id).expanduser().exists():
return True
return repo_id.strip().lower().startswith("unsloth/")
rid = repo_id.strip().lower()
return rid.startswith("unsloth/") or rid in _TRUSTED_NON_GGUF_REPOS
@dataclass(frozen = True)
@ -819,6 +837,16 @@ class DiffusionBackend:
if hf_token:
pipe_kwargs["token"] = hf_token
pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs)
elif kind == "single_file" and fam.single_file_is_pipeline:
# A single-file SDXL-style checkpoint is the WHOLE pipeline
# (U-Net + VAE + both text encoders), not a transformer-only file,
# so load it through the pipeline class. ``config`` points at the
# base repo so diffusers builds the correct structure/scheduler
# around the single-file weights instead of guessing from the file.
sf_pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "config": base}
if hf_token:
sf_pipe_kwargs["token"] = hf_token
pipe = pipeline_cls.from_single_file(single_file_path, **sf_pipe_kwargs)
else:
# Single-file transformer; the VAE / text-encoder / scheduler come
# from the base diffusers repo (the single file is transformer-only).
@ -1228,20 +1256,21 @@ class DiffusionBackend:
return pipe
@staticmethod
def _align_vae_dtype(pipe: Any) -> None:
"""Cast the VAE to the transformer's compute dtype before an image-conditioned
def _align_vae_dtype(pipe: Any, denoiser_attr: str = "transformer") -> None:
"""Cast the VAE to the denoiser's compute dtype before an image-conditioned
call. The img2img/inpaint pipelines VAE-encode the input image at the text-
encoder dtype (bf16), but a prior txt2img DECODE may have left the shared VAE
upcast to fp32 (its ``force_upcast`` path), so the encode would mismatch
(bf16 image vs fp32 VAE). Re-aligning here is safe: our families run bf16 or
fp32 only (the fp16 guard promotes fp16), and a later txt2img decode re-upcasts
as needed. Best-effort; a no-op when already aligned."""
transformer = getattr(pipe, "transformer", None)
as needed. ``denoiser_attr`` is ``pipe.transformer`` for DiT families and
``pipe.unet`` for SDXL. Best-effort; a no-op when already aligned."""
denoiser = getattr(pipe, denoiser_attr, None)
vae = getattr(pipe, "vae", None)
if transformer is None or vae is None:
if denoiser is None or vae is None:
return
try:
target_dtype = transformer.dtype
target_dtype = denoiser.dtype
if next(vae.parameters()).dtype != target_dtype:
vae.to(dtype = target_dtype)
except (StopIteration, AttributeError, RuntimeError):
@ -1509,7 +1538,7 @@ 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)
self._align_vae_dtype(pipe, getattr(state.family, "denoiser_attr", "transformer"))
# 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

View file

@ -29,6 +29,18 @@ class DiffusionFamily:
# Pipeline kwarg carrying the guidance value. Most use "guidance_scale";
# Qwen-Image's distilled guidance is off, so its real CFG is "true_cfg_scale".
cfg_kwarg: str = "guidance_scale"
# The pipe attribute holding the denoiser module. DiT families expose it as
# ``pipe.transformer`` (the default); U-Net families (SDXL) as ``pipe.unet``.
# Read wherever the backend touches the denoiser generically (VAE dtype
# alignment, optimisation guards), so a U-Net family works without assuming a
# ``transformer`` attribute exists.
denoiser_attr: str = "transformer"
# True when a single-file ``.safetensors`` checkpoint is the WHOLE pipeline
# (U-Net + VAE + text encoders), not a transformer-only file. SDXL ships this
# way, so the loader calls ``pipeline_class.from_single_file`` on it directly
# rather than ``transformer_class.from_single_file`` + a companion base repo.
# DiT families leave this False (their single file is transformer-only).
single_file_is_pipeline: bool = False
# Optional diffusers pipeline classes for image-conditioned workflows. The backend
# builds these around the ALREADY-loaded transformer/VAE/text-encoder via
# ``Pipeline.from_pipe`` (no extra weights, no reload), so a family only needs the
@ -243,6 +255,30 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"),
),
),
# SDXL is the one U-Net family here: the denoiser is ``pipe.unet``
# (UNet2DConditionModel), not a DiT ``pipe.transformer``, and a single-file
# ``.safetensors`` is the WHOLE pipeline rather than a transformer-only file.
# So it declares ``denoiser_attr = "unet"`` + ``single_file_is_pipeline = True``
# and loads via the pipeline class (from_pretrained for a repo, from_single_file
# for a single .safetensors). The base repo supplies both CLIP text encoders,
# the VAE and the scheduler on the pipeline path. img2img / inpaint / ControlNet
# are the standard SDXL pipelines, built around the resident modules via
# from_pipe like every other family. There is no GGUF/single-file transformer
# path for SDXL (the whole checkpoint is one file), and no native sd.cpp mapping
# yet, so the no-GPU route falls back to diffusers.
DiffusionFamily(
name = "sdxl",
pipeline_class = "StableDiffusionXLPipeline",
transformer_class = "UNet2DConditionModel",
base_repo = "stabilityai/stable-diffusion-xl-base-1.0",
aliases = ("stable-diffusion-xl", "sd-xl", "sd_xl", "sdxl-turbo", "sdxl-base"),
denoiser_attr = "unet",
single_file_is_pipeline = True,
img2img_pipeline_class = "StableDiffusionXLImg2ImgPipeline",
inpaint_pipeline_class = "StableDiffusionXLInpaintPipeline",
controlnet_pipeline_class = "StableDiffusionXLControlNetPipeline",
controlnet_model_class = "ControlNetModel",
),
)
# Editing / inpaint checkpoints share an arch keyword but need a different

View file

@ -209,12 +209,19 @@ class _FakePipe:
class _FakePipeline:
last: dict = {}
last_single_file: dict = {}
@classmethod
def from_pretrained(cls, base, **kwargs):
_FakePipeline.last = {"base": base, **kwargs}
return _FakePipe()
@classmethod
def from_single_file(cls, path, **kwargs):
# SDXL-style single-file: the WHOLE pipeline comes from one .safetensors file.
_FakePipeline.last_single_file = {"path": path, **kwargs}
return _FakePipe()
class _FakeTransformer:
last: dict = {}
@ -330,6 +337,13 @@ def fake_runtime(monkeypatch):
diffusers.QwenImageInpaintPipeline = _FakeInpaintPipeline
# Instruction-editing pipeline (Qwen-Image-Edit): its own pipeline IS the loaded one.
diffusers.QwenImageEditPlusPipeline = _FakePipeline
# SDXL: a U-Net family. Its single-file checkpoint is the whole pipeline, so the
# pipeline class carries from_single_file; UNet2DConditionModel is the denoiser
# class (fetched but unused on the pipeline/single-file-pipeline paths).
diffusers.StableDiffusionXLPipeline = _FakePipeline
diffusers.UNet2DConditionModel = _FakeTransformer
diffusers.StableDiffusionXLImg2ImgPipeline = _FakeImg2ImgPipeline
diffusers.StableDiffusionXLInpaintPipeline = _FakeInpaintPipeline
monkeypatch.setitem(sys.modules, "torch", torch)
monkeypatch.setitem(sys.modules, "diffusers", diffusers)
@ -337,6 +351,7 @@ def fake_runtime(monkeypatch):
# run real hardware detection against the stubbed torch.
monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: None)
_FakePipeline.last = {}
_FakePipeline.last_single_file = {}
_FakeTransformer.last = {}
_FakeImg2ImgPipeline.built_from = None
_FakeImg2ImgPipe.last_kwargs = {}
@ -847,12 +862,64 @@ def test_load_single_file_safetensors_no_gguf_config(fake_runtime, tmp_path):
assert "transformer" in _FakePipeline.last
def test_load_sdxl_pipeline_from_pretrained(fake_runtime):
"""SDXL as a full pipeline (no single-file name) loads via pipeline_cls.from_pretrained
on the allowlisted official base repo -- no U-Net single-file build, no GGUF config.
A U-Net family must NOT try to build a transformer from a single file."""
backend = DiffusionBackend()
status = backend.load_pipeline("stabilityai/stable-diffusion-xl-base-1.0")
assert status["loaded"] is True
assert status["family"] == "sdxl"
assert _FakePipeline.last["base"] == "stabilityai/stable-diffusion-xl-base-1.0"
assert "transformer" not in _FakePipeline.last
# Neither single-file path (transformer-only nor whole-pipeline) was taken.
assert _FakeTransformer.last == {}
assert _FakePipeline.last_single_file == {}
def test_load_sdxl_single_file_uses_pipeline_from_single_file(fake_runtime, tmp_path):
"""A single-file SDXL *.safetensors is the WHOLE pipeline: it must load via
pipeline_cls.from_single_file(path, config=base), NOT transformer_cls.from_single_file
(UNet2DConditionModel has no companion-transformer assembly here)."""
(tmp_path / "sdxl.safetensors").write_bytes(b"weights")
backend = DiffusionBackend()
status = backend.load_pipeline(
str(tmp_path), gguf_filename = "sdxl.safetensors", family_override = "sdxl"
)
assert status["loaded"] is True
assert status["family"] == "sdxl"
# The whole-pipeline single-file path was taken with the base repo as config.
assert _FakePipeline.last_single_file["path"] == str(
(tmp_path / "sdxl.safetensors").resolve()
)
assert _FakePipeline.last_single_file["config"] == "stabilityai/stable-diffusion-xl-base-1.0"
# The transformer-only single-file build was NOT taken.
assert _FakeTransformer.last == {}
def test_load_sdxl_allowlisted_turbo_repo_is_trusted(fake_runtime):
"""The official sdxl-turbo repo is on the non-GGUF allowlist, so a full-pipeline load
is permitted even though it is not under unsloth/*."""
backend = DiffusionBackend()
status = backend.load_pipeline("stabilityai/sdxl-turbo")
assert status["loaded"] is True
assert status["family"] == "sdxl"
def test_load_pipeline_rejects_non_unsloth_repo(fake_runtime):
backend = DiffusionBackend()
with pytest.raises(ValueError, match = "unsloth"):
backend.load_pipeline("randomorg/Z-Image-bnb-4bit", family_override = "z-image")
def test_load_sdxl_rejects_untrusted_repo(fake_runtime):
"""A random non-allowlisted, non-unsloth repo is still rejected for a full pipeline
load even when it detects as SDXL -- the allowlist is exact-match only."""
backend = DiffusionBackend()
with pytest.raises(ValueError, match = "unsloth"):
backend.load_pipeline("randomorg/my-sdxl-merge", family_override = "sdxl")
def test_detect_family_rejects_layered():
# Qwen-Image-Layered needs a dedicated pipeline (additional_t_cond); it must be
# rejected so it fails fast at load instead of crashing at the first denoise step.

View file

@ -0,0 +1,129 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""CPU-only unit tests for the SDXL diffusion family.
SDXL is the one U-Net family: the denoiser is ``pipe.unet`` (not ``pipe.transformer``)
and a single-file ``.safetensors`` is the whole pipeline (not a transformer-only file).
These tests cover the pure helpers that encode those differences -- family detection,
the ``denoiser_attr`` / ``single_file_is_pipeline`` flags, the non-GGUF trust allowlist,
the VAE-dtype alignment reading the U-Net denoiser, and the LoRA-support gate -- with no
torch/diffusers/GPU needed.
"""
from __future__ import annotations
import types
from core.inference import diffusion_lora
from core.inference.diffusion import (
DiffusionBackend,
_is_trusted_diffusion_repo,
resolve_model_kind,
)
from core.inference.diffusion_families import detect_family, family_sd_cpp_supported
def test_sdxl_family_shape():
fam = detect_family("stabilityai/stable-diffusion-xl-base-1.0")
assert fam is not None and fam.name == "sdxl"
assert fam.pipeline_class == "StableDiffusionXLPipeline"
# The denoiser is a U-Net, addressed via pipe.unet (DiT families use pipe.transformer).
assert fam.denoiser_attr == "unet"
assert fam.transformer_class == "UNet2DConditionModel"
# A single-file SDXL checkpoint is the whole pipeline, loaded via the pipeline class.
assert fam.single_file_is_pipeline is True
# Image-conditioned + ControlNet workflows are the standard SDXL pipelines.
assert fam.img2img_pipeline_class == "StableDiffusionXLImg2ImgPipeline"
assert fam.inpaint_pipeline_class == "StableDiffusionXLInpaintPipeline"
assert fam.controlnet_pipeline_class == "StableDiffusionXLControlNetPipeline"
assert fam.controlnet_model_class == "ControlNetModel"
# Real CFG; SDXL uses guidance_scale, not a distilled true_cfg_scale.
assert fam.cfg_kwarg == "guidance_scale"
def test_sdxl_detection_by_repo_and_override():
assert detect_family("stabilityai/sdxl-turbo").name == "sdxl"
assert detect_family("some-org/My-Cool-SDXL-Merge").name == "sdxl"
assert detect_family("some-org/stable-diffusion-xl-anime").name == "sdxl"
assert detect_family("x", override="sdxl").name == "sdxl"
# A GGUF DiT family must NOT be swallowed by the SDXL match.
assert detect_family("unsloth/FLUX.1-schnell-GGUF").name == "flux.1"
def test_dit_families_keep_transformer_denoiser():
# The generalisation must not change existing DiT families: they stay on
# pipe.transformer and their single file is transformer-only.
for rid in ("unsloth/FLUX.1-schnell-GGUF", "unsloth/Qwen-Image-GGUF", "unsloth/Z-Image-GGUF"):
fam = detect_family(rid)
assert fam.denoiser_attr == "transformer"
assert fam.single_file_is_pipeline is False
def test_sdxl_has_no_native_sd_cpp_mapping():
# No single-file VAE/TE mapping yet, so the no-GPU route falls back to diffusers
# rather than trying to drive sd-cli.
assert family_sd_cpp_supported(detect_family("stabilityai/sdxl-turbo")) is False
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")
# 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.
assert not _is_trusted_diffusion_repo("randomorg/my-sdxl-merge")
assert not _is_trusted_diffusion_repo("stabilityai/sdxl-turbo-evil")
def test_sdxl_model_kind_resolution():
# A full-pipeline load (no single-file name) is "pipeline"; a single .safetensors
# is "single_file" (handled by the whole-pipeline branch for SDXL).
assert resolve_model_kind(None) == "pipeline"
assert resolve_model_kind("sdxl.safetensors") == "single_file"
class _FakeVae:
def __init__(self, dtype):
self._dtype = dtype
self.moved_to = None
def parameters(self):
yield types.SimpleNamespace(dtype=self._dtype)
def to(self, dtype=None):
self.moved_to = dtype
self._dtype = dtype
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.
vae = _FakeVae(dtype="float32")
pipe = types.SimpleNamespace(unet=types.SimpleNamespace(dtype="bfloat16"), vae=vae)
DiffusionBackend._align_vae_dtype(pipe, "unet")
assert vae.moved_to == "bfloat16"
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)
DiffusionBackend._align_vae_dtype(pipe)
assert vae.moved_to == "bfloat16"
# No denoiser attribute -> no-op (does not raise, does not move the VAE).
vae2 = _FakeVae(dtype="float32")
DiffusionBackend._align_vae_dtype(types.SimpleNamespace(vae=vae2), "unet")
assert vae2.moved_to is None
def test_sdxl_lora_supported_on_diffusers():
# SDXL is bf16/bnb-4bit on diffusers -> LoRA is allowed (unlike GGUF-via-diffusers).
assert diffusion_lora.supports_lora(
engine="diffusers", family="sdxl", model_kind="pipeline", transformer_quant=None
)
assert diffusion_lora.supports_lora(
engine="diffusers", family="sdxl", model_kind="single_file", transformer_quant=None
)

View file

@ -90,8 +90,9 @@ const editGguf = (id: string, name: string): ModelOption => ({
// How to load a curated non-GGUF (safetensors) model. "pipeline" = a full diffusers
// repo (from_pretrained, embedded bnb-4bit quant auto-applied); "single_file" = a
// single safetensors transformer (e.g. fp8) assembled onto its base repo. The backend
// gates these to unsloth/* repos. Keyed by repo id so the load handler knows the kind
// (and, for single_file, the exact filename).
// gates these to unsloth/* repos plus a short allowlist of official base repos (SDXL).
// Keyed by repo id so the load handler knows the kind (and, for single_file, the exact
// filename).
type SafetensorsSpec = { kind: "pipeline" | "single_file"; filename?: string };
const SAFETENSORS_MODELS: Record<string, SafetensorsSpec> = {
"unsloth/Z-Image-Turbo-unsloth-bnb-4bit": { kind: "pipeline" },
@ -100,6 +101,10 @@ const SAFETENSORS_MODELS: Record<string, SafetensorsSpec> = {
kind: "single_file",
filename: "qwen-image-2512-fp8.safetensors",
},
// SDXL is a U-Net family loaded as a whole pipeline (from_pretrained). These
// official base repos are on the backend's non-GGUF allowlist.
"stabilityai/sdxl-turbo": { kind: "pipeline" },
"stabilityai/stable-diffusion-xl-base-1.0": { kind: "pipeline" },
};
// Curated non-GGUF picker entries (isGguf:false -> no quant expander, direct load).
const safetensors = (id: string, name: string, label: string): ModelOption => ({
@ -135,6 +140,12 @@ const MODELS: ModelOption[] = [
"Qwen-Image 2512 (FP8)",
"Safetensors · fp8",
),
safetensors("stabilityai/sdxl-turbo", "SDXL Turbo", "Safetensors · SDXL"),
safetensors(
"stabilityai/stable-diffusion-xl-base-1.0",
"SDXL Base 1.0",
"Safetensors · SDXL",
),
];
// Workflow tabs. `requires` is the backend workflow id (status.workflows) that must
@ -204,6 +215,11 @@ const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }>
{ match: "flux.2-dev", steps: 28, guidance: 4 },
{ match: "qwen-image", steps: 20, guidance: 4 },
{ match: "z-image", steps: 20, guidance: 4 },
// SDXL: Turbo is distilled (few steps, no CFG); base/full SDXL wants ~30 steps and
// real CFG (~7). "sdxl-turbo" must precede the generic "sdxl" substring match.
{ match: "sdxl-turbo", steps: 3, guidance: 0 },
{ match: "stable-diffusion-xl", steps: 30, guidance: 7 },
{ match: "sdxl", steps: 30, guidance: 7 },
];
function defaultsFor(repoId: string): { steps: number; guidance: number } {