internal/autogen/encoderpool.go
Rama main · réplica de lectura del repositorio original.
package autogen
import (
"encoding/binary"
"encoding/json"
"errors"
"io"
"io/fs"
"os"
"path/filepath"
"regexp"
"sort"
"strings"
"sync"
"time"
"github.com/quartermaster-labs/quartermaster/internal/config"
)
// Diffusion component auto-discovery.
//
// A bare --diffusion-model GGUF carries no VAE and no text encoder, so sd-server
// needs those wired as separate files. They used to come exclusively from
// settings.encoders: a hand-written path per role, per machine, which is the
// opposite of this project's "no hand tuning" goal and silently broke whenever
// the models tree moved.
//
// Every one of those files announces what it is, structurally, in its header, so
// the pool is discovered instead of declared:
//
// - a VAE has an encoder/decoder conv stack, and the thing that makes two VAEs
// incompatible (latent channel count) is literally decoder.conv_in's input
// width: 4 = SD/SDXL, 16 = flux.1 "ae", 32 = flux.2. Filenames cannot do
// this job: three unrelated files on the dev box are all "ae.safetensors",
// and two of those turned out to be byte-identical copies.
// - a CLIP has text_model.embeddings.token_embedding, whose width separates
// CLIP-L (768) from CLIP-G (1280).
// - a text-encoder LLM reports its own hidden width (gguf embedding_length, or
// safetensors model.embed_tokens), and the DiT that needs it states the width
// it expects in its caption-projection tensor (Metadata.CondHidden). Matching
// those two numbers picks the right encoder with no name table: verified
// against LongCat and Qwen-Image-Edit (3584 = Qwen2.5-VL-7B), Z-Image and
// Krea2 (2560 = Qwen3-4B / Qwen3-VL-4B), ERNIE (3072 = Ministral-3B) and
// flux (4096 = T5-XXL).
//
// settings.encoders still wins wherever it is set, so an existing config keeps
// working and an odd pick stays overridable.
// ComponentRole is what a discovered component file is FOR. One file has exactly
// one role; the family/width fields below say which models can use it.
type ComponentRole string
const (
RoleNone ComponentRole = ""
RoleVae ComponentRole = "vae"
// RoleAudioVae decodes an AUDIO latent, not a picture. Its own role rather
// than a VAE family because it is never a substitute for one: a video model
// that emits a soundtrack loads BOTH (--vae and --audio-vae).
RoleAudioVae ComponentRole = "audio-vae"
RoleClip ComponentRole = "clip"
RoleT5 ComponentRole = "t5"
RoleLlm ComponentRole = "llm"
// RoleRecurrentLlm is an LLM with recurrent (GatedDeltaNet/SSM) layers. It
// is its own role, not RoleLlm with a flag, so that no width-matched pick
// (Llm, LlmHinted, fillEncoderSet) can ever reach one: most conditioners
// are a plain transformer stack and die loading it (see isRecurrentLlm).
// Only a family whose sd.cpp runner ships a recurrent conditioner asks for
// it by name, through RecurrentLlm (MOSS-Image on Qwen3.5-9B).
RoleRecurrentLlm ComponentRole = "recurrent-llm"
// RoleTokenizer is a standalone HF tokenizer.json, which sd.cpp takes as
// --tokenizer for a family whose conditioner will not use the vocab baked
// into the encoder gguf (MOSS-Image: "requires the official Qwen3.5
// tokenizer.json"). Recorded by FILENAME, never parsed: a 20 MB JSON per
// HF folder per scan is not worth the read, and the picker is hint-scoped.
RoleTokenizer ComponentRole = "tokenizer"
// roleProj is a vision projector: paired to an llm, never picked alone.
roleProj ComponentRole = "mmproj"
)
// VAE families, keyed by latent channel count (decoder.conv_in input width).
// The count IS the compatibility class: a 16-channel model cannot decode a
// 4-channel latent, so a family mismatch is a broken image, not a quality knob.
const (
VaeFamilySD = "sd" // 4 ch: SD1/SD2/SDXL
VaeFamilyFlux = "flux" // 16 ch: flux.1 "ae", and every model that reuses it
VaeFamilyFlux2 = "flux2" // 32 ch: flux.2 / ERNIE (AutoencoderKLFlux2)
VaeFamilyWan3D = "wan3d" // Wan 2.1-derived 3D causal VAE (conv1 is 5-dim)
// VaeFamilyVideo3D is a transformer 3D video autoencoder (MiniMax-H3): the
// encoder patchifies RGB with a conv3d but the DECODER is a transformer, so
// it has neither decoder.conv_in nor Wan's bare conv1.
VaeFamilyVideo3D = "video3d"
// VaeFamilyLtx is the LTX-2.x pair of autoencoders. Both halves carry this
// family: they ship as two files (video + audio) that are only ever loaded
// TOGETHER, and giving them one family is what keeps an LTX model from being
// handed MiniMax-H3's soundtrack decoder, which is a different role-mate
// entirely. The video half is a 128-channel 3D conv VAE whose convolutions
// are wrapped one level deeper than anyone else's (decoder.conv_in.CONV.
// weight), which is also why neither arm above sees it.
VaeFamilyLtx = "ltx"
)
// ComponentFile is one classified file on disk.
type ComponentFile struct {
Path string
Role ComponentRole
Family string // VAE family, or the gguf arch for an llm/t5
Width int64 // hidden width: clip 768/1280, llm/t5 embedding length
SizeGB float64
Mmproj string // llm only: the vision projector paired to it (see pairProjectors)
Vision bool // llm only: carries (or is paired to) a vision tower
}
// EncoderPool is the discovered set of component files.
type EncoderPool struct {
Files []ComponentFile
}
// safetensorsHeaderCap bounds the JSON header read. Real headers are tens of KB
// (a few hundred for a full checkpoint); anything past this is a corrupt or
// hostile length field, not a model.
const safetensorsHeaderCap = 64 << 20
type stTensor struct {
DType string `json:"dtype"`
Shape []int64 `json:"shape"`
}
// readSafetensorsHeader parses the tensor table at the head of a .safetensors
// file: an 8-byte little-endian JSON length, then that many bytes of JSON. Only
// the header is read, so the cost is independent of the file size.
func readSafetensorsHeader(path string) (map[string]stTensor, error) {
f, err := os.Open(path)
if err != nil {
return nil, err
}
defer f.Close()
var n uint64
if err := binary.Read(f, binary.LittleEndian, &n); err != nil {
return nil, err
}
if n == 0 || n > safetensorsHeaderCap {
return nil, io.ErrUnexpectedEOF
}
buf := make([]byte, n)
if _, err := io.ReadFull(f, buf); err != nil {
return nil, err
}
var raw map[string]json.RawMessage
if err := json.Unmarshal(buf, &raw); err != nil {
return nil, err
}
out := make(map[string]stTensor, len(raw))
for k, v := range raw {
if k == "__metadata__" {
continue
}
var t stTensor
if json.Unmarshal(v, &t) == nil {
out[k] = t
}
}
return out, nil
}
// classifySafetensors identifies a .safetensors component from its tensor table.
// Returns RoleNone for anything that is not a component (a DiT, a LoRA, a full
// checkpoint) so the caller can ignore it.
func classifySafetensors(h map[string]stTensor, sizeGB float64, path string) ComponentFile {
c := ComponentFile{Path: path, SizeGB: sizeGB}
dim := func(name string, i int) int64 {
t, ok := h[name]
if !ok || i >= len(t.Shape) {
return 0
}
return t.Shape[i]
}
switch {
// 2D AE (SD / flux lineage). decoder.conv_in.weight is [out, latent, 3, 3],
// so shape[1] is the latent channel count = the compatibility family.
case dim("decoder.conv_in.weight", 1) > 0:
c.Role = RoleVae
switch dim("decoder.conv_in.weight", 1) {
case 4:
c.Family = VaeFamilySD
case 16:
c.Family = VaeFamilyFlux
case 32:
// The FLUX.2 "small decoder" (MOSS-Image ships it as its vae/) has
// the same 32-ch latent but a 96-wide decoder (conv_in out 384,
// not 512). sd.cpp builds the FLUX.2 autoencoder at a fixed
// width, so it fails to load; left in the family, the unhinted
// path sort would hand it to every FLUX.2 / ERNIE / MOSS model.
if dim("decoder.conv_in.weight", 0) != 512 {
return ComponentFile{}
}
c.Family = VaeFamilyFlux2
default:
return ComponentFile{}
}
// Wan 2.1-derived 3D causal VAE: conv1.weight is 5-dim. Wan 2.1 and the
// Qwen-Image 20B line have the identical shape table, so those two need a
// path hint to pick between them (see EncoderPool.Vae).
//
// Width is the LATENT channel count, read off the decoder's input projection
// ([out, latent, t, h, w]). It is 16 for Wan 2.1 and Qwen-Image 20B, and 64
// for Qwen-Image 2.1's RGBA autoencoder, which is the first member of this
// family whose latent is a different size. A 64-channel VAE handed to a
// 16-channel DiT decodes to noise rather than failing, so the width is what
// keeps the filename sort in pickByHint from deciding it.
case len(h["conv1.weight"].Shape) == 5:
c.Role = RoleVae
c.Family = VaeFamilyWan3D
c.Width = dim("decoder.conv1.weight", 1)
// Transformer 3D video VAE (MiniMax-H3). Neither arm above sees it: the
// decoder is a transformer stack (decoder.x_embedder), so there is no
// decoder.conv_in and no bare conv1. What remains is the encoder's conv3d
// stem, encoder.conv_in.weight = [out, 3, t, h, w]. Deliberately placed
// AFTER the wan3d arm so a Wan VAE that also carries this tensor keeps its
// own family.
case len(h["encoder.conv_in.weight"].Shape) == 5:
c.Role = RoleVae
c.Family = VaeFamilyVideo3D
// LTX-2.x video VAE. Its conv modules are wrapped (LTXVideoCausalConv3d),
// so the tensor is decoder.conv_in.CONV.weight = [out, latent, t, h, w] and
// neither the 2D arm nor the two 3D arms above match it. shape[1] is the
// latent channel count (128 for LTX-2.5), kept as the width so a future
// second latent size is a family mismatch rather than a silent bad decode.
case len(h["decoder.conv_in.conv.weight"].Shape) == 5:
c.Role = RoleVae
c.Family = VaeFamilyLtx
c.Width = dim("decoder.conv_in.conv.weight", 1)
// LTX-2.x audio VAE. One file, two stacks: the VAE proper under an
// audio_vae. prefix plus the vocoder that turns its mel output into samples.
// The decoder input projection is a 2D conv over the mel spectrogram
// ([width, latent, 3, 3]), NOT the 1D stack H3 uses, so it needs its own arm
// and carries the LTX family so the two audio decoders never cross-wire.
case len(h["audio_vae.decoder.conv_in.conv.weight"].Shape) == 4:
c.Role = RoleAudioVae
c.Family = VaeFamilyLtx
c.Width = dim("audio_vae.decoder.conv_in.conv.weight", 1)
// Audio VAE (MiniMax-H3): a 1D conv stack, so the decoder input projection
// is [width, latent, 1]. Width here is the LATENT channel count, which is
// what has to match the DiT's audio_patch_proj.
case len(h["dec_in_proj.weight"].Shape) == 3:
c.Role = RoleAudioVae
c.Family = VaeFamilyVideo3D
c.Width = dim("dec_in_proj.weight", 1)
// CLIP text tower. token_embedding is [vocab, width]; 768 = CLIP-L, 1280 = G.
case dim("text_model.embeddings.token_embedding.weight", 1) > 0:
c.Role = RoleClip
c.Width = dim("text_model.embeddings.token_embedding.weight", 1)
// A plain decoder LLM used as a text encoder (ERNIE takes Ministral-3 this
// way). Hidden width comes from the embedding table; a vision_tower prefix
// means it can also condition on a reference image.
case dim("model.embed_tokens.weight", 1) > 0:
c.Role = RoleLlm
c.Width = dim("model.embed_tokens.weight", 1)
for k := range h {
if strings.HasPrefix(k, "vision_tower.") || strings.HasPrefix(k, "visual.") {
c.Vision = true
break
}
}
// T5 encoder stack.
case dim("encoder.block.0.layer.0.SelfAttention.q.weight", 0) > 0:
c.Role = RoleT5
c.Width = dim("encoder.block.0.layer.0.SelfAttention.q.weight", 0)
default:
return ComponentFile{}
}
return c
}
// ScanEncoderPool walks the model roots and classifies every diffusion component
// it finds. GGUF headers come from the shared metadata cache, so a component
// GGUF is not parsed twice when discovery has already read it.
func ScanEncoderPool(roots []string) *EncoderPool {
p := &EncoderPool{}
projByDir := map[string]ComponentFile{}
seen := map[string]bool{}
for _, root := range roots {
if strings.TrimSpace(root) == "" {
continue
}
abs, err := filepath.Abs(root)
if err != nil {
continue
}
filepath.WalkDir(abs, func(path string, d fs.DirEntry, err error) error {
if err != nil || d.IsDir() {
return nil
}
key := config.PathKey(path)
if seen[key] {
return nil
}
fi, e := d.Info()
if e != nil {
return nil
}
sizeGB := round(float64(fi.Size())/gib, 2)
if strings.EqualFold(d.Name(), "tokenizer.json") {
seen[key] = true
p.Files = append(p.Files, ComponentFile{Path: path, Role: RoleTokenizer, SizeGB: sizeGB})
return nil
}
switch strings.ToLower(filepath.Ext(path)) {
case ".safetensors":
h, e := readSafetensorsHeader(path)
if e != nil {
return nil
}
// Same exclusion as the gguf ladder below: a PE published as
// safetensors is still not an encoder.
if c := classifySafetensors(h, sizeGB, path); c.Role != RoleNone && !(c.Role == RoleLlm && IsPromptEnhancerFile(path)) {
seen[key] = true
p.Files = append(p.Files, c)
}
case ".gguf":
meta, e := ReadGgufMetadataCached(path)
if e != nil {
return nil
}
arch := strings.ToLower(strings.TrimSpace(meta.Architecture))
switch {
case arch == "clip":
// A vision projector is never a standalone encoder: it is
// paired to the LLM gguf sitting beside it.
projByDir[config.PathKey(filepath.Dir(path))] = ComponentFile{
Path: path, Role: roleProj, SizeGB: sizeGB,
}
case encoderArch[arch]:
seen[key] = true
p.Files = append(p.Files, ComponentFile{
Path: path, Role: RoleT5, Family: arch,
Width: meta.EmbeddingLength, SizeGB: sizeGB,
})
case isImageArch(effectiveImageArch(meta)) || meta.EmbeddingLength == 0:
// A diffusion model, or something with no hidden width:
// not usable as a text encoder.
case IsPromptEnhancerFile(path):
// A prompt rewriter is a chat finetune NAMED AFTER the image
// model it serves, so it lands in that model's own width class
// while being the one file there that was never trained as its
// conditioner. Worse, it is usually the largest: better() breaks
// a same-arch tie on file size, so a BF16 PE outranks the real
// Q8_0 encoder and the model renders confident nonsense.
case isDraftSidecar(filepath.Base(path)):
// A drafter is a reduced head, not an encoder. Its width
// can coincide with a real encoder's (the Gemma-4 MTP
// sidecars are the case), so leaving it in the pool risks
// conditioning a diffusion model on a speculation head.
case IsDecisionModel(meta):
// A decision model is a classifier head over its LM; its
// width can match a real encoder's (kev is a Qwen3), but it
// was never trained to condition anything.
case IsEmbeddingModel(meta):
// A pooled embedder emits one vector for the whole prompt,
// not the per-token sequence a DiT cross-attends to.
case isRecurrentLlm(meta):
// Most of sd.cpp's conditioners are a plain transformer
// stack, so a hybrid cannot load there at all. Left in the
// LLM pool it still wins its width class:
// MiMo-V2.6-Distill-Qwen-9B (a Qwen3.5 hybrid) is 4096
// wide, arch qwen*, has an mmproj and outweighs
// Qwen3-VL-8B, so it took Qwen-Image 2.1 and sd-server
// died on "model metadata validation failed". So it is
// recorded under its own role, which only a family with a
// recurrent conditioner (MOSS-Image) asks for.
seen[key] = true
p.Files = append(p.Files, ComponentFile{
Path: path, Role: RoleRecurrentLlm, Family: arch,
Width: meta.EmbeddingLength, SizeGB: sizeGB,
})
default:
seen[key] = true
p.Files = append(p.Files, ComponentFile{
Path: path, Role: RoleLlm, Family: arch,
Width: meta.EmbeddingLength, SizeGB: sizeGB,
})
}
}
return nil
})
}
p.pairProjectors(projByDir)
sort.Slice(p.Files, func(i, j int) bool { return p.Files[i].Path < p.Files[j].Path })
return p
}
// isRecurrentLlm reports whether an LLM gguf carries recurrent (linear/SSM)
// layers: a GatedDeltaNet hybrid like Qwen3.5/3.6 states full_attention_interval,
// and a Mamba-style model or hybrid states its ssm.* sizes. No diffusion model
// conditions on one, and the same keys already drive the recurrent KV and
// checkpoint sizing, so this is no new identity table.
func isRecurrentLlm(meta Metadata) bool {
return meta.FullAttnInterval > 0 || meta.SsmStateSize > 0 || meta.SsmInnerSize > 0
}
// isDraftSidecar reports whether a gguf filename is one of the speculation
// sidecars discovery already refuses to serve on its own. The pool reuses the
// same patterns so a drafter never becomes a text encoder.
func isDraftSidecar(base string) bool {
return mtpFileRe.MatchString(base) || fastMtpFileRe.MatchString(base) || dflashFileRe.MatchString(base)
}
// pairProjectors attaches each LLM gguf to the vision projector in its own
// directory, reusing the convention discovery already relies on for vision LLMs
// (publishers ship "mmproj-*.gguf" beside the model). This is what makes
// --llm_vision free: the projector for a chosen --llm is simply its neighbour.
func (p *EncoderPool) pairProjectors(projByDir map[string]ComponentFile) {
for i := range p.Files {
if p.Files[i].Role != RoleLlm {
continue
}
if proj, ok := projByDir[config.PathKey(filepath.Dir(p.Files[i].Path))]; ok {
p.Files[i].Mmproj = proj.Path
p.Files[i].Vision = true
}
}
}
// Vae returns the discovered VAE of a family, or "" when none is on disk. hints
// are path substrings preferred when a family holds more than one file: the
// Wan-2.1 and Qwen-Image VAEs have the same shape table (same 194 tensors, every
// dimension equal), so nothing but the path distinguishes them. With several
// candidates and no hint match the pick is the first path in sorted order, which
// is at least stable across regens.
func (p *EncoderPool) Vae(family string, hints ...string) string {
return p.VaeOfWidth(family, 0, hints...)
}
// VaeOfWidth is Vae narrowed to a latent channel count. want 0 means "no
// opinion" and behaves exactly like Vae, which is what every family with a
// single latent size passes.
//
// A want that matches nothing falls back to the full candidate list rather than
// returning "": a family whose files predate the width being recorded (Width 0)
// would otherwise turn a working config into a missing-vae warning, and an
// unfiltered pick is no worse than what this function did before.
func (p *EncoderPool) VaeOfWidth(family string, want int64, hints ...string) string {
if p == nil {
return ""
}
var cands, sized []ComponentFile
for _, f := range p.Files {
if f.Role != RoleVae || f.Family != family {
continue
}
cands = append(cands, f)
if want > 0 && f.Width == want {
sized = append(sized, f)
}
}
if len(sized) > 0 {
return pickByHint(sized, hints)
}
return pickByHint(cands, hints)
}
// pickByHint returns the first candidate whose path contains one of the hints
// (hints in priority order), else the first candidate, else "".
func pickByHint(cands []ComponentFile, hints []string) string {
if len(cands) == 0 {
return ""
}
for _, h := range hints {
h = strings.ToLower(h)
if h == "" {
continue
}
for _, c := range cands {
if strings.Contains(strings.ToLower(filepath.ToSlash(c.Path)), h) {
return c.Path
}
}
}
return cands[0].Path
}
// AudioVae returns the discovered audio VAE of a family, or "" when none is on
// disk. An empty family matches any, which is what a box with a single audio
// model wants; naming one is what keeps a machine holding BOTH an LTX and a
// MiniMax-H3 checkpoint from handing either the other's soundtrack decoder (the
// two are unrelated networks over unrelated latents, so a cross-wire is a load
// failure at best). hints work exactly as in Vae.
func (p *EncoderPool) AudioVae(family string, hints ...string) string {
if p == nil {
return ""
}
var cands []ComponentFile
for _, f := range p.Files {
if f.Role == RoleAudioVae && (family == "" || f.Family == family) {
cands = append(cands, f)
}
}
return pickByHint(cands, hints)
}
// LlmHinted returns the text-encoder LLM whose PATH matches one of the hints,
// or "" when none does. Unlike Llm it never falls back to an arbitrary
// candidate: it exists for a family whose encoder is a bespoke, republished file
// (LTX ships a Gemma-4-12B with the caption projection grafted on), where the
// width alone would also match every stock Gemma-3-12B on the disk. A miss here
// is meant to fall through to the ordinary width-matched pick, not to guess.
func (p *EncoderPool) LlmHinted(hints ...string) string {
if p == nil {
return ""
}
for _, h := range hints {
h = strings.ToLower(strings.TrimSpace(h))
if h == "" {
continue
}
for _, f := range p.Files {
if f.Role != RoleLlm {
continue
}
if strings.Contains(strings.ToLower(filepath.ToSlash(f.Path)), h) {
return f.Path
}
}
}
return ""
}
// RecurrentLlm returns the recurrent (hybrid) text encoder of exactly this
// width, for the one kind of caller that can load one: a family whose sd.cpp
// runner ships its own recurrent conditioner (MOSS-Image's Qwen35DeltaNet).
// Exact width only, no caption concatenation: no recurrent-conditioned DiT
// stacks layers, and width is the only structural gate this role has.
//
// hints are tried first, in order, against the path: a hybrid distill of the
// same base (MiMo-V2.6-Distill-Qwen-9B against Qwen3.5-9B) has the same width
// and arch, and size would hand the DiT the wrong one. Without a hint hit the
// usual better() order decides.
func (p *EncoderPool) RecurrentLlm(hidden int64, hints ...string) string {
if p == nil || hidden <= 0 {
return ""
}
var cands []ComponentFile
for _, f := range p.Files {
if f.Role == RoleRecurrentLlm && f.Width == hidden {
cands = append(cands, f)
}
}
for _, h := range hints {
h = strings.ToLower(strings.TrimSpace(h))
if h == "" {
continue
}
for _, f := range cands {
if strings.Contains(strings.ToLower(filepath.ToSlash(f.Path)), h) {
return f.Path
}
}
}
var best ComponentFile
for _, f := range cands {
if best.Path == "" || better(f, best) {
best = f
}
}
return best.Path
}
// Tokenizer returns the standalone tokenizer.json whose path carries one of
// the hints, tried in order. Hint-scoped with NO fallback, unlike the VAE and
// LLM pickers: every HF model folder on the box ships a tokenizer.json, and a
// wrong vocab is not a crash, it is an encoder reading a prompt as different
// tokens. An empty answer surfaces as the "needs component [tokenizer]"
// warning, which the model browser's companion download can clear.
func (p *EncoderPool) Tokenizer(hints ...string) string {
if p == nil {
return ""
}
for _, h := range hints {
h = strings.ToLower(strings.TrimSpace(h))
if h == "" {
continue
}
for _, f := range p.Files {
if f.Role == RoleTokenizer && strings.Contains(strings.ToLower(filepath.ToSlash(f.Path)), h) {
return f.Path
}
}
}
return ""
}
// Clip returns the discovered CLIP tower of a given width (768 = L, 1280 = G).
func (p *EncoderPool) Clip(width int64) string {
if p == nil {
return ""
}
for _, f := range p.Files {
if f.Role == RoleClip && f.Width == width {
return f.Path
}
}
return ""
}
// T5 returns the discovered T5/UMT5 encoder, preferring the widest: T5-XXL is
// 4096, and a stray smaller T5 would condition a flux model into mush.
func (p *EncoderPool) T5() string {
if p == nil {
return ""
}
best := ComponentFile{}
for _, f := range p.Files {
if f.Role == RoleT5 && f.Width >= best.Width {
best = f
}
}
return best.Path
}
// encoderArchRank ranks a candidate text encoder by how likely its arch is to be
// the one a diffusion model was trained against. Hidden width alone is NOT an
// identity: on the dev box FIVE unrelated 2560-wide models are installed (Qwen3-4B,
// Qwen3.5-4B, Qwen3-VL-4B, Granite-4.2-3B, Gemma-4-E4B), and picking the wrong one
// produces confident garbage rather than an error, so a size tiebreak alone was
// wiring Z-Image to Gemma.
//
// The ranking is shipped knowledge, not per-machine tuning: essentially every
// LLM-conditioned DiT in circulation (Qwen-Image, Z-Image, LongCat, Krea,
// Flux.2-Klein) conditions on a Qwen, and the Mistral rung covers ERNIE-Image and
// Flux.2-dev. Lumina-Image-2.0 proper is the known exception (it conditions on
// Gemma-2), which is why Gemma is ranked rather than excluded, and why
// settings.encoders / textEncoderPath still win outright.
func encoderArchRank(arch string) int {
a := strings.ToLower(arch)
switch {
case strings.HasPrefix(a, "qwen"):
return 3
case strings.HasPrefix(a, "mistral") || strings.HasPrefix(a, "ministral"):
return 2
case strings.HasPrefix(a, "gemma") || strings.HasPrefix(a, "llama"):
return 1
}
return 0
}
// better orders two same-width candidates: known encoder arch first, then the
// larger file (the encoder runs once per generation, usually on the CPU, so a
// higher quant is nearly free), then the path so a regen is reproducible.
func better(f, best ComponentFile) bool {
if r, br := encoderArchRank(f.Family), encoderArchRank(best.Family); r != br {
return r > br
}
if f.SizeGB != best.SizeGB {
return f.SizeGB > best.SizeGB
}
return f.Path < best.Path
}
// Llm returns the text-encoder LLM whose hidden width matches what the DiT's
// caption projection expects, plus its vision projector when one is paired.
// wantVision restricts the pick to encoders that have one: an edit model without
// its vision tower does not fail, it silently conditions on nothing and emits an
// unrelated image, which is worse than not starting at all.
//
// "Matches" is captionFactor, not equality: a DiT that concatenates several
// encoder hidden layers states a caption width that is a MULTIPLE of the encoder
// it wants (flux.2 klein 9B: txt_in 12288 against a 4096-wide Qwen3-8B). An
// exact match always outranks a concatenated one.
//
// prefer is the declared settings.encoders.qwenLlm (empty when none declared).
// A declared file that clears the same width and vision gates wins outright:
// width only narrows the field, it does not name the file a publisher trained
// against. Qwen3-4B-Instruct-2507 and Qwen3-VL-4B-Instruct are both 2560 wide
// and round to the same size, so the tiebreak below will hand one of them to
// either kind of DiT; a user who pinned the right one has answered the question.
// A prefer that fails the gates (wrong width, no vision tower for a vision
// model, a file the scan never saw) is ignored and the scan decides.
//
// Ties break by encoderArchRank, then file size, then path (see better).
func (p *EncoderPool) Llm(hidden int64, wantVision bool, prefer string) (path, mmproj string) {
return p.LlmArch(hidden, wantVision, prefer, "")
}
// LlmArch is Llm restricted to one encoder ARCHITECTURE (the gguf's
// general.architecture, recorded as Family), for a DiT whose trained-against
// encoder shares its width with a different model. "" means no restriction.
//
// Width cannot separate Qwen3-4B (`qwen3`) from Qwen3-VL-4B (`qwen3vl`): both
// are 2560 wide, and FLUX.3 (8 x 2560) wants the VL one while Z-Image and
// FLUX.2 klein want the plain one. Before this, the answer was whichever won
// the tiebreak (the larger file) or the single global qwenLlm pin, so ONE of
// the two families was always wired wrong on a box that had both, with no
// warning: a mismatched encoder renders, just not the asked-for picture.
//
// The restriction applies to the pin too: a pin of the wrong architecture is
// ignored and the scan decides, exactly as for a pin of the wrong width. A
// file of UNKNOWN family (a safetensors encoder, which records no arch) is not
// excluded: only a positively different architecture is.
func (p *EncoderPool) LlmArch(hidden int64, wantVision bool, prefer, arch string) (path, mmproj string) {
if p == nil {
return "", ""
}
if hidden <= 0 {
return "", ""
}
pin := config.PathKey(strings.TrimSpace(prefer))
var best, pinned ComponentFile
var bestFactor, pinnedFactor int64
for _, f := range p.Files {
if f.Role != RoleLlm {
continue
}
factor := captionFactor(hidden, f.Width)
if factor == 0 {
continue
}
if wantVision && !f.Vision {
continue
}
if arch != "" && f.Family != "" && !strings.EqualFold(f.Family, arch) {
continue
}
if pin != "" && config.PathKey(f.Path) == pin {
if pinnedFactor == 0 || factor < pinnedFactor {
pinned, pinnedFactor = f, factor
}
continue
}
if bestFactor == 0 || factor < bestFactor || (factor == bestFactor && better(f, best)) {
best, bestFactor = f, factor
}
}
// The pin wins inside its tier, not across tiers: an exact-width scan hit
// beats a pin that only fits by concatenation, because k>1 is an inference
// about the model's layout and k==1 is not.
if pinnedFactor != 0 && (bestFactor == 0 || pinnedFactor <= bestFactor) {
return pinned.Path, pinned.Mmproj
}
return best.Path, best.Mmproj
}
// captionLayouts are the encoder-layer concatenations a DiT is known to use,
// i.e. the k in txt_in = k x encoder width. 1 is every single-hidden-state DiT;
// 3 is FLUX.2 (klein 9B: 12288 = 3 x 4096, Qwen3-8B); 8 is FLUX.3 klein
// (20480 = 8 x 2560, Qwen3-VL-4B).
//
// A whitelist, not a cap, because the caller prefers the LOWEST factor and the
// two disagree the moment a DiT concatenates more than a few layers. Under the
// old "any k <= 4" rule FLUX.3's 20480 matched nothing (8 > 4); raising the cap
// to 8 would have matched 4 x 5120 (a Qwen3-32B) and 5 x 4096 (any Qwen3-8B on
// the box) ahead of the 8 x 2560 it was trained against, wiring the wrong
// encoder with no warning. Only layouts that exist get to compete.
var captionLayouts = []int64{1, 3, 8}
// captionFactor reports how many encoder hidden states of the given width this
// DiT's caption projection consumes, or 0 when that encoder cannot be the one.
//
// Nearly every DiT projects a SINGLE encoder hidden state, so txt_in is exactly
// the encoder width and the answer is 1. The FLUX.2/3 line is the exception: it
// conditions on several hidden LAYERS of its LLM concatenated. Matched on
// equality that finds nothing, and the caller's pin fallback then substitutes
// an encoder of the wrong width with no warning at all.
//
// Exactness is still preferred by the caller.
func captionFactor(hidden, width int64) int64 {
if hidden <= 0 || width <= 0 || hidden%width != 0 {
return 0
}
k := hidden / width
for _, l := range captionLayouts {
if k == l {
return k
}
}
return 0
}
// knowsLlm reports whether the scan classified this exact path as a text-encoder
// LLM, at ANY width. That is a different question from whether Llm would PICK it
// for a given DiT: a pin the scan measured and rejected on width is a mismatch
// worth reporting, while one the scan never saw (kept outside the models root,
// say) is still the user's explicit choice and should be honoured.
func (p *EncoderPool) knowsLlm(path string) bool {
key := config.PathKey(strings.TrimSpace(path))
if p == nil || key == "" {
return false
}
for _, f := range p.Files {
if f.Role == RoleLlm && config.PathKey(f.Path) == key {
return true
}
}
return false
}
// llmCandidate reports whether a declared pin would actually be HONOURED by Llm
// at this width, which is a different question from whether one was declared.
// settings.encoders.qwenLlm is a single global field shared by every diffusion
// model on the box, and Llm applies it only to candidates that clear the same
// width and vision gates, so a pin aimed at an image model is silently dropped
// for a video one. A family that carries its own path hint has to test for the
// pin APPLYING, or a populated encoders block (every real install has one)
// suppresses the hint without ever using the pin.
//
// The width test is captionFactor, not equality, for the same reason Llm's is:
// a pin this mirror rejected on a width Llm would have accepted by concatenation
// would fire the family's path hint over a pin that was about to be honoured.
func (p *EncoderPool) llmCandidate(path string, hidden int64, wantVision bool) bool {
key := config.PathKey(strings.TrimSpace(path))
if p == nil || key == "" || hidden <= 0 {
return false
}
for _, f := range p.Files {
if f.Role != RoleLlm || captionFactor(hidden, f.Width) == 0 || config.PathKey(f.Path) != key {
continue
}
return !wantVision || f.Vision
}
return false
}
// Pool scans are cached per root set: a regen happens on every settings save and
// on every models-watcher tick, and re-walking the tree (plus re-reading every
// safetensors header) each time would dominate. The TTL keeps a freshly
// downloaded encoder from needing a restart to be seen.
var (
poolCacheMu sync.Mutex
poolCache = map[string]*cachedPool{}
)
const poolCacheTTL = 30 * time.Second
type cachedPool struct {
pool *EncoderPool
at time.Time
}
func encoderPoolFor(roots []string) *EncoderPool {
key := strings.Join(roots, "\x00")
poolCacheMu.Lock()
e, ok := poolCache[key]
poolCacheMu.Unlock()
if ok && time.Since(e.at) < poolCacheTTL {
return e.pool
}
p := ScanEncoderPool(roots)
poolCacheMu.Lock()
poolCache[key] = &cachedPool{pool: p, at: time.Now()}
poolCacheMu.Unlock()
return p
}
// InvalidateEncoderPool drops every cached scan, for a caller that has just put
// a component on disk and is about to regen. The TTL alone is not enough there:
// the model browser queues a DiT and its companions together, a small VAE lands
// seconds after the previous regen scanned, and a cached pool would wire the
// model without it until the next unrelated regen.
func InvalidateEncoderPool() {
poolCacheMu.Lock()
clear(poolCache)
poolCacheMu.Unlock()
}
// fillEncoderSet returns declared with every blank field filled from what the
// scan actually found on disk. Declared paths always win: an explicit setting is
// a deliberate statement about a machine, and silently second-guessing it would
// make a working config drift.
//
// The one exception is a declared file that no longer EXISTS. It is cleared
// first, so the pool fills the slot or the role is reported missing. Kept, it
// read as satisfied: the regen emitted a flag at a path sd-server then failed to
// open, with no warning, and the model browser never offered the file again
// because nothing was "missing". Moving one encoder out of the models folder
// was enough to hit it.
//
// The hints exist because a family can hold more than one file (see Vae).
func fillEncoderSet(declared EncoderSet, p *EncoderPool) EncoderSet {
if p == nil {
return declared
}
for _, dst := range []*string{
&declared.FluxVae, &declared.ClipL, &declared.ClipG, &declared.T5, &declared.SdxlVae,
&declared.ZimageVae, &declared.QwenLlm, &declared.Flux2Vae, &declared.VideoVae, &declared.AudioVae,
} {
if declaredGone(*dst, p) {
*dst = ""
}
}
fill := func(dst *string, v string) {
if strings.TrimSpace(*dst) == "" {
*dst = v
}
}
fill(&declared.FluxVae, p.Vae(VaeFamilyFlux, "flux", "/ae."))
fill(&declared.Flux2Vae, p.Vae(VaeFamilyFlux2))
fill(&declared.SdxlVae, p.Vae(VaeFamilySD, "sdxl"))
// Z-Image ships "its own" ae.safetensors that is byte-for-byte the flux.1
// AE (verified: same size, same hash), so the flux pick is the correct
// fallback and one copy on disk serves both.
fill(&declared.ZimageVae, p.Vae(VaeFamilyFlux, "z-image", "zimage"))
fill(&declared.ZimageVae, declared.FluxVae)
// Video components. Both fields describe the MiniMax-H3 pair: the 3D
// transformer VAE has no image-model consumer, so an unhinted pick within
// the family is safe. LTX's two VAEs are deliberately NOT filled here -
// there is one slot per role and the two families' files are not
// interchangeable, so LTX resolves its own pair structurally in
// videoComponents and leaves these pins meaning what they have always meant.
fill(&declared.VideoVae, p.Vae(VaeFamilyVideo3D))
fill(&declared.AudioVae, p.AudioVae(VaeFamilyVideo3D))
fill(&declared.ClipL, p.Clip(768))
fill(&declared.ClipG, p.Clip(1280))
fill(&declared.T5, p.T5())
return declared
}
// declaredGone reports a declared component path that is provably absent: the
// scan did not find it AND the filesystem says it does not exist. A file the pool
// lists is present by construction, and any stat error other than not-exist
// (permissions, a slow share) keeps the declaration, since dropping a working
// path on a transient error would rewire a model behind the user's back.
func declaredGone(path string, p *EncoderPool) bool {
path = strings.TrimSpace(path)
if path == "" {
return false
}
key := config.PathKey(path)
for _, f := range p.Files {
if config.PathKey(f.Path) == key {
return false
}
}
_, err := os.Stat(path)
return errors.Is(err, fs.ErrNotExist)
}
// projectorBeside returns the vision projector paired to a given text-encoder
// gguf, i.e. the mmproj the pool found in that file's own directory. Empty for a
// safetensors encoder (its vision tower, if any, is inside the same file) or
// when the encoder ships without one.
func projectorBeside(llmPath string, p *EncoderPool) string {
if p == nil || strings.TrimSpace(llmPath) == "" {
return ""
}
key := config.PathKey(llmPath)
for _, f := range p.Files {
if f.Role == RoleLlm && config.PathKey(f.Path) == key {
return f.Mmproj
}
}
return ""
}
// editModelRe marks a diffusion model whose pipeline conditions on a REFERENCE
// IMAGE and therefore needs the text encoder's vision tower (--llm_vision).
//
// This one genuinely cannot be read off the weights. The obvious structural
// tells do not hold: LongCat-Image-Edit and plain LongCat-Image have identical
// img_in/txt_in shapes and differ by two unrelated norm tensors, because the
// reference image enters as extra SEQUENCE tokens rather than extra input
// channels. So the model name is the signal, with an explicit llmVision
// override as the escape hatch when a publisher names something unhelpfully.
var editModelRe = regexp.MustCompile(`(?i)(^|[-_. ])(edit|rapid|kontext|instruct[-_]?pix2pix|inpaint|redux)([-_. ]|$)`)
// unifiedEditRe matches models where ONE checkpoint does both text-to-image and
// reference editing, so the name carries no edit token to detect and the two
// regexes above both miss. Qwen-Image 2.1 is the first of these: upstream
// documents `-r ref.png` against the same weights used for plain generation,
// and every 2.1 checkpoint is edit-capable, so this is a property of the model
// version rather than of the individual file.
//
// Still matched by name, for the reason given above: the reference enters as
// extra sequence tokens, so an edit-capable and a text-only checkpoint have the
// same tensor shapes. The version number is the reliable part of the name here,
// since both upstream releases (leejet's gguf, Comfy-Org's safetensors) spell
// it in the filename, whereas neither says "edit" anywhere.
//
// Getting this wrong is a silent downgrade, not an error: the model would
// declare `in: [text]`, the playground would fall through to img2img, and that
// route scales the step count by the denoise strength and redraws the whole
// frame instead of editing against a reference.
var unifiedEditRe = regexp.MustCompile(`(?i)(^|[-_. ])qwen[-_. ]?image[-_. ]?2\.1([-_. ]|$)`)
// fluxRefEditRe matches the FLUX.2 / FLUX.3 line (dev and klein), which is
// unified the same way: BFL ships one checkpoint for text-to-image, a single
// reference edit and multi-reference composition (FLUX.3 klein's card: "up to
// 10 reference images ... through the same model"). Every converter tags these
// general.architecture "flux", the same as flux.1, so the name is again the
// only signal; it is the same set resolveComponents' klein arm keys on.
//
// Kept OUT of unifiedEditRe on purpose, because that one also drives
// wantsVisionEncoder. FLUX takes its references through the VAE, as latent
// sequence tokens, not through the text encoder's vision tower: klein 4B/9B
// condition on text-only Qwen3, so there is no projector to pair and a
// --llm_vision would be a wasted load at best. Only the capability changes.
var fluxRefEditRe = regexp.MustCompile(`(?i)(^|[-_. ])(flux[-_. ]?[23]|klein)([-_. ]|$)`)
// wantsVisionEncoder reports whether this model should get --llm_vision. Only
// llm-conditioned families are candidates: flux.1 edit models condition through
// T5, which has no vision tower at all.
// refEditRe matches models that take their source image as a REFERENCE (extra
// sequence tokens) rather than as an img2img denoise base. Deliberately
// narrower than editModelRe: inpaint/redux models are edits too, but they want
// the masked img2img route, and routing them through reference images would
// redraw the whole frame instead of honouring the mask.
var refEditRe = regexp.MustCompile(`(?i)(^|[-_. ])(edit|rapid|kontext)([-_. ]|$)`)
// IsReferenceEditModel reports whether a model conditions on a reference image.
// It is what the emitted `in: [text, image]` capability is keyed on, so the
// client can route an edit to the reference path instead of guessing from the
// model id. Like the projector pairing this is name detection, because the
// distinction is not structural: the reference enters as extra sequence tokens,
// not as extra input channels, so a base and an edit checkpoint have identical
// tensor shapes. An explicit llmVision pin doubles as the escape hatch.
func IsReferenceEditModel(name string, ov *Override) bool {
if ov != nil {
// refEdit is the field that means exactly this, so it wins outright.
switch ov.RefEdit {
case "on":
return true
case "off":
return false
}
// An llmVision pin is the older, narrower spelling: a user who pinned a
// vision projector on has said this model consumes a reference image.
switch ov.LlmVision {
case "on":
return true
case "off":
return false
}
}
return refEditRe.MatchString(name) || unifiedEditRe.MatchString(name) || fluxRefEditRe.MatchString(name)
}
func wantsVisionEncoder(arch, name string, ov *Override) bool {
if ov != nil {
switch ov.LlmVision {
case "on":
return true
case "off":
return false
}
// refEdit on is the same statement from the other side: the user has
// said this model consumes a reference image, and for an llm-conditioned
// family the reference is READ by the encoder's vision tower. Pinning
// one without the other is the trap this closes: sd.cpp refuses the job
// outright ("Qwen Image 2.1 editing requires Qwen3-VL vision weights;
// provide --llm_vision or a combined encoder"), so a user who turns on
// reference editing and nothing else gets a model that cannot generate.
//
// Only the "on" direction implies anything. refEdit off does NOT mean no
// vision: an inpaint model is an edit that wants the masked img2img
// route, and it still needs the projector.
if ov.RefEdit == "on" {
return true
}
}
return editModelRe.MatchString(name) || unifiedEditRe.MatchString(name)
}