ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/core/inference/diffusion_families.py
Daniel Han 36df317293 Trim the comments across the diffusion backend 
Comment-only pass over the Python this PR touches: drop what the code already
says, collapse multi-line explanations that still read on one line, and keep
the reasoning that is not recoverable from the code. No code, docstring
semantics or behaviour changes; verified with an AST comparison against the
previous revision, and the backend suite is unchanged (same 37 environment
failures as before: the API integration tests that need a live keyed server,
the flash-attn install hooks, and the GPU memory fields).
2026-07-26 20:31:19 +00:00

676 lines
35 KiB
Python
Raw Blame History

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Pure helpers for diffusion model identification.
No torch/diffusers imports here: everything in this module is a pure function of
its string/path arguments so it can be unit-tested without the heavy runtime.
A diffusion checkpoint published as a single-file GGUF only carries the
transformer weights; the matching VAE / text encoders / scheduler come from a
companion ``diffusers`` base repo. ``DiffusionFamily`` maps a checkpoint to the
diffusers classes and base repo needed to assemble the full pipeline.
"""
from __future__ import annotations
import re
from dataclasses import dataclass, field
from pathlib import Path, PurePosixPath
from typing import Optional
# Runtime->route contract: RuntimeError messages for client-recoverable generate states. The
# /images/generate route matches these EXACTLY for a 409 (vs a 500), so both engines raise them
# verbatim -- named here, not as scattered literals.
DIFFUSION_NOT_LOADED_MSG = "No diffusion model is loaded."
DIFFUSION_CANCELLED_MSG = "Diffusion generation was cancelled."
@dataclass(frozen = True)
class DiffusionFamily:
name: str
pipeline_class: str
transformer_class: str
base_repo: str
# Pipeline kwarg carrying guidance. Most use "guidance_scale"; Qwen-Image's real CFG is
# "true_cfg_scale" (its distilled guidance is off).
cfg_kwarg: str = "guidance_scale"
# The pipe attribute holding the denoiser: ``pipe.transformer`` for DiT families (default),
# ``pipe.unet`` for U-Net families (SDXL).
denoiser_attr: str = "transformer"
# True when a single-file ``.safetensors`` is the WHOLE pipeline (SDXL), so the loader calls
# ``pipeline_class.from_single_file``. DiT families leave this False (transformer-only).
single_file_is_pipeline: bool = False
# True for families needing MULTIPLE denoisers no single file carries (Ideogram 4), so only a
# full ``pipeline`` load is valid; validate_load_request rejects single-file / GGUF up front.
pipeline_only: bool = False
# Optional diffusers pipeline classes for image-conditioned workflows, built around the resident
# modules via ``Pipeline.from_pipe`` (no reload). None = unsupported (UI gates it off).
img2img_pipeline_class: Optional[str] = None
inpaint_pipeline_class: Optional[str] = None
# ControlNet pipeline + model classes: the model loads via from_pretrained and the pipeline via
# ``from_pipe(base, controlnet=model)`` (no reload). None on both = no support.
controlnet_pipeline_class: Optional[str] = None
controlnet_model_class: Optional[str] = None
# True when the inpaint pipeline keeps the canvas size, so it can also drive outpaint. False for
# FLUX.2 (it scales >1MP inputs to ~1MP, shrinking the canvas) -> Inpaint but not Extend.
inpaint_preserves_size: bool = True
# True for instruction-editing families (Qwen-Image-Edit / FLUX Kontext): the OWN pipeline IS
# the edit pipeline (image + instruction, no plain text-to-image), used directly (no from_pipe).
# ``base_repo`` supplies the VAE / text-encoder / processor / scheduler for the GGUF transformer.
edit: bool = False
# True for families whose text-to-image pipeline ALSO accepts reference image(s) (FLUX.2's
# ``image`` arg). Unlike ``edit`` they still do plain text-to-image; unlike img2img the
# conditioning is reference-based (no ``strength``, output size from width/height).
reference: bool = False
# Extra lowercased substrings (besides ``name``) that map a repo id here.
aliases: tuple[str, ...] = field(default_factory = tuple)
# True for families whose activations overflow float16 (-> inf/NaN -> black image); the backend
# promotes a resolved float16 to float32 for these.
fp16_incompatible: bool = False
# False only for a family whose denoiser block doesn't compile cleanly with regional
# torch.compile. Consulted on the GGUF path too; all current families compile.
supports_torch_compile: bool = True
# Optional pre-quantized transformer checkpoints as (scheme, repo_id) pairs: the loader fetches
# already-quantized weights instead of the dense bf16 (lower load VRAM + smaller download).
prequant_repos: tuple[tuple[str, str], ...] = field(default_factory = tuple)
# Hosted checkpoints for NON-DEFAULT bases of the family, as (base_repo, scheme, repo_id)
# triples with base_repo lowercased. One family entry covers several published variants
# (flux.1: schnell/dev/Krea-dev) whose weights differ, so each needs its own baked checkpoint.
# Resolution prefers an exact variant match, then falls back to ``prequant_repos``.
prequant_variant_repos: tuple[tuple[str, str, str], ...] = field(default_factory = tuple)
# Hosted PRE-CAST text-encoder checkpoints as (scheme, component, repo_id) triples. Serves the
# layerwise-fp8 storage scheme only: the cast is deterministic, so the artifact is bit-identical
# to dense-load-then-cast while skipping the multi-GB dense TE download. Empty -> load dense.
te_prequant_repos: tuple[tuple[str, str, str], ...] = field(default_factory = tuple)
# Native (sd.cpp) single-file assets, used only on the no-GPU sd.cpp engine. The transformer
# GGUF is shared with diffusers; sd-cli also needs a single-file VAE + text encoder(s). Each is
# a (repo_id, filename); ``sd_cpp_text_encoders`` carries a trailing SdCppModelFiles field name
# (clip_l / t5xxl / llm / qwen2vl / clip_g) for the sd-cli flag. Empty -> no native mapping.
sd_cpp_vae: Optional[tuple[str, str]] = None
# VAE latent-format override for sd-cli (--vae-format): "flux2" for FLUX.2, None otherwise.
sd_cpp_vae_format: Optional[str] = None
sd_cpp_text_encoders: tuple[tuple[str, str, str], ...] = field(default_factory = tuple)
# Family-specific sd-cli sampler settings so the native output matches the model's supported
# invocation (e.g. Qwen-Image needs euler + flow-shift 3). None leaves sd-cli defaults.
sd_cpp_sampling_method: Optional[str] = None
sd_cpp_flow_shift: Optional[float] = None
# True when Studio can TRAIN a LoRA on this family (a trainer is registered). Opt-in per family;
# the training-start path refuses a non-trainable family up front.
trainable: bool = False
# Recommended base repos to train FROM, most-preferred first (e.g. a QLoRA prequant repo, then
# bf16). Surfaced by the Train UI.
train_base_repos: tuple[str, ...] = field(default_factory = tuple)
# When set, deploying a LoRA trained on this family loads THIS repo instead of the trained-on
# checkpoint (Krea: train on Raw, preview on Turbo). Both sides must be the same precision.
deploy_base_repo: Optional[str] = None
# Keyed by architecture, not per variant: a checkpoint's base repo is read from its HF
# base_model tag at load time, so one entry covers Turbo/full, schnell/dev, etc. Only archs
# whose diffusers transformer supports from_single_file load here.
_FAMILIES: tuple[DiffusionFamily, ...] = (
DiffusionFamily(
name = "flux.1",
pipeline_class = "FluxPipeline",
transformer_class = "FluxTransformer2DModel",
base_repo = "black-forest-labs/FLUX.1-schnell",
# Hosted pre-quantized DiT checkpoints (gate-validated vs same-seed bf16). The loader verifies
# the baked base_model_id against the repo being loaded, so a non-default base safely falls back
# to the dense-quantize path instead of loading schnell weights.
prequant_repos = (
("int8", "unsloth/FLUX.1-schnell-FP8"),
("fp8", "unsloth/FLUX.1-schnell-FP8"),
),
# Gate-validated checkpoints baked from the dev / Krea-dev weights (same arch, different
# weights): without these the default schnell checkpoint is refused and every int8/fp8 load pays
# the dense download + on-the-fly quantise.
prequant_variant_repos = (
("black-forest-labs/flux.1-dev", "int8", "unsloth/FLUX.1-dev-FP8"),
("black-forest-labs/flux.1-dev", "fp8", "unsloth/FLUX.1-dev-FP8"),
("black-forest-labs/flux.1-krea-dev", "int8", "unsloth/FLUX.1-Krea-dev-FP8"),
("black-forest-labs/flux.1-krea-dev", "fp8", "unsloth/FLUX.1-Krea-dev-FP8"),
),
# Pre-cast T5-XXL (9.52 -> 5.90 GB; CLIP-L stays dense). One artifact serves schnell/dev/
# Krea-dev: the T5 shards are byte-identical across all three (verified sha256).
te_prequant_repos = (("fp8", "text_encoder_2", "unsloth/FLUX.1-schnell-FP8"),),
aliases = ("flux1", "flux-1"),
# LoRA training targets FLUX.1-dev via the DiT trainer (QLoRA nf4); the dev repo is gated.
trainable = True,
train_base_repos = ("black-forest-labs/FLUX.1-dev",),
img2img_pipeline_class = "FluxImg2ImgPipeline",
inpaint_pipeline_class = "FluxInpaintPipeline",
controlnet_pipeline_class = "FluxControlNetPipeline",
controlnet_model_class = "FluxControlNetModel",
sd_cpp_vae = ("black-forest-labs/FLUX.1-schnell", "ae.safetensors"),
sd_cpp_text_encoders = (
("comfyanonymous/flux_text_encoders", "clip_l.safetensors", "clip_l"),
("comfyanonymous/flux_text_encoders", "t5xxl_fp16.safetensors", "t5xxl"),
),
),
# FLUX.2-klein is Flux2KleinPipeline (Qwen3 encoder), not the Mistral Flux2Pipeline, so it must
# precede a generic flux match. The Mistral Flux2Pipeline is the flux.2-dev family below.
DiffusionFamily(
name = "flux.2-klein",
pipeline_class = "Flux2KleinPipeline",
transformer_class = "Flux2Transformer2DModel",
base_repo = "black-forest-labs/FLUX.2-klein-4B",
prequant_repos = (
("int8", "unsloth/FLUX.2-klein-4B-FP8"),
("fp8", "unsloth/FLUX.2-klein-4B-FP8"),
),
aliases = ("flux2-klein",),
# LoRA training via the DiT trainer (QLoRA nf4 by default); klein-4B is not gated.
trainable = True,
train_base_repos = ("black-forest-labs/FLUX.2-klein-4B",),
# Flux2KleinPipeline takes reference image(s) via `image`, so it exposes a "reference" workflow
# atop text-to-image. It has an inpaint pipeline (no img2img) -> inpaint + extend.
reference = True,
inpaint_pipeline_class = "Flux2KleinInpaintPipeline",
# FLUX.2 scales >1MP inputs to ~1MP, so outpaint can't grow.
inpaint_preserves_size = False,
# FLUX.2's 32-channel AE needs the latent-format override; the single-file VAE ships in
# Comfy-Org/flux2-dev (klein-4B has only a sharded diffusers VAE). Shares Qwen3-4B with z-image.
sd_cpp_vae = ("Comfy-Org/flux2-dev", "split_files/vae/flux2-vae.safetensors"),
sd_cpp_vae_format = "flux2",
sd_cpp_text_encoders = (
("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"),
),
),
# FLUX.2-dev: full (non-distilled) FLUX.2 on the Mistral Flux2Pipeline, so its own entry. Base
# repo is gated. Text-to-image only (no Flux2 img2img/inpaint in diffusers 0.38). VAE + Mistral
# encoder come from the open Comfy-Org/flux2-dev mirror for sd-cli.
DiffusionFamily(
name = "flux.2-dev",
pipeline_class = "Flux2Pipeline",
transformer_class = "Flux2Transformer2DModel",
base_repo = "black-forest-labs/FLUX.2-dev",
prequant_repos = (
("int8", "unsloth/FLUX.2-dev-FP8"),
("fp8", "unsloth/FLUX.2-dev-FP8"),
),
# Pre-cast Mistral-Small-24B conditioner (bf16 ~48 GB dense, ~24.7 GB pre-cast).
te_prequant_repos = (("fp8", "text_encoder", "unsloth/FLUX.2-dev-FP8"),),
aliases = ("flux2-dev", "flux2dev"),
# LoRA training via the DiT trainer (QLoRA nf4 by default); the base repo is gated, so training
# requires an HF token with the FLUX.2-dev license accepted.
trainable = True,
train_base_repos = ("black-forest-labs/FLUX.2-dev",),
sd_cpp_vae = ("Comfy-Org/flux2-dev", "split_files/vae/flux2-vae.safetensors"),
sd_cpp_vae_format = "flux2",
sd_cpp_text_encoders = (
(
"Comfy-Org/flux2-dev",
"split_files/text_encoders/mistral_3_small_flux2_bf16.safetensors",
"llm",
),
),
),
DiffusionFamily(
# FLUX instruction editing: FluxKontextPipeline takes an image + edit instruction; the GGUF
# transformer is standard FluxTransformer2DModel. Specific aliases first so detect_family prefers
# this over "flux.1" and un-rejects the "kontext" keyword.
name = "flux.1-kontext",
pipeline_class = "FluxKontextPipeline",
transformer_class = "FluxTransformer2DModel",
base_repo = "black-forest-labs/FLUX.1-Kontext-dev",
aliases = ("flux.1-kontext-dev", "flux1-kontext", "flux-kontext", "kontext"),
edit = True,
),
DiffusionFamily(
# Qwen instruction editing: the 2511 checkpoint ships as QwenImageEditPlusPipeline
# (multi-image); the GGUF transformer is standard QwenImageTransformer2DModel. Specific aliases
# first so detect_family prefers this over "qwen-image".
name = "qwen-image-edit",
pipeline_class = "QwenImageEditPlusPipeline",
transformer_class = "QwenImageTransformer2DModel",
base_repo = "Qwen/Qwen-Image-Edit-2511",
cfg_kwarg = "true_cfg_scale",
aliases = (
"qwen-image-edit-2511",
"qwen-image-edit-2509",
"qwen-image-edit",
"qwen_image_edit",
"qwenimageedit",
),
edit = True,
),
DiffusionFamily(
name = "qwen-image",
pipeline_class = "QwenImagePipeline",
transformer_class = "QwenImageTransformer2DModel",
base_repo = "Qwen/Qwen-Image",
# int8 only: fp8 is family-denied (_FAMILY_SCHEME_DENY) so a repo entry would be dead.
prequant_repos = (("int8", "unsloth/Qwen-Image-FP8"),),
# Pre-cast Qwen2.5-VL-7B (bf16 ~16.6 GB dense, ~8.8 GB pre-cast). The DiT fp8 denial is a
# transformer-scheme rule; the layerwise TE cast is unaffected.
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),),
cfg_kwarg = "true_cfg_scale",
aliases = ("qwen_image", "qwenimage"),
# LoRA training via the DiT trainer, defaulting to the prequant nf4 repo (QLoRA).
trainable = True,
train_base_repos = ("unsloth/Qwen-Image-2512-unsloth-bnb-4bit", "Qwen/Qwen-Image"),
img2img_pipeline_class = "QwenImageImg2ImgPipeline",
inpaint_pipeline_class = "QwenImageInpaintPipeline",
controlnet_pipeline_class = "QwenImageControlNetPipeline",
controlnet_model_class = "QwenImageControlNetModel",
sd_cpp_vae = ("Comfy-Org/Qwen-Image_ComfyUI", "split_files/vae/qwen_image_vae.safetensors"),
# Qwen2.5-VL as a Q4_K_M GGUF keeps the CPU RAM win (bf16 encoder is ~15 GB). sd-cli's
# --qwen2vl aliases --llm.
sd_cpp_text_encoders = (
(
"unsloth/Qwen2.5-VL-7B-Instruct-GGUF",
"Qwen2.5-VL-7B-Instruct-Q4_K_M.gguf",
"qwen2vl",
),
),
# Qwen-Image's supported sd.cpp invocation (docs/qwen_image.md).
sd_cpp_sampling_method = "euler",
sd_cpp_flow_shift = 3.0,
),
DiffusionFamily(
name = "z-image",
pipeline_class = "ZImagePipeline",
transformer_class = "ZImageTransformer2DModel",
base_repo = "Tongyi-MAI/Z-Image-Turbo",
prequant_repos = (
("int8", "unsloth/Z-Image-Turbo-FP8"),
("fp8", "unsloth/Z-Image-Turbo-FP8"),
),
# Pre-cast Qwen3-4B TE (8.04 -> 4.41 GB). NOT shared with flux.2-klein-4B: klein's TE retrained
# layer 35's MLP (up/down_proj maxdiff 0.86 vs this checkpoint).
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Z-Image-Turbo-FP8"),),
aliases = ("zimage", "z_image"),
# LoRA training via the DiT trainer (bf16); defaults to the prequant nf4 repo for QLoRA.
trainable = True,
train_base_repos = ("unsloth/Z-Image-Turbo-unsloth-bnb-4bit", "Tongyi-MAI/Z-Image-Turbo"),
img2img_pipeline_class = "ZImageImg2ImgPipeline",
inpaint_pipeline_class = "ZImageInpaintPipeline",
# Z-Image's MLP down-projections peak near 9e5, which overflows float16.
fp16_incompatible = True,
sd_cpp_vae = ("Comfy-Org/z_image_turbo", "split_files/vae/ae.safetensors"),
sd_cpp_text_encoders = (
("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"),
),
),
# Krea 2 (diffusers >= 0.39): a ~12B single-stream DiT with a Qwen3-VL-4B encoder and the
# Qwen-Image VAE. Loaded per-component (diffusion_krea2.py) because the repo ships
# transformers-5.x configs. No GGUF/sd.cpp mapping yet.
DiffusionFamily(
name = "krea-2",
pipeline_class = "Krea2Pipeline",
transformer_class = "Krea2Transformer2DModel",
base_repo = "krea/Krea-2-Turbo",
prequant_repos = (
("int8", "unsloth/Krea-2-Turbo-FP8"),
("fp8", "unsloth/Krea-2-Turbo-FP8"),
),
# Pre-cast Qwen3-VL-4B TE (8.88 -> 4.83 GB); handed into load_krea2_pipeline directly
# (constructor assembly never sees pipe_kwargs).
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Krea-2-Turbo-FP8"),),
aliases = ("krea2",),
# LoRA training via the DiT trainer (no prequant repo yet, so nf4 quantizes on the fly). Krea's
# guidance: train on the undistilled Raw, run adapters on Turbo, so Raw is the default training
# base and Turbo the inference/base repo.
trainable = True,
train_base_repos = ("krea/Krea-2-Raw", "krea/Krea-2-Turbo"),
# Adapters trained on Raw run on Turbo; deploy previews them there (same bf16 precision).
deploy_base_repo = "krea/Krea-2-Turbo",
# Exported bf16-only; fp16 unvalidated upstream, so keep the fp16 fallback off like z-image.
fp16_incompatible = True,
),
# Lumina Image 2.0: a 2.6B single-stream DiT with a Gemma2-2B encoder and a standard 16-channel
# AutoencoderKL, all transformers-4.x-compatible, so the generic from_pretrained path loads it.
# No GGUF/sd.cpp mapping upstream. NOT aliased to bare "lumina": Lumina-Next checkpoints are a
# different arch and must stay unknown rather than crash mid-load.
DiffusionFamily(
name = "lumina-2",
pipeline_class = "Lumina2Pipeline",
transformer_class = "Lumina2Transformer2DModel",
base_repo = "Alpha-VLLM/Lumina-Image-2.0",
# Gate-validated hosted checkpoints (28/28 pairs each; LPIPS mean 0.146 int8 / 0.116 fp8).
prequant_repos = (
("int8", "unsloth/Lumina-Image-2.0-FP8"),
("fp8", "unsloth/Lumina-Image-2.0-FP8"),
),
# Pre-cast Gemma2-2B TE. The Hub stores it fp32 (10.46 GB), so the 3.20 GB artifact is a 3.3x
# download cut even though the model is small.
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Lumina-Image-2.0-FP8"),),
aliases = ("lumina-image-2.0", "lumina-image-2", "lumina2"),
# Published and validated bf16-only upstream; keep the fp16 fallback off like z-image.
fp16_incompatible = True,
),
# HunyuanImage 2.1 (diffusers >= 0.39): a 17B dual-stream DiT with a Qwen2.5-VL text encoder, a
# ByT5 glyph encoder, and the 32x-compression HunyuanImage VAE. The community mirror also ships
# guider/ocr_guider components 0.39 loads natively, so the generic from_pretrained path covers
# the stack. 2K-native; CFG runs inside the repo's guider at its baked scale and the call's own
# knob is distilled_guidance_scale. Distinct from the excluded HunyuanImage-3.0.
DiffusionFamily(
name = "hunyuanimage-2.1",
# Hosted checkpoints, verified bit-identical to on-the-fly quantize (the guider pipeline is not
# run-to-run deterministic, so same-seed LPIPS vs bf16 blends trajectory divergence with harness
# noise; per-case hard checks pass and the drift is compositional, reviewed visually).
prequant_repos = (
("int8", "unsloth/HunyuanImage-2.1-FP8"),
("fp8", "unsloth/HunyuanImage-2.1-FP8"),
),
# The Qwen2.5-VL TE is byte-identical to Qwen-Image's (verified sha256), so the family reuses
# that artifact: zero new hosting, 16.58 -> 8.84 GB download. ByT5 stays dense.
te_prequant_repos = (("fp8", "text_encoder", "unsloth/Qwen-Image-FP8"),),
pipeline_class = "HunyuanImagePipeline",
transformer_class = "HunyuanImageTransformer2DModel",
base_repo = "hunyuanvideo-community/HunyuanImage-2.1-Diffusers",
cfg_kwarg = "distilled_guidance_scale",
aliases = ("hunyuanimage-2.1-diffusers", "hunyuanimage2.1"),
# Exported bf16-only; keep the fp16 fallback off like z-image / krea-2.
fp16_incompatible = True,
),
# HiDream-I1: a 17B MoE DiT (16 double + 32 single layers, 4 routed experts) with FOUR text
# encoders. The repos ship CLIP-L/CLIP-G/T5-XXL but NOT the Llama-3.1-8B text_encoder_4 their
# model_index names: the loader assembles it from the open unsloth mirror. Full / Dev / Fast
# share the arch, so one family covers all three. A GGUF path would need the same TE4 assembly
# for tiny demand, so none is wired yet.
DiffusionFamily(
name = "hidream-i1",
# Hosted checkpoints: 28/28 per-case gate pairs per scheme (LPIPS suite means 0.291 int8 /
# 0.278 fp8, the 50-step trajectory band); int8 verified bit-identical to on-the-fly quantize.
prequant_repos = (
("int8", "unsloth/HiDream-I1-Full-FP8"),
("fp8", "unsloth/HiDream-I1-Full-FP8"),
),
# Pre-cast Llama-3.1-8B TE4 (16.1 GB bf16 -> 8.1 GB). The generic TE pass only covers
# text_encoder.._3, so TE4 engages via hidream_te4_kwargs.
te_prequant_repos = (("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"),),
pipeline_class = "HiDreamImagePipeline",
transformer_class = "HiDreamImageTransformer2DModel",
base_repo = "HiDream-ai/HiDream-I1-Full",
aliases = ("hidream", "hidream-i1-full", "hidream-i1-dev", "hidream-i1-fast"),
# Exported bf16-only; keep the fp16 fallback off like the other modern DiTs.
fp16_incompatible = True,
),
# Ideogram 4 (diffusers >= 0.39): a 34-layer DiT PAIR (conditional + unconditional_transformer
# for dual-branch CFG, both ~9B, so memory planning counts two DiTs) with a Qwen3-VL encoder.
# No bf16 checkpoint: ideogram-4-fp8 (raw float8, upcast on load) is the family base; the -nf4
# repos carry bnb-4bit configs. All gated, no GGUF/sd.cpp mapping. CFG quirk: the pipeline takes
# guidance_scale OR a per-step guidance_schedule (see the loader's IDEOGRAM4 branch).
DiffusionFamily(
name = "ideogram-4",
pipeline_class = "Ideogram4Pipeline",
transformer_class = "Ideogram4Transformer2DModel",
base_repo = "ideogram-ai/ideogram-4-fp8",
aliases = ("ideogram4", "ideogram-v4", "ideogram"),
# Two DiTs assembled per-component, so no transformer-only single-file / GGUF load.
pipeline_only = True,
),
# SDXL is the one U-Net family: the denoiser is ``pipe.unet`` and a single-file ``.safetensors``
# is the WHOLE pipeline, so it sets ``denoiser_attr="unet"`` + ``single_file_is_pipeline=True``.
# img2img / inpaint / ControlNet are the standard SDXL pipelines via from_pipe. No GGUF path.
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",
# SDXL uses the U-Net LoRA trainer.
trainable = True,
train_base_repos = (
"stabilityai/stable-diffusion-xl-base-1.0",
"stabilityai/sdxl-turbo",
),
),
)
def trainable_family_names() -> tuple[str, ...]:
"""Names of families Studio can train a LoRA on, in registry order."""
return tuple(fam.name for fam in _FAMILIES if fam.trainable)
# The family whose CFG uses a guidance_scale/guidance_schedule pair (the loader special-cases
# the call). Named here so the two modules can't drift.
IDEOGRAM4_FAMILY_NAME = "ideogram-4"
# The family whose generate call carries the card's CFG-truncation ratio (the loader
# special-cases the call). Named here so the two modules can't drift.
LUMINA2_FAMILY_NAME = "lumina-2"
# Models Studio deliberately does NOT support, reason surfaced verbatim in the load error.
# Keyed by a lowercase repo-id substring. The bar is a diffusers pipeline: HunyuanImage-3.0 is
# an 80B MoE needing AutoModelForCausalLM + trust_remote_code (RCE out of the question).
_EXCLUDED_MODELS: tuple[tuple[str, str], ...] = (
(
# "-3" scoped so a future HunyuanImage 2.x with a diffusers pipeline falls through normally.
"hunyuanimage-3",
"HunyuanImage-3.0 has no diffusers pipeline (it is an 80B autoregressive MoE "
"that requires trust_remote_code), so Studio does not support it.",
),
)
def excluded_model_reason(repo_id: str) -> Optional[str]:
"""The stated reason ``repo_id`` is unsupported, or None when it is simply unknown."""
needle = (repo_id or "").lower()
for token, reason in _EXCLUDED_MODELS:
if _token_in_needle(token, needle):
return reason
return None
# Editing / inpaint checkpoints share an arch keyword but need a different pipeline + input
# image. "layered" rejects Qwen-Image-Layered, whose transformer expects an extra
# addition_t_cond input the standard pipeline never supplies. Fails the load fast.
_EDIT_KEYWORDS = ("edit", "kontext", "inpaint", "layered")
def _token_in_needle(token: str, needle: str) -> bool:
"""True when ``token`` appears in ``needle`` as a whole segment (delimited by ``- _ . / \\`` or
a boundary), not a raw substring, so 'qwen-image-edit' matches '...-2511' but 'kontext' doesn't
match 'kontextual'."""
return re.search(r"(?:^|[-_./\\])" + re.escape(token) + r"(?:$|[-_./\\])", needle) is not None
def _best_family_match(needle: str) -> Optional[DiffusionFamily]:
"""The family whose name/alias is the LONGEST whole-segment token of ``needle`` (longest = most
specific, so '...qwen-image-edit-2511...' matches 'qwen-image-edit', not 'qwen-image')."""
best: Optional[tuple[DiffusionFamily, int]] = None
for fam in _FAMILIES:
for token in (fam.name, *fam.aliases):
if _token_in_needle(token, needle) and (best is None or len(token) > best[1]):
best = (fam, len(token))
return best[0] if best else None
def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[DiffusionFamily]:
"""Resolve a ``DiffusionFamily`` from a repo id, or an explicit override.
``override`` matches a family ``name``/alias exactly; otherwise the most-specific family whose
name/alias is a substring of the repo id wins. Supported editing families match here;
unsupported editing/inpaint/layered checkpoints sharing only an arch keyword are rejected (None).
"""
if override:
key = override.strip().lower()
for fam in _FAMILIES:
if key == fam.name or key in fam.aliases:
return fam
return None
needle = repo_id.lower()
match = _best_family_match(needle)
if match is not None:
# Don't let a generic family (qwen-image) swallow a variant it can't run (qwen-image-LAYERED):
# if the id carries a reject keyword the matched family doesn't declare, reject. Scoped to the
# LAST path component so a parent folder named `edit` doesn't reject a valid file.
basename = re.split(r"[/\\]+", needle)[-1]
matched_tokens = (match.name, *match.aliases)
if any(
_token_in_needle(kw, basename) and not any(kw in tok for tok in matched_tokens)
for kw in _EDIT_KEYWORDS
):
return None
return match
return None
def supported_family_names() -> tuple[str, ...]:
"""Family names accepted as ``family_override`` and shown in the unknown-model error (registry
order)."""
return tuple(fam.name for fam in _FAMILIES)
def detect_family_for_pick(
repo_id: str,
gguf_filename: Optional[str] = None,
override: Optional[str] = None,
) -> Optional[DiffusionFamily]:
"""``detect_family``, falling back to the combined path/filename for a local ``.gguf`` pick
where the family keyword lives only in the filename. Only a fallback, so remote picks and
overrides behave exactly as ``detect_family``. Shared by both engines."""
fam = detect_family(repo_id, override)
if fam is None and gguf_filename and not override:
fam = detect_family(f"{repo_id}/{gguf_filename}", override)
return fam
def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
"""The companion diffusers repo: caller-supplied if given, else the family fallback."""
base = (base_repo or "").strip()
return base or fam.base_repo
# Default (steps, guidance) per model for callers that can't pass them (the OpenAI
# /v1/images/generations endpoint has no such knobs). Matched by substring, most specific first;
# same values as the UI's MODEL_DEFAULTS table (images-page.tsx), keep in sync.
_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
("z-image-turbo", 9, 0.0),
# FLUX.1 Krea dev is a FLUX.1-dev finetune (flux.1 family), NOT a Krea-2: its card runs 28 steps
# at guidance 4.5. Must precede the generic "krea" key, which would hand it Turbo's recipe.
("flux.1-krea", 28, 4.5),
# Krea 2 Raw (undistilled): 52 steps / guidance 3.5. Must precede the generic "krea" key.
("krea-2-raw", 52, 3.5),
# Krea 2 Turbo (distilled): 8 steps, no CFG. "krea" then covers Turbo and other krea ids but Raw.
("krea", 8, 0.0),
("flux.1-schnell", 4, 0.0),
("kontext", 28, 2.5), # editing: before the generic flux.1
("flux.1", 28, 3.5),
("flux.2-klein", 4, 0.0),
("flux.2-dev", 28, 4.0), # full (non-distilled)
("qwen-image", 20, 4.0),
("z-image", 20, 4.0),
# Lumina Image 2.0 model-card: 50 steps, guidance 4 (plus cfg_trunc_ratio 0.25, which the loader
# passes itself; see LUMINA2_FAMILY_NAME).
("lumina", 50, 4.0),
# HunyuanImage 2.1 model-card: 50 steps; the guidance value feeds distilled_guidance_scale
# (default 3.25), while real CFG runs inside the repo guiders.
("hunyuanimage", 50, 3.25),
# HiDream-I1 upstream inference.py: Full 50 steps / guidance 5; the distilled Dev (28) and Fast
# (16) run guidance-free. Specific keys precede the generic "hidream".
("hidream-i1-dev", 28, 0.0),
("hidream-i1-fast", 16, 0.0),
("hidream", 50, 5.0),
# Ideogram 4 model-card: 48 steps, guidance 7 (its schedule tapers the last 3 steps to 3.0; the
# loader keeps that taper when the request matches these defaults exactly).
("ideogram", 48, 7.0),
# SDXL: Turbo distilled; base wants ~30 steps + CFG ~7. "sdxl-turbo" precedes "sdxl".
("sdxl-turbo", 3, 0.0),
("stable-diffusion-xl", 30, 7.0),
("sdxl", 30, 7.0),
)
# Unrecognised model: distilled few-step / no-CFG shape, matching the UI fallback.
_GENERATION_DEFAULT_FALLBACK = (9, 0.0)
def default_generation_params(*identifiers: Optional[str]) -> tuple[int, float]:
"""Default ``(steps, guidance)`` for a loaded model. The first identifier naming a known model
wins (repo id, then resolved base repo), so a local-path load still resolves via its base repo.
Keys matched as substrings, most specific first."""
for identifier in identifiers:
needle = (identifier or "").lower()
for key, steps, guidance in _GENERATION_DEFAULTS:
if key in needle:
return steps, guidance
return _GENERATION_DEFAULT_FALLBACK
def family_prequant_repo(
fam: DiffusionFamily,
scheme: str,
base_repo: Optional[str] = None,
) -> Optional[str]:
"""The hosted pre-quantized transformer repo for ``scheme`` in this family, or None.
``base_repo`` (when known) selects a variant-specific checkpoint first: a checkpoint is
baked from ONE base's weights and the loader refuses it for any other base, so a variant
without its own entry still returns the family default (harmless: the base_model_id
validation then falls back to dense-quantise, exactly as before this table existed)."""
base = (base_repo or "").strip().lower()
if base:
for entry_base, entry_scheme, repo_id in fam.prequant_variant_repos:
if entry_base == base and entry_scheme == scheme:
return repo_id
for entry_scheme, repo_id in fam.prequant_repos:
if entry_scheme == scheme:
return repo_id
return None
def family_sd_cpp_supported(fam: DiffusionFamily) -> bool:
"""True when the family has the single-file VAE + text-encoder mapping sd.cpp needs; without it
the no-GPU route falls back to diffusers."""
return bool(fam.sd_cpp_vae and fam.sd_cpp_text_encoders)
# FLUX.2-klein's 9B transformer pairs with Qwen3-8B, the 4B with the family-default Qwen3-4B
# (a mismatched encoder fails deep in sd-cli), so the encoder is picked per variant below.
_FLUX2_KLEIN_9B_SD_CPP_TEXT_ENCODERS = (
(
"Comfy-Org/vae-text-encorder-for-flux-klein-9b",
"split_files/text_encoders/qwen_3_8b.safetensors",
"llm",
),
)
def sd_cpp_text_encoders_for(
fam: DiffusionFamily,
repo_id: Optional[str] = None,
gguf_filename: Optional[str] = None,
) -> tuple[tuple[str, str, str], ...]:
"""The sd.cpp text encoders for a specific load.
FLUX.2-klein picks by variant (9B needs Qwen3-8B, 4B the family default) keyed on the load
identity (repo id + GGUF filename); every other family returns its static table."""
if fam.name == "flux.2-klein":
identity = f"{repo_id or ''}/{gguf_filename or ''}".lower()
if "klein-9b" in identity or "klein_9b" in identity:
return _FLUX2_KLEIN_9B_SD_CPP_TEXT_ENCODERS
return fam.sd_cpp_text_encoders
def resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path:
"""Resolve ``gguf_filename`` (user-supplied) to a file under ``repo_root``, rejecting absolute
paths and ``..`` escapes."""
if (
Path(gguf_filename).is_absolute()
or PurePosixPath(gguf_filename).is_absolute()
or gguf_filename.startswith(("/", "\\"))
or "\\" in gguf_filename
):
raise ValueError("gguf_filename must be a relative path inside the repo.")
rel = PurePosixPath(gguf_filename)
if any(part in ("", ".", "..") for part in rel.parts):
raise ValueError("gguf_filename must not contain '', '.', or '..' segments.")
# Resolve symlinks before the containment check (the lexical guards miss a symlink escape).
repo_real = repo_root.resolve()
child = repo_root.joinpath(*rel.parts).resolve()
if child != repo_real and repo_real not in child.parents:
raise ValueError("gguf_filename must resolve to a file inside the repo.")
if not child.is_file():
raise FileNotFoundError(f"'{gguf_filename}' is not a file under {repo_root}.")
return child
Reference in a new issue View git blame Copy permalink