|
| 1 | +package importers |
| 2 | + |
| 3 | +import ( |
| 4 | + "encoding/json" |
| 5 | + "path/filepath" |
| 6 | + "strings" |
| 7 | + |
| 8 | + "github.com/mudler/LocalAI/core/config" |
| 9 | + "github.com/mudler/LocalAI/core/gallery" |
| 10 | + "github.com/mudler/LocalAI/core/schema" |
| 11 | + hfapi "github.com/mudler/LocalAI/pkg/huggingface-api" |
| 12 | + "go.yaml.in/yaml/v2" |
| 13 | +) |
| 14 | + |
| 15 | +var _ Importer = &VibeVoiceCppImporter{} |
| 16 | + |
| 17 | +// VibeVoiceCppImporter recognises the GGUF bundle that the vibevoice.cpp |
| 18 | +// backend consumes — primary model file (vibevoice-realtime-*.gguf for TTS or |
| 19 | +// vibevoice-asr-*.gguf for ASR), a sibling tokenizer.gguf (always required), |
| 20 | +// and optional voice-*.gguf prompts for TTS voice cloning. Detection fires on |
| 21 | +// the HF repo name containing "vibevoice.cpp"/"vibevoice-cpp", or on the |
| 22 | +// presence of a vibevoice-*.gguf + tokenizer.gguf pair. preferences.backend |
| 23 | +// ="vibevoice-cpp" forces the importer regardless of artefacts. |
| 24 | +// |
| 25 | +// Role pick: defaults to TTS (the realtime model is small and the common |
| 26 | +// case). preferences.usecase="asr" routes to the ASR/diarization model. If a |
| 27 | +// repo only ships one of the two roles, that role wins automatically. |
| 28 | +// |
| 29 | +// MUST be registered ahead of VibeVoiceImporter — the older Python-backed |
| 30 | +// importer matches any repo with "vibevoice" in the name, which would |
| 31 | +// otherwise swallow the C++ bundle. |
| 32 | +type VibeVoiceCppImporter struct{} |
| 33 | + |
| 34 | +func (i *VibeVoiceCppImporter) Name() string { return "vibevoice-cpp" } |
| 35 | +func (i *VibeVoiceCppImporter) Modality() string { return "tts" } |
| 36 | +func (i *VibeVoiceCppImporter) AutoDetects() bool { return true } |
| 37 | + |
| 38 | +func (i *VibeVoiceCppImporter) Match(details Details) bool { |
| 39 | + preferencesMap := unmarshalPreferences(details.Preferences) |
| 40 | + if b, ok := preferencesMap["backend"].(string); ok && b == "vibevoice-cpp" { |
| 41 | + return true |
| 42 | + } |
| 43 | + |
| 44 | + // Repo-name signal: anything carrying "vibevoice.cpp" or "vibevoice-cpp" |
| 45 | + // — the canonical naming for the C++ port bundles. |
| 46 | + repoSignals := []string{strings.ToLower(repoNameOnly(details))} |
| 47 | + if _, repo, ok := HFOwnerRepoFromURI(details.URI); ok { |
| 48 | + repoSignals = append(repoSignals, strings.ToLower(repo)) |
| 49 | + } |
| 50 | + for _, s := range repoSignals { |
| 51 | + if strings.Contains(s, "vibevoice.cpp") || strings.Contains(s, "vibevoice-cpp") { |
| 52 | + return true |
| 53 | + } |
| 54 | + } |
| 55 | + |
| 56 | + // File-listing signal: a vibevoice-*.gguf primary + tokenizer.gguf is |
| 57 | + // only what the C++ backend ships — the Python VibeVoice fork distributes |
| 58 | + // safetensors, never GGUF. |
| 59 | + if details.HuggingFace != nil && |
| 60 | + HasFile(details.HuggingFace.Files, "tokenizer.gguf") && |
| 61 | + hasVibeVoiceGGUF(details.HuggingFace.Files) { |
| 62 | + return true |
| 63 | + } |
| 64 | + |
| 65 | + return false |
| 66 | +} |
| 67 | + |
| 68 | +func (i *VibeVoiceCppImporter) Import(details Details) (gallery.ModelConfig, error) { |
| 69 | + preferencesMap := unmarshalPreferences(details.Preferences) |
| 70 | + |
| 71 | + name, ok := preferencesMap["name"].(string) |
| 72 | + if !ok { |
| 73 | + name = filepath.Base(details.URI) |
| 74 | + } |
| 75 | + |
| 76 | + description, ok := preferencesMap["description"].(string) |
| 77 | + if !ok { |
| 78 | + description = "Imported from " + details.URI |
| 79 | + } |
| 80 | + |
| 81 | + // Quant preference — default order matches what mudler/vibevoice.cpp-models |
| 82 | + // ships today. Same comma-separated convention as whisper / llama-cpp. |
| 83 | + quants := []string{"q8_0", "q4_k", "q5_k", "q4_0"} |
| 84 | + if preferred, ok := preferencesMap["quantizations"].(string); ok && preferred != "" { |
| 85 | + quants = strings.Split(preferred, ",") |
| 86 | + } |
| 87 | + |
| 88 | + usecase := strings.ToLower(stringPref(preferencesMap, "usecase")) |
| 89 | + |
| 90 | + cfg := gallery.ModelConfig{ |
| 91 | + Name: name, |
| 92 | + Description: description, |
| 93 | + } |
| 94 | + |
| 95 | + modelConfig := config.ModelConfig{ |
| 96 | + Name: name, |
| 97 | + Description: description, |
| 98 | + Backend: "vibevoice-cpp", |
| 99 | + } |
| 100 | + |
| 101 | + // Without HF metadata we can only emit a skeleton config — the user must |
| 102 | + // edit it post-import to point at real files. Mirrors whisper's bare-URI |
| 103 | + // fallback so preference-only invocations still produce something usable. |
| 104 | + if details.HuggingFace == nil { |
| 105 | + modelConfig.PredictionOptions = schema.PredictionOptions{ |
| 106 | + BasicModelRequest: schema.BasicModelRequest{Model: filepath.Base(details.URI)}, |
| 107 | + } |
| 108 | + if usecase == "asr" { |
| 109 | + modelConfig.KnownUsecaseStrings = []string{"transcript"} |
| 110 | + modelConfig.Options = []string{"type=asr", "tokenizer=tokenizer.gguf"} |
| 111 | + } else { |
| 112 | + modelConfig.KnownUsecaseStrings = []string{"tts"} |
| 113 | + modelConfig.Options = []string{"tokenizer=tokenizer.gguf"} |
| 114 | + } |
| 115 | + data, err := yaml.Marshal(modelConfig) |
| 116 | + if err != nil { |
| 117 | + return gallery.ModelConfig{}, err |
| 118 | + } |
| 119 | + cfg.ConfigFile = string(data) |
| 120 | + return cfg, nil |
| 121 | + } |
| 122 | + |
| 123 | + files := details.HuggingFace.Files |
| 124 | + ttsFiles := filterByPrefix(files, "vibevoice-realtime-") |
| 125 | + asrFiles := filterByPrefix(files, "vibevoice-asr-") |
| 126 | + |
| 127 | + // Auto-pick role when the repo only ships one. Explicit usecase wins. |
| 128 | + role := usecase |
| 129 | + if role == "" { |
| 130 | + switch { |
| 131 | + case len(ttsFiles) > 0 && len(asrFiles) == 0: |
| 132 | + role = "tts" |
| 133 | + case len(asrFiles) > 0 && len(ttsFiles) == 0: |
| 134 | + role = "asr" |
| 135 | + default: |
| 136 | + role = "tts" // default: realtime TTS is the smaller, more common case |
| 137 | + } |
| 138 | + } |
| 139 | + |
| 140 | + // Layout under <models>/vibevoice-cpp/<name>/ — same pattern as whisper's |
| 141 | + // nesting so multiple imports of the same upstream repo (with different |
| 142 | + // quants) don't collide on disk. Options[] paths are emitted relative to |
| 143 | + // opts.ModelPath, which the backend resolves against the LocalAI models |
| 144 | + // root in govibevoicecpp.go:resolvePath. |
| 145 | + relDir := filepath.Join("vibevoice-cpp", name) |
| 146 | + |
| 147 | + var primary []hfapi.ModelFile |
| 148 | + switch role { |
| 149 | + case "asr", "transcript", "stt", "speech-to-text": |
| 150 | + primary = asrFiles |
| 151 | + modelConfig.KnownUsecaseStrings = []string{"transcript"} |
| 152 | + default: |
| 153 | + primary = ttsFiles |
| 154 | + modelConfig.KnownUsecaseStrings = []string{"tts"} |
| 155 | + } |
| 156 | + // If the requested role has no matching files, fall back to any |
| 157 | + // vibevoice-*.gguf so the import still produces something runnable. |
| 158 | + if len(primary) == 0 { |
| 159 | + primary = filterByPrefix(files, "vibevoice-") |
| 160 | + } |
| 161 | + |
| 162 | + chosen, ok := pickPreferredGGUFFile(primary, quants) |
| 163 | + if !ok { |
| 164 | + // Nothing to download. Emit the skeleton — same shape as the |
| 165 | + // no-HF-metadata branch above, just with a sensible default name. |
| 166 | + modelConfig.PredictionOptions = schema.PredictionOptions{ |
| 167 | + BasicModelRequest: schema.BasicModelRequest{Model: name + ".gguf"}, |
| 168 | + } |
| 169 | + if role == "asr" { |
| 170 | + modelConfig.Options = []string{"type=asr", "tokenizer=" + filepath.Join(relDir, "tokenizer.gguf")} |
| 171 | + } else { |
| 172 | + modelConfig.Options = []string{"tokenizer=" + filepath.Join(relDir, "tokenizer.gguf")} |
| 173 | + } |
| 174 | + data, err := yaml.Marshal(modelConfig) |
| 175 | + if err != nil { |
| 176 | + return gallery.ModelConfig{}, err |
| 177 | + } |
| 178 | + cfg.ConfigFile = string(data) |
| 179 | + return cfg, nil |
| 180 | + } |
| 181 | + |
| 182 | + modelTarget := filepath.Join(relDir, filepath.Base(chosen.Path)) |
| 183 | + cfg.Files = append(cfg.Files, gallery.File{ |
| 184 | + URI: chosen.URL, |
| 185 | + Filename: modelTarget, |
| 186 | + SHA256: chosen.SHA256, |
| 187 | + }) |
| 188 | + modelConfig.PredictionOptions = schema.PredictionOptions{ |
| 189 | + BasicModelRequest: schema.BasicModelRequest{Model: modelTarget}, |
| 190 | + } |
| 191 | + |
| 192 | + // tokenizer.gguf is mandatory — Load() rejects without it. Always pull |
| 193 | + // it when the repo provides one (every official vibevoice.cpp bundle does). |
| 194 | + options := []string{} |
| 195 | + if role == "asr" { |
| 196 | + options = append(options, "type=asr") |
| 197 | + } |
| 198 | + if tok, ok := findFile(files, "tokenizer.gguf"); ok { |
| 199 | + tokTarget := filepath.Join(relDir, "tokenizer.gguf") |
| 200 | + cfg.Files = append(cfg.Files, gallery.File{ |
| 201 | + URI: tok.URL, |
| 202 | + Filename: tokTarget, |
| 203 | + SHA256: tok.SHA256, |
| 204 | + }) |
| 205 | + options = append(options, "tokenizer="+tokTarget) |
| 206 | + } |
| 207 | + |
| 208 | + // For TTS, ship the first voice-*.gguf as a default — the backend needs |
| 209 | + // a reference voice to clone from. ASR doesn't use voice prompts. |
| 210 | + if role != "asr" { |
| 211 | + if voice, ok := pickVoicePrompt(files, stringPref(preferencesMap, "voice")); ok { |
| 212 | + voiceTarget := filepath.Join(relDir, filepath.Base(voice.Path)) |
| 213 | + cfg.Files = append(cfg.Files, gallery.File{ |
| 214 | + URI: voice.URL, |
| 215 | + Filename: voiceTarget, |
| 216 | + SHA256: voice.SHA256, |
| 217 | + }) |
| 218 | + options = append(options, "voice="+voiceTarget) |
| 219 | + } |
| 220 | + } |
| 221 | + modelConfig.Options = options |
| 222 | + |
| 223 | + data, err := yaml.Marshal(modelConfig) |
| 224 | + if err != nil { |
| 225 | + return gallery.ModelConfig{}, err |
| 226 | + } |
| 227 | + cfg.ConfigFile = string(data) |
| 228 | + return cfg, nil |
| 229 | +} |
| 230 | + |
| 231 | +// hasVibeVoiceGGUF returns true when any file matches "vibevoice-*.gguf" |
| 232 | +// (case-insensitive). Narrow on purpose — third-party GGUF mirrors that |
| 233 | +// re-pack the model under different filenames will be missed, but those |
| 234 | +// users can pass preferences.backend="vibevoice-cpp" to force the importer. |
| 235 | +func hasVibeVoiceGGUF(files []hfapi.ModelFile) bool { |
| 236 | + for _, f := range files { |
| 237 | + name := strings.ToLower(filepath.Base(f.Path)) |
| 238 | + if strings.HasPrefix(name, "vibevoice-") && strings.HasSuffix(name, ".gguf") { |
| 239 | + return true |
| 240 | + } |
| 241 | + } |
| 242 | + return false |
| 243 | +} |
| 244 | + |
| 245 | +// filterByPrefix returns every file whose basename starts with prefix and |
| 246 | +// ends in .gguf (case-insensitive on the suffix, exact on the prefix). |
| 247 | +func filterByPrefix(files []hfapi.ModelFile, prefix string) []hfapi.ModelFile { |
| 248 | + var out []hfapi.ModelFile |
| 249 | + for _, f := range files { |
| 250 | + base := filepath.Base(f.Path) |
| 251 | + if !strings.HasPrefix(base, prefix) { |
| 252 | + continue |
| 253 | + } |
| 254 | + if !strings.HasSuffix(strings.ToLower(base), ".gguf") { |
| 255 | + continue |
| 256 | + } |
| 257 | + out = append(out, f) |
| 258 | + } |
| 259 | + return out |
| 260 | +} |
| 261 | + |
| 262 | +// findFile is HasFile's lookup-returning sibling. Returns the first file |
| 263 | +// whose basename equals name (exact match), or false when none exists. |
| 264 | +func findFile(files []hfapi.ModelFile, name string) (hfapi.ModelFile, bool) { |
| 265 | + for _, f := range files { |
| 266 | + if filepath.Base(f.Path) == name { |
| 267 | + return f, true |
| 268 | + } |
| 269 | + } |
| 270 | + return hfapi.ModelFile{}, false |
| 271 | +} |
| 272 | + |
| 273 | +// pickPreferredGGUFFile mirrors pickPreferredGGMLFile but operates on .gguf |
| 274 | +// files: walks prefs in order, returns the first file whose basename contains |
| 275 | +// any preference token (case-insensitive). On no match, falls back to the |
| 276 | +// last file so a missing quant still yields a runnable import. |
| 277 | +func pickPreferredGGUFFile(files []hfapi.ModelFile, prefs []string) (hfapi.ModelFile, bool) { |
| 278 | + if len(files) == 0 { |
| 279 | + return hfapi.ModelFile{}, false |
| 280 | + } |
| 281 | + for _, pref := range prefs { |
| 282 | + lower := strings.ToLower(strings.TrimSpace(pref)) |
| 283 | + if lower == "" { |
| 284 | + continue |
| 285 | + } |
| 286 | + for _, f := range files { |
| 287 | + if strings.Contains(strings.ToLower(filepath.Base(f.Path)), lower) { |
| 288 | + return f, true |
| 289 | + } |
| 290 | + } |
| 291 | + } |
| 292 | + return files[len(files)-1], true |
| 293 | +} |
| 294 | + |
| 295 | +// pickVoicePrompt selects a voice-*.gguf to bundle with a TTS import. |
| 296 | +// Honours an explicit preferences.voice substring (e.g. "Emma" picks |
| 297 | +// voice-en-Emma.gguf); otherwise returns the first voice file in listing |
| 298 | +// order so the choice is stable across imports of the same repo. |
| 299 | +func pickVoicePrompt(files []hfapi.ModelFile, hint string) (hfapi.ModelFile, bool) { |
| 300 | + hint = strings.ToLower(strings.TrimSpace(hint)) |
| 301 | + var voices []hfapi.ModelFile |
| 302 | + for _, f := range files { |
| 303 | + base := strings.ToLower(filepath.Base(f.Path)) |
| 304 | + if strings.HasPrefix(base, "voice-") && strings.HasSuffix(base, ".gguf") { |
| 305 | + voices = append(voices, f) |
| 306 | + } |
| 307 | + } |
| 308 | + if len(voices) == 0 { |
| 309 | + return hfapi.ModelFile{}, false |
| 310 | + } |
| 311 | + if hint != "" { |
| 312 | + for _, v := range voices { |
| 313 | + if strings.Contains(strings.ToLower(filepath.Base(v.Path)), hint) { |
| 314 | + return v, true |
| 315 | + } |
| 316 | + } |
| 317 | + } |
| 318 | + return voices[0], true |
| 319 | +} |
| 320 | + |
| 321 | +// repoNameOnly extracts the repo basename (everything after the last "/") |
| 322 | +// from HF metadata or, failing that, the URI. Empty when neither is set. |
| 323 | +func repoNameOnly(details Details) string { |
| 324 | + if details.HuggingFace != nil { |
| 325 | + id := details.HuggingFace.ModelID |
| 326 | + if idx := strings.Index(id, "/"); idx >= 0 { |
| 327 | + return id[idx+1:] |
| 328 | + } |
| 329 | + return id |
| 330 | + } |
| 331 | + return "" |
| 332 | +} |
| 333 | + |
| 334 | +// unmarshalPreferences decodes details.Preferences into a generic map. Returns |
| 335 | +// an empty map (never nil) on any failure so callers can index without nil |
| 336 | +// checks. Bad JSON is silently ignored — every importer here treats |
| 337 | +// preferences as best-effort hints. |
| 338 | +func unmarshalPreferences(raw json.RawMessage) map[string]any { |
| 339 | + out := map[string]any{} |
| 340 | + b, err := raw.MarshalJSON() |
| 341 | + if err != nil || len(b) == 0 { |
| 342 | + return out |
| 343 | + } |
| 344 | + _ = json.Unmarshal(b, &out) |
| 345 | + return out |
| 346 | +} |
| 347 | + |
| 348 | +// stringPref reads a string preference by key, returning "" when missing or |
| 349 | +// of the wrong type. |
| 350 | +func stringPref(m map[string]any, key string) string { |
| 351 | + if v, ok := m[key].(string); ok { |
| 352 | + return v |
| 353 | + } |
| 354 | + return "" |
| 355 | +} |
0 commit comments