internal/autogen/discover.go
Rama main · réplica de lectura del repositorio original.
package autogen
import (
"fmt"
"io/fs"
"os"
"path/filepath"
"regexp"
"strings"
"github.com/quartermaster-labs/quartermaster/internal/config"
"github.com/quartermaster-labs/quartermaster/internal/quant"
)
// GgufRow describes one discovered served model (one .gguf, or the first shard
// of a split set). Mirrors the PowerShell Get-LocalGgufModels output row.
type GgufRow struct {
ID string // BaseID + "-<quant>" (lowercased); BaseID when no quant
BaseID string // model name only, no publisher prefix, no quant
FullPath string // real file llama-server is given (first shard)
FileName string // shard-stripped file name
Quant string // detected quant token, upper-case ("" if none)
SizeGB float64
Publisher string
Repo string
// Separate MTP or DFlash draft gguf sitting in the same dir, "" if none.
// MTP sidecars are named "mtp-*.gguf"; DFlash drafters carry "dflash"
// anywhere in the name (e.g. "Qwen3.6-35B-A3B-DFlash-Q8_0.gguf").
// When this dir has none, a header-compatible one from another quant of the
// same model or from a finetune sibling is inherited (family.go).
DraftPath string
DraftKind string // "mtp" | "dflash"; "" when DraftPath == ""
DraftSizeGB float64 // on-disk size of DraftPath; 0 when DraftPath == ""
// Vision projector (clip-arch gguf) sitting in the same dir, if any. Loaded via
// --mmproj to enable image input; drives the auto-generated "-vision" variant.
// Inherited from a family sibling the same way DraftPath is when this dir
// ships none.
MmprojPath string
MmprojSizeGB float64
// Qwen3-TTS audio codec (the "qwen-tokenizer-*hz" gguf) sitting in the same
// dir as a talker, if any. Loaded via tts-server --codec, never served alone.
CodecPath string
CodecSizeGB float64
// IsSam marks a SAM (Segment Anything) model: a *.ggml file (NOT gguf — sam3.cpp
// uses the raw ggml format, no metadata header) served by sam3_server. Routed
// before the gguf metadata read in emitModel, since ReadGgufMetadata can't parse it.
IsSam bool
// IsTrellis marks a TRELLIS.2 package: a DIRECTORY (pipeline.json plus
// ckpts/*.safetensors) served by trellis2-server, with FullPath naming the
// directory instead of a file. Like IsSam it routes before the gguf metadata
// read, which has no file to open here at all.
IsTrellis bool
// IsHF marks a Hugging Face model folder (config.json + *.safetensors, see
// hf.go): FullPath is the DIRECTORY `vllm serve` takes, and metadata comes
// from config.json. Only vllm can load one, so it routes before the gguf
// metadata read and never reaches the llama path.
IsHF bool
}
// skipsGgufPipeline reports whether the row is served by a provider binary
// rather than llama-server, so the gguf-only passes — the metadata read, family
// pairing, sidecar inheritance — must leave it alone: SAM's *.ggml has no header
// at all, and a TRELLIS.2 package or an HF folder is a directory.
func (r GgufRow) skipsGgufPipeline() bool { return r.IsSam || r.IsTrellis || r.IsHF }
var (
shardRe = regexp.MustCompile(`-(\d{5})-of-(\d{5})\.gguf$`)
ggufSuffixRe = regexp.MustCompile(`(?i)-GGUF$`)
// Separate MTP/draft sidecar (e.g. Gemma-4 ships "mtp-gemma-4-12B-it.gguf"
// alongside the main model). Loaded via -md + --spec-type draft-mtp, not served alone.
mtpFileRe = regexp.MustCompile(`(?i)^mtp[-_.]`)
// "FastMTP" heads ("...-Aggressive-FastMTP-32K.gguf") are skipped outright:
// neither served nor paired as a draft.
//
// The file is a well-formed reduced-vocab draft, not a broken one. It holds
// exactly the 19 tensors of the single nextn block (blk.64.nextn.*) with its
// output.weight shrunk to [n_embd, 32768] plus a `d2t` table remapping those
// draft ids back onto the parent's 248320-token vocab: the EAGLE-3 style
// convention, where the small head is the whole point of the speedup.
// Two independent reasons it must not become a -md sidecar:
// 1. No llama-server build we ship references `d2t` at all, so the loader
// validates output.weight against the tokenizer's n_vocab and hard-fails
// ("tensor 'output.weight' has wrong shape; expected 5120, 248320, got
// 5120, 32768"), taking the whole launch down with it.
// 2. It is redundant regardless: these repos bake the same nextn block into
// the main gguf at full head size, so IsMTP is already true and
// effectiveSpec drives draft-mtp off the baked layer with no -md.
//
// Revisit only if a build starts advertising reduced-vocab drafts. Kept
// distinct from mtpFileRe so a genuine sidecar still pairs. Deliberately
// pinned to "fastmtp", not a bare "mtp" infix, which would also swallow real
// models naming a baked-in head ("...-Native-MTP-Preserved-Q4_K_M.gguf").
fastMtpFileRe = regexp.MustCompile(`(?i)(^|[-_.])fast[-_.]?mtp([-_.]|\.gguf$)`)
// Separate DFlash draft sidecar (block-diffusion drafter, e.g.
// "Qwen3.6-35B-A3B-DFlash-Q8_0.gguf"). Unlike the MTP prefix convention,
// publishers embed "dflash" as an infix, so match anywhere in the name.
// Loaded via -md + --spec-type draft-dflash, not served alone.
dflashFileRe = regexp.MustCompile(`(?i)dflash`)
// Qwen3-TTS audio codec / tokenizer sidecar (e.g. "qwen-tokenizer-12hz-Q8_0.gguf").
// Loaded via tts-server --codec alongside the talker, never served as its own
// model. "hz" scopes the tokenizer match so a normal model file can't trip it.
ttsCodecFileRe = regexp.MustCompile(`(?i)tokenizer-?\d+hz|codec`)
// Diffusion text encoders / VAEs (T5-XXL, UMT5, CLIP-L/G, standalone VAE).
// These are COMPONENTS of an image model — wired in via settings.encoders
// (see image.go resolveComponents), never served on their own. Without this
// they parse as ordinary ggufs and get emitted as llama-server "LLM" rows
// (a t5 encoder has no chat template and no decoder; it can't generate).
// Narrow on purpose: an "-encoder" tail or a known encoder/VAE stem only, so
// a real seq2seq LLM (flan-t5-*) is untouched.
//
// "with-proj" is the odd one out and it earns its place: LTX-2.x republishes
// an ordinary Gemma-4-12B with the DiT's caption projection grafted onto it
// and names the result *-with-proj-*. It parses as a perfectly good chat
// model, so without this it is SERVED as one, and the extra projection head
// makes it a strictly worse Gemma than the stock file beside it.
encoderFileRe = regexp.MustCompile(`(?i)(^|[-_.])(t5xxl|t5[-_]?v1[-_.]?1|umt5|clip[-_]?[lg]|text[-_]?encoder|ae|vae|taesd\w*)([-_.]|\.gguf$)|[-_.](encoder|with[-_]?proj)([-_.]|\.gguf$)`)
// GGUF architectures that are encoder-only diffusion components. "clip" is
// handled separately (it doubles as a vision projector and IS paired).
// Plain "t5" is deliberately absent: flan-t5 is a real seq2seq LLM
// llama.cpp serves; only the encoder-only archs are components.
encoderArch = map[string]bool{"t5encoder": true, "umt5": true}
// SAM (Segment Anything) model files, served by sam3_server. These are raw
// *.ggml (not gguf), so they're matched by name — ".ggml" alone is too generic
// (other ggml projects use it). Add new families here (mobilesam, etc.).
samFileRe = regexp.MustCompile(`(?i)(sam2|sam3|edgetam|hiera)`)
)
// quantFromName extracts the quant token (upper-cased) from a gguf file name,
// or "" when none is present. Thin alias for the shared matcher in
// internal/quant, which is where the token shape itself lives.
func quantFromName(name string) string { return quant.FromName(name) }
// hasQuantPart reports whether an already-lowercased base id carries quant as a
// whole separator-bounded part. The quant is stripped from the base name only
// when it TRAILS it, so a publisher who writes the build tag last
// ("Qwen3.8-27B-NVFP4-MTP-MID-HIGH") leaves it embedded - and appending it again
// would mint "…-nvfp4-mtp-mid-high-nvfp4". The id is the key every override,
// favourite and client request is written against, so it stays the name the file
// already spells. Bounded on both sides because the token itself contains "_"
// (Q4_K_M), which rules out splitting the id into parts.
func hasQuantPart(baseKey, quant string) bool {
re := regexp.MustCompile(`(?i)(^|[-_.])` + regexp.QuoteMeta(quant) + `([-_.]|$)`)
return re.MatchString(baseKey)
}
// DiscoverGgufModelsMulti walks each root and returns the union of rows,
// de-duplicated by full path (a model reachable from two overlapping roots is
// served once). Roots are scanned in order; the first occurrence wins.
func DiscoverGgufModelsMulti(roots []string, skipPatterns ...string) ([]GgufRow, error) {
seen := map[string]bool{}
var all []GgufRow
for _, root := range roots {
rows, err := DiscoverGgufModels(root, skipPatterns...)
if err != nil {
return nil, err
}
for _, row := range rows {
key := config.PathKey(row.FullPath)
if seen[key] {
continue
}
seen[key] = true
all = append(all, row)
}
}
// Re-run inheritance over the union: each root's own pass could only see its
// own rows, so a model kept in one root and its finetune's drafter in another
// would never have met. Blanks only, so per-root results are preserved.
inheritSidecars(all)
return all, nil
}
// DiscoverGgufModels walks modelsRoot for *.gguf files and returns one row per
// served model. Vision projectors (clip-arch ggufs) and non-first split shards
// are skipped (projectors are paired to their model dir instead). skipPatterns
// are optional filename globs to exclude.
func DiscoverGgufModels(modelsRoot string, skipPatterns ...string) ([]GgufRow, error) {
if strings.TrimSpace(modelsRoot) == "" {
return nil, nil // no models root configured yet => empty catalog
}
info, err := filepath.Abs(modelsRoot)
if err != nil {
return nil, err
}
modelsRoot = info
var rows []GgufRow
draftByDir := map[string]draftSidecar{} // dir -> paired MTP/DFlash draft gguf
mmprojByDir := map[string]mmprojSidecar{} // dir -> vision projector gguf
codecByDir := map[string]codecSidecar{} // dir -> Qwen3-TTS audio codec gguf
walkErr := filepath.WalkDir(modelsRoot, func(path string, d fs.DirEntry, err error) error {
if err != nil {
return nil // skip unreadable entries, matching -ErrorAction SilentlyContinue
}
// A TRELLIS.2 package is a directory, and nothing inside it is a model in
// its own right (the BiRefNet remover may well live under it): emit the
// package and stop, so ckpts/*.gguf or a vendored gguf cannot turn into
// separate models. The root itself never counts — pointing a scan root at
// a package must not swallow the walk.
if d.IsDir() {
if path != modelsRoot {
if IsTrellisPackageDir(path) {
rows = append(rows, trellisRow(path))
return filepath.SkipDir
}
}
// An HF folder is a row, but NOT a stop: unlike a TRELLIS package its
// subfolders are ordinary (a repo that also ships a gguf/ subfolder keeps
// those rows), and the root may itself be the folder.
if c, ok := hfModelDir(path); ok {
rows = append(rows, hfRow(path, c))
}
return nil
}
// SAM models are raw *.ggml (no gguf header); served by sam3_server. Emit a
// row by filename and skip the gguf pipeline (metadata/quant/shard logic).
if !d.IsDir() && strings.EqualFold(filepath.Ext(path), ".ggml") && samFileRe.MatchString(d.Name()) {
fi, e := d.Info()
if e != nil {
return nil
}
base := strings.TrimSuffix(d.Name(), filepath.Ext(d.Name()))
id := strings.ToLower(strings.ReplaceAll(base, "_", "-"))
rows = append(rows, GgufRow{
ID: id,
BaseID: id,
FullPath: path,
FileName: d.Name(),
SizeGB: round(float64(fi.Size())/gib, 2),
IsSam: true,
})
return nil
}
if d.IsDir() || !strings.EqualFold(filepath.Ext(path), ".gguf") {
return nil
}
name := d.Name()
// BiRefNet is TRELLIS.2's background remover: a gguf, but not a language
// model, and llama-server cannot load it. Skipped rather than served so a
// scan over the weights folder does not offer it as a chat model; the
// trellis emitter picks the same file up as --birefnet instead.
if trellisBirefRe.MatchString(name) {
return nil
}
// Unloadable FastMTP head: not served, not paired. Checked ahead of the
// draft rules so it can never become a -md sidecar.
if fastMtpFileRe.MatchString(name) {
return nil
}
// Separate MTP or DFlash draft: record it for pairing, don't serve it as
// its own model. Checked in this order since an MTP sidecar never also
// matches the dflash infix.
if mtpFileRe.MatchString(name) {
if fi, e := d.Info(); e == nil {
draftByDir[filepath.Dir(path)] = draftSidecar{path: path, sizeGB: round(float64(fi.Size())/gib, 2), kind: "mtp"}
}
return nil
}
if dflashFileRe.MatchString(name) {
if fi, e := d.Info(); e == nil {
draftByDir[filepath.Dir(path)] = draftSidecar{path: path, sizeGB: round(float64(fi.Size())/gib, 2), kind: "dflash"}
}
return nil
}
// Qwen3-TTS audio codec: pair it to the talker in its dir, don't serve it.
if ttsCodecFileRe.MatchString(name) {
if fi, e := d.Info(); e == nil {
codecByDir[filepath.Dir(path)] = codecSidecar{path: path, sizeGB: round(float64(fi.Size())/gib, 2)}
}
return nil
}
// Diffusion text encoder / VAE: an image-model component, not a model.
// Dropped outright (unlike mmproj/codec sidecars, which are paired) — the
// image emitter resolves these from settings.encoders by explicit path.
if encoderFileRe.MatchString(name) {
return nil
}
for _, p := range skipPatterns {
if ok, _ := filepath.Match(p, name); ok {
return nil
}
}
// Split shards "<model>-00001-of-00003.gguf": represent the set by shard 1
// only and strip the shard token before id derivation. FullPath stays the
// real first-shard file.
if mm := shardRe.FindStringSubmatch(name); mm != nil {
if mm[1] != "00001" {
return nil
}
name = shardRe.ReplaceAllString(name, ".gguf")
}
fi, statErr := d.Info()
if statErr != nil {
return nil
}
// Vision projector detection by header arch, not filename — projector
// naming varies by publisher (mmproj-*, <model>-mmproj-*, clip-*, bare).
// A clip-arch gguf is paired to its model dir, never served alone. The
// cached read is reused by the later VRAM planning pass, so it's not an
// extra parse. A parse error falls through and the file is treated as a
// normal model candidate.
if meta, e := ReadGgufMetadataCached(path); e == nil {
if meta.Architecture == "clip" {
mmprojByDir[filepath.Dir(path)] = mmprojSidecar{
path: path,
sizeGB: round(float64(fi.Size())/gib, 2),
}
return nil
}
// Encoder-only arch (t5encoder/umt5): a diffusion text encoder that
// escaped the filename rule. No decoder → nothing to serve.
if encoderArch[strings.ToLower(meta.Architecture)] {
return nil
}
}
quant := quantFromName(name)
repoDir := filepath.Dir(path)
repoName := filepath.Base(repoDir)
pubName := filepath.Base(filepath.Dir(repoDir))
base := strings.TrimSuffix(name, filepath.Ext(name)) // drop .gguf
if quant != "" {
// Strip a trailing [-_.]<quant> from the base name.
trailRe := regexp.MustCompile(`(?i)[-_.]` + regexp.QuoteMeta(quant) + `$`)
base = trailRe.ReplaceAllString(base, "")
}
baseKey := strings.ToLower(ggufSuffixRe.ReplaceAllString(base, ""))
idKey := baseKey
if quant != "" && !hasQuantPart(baseKey, quant) {
idKey = fmt.Sprintf("%s-%s", baseKey, strings.ToLower(quant))
}
rows = append(rows, GgufRow{
ID: idKey,
BaseID: baseKey,
FullPath: path,
FileName: name,
Quant: quant,
// Split set: the whole set's bytes, not shard 1's. See ggufSetSizeBytes.
SizeGB: round(float64(ggufSetSizeBytes(path, fi.Size()))/gib, 2),
Publisher: pubName,
Repo: repoName,
})
return nil
})
if walkErr != nil {
return nil, fmt.Errorf("walking %s: %w", modelsRoot, walkErr)
}
// Pair each model with the MTP/DFlash draft sitting in its own dir (typically
// one model per dir). Enables --spec-type draft-mtp/draft-dflash + -md
// without hand config.
for i := range rows {
// A directory row's filepath.Dir is its PARENT, so a sidecar found there
// belongs to some other model; and none of these backends load one anyway.
if rows[i].skipsGgufPipeline() {
continue
}
dir := filepath.Dir(rows[i].FullPath)
if d, ok := draftByDir[dir]; ok {
rows[i].DraftPath = d.path
rows[i].DraftKind = d.kind
rows[i].DraftSizeGB = d.sizeGB
}
if m, ok := mmprojByDir[dir]; ok {
rows[i].MmprojPath = m.path
rows[i].MmprojSizeGB = m.sizeGB
}
if c, ok := codecByDir[dir]; ok {
rows[i].CodecPath = c.path
rows[i].CodecSizeGB = c.sizeGB
}
}
// Then let the models that got nothing borrow a drafter/projector from a
// compatible sibling: another quant of the same model, or a finetune of the
// same base. Dir-local always wins - this only fills blanks. See family.go.
// Deliberately NOT extended to the TTS codec: that one is sample-rate bound
// to its talker, and no second copy of it exists to borrow from.
inheritSidecars(rows)
return rows, nil
}
// DraftSidecarForDir scans one directory (non-recursive) for a paired MTP or
// DFlash draft gguf, same convention as DiscoverGgufModels' walk. Used by the
// config-editor API to answer "is draft-dflash/draft-mtp available for this
// model" without a full models-root walk.
func DraftSidecarForDir(dir string) (path, kind string, sizeGB float64) {
entries, err := os.ReadDir(dir)
if err != nil {
return "", "", 0
}
for _, ent := range entries {
if ent.IsDir() || !strings.EqualFold(filepath.Ext(ent.Name()), ".gguf") {
continue
}
name := ent.Name()
k := ""
switch {
case mtpFileRe.MatchString(name):
k = "mtp"
case dflashFileRe.MatchString(name):
k = "dflash"
default:
continue
}
fi, err := ent.Info()
if err != nil {
continue
}
return filepath.Join(dir, name), k, round(float64(fi.Size())/gib, 2)
}
return "", "", 0
}
// mmprojSidecar is a discovered vision projector paired to a model dir.
type mmprojSidecar struct {
path string
sizeGB float64
}
// codecSidecar is a discovered Qwen3-TTS audio codec paired to a talker's dir.
type codecSidecar struct {
path string
sizeGB float64
}
// draftSidecar is a discovered speculative-decoding draft gguf paired to a
// model dir (either an MTP nextn sidecar or a DFlash block-diffusion drafter).
type draftSidecar struct {
path string
sizeGB float64
kind string // "mtp" | "dflash"
}