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:
parent
355f54f6e4
commit
977c9b7dd0
5 changed files with 290 additions and 13 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
129
studio/backend/tests/test_diffusion_sdxl.py
Normal file
129
studio/backend/tests/test_diffusion_sdxl.py
Normal 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
|
||||
)
|
||||
|
|
@ -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 } {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue