diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index a396667bce..816e3a0e01 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index a8e5f7e683..47ab01401e 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -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 diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 88639b3721..0b0e1e11cb 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -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. diff --git a/studio/backend/tests/test_diffusion_sdxl.py b/studio/backend/tests/test_diffusion_sdxl.py new file mode 100644 index 0000000000..a0b3044952 --- /dev/null +++ b/studio/backend/tests/test_diffusion_sdxl.py @@ -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 + ) diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index f58c1a565e..0942aa184d 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 } {