diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 17ae1c5687..53d83ac4a4 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -181,26 +181,47 @@ 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. +# 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/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. A bare ``owner/name`` HF id is never a real filesystem path, and an id with invalid characters makes ``Path.exists()`` raise OSError; treat any such failure as "not a - local path" so the trust decision falls through to the unsloth/ check (the loader's - validate_load_request raises the clear FileNotFoundError for a genuinely missing - local pick).""" + local path" so the trust decision falls through to the org/allowlist checks (the + loader's validate_load_request raises the clear FileNotFoundError for a genuinely + missing local pick).""" try: if Path(repo_id).expanduser().exists(): return True except OSError: pass - 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) @@ -450,6 +471,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. @@ -614,6 +645,7 @@ class DiffusionBackend: base, kwargs.get("hf_token"), kind = kind, + single_file_is_pipeline = bool(fam and fam.single_file_is_pipeline), # The dense transformer-quant path downloads the base repo's # transformer/ shards via from_pretrained(subfolder="transformer") # INSIDE the locked finalize phase, where unload/cancellation cannot @@ -700,6 +732,7 @@ class DiffusionBackend: hf_token: Optional[str], *, kind: str = "gguf", + single_file_is_pipeline: bool = False, include_transformer: bool = False, ) -> tuple[int, list[str]]: """Total download size for the progress bar, plus the base-repo files to @@ -708,7 +741,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() @@ -717,10 +752,19 @@ class DiffusionBackend: try: if kind == "pipeline": info = api.model_info(repo_id, files_metadata = True, token = hf_token) - for s in info.siblings: - if _pipeline_file_downloaded(s.rfilename): - base_files.append(s.rfilename) - total += s.size or 0 + picked = [s for s in info.siblings if _pipeline_file_downloaded(s.rfilename)] + # diffusers prefers safetensors per component: drop a .bin whose + # directory also carries a picked .safetensors weight. + st_dirs = { + s.rfilename.rsplit("/", 1)[0] + for s in picked + if s.rfilename.endswith(".safetensors") + } + for s in picked: + if s.rfilename.endswith(".bin") and s.rfilename.rsplit("/", 1)[0] in st_dirs: + continue + base_files.append(s.rfilename) + total += s.size or 0 return total, base_files # Skip the Hub size lookup for a LOCAL gguf path: model_info(repo_id) would # raise on a filesystem path and (caught below) skip the base-repo lookup too, @@ -729,9 +773,18 @@ 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. + if kind == "single_file" and single_file_is_pipeline: + base_filter = _base_config_file_downloaded + else: + + def base_filter(rfilename: str) -> bool: + return _base_file_downloaded(rfilename, include_transformer = include_transformer) + 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, include_transformer = include_transformer): + 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 @@ -815,6 +868,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. Sanitize the # token here too (direct callers bypass begin_load): a blank string must @@ -941,6 +1001,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). @@ -1453,23 +1523,33 @@ 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 + # Read the dtype from the parameters (not denoiser.dtype): a plain nn.Module + # has no .dtype, and a torch.compile'd / wrapped denoiser can obscure it. Take + # the first FLOATING dtype: a GGUF-quantized transformer's leading params are + # packed uint8 storage, and nn.Module.to() rejects integer dtypes outright. + target_dtype = next( + (p.dtype for p in denoiser.parameters() if p.dtype.is_floating_point), + None, + ) + if target_dtype is None: + return if next(vae.parameters()).dtype != target_dtype: vae.to(dtype = target_dtype) - except (StopIteration, AttributeError, RuntimeError): + except (StopIteration, AttributeError, RuntimeError, TypeError): pass def _apply_loras( @@ -1804,7 +1884,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) + # 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 @@ -2168,17 +2249,57 @@ def _base_file_downloaded(rfilename: str, *, include_transformer: bool = False) 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. Like ``_base_file_downloaded`` but for the ``pipeline`` kind, where the repo supplies its OWN transformer weights, so the ``transformer/`` subfolder is kept. - Top-level docs (README/PDF/images) and ``assets/`` are still skipped so the - progress estimate matches what actually lands on disk. + Top-level docs (README/PDF/images) and ``assets/`` are skipped, and so are + artifacts the torch loader never touches -- ONNX / OpenVINO / Flax exports and + dtype-variant twins (``*.fp16.safetensors``: the loader requests the default + variant) -- so an official repo that ships many formats (e.g. SDXL Base) does + not prefetch tens of GB it will not load. """ if "/" not in rfilename: # top-level: only the pipeline manifest is fetched return rfilename == "model_index.json" - return not rfilename.startswith("assets/") + lower = rfilename.lower() + if lower.startswith(("assets/", "onnx/", "openvino/")): + return False + name = lower.rsplit("/", 1)[1] + if name.startswith(("openvino_", "flax_")): + return False + if name.endswith((".onnx", ".onnx_data", ".pb", ".msgpack", ".h5", ".ckpt")): + return False + if ".fp16." in name or ".bf16." in name or ".non_ema." in name: + return False + return True def _progress( diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 86493a3070..3c56bed4a9 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -37,6 +37,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 @@ -251,6 +263,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 diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 2b258df003..d36e62f3cf 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -247,12 +247,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 = {} @@ -368,6 +375,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) @@ -375,6 +389,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 = {} @@ -918,12 +933,62 @@ 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. diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py new file mode 100644 index 0000000000..6bec0a4100 --- /dev/null +++ b/studio/backend/tests/test_diffusion_sdxl.py @@ -0,0 +1,221 @@ +# 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 + +import pytest + +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") + # 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. + 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. The + # dtype is read from a parameter (denoiser has no .dtype), so use a _FakeVae denoiser. + import torch + + vae = _FakeVae(dtype = torch.float32) + unet = _FakeVae(dtype = torch.bfloat16) + pipe = types.SimpleNamespace(unet = unet, vae = vae) + DiffusionBackend._align_vae_dtype(pipe, "unet") + assert vae.moved_to == torch.bfloat16 + + +def test_align_vae_dtype_transformer_default_unchanged(): + # DiT default: reads pipe.transformer; a pipe with no transformer is a safe no-op. + import torch + + vae = _FakeVae(dtype = torch.float32) + transformer = _FakeVae(dtype = torch.bfloat16) + pipe = types.SimpleNamespace(transformer = transformer, vae = vae) + DiffusionBackend._align_vae_dtype(pipe) + assert vae.moved_to == torch.bfloat16 + # No denoiser attribute -> no-op (does not raise, does not move the VAE). + vae2 = _FakeVae(dtype = torch.float32) + DiffusionBackend._align_vae_dtype(types.SimpleNamespace(vae = vae2), "unet") + assert vae2.moved_to is None + + +def test_align_vae_dtype_skips_gguf_packed_uint8_params(): + # A GGUF-quantized transformer's leading parameters are packed uint8 storage; the + # dtype probe must skip them and use the first FLOATING dtype, or nn.Module.to() + # rejects the integer dtype and an Edit/img2img call 500s (regression: Qwen-Image- + # Edit GGUF). All-integer params (no floating dtype at all) must be a clean no-op. + import torch + + class _GgufDenoiser: + def parameters(self): + yield types.SimpleNamespace(dtype = torch.uint8) # packed GGUF block + yield types.SimpleNamespace(dtype = torch.bfloat16) # compute dtype + + vae = _FakeVae(dtype = torch.float32) + pipe = types.SimpleNamespace(transformer = _GgufDenoiser(), vae = vae) + DiffusionBackend._align_vae_dtype(pipe) + assert vae.moved_to == torch.bfloat16 + + class _AllPacked: + def parameters(self): + yield types.SimpleNamespace(dtype = torch.uint8) + + vae2 = _FakeVae(dtype = torch.float32) + DiffusionBackend._align_vae_dtype(types.SimpleNamespace(transformer = _AllPacked(), vae = vae2)) + 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 + ) + + +def test_pipeline_prefetch_skips_non_torch_artifacts(): + # The SDXL Base repo ships fp16 variants, ONNX, OpenVINO and Flax exports next to + # the default safetensors; from_pretrained (no variant kwarg) loads only the + # default torch weights, so the prefetch filter must skip everything else or a + # catalog load pulls tens of GB of unused artifacts. + from core.inference.diffusion import _pipeline_file_downloaded as keep + + assert keep("model_index.json") + assert keep("unet/diffusion_pytorch_model.safetensors") + assert keep("text_encoder/model.safetensors") + assert keep("scheduler/scheduler_config.json") + assert not keep("sd_xl_base_1.0.safetensors") # top-level single-file twin + assert not keep("unet/diffusion_pytorch_model.fp16.safetensors") + assert not keep("text_encoder/model.onnx") + assert not keep("text_encoder/openvino_model.bin") + 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") diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index d9f9004a27..1104b40e74 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -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 = { "unsloth/Z-Image-Turbo-unsloth-bnb-4bit": { kind: "pipeline" }, @@ -100,6 +101,10 @@ const SAFETENSORS_MODELS: Record = { 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 } {