mirror of
https://github.com/mudler/LocalAI.git
synced 2026-02-03 11:13:31 -05:00
Compare commits
370 Commits
feat/trans
...
btree
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
15da375c4d | ||
|
|
ba1b8e7757 | ||
|
|
79b68fdc25 | ||
|
|
a946cb08b5 | ||
|
|
d95d4992fe | ||
|
|
e13cb8346d | ||
|
|
615c56503e | ||
|
|
79a8edd8b9 | ||
|
|
8d138dd68f | ||
|
|
2b33844562 | ||
|
|
63e6721c2f | ||
|
|
4859d809aa | ||
|
|
be027b1ccd | ||
|
|
3ecadeeb93 | ||
|
|
4af3348f91 | ||
|
|
dde08845bf | ||
|
|
76d1ba168d | ||
|
|
80605e4f66 | ||
|
|
5b99584a31 | ||
|
|
fc134b18fe | ||
|
|
c2006273c5 | ||
|
|
5343889098 | ||
|
|
c42afc56d9 | ||
|
|
53f44dac89 | ||
|
|
0468456fad | ||
|
|
df899ee26a | ||
|
|
93fe25468f | ||
|
|
238aad666e | ||
|
|
4408ed4f88 | ||
|
|
5df1f59a3c | ||
|
|
8225697139 | ||
|
|
0c0186d866 | ||
|
|
ce2f8828f9 | ||
|
|
7a8565a45e | ||
|
|
192589a17f | ||
|
|
28ab73d4a1 | ||
|
|
ed4ac0b61e | ||
|
|
e41d8b65ce | ||
|
|
c28e5b39d6 | ||
|
|
b66bd2706f | ||
|
|
fa7a9d96f8 | ||
|
|
61d972a2ef | ||
|
|
fffdbc31c6 | ||
|
|
32c0ab3a7f | ||
|
|
24ce79a67c | ||
|
|
bfa8530088 | ||
|
|
4278144dd5 | ||
|
|
79fa4d691e | ||
|
|
7a3d9ee5c1 | ||
|
|
22923d3b23 | ||
|
|
d32a459209 | ||
|
|
47b2a502dd | ||
|
|
b85f339eb4 | ||
|
|
8821865eac | ||
|
|
4b30846d57 | ||
|
|
7a35986407 | ||
|
|
ee34aa7bd5 | ||
|
|
40cf798dfe | ||
|
|
18810038f5 | ||
|
|
8fb79bc6f6 | ||
|
|
4b5ad1405f | ||
|
|
4493078cdd | ||
|
|
7f68c89cbe | ||
|
|
69adc46936 | ||
|
|
d22439918f | ||
|
|
103d4e87e5 | ||
|
|
8c5ba9e0d7 | ||
|
|
f1b713df08 | ||
|
|
f94b89c1b5 | ||
|
|
a1b056737a | ||
|
|
a22f6a499d | ||
|
|
e5bf2a9a11 | ||
|
|
05aba5a311 | ||
|
|
354bf5debb | ||
|
|
7f88abb3b1 | ||
|
|
36b3a538f8 | ||
|
|
e293b65ad9 | ||
|
|
cce185b345 | ||
|
|
03ed4382c7 | ||
|
|
1c73e10676 | ||
|
|
4ade65f959 | ||
|
|
c54f5cdf12 | ||
|
|
33c48164d7 | ||
|
|
7aed3b3bac | ||
|
|
9e349c715e | ||
|
|
639ecb59b3 | ||
|
|
bfb0794f87 | ||
|
|
05f1e9e757 | ||
|
|
1ca6f6dada | ||
|
|
bc5397bcfc | ||
|
|
f452a027a2 | ||
|
|
7bac49fb87 | ||
|
|
02300cfbd1 | ||
|
|
17c5c732c7 | ||
|
|
10a66938f9 | ||
|
|
f0245fa36c | ||
|
|
83534f8e00 | ||
|
|
75eaf8c853 | ||
|
|
03096154d4 | ||
|
|
22c9e8c09e | ||
|
|
da16727ad6 | ||
|
|
ad44df6d83 | ||
|
|
276c552583 | ||
|
|
9109e5c149 | ||
|
|
71a84b91e3 | ||
|
|
209d40be71 | ||
|
|
bfd76805e8 | ||
|
|
561aa5e443 | ||
|
|
b0eb1ab2a1 | ||
|
|
1208fb6fa1 | ||
|
|
f98fe85c42 | ||
|
|
167c183c84 | ||
|
|
244e47e1e0 | ||
|
|
9680a0b0fe | ||
|
|
acbd10a661 | ||
|
|
c6b989be13 | ||
|
|
670103705c | ||
|
|
cb90bd226e | ||
|
|
df9b2abf84 | ||
|
|
582114bda9 | ||
|
|
91ffe5ac38 | ||
|
|
8a58d76254 | ||
|
|
c3442fe574 | ||
|
|
1087bd217e | ||
|
|
7ed3666d2e | ||
|
|
2e2e89e499 | ||
|
|
13c9c20f42 | ||
|
|
b3d3988d85 | ||
|
|
0529c7d0a0 | ||
|
|
af31a77061 | ||
|
|
2d8956167f | ||
|
|
509f85f82c | ||
|
|
bb2b377b18 | ||
|
|
48917889ce | ||
|
|
ef754259b0 | ||
|
|
7e26f28113 | ||
|
|
d7c8129549 | ||
|
|
3a8fbb698e | ||
|
|
b1ef34ef9f | ||
|
|
b7822250fe | ||
|
|
05055f7e95 | ||
|
|
c856d7dc73 | ||
|
|
69d565e55d | ||
|
|
fa6bbd9fa2 | ||
|
|
3f767121d2 | ||
|
|
e963e16bc5 | ||
|
|
1e9b115251 | ||
|
|
cd1e1124ea | ||
|
|
81b31b4283 | ||
|
|
d763bce46d | ||
|
|
4aac0ef42e | ||
|
|
7a36e8d967 | ||
|
|
dc2be93412 | ||
|
|
69a2b91495 | ||
|
|
791bc769c1 | ||
|
|
a15a1f07e3 | ||
|
|
c6f0b44228 | ||
|
|
cb0ed55d89 | ||
|
|
2fe97110fd | ||
|
|
fa8037b21d | ||
|
|
99a72a4b11 | ||
|
|
1a52ce1bd4 | ||
|
|
925d752f8d | ||
|
|
c0b9d00f35 | ||
|
|
fcf8d41a00 | ||
|
|
27c4161401 | ||
|
|
459b6ab86d | ||
|
|
336257cc3c | ||
|
|
df46a438b8 | ||
|
|
5e1d809904 | ||
|
|
a9c7ce7275 | ||
|
|
8c47c8c8ed | ||
|
|
8e8d427549 | ||
|
|
ee251115f4 | ||
|
|
661e66090c | ||
|
|
c38564e22c | ||
|
|
20f1e842b3 | ||
|
|
aa8965b634 | ||
|
|
35c676188b | ||
|
|
183559bb98 | ||
|
|
1123a5c49c | ||
|
|
6f17c260a7 | ||
|
|
da6278aae9 | ||
|
|
2e51871ad5 | ||
|
|
8067d25710 | ||
|
|
cb2df6c5bf | ||
|
|
07e1519b3f | ||
|
|
8fc41673fa | ||
|
|
fff0e5911b | ||
|
|
09346bdc06 | ||
|
|
d4d42740c8 | ||
|
|
5de7a43319 | ||
|
|
85e27ec74c | ||
|
|
698205a2f3 | ||
|
|
3ed582b091 | ||
|
|
752e33f676 | ||
|
|
930553ef60 | ||
|
|
fc8d5c9198 | ||
|
|
60b6472fa0 | ||
|
|
6b2c8277c2 | ||
|
|
6d5d3ebcf6 | ||
|
|
530c174fd3 | ||
|
|
8fb95686af | ||
|
|
4132085c01 | ||
|
|
c14f1ffcfd | ||
|
|
07cca4b69a | ||
|
|
dd927c36f6 | ||
|
|
052f42e926 | ||
|
|
30d43588ab | ||
|
|
d21ec22f74 | ||
|
|
04fecd634a | ||
|
|
33c14198db | ||
|
|
967c2727e3 | ||
|
|
f41f30ad92 | ||
|
|
e77340e8a5 | ||
|
|
d51a3090f7 | ||
|
|
1bf3bc932c | ||
|
|
564a47da4e | ||
|
|
c37ee93ff2 | ||
|
|
f4b65db4e7 | ||
|
|
f5fa8e6649 | ||
|
|
570e39bdcf | ||
|
|
2ebe37b671 | ||
|
|
dca685f784 | ||
|
|
84ebf2a2c9 | ||
|
|
ce5662ba90 | ||
|
|
9878f27813 | ||
|
|
f2b9452ec4 | ||
|
|
585da99c52 | ||
|
|
fd4f432079 | ||
|
|
238c68c57b | ||
|
|
04fbf5cb82 | ||
|
|
c85d559919 | ||
|
|
b5efc4f89e | ||
|
|
3f9c09a4c5 | ||
|
|
4a84660475 | ||
|
|
737248256e | ||
|
|
0ae334fc62 | ||
|
|
36c373b7c9 | ||
|
|
6afcb932b7 | ||
|
|
357bf571a3 | ||
|
|
e74ade9ebb | ||
|
|
f7f26b8efa | ||
|
|
75eb98f8bd | ||
|
|
c337e7baf7 | ||
|
|
660bd45be8 | ||
|
|
c27da0a0f6 | ||
|
|
ac043ed9ba | ||
|
|
2e0d66a1c8 | ||
|
|
41a0f361eb | ||
|
|
d3c5c02837 | ||
|
|
ae3d8fb0c4 | ||
|
|
902e47f0b0 | ||
|
|
50bb78fd24 | ||
|
|
542f07ab2d | ||
|
|
77c5acb9db | ||
|
|
44bbf4d778 | ||
|
|
633c12f93d | ||
|
|
6f24135f1d | ||
|
|
b72aa7b4fa | ||
|
|
e94e725479 | ||
|
|
e4ac7b14a3 | ||
|
|
ddb39c73f2 | ||
|
|
264b09fb1e | ||
|
|
36dd45df51 | ||
|
|
e5599f87b8 | ||
|
|
e89b5cc0e3 | ||
|
|
10bf1084cc | ||
|
|
b08ae559b3 | ||
|
|
aa7cb7e18c | ||
|
|
eadd3d4e46 | ||
|
|
2a18206033 | ||
|
|
39798d734e | ||
|
|
d0e99562af | ||
|
|
6410c99bf2 | ||
|
|
55766d269b | ||
|
|
ffa0ad1eac | ||
|
|
623789a29e | ||
|
|
2b9a3d32c9 | ||
|
|
f8b71dc5d0 | ||
|
|
1d3331b5cb | ||
|
|
2c0b9c6349 | ||
|
|
3c6c976755 | ||
|
|
ebbcba342a | ||
|
|
0de75519dc | ||
|
|
37f5e4f5c1 | ||
|
|
ffa934b959 | ||
|
|
59311d8b1e | ||
|
|
d9e25af7b5 | ||
|
|
e4f8b63b40 | ||
|
|
1364ae9be6 | ||
|
|
cfd6a9150d | ||
|
|
cd352d0c5f | ||
|
|
8d47309695 | ||
|
|
5f6fc02a55 | ||
|
|
0b528458d8 | ||
|
|
caab380c5d | ||
|
|
8a3a362504 | ||
|
|
07238eb743 | ||
|
|
e905e90dd7 | ||
|
|
08432d49e5 | ||
|
|
e51e2aacb9 | ||
|
|
9c3d85fc28 | ||
|
|
007ca647a7 | ||
|
|
59af928379 | ||
|
|
dbc2bb561b | ||
|
|
c72c85dcac | ||
|
|
ef984901e6 | ||
|
|
9911ec84a3 | ||
|
|
1956681d4c | ||
|
|
326f6e5ccb | ||
|
|
302958efd6 | ||
|
|
3dc86b247d | ||
|
|
5ec724af06 | ||
|
|
1f1e156bf0 | ||
|
|
df625e366a | ||
|
|
9e6685ac9c | ||
|
|
90c818aa71 | ||
|
|
034b9b691b | ||
|
|
ba52822e5c | ||
|
|
eb30f6c090 | ||
|
|
caba098959 | ||
|
|
3c75ea1e0e | ||
|
|
c5f911812f | ||
|
|
d82922786a | ||
|
|
d9e9bb4c0e | ||
|
|
657027bec6 | ||
|
|
2f5635308d | ||
|
|
63b5338dbd | ||
|
|
3150174962 | ||
|
|
4330fdce33 | ||
|
|
fef8583144 | ||
|
|
d4d6a56a4f | ||
|
|
2900a601a0 | ||
|
|
43e0437db6 | ||
|
|
976c159fdb | ||
|
|
969922ffec | ||
|
|
739573e41b | ||
|
|
dbdf2908ad | ||
|
|
317f8641dc | ||
|
|
54ff70e451 | ||
|
|
723f01c87e | ||
|
|
79a41a5e07 | ||
|
|
d0b6aa3f7d | ||
|
|
ad99399c6e | ||
|
|
e6ebfd3ba1 | ||
|
|
ead00a28b9 | ||
|
|
9621edb4c5 | ||
|
|
7ce92f0646 | ||
|
|
6a4ab3c1e0 | ||
|
|
83b85494c1 | ||
|
|
df6a80b38d | ||
|
|
21faa4114b | ||
|
|
e35ad56602 | ||
|
|
3be8b2d8e1 | ||
|
|
900745bb4d | ||
|
|
15a7fc7e9a | ||
|
|
03dddec538 | ||
|
|
3d34386712 | ||
|
|
1b3f66018b | ||
|
|
4381e892b8 | ||
|
|
3c3f477854 | ||
|
|
f8a8cf3e95 | ||
|
|
0fc88b3cdf | ||
|
|
4993df81c3 | ||
|
|
599bc88c6c | ||
|
|
1a0d06f3db | ||
|
|
5e1a8b3621 | ||
|
|
960e51e527 | ||
|
|
195aa22e77 |
@@ -6,6 +6,10 @@ models
|
||||
backends
|
||||
examples/chatbot-ui/models
|
||||
backend/go/image/stablediffusion-ggml/build/
|
||||
backend/go/*/build
|
||||
backend/go/*/.cache
|
||||
backend/go/*/sources
|
||||
backend/go/*/package
|
||||
examples/rwkv/models
|
||||
examples/**/models
|
||||
Dockerfile*
|
||||
|
||||
288
.github/gallery-agent/agent.go
vendored
Normal file
288
.github/gallery-agent/agent.go
vendored
Normal file
@@ -0,0 +1,288 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/go-skynet/LocalAI/.github/gallery-agent/hfapi"
|
||||
"github.com/mudler/cogito"
|
||||
|
||||
"github.com/mudler/cogito/structures"
|
||||
"github.com/sashabaranov/go-openai/jsonschema"
|
||||
)
|
||||
|
||||
var (
|
||||
openAIModel = os.Getenv("OPENAI_MODEL")
|
||||
openAIKey = os.Getenv("OPENAI_KEY")
|
||||
openAIBaseURL = os.Getenv("OPENAI_BASE_URL")
|
||||
galleryIndexPath = os.Getenv("GALLERY_INDEX_PATH")
|
||||
//defaultclient
|
||||
llm = cogito.NewOpenAILLM(openAIModel, openAIKey, openAIBaseURL)
|
||||
)
|
||||
|
||||
// cleanTextContent removes trailing spaces, tabs, and normalizes line endings
|
||||
// to prevent YAML linting issues like trailing spaces and multiple empty lines
|
||||
func cleanTextContent(text string) string {
|
||||
lines := strings.Split(text, "\n")
|
||||
var cleanedLines []string
|
||||
var prevEmpty bool
|
||||
for _, line := range lines {
|
||||
// Remove all trailing whitespace (spaces, tabs, etc.)
|
||||
trimmed := strings.TrimRight(line, " \t\r")
|
||||
// Avoid multiple consecutive empty lines
|
||||
if trimmed == "" {
|
||||
if !prevEmpty {
|
||||
cleanedLines = append(cleanedLines, "")
|
||||
}
|
||||
prevEmpty = true
|
||||
} else {
|
||||
cleanedLines = append(cleanedLines, trimmed)
|
||||
prevEmpty = false
|
||||
}
|
||||
}
|
||||
// Remove trailing empty lines from the result
|
||||
result := strings.Join(cleanedLines, "\n")
|
||||
return strings.TrimRight(result, "\n")
|
||||
}
|
||||
|
||||
// isModelExisting checks if a specific model ID exists in the gallery using text search
|
||||
func isModelExisting(modelID string) (bool, error) {
|
||||
indexPath := getGalleryIndexPath()
|
||||
content, err := os.ReadFile(indexPath)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("failed to read %s: %w", indexPath, err)
|
||||
}
|
||||
|
||||
contentStr := string(content)
|
||||
// Simple text search - if the model ID appears anywhere in the file, it exists
|
||||
return strings.Contains(contentStr, modelID), nil
|
||||
}
|
||||
|
||||
// filterExistingModels removes models that already exist in the gallery
|
||||
func filterExistingModels(models []ProcessedModel) ([]ProcessedModel, error) {
|
||||
var filteredModels []ProcessedModel
|
||||
for _, model := range models {
|
||||
exists, err := isModelExisting(model.ModelID)
|
||||
if err != nil {
|
||||
fmt.Printf("Error checking if model %s exists: %v, skipping\n", model.ModelID, err)
|
||||
continue
|
||||
}
|
||||
|
||||
if !exists {
|
||||
filteredModels = append(filteredModels, model)
|
||||
} else {
|
||||
fmt.Printf("Skipping existing model: %s\n", model.ModelID)
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("Filtered out %d existing models, %d new models remaining\n",
|
||||
len(models)-len(filteredModels), len(filteredModels))
|
||||
|
||||
return filteredModels, nil
|
||||
}
|
||||
|
||||
// getGalleryIndexPath returns the gallery index file path, with a default fallback
|
||||
func getGalleryIndexPath() string {
|
||||
if galleryIndexPath != "" {
|
||||
return galleryIndexPath
|
||||
}
|
||||
return "gallery/index.yaml"
|
||||
}
|
||||
|
||||
func getRealReadme(ctx context.Context, repository string) (string, error) {
|
||||
// Create a conversation fragment
|
||||
fragment := cogito.NewEmptyFragment().
|
||||
AddMessage("user",
|
||||
`Your task is to get a clear description of a large language model from huggingface by using the provided tool. I will share with you a repository that might be quantized, and as such probably not by the original model author. We need to get the real description of the model, and not the one that might be quantized. You will have to call the tool to get the readme more than once by figuring out from the quantized readme which is the base model readme. This is the repository: `+repository)
|
||||
|
||||
// Execute with tools
|
||||
result, err := cogito.ExecuteTools(llm, fragment,
|
||||
cogito.WithIterations(3),
|
||||
cogito.WithMaxAttempts(3),
|
||||
cogito.WithTools(&HFReadmeTool{client: hfapi.NewClient()}))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
result = result.AddMessage("user", "Describe the model in a clear and concise way that can be shared in a model gallery.")
|
||||
|
||||
// Get a response
|
||||
newFragment, err := llm.Ask(ctx, result)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
content := newFragment.LastMessage().Content
|
||||
return cleanTextContent(content), nil
|
||||
}
|
||||
|
||||
func selectMostInterestingModels(ctx context.Context, searchResult *SearchResult) ([]ProcessedModel, error) {
|
||||
// Create a conversation fragment
|
||||
fragment := cogito.NewEmptyFragment().
|
||||
AddMessage("user",
|
||||
`Your task is to analyze a list of AI models and select the most interesting ones for a model gallery. You will be given detailed information about multiple models including their metadata, file information, and README content.
|
||||
|
||||
Consider the following criteria when selecting models:
|
||||
1. Model popularity (download count)
|
||||
2. Model recency (last modified date)
|
||||
3. Model completeness (has preferred model file, README, etc.)
|
||||
4. Model uniqueness (not duplicates or very similar models)
|
||||
5. Model quality (based on README content and description)
|
||||
6. Model utility (practical applications)
|
||||
|
||||
You should select models that would be most valuable for users browsing a model gallery. Prioritize models that are:
|
||||
- Well-documented with clear READMEs
|
||||
- Recently updated
|
||||
- Popular (high download count)
|
||||
- Have the preferred quantization format available
|
||||
- Offer unique capabilities or are from reputable authors
|
||||
|
||||
Return your analysis and selection reasoning.`)
|
||||
|
||||
// Add the search results as context
|
||||
modelsInfo := fmt.Sprintf("Found %d models matching '%s' with quantization preference '%s':\n\n",
|
||||
searchResult.TotalModelsFound, searchResult.SearchTerm, searchResult.Quantization)
|
||||
|
||||
for i, model := range searchResult.Models {
|
||||
modelsInfo += fmt.Sprintf("Model %d:\n", i+1)
|
||||
modelsInfo += fmt.Sprintf(" ID: %s\n", model.ModelID)
|
||||
modelsInfo += fmt.Sprintf(" Author: %s\n", model.Author)
|
||||
modelsInfo += fmt.Sprintf(" Downloads: %d\n", model.Downloads)
|
||||
modelsInfo += fmt.Sprintf(" Last Modified: %s\n", model.LastModified)
|
||||
modelsInfo += fmt.Sprintf(" Files: %d files\n", len(model.Files))
|
||||
|
||||
if model.PreferredModelFile != nil {
|
||||
modelsInfo += fmt.Sprintf(" Preferred Model File: %s (%d bytes)\n",
|
||||
model.PreferredModelFile.Path, model.PreferredModelFile.Size)
|
||||
} else {
|
||||
modelsInfo += " No preferred model file found\n"
|
||||
}
|
||||
|
||||
if model.ReadmeContent != "" {
|
||||
modelsInfo += fmt.Sprintf(" README: %s\n", model.ReadmeContent)
|
||||
}
|
||||
|
||||
if model.ProcessingError != "" {
|
||||
modelsInfo += fmt.Sprintf(" Processing Error: %s\n", model.ProcessingError)
|
||||
}
|
||||
|
||||
modelsInfo += "\n"
|
||||
}
|
||||
|
||||
fragment = fragment.AddMessage("user", modelsInfo)
|
||||
|
||||
fragment = fragment.AddMessage("user", "Based on your analysis, select the top 5 most interesting models and provide a brief explanation for each selection. Also, create a filtered SearchResult with only the selected models. Return just a list of repositories IDs, you will later be asked to output it as a JSON array with the json tool.")
|
||||
|
||||
// Get a response
|
||||
newFragment, err := llm.Ask(ctx, fragment)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
fmt.Println(newFragment.LastMessage().Content)
|
||||
repositories := struct {
|
||||
Repositories []string `json:"repositories"`
|
||||
}{}
|
||||
|
||||
s := structures.Structure{
|
||||
Schema: jsonschema.Definition{
|
||||
Type: jsonschema.Object,
|
||||
AdditionalProperties: false,
|
||||
Properties: map[string]jsonschema.Definition{
|
||||
"repositories": {
|
||||
Type: jsonschema.Array,
|
||||
Items: &jsonschema.Definition{Type: jsonschema.String},
|
||||
Description: "The trending repositories IDs",
|
||||
},
|
||||
},
|
||||
Required: []string{"repositories"},
|
||||
},
|
||||
Object: &repositories,
|
||||
}
|
||||
|
||||
err = newFragment.ExtractStructure(ctx, llm, s)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
filteredModels := []ProcessedModel{}
|
||||
for _, m := range searchResult.Models {
|
||||
if slices.Contains(repositories.Repositories, m.ModelID) {
|
||||
filteredModels = append(filteredModels, m)
|
||||
}
|
||||
}
|
||||
|
||||
return filteredModels, nil
|
||||
}
|
||||
|
||||
// ModelFamily represents a YAML anchor/family
|
||||
type ModelFamily struct {
|
||||
Anchor string `json:"anchor"`
|
||||
Name string `json:"name"`
|
||||
}
|
||||
|
||||
// selectModelFamily selects the appropriate model family/anchor for a given model
|
||||
func selectModelFamily(ctx context.Context, model ProcessedModel, availableFamilies []ModelFamily) (string, error) {
|
||||
// Create a conversation fragment
|
||||
fragment := cogito.NewEmptyFragment().
|
||||
AddMessage("user",
|
||||
`Your task is to select the most appropriate model family/anchor for a given AI model. You will be provided with:
|
||||
1. Information about the model (name, description, etc.)
|
||||
2. A list of available model families/anchors
|
||||
|
||||
You need to select the family that best matches the model's architecture, capabilities, or characteristics. Consider:
|
||||
- Model architecture (e.g., Llama, Qwen, Mistral, etc.)
|
||||
- Model capabilities (e.g., vision, coding, chat, etc.)
|
||||
- Model size/type (e.g., small, medium, large)
|
||||
- Model purpose (e.g., general purpose, specialized, etc.)
|
||||
|
||||
Return the anchor name that best fits the model.`)
|
||||
|
||||
// Add model information
|
||||
modelInfo := "Model Information:\n"
|
||||
modelInfo += fmt.Sprintf(" ID: %s\n", model.ModelID)
|
||||
modelInfo += fmt.Sprintf(" Author: %s\n", model.Author)
|
||||
modelInfo += fmt.Sprintf(" Downloads: %d\n", model.Downloads)
|
||||
modelInfo += fmt.Sprintf(" Description: %s\n", model.ReadmeContentPreview)
|
||||
|
||||
fragment = fragment.AddMessage("user", modelInfo)
|
||||
|
||||
// Add available families
|
||||
familiesInfo := "Available Model Families:\n"
|
||||
for _, family := range availableFamilies {
|
||||
familiesInfo += fmt.Sprintf(" - %s (%s)\n", family.Anchor, family.Name)
|
||||
}
|
||||
|
||||
fragment = fragment.AddMessage("user", familiesInfo)
|
||||
fragment = fragment.AddMessage("user", "Select the most appropriate family anchor for this model. Return just the anchor name.")
|
||||
|
||||
// Get a response
|
||||
newFragment, err := llm.Ask(ctx, fragment)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Extract the selected family
|
||||
selectedFamily := strings.TrimSpace(newFragment.LastMessage().Content)
|
||||
|
||||
// Validate that the selected family exists in our list
|
||||
for _, family := range availableFamilies {
|
||||
if family.Anchor == selectedFamily {
|
||||
return selectedFamily, nil
|
||||
}
|
||||
}
|
||||
|
||||
// If no exact match, try to find a close match
|
||||
for _, family := range availableFamilies {
|
||||
if strings.Contains(strings.ToLower(family.Anchor), strings.ToLower(selectedFamily)) ||
|
||||
strings.Contains(strings.ToLower(selectedFamily), strings.ToLower(family.Anchor)) {
|
||||
return family.Anchor, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Default fallback
|
||||
return "llama3", nil
|
||||
}
|
||||
203
.github/gallery-agent/gallery.go
vendored
Normal file
203
.github/gallery-agent/gallery.go
vendored
Normal file
@@ -0,0 +1,203 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// generateYAMLEntry generates a YAML entry for a model using the specified anchor
|
||||
func generateYAMLEntry(model ProcessedModel, familyAnchor string) string {
|
||||
// Extract model name from ModelID
|
||||
parts := strings.Split(model.ModelID, "/")
|
||||
modelName := model.ModelID
|
||||
if len(parts) > 0 {
|
||||
modelName = strings.ToLower(parts[len(parts)-1])
|
||||
}
|
||||
// Remove common suffixes
|
||||
modelName = strings.ReplaceAll(modelName, "-gguf", "")
|
||||
modelName = strings.ReplaceAll(modelName, "-q4_k_m", "")
|
||||
modelName = strings.ReplaceAll(modelName, "-q4_k_s", "")
|
||||
modelName = strings.ReplaceAll(modelName, "-q3_k_m", "")
|
||||
modelName = strings.ReplaceAll(modelName, "-q2_k", "")
|
||||
|
||||
fileName := ""
|
||||
checksum := ""
|
||||
if model.PreferredModelFile != nil {
|
||||
fileParts := strings.Split(model.PreferredModelFile.Path, "/")
|
||||
if len(fileParts) > 0 {
|
||||
fileName = fileParts[len(fileParts)-1]
|
||||
}
|
||||
checksum = model.PreferredModelFile.SHA256
|
||||
} else {
|
||||
fileName = model.ModelID
|
||||
}
|
||||
|
||||
description := model.ReadmeContent
|
||||
if description == "" {
|
||||
description = fmt.Sprintf("AI model: %s", modelName)
|
||||
}
|
||||
|
||||
// Clean up description to prevent YAML linting issues
|
||||
description = cleanTextContent(description)
|
||||
|
||||
// Format description for YAML (indent each line and ensure no trailing spaces)
|
||||
lines := strings.Split(description, "\n")
|
||||
var formattedLines []string
|
||||
for _, line := range lines {
|
||||
if strings.TrimSpace(line) == "" {
|
||||
// Keep empty lines as empty (no indentation)
|
||||
formattedLines = append(formattedLines, "")
|
||||
} else {
|
||||
// Add indentation to non-empty lines
|
||||
formattedLines = append(formattedLines, " "+line)
|
||||
}
|
||||
}
|
||||
formattedDescription := strings.Join(formattedLines, "\n")
|
||||
// Remove any trailing spaces from the formatted description
|
||||
formattedDescription = strings.TrimRight(formattedDescription, " \t")
|
||||
yamlTemplate := ""
|
||||
if checksum != "" {
|
||||
yamlTemplate = `- !!merge <<: *%s
|
||||
name: "%s"
|
||||
urls:
|
||||
- https://huggingface.co/%s
|
||||
description: |
|
||||
%s
|
||||
overrides:
|
||||
parameters:
|
||||
model: %s
|
||||
files:
|
||||
- filename: %s
|
||||
sha256: %s
|
||||
uri: huggingface://%s/%s`
|
||||
return fmt.Sprintf(yamlTemplate,
|
||||
familyAnchor,
|
||||
modelName,
|
||||
model.ModelID,
|
||||
formattedDescription,
|
||||
fileName,
|
||||
fileName,
|
||||
checksum,
|
||||
model.ModelID,
|
||||
fileName,
|
||||
)
|
||||
} else {
|
||||
yamlTemplate = `- !!merge <<: *%s
|
||||
name: "%s"
|
||||
urls:
|
||||
- https://huggingface.co/%s
|
||||
description: |
|
||||
%s
|
||||
overrides:
|
||||
parameters:
|
||||
model: %s`
|
||||
return fmt.Sprintf(yamlTemplate,
|
||||
familyAnchor,
|
||||
modelName,
|
||||
model.ModelID,
|
||||
formattedDescription,
|
||||
fileName,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// extractModelFamilies extracts all YAML anchors from the gallery index.yaml file
|
||||
func extractModelFamilies() ([]ModelFamily, error) {
|
||||
// Read the index.yaml file
|
||||
indexPath := getGalleryIndexPath()
|
||||
content, err := os.ReadFile(indexPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read %s: %w", indexPath, err)
|
||||
}
|
||||
|
||||
lines := strings.Split(string(content), "\n")
|
||||
var families []ModelFamily
|
||||
|
||||
for _, line := range lines {
|
||||
line = strings.TrimSpace(line)
|
||||
// Look for YAML anchors (lines starting with "- &")
|
||||
if strings.HasPrefix(line, "- &") {
|
||||
// Extract the anchor name (everything after "- &")
|
||||
anchor := strings.TrimPrefix(line, "- &")
|
||||
// Remove any trailing colon or other characters
|
||||
anchor = strings.Split(anchor, ":")[0]
|
||||
anchor = strings.Split(anchor, " ")[0]
|
||||
|
||||
if anchor != "" {
|
||||
families = append(families, ModelFamily{
|
||||
Anchor: anchor,
|
||||
Name: anchor, // Use anchor as name for now
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return families, nil
|
||||
}
|
||||
|
||||
// generateYAMLForModels generates YAML entries for selected models and appends to index.yaml
|
||||
func generateYAMLForModels(ctx context.Context, models []ProcessedModel) error {
|
||||
// Extract available model families
|
||||
families, err := extractModelFamilies()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to extract model families: %w", err)
|
||||
}
|
||||
|
||||
fmt.Printf("Found %d model families: %v\n", len(families),
|
||||
func() []string {
|
||||
var names []string
|
||||
for _, f := range families {
|
||||
names = append(names, f.Anchor)
|
||||
}
|
||||
return names
|
||||
}())
|
||||
|
||||
// Generate YAML entries for each model
|
||||
var yamlEntries []string
|
||||
for _, model := range models {
|
||||
fmt.Printf("Selecting family for model: %s\n", model.ModelID)
|
||||
|
||||
// Select appropriate family for this model
|
||||
familyAnchor, err := selectModelFamily(ctx, model, families)
|
||||
if err != nil {
|
||||
fmt.Printf("Error selecting family for %s: %v, using default\n", model.ModelID, err)
|
||||
familyAnchor = "llama3" // Default fallback
|
||||
}
|
||||
|
||||
fmt.Printf("Selected family '%s' for model %s\n", familyAnchor, model.ModelID)
|
||||
|
||||
// Generate YAML entry
|
||||
yamlEntry := generateYAMLEntry(model, familyAnchor)
|
||||
yamlEntries = append(yamlEntries, yamlEntry)
|
||||
}
|
||||
|
||||
// Append to index.yaml
|
||||
if len(yamlEntries) > 0 {
|
||||
indexPath := getGalleryIndexPath()
|
||||
fmt.Printf("Appending YAML entries to %s...\n", indexPath)
|
||||
|
||||
// Read current content
|
||||
content, err := os.ReadFile(indexPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read %s: %w", indexPath, err)
|
||||
}
|
||||
|
||||
// Append new entries
|
||||
// Remove trailing whitespace from existing content and join entries without extra newlines
|
||||
existingContent := strings.TrimRight(string(content), " \t\n\r")
|
||||
yamlBlock := strings.Join(yamlEntries, "\n")
|
||||
newContent := existingContent + "\n" + yamlBlock + "\n"
|
||||
|
||||
// Write back to file
|
||||
err = os.WriteFile(indexPath, []byte(newContent), 0644)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to write %s: %w", indexPath, err)
|
||||
}
|
||||
|
||||
fmt.Printf("Successfully added %d models to %s\n", len(yamlEntries), indexPath)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
39
.github/gallery-agent/go.mod
vendored
Normal file
39
.github/gallery-agent/go.mod
vendored
Normal file
@@ -0,0 +1,39 @@
|
||||
module github.com/go-skynet/LocalAI/.github/gallery-agent
|
||||
|
||||
go 1.24.1
|
||||
|
||||
require (
|
||||
github.com/mudler/cogito v0.3.0
|
||||
github.com/onsi/ginkgo/v2 v2.25.3
|
||||
github.com/onsi/gomega v1.38.2
|
||||
github.com/sashabaranov/go-openai v1.41.2
|
||||
github.com/tmc/langchaingo v0.1.13
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
|
||||
require (
|
||||
dario.cat/mergo v1.0.1 // indirect
|
||||
github.com/Masterminds/goutils v1.1.1 // indirect
|
||||
github.com/Masterminds/semver/v3 v3.4.0 // indirect
|
||||
github.com/Masterminds/sprig/v3 v3.3.0 // indirect
|
||||
github.com/go-logr/logr v1.4.3 // indirect
|
||||
github.com/go-task/slim-sprig/v3 v3.0.0 // indirect
|
||||
github.com/google/go-cmp v0.7.0 // indirect
|
||||
github.com/google/jsonschema-go v0.3.0 // indirect
|
||||
github.com/google/pprof v0.0.0-20250403155104-27863c87afa6 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/huandu/xstrings v1.5.0 // indirect
|
||||
github.com/mitchellh/copystructure v1.2.0 // indirect
|
||||
github.com/mitchellh/reflectwalk v1.0.2 // indirect
|
||||
github.com/modelcontextprotocol/go-sdk v1.0.0 // indirect
|
||||
github.com/shopspring/decimal v1.4.0 // indirect
|
||||
github.com/spf13/cast v1.7.0 // indirect
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
|
||||
go.uber.org/automaxprocs v1.6.0 // indirect
|
||||
go.yaml.in/yaml/v3 v3.0.4 // indirect
|
||||
golang.org/x/crypto v0.41.0 // indirect
|
||||
golang.org/x/net v0.43.0 // indirect
|
||||
golang.org/x/sys v0.35.0 // indirect
|
||||
golang.org/x/text v0.28.0 // indirect
|
||||
golang.org/x/tools v0.36.0 // indirect
|
||||
)
|
||||
168
.github/gallery-agent/go.sum
vendored
Normal file
168
.github/gallery-agent/go.sum
vendored
Normal file
@@ -0,0 +1,168 @@
|
||||
dario.cat/mergo v1.0.1 h1:Ra4+bf83h2ztPIQYNP99R6m+Y7KfnARDfID+a+vLl4s=
|
||||
dario.cat/mergo v1.0.1/go.mod h1:uNxQE+84aUszobStD9th8a29P2fMDhsBdgRYvZOxGmk=
|
||||
github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 h1:L/gRVlceqvL25UVaW/CKtUDjefjrs0SPonmDGUVOYP0=
|
||||
github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
|
||||
github.com/Masterminds/goutils v1.1.1 h1:5nUrii3FMTL5diU80unEVvNevw1nH4+ZV4DSLVJLSYI=
|
||||
github.com/Masterminds/goutils v1.1.1/go.mod h1:8cTjp+g8YejhMuvIA5y2vz3BpJxksy863GQaJW2MFNU=
|
||||
github.com/Masterminds/semver/v3 v3.4.0 h1:Zog+i5UMtVoCU8oKka5P7i9q9HgrJeGzI9SA1Xbatp0=
|
||||
github.com/Masterminds/semver/v3 v3.4.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM=
|
||||
github.com/Masterminds/sprig/v3 v3.3.0 h1:mQh0Yrg1XPo6vjYXgtf5OtijNAKJRNcTdOOGZe3tPhs=
|
||||
github.com/Masterminds/sprig/v3 v3.3.0/go.mod h1:Zy1iXRYNqNLUolqCpL4uhk6SHUMAOSCzdgBfDb35Lz0=
|
||||
github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
|
||||
github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
|
||||
github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
|
||||
github.com/cenkalti/backoff/v4 v4.2.1 h1:y4OZtCnogmCPw98Zjyt5a6+QwPLGkiQsYW5oUqylYbM=
|
||||
github.com/cenkalti/backoff/v4 v4.2.1/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
|
||||
github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI=
|
||||
github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M=
|
||||
github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE=
|
||||
github.com/containerd/errdefs/pkg v0.3.0/go.mod h1:NJw6s9HwNuRhnjJhM7pylWwMyAkmCQvQ4GpJHEqRLVk=
|
||||
github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I=
|
||||
github.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo=
|
||||
github.com/containerd/platforms v0.2.1 h1:zvwtM3rz2YHPQsF2CHYM8+KtB5dvhISiXh5ZpSBQv6A=
|
||||
github.com/containerd/platforms v0.2.1/go.mod h1:XHCb+2/hzowdiut9rkudds9bE5yJ7npe7dG/wG+uFPw=
|
||||
github.com/cpuguy83/dockercfg v0.3.2 h1:DlJTyZGBDlXqUZ2Dk2Q3xHs/FtnooJJVaad2S9GKorA=
|
||||
github.com/cpuguy83/dockercfg v0.3.2/go.mod h1:sugsbF4//dDlL/i+S+rtpIWp+5h0BHJHfjj5/jFyUJc=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk=
|
||||
github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E=
|
||||
github.com/docker/docker v28.2.2+incompatible h1:CjwRSksz8Yo4+RmQ339Dp/D2tGO5JxwYeqtMOEe0LDw=
|
||||
github.com/docker/docker v28.2.2+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
|
||||
github.com/docker/go-connections v0.5.0 h1:USnMq7hx7gwdVZq1L49hLXaFtUdTADjXGp+uj1Br63c=
|
||||
github.com/docker/go-connections v0.5.0/go.mod h1:ov60Kzw0kKElRwhNs9UlUHAE/F9Fe6GLaXnqyDdmEXc=
|
||||
github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4=
|
||||
github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
|
||||
github.com/ebitengine/purego v0.8.4 h1:CF7LEKg5FFOsASUj0+QwaXf8Ht6TlFxg09+S9wz0omw=
|
||||
github.com/ebitengine/purego v0.8.4/go.mod h1:iIjxzd6CiRiOG0UyXP+V1+jWqUXVjPKLAI0mRfJZTmQ=
|
||||
github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
|
||||
github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
|
||||
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
|
||||
github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
|
||||
github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
|
||||
github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
|
||||
github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
|
||||
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
|
||||
github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY=
|
||||
github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0=
|
||||
github.com/go-task/slim-sprig/v3 v3.0.0 h1:sUs3vkvUymDpBKi3qH1YSqBQk9+9D/8M2mN1vB6EwHI=
|
||||
github.com/go-task/slim-sprig/v3 v3.0.0/go.mod h1:W848ghGpv3Qj3dhTPRyJypKRiqCdHZiAzKg9hl15HA8=
|
||||
github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
|
||||
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/jsonschema-go v0.3.0 h1:6AH2TxVNtk3IlvkkhjrtbUc4S8AvO0Xii0DxIygDg+Q=
|
||||
github.com/google/jsonschema-go v0.3.0/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE=
|
||||
github.com/google/pprof v0.0.0-20250403155104-27863c87afa6 h1:BHT72Gu3keYf3ZEu2J0b1vyeLSOYI8bm5wbJM/8yDe8=
|
||||
github.com/google/pprof v0.0.0-20250403155104-27863c87afa6/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/huandu/xstrings v1.5.0 h1:2ag3IFq9ZDANvthTwTiqSSZLjDc+BedvHPAp5tJy2TI=
|
||||
github.com/huandu/xstrings v1.5.0/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE=
|
||||
github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
|
||||
github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
|
||||
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
|
||||
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
|
||||
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
|
||||
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
|
||||
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4=
|
||||
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I=
|
||||
github.com/magiconair/properties v1.8.10 h1:s31yESBquKXCV9a/ScB3ESkOjUYYv+X0rg8SYxI99mE=
|
||||
github.com/magiconair/properties v1.8.10/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=
|
||||
github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw=
|
||||
github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s=
|
||||
github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ=
|
||||
github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw=
|
||||
github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0=
|
||||
github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo=
|
||||
github.com/moby/go-archive v0.1.0 h1:Kk/5rdW/g+H8NHdJW2gsXyZ7UnzvJNOy6VKJqueWdcQ=
|
||||
github.com/moby/go-archive v0.1.0/go.mod h1:G9B+YoujNohJmrIYFBpSd54GTUB4lt9S+xVQvsJyFuo=
|
||||
github.com/moby/patternmatcher v0.6.0 h1:GmP9lR19aU5GqSSFko+5pRqHi+Ohk1O69aFiKkVGiPk=
|
||||
github.com/moby/patternmatcher v0.6.0/go.mod h1:hDPoyOpDY7OrrMDLaYoY3hf52gNCR/YOUYxkhApJIxc=
|
||||
github.com/moby/sys/sequential v0.6.0 h1:qrx7XFUd/5DxtqcoH1h438hF5TmOvzC/lspjy7zgvCU=
|
||||
github.com/moby/sys/sequential v0.6.0/go.mod h1:uyv8EUTrca5PnDsdMGXhZe6CCe8U/UiTWd+lL+7b/Ko=
|
||||
github.com/moby/sys/user v0.4.0 h1:jhcMKit7SA80hivmFJcbB1vqmw//wU61Zdui2eQXuMs=
|
||||
github.com/moby/sys/user v0.4.0/go.mod h1:bG+tYYYJgaMtRKgEmuueC0hJEAZWwtIbZTB+85uoHjs=
|
||||
github.com/moby/sys/userns v0.1.0 h1:tVLXkFOxVu9A64/yh59slHVv9ahO9UIev4JZusOLG/g=
|
||||
github.com/moby/sys/userns v0.1.0/go.mod h1:IHUYgu/kao6N8YZlp9Cf444ySSvCmDlmzUcYfDHOl28=
|
||||
github.com/moby/term v0.5.0 h1:xt8Q1nalod/v7BqbG21f8mQPqH+xAaC9C3N3wfWbVP0=
|
||||
github.com/moby/term v0.5.0/go.mod h1:8FzsFHVUBGZdbDsJw/ot+X+d5HLUbvklYLJ9uGfcI3Y=
|
||||
github.com/modelcontextprotocol/go-sdk v1.0.0 h1:Z4MSjLi38bTgLrd/LjSmofqRqyBiVKRyQSJgw8q8V74=
|
||||
github.com/modelcontextprotocol/go-sdk v1.0.0/go.mod h1:nYtYQroQ2KQiM0/SbyEPUWQ6xs4B95gJjEalc9AQyOs=
|
||||
github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
|
||||
github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
|
||||
github.com/mudler/cogito v0.3.0 h1:NbVAO3bLkK5oGSY0xq87jlz8C9OIsLW55s+8Hfzeu9s=
|
||||
github.com/mudler/cogito v0.3.0/go.mod h1:abMwl+CUjCp87IufA2quZdZt0bbLaHHN79o17HbUKxU=
|
||||
github.com/onsi/ginkgo/v2 v2.25.3 h1:Ty8+Yi/ayDAGtk4XxmmfUy4GabvM+MegeB4cDLRi6nw=
|
||||
github.com/onsi/ginkgo/v2 v2.25.3/go.mod h1:43uiyQC4Ed2tkOzLsEYm7hnrb7UJTWHYNsuy3bG/snE=
|
||||
github.com/onsi/gomega v1.38.2 h1:eZCjf2xjZAqe+LeWvKb5weQ+NcPwX84kqJ0cZNxok2A=
|
||||
github.com/onsi/gomega v1.38.2/go.mod h1:W2MJcYxRGV63b418Ai34Ud0hEdTVXq9NW9+Sx6uXf3k=
|
||||
github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
|
||||
github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
|
||||
github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040=
|
||||
github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M=
|
||||
github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
|
||||
github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw=
|
||||
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=
|
||||
github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
|
||||
github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U=
|
||||
github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=
|
||||
github.com/rogpeppe/go-internal v1.11.0/go.mod h1:ddIwULY96R17DhadqLgMfk9H9tvdUzkipdSkR5nkCZA=
|
||||
github.com/sashabaranov/go-openai v1.41.2 h1:vfPRBZNMpnqu8ELsclWcAvF19lDNgh1t6TVfFFOPiSM=
|
||||
github.com/sashabaranov/go-openai v1.41.2/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
|
||||
github.com/shirou/gopsutil/v4 v4.25.5 h1:rtd9piuSMGeU8g1RMXjZs9y9luK5BwtnG7dZaQUJAsc=
|
||||
github.com/shirou/gopsutil/v4 v4.25.5/go.mod h1:PfybzyydfZcN+JMMjkF6Zb8Mq1A/VcogFFg7hj50W9c=
|
||||
github.com/shopspring/decimal v1.4.0 h1:bxl37RwXBklmTi0C79JfXCEBD1cqqHt0bbgBAGFp81k=
|
||||
github.com/shopspring/decimal v1.4.0/go.mod h1:gawqmDU56v4yIKSwfBSFip1HdCCXN8/+DMd9qYNcwME=
|
||||
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
|
||||
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
|
||||
github.com/spf13/cast v1.7.0 h1:ntdiHjuueXFgm5nzDRdOS4yfT43P5Fnud6DH50rz/7w=
|
||||
github.com/spf13/cast v1.7.0/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
|
||||
github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
|
||||
github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
|
||||
github.com/testcontainers/testcontainers-go v0.38.0 h1:d7uEapLcv2P8AvH8ahLqDMMxda2W9gQN1nRbHS28HBw=
|
||||
github.com/testcontainers/testcontainers-go v0.38.0/go.mod h1:C52c9MoHpWO+C4aqmgSU+hxlR5jlEayWtgYrb8Pzz1w=
|
||||
github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU=
|
||||
github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI=
|
||||
github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk=
|
||||
github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY=
|
||||
github.com/tmc/langchaingo v0.1.13 h1:rcpMWBIi2y3B90XxfE4Ao8dhCQPVDMaNPnN5cGB1CaA=
|
||||
github.com/tmc/langchaingo v0.1.13/go.mod h1:vpQ5NOIhpzxDfTZK9B6tf2GM/MoaHewPWM5KXXGh7hg=
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
|
||||
github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
|
||||
github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
|
||||
github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
|
||||
go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.51.0 h1:Xs2Ncz0gNihqu9iosIZ5SkBbWo5T8JhhLJFMQL1qmLI=
|
||||
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.51.0/go.mod h1:vy+2G/6NvVMpwGX/NyLqcC41fxepnuKHk16E6IZUcJc=
|
||||
go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8=
|
||||
go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM=
|
||||
go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA=
|
||||
go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI=
|
||||
go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE=
|
||||
go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs=
|
||||
go.uber.org/automaxprocs v1.6.0 h1:O3y2/QNTOdbF+e/dpXNNW7Rx2hZ4sTIPyybbxyNqTUs=
|
||||
go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwEFJ8r8=
|
||||
go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc=
|
||||
go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
|
||||
golang.org/x/crypto v0.41.0 h1:WKYxWedPGCTVVl5+WHSSrOBT0O8lx32+zxmHxijgXp4=
|
||||
golang.org/x/crypto v0.41.0/go.mod h1:pO5AFd7FA68rFak7rOAGVuygIISepHftHnr8dr6+sUc=
|
||||
golang.org/x/net v0.43.0 h1:lat02VYK2j4aLzMzecihNvTlJNQUq316m2Mr9rnM6YE=
|
||||
golang.org/x/net v0.43.0/go.mod h1:vhO1fvI4dGsIjh73sWfUVjj3N7CA9WkKJNQm2svM6Jg=
|
||||
golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI=
|
||||
golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
|
||||
golang.org/x/text v0.28.0 h1:rhazDwis8INMIwQ4tpjLDzUhx6RlXqZNPEM0huQojng=
|
||||
golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU=
|
||||
golang.org/x/tools v0.36.0 h1:kWS0uv/zsvHEle1LbV5LE8QujrxB3wfQyxHfhOk0Qkg=
|
||||
golang.org/x/tools v0.36.0/go.mod h1:WBDiHKJK8YgLHlcQPYQzNCkUxUypCaa5ZegCVutKm+s=
|
||||
google.golang.org/protobuf v1.36.8 h1:xHScyCOEuuwZEc6UtSOvPbAT4zRh0xcNRYekJwfqyMc=
|
||||
google.golang.org/protobuf v1.36.8/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
|
||||
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
299
.github/gallery-agent/hfapi/client.go
vendored
Normal file
299
.github/gallery-agent/hfapi/client.go
vendored
Normal file
@@ -0,0 +1,299 @@
|
||||
package hfapi
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Model represents a model from the Hugging Face API
|
||||
type Model struct {
|
||||
ModelID string `json:"modelId"`
|
||||
Author string `json:"author"`
|
||||
Downloads int `json:"downloads"`
|
||||
LastModified string `json:"lastModified"`
|
||||
PipelineTag string `json:"pipelineTag"`
|
||||
Private bool `json:"private"`
|
||||
Tags []string `json:"tags"`
|
||||
CreatedAt string `json:"createdAt"`
|
||||
UpdatedAt string `json:"updatedAt"`
|
||||
Sha string `json:"sha"`
|
||||
Config map[string]interface{} `json:"config"`
|
||||
ModelIndex string `json:"model_index"`
|
||||
LibraryName string `json:"library_name"`
|
||||
MaskToken string `json:"mask_token"`
|
||||
TokenizerClass string `json:"tokenizer_class"`
|
||||
}
|
||||
|
||||
// FileInfo represents file information from HuggingFace
|
||||
type FileInfo struct {
|
||||
Type string `json:"type"`
|
||||
Oid string `json:"oid"`
|
||||
Size int64 `json:"size"`
|
||||
Path string `json:"path"`
|
||||
LFS *LFSInfo `json:"lfs,omitempty"`
|
||||
XetHash string `json:"xetHash,omitempty"`
|
||||
}
|
||||
|
||||
// LFSInfo represents LFS (Large File Storage) information
|
||||
type LFSInfo struct {
|
||||
Oid string `json:"oid"`
|
||||
Size int64 `json:"size"`
|
||||
PointerSize int `json:"pointerSize"`
|
||||
}
|
||||
|
||||
// ModelFile represents a file in a model repository
|
||||
type ModelFile struct {
|
||||
Path string
|
||||
Size int64
|
||||
SHA256 string
|
||||
IsReadme bool
|
||||
}
|
||||
|
||||
// ModelDetails represents detailed information about a model
|
||||
type ModelDetails struct {
|
||||
ModelID string
|
||||
Author string
|
||||
Files []ModelFile
|
||||
ReadmeFile *ModelFile
|
||||
ReadmeContent string
|
||||
}
|
||||
|
||||
// SearchParams represents the parameters for searching models
|
||||
type SearchParams struct {
|
||||
Sort string `json:"sort"`
|
||||
Direction int `json:"direction"`
|
||||
Limit int `json:"limit"`
|
||||
Search string `json:"search"`
|
||||
}
|
||||
|
||||
// Client represents a Hugging Face API client
|
||||
type Client struct {
|
||||
baseURL string
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
// NewClient creates a new Hugging Face API client
|
||||
func NewClient() *Client {
|
||||
return &Client{
|
||||
baseURL: "https://huggingface.co/api/models",
|
||||
client: &http.Client{},
|
||||
}
|
||||
}
|
||||
|
||||
// SearchModels searches for models using the Hugging Face API
|
||||
func (c *Client) SearchModels(params SearchParams) ([]Model, error) {
|
||||
req, err := http.NewRequest("GET", c.baseURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
// Add query parameters
|
||||
q := req.URL.Query()
|
||||
q.Add("sort", params.Sort)
|
||||
q.Add("direction", fmt.Sprintf("%d", params.Direction))
|
||||
q.Add("limit", fmt.Sprintf("%d", params.Limit))
|
||||
q.Add("search", params.Search)
|
||||
req.URL.RawQuery = q.Encode()
|
||||
|
||||
// Make the HTTP request
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to make request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("failed to fetch models. Status code: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
// Read the response body
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response body: %w", err)
|
||||
}
|
||||
|
||||
// Parse the JSON response
|
||||
var models []Model
|
||||
if err := json.Unmarshal(body, &models); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse JSON response: %w", err)
|
||||
}
|
||||
|
||||
return models, nil
|
||||
}
|
||||
|
||||
// GetLatest fetches the latest GGUF models
|
||||
func (c *Client) GetLatest(searchTerm string, limit int) ([]Model, error) {
|
||||
params := SearchParams{
|
||||
Sort: "lastModified",
|
||||
Direction: -1,
|
||||
Limit: limit,
|
||||
Search: searchTerm,
|
||||
}
|
||||
|
||||
return c.SearchModels(params)
|
||||
}
|
||||
|
||||
// BaseURL returns the current base URL
|
||||
func (c *Client) BaseURL() string {
|
||||
return c.baseURL
|
||||
}
|
||||
|
||||
// SetBaseURL sets a new base URL (useful for testing)
|
||||
func (c *Client) SetBaseURL(url string) {
|
||||
c.baseURL = url
|
||||
}
|
||||
|
||||
// ListFiles lists all files in a HuggingFace repository
|
||||
func (c *Client) ListFiles(repoID string) ([]FileInfo, error) {
|
||||
baseURL := strings.TrimSuffix(c.baseURL, "/api/models")
|
||||
url := fmt.Sprintf("%s/api/models/%s/tree/main", baseURL, repoID)
|
||||
|
||||
req, err := http.NewRequest("GET", url, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to make request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("failed to fetch files. Status code: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response body: %w", err)
|
||||
}
|
||||
|
||||
var files []FileInfo
|
||||
if err := json.Unmarshal(body, &files); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse JSON response: %w", err)
|
||||
}
|
||||
|
||||
return files, nil
|
||||
}
|
||||
|
||||
// GetFileSHA gets the SHA256 checksum for a specific file by searching through the file list
|
||||
func (c *Client) GetFileSHA(repoID, fileName string) (string, error) {
|
||||
files, err := c.ListFiles(repoID)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to list files: %w", err)
|
||||
}
|
||||
|
||||
for _, file := range files {
|
||||
if filepath.Base(file.Path) == fileName {
|
||||
if file.LFS != nil && file.LFS.Oid != "" {
|
||||
// The LFS OID contains the SHA256 hash
|
||||
return file.LFS.Oid, nil
|
||||
}
|
||||
// If no LFS, return the regular OID
|
||||
return file.Oid, nil
|
||||
}
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("file %s not found", fileName)
|
||||
}
|
||||
|
||||
// GetModelDetails gets detailed information about a model including files and checksums
|
||||
func (c *Client) GetModelDetails(repoID string) (*ModelDetails, error) {
|
||||
files, err := c.ListFiles(repoID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to list files: %w", err)
|
||||
}
|
||||
|
||||
details := &ModelDetails{
|
||||
ModelID: repoID,
|
||||
Author: strings.Split(repoID, "/")[0],
|
||||
Files: make([]ModelFile, 0, len(files)),
|
||||
}
|
||||
|
||||
// Process each file
|
||||
for _, file := range files {
|
||||
fileName := filepath.Base(file.Path)
|
||||
isReadme := strings.Contains(strings.ToLower(fileName), "readme")
|
||||
|
||||
// Extract SHA256 from LFS or use OID
|
||||
sha256 := ""
|
||||
if file.LFS != nil && file.LFS.Oid != "" {
|
||||
sha256 = file.LFS.Oid
|
||||
} else {
|
||||
sha256 = file.Oid
|
||||
}
|
||||
|
||||
modelFile := ModelFile{
|
||||
Path: file.Path,
|
||||
Size: file.Size,
|
||||
SHA256: sha256,
|
||||
IsReadme: isReadme,
|
||||
}
|
||||
|
||||
details.Files = append(details.Files, modelFile)
|
||||
|
||||
// Set the readme file
|
||||
if isReadme && details.ReadmeFile == nil {
|
||||
details.ReadmeFile = &modelFile
|
||||
}
|
||||
}
|
||||
|
||||
return details, nil
|
||||
}
|
||||
|
||||
// GetReadmeContent gets the content of a README file
|
||||
func (c *Client) GetReadmeContent(repoID, readmePath string) (string, error) {
|
||||
baseURL := strings.TrimSuffix(c.baseURL, "/api/models")
|
||||
url := fmt.Sprintf("%s/%s/raw/main/%s", baseURL, repoID, readmePath)
|
||||
|
||||
req, err := http.NewRequest("GET", url, nil)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to make request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("failed to fetch readme content. Status code: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to read response body: %w", err)
|
||||
}
|
||||
|
||||
return string(body), nil
|
||||
}
|
||||
|
||||
// FilterFilesByQuantization filters files by quantization type
|
||||
func FilterFilesByQuantization(files []ModelFile, quantization string) []ModelFile {
|
||||
var filtered []ModelFile
|
||||
for _, file := range files {
|
||||
fileName := filepath.Base(file.Path)
|
||||
if strings.Contains(strings.ToLower(fileName), strings.ToLower(quantization)) {
|
||||
filtered = append(filtered, file)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
// FindPreferredModelFile finds the preferred model file based on quantization preferences
|
||||
func FindPreferredModelFile(files []ModelFile, preferences []string) *ModelFile {
|
||||
for _, preference := range preferences {
|
||||
for i := range files {
|
||||
fileName := filepath.Base(files[i].Path)
|
||||
if strings.Contains(strings.ToLower(fileName), strings.ToLower(preference)) {
|
||||
return &files[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
511
.github/gallery-agent/hfapi/client_test.go
vendored
Normal file
511
.github/gallery-agent/hfapi/client_test.go
vendored
Normal file
@@ -0,0 +1,511 @@
|
||||
package hfapi_test
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/go-skynet/LocalAI/.github/gallery-agent/hfapi"
|
||||
)
|
||||
|
||||
var _ = Describe("HuggingFace API Client", func() {
|
||||
var (
|
||||
client *hfapi.Client
|
||||
server *httptest.Server
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
client = hfapi.NewClient()
|
||||
})
|
||||
|
||||
AfterEach(func() {
|
||||
if server != nil {
|
||||
server.Close()
|
||||
}
|
||||
})
|
||||
|
||||
Context("when creating a new client", func() {
|
||||
It("should initialize with correct base URL", func() {
|
||||
Expect(client).ToNot(BeNil())
|
||||
Expect(client.BaseURL()).To(Equal("https://huggingface.co/api/models"))
|
||||
})
|
||||
})
|
||||
|
||||
Context("when searching for models", func() {
|
||||
BeforeEach(func() {
|
||||
// Mock response data
|
||||
mockResponse := `[
|
||||
{
|
||||
"modelId": "test-model-1",
|
||||
"author": "test-author",
|
||||
"downloads": 1000,
|
||||
"lastModified": "2024-01-01T00:00:00.000Z",
|
||||
"pipelineTag": "text-generation",
|
||||
"private": false,
|
||||
"tags": ["gguf", "llama"],
|
||||
"createdAt": "2024-01-01T00:00:00.000Z",
|
||||
"updatedAt": "2024-01-01T00:00:00.000Z",
|
||||
"sha": "abc123",
|
||||
"config": {},
|
||||
"model_index": "test-index",
|
||||
"library_name": "transformers",
|
||||
"mask_token": null,
|
||||
"tokenizer_class": "LlamaTokenizer"
|
||||
},
|
||||
{
|
||||
"modelId": "test-model-2",
|
||||
"author": "test-author-2",
|
||||
"downloads": 2000,
|
||||
"lastModified": "2024-01-02T00:00:00.000Z",
|
||||
"pipelineTag": "text-generation",
|
||||
"private": false,
|
||||
"tags": ["gguf", "mistral"],
|
||||
"createdAt": "2024-01-02T00:00:00.000Z",
|
||||
"updatedAt": "2024-01-02T00:00:00.000Z",
|
||||
"sha": "def456",
|
||||
"config": {},
|
||||
"model_index": "test-index-2",
|
||||
"library_name": "transformers",
|
||||
"mask_token": null,
|
||||
"tokenizer_class": "MistralTokenizer"
|
||||
}
|
||||
]`
|
||||
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Verify request parameters
|
||||
Expect(r.URL.Query().Get("sort")).To(Equal("lastModified"))
|
||||
Expect(r.URL.Query().Get("direction")).To(Equal("-1"))
|
||||
Expect(r.URL.Query().Get("limit")).To(Equal("30"))
|
||||
Expect(r.URL.Query().Get("search")).To(Equal("GGUF"))
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(mockResponse))
|
||||
}))
|
||||
|
||||
// Override the client's base URL to use our mock server
|
||||
client.SetBaseURL(server.URL)
|
||||
})
|
||||
|
||||
It("should successfully search for models", func() {
|
||||
params := hfapi.SearchParams{
|
||||
Sort: "lastModified",
|
||||
Direction: -1,
|
||||
Limit: 30,
|
||||
Search: "GGUF",
|
||||
}
|
||||
|
||||
models, err := client.SearchModels(params)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(models).To(HaveLen(2))
|
||||
|
||||
// Verify first model
|
||||
Expect(models[0].ModelID).To(Equal("test-model-1"))
|
||||
Expect(models[0].Author).To(Equal("test-author"))
|
||||
Expect(models[0].Downloads).To(Equal(1000))
|
||||
Expect(models[0].PipelineTag).To(Equal("text-generation"))
|
||||
Expect(models[0].Private).To(BeFalse())
|
||||
Expect(models[0].Tags).To(ContainElements("gguf", "llama"))
|
||||
|
||||
// Verify second model
|
||||
Expect(models[1].ModelID).To(Equal("test-model-2"))
|
||||
Expect(models[1].Author).To(Equal("test-author-2"))
|
||||
Expect(models[1].Downloads).To(Equal(2000))
|
||||
Expect(models[1].Tags).To(ContainElements("gguf", "mistral"))
|
||||
})
|
||||
|
||||
It("should handle empty search results", func() {
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("[]"))
|
||||
}))
|
||||
|
||||
client.SetBaseURL(server.URL)
|
||||
|
||||
params := hfapi.SearchParams{
|
||||
Sort: "lastModified",
|
||||
Direction: -1,
|
||||
Limit: 30,
|
||||
Search: "nonexistent",
|
||||
}
|
||||
|
||||
models, err := client.SearchModels(params)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(models).To(HaveLen(0))
|
||||
})
|
||||
|
||||
It("should handle HTTP errors", func() {
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
w.Write([]byte("Internal Server Error"))
|
||||
}))
|
||||
|
||||
client.SetBaseURL(server.URL)
|
||||
|
||||
params := hfapi.SearchParams{
|
||||
Sort: "lastModified",
|
||||
Direction: -1,
|
||||
Limit: 30,
|
||||
Search: "GGUF",
|
||||
}
|
||||
|
||||
models, err := client.SearchModels(params)
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("Status code: 500"))
|
||||
Expect(models).To(BeNil())
|
||||
})
|
||||
|
||||
It("should handle malformed JSON response", func() {
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("invalid json"))
|
||||
}))
|
||||
|
||||
client.SetBaseURL(server.URL)
|
||||
|
||||
params := hfapi.SearchParams{
|
||||
Sort: "lastModified",
|
||||
Direction: -1,
|
||||
Limit: 30,
|
||||
Search: "GGUF",
|
||||
}
|
||||
|
||||
models, err := client.SearchModels(params)
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("failed to parse JSON response"))
|
||||
Expect(models).To(BeNil())
|
||||
})
|
||||
})
|
||||
|
||||
Context("when getting latest GGUF models", func() {
|
||||
BeforeEach(func() {
|
||||
mockResponse := `[
|
||||
{
|
||||
"modelId": "latest-gguf-model",
|
||||
"author": "gguf-author",
|
||||
"downloads": 5000,
|
||||
"lastModified": "2024-01-03T00:00:00.000Z",
|
||||
"pipelineTag": "text-generation",
|
||||
"private": false,
|
||||
"tags": ["gguf", "latest"],
|
||||
"createdAt": "2024-01-03T00:00:00.000Z",
|
||||
"updatedAt": "2024-01-03T00:00:00.000Z",
|
||||
"sha": "latest123",
|
||||
"config": {},
|
||||
"model_index": "latest-index",
|
||||
"library_name": "transformers",
|
||||
"mask_token": null,
|
||||
"tokenizer_class": "LlamaTokenizer"
|
||||
}
|
||||
]`
|
||||
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Verify the search parameters are correct for GGUF search
|
||||
Expect(r.URL.Query().Get("search")).To(Equal("GGUF"))
|
||||
Expect(r.URL.Query().Get("sort")).To(Equal("lastModified"))
|
||||
Expect(r.URL.Query().Get("direction")).To(Equal("-1"))
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(mockResponse))
|
||||
}))
|
||||
|
||||
client.SetBaseURL(server.URL)
|
||||
})
|
||||
|
||||
It("should fetch latest GGUF models with correct parameters", func() {
|
||||
models, err := client.GetLatest("GGUF", 10)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(models).To(HaveLen(1))
|
||||
Expect(models[0].ModelID).To(Equal("latest-gguf-model"))
|
||||
Expect(models[0].Author).To(Equal("gguf-author"))
|
||||
Expect(models[0].Downloads).To(Equal(5000))
|
||||
Expect(models[0].Tags).To(ContainElements("gguf", "latest"))
|
||||
})
|
||||
|
||||
It("should use custom search term", func() {
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
Expect(r.URL.Query().Get("search")).To(Equal("custom-search"))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("[]"))
|
||||
}))
|
||||
|
||||
client.SetBaseURL(server.URL)
|
||||
|
||||
models, err := client.GetLatest("custom-search", 5)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(models).To(HaveLen(0))
|
||||
})
|
||||
})
|
||||
|
||||
Context("when handling network errors", func() {
|
||||
It("should handle connection failures gracefully", func() {
|
||||
// Use an invalid URL to simulate connection failure
|
||||
client.SetBaseURL("http://invalid-url-that-does-not-exist")
|
||||
|
||||
params := hfapi.SearchParams{
|
||||
Sort: "lastModified",
|
||||
Direction: -1,
|
||||
Limit: 30,
|
||||
Search: "GGUF",
|
||||
}
|
||||
|
||||
models, err := client.SearchModels(params)
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("failed to make request"))
|
||||
Expect(models).To(BeNil())
|
||||
})
|
||||
})
|
||||
|
||||
Context("when listing files", func() {
|
||||
BeforeEach(func() {
|
||||
mockFilesResponse := `[
|
||||
{
|
||||
"type": "file",
|
||||
"path": "model-Q4_K_M.gguf",
|
||||
"size": 1000000,
|
||||
"oid": "abc123",
|
||||
"lfs": {
|
||||
"oid": "def456789",
|
||||
"size": 1000000,
|
||||
"pointerSize": 135
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "file",
|
||||
"path": "README.md",
|
||||
"size": 5000,
|
||||
"oid": "readme123"
|
||||
},
|
||||
{
|
||||
"type": "file",
|
||||
"path": "config.json",
|
||||
"size": 1000,
|
||||
"oid": "config123"
|
||||
}
|
||||
]`
|
||||
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, "/tree/main") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(mockFilesResponse))
|
||||
} else {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}
|
||||
}))
|
||||
|
||||
client.SetBaseURL(server.URL)
|
||||
})
|
||||
|
||||
It("should list files successfully", func() {
|
||||
files, err := client.ListFiles("test/model")
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(files).To(HaveLen(3))
|
||||
|
||||
Expect(files[0].Path).To(Equal("model-Q4_K_M.gguf"))
|
||||
Expect(files[0].Size).To(Equal(int64(1000000)))
|
||||
Expect(files[0].LFS).ToNot(BeNil())
|
||||
Expect(files[0].LFS.Oid).To(Equal("def456789"))
|
||||
|
||||
Expect(files[1].Path).To(Equal("README.md"))
|
||||
Expect(files[1].Size).To(Equal(int64(5000)))
|
||||
})
|
||||
})
|
||||
|
||||
Context("when getting file SHA", func() {
|
||||
BeforeEach(func() {
|
||||
mockFileInfoResponse := `{
|
||||
"path": "model-Q4_K_M.gguf",
|
||||
"size": 1000000,
|
||||
"oid": "abc123",
|
||||
"lfs": {
|
||||
"oid": "sha256:def456",
|
||||
"size": 1000000,
|
||||
"pointer": "version https://git-lfs.github.com/spec/v1",
|
||||
"sha256": "def456789"
|
||||
}
|
||||
}`
|
||||
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, "/paths-info") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(mockFileInfoResponse))
|
||||
} else {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}
|
||||
}))
|
||||
|
||||
client.SetBaseURL(server.URL)
|
||||
})
|
||||
|
||||
It("should get file SHA successfully", func() {
|
||||
sha, err := client.GetFileSHA("test/model", "model-Q4_K_M.gguf")
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(sha).To(Equal("def456789"))
|
||||
})
|
||||
|
||||
It("should handle missing SHA gracefully", func() {
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(`{"path": "file.txt", "size": 100}`))
|
||||
}))
|
||||
|
||||
client.SetBaseURL(server.URL)
|
||||
|
||||
sha, err := client.GetFileSHA("test/model", "file.txt")
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("no SHA256 found"))
|
||||
Expect(sha).To(Equal(""))
|
||||
})
|
||||
})
|
||||
|
||||
Context("when getting model details", func() {
|
||||
BeforeEach(func() {
|
||||
mockFilesResponse := `[
|
||||
{
|
||||
"path": "model-Q4_K_M.gguf",
|
||||
"size": 1000000,
|
||||
"oid": "abc123",
|
||||
"lfs": {
|
||||
"oid": "sha256:def456",
|
||||
"size": 1000000,
|
||||
"pointer": "version https://git-lfs.github.com/spec/v1",
|
||||
"sha256": "def456789"
|
||||
}
|
||||
},
|
||||
{
|
||||
"path": "README.md",
|
||||
"size": 5000,
|
||||
"oid": "readme123"
|
||||
}
|
||||
]`
|
||||
|
||||
mockFileInfoResponse := `{
|
||||
"path": "model-Q4_K_M.gguf",
|
||||
"size": 1000000,
|
||||
"oid": "abc123",
|
||||
"lfs": {
|
||||
"oid": "sha256:def456",
|
||||
"size": 1000000,
|
||||
"pointer": "version https://git-lfs.github.com/spec/v1",
|
||||
"sha256": "def456789"
|
||||
}
|
||||
}`
|
||||
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, "/tree/main") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(mockFilesResponse))
|
||||
} else if strings.Contains(r.URL.Path, "/paths-info") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(mockFileInfoResponse))
|
||||
} else {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}
|
||||
}))
|
||||
|
||||
client.SetBaseURL(server.URL)
|
||||
})
|
||||
|
||||
It("should get model details successfully", func() {
|
||||
details, err := client.GetModelDetails("test/model")
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(details.ModelID).To(Equal("test/model"))
|
||||
Expect(details.Author).To(Equal("test"))
|
||||
Expect(details.Files).To(HaveLen(2))
|
||||
|
||||
Expect(details.ReadmeFile).ToNot(BeNil())
|
||||
Expect(details.ReadmeFile.Path).To(Equal("README.md"))
|
||||
Expect(details.ReadmeFile.IsReadme).To(BeTrue())
|
||||
})
|
||||
})
|
||||
|
||||
Context("when getting README content", func() {
|
||||
BeforeEach(func() {
|
||||
mockReadmeContent := "# Test Model\n\nThis is a test model for demonstration purposes."
|
||||
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, "/raw/main/") {
|
||||
w.Header().Set("Content-Type", "text/plain")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte(mockReadmeContent))
|
||||
} else {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}
|
||||
}))
|
||||
|
||||
client.SetBaseURL(server.URL)
|
||||
})
|
||||
|
||||
It("should get README content successfully", func() {
|
||||
content, err := client.GetReadmeContent("test/model", "README.md")
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(content).To(Equal("# Test Model\n\nThis is a test model for demonstration purposes."))
|
||||
})
|
||||
})
|
||||
|
||||
Context("when filtering files", func() {
|
||||
It("should filter files by quantization", func() {
|
||||
files := []hfapi.ModelFile{
|
||||
{Path: "model-Q4_K_M.gguf"},
|
||||
{Path: "model-Q3_K_M.gguf"},
|
||||
{Path: "README.md", IsReadme: true},
|
||||
}
|
||||
|
||||
filtered := hfapi.FilterFilesByQuantization(files, "Q4_K_M")
|
||||
|
||||
Expect(filtered).To(HaveLen(1))
|
||||
Expect(filtered[0].Path).To(Equal("model-Q4_K_M.gguf"))
|
||||
})
|
||||
|
||||
It("should find preferred model file", func() {
|
||||
files := []hfapi.ModelFile{
|
||||
{Path: "model-Q3_K_M.gguf"},
|
||||
{Path: "model-Q4_K_M.gguf"},
|
||||
{Path: "README.md", IsReadme: true},
|
||||
}
|
||||
|
||||
preferences := []string{"Q4_K_M", "Q3_K_M"}
|
||||
preferred := hfapi.FindPreferredModelFile(files, preferences)
|
||||
|
||||
Expect(preferred).ToNot(BeNil())
|
||||
Expect(preferred.Path).To(Equal("model-Q4_K_M.gguf"))
|
||||
Expect(preferred.IsReadme).To(BeFalse())
|
||||
})
|
||||
|
||||
It("should return nil if no preferred file found", func() {
|
||||
files := []hfapi.ModelFile{
|
||||
{Path: "model-Q2_K.gguf"},
|
||||
{Path: "README.md", IsReadme: true},
|
||||
}
|
||||
|
||||
preferences := []string{"Q4_K_M", "Q3_K_M"}
|
||||
preferred := hfapi.FindPreferredModelFile(files, preferences)
|
||||
|
||||
Expect(preferred).To(BeNil())
|
||||
})
|
||||
})
|
||||
})
|
||||
13
.github/gallery-agent/hfapi/hfapi_suite_test.go
vendored
Normal file
13
.github/gallery-agent/hfapi/hfapi_suite_test.go
vendored
Normal file
@@ -0,0 +1,13 @@
|
||||
package hfapi_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
func TestHfapi(t *testing.T) {
|
||||
RegisterFailHandler(Fail)
|
||||
RunSpecs(t, "HuggingFace API Suite")
|
||||
}
|
||||
351
.github/gallery-agent/main.go
vendored
Normal file
351
.github/gallery-agent/main.go
vendored
Normal file
@@ -0,0 +1,351 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/go-skynet/LocalAI/.github/gallery-agent/hfapi"
|
||||
)
|
||||
|
||||
// ProcessedModelFile represents a processed model file with additional metadata
|
||||
type ProcessedModelFile struct {
|
||||
Path string `json:"path"`
|
||||
Size int64 `json:"size"`
|
||||
SHA256 string `json:"sha256"`
|
||||
IsReadme bool `json:"is_readme"`
|
||||
FileType string `json:"file_type"` // "model", "readme", "other"
|
||||
}
|
||||
|
||||
// ProcessedModel represents a processed model with all gathered metadata
|
||||
type ProcessedModel struct {
|
||||
ModelID string `json:"model_id"`
|
||||
Author string `json:"author"`
|
||||
Downloads int `json:"downloads"`
|
||||
LastModified string `json:"last_modified"`
|
||||
Files []ProcessedModelFile `json:"files"`
|
||||
PreferredModelFile *ProcessedModelFile `json:"preferred_model_file,omitempty"`
|
||||
ReadmeFile *ProcessedModelFile `json:"readme_file,omitempty"`
|
||||
ReadmeContent string `json:"readme_content,omitempty"`
|
||||
ReadmeContentPreview string `json:"readme_content_preview,omitempty"`
|
||||
QuantizationPreferences []string `json:"quantization_preferences"`
|
||||
ProcessingError string `json:"processing_error,omitempty"`
|
||||
}
|
||||
|
||||
// SearchResult represents the complete result of searching and processing models
|
||||
type SearchResult struct {
|
||||
SearchTerm string `json:"search_term"`
|
||||
Limit int `json:"limit"`
|
||||
Quantization string `json:"quantization"`
|
||||
TotalModelsFound int `json:"total_models_found"`
|
||||
Models []ProcessedModel `json:"models"`
|
||||
FormattedOutput string `json:"formatted_output"`
|
||||
}
|
||||
|
||||
// AddedModelSummary represents a summary of models added to the gallery
|
||||
type AddedModelSummary struct {
|
||||
SearchTerm string `json:"search_term"`
|
||||
TotalFound int `json:"total_found"`
|
||||
ModelsAdded int `json:"models_added"`
|
||||
AddedModelIDs []string `json:"added_model_ids"`
|
||||
AddedModelURLs []string `json:"added_model_urls"`
|
||||
Quantization string `json:"quantization"`
|
||||
ProcessingTime string `json:"processing_time"`
|
||||
}
|
||||
|
||||
func main() {
|
||||
startTime := time.Now()
|
||||
|
||||
// Check for synthetic mode
|
||||
syntheticMode := os.Getenv("SYNTHETIC_MODE")
|
||||
if syntheticMode == "true" || syntheticMode == "1" {
|
||||
fmt.Println("Running in SYNTHETIC MODE - generating random test data")
|
||||
err := runSyntheticMode()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error in synthetic mode: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Get configuration from environment variables
|
||||
searchTerm := os.Getenv("SEARCH_TERM")
|
||||
if searchTerm == "" {
|
||||
searchTerm = "GGUF"
|
||||
}
|
||||
|
||||
limitStr := os.Getenv("LIMIT")
|
||||
if limitStr == "" {
|
||||
limitStr = "5"
|
||||
}
|
||||
limit, err := strconv.Atoi(limitStr)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error parsing LIMIT: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
quantization := os.Getenv("QUANTIZATION")
|
||||
|
||||
maxModels := os.Getenv("MAX_MODELS")
|
||||
if maxModels == "" {
|
||||
maxModels = "1"
|
||||
}
|
||||
maxModelsInt, err := strconv.Atoi(maxModels)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error parsing MAX_MODELS: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Print configuration
|
||||
fmt.Printf("Gallery Agent Configuration:\n")
|
||||
fmt.Printf(" Search Term: %s\n", searchTerm)
|
||||
fmt.Printf(" Limit: %d\n", limit)
|
||||
fmt.Printf(" Quantization: %s\n", quantization)
|
||||
fmt.Printf(" Max Models to Add: %d\n", maxModelsInt)
|
||||
fmt.Printf(" Gallery Index Path: %s\n", os.Getenv("GALLERY_INDEX_PATH"))
|
||||
fmt.Println()
|
||||
|
||||
result, err := searchAndProcessModels(searchTerm, limit, quantization)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
fmt.Println(result.FormattedOutput)
|
||||
|
||||
// Use AI agent to select the most interesting models
|
||||
fmt.Println("Using AI agent to select the most interesting models...")
|
||||
models, err := selectMostInterestingModels(context.Background(), result)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error in model selection: %v\n", err)
|
||||
// Continue with original result if selection fails
|
||||
models = result.Models
|
||||
}
|
||||
|
||||
fmt.Print(models)
|
||||
|
||||
// Filter out models that already exist in the gallery
|
||||
fmt.Println("Filtering out existing models...")
|
||||
models, err = filterExistingModels(models)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error filtering existing models: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Limit to maxModelsInt after filtering
|
||||
if len(models) > maxModelsInt {
|
||||
models = models[:maxModelsInt]
|
||||
}
|
||||
|
||||
// Track added models for summary
|
||||
var addedModelIDs []string
|
||||
var addedModelURLs []string
|
||||
|
||||
// Generate YAML entries and append to gallery/index.yaml
|
||||
if len(models) > 0 {
|
||||
for _, model := range models {
|
||||
addedModelIDs = append(addedModelIDs, model.ModelID)
|
||||
// Generate Hugging Face URL for the model
|
||||
modelURL := fmt.Sprintf("https://huggingface.co/%s", model.ModelID)
|
||||
addedModelURLs = append(addedModelURLs, modelURL)
|
||||
}
|
||||
fmt.Println("Generating YAML entries for selected models...")
|
||||
err = generateYAMLForModels(context.Background(), models)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error generating YAML entries: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
} else {
|
||||
fmt.Println("No new models to add to the gallery.")
|
||||
}
|
||||
|
||||
// Create and write summary
|
||||
processingTime := time.Since(startTime).String()
|
||||
summary := AddedModelSummary{
|
||||
SearchTerm: searchTerm,
|
||||
TotalFound: result.TotalModelsFound,
|
||||
ModelsAdded: len(addedModelIDs),
|
||||
AddedModelIDs: addedModelIDs,
|
||||
AddedModelURLs: addedModelURLs,
|
||||
Quantization: quantization,
|
||||
ProcessingTime: processingTime,
|
||||
}
|
||||
|
||||
// Write summary to file
|
||||
summaryData, err := json.MarshalIndent(summary, "", " ")
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error marshaling summary: %v\n", err)
|
||||
} else {
|
||||
err = os.WriteFile("gallery-agent-summary.json", summaryData, 0644)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error writing summary file: %v\n", err)
|
||||
} else {
|
||||
fmt.Printf("Summary written to gallery-agent-summary.json\n")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func searchAndProcessModels(searchTerm string, limit int, quantization string) (*SearchResult, error) {
|
||||
client := hfapi.NewClient()
|
||||
var outputBuilder strings.Builder
|
||||
|
||||
fmt.Println("Searching for models...")
|
||||
// Initialize the result struct
|
||||
result := &SearchResult{
|
||||
SearchTerm: searchTerm,
|
||||
Limit: limit,
|
||||
Quantization: quantization,
|
||||
Models: []ProcessedModel{},
|
||||
}
|
||||
|
||||
models, err := client.GetLatest(searchTerm, limit)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to fetch models: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("Models found:", len(models))
|
||||
result.TotalModelsFound = len(models)
|
||||
|
||||
if len(models) == 0 {
|
||||
outputBuilder.WriteString("No models found.\n")
|
||||
result.FormattedOutput = outputBuilder.String()
|
||||
return result, nil
|
||||
}
|
||||
|
||||
outputBuilder.WriteString(fmt.Sprintf("Found %d models matching '%s':\n\n", len(models), searchTerm))
|
||||
|
||||
// Process each model
|
||||
for i, model := range models {
|
||||
outputBuilder.WriteString(fmt.Sprintf("%d. Processing Model: %s\n", i+1, model.ModelID))
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Author: %s\n", model.Author))
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Downloads: %d\n", model.Downloads))
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Last Modified: %s\n", model.LastModified))
|
||||
|
||||
// Initialize processed model struct
|
||||
processedModel := ProcessedModel{
|
||||
ModelID: model.ModelID,
|
||||
Author: model.Author,
|
||||
Downloads: model.Downloads,
|
||||
LastModified: model.LastModified,
|
||||
QuantizationPreferences: []string{quantization, "Q4_K_M", "Q4_K_S", "Q3_K_M", "Q2_K"},
|
||||
}
|
||||
|
||||
// Get detailed model information
|
||||
details, err := client.GetModelDetails(model.ModelID)
|
||||
if err != nil {
|
||||
errorMsg := fmt.Sprintf(" Error getting model details: %v\n", err)
|
||||
outputBuilder.WriteString(errorMsg)
|
||||
processedModel.ProcessingError = err.Error()
|
||||
result.Models = append(result.Models, processedModel)
|
||||
continue
|
||||
}
|
||||
|
||||
// Define quantization preferences (in order of preference)
|
||||
quantizationPreferences := []string{quantization, "Q4_K_M", "Q4_K_S", "Q3_K_M", "Q2_K"}
|
||||
|
||||
// Find preferred model file
|
||||
preferredModelFile := hfapi.FindPreferredModelFile(details.Files, quantizationPreferences)
|
||||
|
||||
// Process files
|
||||
processedFiles := make([]ProcessedModelFile, len(details.Files))
|
||||
for j, file := range details.Files {
|
||||
fileType := "other"
|
||||
if file.IsReadme {
|
||||
fileType = "readme"
|
||||
} else if preferredModelFile != nil && file.Path == preferredModelFile.Path {
|
||||
fileType = "model"
|
||||
}
|
||||
|
||||
processedFiles[j] = ProcessedModelFile{
|
||||
Path: file.Path,
|
||||
Size: file.Size,
|
||||
SHA256: file.SHA256,
|
||||
IsReadme: file.IsReadme,
|
||||
FileType: fileType,
|
||||
}
|
||||
}
|
||||
|
||||
processedModel.Files = processedFiles
|
||||
|
||||
// Set preferred model file
|
||||
if preferredModelFile != nil {
|
||||
for _, file := range processedFiles {
|
||||
if file.Path == preferredModelFile.Path {
|
||||
processedModel.PreferredModelFile = &file
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Print file information
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Files found: %d\n", len(details.Files)))
|
||||
|
||||
if preferredModelFile != nil {
|
||||
outputBuilder.WriteString(fmt.Sprintf(" Preferred Model File: %s (SHA256: %s)\n",
|
||||
preferredModelFile.Path,
|
||||
preferredModelFile.SHA256))
|
||||
} else {
|
||||
outputBuilder.WriteString(fmt.Sprintf(" No model file found with quantization preferences: %v\n", quantizationPreferences))
|
||||
}
|
||||
|
||||
if details.ReadmeFile != nil {
|
||||
outputBuilder.WriteString(fmt.Sprintf(" README File: %s\n", details.ReadmeFile.Path))
|
||||
|
||||
// Find and set readme file
|
||||
for _, file := range processedFiles {
|
||||
if file.IsReadme {
|
||||
processedModel.ReadmeFile = &file
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Println("Getting real readme for", model.ModelID, "waiting...")
|
||||
// Use agent to get the real readme and prepare the model description
|
||||
readmeContent, err := getRealReadme(context.Background(), model.ModelID)
|
||||
if err == nil {
|
||||
processedModel.ReadmeContent = readmeContent
|
||||
processedModel.ReadmeContentPreview = truncateString(readmeContent, 200)
|
||||
outputBuilder.WriteString(fmt.Sprintf(" README Content Preview: %s\n",
|
||||
processedModel.ReadmeContentPreview))
|
||||
} else {
|
||||
continue
|
||||
}
|
||||
fmt.Println("Real readme got", readmeContent)
|
||||
// Get README content
|
||||
// readmeContent, err := client.GetReadmeContent(model.ModelID, details.ReadmeFile.Path)
|
||||
// if err == nil {
|
||||
// processedModel.ReadmeContent = readmeContent
|
||||
// processedModel.ReadmeContentPreview = truncateString(readmeContent, 200)
|
||||
// outputBuilder.WriteString(fmt.Sprintf(" README Content Preview: %s\n",
|
||||
// processedModel.ReadmeContentPreview))
|
||||
// }
|
||||
}
|
||||
|
||||
// Print all files with their checksums
|
||||
outputBuilder.WriteString(" All Files:\n")
|
||||
for _, file := range processedFiles {
|
||||
outputBuilder.WriteString(fmt.Sprintf(" - %s (%s, %d bytes", file.Path, file.FileType, file.Size))
|
||||
if file.SHA256 != "" {
|
||||
outputBuilder.WriteString(fmt.Sprintf(", SHA256: %s", file.SHA256))
|
||||
}
|
||||
outputBuilder.WriteString(")\n")
|
||||
}
|
||||
|
||||
outputBuilder.WriteString("\n")
|
||||
result.Models = append(result.Models, processedModel)
|
||||
}
|
||||
|
||||
result.FormattedOutput = outputBuilder.String()
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func truncateString(s string, maxLen int) string {
|
||||
if len(s) <= maxLen {
|
||||
return s
|
||||
}
|
||||
return s[:maxLen] + "..."
|
||||
}
|
||||
190
.github/gallery-agent/testing.go
vendored
Normal file
190
.github/gallery-agent/testing.go
vendored
Normal file
@@ -0,0 +1,190 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// runSyntheticMode generates synthetic test data and appends it to the gallery
|
||||
func runSyntheticMode() error {
|
||||
generator := NewSyntheticDataGenerator()
|
||||
|
||||
// Generate a random number of synthetic models (1-3)
|
||||
numModels := generator.rand.Intn(3) + 1
|
||||
fmt.Printf("Generating %d synthetic models for testing...\n", numModels)
|
||||
|
||||
var models []ProcessedModel
|
||||
for i := 0; i < numModels; i++ {
|
||||
model := generator.GenerateProcessedModel()
|
||||
models = append(models, model)
|
||||
fmt.Printf("Generated synthetic model: %s\n", model.ModelID)
|
||||
}
|
||||
|
||||
// Generate YAML entries and append to gallery/index.yaml
|
||||
fmt.Println("Generating YAML entries for synthetic models...")
|
||||
err := generateYAMLForModels(context.Background(), models)
|
||||
if err != nil {
|
||||
return fmt.Errorf("error generating YAML entries: %w", err)
|
||||
}
|
||||
|
||||
fmt.Printf("Successfully added %d synthetic models to the gallery for testing!\n", len(models))
|
||||
return nil
|
||||
}
|
||||
|
||||
// SyntheticDataGenerator provides methods to generate synthetic test data
|
||||
type SyntheticDataGenerator struct {
|
||||
rand *rand.Rand
|
||||
}
|
||||
|
||||
// NewSyntheticDataGenerator creates a new synthetic data generator
|
||||
func NewSyntheticDataGenerator() *SyntheticDataGenerator {
|
||||
return &SyntheticDataGenerator{
|
||||
rand: rand.New(rand.NewSource(time.Now().UnixNano())),
|
||||
}
|
||||
}
|
||||
|
||||
// GenerateProcessedModelFile creates a synthetic ProcessedModelFile
|
||||
func (g *SyntheticDataGenerator) GenerateProcessedModelFile() ProcessedModelFile {
|
||||
fileTypes := []string{"model", "readme", "other"}
|
||||
fileType := fileTypes[g.rand.Intn(len(fileTypes))]
|
||||
|
||||
var path string
|
||||
var isReadme bool
|
||||
|
||||
switch fileType {
|
||||
case "model":
|
||||
path = fmt.Sprintf("model-%s.gguf", g.randomString(8))
|
||||
isReadme = false
|
||||
case "readme":
|
||||
path = "README.md"
|
||||
isReadme = true
|
||||
default:
|
||||
path = fmt.Sprintf("file-%s.txt", g.randomString(6))
|
||||
isReadme = false
|
||||
}
|
||||
|
||||
return ProcessedModelFile{
|
||||
Path: path,
|
||||
Size: int64(g.rand.Intn(1000000000) + 1000000), // 1MB to 1GB
|
||||
SHA256: g.randomSHA256(),
|
||||
IsReadme: isReadme,
|
||||
FileType: fileType,
|
||||
}
|
||||
}
|
||||
|
||||
// GenerateProcessedModel creates a synthetic ProcessedModel
|
||||
func (g *SyntheticDataGenerator) GenerateProcessedModel() ProcessedModel {
|
||||
authors := []string{"microsoft", "meta", "google", "openai", "anthropic", "mistralai", "huggingface"}
|
||||
modelNames := []string{"llama", "gpt", "claude", "mistral", "gemma", "phi", "qwen", "codellama"}
|
||||
|
||||
author := authors[g.rand.Intn(len(authors))]
|
||||
modelName := modelNames[g.rand.Intn(len(modelNames))]
|
||||
modelID := fmt.Sprintf("%s/%s-%s", author, modelName, g.randomString(6))
|
||||
|
||||
// Generate files
|
||||
numFiles := g.rand.Intn(5) + 2 // 2-6 files
|
||||
files := make([]ProcessedModelFile, numFiles)
|
||||
|
||||
// Ensure at least one model file and one readme
|
||||
hasModelFile := false
|
||||
hasReadme := false
|
||||
|
||||
for i := 0; i < numFiles; i++ {
|
||||
files[i] = g.GenerateProcessedModelFile()
|
||||
if files[i].FileType == "model" {
|
||||
hasModelFile = true
|
||||
}
|
||||
if files[i].FileType == "readme" {
|
||||
hasReadme = true
|
||||
}
|
||||
}
|
||||
|
||||
// Add required files if missing
|
||||
if !hasModelFile {
|
||||
modelFile := g.GenerateProcessedModelFile()
|
||||
modelFile.FileType = "model"
|
||||
modelFile.Path = fmt.Sprintf("%s-Q4_K_M.gguf", modelName)
|
||||
files = append(files, modelFile)
|
||||
}
|
||||
|
||||
if !hasReadme {
|
||||
readmeFile := g.GenerateProcessedModelFile()
|
||||
readmeFile.FileType = "readme"
|
||||
readmeFile.Path = "README.md"
|
||||
readmeFile.IsReadme = true
|
||||
files = append(files, readmeFile)
|
||||
}
|
||||
|
||||
// Find preferred model file
|
||||
var preferredModelFile *ProcessedModelFile
|
||||
for i := range files {
|
||||
if files[i].FileType == "model" {
|
||||
preferredModelFile = &files[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// Find readme file
|
||||
var readmeFile *ProcessedModelFile
|
||||
for i := range files {
|
||||
if files[i].FileType == "readme" {
|
||||
readmeFile = &files[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
readmeContent := g.generateReadmeContent(modelName, author)
|
||||
|
||||
return ProcessedModel{
|
||||
ModelID: modelID,
|
||||
Author: author,
|
||||
Downloads: g.rand.Intn(1000000) + 1000,
|
||||
LastModified: g.randomDate(),
|
||||
Files: files,
|
||||
PreferredModelFile: preferredModelFile,
|
||||
ReadmeFile: readmeFile,
|
||||
ReadmeContent: readmeContent,
|
||||
ReadmeContentPreview: truncateString(readmeContent, 200),
|
||||
QuantizationPreferences: []string{"Q4_K_M", "Q4_K_S", "Q3_K_M", "Q2_K"},
|
||||
ProcessingError: "",
|
||||
}
|
||||
}
|
||||
|
||||
// Helper methods for synthetic data generation
|
||||
func (g *SyntheticDataGenerator) randomString(length int) string {
|
||||
const charset = "abcdefghijklmnopqrstuvwxyz0123456789"
|
||||
b := make([]byte, length)
|
||||
for i := range b {
|
||||
b[i] = charset[g.rand.Intn(len(charset))]
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func (g *SyntheticDataGenerator) randomSHA256() string {
|
||||
const charset = "0123456789abcdef"
|
||||
b := make([]byte, 64)
|
||||
for i := range b {
|
||||
b[i] = charset[g.rand.Intn(len(charset))]
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func (g *SyntheticDataGenerator) randomDate() string {
|
||||
now := time.Now()
|
||||
daysAgo := g.rand.Intn(365) // Random date within last year
|
||||
pastDate := now.AddDate(0, 0, -daysAgo)
|
||||
return pastDate.Format("2006-01-02T15:04:05.000Z")
|
||||
}
|
||||
|
||||
func (g *SyntheticDataGenerator) generateReadmeContent(modelName, author string) string {
|
||||
templates := []string{
|
||||
fmt.Sprintf("# %s Model\n\nThis is a %s model developed by %s. It's designed for various natural language processing tasks including text generation, question answering, and conversation.\n\n## Features\n\n- High-quality text generation\n- Efficient inference\n- Multiple quantization options\n- Easy to use with LocalAI\n\n## Usage\n\nUse this model with LocalAI for various AI tasks.", strings.Title(modelName), modelName, author),
|
||||
fmt.Sprintf("# %s\n\nA powerful language model from %s. This model excels at understanding and generating human-like text across multiple domains.\n\n## Capabilities\n\n- Text completion\n- Code generation\n- Creative writing\n- Technical documentation\n\n## Model Details\n\n- Architecture: Transformer-based\n- Training: Large-scale supervised learning\n- Quantization: Available in multiple formats", strings.Title(modelName), author),
|
||||
fmt.Sprintf("# %s Language Model\n\nDeveloped by %s, this model represents state-of-the-art performance in natural language understanding and generation.\n\n## Key Features\n\n- Multilingual support\n- Context-aware responses\n- Efficient memory usage\n- Fast inference speed\n\n## Applications\n\n- Chatbots and virtual assistants\n- Content generation\n- Code completion\n- Educational tools", strings.Title(modelName), author),
|
||||
}
|
||||
|
||||
return templates[g.rand.Intn(len(templates))]
|
||||
}
|
||||
46
.github/gallery-agent/tools.go
vendored
Normal file
46
.github/gallery-agent/tools.go
vendored
Normal file
@@ -0,0 +1,46 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/go-skynet/LocalAI/.github/gallery-agent/hfapi"
|
||||
"github.com/sashabaranov/go-openai"
|
||||
"github.com/tmc/langchaingo/jsonschema"
|
||||
)
|
||||
|
||||
// Get repository README from HF
|
||||
type HFReadmeTool struct {
|
||||
client *hfapi.Client
|
||||
}
|
||||
|
||||
func (s *HFReadmeTool) Run(args map[string]any) (string, error) {
|
||||
q, ok := args["repository"].(string)
|
||||
if !ok {
|
||||
return "", fmt.Errorf("no query")
|
||||
}
|
||||
readme, err := s.client.GetReadmeContent(q, "README.md")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return readme, nil
|
||||
}
|
||||
|
||||
func (s *HFReadmeTool) Tool() openai.Tool {
|
||||
return openai.Tool{
|
||||
Type: openai.ToolTypeFunction,
|
||||
Function: &openai.FunctionDefinition{
|
||||
Name: "hf_readme",
|
||||
Description: "A tool to get the README content of a huggingface repository",
|
||||
Parameters: jsonschema.Definition{
|
||||
Type: jsonschema.Object,
|
||||
Properties: map[string]jsonschema.Definition{
|
||||
"repository": {
|
||||
Type: jsonschema.String,
|
||||
Description: "The huggingface repository to get the README content of",
|
||||
},
|
||||
},
|
||||
Required: []string{"repository"},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
162
.github/workflows/backend.yml
vendored
162
.github/workflows/backend.yml
vendored
@@ -111,6 +111,18 @@ jobs:
|
||||
backend: "diffusers"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-chatterbox'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:22.04"
|
||||
skip-drivers: 'true'
|
||||
backend: "chatterbox"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
# CUDA 11 additional backends
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "11"
|
||||
@@ -230,7 +242,7 @@ jobs:
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:22.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "diffusers"
|
||||
backend: "diffusers"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
# CUDA 12 additional backends
|
||||
@@ -477,6 +489,18 @@ jobs:
|
||||
backend: "diffusers"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
- build-type: 'l4t'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-l4t-kokoro'
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
base-image: "nvcr.io/nvidia/l4t-jetpack:r36.4.0"
|
||||
skip-drivers: 'true'
|
||||
backend: "kokoro"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
# SYCL additional backends
|
||||
- build-type: 'intel'
|
||||
cuda-major-version: ""
|
||||
@@ -763,7 +787,7 @@ jobs:
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-hipblas-whisper'
|
||||
tag-suffix: '-gpu-rocm-hipblas-whisper'
|
||||
base-image: "rocm/dev-ubuntu-22.04:6.4.3"
|
||||
runs-on: 'ubuntu-latest'
|
||||
skip-drivers: 'false'
|
||||
@@ -858,7 +882,7 @@ jobs:
|
||||
backend: "rfdetr"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
- build-type: 'cublas'
|
||||
- build-type: 'l4t'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
@@ -931,6 +955,18 @@ jobs:
|
||||
backend: "exllama2"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
- build-type: 'l4t'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'true'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-l4t-arm64-chatterbox'
|
||||
base-image: "nvcr.io/nvidia/l4t-jetpack:r36.4.0"
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
backend: "chatterbox"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
# runs out of space on the runner
|
||||
# - build-type: 'hipblas'
|
||||
# cuda-major-version: ""
|
||||
@@ -957,40 +993,90 @@ jobs:
|
||||
backend: "kitten-tts"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
diffusers-darwin:
|
||||
# neutts
|
||||
- build-type: ''
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64,linux/arm64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-cpu-neutts'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:22.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "neutts"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
- build-type: 'cublas'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-nvidia-cuda-12-neutts'
|
||||
runs-on: 'ubuntu-latest'
|
||||
base-image: "ubuntu:22.04"
|
||||
skip-drivers: 'false'
|
||||
backend: "neutts"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
- build-type: 'hipblas'
|
||||
cuda-major-version: ""
|
||||
cuda-minor-version: ""
|
||||
platforms: 'linux/amd64'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-gpu-rocm-hipblas-neutts'
|
||||
runs-on: 'arc-runner-set'
|
||||
base-image: "rocm/dev-ubuntu-22.04:6.4.3"
|
||||
skip-drivers: 'false'
|
||||
backend: "neutts"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
- build-type: 'l4t'
|
||||
cuda-major-version: "12"
|
||||
cuda-minor-version: "0"
|
||||
platforms: 'linux/arm64'
|
||||
skip-drivers: 'true'
|
||||
tag-latest: 'auto'
|
||||
tag-suffix: '-nvidia-l4t-arm64-neutts'
|
||||
base-image: "nvcr.io/nvidia/l4t-jetpack:r36.4.0"
|
||||
runs-on: 'ubuntu-24.04-arm'
|
||||
backend: "neutts"
|
||||
dockerfile: "./backend/Dockerfile.python"
|
||||
context: "./backend"
|
||||
backend-jobs-darwin:
|
||||
uses: ./.github/workflows/backend_build_darwin.yml
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- backend: "diffusers"
|
||||
tag-suffix: "-metal-darwin-arm64-diffusers"
|
||||
build-type: "mps"
|
||||
- backend: "mlx"
|
||||
tag-suffix: "-metal-darwin-arm64-mlx"
|
||||
build-type: "mps"
|
||||
- backend: "chatterbox"
|
||||
tag-suffix: "-metal-darwin-arm64-chatterbox"
|
||||
build-type: "mps"
|
||||
- backend: "mlx-vlm"
|
||||
tag-suffix: "-metal-darwin-arm64-mlx-vlm"
|
||||
build-type: "mps"
|
||||
- backend: "mlx-audio"
|
||||
tag-suffix: "-metal-darwin-arm64-mlx-audio"
|
||||
build-type: "mps"
|
||||
- backend: "stablediffusion-ggml"
|
||||
tag-suffix: "-metal-darwin-arm64-stablediffusion-ggml"
|
||||
build-type: "metal"
|
||||
lang: "go"
|
||||
- backend: "whisper"
|
||||
tag-suffix: "-metal-darwin-arm64-whisper"
|
||||
build-type: "metal"
|
||||
lang: "go"
|
||||
with:
|
||||
backend: "diffusers"
|
||||
build-type: "mps"
|
||||
backend: ${{ matrix.backend }}
|
||||
build-type: ${{ matrix.build-type }}
|
||||
go-version: "1.24.x"
|
||||
tag-suffix: "-metal-darwin-arm64-diffusers"
|
||||
use-pip: true
|
||||
runs-on: "macOS-14"
|
||||
secrets:
|
||||
dockerUsername: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
dockerPassword: ${{ secrets.DOCKERHUB_PASSWORD }}
|
||||
quayUsername: ${{ secrets.LOCALAI_REGISTRY_USERNAME }}
|
||||
quayPassword: ${{ secrets.LOCALAI_REGISTRY_PASSWORD }}
|
||||
mlx-darwin:
|
||||
uses: ./.github/workflows/backend_build_darwin.yml
|
||||
with:
|
||||
backend: "mlx"
|
||||
build-type: "mps"
|
||||
go-version: "1.24.x"
|
||||
tag-suffix: "-metal-darwin-arm64-mlx"
|
||||
runs-on: "macOS-14"
|
||||
secrets:
|
||||
dockerUsername: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
dockerPassword: ${{ secrets.DOCKERHUB_PASSWORD }}
|
||||
quayUsername: ${{ secrets.LOCALAI_REGISTRY_USERNAME }}
|
||||
quayPassword: ${{ secrets.LOCALAI_REGISTRY_PASSWORD }}
|
||||
mlx-vlm-darwin:
|
||||
uses: ./.github/workflows/backend_build_darwin.yml
|
||||
with:
|
||||
backend: "mlx-vlm"
|
||||
build-type: "mps"
|
||||
go-version: "1.24.x"
|
||||
tag-suffix: "-metal-darwin-arm64-mlx-vlm"
|
||||
tag-suffix: ${{ matrix.tag-suffix }}
|
||||
lang: ${{ matrix.lang || 'python' }}
|
||||
use-pip: ${{ matrix.backend == 'diffusers' }}
|
||||
runs-on: "macOS-14"
|
||||
secrets:
|
||||
dockerUsername: ${{ secrets.DOCKERHUB_USERNAME }}
|
||||
@@ -1023,7 +1109,7 @@ jobs:
|
||||
make protogen-go
|
||||
make backends/llama-cpp-darwin
|
||||
- name: Upload llama-cpp.tar
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v5
|
||||
with:
|
||||
name: llama-cpp-tar
|
||||
path: backend-images/llama-cpp.tar
|
||||
@@ -1033,7 +1119,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Download llama-cpp.tar
|
||||
uses: actions/download-artifact@v5
|
||||
uses: actions/download-artifact@v6
|
||||
with:
|
||||
name: llama-cpp-tar
|
||||
path: .
|
||||
@@ -1111,7 +1197,7 @@ jobs:
|
||||
export PLATFORMARCH=darwin/amd64
|
||||
make backends/llama-cpp-darwin
|
||||
- name: Upload llama-cpp.tar
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v5
|
||||
with:
|
||||
name: llama-cpp-tar-x86
|
||||
path: backend-images/llama-cpp.tar
|
||||
@@ -1121,7 +1207,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Download llama-cpp.tar
|
||||
uses: actions/download-artifact@v5
|
||||
uses: actions/download-artifact@v6
|
||||
with:
|
||||
name: llama-cpp-tar-x86
|
||||
path: .
|
||||
|
||||
34
.github/workflows/backend_build_darwin.yml
vendored
34
.github/workflows/backend_build_darwin.yml
vendored
@@ -16,6 +16,10 @@ on:
|
||||
description: 'Use pip to install dependencies'
|
||||
default: false
|
||||
type: boolean
|
||||
lang:
|
||||
description: 'Programming language (e.g. go)'
|
||||
default: 'python'
|
||||
type: string
|
||||
go-version:
|
||||
description: 'Go version to use'
|
||||
default: '1.24.x'
|
||||
@@ -49,28 +53,28 @@ jobs:
|
||||
uses: actions/checkout@v5
|
||||
with:
|
||||
submodules: true
|
||||
|
||||
|
||||
- name: Setup Go ${{ matrix.go-version }}
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: ${{ matrix.go-version }}
|
||||
cache: false
|
||||
|
||||
|
||||
# You can test your matrix by printing the current Go version
|
||||
- name: Display Go version
|
||||
run: go version
|
||||
|
||||
|
||||
- name: Dependencies
|
||||
run: |
|
||||
brew install protobuf grpc make protoc-gen-go protoc-gen-go-grpc libomp llvm
|
||||
|
||||
|
||||
- name: Build ${{ inputs.backend }}-darwin
|
||||
run: |
|
||||
make protogen-go
|
||||
BACKEND=${{ inputs.backend }} BUILD_TYPE=${{ inputs.build-type }} USE_PIP=${{ inputs.use-pip }} make build-darwin-python-backend
|
||||
|
||||
BACKEND=${{ inputs.backend }} BUILD_TYPE=${{ inputs.build-type }} USE_PIP=${{ inputs.use-pip }} make build-darwin-${{ inputs.lang }}-backend
|
||||
|
||||
- name: Upload ${{ inputs.backend }}.tar
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v5
|
||||
with:
|
||||
name: ${{ inputs.backend }}-tar
|
||||
path: backend-images/${{ inputs.backend }}.tar
|
||||
@@ -81,24 +85,24 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Download ${{ inputs.backend }}.tar
|
||||
uses: actions/download-artifact@v5
|
||||
uses: actions/download-artifact@v6
|
||||
with:
|
||||
name: ${{ inputs.backend }}-tar
|
||||
path: .
|
||||
|
||||
|
||||
- name: Install crane
|
||||
run: |
|
||||
curl -L https://github.com/google/go-containerregistry/releases/latest/download/go-containerregistry_Linux_x86_64.tar.gz | tar -xz
|
||||
sudo mv crane /usr/local/bin/
|
||||
|
||||
|
||||
- name: Log in to DockerHub
|
||||
run: |
|
||||
echo "${{ secrets.dockerPassword }}" | crane auth login docker.io -u "${{ secrets.dockerUsername }}" --password-stdin
|
||||
|
||||
|
||||
- name: Log in to quay.io
|
||||
run: |
|
||||
echo "${{ secrets.quayPassword }}" | crane auth login quay.io -u "${{ secrets.quayUsername }}" --password-stdin
|
||||
|
||||
|
||||
- name: Docker meta
|
||||
id: meta
|
||||
uses: docker/metadata-action@v5
|
||||
@@ -112,7 +116,7 @@ jobs:
|
||||
flavor: |
|
||||
latest=auto
|
||||
suffix=${{ inputs.tag-suffix }},onlatest=true
|
||||
|
||||
|
||||
- name: Docker meta
|
||||
id: quaymeta
|
||||
uses: docker/metadata-action@v5
|
||||
@@ -126,13 +130,13 @@ jobs:
|
||||
flavor: |
|
||||
latest=auto
|
||||
suffix=${{ inputs.tag-suffix }},onlatest=true
|
||||
|
||||
|
||||
- name: Push Docker image (DockerHub)
|
||||
run: |
|
||||
for tag in $(echo "${{ steps.meta.outputs.tags }}" | tr ',' '\n'); do
|
||||
crane push ${{ inputs.backend }}.tar $tag
|
||||
done
|
||||
|
||||
|
||||
- name: Push Docker image (Quay)
|
||||
run: |
|
||||
for tag in $(echo "${{ steps.quaymeta.outputs.tags }}" | tr ',' '\n'); do
|
||||
|
||||
20
.github/workflows/backend_pr.yml
vendored
20
.github/workflows/backend_pr.yml
vendored
@@ -12,7 +12,9 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
outputs:
|
||||
matrix: ${{ steps.set-matrix.outputs.matrix }}
|
||||
matrix-darwin: ${{ steps.set-matrix.outputs.matrix-darwin }}
|
||||
has-backends: ${{ steps.set-matrix.outputs.has-backends }}
|
||||
has-backends-darwin: ${{ steps.set-matrix.outputs.has-backends-darwin }}
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v5
|
||||
@@ -56,3 +58,21 @@ jobs:
|
||||
strategy:
|
||||
fail-fast: true
|
||||
matrix: ${{ fromJson(needs.generate-matrix.outputs.matrix) }}
|
||||
backend-jobs-darwin:
|
||||
needs: generate-matrix
|
||||
uses: ./.github/workflows/backend_build_darwin.yml
|
||||
if: needs.generate-matrix.outputs.has-backends-darwin == 'true'
|
||||
with:
|
||||
backend: ${{ matrix.backend }}
|
||||
build-type: ${{ matrix.build-type }}
|
||||
go-version: "1.24.x"
|
||||
tag-suffix: ${{ matrix.tag-suffix }}
|
||||
lang: ${{ matrix.lang || 'python' }}
|
||||
use-pip: ${{ matrix.backend == 'diffusers' }}
|
||||
runs-on: "macOS-14"
|
||||
secrets:
|
||||
quayUsername: ${{ secrets.LOCALAI_REGISTRY_USERNAME }}
|
||||
quayPassword: ${{ secrets.LOCALAI_REGISTRY_PASSWORD }}
|
||||
strategy:
|
||||
fail-fast: true
|
||||
matrix: ${{ fromJson(needs.generate-matrix.outputs.matrix-darwin) }}
|
||||
|
||||
46
.github/workflows/build-test.yaml
vendored
46
.github/workflows/build-test.yaml
vendored
@@ -17,7 +17,51 @@ jobs:
|
||||
- name: Set up Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: 1.23
|
||||
go-version: 1.25
|
||||
- name: Run GoReleaser
|
||||
run: |
|
||||
make dev-dist
|
||||
launcher-build-darwin:
|
||||
runs-on: macos-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v5
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- name: Set up Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: 1.25
|
||||
- name: Build launcher for macOS ARM64
|
||||
run: |
|
||||
make build-launcher-darwin
|
||||
ls -liah dist
|
||||
- name: Upload macOS launcher artifacts
|
||||
uses: actions/upload-artifact@v5
|
||||
with:
|
||||
name: launcher-macos
|
||||
path: dist/
|
||||
retention-days: 30
|
||||
|
||||
launcher-build-linux:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v5
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- name: Set up Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: 1.25
|
||||
- name: Build launcher for Linux
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install golang gcc libgl1-mesa-dev xorg-dev libxkbcommon-dev
|
||||
make build-launcher-linux
|
||||
- name: Upload Linux launcher artifacts
|
||||
uses: actions/upload-artifact@v5
|
||||
with:
|
||||
name: launcher-linux
|
||||
path: local-ai-launcher-linux.tar.xz
|
||||
retention-days: 30
|
||||
126
.github/workflows/gallery-agent.yaml
vendored
Normal file
126
.github/workflows/gallery-agent.yaml
vendored
Normal file
@@ -0,0 +1,126 @@
|
||||
name: Gallery Agent
|
||||
on:
|
||||
|
||||
schedule:
|
||||
- cron: '0 */1 * * *' # Run every 4 hours
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
search_term:
|
||||
description: 'Search term for models'
|
||||
required: false
|
||||
default: 'GGUF'
|
||||
type: string
|
||||
limit:
|
||||
description: 'Maximum number of models to process'
|
||||
required: false
|
||||
default: '15'
|
||||
type: string
|
||||
quantization:
|
||||
description: 'Preferred quantization format'
|
||||
required: false
|
||||
default: 'Q4_K_M'
|
||||
type: string
|
||||
max_models:
|
||||
description: 'Maximum number of models to add to the gallery'
|
||||
required: false
|
||||
default: '1'
|
||||
type: string
|
||||
jobs:
|
||||
gallery-agent:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v5
|
||||
with:
|
||||
token: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- name: Set up Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: '1.21'
|
||||
|
||||
- name: Build gallery agent
|
||||
run: |
|
||||
cd .github/gallery-agent
|
||||
go mod download
|
||||
go build -o gallery-agent .
|
||||
|
||||
- name: Run gallery agent
|
||||
env:
|
||||
OPENAI_MODEL: ${{ secrets.OPENAI_MODEL }}
|
||||
OPENAI_KEY: ${{ secrets.OPENAI_KEY }}
|
||||
OPENAI_BASE_URL: ${{ secrets.OPENAI_BASE_URL }}
|
||||
SEARCH_TERM: ${{ github.event.inputs.search_term || 'GGUF' }}
|
||||
LIMIT: ${{ github.event.inputs.limit || '15' }}
|
||||
QUANTIZATION: ${{ github.event.inputs.quantization || 'Q4_K_M' }}
|
||||
MAX_MODELS: ${{ github.event.inputs.max_models || '1' }}
|
||||
run: |
|
||||
export GALLERY_INDEX_PATH=$PWD/gallery/index.yaml
|
||||
cd .github/gallery-agent
|
||||
./gallery-agent
|
||||
rm -rf gallery-agent
|
||||
|
||||
- name: Check for changes
|
||||
id: check_changes
|
||||
run: |
|
||||
if git diff --quiet gallery/index.yaml; then
|
||||
echo "changes=false" >> $GITHUB_OUTPUT
|
||||
echo "No changes detected in gallery/index.yaml"
|
||||
else
|
||||
echo "changes=true" >> $GITHUB_OUTPUT
|
||||
echo "Changes detected in gallery/index.yaml"
|
||||
git diff gallery/index.yaml
|
||||
fi
|
||||
|
||||
- name: Read gallery agent summary
|
||||
id: read_summary
|
||||
if: steps.check_changes.outputs.changes == 'true'
|
||||
run: |
|
||||
if [ -f ".github/gallery-agent/gallery-agent-summary.json" ]; then
|
||||
echo "summary_exists=true" >> $GITHUB_OUTPUT
|
||||
# Extract summary data using jq
|
||||
echo "search_term=$(jq -r '.search_term' .github/gallery-agent/gallery-agent-summary.json)" >> $GITHUB_OUTPUT
|
||||
echo "total_found=$(jq -r '.total_found' .github/gallery-agent/gallery-agent-summary.json)" >> $GITHUB_OUTPUT
|
||||
echo "models_added=$(jq -r '.models_added' .github/gallery-agent/gallery-agent-summary.json)" >> $GITHUB_OUTPUT
|
||||
echo "quantization=$(jq -r '.quantization' .github/gallery-agent/gallery-agent-summary.json)" >> $GITHUB_OUTPUT
|
||||
echo "processing_time=$(jq -r '.processing_time' .github/gallery-agent/gallery-agent-summary.json)" >> $GITHUB_OUTPUT
|
||||
|
||||
# Create a formatted list of added models with URLs
|
||||
added_models=$(jq -r 'range(0; .added_model_ids | length) as $i | "- [\(.added_model_ids[$i])](\(.added_model_urls[$i]))"' .github/gallery-agent/gallery-agent-summary.json | tr '\n' '\n')
|
||||
echo "added_models<<EOF" >> $GITHUB_OUTPUT
|
||||
echo "$added_models" >> $GITHUB_OUTPUT
|
||||
echo "EOF" >> $GITHUB_OUTPUT
|
||||
rm -f .github/gallery-agent/gallery-agent-summary.json
|
||||
else
|
||||
echo "summary_exists=false" >> $GITHUB_OUTPUT
|
||||
fi
|
||||
|
||||
- name: Create Pull Request
|
||||
if: steps.check_changes.outputs.changes == 'true'
|
||||
uses: peter-evans/create-pull-request@v7
|
||||
with:
|
||||
token: ${{ secrets.UPDATE_BOT_TOKEN }}
|
||||
push-to-fork: ci-forks/LocalAI
|
||||
commit-message: 'chore(model gallery): :robot: add new models via gallery agent'
|
||||
title: 'chore(model gallery): :robot: add ${{ steps.read_summary.outputs.models_added || 0 }} new models via gallery agent'
|
||||
# Branch has to be unique so PRs are not overriding each other
|
||||
branch-suffix: timestamp
|
||||
body: |
|
||||
This PR was automatically created by the gallery agent workflow.
|
||||
|
||||
**Summary:**
|
||||
- **Search Term:** ${{ steps.read_summary.outputs.search_term || github.event.inputs.search_term || 'GGUF' }}
|
||||
- **Models Found:** ${{ steps.read_summary.outputs.total_found || 'N/A' }}
|
||||
- **Models Added:** ${{ steps.read_summary.outputs.models_added || '0' }}
|
||||
- **Quantization:** ${{ steps.read_summary.outputs.quantization || github.event.inputs.quantization || 'Q4_K_M' }}
|
||||
- **Processing Time:** ${{ steps.read_summary.outputs.processing_time || 'N/A' }}
|
||||
|
||||
**Added Models:**
|
||||
${{ steps.read_summary.outputs.added_models || '- No models added' }}
|
||||
|
||||
**Workflow Details:**
|
||||
- Triggered by: `${{ github.event_name }}`
|
||||
- Run ID: `${{ github.run_id }}`
|
||||
- Commit: `${{ github.sha }}`
|
||||
signoff: true
|
||||
delete-branch: true
|
||||
2
.github/workflows/labeler.yml
vendored
2
.github/workflows/labeler.yml
vendored
@@ -9,4 +9,4 @@ jobs:
|
||||
pull-requests: write
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/labeler@v5
|
||||
- uses: actions/labeler@v6
|
||||
5
.github/workflows/localaibot_automerge.yml
vendored
5
.github/workflows/localaibot_automerge.yml
vendored
@@ -6,11 +6,12 @@ permissions:
|
||||
contents: write
|
||||
pull-requests: write
|
||||
packages: read
|
||||
|
||||
issues: write # for Homebrew/actions/post-comment
|
||||
actions: write # to dispatch publish workflow
|
||||
jobs:
|
||||
dependabot:
|
||||
runs-on: ubuntu-latest
|
||||
if: ${{ github.actor == 'localai-bot' }}
|
||||
if: ${{ github.actor == 'localai-bot' && !contains(github.event.pull_request.title, 'chore(model gallery):') }}
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v5
|
||||
|
||||
18
.github/workflows/notify-models.yaml
vendored
18
.github/workflows/notify-models.yaml
vendored
@@ -1,22 +1,27 @@
|
||||
name: Notifications for new models
|
||||
on:
|
||||
pull_request:
|
||||
pull_request_target:
|
||||
types:
|
||||
- closed
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
|
||||
jobs:
|
||||
notify-discord:
|
||||
if: ${{ (github.event.pull_request.merged == true) && (contains(github.event.pull_request.labels.*.name, 'area/ai-model')) }}
|
||||
env:
|
||||
MODEL_NAME: gemma-3-12b-it
|
||||
MODEL_NAME: gemma-3-12b-it-qat
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v5
|
||||
with:
|
||||
fetch-depth: 0 # needed to checkout all branches for this Action to work
|
||||
ref: ${{ github.event.pull_request.head.sha }} # Checkout the PR head to get the actual changes
|
||||
- uses: mudler/localai-github-action@v1
|
||||
with:
|
||||
model: 'gemma-3-12b-it' # Any from models.localai.io, or from huggingface.com with: "huggingface://<repository>/file"
|
||||
model: 'gemma-3-12b-it-qat' # Any from models.localai.io, or from huggingface.com with: "huggingface://<repository>/file"
|
||||
# Check the PR diff using the current branch and the base branch of the PR
|
||||
- uses: GrantBirki/git-diff-action@v2.8.1
|
||||
id: git-diff-action
|
||||
@@ -79,7 +84,7 @@ jobs:
|
||||
args: ${{ steps.summarize.outputs.message }}
|
||||
- name: Setup tmate session if fails
|
||||
if: ${{ failure() }}
|
||||
uses: mxschmitt/action-tmate@v3.22
|
||||
uses: mxschmitt/action-tmate@v3.23
|
||||
with:
|
||||
detached: true
|
||||
connect-timeout-seconds: 180
|
||||
@@ -87,12 +92,13 @@ jobs:
|
||||
notify-twitter:
|
||||
if: ${{ (github.event.pull_request.merged == true) && (contains(github.event.pull_request.labels.*.name, 'area/ai-model')) }}
|
||||
env:
|
||||
MODEL_NAME: gemma-3-12b-it
|
||||
MODEL_NAME: gemma-3-12b-it-qat
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v5
|
||||
with:
|
||||
fetch-depth: 0 # needed to checkout all branches for this Action to work
|
||||
ref: ${{ github.event.pull_request.head.sha }} # Checkout the PR head to get the actual changes
|
||||
- name: Start LocalAI
|
||||
run: |
|
||||
echo "Starting LocalAI..."
|
||||
@@ -161,7 +167,7 @@ jobs:
|
||||
TWITTER_ACCESS_TOKEN_SECRET: ${{ secrets.TWITTER_ACCESS_TOKEN_SECRET }}
|
||||
- name: Setup tmate session if fails
|
||||
if: ${{ failure() }}
|
||||
uses: mxschmitt/action-tmate@v3.22
|
||||
uses: mxschmitt/action-tmate@v3.23
|
||||
with:
|
||||
detached: true
|
||||
connect-timeout-seconds: 180
|
||||
|
||||
3
.github/workflows/notify-releases.yaml
vendored
3
.github/workflows/notify-releases.yaml
vendored
@@ -11,10 +11,11 @@ jobs:
|
||||
RELEASE_BODY: ${{ github.event.release.body }}
|
||||
RELEASE_TITLE: ${{ github.event.release.name }}
|
||||
RELEASE_TAG_NAME: ${{ github.event.release.tag_name }}
|
||||
MODEL_NAME: gemma-3-12b-it-qat
|
||||
steps:
|
||||
- uses: mudler/localai-github-action@v1
|
||||
with:
|
||||
model: 'gemma-3-12b-it' # Any from models.localai.io, or from huggingface.com with: "huggingface://<repository>/file"
|
||||
model: 'gemma-3-12b-it-qat' # Any from models.localai.io, or from huggingface.com with: "huggingface://<repository>/file"
|
||||
- name: Summarize
|
||||
id: summarize
|
||||
run: |
|
||||
|
||||
40
.github/workflows/release.yaml
vendored
40
.github/workflows/release.yaml
vendored
@@ -23,4 +23,42 @@ jobs:
|
||||
version: v2.11.0
|
||||
args: release --clean
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
launcher-build-darwin:
|
||||
runs-on: macos-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v5
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- name: Set up Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: 1.23
|
||||
- name: Build launcher for macOS ARM64
|
||||
run: |
|
||||
make build-launcher-darwin
|
||||
- name: Upload DMG to Release
|
||||
uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
files: ./dist/LocalAI.dmg
|
||||
launcher-build-linux:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v5
|
||||
with:
|
||||
fetch-depth: 0
|
||||
- name: Set up Go
|
||||
uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version: 1.23
|
||||
- name: Build launcher for Linux
|
||||
run: |
|
||||
sudo apt-get update
|
||||
sudo apt-get install golang gcc libgl1-mesa-dev xorg-dev libxkbcommon-dev
|
||||
make build-launcher-linux
|
||||
- name: Upload Linux launcher artifacts
|
||||
uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
files: ./local-ai-launcher-linux.tar.xz
|
||||
|
||||
4
.github/workflows/secscan.yaml
vendored
4
.github/workflows/secscan.yaml
vendored
@@ -18,13 +18,13 @@ jobs:
|
||||
if: ${{ github.actor != 'dependabot[bot]' }}
|
||||
- name: Run Gosec Security Scanner
|
||||
if: ${{ github.actor != 'dependabot[bot]' }}
|
||||
uses: securego/gosec@v2.22.8
|
||||
uses: securego/gosec@v2.22.9
|
||||
with:
|
||||
# we let the report trigger content trigger a failure using the GitHub Security features.
|
||||
args: '-no-fail -fmt sarif -out results.sarif ./...'
|
||||
- name: Upload SARIF file
|
||||
if: ${{ github.actor != 'dependabot[bot]' }}
|
||||
uses: github/codeql-action/upload-sarif@v3
|
||||
uses: github/codeql-action/upload-sarif@v4
|
||||
with:
|
||||
# Path to SARIF file relative to the root of the repository
|
||||
sarif_file: results.sarif
|
||||
|
||||
2
.github/workflows/stalebot.yml
vendored
2
.github/workflows/stalebot.yml
vendored
@@ -10,7 +10,7 @@ jobs:
|
||||
stale:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/stale@5bef64f19d7facfb25b37b414482c7164d639639 # v9
|
||||
- uses: actions/stale@5f858e3efba33a5ca4407a664cc011ad407f2008 # v9
|
||||
with:
|
||||
stale-issue-message: 'This issue is stale because it has been open 90 days with no activity. Remove stale label or comment or this will be closed in 5 days.'
|
||||
stale-pr-message: 'This PR is stale because it has been open 90 days with no activity. Remove stale label or comment or this will be closed in 10 days.'
|
||||
|
||||
10
.github/workflows/test.yml
vendored
10
.github/workflows/test.yml
vendored
@@ -21,7 +21,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
go-version: ['1.21.x']
|
||||
go-version: ['1.25.x']
|
||||
steps:
|
||||
- name: Free Disk Space (Ubuntu)
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
@@ -124,7 +124,7 @@ jobs:
|
||||
PATH="$PATH:/root/go/bin" GO_TAGS="tts" make --jobs 5 --output-sync=target test
|
||||
- name: Setup tmate session if tests fail
|
||||
if: ${{ failure() }}
|
||||
uses: mxschmitt/action-tmate@v3.22
|
||||
uses: mxschmitt/action-tmate@v3.23
|
||||
with:
|
||||
detached: true
|
||||
connect-timeout-seconds: 180
|
||||
@@ -183,7 +183,7 @@ jobs:
|
||||
PATH="$PATH:$HOME/go/bin" make backends/local-store backends/silero-vad backends/llama-cpp backends/whisper backends/piper backends/stablediffusion-ggml docker-build-aio e2e-aio
|
||||
- name: Setup tmate session if tests fail
|
||||
if: ${{ failure() }}
|
||||
uses: mxschmitt/action-tmate@v3.22
|
||||
uses: mxschmitt/action-tmate@v3.23
|
||||
with:
|
||||
detached: true
|
||||
connect-timeout-seconds: 180
|
||||
@@ -193,7 +193,7 @@ jobs:
|
||||
runs-on: macOS-14
|
||||
strategy:
|
||||
matrix:
|
||||
go-version: ['1.21.x']
|
||||
go-version: ['1.25.x']
|
||||
steps:
|
||||
- name: Clone
|
||||
uses: actions/checkout@v5
|
||||
@@ -226,7 +226,7 @@ jobs:
|
||||
PATH="$PATH:$HOME/go/bin" BUILD_TYPE="GITHUB_CI_HAS_BROKEN_METAL" CMAKE_ARGS="-DGGML_F16C=OFF -DGGML_AVX512=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF" make --jobs 4 --output-sync=target test
|
||||
- name: Setup tmate session if tests fail
|
||||
if: ${{ failure() }}
|
||||
uses: mxschmitt/action-tmate@v3.22
|
||||
uses: mxschmitt/action-tmate@v3.23
|
||||
with:
|
||||
detached: true
|
||||
connect-timeout-seconds: 180
|
||||
|
||||
2
.gitignore
vendored
2
.gitignore
vendored
@@ -24,7 +24,7 @@ go-bert
|
||||
|
||||
# LocalAI build binary
|
||||
LocalAI
|
||||
local-ai
|
||||
/local-ai
|
||||
# prevent above rules from omitting the helm chart
|
||||
!charts/*
|
||||
# prevent above rules from omitting the api/localai folder
|
||||
|
||||
@@ -8,7 +8,7 @@ source:
|
||||
enabled: true
|
||||
name_template: '{{ .ProjectName }}-{{ .Tag }}-source'
|
||||
builds:
|
||||
-
|
||||
- main: ./cmd/local-ai
|
||||
env:
|
||||
- CGO_ENABLED=0
|
||||
ldflags:
|
||||
|
||||
14
Dockerfile
14
Dockerfile
@@ -78,6 +78,16 @@ RUN <<EOT bash
|
||||
fi
|
||||
EOT
|
||||
|
||||
# https://github.com/NVIDIA/Isaac-GR00T/issues/343
|
||||
RUN <<EOT bash
|
||||
if [ "${BUILD_TYPE}" = "cublas" ] && [ "${TARGETARCH}" = "arm64" ]; then
|
||||
wget https://developer.download.nvidia.com/compute/cudss/0.6.0/local_installers/cudss-local-tegra-repo-ubuntu2204-0.6.0_0.6.0-1_arm64.deb && \
|
||||
dpkg -i cudss-local-tegra-repo-ubuntu2204-0.6.0_0.6.0-1_arm64.deb && \
|
||||
cp /var/cudss-local-tegra-repo-ubuntu2204-0.6.0/cudss-*-keyring.gpg /usr/share/keyrings/ && \
|
||||
apt-get update && apt-get -y install cudss
|
||||
fi
|
||||
EOT
|
||||
|
||||
# If we are building with clblas support, we need the libraries for the builds
|
||||
RUN if [ "${BUILD_TYPE}" = "clblas" ] && [ "${SKIP_DRIVERS}" = "false" ]; then \
|
||||
apt-get update && \
|
||||
@@ -100,6 +110,10 @@ RUN if [ "${BUILD_TYPE}" = "hipblas" ] && [ "${SKIP_DRIVERS}" = "false" ]; then
|
||||
ldconfig \
|
||||
; fi
|
||||
|
||||
RUN if [ "${BUILD_TYPE}" = "hipblas" ]; then \
|
||||
ln -s /opt/rocm-**/lib/llvm/lib/libomp.so /usr/lib/libomp.so \
|
||||
; fi
|
||||
|
||||
RUN expr "${BUILD_TYPE}" = intel && echo "intel" > /run/localai/capability || echo "not intel"
|
||||
|
||||
# Cuda
|
||||
|
||||
61
Makefile
61
Makefile
@@ -2,6 +2,7 @@ GOCMD=go
|
||||
GOTEST=$(GOCMD) test
|
||||
GOVET=$(GOCMD) vet
|
||||
BINARY_NAME=local-ai
|
||||
LAUNCHER_BINARY_NAME=local-ai-launcher
|
||||
|
||||
GORELEASER?=
|
||||
|
||||
@@ -90,7 +91,17 @@ build: protogen-go install-go-tools ## Build the project
|
||||
$(info ${GREEN}I LD_FLAGS: ${YELLOW}$(LD_FLAGS)${RESET})
|
||||
$(info ${GREEN}I UPX: ${YELLOW}$(UPX)${RESET})
|
||||
rm -rf $(BINARY_NAME) || true
|
||||
CGO_LDFLAGS="$(CGO_LDFLAGS)" $(GOCMD) build -ldflags "$(LD_FLAGS)" -tags "$(GO_TAGS)" -o $(BINARY_NAME) ./
|
||||
CGO_LDFLAGS="$(CGO_LDFLAGS)" $(GOCMD) build -ldflags "$(LD_FLAGS)" -tags "$(GO_TAGS)" -o $(BINARY_NAME) ./cmd/local-ai
|
||||
|
||||
build-launcher: ## Build the launcher application
|
||||
$(info ${GREEN}I local-ai launcher build info:${RESET})
|
||||
$(info ${GREEN}I BUILD_TYPE: ${YELLOW}$(BUILD_TYPE)${RESET})
|
||||
$(info ${GREEN}I GO_TAGS: ${YELLOW}$(GO_TAGS)${RESET})
|
||||
$(info ${GREEN}I LD_FLAGS: ${YELLOW}$(LD_FLAGS)${RESET})
|
||||
rm -rf $(LAUNCHER_BINARY_NAME) || true
|
||||
CGO_LDFLAGS="$(CGO_LDFLAGS)" $(GOCMD) build -ldflags "$(LD_FLAGS)" -tags "$(GO_TAGS)" -o $(LAUNCHER_BINARY_NAME) ./cmd/launcher
|
||||
|
||||
build-all: build build-launcher ## Build both server and launcher
|
||||
|
||||
dev-dist:
|
||||
$(GORELEASER) build --snapshot --clean
|
||||
@@ -106,8 +117,8 @@ run: ## run local-ai
|
||||
CGO_LDFLAGS="$(CGO_LDFLAGS)" $(GOCMD) run ./
|
||||
|
||||
test-models/testmodel.ggml:
|
||||
mkdir test-models
|
||||
mkdir test-dir
|
||||
mkdir -p test-models
|
||||
mkdir -p test-dir
|
||||
wget -q https://huggingface.co/mradermacher/gpt2-alpaca-gpt4-GGUF/resolve/main/gpt2-alpaca-gpt4.Q4_K_M.gguf -O test-models/testmodel.ggml
|
||||
wget -q https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-base.en.bin -O test-models/whisper-en
|
||||
wget -q https://huggingface.co/mudler/all-MiniLM-L6-v2/resolve/main/ggml-model-q4_0.bin -O test-models/bert
|
||||
@@ -358,13 +369,22 @@ backends/kitten-tts: docker-build-kitten-tts docker-save-kitten-tts build
|
||||
backends/kokoro: docker-build-kokoro docker-save-kokoro build
|
||||
./local-ai backends install "ocifile://$(abspath ./backend-images/kokoro.tar)"
|
||||
|
||||
backends/chatterbox: docker-build-chatterbox docker-save-chatterbox build
|
||||
./local-ai backends install "ocifile://$(abspath ./backend-images/chatterbox.tar)"
|
||||
|
||||
backends/llama-cpp-darwin: build
|
||||
bash ./scripts/build/llama-cpp-darwin.sh
|
||||
./local-ai backends install "ocifile://$(abspath ./backend-images/llama-cpp.tar)"
|
||||
|
||||
backends/neutts: docker-build-neutts docker-save-neutts build
|
||||
./local-ai backends install "ocifile://$(abspath ./backend-images/neutts.tar)"
|
||||
|
||||
build-darwin-python-backend: build
|
||||
bash ./scripts/build/python-darwin.sh
|
||||
|
||||
build-darwin-go-backend: build
|
||||
bash ./scripts/build/golang-darwin.sh
|
||||
|
||||
backends/mlx:
|
||||
BACKEND=mlx $(MAKE) build-darwin-python-backend
|
||||
./local-ai backends install "ocifile://$(abspath ./backend-images/mlx.tar)"
|
||||
@@ -377,6 +397,14 @@ backends/mlx-vlm:
|
||||
BACKEND=mlx-vlm $(MAKE) build-darwin-python-backend
|
||||
./local-ai backends install "ocifile://$(abspath ./backend-images/mlx-vlm.tar)"
|
||||
|
||||
backends/mlx-audio:
|
||||
BACKEND=mlx-audio $(MAKE) build-darwin-python-backend
|
||||
./local-ai backends install "ocifile://$(abspath ./backend-images/mlx-audio.tar)"
|
||||
|
||||
backends/stablediffusion-ggml-darwin:
|
||||
BACKEND=stablediffusion-ggml BUILD_TYPE=metal $(MAKE) build-darwin-go-backend
|
||||
./local-ai backends install "ocifile://$(abspath ./backend-images/stablediffusion-ggml.tar)"
|
||||
|
||||
backend-images:
|
||||
mkdir -p backend-images
|
||||
|
||||
@@ -404,6 +432,15 @@ docker-build-kitten-tts:
|
||||
docker-save-kitten-tts: backend-images
|
||||
docker save local-ai-backend:kitten-tts -o backend-images/kitten-tts.tar
|
||||
|
||||
docker-save-chatterbox: backend-images
|
||||
docker save local-ai-backend:chatterbox -o backend-images/chatterbox.tar
|
||||
|
||||
docker-build-neutts:
|
||||
docker build --build-arg BUILD_TYPE=$(BUILD_TYPE) --build-arg BASE_IMAGE=$(BASE_IMAGE) -t local-ai-backend:neutts -f backend/Dockerfile.python --build-arg BACKEND=neutts ./backend
|
||||
|
||||
docker-save-neutts: backend-images
|
||||
docker save local-ai-backend:neutts -o backend-images/neutts.tar
|
||||
|
||||
docker-build-kokoro:
|
||||
docker build --build-arg BUILD_TYPE=$(BUILD_TYPE) --build-arg BASE_IMAGE=$(BASE_IMAGE) -t local-ai-backend:kokoro -f backend/Dockerfile.python --build-arg BACKEND=kokoro ./backend
|
||||
|
||||
@@ -471,7 +508,7 @@ docker-build-bark:
|
||||
docker build --build-arg BUILD_TYPE=$(BUILD_TYPE) --build-arg BASE_IMAGE=$(BASE_IMAGE) -t local-ai-backend:bark -f backend/Dockerfile.python --build-arg BACKEND=bark .
|
||||
|
||||
docker-build-chatterbox:
|
||||
docker build --build-arg BUILD_TYPE=$(BUILD_TYPE) --build-arg BASE_IMAGE=$(BASE_IMAGE) -t local-ai-backend:chatterbox -f backend/Dockerfile.python --build-arg BACKEND=chatterbox .
|
||||
docker build --build-arg BUILD_TYPE=$(BUILD_TYPE) --build-arg BASE_IMAGE=$(BASE_IMAGE) -t local-ai-backend:chatterbox -f backend/Dockerfile.python --build-arg BACKEND=chatterbox ./backend
|
||||
|
||||
docker-build-exllama2:
|
||||
docker build --build-arg BUILD_TYPE=$(BUILD_TYPE) --build-arg BASE_IMAGE=$(BASE_IMAGE) -t local-ai-backend:exllama2 -f backend/Dockerfile.python --build-arg BACKEND=exllama2 .
|
||||
@@ -507,3 +544,19 @@ docs-clean:
|
||||
.PHONY: docs
|
||||
docs: docs/static/gallery.html
|
||||
cd docs && hugo serve
|
||||
|
||||
########################################################
|
||||
## Platform-specific builds
|
||||
########################################################
|
||||
|
||||
## fyne cross-platform build
|
||||
build-launcher-darwin: build-launcher
|
||||
go run github.com/tiagomelo/macos-dmg-creator/cmd/createdmg@latest \
|
||||
--appName "LocalAI" \
|
||||
--appBinaryPath "$(LAUNCHER_BINARY_NAME)" \
|
||||
--bundleIdentifier "com.localai.launcher" \
|
||||
--iconPath "core/http/static/logo.png" \
|
||||
--outputDir "dist/"
|
||||
|
||||
build-launcher-linux:
|
||||
cd cmd/launcher && go run fyne.io/tools/cmd/fyne@latest package -os linux -icon ../../core/http/static/logo.png --executable $(LAUNCHER_BINARY_NAME)-linux && mv launcher.tar.xz ../../$(LAUNCHER_BINARY_NAME)-linux.tar.xz
|
||||
|
||||
83
README.md
83
README.md
@@ -43,7 +43,7 @@
|
||||
|
||||
> :bulb: Get help - [❓FAQ](https://localai.io/faq/) [💭Discussions](https://github.com/go-skynet/LocalAI/discussions) [:speech_balloon: Discord](https://discord.gg/uJAeKSAGDy) [:book: Documentation website](https://localai.io/)
|
||||
>
|
||||
> [💻 Quickstart](https://localai.io/basics/getting_started/) [🖼️ Models](https://models.localai.io/) [🚀 Roadmap](https://github.com/mudler/LocalAI/issues?q=is%3Aissue+is%3Aopen+label%3Aroadmap) [🥽 Demo](https://demo.localai.io) [🌍 Explorer](https://explorer.localai.io) [🛫 Examples](https://github.com/mudler/LocalAI-examples) Try on
|
||||
> [💻 Quickstart](https://localai.io/basics/getting_started/) [🖼️ Models](https://models.localai.io/) [🚀 Roadmap](https://github.com/mudler/LocalAI/issues?q=is%3Aissue+is%3Aopen+label%3Aroadmap) [🌍 Explorer](https://explorer.localai.io) [🛫 Examples](https://github.com/mudler/LocalAI-examples) Try on
|
||||
[](https://t.me/localaiofficial_bot)
|
||||
|
||||
[](https://github.com/go-skynet/LocalAI/actions/workflows/test.yml)[](https://github.com/go-skynet/LocalAI/actions/workflows/release.yaml)[](https://github.com/go-skynet/LocalAI/actions/workflows/image.yml)[](https://github.com/go-skynet/LocalAI/actions/workflows/bump_deps.yaml)[](https://artifacthub.io/packages/search?repo=localai)
|
||||
@@ -110,8 +110,21 @@ curl https://localai.io/install.sh | sh
|
||||
|
||||
For more installation options, see [Installer Options](https://localai.io/docs/advanced/installer/).
|
||||
|
||||
### macOS Download:
|
||||
|
||||
<a href="https://github.com/mudler/LocalAI/releases/latest/download/LocalAI.dmg">
|
||||
<img src="https://img.shields.io/badge/Download-macOS-blue?style=for-the-badge&logo=apple&logoColor=white" alt="Download LocalAI for macOS"/>
|
||||
</a>
|
||||
|
||||
Or run with docker:
|
||||
|
||||
> **💡 Docker Run vs Docker Start**
|
||||
>
|
||||
> - `docker run` creates and starts a new container. If a container with the same name already exists, this command will fail.
|
||||
> - `docker start` starts an existing container that was previously created with `docker run`.
|
||||
>
|
||||
> If you've already run LocalAI before and want to start it again, use: `docker start -i local-ai`
|
||||
|
||||
### CPU only image:
|
||||
|
||||
```bash
|
||||
@@ -191,6 +204,8 @@ For more information, see [💻 Getting started](https://localai.io/basics/getti
|
||||
|
||||
## 📰 Latest project news
|
||||
|
||||
- October 2025: 🔌 [Model Context Protocol (MCP)](https://localai.io/docs/features/mcp/) support added for agentic capabilities with external tools
|
||||
- September 2025: New Launcher application for MacOS and Linux, extended support to many backends for Mac and Nvidia L4T devices. Models: Added MLX-Audio, WAN 2.2. WebUI improvements and Python-based backends now ships portable python environments.
|
||||
- August 2025: MLX, MLX-VLM, Diffusers and llama.cpp are now supported on Mac M1/M2/M3+ chips ( with `development` suffix in the gallery ): https://github.com/mudler/LocalAI/pull/6049 https://github.com/mudler/LocalAI/pull/6119 https://github.com/mudler/LocalAI/pull/6121 https://github.com/mudler/LocalAI/pull/6060
|
||||
- July/August 2025: 🔍 [Object Detection](https://localai.io/features/object-detection/) added to the API featuring [rf-detr](https://github.com/roboflow/rf-detr)
|
||||
- July 2025: All backends migrated outside of the main binary. LocalAI is now more lightweight, small, and automatically downloads the required backend to run the model. [Read the release notes](https://github.com/mudler/LocalAI/releases/tag/v3.2.0)
|
||||
@@ -229,10 +244,65 @@ Roadmap items: [List of issues](https://github.com/mudler/LocalAI/issues?q=is%3A
|
||||
- 🔍 [Object Detection](https://localai.io/features/object-detection/)
|
||||
- 📈 [Reranker API](https://localai.io/features/reranker/)
|
||||
- 🆕🖧 [P2P Inferencing](https://localai.io/features/distribute/)
|
||||
- [Agentic capabilities](https://github.com/mudler/LocalAGI)
|
||||
- 🆕🔌 [Model Context Protocol (MCP)](https://localai.io/docs/features/mcp/) - Agentic capabilities with external tools and [LocalAGI's Agentic capabilities](https://github.com/mudler/LocalAGI)
|
||||
- 🔊 Voice activity detection (Silero-VAD support)
|
||||
- 🌍 Integrated WebUI!
|
||||
|
||||
## 🧩 Supported Backends & Acceleration
|
||||
|
||||
LocalAI supports a comprehensive range of AI backends with multiple acceleration options:
|
||||
|
||||
### Text Generation & Language Models
|
||||
| Backend | Description | Acceleration Support |
|
||||
|---------|-------------|---------------------|
|
||||
| **llama.cpp** | LLM inference in C/C++ | CUDA 11/12, ROCm, Intel SYCL, Vulkan, Metal, CPU |
|
||||
| **vLLM** | Fast LLM inference with PagedAttention | CUDA 12, ROCm, Intel |
|
||||
| **transformers** | HuggingFace transformers framework | CUDA 11/12, ROCm, Intel, CPU |
|
||||
| **exllama2** | GPTQ inference library | CUDA 12 |
|
||||
| **MLX** | Apple Silicon LLM inference | Metal (M1/M2/M3+) |
|
||||
| **MLX-VLM** | Apple Silicon Vision-Language Models | Metal (M1/M2/M3+) |
|
||||
|
||||
### Audio & Speech Processing
|
||||
| Backend | Description | Acceleration Support |
|
||||
|---------|-------------|---------------------|
|
||||
| **whisper.cpp** | OpenAI Whisper in C/C++ | CUDA 12, ROCm, Intel SYCL, Vulkan, CPU |
|
||||
| **faster-whisper** | Fast Whisper with CTranslate2 | CUDA 12, ROCm, Intel, CPU |
|
||||
| **bark** | Text-to-audio generation | CUDA 12, ROCm, Intel |
|
||||
| **bark-cpp** | C++ implementation of Bark | CUDA, Metal, CPU |
|
||||
| **coqui** | Advanced TTS with 1100+ languages | CUDA 12, ROCm, Intel, CPU |
|
||||
| **kokoro** | Lightweight TTS model | CUDA 12, ROCm, Intel, CPU |
|
||||
| **chatterbox** | Production-grade TTS | CUDA 11/12, CPU |
|
||||
| **piper** | Fast neural TTS system | CPU |
|
||||
| **kitten-tts** | Kitten TTS models | CPU |
|
||||
| **silero-vad** | Voice Activity Detection | CPU |
|
||||
| **neutts** | Text-to-speech with voice cloning | CUDA 12, ROCm, CPU |
|
||||
|
||||
### Image & Video Generation
|
||||
| Backend | Description | Acceleration Support |
|
||||
|---------|-------------|---------------------|
|
||||
| **stablediffusion.cpp** | Stable Diffusion in C/C++ | CUDA 12, Intel SYCL, Vulkan, CPU |
|
||||
| **diffusers** | HuggingFace diffusion models | CUDA 11/12, ROCm, Intel, Metal, CPU |
|
||||
|
||||
### Specialized AI Tasks
|
||||
| Backend | Description | Acceleration Support |
|
||||
|---------|-------------|---------------------|
|
||||
| **rfdetr** | Real-time object detection | CUDA 12, Intel, CPU |
|
||||
| **rerankers** | Document reranking API | CUDA 11/12, ROCm, Intel, CPU |
|
||||
| **local-store** | Vector database | CPU |
|
||||
| **huggingface** | HuggingFace API integration | API-based |
|
||||
|
||||
### Hardware Acceleration Matrix
|
||||
|
||||
| Acceleration Type | Supported Backends | Hardware Support |
|
||||
|-------------------|-------------------|------------------|
|
||||
| **NVIDIA CUDA 11** | llama.cpp, whisper, stablediffusion, diffusers, rerankers, bark, chatterbox | Nvidia hardware |
|
||||
| **NVIDIA CUDA 12** | All CUDA-compatible backends | Nvidia hardware |
|
||||
| **AMD ROCm** | llama.cpp, whisper, vllm, transformers, diffusers, rerankers, coqui, kokoro, bark, neutts | AMD Graphics |
|
||||
| **Intel oneAPI** | llama.cpp, whisper, stablediffusion, vllm, transformers, diffusers, rfdetr, rerankers, exllama2, coqui, kokoro, bark | Intel Arc, Intel iGPUs |
|
||||
| **Apple Metal** | llama.cpp, whisper, diffusers, MLX, MLX-VLM, bark-cpp | Apple M1/M2/M3+ |
|
||||
| **Vulkan** | llama.cpp, whisper, stablediffusion | Cross-platform GPUs |
|
||||
| **NVIDIA Jetson** | llama.cpp, whisper, stablediffusion, diffusers, rfdetr | ARM64 embedded AI |
|
||||
| **CPU Optimized** | All backends | AVX/AVX2/AVX512, quantization support |
|
||||
|
||||
### 🔗 Community and integrations
|
||||
|
||||
@@ -244,9 +314,18 @@ WebUIs:
|
||||
- https://github.com/go-skynet/LocalAI-frontend
|
||||
- QA-Pilot(An interactive chat project that leverages LocalAI LLMs for rapid understanding and navigation of GitHub code repository) https://github.com/reid41/QA-Pilot
|
||||
|
||||
Agentic Libraries:
|
||||
- https://github.com/mudler/cogito
|
||||
|
||||
MCPs:
|
||||
- https://github.com/mudler/MCPs
|
||||
|
||||
Model galleries
|
||||
- https://github.com/go-skynet/model-gallery
|
||||
|
||||
Voice:
|
||||
- https://github.com/richiejp/VoxInput
|
||||
|
||||
Other:
|
||||
- Helm chart https://github.com/go-skynet/helm-charts
|
||||
- VSCode extension https://github.com/badgooooor/localai-vscode-plugin
|
||||
|
||||
@@ -2,10 +2,10 @@ context_size: 4096
|
||||
f16: true
|
||||
backend: llama-cpp
|
||||
mmap: true
|
||||
mmproj: minicpm-v-2_6-mmproj-f16.gguf
|
||||
mmproj: minicpm-v-4_5-mmproj-f16.gguf
|
||||
name: gpt-4o
|
||||
parameters:
|
||||
model: minicpm-v-2_6-Q4_K_M.gguf
|
||||
model: minicpm-v-4_5-Q4_K_M.gguf
|
||||
stopwords:
|
||||
- <|im_end|>
|
||||
- <dummy32000>
|
||||
@@ -42,9 +42,9 @@ template:
|
||||
<|im_start|>assistant
|
||||
|
||||
download_files:
|
||||
- filename: minicpm-v-2_6-Q4_K_M.gguf
|
||||
sha256: 3a4078d53b46f22989adbf998ce5a3fd090b6541f112d7e936eb4204a04100b1
|
||||
uri: huggingface://openbmb/MiniCPM-V-2_6-gguf/ggml-model-Q4_K_M.gguf
|
||||
- filename: minicpm-v-2_6-mmproj-f16.gguf
|
||||
uri: huggingface://openbmb/MiniCPM-V-2_6-gguf/mmproj-model-f16.gguf
|
||||
sha256: 4485f68a0f1aa404c391e788ea88ea653c100d8e98fe572698f701e5809711fd
|
||||
- filename: minicpm-v-4_5-Q4_K_M.gguf
|
||||
sha256: c1c3c33100b15b4caf7319acce4e23c0eb0ce1cbd12f70e8d24f05aa67b7512f
|
||||
uri: huggingface://openbmb/MiniCPM-V-4_5-gguf/ggml-model-Q4_K_M.gguf
|
||||
- filename: minicpm-v-4_5-mmproj-f16.gguf
|
||||
uri: huggingface://openbmb/MiniCPM-V-4_5-gguf/mmproj-model-f16.gguf
|
||||
sha256: 7a7225a32e8d453aaa3d22d8c579b5bf833c253f784cdb05c99c9a76fd616df8
|
||||
@@ -2,10 +2,10 @@ context_size: 4096
|
||||
backend: llama-cpp
|
||||
f16: true
|
||||
mmap: true
|
||||
mmproj: minicpm-v-2_6-mmproj-f16.gguf
|
||||
mmproj: minicpm-v-4_5-mmproj-f16.gguf
|
||||
name: gpt-4o
|
||||
parameters:
|
||||
model: minicpm-v-2_6-Q4_K_M.gguf
|
||||
model: minicpm-v-4_5-Q4_K_M.gguf
|
||||
stopwords:
|
||||
- <|im_end|>
|
||||
- <dummy32000>
|
||||
@@ -42,9 +42,9 @@ template:
|
||||
<|im_start|>assistant
|
||||
|
||||
download_files:
|
||||
- filename: minicpm-v-2_6-Q4_K_M.gguf
|
||||
sha256: 3a4078d53b46f22989adbf998ce5a3fd090b6541f112d7e936eb4204a04100b1
|
||||
uri: huggingface://openbmb/MiniCPM-V-2_6-gguf/ggml-model-Q4_K_M.gguf
|
||||
- filename: minicpm-v-2_6-mmproj-f16.gguf
|
||||
uri: huggingface://openbmb/MiniCPM-V-2_6-gguf/mmproj-model-f16.gguf
|
||||
sha256: 4485f68a0f1aa404c391e788ea88ea653c100d8e98fe572698f701e5809711fd
|
||||
- filename: minicpm-v-4_5-Q4_K_M.gguf
|
||||
sha256: c1c3c33100b15b4caf7319acce4e23c0eb0ce1cbd12f70e8d24f05aa67b7512f
|
||||
uri: huggingface://openbmb/MiniCPM-V-4_5-gguf/ggml-model-Q4_K_M.gguf
|
||||
- filename: minicpm-v-4_5-mmproj-f16.gguf
|
||||
uri: huggingface://openbmb/MiniCPM-V-4_5-gguf/mmproj-model-f16.gguf
|
||||
sha256: 7a7225a32e8d453aaa3d22d8c579b5bf833c253f784cdb05c99c9a76fd616df8
|
||||
@@ -2,10 +2,10 @@ context_size: 4096
|
||||
backend: llama-cpp
|
||||
f16: true
|
||||
mmap: true
|
||||
mmproj: minicpm-v-2_6-mmproj-f16.gguf
|
||||
mmproj: minicpm-v-4_5-mmproj-f16.gguf
|
||||
name: gpt-4o
|
||||
parameters:
|
||||
model: minicpm-v-2_6-Q4_K_M.gguf
|
||||
model: minicpm-v-4_5-Q4_K_M.gguf
|
||||
stopwords:
|
||||
- <|im_end|>
|
||||
- <dummy32000>
|
||||
@@ -43,9 +43,9 @@ template:
|
||||
|
||||
|
||||
download_files:
|
||||
- filename: minicpm-v-2_6-Q4_K_M.gguf
|
||||
sha256: 3a4078d53b46f22989adbf998ce5a3fd090b6541f112d7e936eb4204a04100b1
|
||||
uri: huggingface://openbmb/MiniCPM-V-2_6-gguf/ggml-model-Q4_K_M.gguf
|
||||
- filename: minicpm-v-2_6-mmproj-f16.gguf
|
||||
uri: huggingface://openbmb/MiniCPM-V-2_6-gguf/mmproj-model-f16.gguf
|
||||
sha256: 4485f68a0f1aa404c391e788ea88ea653c100d8e98fe572698f701e5809711fd
|
||||
- filename: minicpm-v-4_5-Q4_K_M.gguf
|
||||
sha256: c1c3c33100b15b4caf7319acce4e23c0eb0ce1cbd12f70e8d24f05aa67b7512f
|
||||
uri: huggingface://openbmb/MiniCPM-V-4_5-gguf/ggml-model-Q4_K_M.gguf
|
||||
- filename: minicpm-v-4_5-mmproj-f16.gguf
|
||||
uri: huggingface://openbmb/MiniCPM-V-4_5-gguf/mmproj-model-f16.gguf
|
||||
sha256: 7a7225a32e8d453aaa3d22d8c579b5bf833c253f784cdb05c99c9a76fd616df8
|
||||
@@ -197,7 +197,7 @@ EOT
|
||||
|
||||
|
||||
# Copy libraries using a script to handle architecture differences
|
||||
RUN make -C /LocalAI/backend/cpp/llama-cpp package
|
||||
RUN make -BC /LocalAI/backend/cpp/llama-cpp package
|
||||
|
||||
|
||||
FROM scratch
|
||||
|
||||
@@ -28,7 +28,7 @@ RUN apt-get update && \
|
||||
curl python3-pip \
|
||||
python-is-python3 \
|
||||
python3-dev llvm \
|
||||
python3-venv make && \
|
||||
python3-venv make cmake && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/* && \
|
||||
pip install --upgrade pip
|
||||
|
||||
213
backend/README.md
Normal file
213
backend/README.md
Normal file
@@ -0,0 +1,213 @@
|
||||
# LocalAI Backend Architecture
|
||||
|
||||
This directory contains the core backend infrastructure for LocalAI, including the gRPC protocol definition, multi-language Dockerfiles, and language-specific backend implementations.
|
||||
|
||||
## Overview
|
||||
|
||||
LocalAI uses a unified gRPC-based architecture that allows different programming languages to implement AI backends while maintaining consistent interfaces and capabilities. The backend system supports multiple hardware acceleration targets and provides a standardized way to integrate various AI models and frameworks.
|
||||
|
||||
## Architecture Components
|
||||
|
||||
### 1. Protocol Definition (`backend.proto`)
|
||||
|
||||
The `backend.proto` file defines the gRPC service interface that all backends must implement. This ensures consistency across different language implementations and provides a contract for communication between LocalAI core and backend services.
|
||||
|
||||
#### Core Services
|
||||
|
||||
- **Text Generation**: `Predict`, `PredictStream` for LLM inference
|
||||
- **Embeddings**: `Embedding` for text vectorization
|
||||
- **Image Generation**: `GenerateImage` for stable diffusion and image models
|
||||
- **Audio Processing**: `AudioTranscription`, `TTS`, `SoundGeneration`
|
||||
- **Video Generation**: `GenerateVideo` for video synthesis
|
||||
- **Object Detection**: `Detect` for computer vision tasks
|
||||
- **Vector Storage**: `StoresSet`, `StoresGet`, `StoresFind` for RAG operations
|
||||
- **Reranking**: `Rerank` for document relevance scoring
|
||||
- **Voice Activity Detection**: `VAD` for audio segmentation
|
||||
|
||||
#### Key Message Types
|
||||
|
||||
- **`PredictOptions`**: Comprehensive configuration for text generation
|
||||
- **`ModelOptions`**: Model loading and configuration parameters
|
||||
- **`Result`**: Standardized response format
|
||||
- **`StatusResponse`**: Backend health and memory usage information
|
||||
|
||||
### 2. Multi-Language Dockerfiles
|
||||
|
||||
The backend system provides language-specific Dockerfiles that handle the build environment and dependencies for different programming languages:
|
||||
|
||||
- `Dockerfile.python`
|
||||
- `Dockerfile.golang`
|
||||
- `Dockerfile.llama-cpp`
|
||||
|
||||
### 3. Language-Specific Implementations
|
||||
|
||||
#### Python Backends (`python/`)
|
||||
- **transformers**: Hugging Face Transformers framework
|
||||
- **vllm**: High-performance LLM inference
|
||||
- **mlx**: Apple Silicon optimization
|
||||
- **diffusers**: Stable Diffusion models
|
||||
- **Audio**: bark, coqui, faster-whisper, kitten-tts
|
||||
- **Vision**: mlx-vlm, rfdetr
|
||||
- **Specialized**: rerankers, chatterbox, kokoro
|
||||
|
||||
#### Go Backends (`go/`)
|
||||
- **whisper**: OpenAI Whisper speech recognition in Go with GGML cpp backend (whisper.cpp)
|
||||
- **stablediffusion-ggml**: Stable Diffusion in Go with GGML Cpp backend
|
||||
- **huggingface**: Hugging Face model integration
|
||||
- **piper**: Text-to-speech synthesis Golang with C bindings using rhaspy/piper
|
||||
- **bark-cpp**: Bark TTS models Golang with Cpp bindings
|
||||
- **local-store**: Vector storage backend
|
||||
|
||||
#### C++ Backends (`cpp/`)
|
||||
- **llama-cpp**: Llama.cpp integration
|
||||
- **grpc**: GRPC utilities and helpers
|
||||
|
||||
## Hardware Acceleration Support
|
||||
|
||||
### CUDA (NVIDIA)
|
||||
- **Versions**: CUDA 11.x, 12.x
|
||||
- **Features**: cuBLAS, cuDNN, TensorRT optimization
|
||||
- **Targets**: x86_64, ARM64 (Jetson)
|
||||
|
||||
### ROCm (AMD)
|
||||
- **Features**: HIP, rocBLAS, MIOpen
|
||||
- **Targets**: AMD GPUs with ROCm support
|
||||
|
||||
### Intel
|
||||
- **Features**: oneAPI, Intel Extension for PyTorch
|
||||
- **Targets**: Intel GPUs, XPUs, CPUs
|
||||
|
||||
### Vulkan
|
||||
- **Features**: Cross-platform GPU acceleration
|
||||
- **Targets**: Windows, Linux, Android, macOS
|
||||
|
||||
### Apple Silicon
|
||||
- **Features**: MLX framework, Metal Performance Shaders
|
||||
- **Targets**: M1/M2/M3 Macs
|
||||
|
||||
## Backend Registry (`index.yaml`)
|
||||
|
||||
The `index.yaml` file serves as a central registry for all available backends, providing:
|
||||
|
||||
- **Metadata**: Name, description, license, icons
|
||||
- **Capabilities**: Hardware targets and optimization profiles
|
||||
- **Tags**: Categorization for discovery
|
||||
- **URLs**: Source code and documentation links
|
||||
|
||||
## Building Backends
|
||||
|
||||
### Prerequisites
|
||||
- Docker with multi-architecture support
|
||||
- Appropriate hardware drivers (CUDA, ROCm, etc.)
|
||||
- Build tools (make, cmake, compilers)
|
||||
|
||||
### Build Commands
|
||||
|
||||
Example of build commands with Docker
|
||||
|
||||
```bash
|
||||
# Build Python backend
|
||||
docker build -f backend/Dockerfile.python \
|
||||
--build-arg BACKEND=transformers \
|
||||
--build-arg BUILD_TYPE=cublas12 \
|
||||
--build-arg CUDA_MAJOR_VERSION=12 \
|
||||
--build-arg CUDA_MINOR_VERSION=0 \
|
||||
-t localai-backend-transformers .
|
||||
|
||||
# Build Go backend
|
||||
docker build -f backend/Dockerfile.golang \
|
||||
--build-arg BACKEND=whisper \
|
||||
--build-arg BUILD_TYPE=cpu \
|
||||
-t localai-backend-whisper .
|
||||
|
||||
# Build C++ backend
|
||||
docker build -f backend/Dockerfile.llama-cpp \
|
||||
--build-arg BACKEND=llama-cpp \
|
||||
--build-arg BUILD_TYPE=cublas12 \
|
||||
-t localai-backend-llama-cpp .
|
||||
```
|
||||
|
||||
For ARM64/Mac builds, docker can't be used, and the makefile in the respective backend has to be used.
|
||||
|
||||
### Build Types
|
||||
|
||||
- **`cpu`**: CPU-only optimization
|
||||
- **`cublas11`**: CUDA 11.x with cuBLAS
|
||||
- **`cublas12`**: CUDA 12.x with cuBLAS
|
||||
- **`hipblas`**: ROCm with rocBLAS
|
||||
- **`intel`**: Intel oneAPI optimization
|
||||
- **`vulkan`**: Vulkan-based acceleration
|
||||
- **`metal`**: Apple Metal optimization
|
||||
|
||||
## Backend Development
|
||||
|
||||
### Creating a New Backend
|
||||
|
||||
1. **Choose Language**: Select Python, Go, or C++ based on requirements
|
||||
2. **Implement Interface**: Implement the gRPC service defined in `backend.proto`
|
||||
3. **Add Dependencies**: Create appropriate requirements files
|
||||
4. **Configure Build**: Set up Dockerfile and build scripts
|
||||
5. **Register Backend**: Add entry to `index.yaml`
|
||||
6. **Test Integration**: Verify gRPC communication and functionality
|
||||
|
||||
### Backend Structure
|
||||
|
||||
```
|
||||
backend-name/
|
||||
├── backend.py/go/cpp # Main implementation
|
||||
├── requirements.txt # Dependencies
|
||||
├── Dockerfile # Build configuration
|
||||
├── install.sh # Installation script
|
||||
├── run.sh # Execution script
|
||||
├── test.sh # Test script
|
||||
└── README.md # Backend documentation
|
||||
```
|
||||
|
||||
### Required gRPC Methods
|
||||
|
||||
At minimum, backends must implement:
|
||||
- `Health()` - Service health check
|
||||
- `LoadModel()` - Model loading and initialization
|
||||
- `Predict()` - Main inference endpoint
|
||||
- `Status()` - Backend status and metrics
|
||||
|
||||
## Integration with LocalAI Core
|
||||
|
||||
Backends communicate with LocalAI core through gRPC:
|
||||
|
||||
1. **Service Discovery**: Core discovers available backends
|
||||
2. **Model Loading**: Core requests model loading via `LoadModel`
|
||||
3. **Inference**: Core sends requests via `Predict` or specialized endpoints
|
||||
4. **Streaming**: Core handles streaming responses for real-time generation
|
||||
5. **Monitoring**: Core tracks backend health and performance
|
||||
|
||||
## Performance Optimization
|
||||
|
||||
### Memory Management
|
||||
- **Model Caching**: Efficient model loading and caching
|
||||
- **Batch Processing**: Optimize for multiple concurrent requests
|
||||
- **Memory Pinning**: GPU memory optimization for CUDA/ROCm
|
||||
|
||||
### Hardware Utilization
|
||||
- **Multi-GPU**: Support for tensor parallelism
|
||||
- **Mixed Precision**: FP16/BF16 for memory efficiency
|
||||
- **Kernel Fusion**: Optimized CUDA/ROCm kernels
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
1. **GRPC Connection**: Verify backend service is running and accessible
|
||||
2. **Model Loading**: Check model paths and dependencies
|
||||
3. **Hardware Detection**: Ensure appropriate drivers and libraries
|
||||
4. **Memory Issues**: Monitor GPU memory usage and model sizes
|
||||
|
||||
## Contributing
|
||||
|
||||
When contributing to the backend system:
|
||||
|
||||
1. **Follow Protocol**: Implement the exact gRPC interface
|
||||
2. **Add Tests**: Include comprehensive test coverage
|
||||
3. **Document**: Provide clear usage examples
|
||||
4. **Optimize**: Consider performance and resource usage
|
||||
5. **Validate**: Test across different hardware targets
|
||||
@@ -242,7 +242,7 @@ message ModelOptions {
|
||||
|
||||
string Type = 49;
|
||||
|
||||
bool FlashAttention = 56;
|
||||
string FlashAttention = 56;
|
||||
bool NoKVOffload = 57;
|
||||
|
||||
string ModelPath = 59;
|
||||
@@ -276,6 +276,7 @@ message TranscriptRequest {
|
||||
string language = 3;
|
||||
uint32 threads = 4;
|
||||
bool translate = 5;
|
||||
bool diarize = 6;
|
||||
}
|
||||
|
||||
message TranscriptResult {
|
||||
@@ -305,22 +306,24 @@ message GenerateImageRequest {
|
||||
// Diffusers
|
||||
string EnableParameters = 10;
|
||||
int32 CLIPSkip = 11;
|
||||
|
||||
|
||||
// Reference images for models that support them (e.g., Flux Kontext)
|
||||
repeated string ref_images = 12;
|
||||
}
|
||||
|
||||
message GenerateVideoRequest {
|
||||
string prompt = 1;
|
||||
string start_image = 2; // Path or base64 encoded image for the start frame
|
||||
string end_image = 3; // Path or base64 encoded image for the end frame
|
||||
int32 width = 4;
|
||||
int32 height = 5;
|
||||
int32 num_frames = 6; // Number of frames to generate
|
||||
int32 fps = 7; // Frames per second
|
||||
int32 seed = 8;
|
||||
float cfg_scale = 9; // Classifier-free guidance scale
|
||||
string dst = 10; // Output path for the generated video
|
||||
string negative_prompt = 2; // Negative prompt for video generation
|
||||
string start_image = 3; // Path or base64 encoded image for the start frame
|
||||
string end_image = 4; // Path or base64 encoded image for the end frame
|
||||
int32 width = 5;
|
||||
int32 height = 6;
|
||||
int32 num_frames = 7; // Number of frames to generate
|
||||
int32 fps = 8; // Frames per second
|
||||
int32 seed = 9;
|
||||
float cfg_scale = 10; // Classifier-free guidance scale
|
||||
int32 step = 11; // Number of inference steps
|
||||
string dst = 12; // Output path for the generated video
|
||||
}
|
||||
|
||||
message TTSRequest {
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
LLAMA_VERSION?=710dfc465a68f7443b87d9f792cffba00ed739fe
|
||||
LLAMA_VERSION?=5a4ff43e7dd049e35942bc3d12361dab2f155544
|
||||
LLAMA_REPO?=https://github.com/ggerganov/llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
@@ -14,7 +14,7 @@ CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF -DLLAMA_CURL=OFF
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
ifeq ($(NATIVE),false)
|
||||
CMAKE_ARGS+=-DGGML_NATIVE=OFF
|
||||
CMAKE_ARGS+=-DGGML_NATIVE=OFF -DLLAMA_OPENSSL=OFF
|
||||
endif
|
||||
# If build type is cublas, then we set -DGGML_CUDA=ON to CMAKE_ARGS automatically
|
||||
ifeq ($(BUILD_TYPE),cublas)
|
||||
|
||||
@@ -92,7 +92,7 @@ static void start_llama_server(server_context& ctx_server) {
|
||||
ctx_server.queue_tasks.start_loop();
|
||||
}
|
||||
|
||||
json parse_options(bool streaming, const backend::PredictOptions* predict)
|
||||
json parse_options(bool streaming, const backend::PredictOptions* predict, const server_context& ctx_server)
|
||||
{
|
||||
|
||||
// Create now a json data from the prediction options instead
|
||||
@@ -147,6 +147,28 @@ json parse_options(bool streaming, const backend::PredictOptions* predict)
|
||||
// data["n_probs"] = predict->nprobs();
|
||||
//TODO: images,
|
||||
|
||||
// Serialize grammar triggers from server context to JSON array
|
||||
if (!ctx_server.params_base.sampling.grammar_triggers.empty()) {
|
||||
json grammar_triggers = json::array();
|
||||
for (const auto& trigger : ctx_server.params_base.sampling.grammar_triggers) {
|
||||
json trigger_json;
|
||||
trigger_json["value"] = trigger.value;
|
||||
// Always serialize as WORD type since upstream converts WORD to TOKEN internally
|
||||
trigger_json["type"] = static_cast<int>(COMMON_GRAMMAR_TRIGGER_TYPE_WORD);
|
||||
grammar_triggers.push_back(trigger_json);
|
||||
}
|
||||
data["grammar_triggers"] = grammar_triggers;
|
||||
}
|
||||
|
||||
// Serialize preserved tokens from server context to JSON array
|
||||
if (!ctx_server.params_base.sampling.preserved_tokens.empty()) {
|
||||
json preserved_tokens = json::array();
|
||||
for (const auto& token : ctx_server.params_base.sampling.preserved_tokens) {
|
||||
preserved_tokens.push_back(common_token_to_piece(ctx_server.ctx, token));
|
||||
}
|
||||
data["preserved_tokens"] = preserved_tokens;
|
||||
}
|
||||
|
||||
return data;
|
||||
}
|
||||
|
||||
@@ -207,7 +229,7 @@ static void add_rpc_devices(std::string servers) {
|
||||
}
|
||||
}
|
||||
|
||||
static void params_parse(const backend::ModelOptions* request,
|
||||
static void params_parse(server_context& ctx_server, const backend::ModelOptions* request,
|
||||
common_params & params) {
|
||||
|
||||
// this is comparable to: https://github.com/ggerganov/llama.cpp/blob/d9b33fe95bd257b36c84ee5769cc048230067d6f/examples/server/server.cpp#L1809
|
||||
@@ -231,6 +253,7 @@ static void params_parse(const backend::ModelOptions* request,
|
||||
params.cpuparams.n_threads = request->threads();
|
||||
params.n_gpu_layers = request->ngpulayers();
|
||||
params.n_batch = request->nbatch();
|
||||
params.n_ubatch = request->nbatch(); // fixes issue with reranking models being limited to 512 tokens (the default n_ubatch size); allows for setting the maximum input amount of tokens thereby avoiding this error "input is too large to process. increase the physical batch size"
|
||||
// Set params.n_parallel by environment variable (LLAMA_PARALLEL), defaults to 1
|
||||
//params.n_parallel = 1;
|
||||
const char *env_parallel = std::getenv("LLAMACPP_PARALLEL");
|
||||
@@ -268,6 +291,11 @@ static void params_parse(const backend::ModelOptions* request,
|
||||
}
|
||||
}
|
||||
|
||||
if (!params.kv_overrides.empty()) {
|
||||
params.kv_overrides.emplace_back();
|
||||
params.kv_overrides.back().key[0] = 0;
|
||||
}
|
||||
|
||||
// TODO: Add yarn
|
||||
|
||||
if (!request->tensorsplit().empty()) {
|
||||
@@ -304,7 +332,15 @@ static void params_parse(const backend::ModelOptions* request,
|
||||
}
|
||||
params.use_mlock = request->mlock();
|
||||
params.use_mmap = request->mmap();
|
||||
params.flash_attn = request->flashattention();
|
||||
|
||||
if (request->flashattention() == "on" || request->flashattention() == "enabled") {
|
||||
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
|
||||
} else if (request->flashattention() == "off" || request->flashattention() == "disabled") {
|
||||
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
|
||||
} else if (request->flashattention() == "auto") {
|
||||
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO;
|
||||
}
|
||||
|
||||
params.no_kv_offload = request->nokvoffload();
|
||||
params.ctx_shift = false; // We control context-shifting in any case (and we disable it as it could just lead to infinite loops)
|
||||
|
||||
@@ -338,14 +374,14 @@ static void params_parse(const backend::ModelOptions* request,
|
||||
}
|
||||
|
||||
if (request->grammartriggers_size() > 0) {
|
||||
params.sampling.grammar_lazy = true;
|
||||
//params.sampling.grammar_lazy = true;
|
||||
// Store grammar trigger words for processing after model is loaded
|
||||
for (int i = 0; i < request->grammartriggers_size(); i++) {
|
||||
const auto & word = request->grammartriggers(i).word();
|
||||
common_grammar_trigger trigger;
|
||||
trigger.type = COMMON_GRAMMAR_TRIGGER_TYPE_WORD;
|
||||
trigger.value = request->grammartriggers(i).word();
|
||||
// trigger.at_start = request->grammartriggers(i).at_start();
|
||||
params.sampling.grammar_triggers.push_back(trigger);
|
||||
|
||||
trigger.type = COMMON_GRAMMAR_TRIGGER_TYPE_WORD;
|
||||
trigger.value = word;
|
||||
params.sampling.grammar_triggers.push_back(std::move(trigger));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -368,7 +404,7 @@ public:
|
||||
grpc::Status LoadModel(ServerContext* context, const backend::ModelOptions* request, backend::Result* result) {
|
||||
// Implement LoadModel RPC
|
||||
common_params params;
|
||||
params_parse(request, params);
|
||||
params_parse(ctx_server, request, params);
|
||||
|
||||
common_init();
|
||||
|
||||
@@ -387,6 +423,39 @@ public:
|
||||
return Status::CANCELLED;
|
||||
}
|
||||
|
||||
// Process grammar triggers now that vocab is available
|
||||
if (!params.sampling.grammar_triggers.empty()) {
|
||||
std::vector<common_grammar_trigger> processed_triggers;
|
||||
for (const auto& trigger : params.sampling.grammar_triggers) {
|
||||
if (trigger.type == COMMON_GRAMMAR_TRIGGER_TYPE_WORD) {
|
||||
auto ids = common_tokenize(ctx_server.vocab, trigger.value, /* add_special= */ false, /* parse_special= */ true);
|
||||
if (ids.size() == 1) {
|
||||
auto token = ids[0];
|
||||
// Add the token to preserved_tokens if not already present
|
||||
if (params.sampling.preserved_tokens.find(token) == params.sampling.preserved_tokens.end()) {
|
||||
params.sampling.preserved_tokens.insert(token);
|
||||
LOG_INF("Added grammar trigger token to preserved tokens: %d (`%s`)\n", token, trigger.value.c_str());
|
||||
}
|
||||
LOG_INF("Grammar trigger token: %d (`%s`)\n", token, trigger.value.c_str());
|
||||
common_grammar_trigger processed_trigger;
|
||||
processed_trigger.type = COMMON_GRAMMAR_TRIGGER_TYPE_TOKEN;
|
||||
processed_trigger.value = trigger.value;
|
||||
processed_trigger.token = token;
|
||||
processed_triggers.push_back(std::move(processed_trigger));
|
||||
} else {
|
||||
LOG_INF("Grammar trigger word: `%s`\n", trigger.value.c_str());
|
||||
processed_triggers.push_back(trigger);
|
||||
}
|
||||
} else {
|
||||
processed_triggers.push_back(trigger);
|
||||
}
|
||||
}
|
||||
// Update the grammar triggers in params_base
|
||||
ctx_server.params_base.sampling.grammar_triggers = std::move(processed_triggers);
|
||||
// Also update preserved_tokens in params_base
|
||||
ctx_server.params_base.sampling.preserved_tokens = params.sampling.preserved_tokens;
|
||||
}
|
||||
|
||||
//ctx_server.init();
|
||||
result->set_message("Loading succeeded");
|
||||
result->set_success(true);
|
||||
@@ -397,7 +466,7 @@ public:
|
||||
}
|
||||
|
||||
grpc::Status PredictStream(grpc::ServerContext* context, const backend::PredictOptions* request, grpc::ServerWriter<backend::Reply>* writer) override {
|
||||
json data = parse_options(true, request);
|
||||
json data = parse_options(true, request, ctx_server);
|
||||
|
||||
|
||||
//Raise error if embeddings is set to true
|
||||
@@ -460,12 +529,12 @@ public:
|
||||
task.id = ctx_server.queue_tasks.get_new_id();
|
||||
task.index = i;
|
||||
|
||||
task.prompt_tokens = std::move(inputs[i]);
|
||||
task.tokens = std::move(inputs[i]);
|
||||
task.params = server_task::params_from_json_cmpl(
|
||||
ctx_server.ctx,
|
||||
ctx_server.params_base,
|
||||
data);
|
||||
task.id_selected_slot = json_value(data, "id_slot", -1);
|
||||
task.id_slot = json_value(data, "id_slot", -1);
|
||||
|
||||
// OAI-compat
|
||||
task.params.oaicompat = OAICOMPAT_TYPE_NONE;
|
||||
@@ -547,7 +616,7 @@ public:
|
||||
}
|
||||
|
||||
grpc::Status Predict(ServerContext* context, const backend::PredictOptions* request, backend::Reply* reply) {
|
||||
json data = parse_options(true, request);
|
||||
json data = parse_options(true, request, ctx_server);
|
||||
|
||||
data["stream"] = false;
|
||||
//Raise error if embeddings is set to true
|
||||
@@ -615,12 +684,12 @@ public:
|
||||
task.id = ctx_server.queue_tasks.get_new_id();
|
||||
task.index = i;
|
||||
|
||||
task.prompt_tokens = std::move(inputs[i]);
|
||||
task.tokens = std::move(inputs[i]);
|
||||
task.params = server_task::params_from_json_cmpl(
|
||||
ctx_server.ctx,
|
||||
ctx_server.params_base,
|
||||
data);
|
||||
task.id_selected_slot = json_value(data, "id_slot", -1);
|
||||
task.id_slot = json_value(data, "id_slot", -1);
|
||||
|
||||
// OAI-compat
|
||||
task.params.oaicompat = OAICOMPAT_TYPE_NONE;
|
||||
@@ -682,7 +751,7 @@ public:
|
||||
|
||||
grpc::Status Embedding(ServerContext* context, const backend::PredictOptions* request, backend::EmbeddingResult* embeddingResult) {
|
||||
|
||||
json body = parse_options(false, request);
|
||||
json body = parse_options(false, request, ctx_server);
|
||||
|
||||
body["stream"] = false;
|
||||
|
||||
@@ -693,7 +762,7 @@ public:
|
||||
*/
|
||||
|
||||
// for the shape of input/content, see tokenize_input_prompts()
|
||||
json prompt = body.at("prompt");
|
||||
json prompt = body.at("embeddings");
|
||||
|
||||
|
||||
auto tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true);
|
||||
@@ -704,6 +773,7 @@ public:
|
||||
}
|
||||
}
|
||||
|
||||
int embd_normalize = 2; // default to Euclidean/L2 norm
|
||||
// create and queue the task
|
||||
json responses = json::array();
|
||||
bool error = false;
|
||||
@@ -715,11 +785,10 @@ public:
|
||||
|
||||
task.id = ctx_server.queue_tasks.get_new_id();
|
||||
task.index = i;
|
||||
task.prompt_tokens = std::move(tokenized_prompts[i]);
|
||||
|
||||
// OAI-compat
|
||||
task.params.oaicompat = OAICOMPAT_TYPE_EMBEDDING;
|
||||
task.tokens = std::move(tokenized_prompts[i]);
|
||||
|
||||
task.params.oaicompat = OAICOMPAT_TYPE_NONE;
|
||||
task.params.embd_normalize = embd_normalize;
|
||||
tasks.push_back(std::move(task));
|
||||
}
|
||||
|
||||
@@ -735,9 +804,8 @@ public:
|
||||
responses.push_back(res->to_json());
|
||||
}
|
||||
}, [&](const json & error_data) {
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, error_data.value("content", ""));
|
||||
error = true;
|
||||
}, [&]() {
|
||||
// NOTE: we should try to check when the writer is closed here
|
||||
return false;
|
||||
});
|
||||
|
||||
@@ -747,12 +815,36 @@ public:
|
||||
return grpc::Status(grpc::StatusCode::INTERNAL, "Error in receiving results");
|
||||
}
|
||||
|
||||
std::vector<float> embeddings = responses[0].value("embedding", std::vector<float>());
|
||||
// loop the vector and set the embeddings results
|
||||
for (int i = 0; i < embeddings.size(); i++) {
|
||||
embeddingResult->add_embeddings(embeddings[i]);
|
||||
std::cout << "[DEBUG] Responses size: " << responses.size() << std::endl;
|
||||
|
||||
// Process the responses and extract embeddings
|
||||
for (const auto & response_elem : responses) {
|
||||
// Check if the response has an "embedding" field
|
||||
if (response_elem.contains("embedding")) {
|
||||
json embedding_data = json_value(response_elem, "embedding", json::array());
|
||||
|
||||
if (embedding_data.is_array() && !embedding_data.empty()) {
|
||||
for (const auto & embedding_vector : embedding_data) {
|
||||
if (embedding_vector.is_array()) {
|
||||
for (const auto & embedding_value : embedding_vector) {
|
||||
embeddingResult->add_embeddings(embedding_value.get<float>());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Check if the response itself contains the embedding data directly
|
||||
if (response_elem.is_array()) {
|
||||
for (const auto & embedding_value : response_elem) {
|
||||
embeddingResult->add_embeddings(embedding_value.get<float>());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
return grpc::Status::OK;
|
||||
}
|
||||
|
||||
@@ -770,11 +862,6 @@ public:
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "\"documents\" must be a non-empty string array");
|
||||
}
|
||||
|
||||
// Tokenize the query
|
||||
auto tokenized_query = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, request->query(), /* add_special */ false, true);
|
||||
if (tokenized_query.size() != 1) {
|
||||
return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "\"query\" must contain only a single prompt");
|
||||
}
|
||||
// Create and queue the task
|
||||
json responses = json::array();
|
||||
bool error = false;
|
||||
@@ -786,14 +873,13 @@ public:
|
||||
documents.push_back(request->documents(i));
|
||||
}
|
||||
|
||||
auto tokenized_docs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, documents, /* add_special */ false, true);
|
||||
tasks.reserve(tokenized_docs.size());
|
||||
for (size_t i = 0; i < tokenized_docs.size(); i++) {
|
||||
auto tmp = format_rerank(ctx_server.vocab, tokenized_query[0], tokenized_docs[i]);
|
||||
tasks.reserve(documents.size());
|
||||
for (size_t i = 0; i < documents.size(); i++) {
|
||||
auto tmp = format_rerank(ctx_server.model, ctx_server.vocab, ctx_server.mctx, request->query(), documents[i]);
|
||||
server_task task = server_task(SERVER_TASK_TYPE_RERANK);
|
||||
task.id = ctx_server.queue_tasks.get_new_id();
|
||||
task.index = i;
|
||||
task.prompt_tokens = std::move(tmp);
|
||||
task.tokens = std::move(tmp);
|
||||
tasks.push_back(std::move(task));
|
||||
}
|
||||
|
||||
@@ -846,7 +932,7 @@ public:
|
||||
}
|
||||
|
||||
grpc::Status TokenizeString(ServerContext* context, const backend::PredictOptions* request, backend::TokenizationResponse* response) {
|
||||
json body = parse_options(false, request);
|
||||
json body = parse_options(false, request, ctx_server);
|
||||
body["stream"] = false;
|
||||
|
||||
json tokens_response = json::array();
|
||||
|
||||
6
backend/go/stablediffusion-ggml/.gitignore
vendored
Normal file
6
backend/go/stablediffusion-ggml/.gitignore
vendored
Normal file
@@ -0,0 +1,6 @@
|
||||
package/
|
||||
sources/
|
||||
.cache/
|
||||
build/
|
||||
libgosd.so
|
||||
stablediffusion-ggml
|
||||
@@ -5,7 +5,11 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
|
||||
add_subdirectory(./sources/stablediffusion-ggml.cpp)
|
||||
|
||||
add_library(gosd MODULE gosd.cpp)
|
||||
target_link_libraries(gosd PRIVATE stable-diffusion ggml stdc++fs)
|
||||
target_link_libraries(gosd PRIVATE stable-diffusion ggml)
|
||||
|
||||
if(CMAKE_CXX_COMPILER_ID MATCHES "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS 9.0)
|
||||
target_link_libraries(gosd PRIVATE stdc++fs)
|
||||
endif()
|
||||
|
||||
target_include_directories(gosd PUBLIC
|
||||
stable-diffusion.cpp
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# stablediffusion.cpp (ggml)
|
||||
STABLEDIFFUSION_GGML_REPO?=https://github.com/leejet/stable-diffusion.cpp
|
||||
STABLEDIFFUSION_GGML_VERSION?=5900ef6605c6fbf7934239f795c13c97bc993853
|
||||
STABLEDIFFUSION_GGML_VERSION?=0ebe6fe118f125665939b27c89f34ed38716bff8
|
||||
|
||||
CMAKE_ARGS+=-DGGML_MAX_NAME=128
|
||||
|
||||
@@ -29,8 +29,6 @@ else ifeq ($(BUILD_TYPE),clblas)
|
||||
# If it's hipblas we do have also to set CC=/opt/rocm/llvm/bin/clang CXX=/opt/rocm/llvm/bin/clang++
|
||||
else ifeq ($(BUILD_TYPE),hipblas)
|
||||
CMAKE_ARGS+=-DSD_HIPBLAS=ON -DGGML_HIPBLAS=ON
|
||||
# If it's OSX, DO NOT embed the metal library - -DGGML_METAL_EMBED_LIBRARY=ON requires further investigation
|
||||
# But if it's OSX without metal, disable it here
|
||||
else ifeq ($(BUILD_TYPE),vulkan)
|
||||
CMAKE_ARGS+=-DSD_VULKAN=ON -DGGML_VULKAN=ON
|
||||
else ifeq ($(OS),Darwin)
|
||||
@@ -74,10 +72,10 @@ libgosd.so: sources/stablediffusion-ggml.cpp CMakeLists.txt gosd.cpp gosd.h
|
||||
stablediffusion-ggml: main.go gosd.go libgosd.so
|
||||
CGO_ENABLED=0 $(GOCMD) build -tags "$(GO_TAGS)" -o stablediffusion-ggml ./
|
||||
|
||||
package:
|
||||
package: stablediffusion-ggml
|
||||
bash package.sh
|
||||
|
||||
build: stablediffusion-ggml package
|
||||
build: package
|
||||
|
||||
clean:
|
||||
rm -rf libgosd.o build stablediffusion-ggml
|
||||
rm -rf libgosd.so build stablediffusion-ggml package sources
|
||||
|
||||
@@ -4,17 +4,11 @@
|
||||
#include <stdio.h>
|
||||
#include <string.h>
|
||||
#include <time.h>
|
||||
#include <iostream>
|
||||
#include <random>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <filesystem>
|
||||
#include "gosd.h"
|
||||
|
||||
// #include "preprocessing.hpp"
|
||||
#include "flux.hpp"
|
||||
#include "stable-diffusion.h"
|
||||
|
||||
#define STB_IMAGE_IMPLEMENTATION
|
||||
#define STB_IMAGE_STATIC
|
||||
#include "stb_image.h"
|
||||
@@ -29,7 +23,7 @@
|
||||
|
||||
// Names of the sampler method, same order as enum sample_method in stable-diffusion.h
|
||||
const char* sample_method_str[] = {
|
||||
"euler_a",
|
||||
"default",
|
||||
"euler",
|
||||
"heun",
|
||||
"dpm2",
|
||||
@@ -41,19 +35,27 @@ const char* sample_method_str[] = {
|
||||
"lcm",
|
||||
"ddim_trailing",
|
||||
"tcd",
|
||||
"euler_a",
|
||||
};
|
||||
|
||||
static_assert(std::size(sample_method_str) == SAMPLE_METHOD_COUNT, "sample method mismatch");
|
||||
|
||||
// Names of the sigma schedule overrides, same order as sample_schedule in stable-diffusion.h
|
||||
const char* schedule_str[] = {
|
||||
const char* schedulers[] = {
|
||||
"default",
|
||||
"discrete",
|
||||
"karras",
|
||||
"exponential",
|
||||
"ays",
|
||||
"gits",
|
||||
"smoothstep",
|
||||
};
|
||||
|
||||
static_assert(std::size(schedulers) == SCHEDULE_COUNT, "schedulers mismatch");
|
||||
|
||||
sd_ctx_t* sd_c;
|
||||
// Moved from the context (load time) to generation time params
|
||||
scheduler_t scheduler = scheduler_t::DEFAULT;
|
||||
|
||||
sample_method_t sample_method;
|
||||
|
||||
@@ -105,7 +107,7 @@ int load_model(const char *model, char *model_path, char* options[], int threads
|
||||
const char *clip_g_path = "";
|
||||
const char *t5xxl_path = "";
|
||||
const char *vae_path = "";
|
||||
const char *scheduler = "";
|
||||
const char *scheduler_str = "";
|
||||
const char *sampler = "";
|
||||
char *lora_dir = model_path;
|
||||
bool lora_dir_allocated = false;
|
||||
@@ -133,7 +135,7 @@ int load_model(const char *model, char *model_path, char* options[], int threads
|
||||
vae_path = optval;
|
||||
}
|
||||
if (!strcmp(optname, "scheduler")) {
|
||||
scheduler = optval;
|
||||
scheduler_str = optval;
|
||||
}
|
||||
if (!strcmp(optname, "sampler")) {
|
||||
sampler = optval;
|
||||
@@ -166,26 +168,17 @@ int load_model(const char *model, char *model_path, char* options[], int threads
|
||||
}
|
||||
if (sample_method_found == -1) {
|
||||
fprintf(stderr, "Invalid sample method, default to EULER_A!\n");
|
||||
sample_method_found = EULER_A;
|
||||
sample_method_found = sample_method_t::SAMPLE_METHOD_DEFAULT;
|
||||
}
|
||||
sample_method = (sample_method_t)sample_method_found;
|
||||
|
||||
int schedule_found = -1;
|
||||
for (int d = 0; d < SCHEDULE_COUNT; d++) {
|
||||
if (!strcmp(scheduler, schedule_str[d])) {
|
||||
schedule_found = d;
|
||||
fprintf (stderr, "Found scheduler: %s\n", scheduler);
|
||||
|
||||
if (!strcmp(scheduler_str, schedulers[d])) {
|
||||
scheduler = (scheduler_t)d;
|
||||
fprintf (stderr, "Found scheduler: %s\n", scheduler_str);
|
||||
}
|
||||
}
|
||||
|
||||
if (schedule_found == -1) {
|
||||
fprintf (stderr, "Invalid scheduler! using DEFAULT\n");
|
||||
schedule_found = DEFAULT;
|
||||
}
|
||||
|
||||
schedule_t schedule = (schedule_t)schedule_found;
|
||||
|
||||
fprintf (stderr, "Creating context\n");
|
||||
sd_ctx_params_t ctx_params;
|
||||
sd_ctx_params_init(&ctx_params);
|
||||
@@ -199,13 +192,10 @@ int load_model(const char *model, char *model_path, char* options[], int threads
|
||||
ctx_params.control_net_path = "";
|
||||
ctx_params.lora_model_dir = lora_dir;
|
||||
ctx_params.embedding_dir = "";
|
||||
ctx_params.stacked_id_embed_dir = "";
|
||||
ctx_params.vae_decode_only = false;
|
||||
ctx_params.vae_tiling = false;
|
||||
ctx_params.free_params_immediately = false;
|
||||
ctx_params.n_threads = threads;
|
||||
ctx_params.rng_type = STD_DEFAULT_RNG;
|
||||
ctx_params.schedule = schedule;
|
||||
sd_ctx_t* sd_ctx = new_sd_ctx(&ctx_params);
|
||||
|
||||
if (sd_ctx == NULL) {
|
||||
@@ -228,7 +218,49 @@ int load_model(const char *model, char *model_path, char* options[], int threads
|
||||
return 0;
|
||||
}
|
||||
|
||||
int gen_image(char *text, char *negativeText, int width, int height, int steps, int64_t seed, char *dst, float cfg_scale, char *src_image, float strength, char *mask_image, char **ref_images, int ref_images_count) {
|
||||
void sd_tiling_params_set_enabled(sd_tiling_params_t *params, bool enabled) {
|
||||
params->enabled = enabled;
|
||||
}
|
||||
|
||||
void sd_tiling_params_set_tile_sizes(sd_tiling_params_t *params, int tile_size_x, int tile_size_y) {
|
||||
params->tile_size_x = tile_size_x;
|
||||
params->tile_size_y = tile_size_y;
|
||||
}
|
||||
|
||||
void sd_tiling_params_set_rel_sizes(sd_tiling_params_t *params, float rel_size_x, float rel_size_y) {
|
||||
params->rel_size_x = rel_size_x;
|
||||
params->rel_size_y = rel_size_y;
|
||||
}
|
||||
|
||||
void sd_tiling_params_set_target_overlap(sd_tiling_params_t *params, float target_overlap) {
|
||||
params->target_overlap = target_overlap;
|
||||
}
|
||||
|
||||
sd_tiling_params_t* sd_img_gen_params_get_vae_tiling_params(sd_img_gen_params_t *params) {
|
||||
return ¶ms->vae_tiling_params;
|
||||
}
|
||||
|
||||
sd_img_gen_params_t* sd_img_gen_params_new(void) {
|
||||
sd_img_gen_params_t *params = (sd_img_gen_params_t *)std::malloc(sizeof(sd_img_gen_params_t));
|
||||
sd_img_gen_params_init(params);
|
||||
return params;
|
||||
}
|
||||
|
||||
void sd_img_gen_params_set_prompts(sd_img_gen_params_t *params, const char *prompt, const char *negative_prompt) {
|
||||
params->prompt = prompt;
|
||||
params->negative_prompt = negative_prompt;
|
||||
}
|
||||
|
||||
void sd_img_gen_params_set_dimensions(sd_img_gen_params_t *params, int width, int height) {
|
||||
params->width = width;
|
||||
params->height = height;
|
||||
}
|
||||
|
||||
void sd_img_gen_params_set_seed(sd_img_gen_params_t *params, int64_t seed) {
|
||||
params->seed = seed;
|
||||
}
|
||||
|
||||
int gen_image(sd_img_gen_params_t *p, int steps, char *dst, float cfg_scale, char *src_image, float strength, char *mask_image, char **ref_images, int ref_images_count) {
|
||||
|
||||
sd_image_t* results;
|
||||
|
||||
@@ -236,20 +268,15 @@ int gen_image(char *text, char *negativeText, int width, int height, int steps,
|
||||
|
||||
fprintf (stderr, "Generating image\n");
|
||||
|
||||
sd_img_gen_params_t p;
|
||||
sd_img_gen_params_init(&p);
|
||||
p->sample_params.guidance.txt_cfg = cfg_scale;
|
||||
p->sample_params.guidance.slg.layers = skip_layers.data();
|
||||
p->sample_params.guidance.slg.layer_count = skip_layers.size();
|
||||
p->sample_params.sample_method = sample_method;
|
||||
p->sample_params.sample_steps = steps;
|
||||
p->sample_params.scheduler = scheduler;
|
||||
|
||||
p.prompt = text;
|
||||
p.negative_prompt = negativeText;
|
||||
p.guidance.txt_cfg = cfg_scale;
|
||||
p.guidance.slg.layers = skip_layers.data();
|
||||
p.guidance.slg.layer_count = skip_layers.size();
|
||||
p.width = width;
|
||||
p.height = height;
|
||||
p.sample_method = sample_method;
|
||||
p.sample_steps = steps;
|
||||
p.seed = seed;
|
||||
p.input_id_images_path = "";
|
||||
int width = p->width;
|
||||
int height = p->height;
|
||||
|
||||
// Handle input image for img2img
|
||||
bool has_input_image = (src_image != NULL && strlen(src_image) > 0);
|
||||
@@ -298,13 +325,13 @@ int gen_image(char *text, char *negativeText, int width, int height, int steps,
|
||||
input_image_buffer = resized_image_buffer;
|
||||
}
|
||||
|
||||
p.init_image = {(uint32_t)width, (uint32_t)height, 3, input_image_buffer};
|
||||
p.strength = strength;
|
||||
p->init_image = {(uint32_t)width, (uint32_t)height, 3, input_image_buffer};
|
||||
p->strength = strength;
|
||||
fprintf(stderr, "Using img2img with strength: %.2f\n", strength);
|
||||
} else {
|
||||
// No input image, use empty image for text-to-image
|
||||
p.init_image = {(uint32_t)width, (uint32_t)height, 3, NULL};
|
||||
p.strength = 0.0f;
|
||||
p->init_image = {(uint32_t)width, (uint32_t)height, 3, NULL};
|
||||
p->strength = 0.0f;
|
||||
}
|
||||
|
||||
// Handle mask image for inpainting
|
||||
@@ -344,12 +371,12 @@ int gen_image(char *text, char *negativeText, int width, int height, int steps,
|
||||
mask_image_buffer = resized_mask_buffer;
|
||||
}
|
||||
|
||||
p.mask_image = {(uint32_t)width, (uint32_t)height, 1, mask_image_buffer};
|
||||
p->mask_image = {(uint32_t)width, (uint32_t)height, 1, mask_image_buffer};
|
||||
fprintf(stderr, "Using inpainting with mask\n");
|
||||
} else {
|
||||
// No mask image, create default full mask
|
||||
default_mask_image_vec.resize(width * height, 255);
|
||||
p.mask_image = {(uint32_t)width, (uint32_t)height, 1, default_mask_image_vec.data()};
|
||||
p->mask_image = {(uint32_t)width, (uint32_t)height, 1, default_mask_image_vec.data()};
|
||||
}
|
||||
|
||||
// Handle reference images
|
||||
@@ -407,13 +434,15 @@ int gen_image(char *text, char *negativeText, int width, int height, int steps,
|
||||
}
|
||||
|
||||
if (!ref_images_vec.empty()) {
|
||||
p.ref_images = ref_images_vec.data();
|
||||
p.ref_images_count = ref_images_vec.size();
|
||||
p->ref_images = ref_images_vec.data();
|
||||
p->ref_images_count = ref_images_vec.size();
|
||||
fprintf(stderr, "Using %zu reference images\n", ref_images_vec.size());
|
||||
}
|
||||
}
|
||||
|
||||
results = generate_image(sd_c, &p);
|
||||
results = generate_image(sd_c, p);
|
||||
|
||||
std::free(p);
|
||||
|
||||
if (results == NULL) {
|
||||
fprintf (stderr, "NO results\n");
|
||||
|
||||
@@ -22,7 +22,18 @@ type SDGGML struct {
|
||||
|
||||
var (
|
||||
LoadModel func(model, model_apth string, options []uintptr, threads int32, diff int) int
|
||||
GenImage func(text, negativeText string, width, height, steps int, seed int64, dst string, cfgScale float32, srcImage string, strength float32, maskImage string, refImages []string, refImagesCount int) int
|
||||
GenImage func(params uintptr, steps int, dst string, cfgScale float32, srcImage string, strength float32, maskImage string, refImages []string, refImagesCount int) int
|
||||
|
||||
TilingParamsSetEnabled func(params uintptr, enabled bool)
|
||||
TilingParamsSetTileSizes func(params uintptr, tileSizeX int, tileSizeY int)
|
||||
TilingParamsSetRelSizes func(params uintptr, relSizeX float32, relSizeY float32)
|
||||
TilingParamsSetTargetOverlap func(params uintptr, targetOverlap float32)
|
||||
|
||||
ImgGenParamsNew func() uintptr
|
||||
ImgGenParamsSetPrompts func(params uintptr, prompt string, negativePrompt string)
|
||||
ImgGenParamsSetDimensions func(params uintptr, width int, height int)
|
||||
ImgGenParamsSetSeed func(params uintptr, seed int64)
|
||||
ImgGenParamsGetVaeTilingParams func(params uintptr) uintptr
|
||||
)
|
||||
|
||||
// Copied from Purego internal/strings
|
||||
@@ -120,7 +131,15 @@ func (sd *SDGGML) GenerateImage(opts *pb.GenerateImageRequest) error {
|
||||
// Default strength for img2img (0.75 is a good default)
|
||||
strength := float32(0.75)
|
||||
|
||||
ret := GenImage(t, negative, int(opts.Width), int(opts.Height), int(opts.Step), int64(opts.Seed), dst, sd.cfgScale, srcImage, strength, maskImage, refImages, refImagesCount)
|
||||
// free'd by GenImage
|
||||
p := ImgGenParamsNew()
|
||||
ImgGenParamsSetPrompts(p, t, negative)
|
||||
ImgGenParamsSetDimensions(p, int(opts.Width), int(opts.Height))
|
||||
ImgGenParamsSetSeed(p, int64(opts.Seed))
|
||||
vaep := ImgGenParamsGetVaeTilingParams(p)
|
||||
TilingParamsSetEnabled(vaep, false)
|
||||
|
||||
ret := GenImage(p, int(opts.Step), dst, sd.cfgScale, srcImage, strength, maskImage, refImages, refImagesCount)
|
||||
if ret != 0 {
|
||||
return fmt.Errorf("inference failed")
|
||||
}
|
||||
|
||||
@@ -1,8 +1,23 @@
|
||||
#include <cstdint>
|
||||
#include "stable-diffusion.h"
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
void sd_tiling_params_set_enabled(sd_tiling_params_t *params, bool enabled);
|
||||
void sd_tiling_params_set_tile_sizes(sd_tiling_params_t *params, int tile_size_x, int tile_size_y);
|
||||
void sd_tiling_params_set_rel_sizes(sd_tiling_params_t *params, float rel_size_x, float rel_size_y);
|
||||
void sd_tiling_params_set_target_overlap(sd_tiling_params_t *params, float target_overlap);
|
||||
sd_tiling_params_t* sd_img_gen_params_get_vae_tiling_params(sd_img_gen_params_t *params);
|
||||
|
||||
sd_img_gen_params_t* sd_img_gen_params_new(void);
|
||||
void sd_img_gen_params_set_prompts(sd_img_gen_params_t *params, const char *prompt, const char *negative_prompt);
|
||||
void sd_img_gen_params_set_dimensions(sd_img_gen_params_t *params, int width, int height);
|
||||
void sd_img_gen_params_set_seed(sd_img_gen_params_t *params, int64_t seed);
|
||||
|
||||
int load_model(const char *model, char *model_path, char* options[], int threads, int diffusionModel);
|
||||
int gen_image(char *text, char *negativeText, int width, int height, int steps, int64_t seed, char *dst, float cfg_scale, char *src_image, float strength, char *mask_image, char **ref_images, int ref_images_count);
|
||||
int gen_image(sd_img_gen_params_t *p, int steps, char *dst, float cfg_scale, char *src_image, float strength, char *mask_image, char **ref_images, int ref_images_count);
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -11,14 +11,35 @@ var (
|
||||
addr = flag.String("addr", "localhost:50051", "the address to connect to")
|
||||
)
|
||||
|
||||
type LibFuncs struct {
|
||||
FuncPtr any
|
||||
Name string
|
||||
}
|
||||
|
||||
func main() {
|
||||
gosd, err := purego.Dlopen("./libgosd.so", purego.RTLD_NOW|purego.RTLD_GLOBAL)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
purego.RegisterLibFunc(&LoadModel, gosd, "load_model")
|
||||
purego.RegisterLibFunc(&GenImage, gosd, "gen_image")
|
||||
libFuncs := []LibFuncs{
|
||||
{&LoadModel, "load_model"},
|
||||
{&GenImage, "gen_image"},
|
||||
{&TilingParamsSetEnabled, "sd_tiling_params_set_enabled"},
|
||||
{&TilingParamsSetTileSizes, "sd_tiling_params_set_tile_sizes"},
|
||||
{&TilingParamsSetRelSizes, "sd_tiling_params_set_rel_sizes"},
|
||||
{&TilingParamsSetTargetOverlap, "sd_tiling_params_set_target_overlap"},
|
||||
|
||||
{&ImgGenParamsNew, "sd_img_gen_params_new"},
|
||||
{&ImgGenParamsSetPrompts, "sd_img_gen_params_set_prompts"},
|
||||
{&ImgGenParamsSetDimensions, "sd_img_gen_params_set_dimensions"},
|
||||
{&ImgGenParamsSetSeed, "sd_img_gen_params_set_seed"},
|
||||
{&ImgGenParamsGetVaeTilingParams, "sd_img_gen_params_get_vae_tiling_params"},
|
||||
}
|
||||
|
||||
for _, lf := range libFuncs {
|
||||
purego.RegisterLibFunc(lf.FuncPtr, gosd, lf.Name)
|
||||
}
|
||||
|
||||
flag.Parse()
|
||||
|
||||
|
||||
@@ -10,9 +10,9 @@ CURDIR=$(dirname "$(realpath $0)")
|
||||
# Create lib directory
|
||||
mkdir -p $CURDIR/package/lib
|
||||
|
||||
cp -avrf $CURDIR/libgosd.so $CURDIR/package/
|
||||
cp -avrf $CURDIR/stablediffusion-ggml $CURDIR/package/
|
||||
cp -rfv $CURDIR/run.sh $CURDIR/package/
|
||||
cp -avf $CURDIR/libgosd.so $CURDIR/package/
|
||||
cp -avf $CURDIR/stablediffusion-ggml $CURDIR/package/
|
||||
cp -fv $CURDIR/run.sh $CURDIR/package/
|
||||
|
||||
# Detect architecture and copy appropriate libraries
|
||||
if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then
|
||||
@@ -43,6 +43,8 @@ elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libdl.so.2 $CURDIR/package/lib/libdl.so.2
|
||||
cp -arfLv /lib/aarch64-linux-gnu/librt.so.1 $CURDIR/package/lib/librt.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
elif [ $(uname -s) = "Darwin" ]; then
|
||||
echo "Detected Darwin"
|
||||
else
|
||||
echo "Error: Could not detect architecture"
|
||||
exit 1
|
||||
|
||||
7
backend/go/whisper/.gitignore
vendored
Normal file
7
backend/go/whisper/.gitignore
vendored
Normal file
@@ -0,0 +1,7 @@
|
||||
.cache/
|
||||
sources/
|
||||
build/
|
||||
package/
|
||||
whisper
|
||||
libgowhisper.so
|
||||
|
||||
16
backend/go/whisper/CMakeLists.txt
Normal file
16
backend/go/whisper/CMakeLists.txt
Normal file
@@ -0,0 +1,16 @@
|
||||
cmake_minimum_required(VERSION 3.12)
|
||||
project(gowhisper LANGUAGES C CXX)
|
||||
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
|
||||
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
|
||||
|
||||
add_subdirectory(./sources/whisper.cpp)
|
||||
|
||||
add_library(gowhisper MODULE gowhisper.cpp)
|
||||
target_link_libraries(gowhisper PRIVATE whisper ggml)
|
||||
|
||||
if(CMAKE_CXX_COMPILER_ID MATCHES "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_LESS 9.0)
|
||||
target_link_libraries(gosd PRIVATE stdc++fs)
|
||||
endif()
|
||||
|
||||
set_property(TARGET gowhisper PROPERTY CXX_STANDARD 17)
|
||||
set_target_properties(gowhisper PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
|
||||
@@ -1,110 +1,54 @@
|
||||
GOCMD=go
|
||||
CMAKE_ARGS?=
|
||||
BUILD_TYPE?=
|
||||
NATIVE?=false
|
||||
|
||||
BUILD_TYPE?=
|
||||
CMAKE_ARGS?=
|
||||
GOCMD?=go
|
||||
GO_TAGS?=
|
||||
JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# whisper.cpp version
|
||||
WHISPER_REPO?=https://github.com/ggml-org/whisper.cpp
|
||||
WHISPER_CPP_VERSION?=fc45bb86251f774ef817e89878bb4c2636c8a58f
|
||||
WHISPER_CPP_VERSION?=f16c12f3f55f5bd3d6ac8cf2f31ab90a42c884d5
|
||||
SO_TARGET?=libgowhisper.so
|
||||
|
||||
export WHISPER_CMAKE_ARGS?=-DBUILD_SHARED_LIBS=OFF
|
||||
export WHISPER_DIR=$(abspath ./sources/whisper.cpp)
|
||||
export WHISPER_INCLUDE_PATH=$(WHISPER_DIR)/include:$(WHISPER_DIR)/ggml/include
|
||||
export WHISPER_LIBRARY_PATH=$(WHISPER_DIR)/build/src/:$(WHISPER_DIR)/build/ggml/src
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
CGO_LDFLAGS_WHISPER?=
|
||||
CGO_LDFLAGS_WHISPER+=-lggml
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF -DLLAMA_CURL=OFF
|
||||
CUDA_LIBPATH?=/usr/local/cuda/lib64/
|
||||
|
||||
ONEAPI_VERSION?=2025.2
|
||||
|
||||
# IF native is false, we add -DGGML_NATIVE=OFF to CMAKE_ARGS
|
||||
ifeq ($(NATIVE),false)
|
||||
CMAKE_ARGS+=-DGGML_NATIVE=OFF
|
||||
WHISPER_CMAKE_ARGS+=-DGGML_NATIVE=OFF
|
||||
endif
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
ifeq ($(NATIVE),false)
|
||||
CMAKE_ARGS+=-DGGML_NATIVE=OFF
|
||||
endif
|
||||
# If build type is cublas, then we set -DGGML_CUDA=ON to CMAKE_ARGS automatically
|
||||
|
||||
ifeq ($(BUILD_TYPE),cublas)
|
||||
CGO_LDFLAGS+=-lcublas -lcudart -L$(CUDA_LIBPATH) -L$(CUDA_LIBPATH)/stubs/ -lcuda
|
||||
CMAKE_ARGS+=-DGGML_CUDA=ON
|
||||
CGO_LDFLAGS_WHISPER+=-lcufft -lggml-cuda
|
||||
export WHISPER_LIBRARY_PATH:=$(WHISPER_LIBRARY_PATH):$(WHISPER_DIR)/build/ggml/src/ggml-cuda/
|
||||
# If build type is openblas then we set -DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS
|
||||
# to CMAKE_ARGS automatically
|
||||
else ifeq ($(BUILD_TYPE),openblas)
|
||||
CMAKE_ARGS+=-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS
|
||||
# If build type is clblas (openCL) we set -DGGML_CLBLAST=ON -DCLBlast_DIR=/some/path
|
||||
else ifeq ($(BUILD_TYPE),clblas)
|
||||
CMAKE_ARGS+=-DGGML_CLBLAST=ON -DCLBlast_DIR=/some/path
|
||||
# If it's hipblas we do have also to set CC=/opt/rocm/llvm/bin/clang CXX=/opt/rocm/llvm/bin/clang++
|
||||
else ifeq ($(BUILD_TYPE),hipblas)
|
||||
ROCM_HOME ?= /opt/rocm
|
||||
ROCM_PATH ?= /opt/rocm
|
||||
LD_LIBRARY_PATH ?= /opt/rocm/lib:/opt/rocm/llvm/lib
|
||||
export STABLE_BUILD_TYPE=
|
||||
export CXX=$(ROCM_HOME)/llvm/bin/clang++
|
||||
export CC=$(ROCM_HOME)/llvm/bin/clang
|
||||
# GPU_TARGETS ?= gfx803,gfx900,gfx906,gfx908,gfx90a,gfx942,gfx1010,gfx1030,gfx1032,gfx1100,gfx1101,gfx1102
|
||||
# AMDGPU_TARGETS ?= "$(GPU_TARGETS)"
|
||||
CMAKE_ARGS+=-DGGML_HIP=ON
|
||||
CGO_LDFLAGS += -O3 --rtlib=compiler-rt -unwindlib=libgcc -lhipblas -lrocblas --hip-link -L${ROCM_HOME}/lib/llvm/lib -L$(CURRENT_MAKEFILE_DIR)/sources/whisper.cpp/build/ggml/src/ggml-hip/ -lggml-hip
|
||||
# CMAKE_ARGS+=-DGGML_HIP=ON -DAMDGPU_TARGETS="$(AMDGPU_TARGETS)" -DGPU_TARGETS="$(GPU_TARGETS)"
|
||||
CMAKE_ARGS+=-DGGML_HIPBLAS=ON
|
||||
else ifeq ($(BUILD_TYPE),vulkan)
|
||||
CMAKE_ARGS+=-DGGML_VULKAN=1
|
||||
CGO_LDFLAGS_WHISPER+=-lggml-vulkan -lvulkan
|
||||
export WHISPER_LIBRARY_PATH:=$(WHISPER_LIBRARY_PATH):$(WHISPER_DIR)/build/ggml/src/ggml-vulkan/
|
||||
CMAKE_ARGS+=-DGGML_VULKAN=ON
|
||||
else ifeq ($(OS),Darwin)
|
||||
ifeq ($(BUILD_TYPE),)
|
||||
BUILD_TYPE=metal
|
||||
endif
|
||||
ifneq ($(BUILD_TYPE),metal)
|
||||
CMAKE_ARGS+=-DGGML_METAL=OFF
|
||||
CGO_LDFLAGS_WHISPER+=-lggml-blas
|
||||
export WHISPER_LIBRARY_PATH:=$(WHISPER_LIBRARY_PATH):$(WHISPER_DIR)/build/ggml/src/ggml-blas
|
||||
else
|
||||
CMAKE_ARGS+=-DGGML_METAL=ON
|
||||
CMAKE_ARGS+=-DGGML_METAL_EMBED_LIBRARY=ON
|
||||
CMAKE_ARGS+=-DGGML_METAL_USE_BF16=ON
|
||||
CMAKE_ARGS+=-DGGML_OPENMP=OFF
|
||||
CMAKE_ARGS+=-DWHISPER_BUILD_EXAMPLES=OFF
|
||||
CMAKE_ARGS+=-DWHISPER_BUILD_TESTS=OFF
|
||||
CMAKE_ARGS+=-DWHISPER_BUILD_SERVER=OFF
|
||||
CGO_LDFLAGS += -framework Accelerate
|
||||
CGO_LDFLAGS_WHISPER+=-lggml-metal -lggml-blas
|
||||
export WHISPER_LIBRARY_PATH:=$(WHISPER_LIBRARY_PATH):$(WHISPER_DIR)/build/ggml/src/ggml-metal/:$(WHISPER_DIR)/build/ggml/src/ggml-blas
|
||||
endif
|
||||
TARGET+=--target ggml-metal
|
||||
endif
|
||||
|
||||
ifneq (,$(findstring sycl,$(BUILD_TYPE)))
|
||||
export CC=icx
|
||||
export CXX=icpx
|
||||
CGO_LDFLAGS_WHISPER += -fsycl -L${DNNLROOT}/lib -rpath ${ONEAPI_ROOT}/${ONEAPI_VERSION}/lib -ldnnl ${MKLROOT}/lib/intel64/libmkl_sycl.a -fiopenmp -fopenmp-targets=spir64 -lOpenCL -lggml-sycl
|
||||
CGO_LDFLAGS_WHISPER += $(shell pkg-config --libs mkl-static-lp64-gomp)
|
||||
CGO_CXXFLAGS_WHISPER += -fiopenmp -fopenmp-targets=spir64
|
||||
CGO_CXXFLAGS_WHISPER += $(shell pkg-config --cflags mkl-static-lp64-gomp )
|
||||
export WHISPER_LIBRARY_PATH:=$(WHISPER_LIBRARY_PATH):$(WHISPER_DIR)/build/ggml/src/ggml-sycl/
|
||||
CMAKE_ARGS+=-DGGML_SYCL=ON \
|
||||
-DCMAKE_C_COMPILER=icx \
|
||||
-DCMAKE_CXX_COMPILER=icpx \
|
||||
-DCMAKE_CXX_FLAGS="-fsycl"
|
||||
endif
|
||||
|
||||
ifeq ($(BUILD_TYPE),sycl_f16)
|
||||
CMAKE_ARGS+=-DGGML_SYCL_F16=ON
|
||||
CMAKE_ARGS+=-DGGML_SYCL=ON \
|
||||
-DCMAKE_C_COMPILER=icx \
|
||||
-DCMAKE_CXX_COMPILER=icpx \
|
||||
-DGGML_SYCL_F16=ON
|
||||
endif
|
||||
|
||||
ifneq ($(OS),Darwin)
|
||||
CGO_LDFLAGS_WHISPER+=-lgomp
|
||||
ifeq ($(BUILD_TYPE),sycl_f32)
|
||||
CMAKE_ARGS+=-DGGML_SYCL=ON \
|
||||
-DCMAKE_C_COMPILER=icx \
|
||||
-DCMAKE_CXX_COMPILER=icpx
|
||||
endif
|
||||
|
||||
## whisper
|
||||
sources/whisper.cpp:
|
||||
mkdir -p sources/whisper.cpp
|
||||
cd sources/whisper.cpp && \
|
||||
@@ -114,18 +58,65 @@ sources/whisper.cpp:
|
||||
git checkout $(WHISPER_CPP_VERSION) && \
|
||||
git submodule update --init --recursive --depth 1 --single-branch
|
||||
|
||||
sources/whisper.cpp/build/src/libwhisper.a: sources/whisper.cpp
|
||||
cd sources/whisper.cpp && cmake $(CMAKE_ARGS) $(WHISPER_CMAKE_ARGS) . -B ./build
|
||||
cd sources/whisper.cpp/build && cmake --build . --config Release
|
||||
# Detect OS
|
||||
UNAME_S := $(shell uname -s)
|
||||
|
||||
whisper: sources/whisper.cpp sources/whisper.cpp/build/src/libwhisper.a
|
||||
$(GOCMD) mod edit -replace github.com/ggerganov/whisper.cpp=$(CURDIR)/sources/whisper.cpp
|
||||
$(GOCMD) mod edit -replace github.com/ggerganov/whisper.cpp/bindings/go=$(CURDIR)/sources/whisper.cpp/bindings/go
|
||||
CGO_LDFLAGS="$(CGO_LDFLAGS) $(CGO_LDFLAGS_WHISPER)" C_INCLUDE_PATH="${WHISPER_INCLUDE_PATH}" LIBRARY_PATH="${WHISPER_LIBRARY_PATH}" LD_LIBRARY_PATH="${WHISPER_LIBRARY_PATH}" \
|
||||
CGO_CXXFLAGS="$(CGO_CXXFLAGS_WHISPER)" \
|
||||
$(GOCMD) build -ldflags "$(LD_FLAGS)" -tags "$(GO_TAGS)" -o whisper ./
|
||||
# Only build CPU variants on Linux
|
||||
ifeq ($(UNAME_S),Linux)
|
||||
VARIANT_TARGETS = libgowhisper-avx.so libgowhisper-avx2.so libgowhisper-avx512.so libgowhisper-fallback.so
|
||||
else
|
||||
# On non-Linux (e.g., Darwin), build only fallback variant
|
||||
VARIANT_TARGETS = libgowhisper-fallback.so
|
||||
endif
|
||||
|
||||
package:
|
||||
whisper: main.go gowhisper.go $(VARIANT_TARGETS)
|
||||
CGO_ENABLED=0 $(GOCMD) build -tags "$(GO_TAGS)" -o whisper ./
|
||||
|
||||
package: whisper
|
||||
bash package.sh
|
||||
|
||||
build: whisper package
|
||||
build: package
|
||||
|
||||
clean: purge
|
||||
rm -rf libgowhisper*.so sources/whisper.cpp whisper
|
||||
|
||||
purge:
|
||||
rm -rf build*
|
||||
|
||||
# Build all variants (Linux only)
|
||||
ifeq ($(UNAME_S),Linux)
|
||||
libgowhisper-avx.so: sources/whisper.cpp
|
||||
$(MAKE) purge
|
||||
$(info ${GREEN}I whisper build info:avx${RESET})
|
||||
SO_TARGET=libgowhisper-avx.so CMAKE_ARGS="$(CMAKE_ARGS) -DGGML_AVX=on -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off" $(MAKE) libgowhisper-custom
|
||||
rm -rfv build*
|
||||
|
||||
libgowhisper-avx2.so: sources/whisper.cpp
|
||||
$(MAKE) purge
|
||||
$(info ${GREEN}I whisper build info:avx2${RESET})
|
||||
SO_TARGET=libgowhisper-avx2.so CMAKE_ARGS="$(CMAKE_ARGS) -DGGML_AVX=on -DGGML_AVX2=on -DGGML_AVX512=off -DGGML_FMA=on -DGGML_F16C=on" $(MAKE) libgowhisper-custom
|
||||
rm -rfv build*
|
||||
|
||||
libgowhisper-avx512.so: sources/whisper.cpp
|
||||
$(MAKE) purge
|
||||
$(info ${GREEN}I whisper build info:avx512${RESET})
|
||||
SO_TARGET=libgowhisper-avx512.so CMAKE_ARGS="$(CMAKE_ARGS) -DGGML_AVX=on -DGGML_AVX2=off -DGGML_AVX512=on -DGGML_FMA=on -DGGML_F16C=on" $(MAKE) libgowhisper-custom
|
||||
rm -rfv build*
|
||||
endif
|
||||
|
||||
# Build fallback variant (all platforms)
|
||||
libgowhisper-fallback.so: sources/whisper.cpp
|
||||
$(MAKE) purge
|
||||
$(info ${GREEN}I whisper build info:fallback${RESET})
|
||||
SO_TARGET=libgowhisper-fallback.so CMAKE_ARGS="$(CMAKE_ARGS) -DGGML_AVX=off -DGGML_AVX2=off -DGGML_AVX512=off -DGGML_FMA=off -DGGML_F16C=off" $(MAKE) libgowhisper-custom
|
||||
rm -rfv build*
|
||||
|
||||
libgowhisper-custom: CMakeLists.txt gowhisper.cpp gowhisper.h
|
||||
mkdir -p build-$(SO_TARGET) && \
|
||||
cd build-$(SO_TARGET) && \
|
||||
cmake .. $(CMAKE_ARGS) && \
|
||||
cmake --build . --config Release -j$(JOBS) && \
|
||||
cd .. && \
|
||||
mv build-$(SO_TARGET)/libgowhisper.so ./$(SO_TARGET)
|
||||
|
||||
all: whisper package
|
||||
|
||||
154
backend/go/whisper/gowhisper.cpp
Normal file
154
backend/go/whisper/gowhisper.cpp
Normal file
@@ -0,0 +1,154 @@
|
||||
#include "gowhisper.h"
|
||||
#include "ggml-backend.h"
|
||||
#include "whisper.h"
|
||||
#include <vector>
|
||||
|
||||
static struct whisper_vad_context *vctx;
|
||||
static struct whisper_context *ctx;
|
||||
static std::vector<float> flat_segs;
|
||||
|
||||
static void ggml_log_cb(enum ggml_log_level level, const char *log,
|
||||
void *data) {
|
||||
const char *level_str;
|
||||
|
||||
if (!log) {
|
||||
return;
|
||||
}
|
||||
|
||||
switch (level) {
|
||||
case GGML_LOG_LEVEL_DEBUG:
|
||||
level_str = "DEBUG";
|
||||
break;
|
||||
case GGML_LOG_LEVEL_INFO:
|
||||
level_str = "INFO";
|
||||
break;
|
||||
case GGML_LOG_LEVEL_WARN:
|
||||
level_str = "WARN";
|
||||
break;
|
||||
case GGML_LOG_LEVEL_ERROR:
|
||||
level_str = "ERROR";
|
||||
break;
|
||||
default: /* Potential future-proofing */
|
||||
level_str = "?????";
|
||||
break;
|
||||
}
|
||||
|
||||
fprintf(stderr, "[%-5s] ", level_str);
|
||||
fputs(log, stderr);
|
||||
fflush(stderr);
|
||||
}
|
||||
|
||||
int load_model(const char *const model_path) {
|
||||
whisper_log_set(ggml_log_cb, nullptr);
|
||||
ggml_backend_load_all();
|
||||
|
||||
struct whisper_context_params cparams = whisper_context_default_params();
|
||||
|
||||
ctx = whisper_init_from_file_with_params(model_path, cparams);
|
||||
if (ctx == nullptr) {
|
||||
fprintf(stderr, "error: Also failed to init model as transcriber\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int load_model_vad(const char *const model_path) {
|
||||
whisper_log_set(ggml_log_cb, nullptr);
|
||||
ggml_backend_load_all();
|
||||
|
||||
struct whisper_vad_context_params vcparams =
|
||||
whisper_vad_default_context_params();
|
||||
|
||||
// XXX: Overridden to false in upstream due to performance?
|
||||
// vcparams.use_gpu = true;
|
||||
|
||||
vctx = whisper_vad_init_from_file_with_params(model_path, vcparams);
|
||||
if (vctx == nullptr) {
|
||||
fprintf(stderr, "error: Failed to init model as VAD\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int vad(float pcmf32[], size_t pcmf32_len, float **segs_out,
|
||||
size_t *segs_out_len) {
|
||||
if (!whisper_vad_detect_speech(vctx, pcmf32, pcmf32_len)) {
|
||||
fprintf(stderr, "error: failed to detect speech\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
struct whisper_vad_params params = whisper_vad_default_params();
|
||||
struct whisper_vad_segments *segs =
|
||||
whisper_vad_segments_from_probs(vctx, params);
|
||||
size_t segn = whisper_vad_segments_n_segments(segs);
|
||||
|
||||
// fprintf(stderr, "Got segments %zd\n", segn);
|
||||
|
||||
flat_segs.clear();
|
||||
|
||||
for (int i = 0; i < segn; i++) {
|
||||
flat_segs.push_back(whisper_vad_segments_get_segment_t0(segs, i));
|
||||
flat_segs.push_back(whisper_vad_segments_get_segment_t1(segs, i));
|
||||
}
|
||||
|
||||
// fprintf(stderr, "setting out variables: %p=%p -> %p, %p=%zx -> %zx\n",
|
||||
// segs_out, *segs_out, flat_segs.data(), segs_out_len, *segs_out_len,
|
||||
// flat_segs.size());
|
||||
*segs_out = flat_segs.data();
|
||||
*segs_out_len = flat_segs.size();
|
||||
|
||||
// fprintf(stderr, "freeing segs\n");
|
||||
whisper_vad_free_segments(segs);
|
||||
|
||||
// fprintf(stderr, "returning\n");
|
||||
return 0;
|
||||
}
|
||||
|
||||
int transcribe(uint32_t threads, char *lang, bool translate, bool tdrz,
|
||||
float pcmf32[], size_t pcmf32_len, size_t *segs_out_len) {
|
||||
whisper_full_params wparams =
|
||||
whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
|
||||
|
||||
wparams.n_threads = threads;
|
||||
if (*lang != '\0')
|
||||
wparams.language = lang;
|
||||
else {
|
||||
wparams.language = nullptr;
|
||||
}
|
||||
|
||||
wparams.translate = translate;
|
||||
wparams.debug_mode = true;
|
||||
wparams.print_progress = true;
|
||||
wparams.tdrz_enable = tdrz;
|
||||
|
||||
fprintf(stderr, "info: Enable tdrz: %d\n", tdrz);
|
||||
|
||||
if (whisper_full(ctx, wparams, pcmf32, pcmf32_len)) {
|
||||
fprintf(stderr, "error: transcription failed\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
*segs_out_len = whisper_full_n_segments(ctx);
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
const char *get_segment_text(int i) {
|
||||
return whisper_full_get_segment_text(ctx, i);
|
||||
}
|
||||
|
||||
int64_t get_segment_t0(int i) { return whisper_full_get_segment_t0(ctx, i); }
|
||||
|
||||
int64_t get_segment_t1(int i) { return whisper_full_get_segment_t1(ctx, i); }
|
||||
|
||||
int n_tokens(int i) { return whisper_full_n_tokens(ctx, i); }
|
||||
|
||||
int32_t get_token_id(int i, int j) {
|
||||
return whisper_full_get_token_id(ctx, i, j);
|
||||
}
|
||||
|
||||
bool get_segment_speaker_turn_next(int i) {
|
||||
return whisper_full_get_segment_speaker_turn_next(ctx, i);
|
||||
}
|
||||
161
backend/go/whisper/gowhisper.go
Normal file
161
backend/go/whisper/gowhisper.go
Normal file
@@ -0,0 +1,161 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"unsafe"
|
||||
|
||||
"github.com/go-audio/wav"
|
||||
"github.com/mudler/LocalAI/pkg/grpc/base"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/LocalAI/pkg/utils"
|
||||
)
|
||||
|
||||
var (
|
||||
CppLoadModel func(modelPath string) int
|
||||
CppLoadModelVAD func(modelPath string) int
|
||||
CppVAD func(pcmf32 []float32, pcmf32Size uintptr, segsOut unsafe.Pointer, segsOutLen unsafe.Pointer) int
|
||||
CppTranscribe func(threads uint32, lang string, translate bool, diarize bool, pcmf32 []float32, pcmf32Len uintptr, segsOutLen unsafe.Pointer) int
|
||||
CppGetSegmentText func(i int) string
|
||||
CppGetSegmentStart func(i int) int64
|
||||
CppGetSegmentEnd func(i int) int64
|
||||
CppNTokens func(i int) int
|
||||
CppGetTokenID func(i int, j int) int
|
||||
CppGetSegmentSpeakerTurnNext func(i int) bool
|
||||
)
|
||||
|
||||
type Whisper struct {
|
||||
base.SingleThread
|
||||
}
|
||||
|
||||
func (w *Whisper) Load(opts *pb.ModelOptions) error {
|
||||
vadOnly := false
|
||||
|
||||
for _, oo := range opts.Options {
|
||||
if oo == "vad_only" {
|
||||
vadOnly = true
|
||||
} else {
|
||||
fmt.Fprintf(os.Stderr, "Unrecognized option: %v\n", oo)
|
||||
}
|
||||
}
|
||||
|
||||
if vadOnly {
|
||||
if ret := CppLoadModelVAD(opts.ModelFile); ret != 0 {
|
||||
return fmt.Errorf("Failed to load Whisper VAD model")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
if ret := CppLoadModel(opts.ModelFile); ret != 0 {
|
||||
return fmt.Errorf("Failed to load Whisper transcription model")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *Whisper) VAD(req *pb.VADRequest) (pb.VADResponse, error) {
|
||||
audio := req.Audio
|
||||
// We expect 0xdeadbeef to be overwritten and if we see it in a stack trace we know it wasn't
|
||||
segsPtr, segsLen := uintptr(0xdeadbeef), uintptr(0xdeadbeef)
|
||||
segsPtrPtr, segsLenPtr := unsafe.Pointer(&segsPtr), unsafe.Pointer(&segsLen)
|
||||
|
||||
if ret := CppVAD(audio, uintptr(len(audio)), segsPtrPtr, segsLenPtr); ret != 0 {
|
||||
return pb.VADResponse{}, fmt.Errorf("Failed VAD")
|
||||
}
|
||||
|
||||
// Happens when CPP vector has not had any elements pushed to it
|
||||
if segsPtr == 0 {
|
||||
return pb.VADResponse{
|
||||
Segments: []*pb.VADSegment{},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// unsafeptr warning is caused by segsPtr being on the stack and therefor being subject to stack copying AFAICT
|
||||
// however the stack shouldn't have grown between setting segsPtr and now, also the memory pointed to is allocated by C++
|
||||
segs := unsafe.Slice((*float32)(unsafe.Pointer(segsPtr)), segsLen)
|
||||
|
||||
vadSegments := []*pb.VADSegment{}
|
||||
for i := range len(segs) >> 1 {
|
||||
s := segs[2*i] / 100
|
||||
t := segs[2*i+1] / 100
|
||||
vadSegments = append(vadSegments, &pb.VADSegment{
|
||||
Start: s,
|
||||
End: t,
|
||||
})
|
||||
}
|
||||
|
||||
return pb.VADResponse{
|
||||
Segments: vadSegments,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (w *Whisper) AudioTranscription(opts *pb.TranscriptRequest) (pb.TranscriptResult, error) {
|
||||
dir, err := os.MkdirTemp("", "whisper")
|
||||
if err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
convertedPath := filepath.Join(dir, "converted.wav")
|
||||
|
||||
if err := utils.AudioToWav(opts.Dst, convertedPath); err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
}
|
||||
|
||||
// Open samples
|
||||
fh, err := os.Open(convertedPath)
|
||||
if err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
}
|
||||
defer fh.Close()
|
||||
|
||||
// Read samples
|
||||
d := wav.NewDecoder(fh)
|
||||
buf, err := d.FullPCMBuffer()
|
||||
if err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
}
|
||||
|
||||
data := buf.AsFloat32Buffer().Data
|
||||
segsLen := uintptr(0xdeadbeef)
|
||||
segsLenPtr := unsafe.Pointer(&segsLen)
|
||||
|
||||
if ret := CppTranscribe(opts.Threads, opts.Language, opts.Translate, opts.Diarize, data, uintptr(len(data)), segsLenPtr); ret != 0 {
|
||||
return pb.TranscriptResult{}, fmt.Errorf("Failed Transcribe")
|
||||
}
|
||||
|
||||
segments := []*pb.TranscriptSegment{}
|
||||
text := ""
|
||||
for i := range int(segsLen) {
|
||||
s := CppGetSegmentStart(i)
|
||||
t := CppGetSegmentEnd(i)
|
||||
txt := strings.Clone(CppGetSegmentText(i))
|
||||
tokens := make([]int32, CppNTokens(i))
|
||||
|
||||
if opts.Diarize && CppGetSegmentSpeakerTurnNext(i) {
|
||||
txt += " [SPEAKER_TURN]"
|
||||
}
|
||||
|
||||
for j := range tokens {
|
||||
tokens[j] = int32(CppGetTokenID(i, j))
|
||||
}
|
||||
segment := &pb.TranscriptSegment{
|
||||
Id: int32(i),
|
||||
Text: txt,
|
||||
Start: s, End: t,
|
||||
Tokens: tokens,
|
||||
}
|
||||
|
||||
segments = append(segments, segment)
|
||||
|
||||
text += " " + strings.TrimSpace(txt)
|
||||
}
|
||||
|
||||
return pb.TranscriptResult{
|
||||
Segments: segments,
|
||||
Text: strings.TrimSpace(text),
|
||||
}, nil
|
||||
}
|
||||
17
backend/go/whisper/gowhisper.h
Normal file
17
backend/go/whisper/gowhisper.h
Normal file
@@ -0,0 +1,17 @@
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
extern "C" {
|
||||
int load_model(const char *const model_path);
|
||||
int load_model_vad(const char *const model_path);
|
||||
int vad(float pcmf32[], size_t pcmf32_size, float **segs_out,
|
||||
size_t *segs_out_len);
|
||||
int transcribe(uint32_t threads, char *lang, bool translate, bool tdrz,
|
||||
float pcmf32[], size_t pcmf32_len, size_t *segs_out_len);
|
||||
const char *get_segment_text(int i);
|
||||
int64_t get_segment_t0(int i);
|
||||
int64_t get_segment_t1(int i);
|
||||
int n_tokens(int i);
|
||||
int32_t get_token_id(int i, int j);
|
||||
bool get_segment_speaker_turn_next(int i);
|
||||
}
|
||||
@@ -1,10 +1,11 @@
|
||||
package main
|
||||
|
||||
// Note: this is started internally by LocalAI and a server is allocated for each model
|
||||
|
||||
import (
|
||||
"flag"
|
||||
"os"
|
||||
|
||||
"github.com/ebitengine/purego"
|
||||
grpc "github.com/mudler/LocalAI/pkg/grpc"
|
||||
)
|
||||
|
||||
@@ -12,7 +13,40 @@ var (
|
||||
addr = flag.String("addr", "localhost:50051", "the address to connect to")
|
||||
)
|
||||
|
||||
type LibFuncs struct {
|
||||
FuncPtr any
|
||||
Name string
|
||||
}
|
||||
|
||||
func main() {
|
||||
// Get library name from environment variable, default to fallback
|
||||
libName := os.Getenv("WHISPER_LIBRARY")
|
||||
if libName == "" {
|
||||
libName = "./libgowhisper-fallback.so"
|
||||
}
|
||||
|
||||
gosd, err := purego.Dlopen(libName, purego.RTLD_NOW|purego.RTLD_GLOBAL)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
libFuncs := []LibFuncs{
|
||||
{&CppLoadModel, "load_model"},
|
||||
{&CppLoadModelVAD, "load_model_vad"},
|
||||
{&CppVAD, "vad"},
|
||||
{&CppTranscribe, "transcribe"},
|
||||
{&CppGetSegmentText, "get_segment_text"},
|
||||
{&CppGetSegmentStart, "get_segment_t0"},
|
||||
{&CppGetSegmentEnd, "get_segment_t1"},
|
||||
{&CppNTokens, "n_tokens"},
|
||||
{&CppGetTokenID, "get_token_id"},
|
||||
{&CppGetSegmentSpeakerTurnNext, "get_segment_speaker_turn_next"},
|
||||
}
|
||||
|
||||
for _, lf := range libFuncs {
|
||||
purego.RegisterLibFunc(lf.FuncPtr, gosd, lf.Name)
|
||||
}
|
||||
|
||||
flag.Parse()
|
||||
|
||||
if err := grpc.StartServer(*addr, &Whisper{}); err != nil {
|
||||
|
||||
@@ -10,8 +10,9 @@ CURDIR=$(dirname "$(realpath $0)")
|
||||
# Create lib directory
|
||||
mkdir -p $CURDIR/package/lib
|
||||
|
||||
cp -avrf $CURDIR/whisper $CURDIR/package/
|
||||
cp -rfv $CURDIR/run.sh $CURDIR/package/
|
||||
cp -avf $CURDIR/whisper $CURDIR/package/
|
||||
cp -fv $CURDIR/libgowhisper-*.so $CURDIR/package/
|
||||
cp -fv $CURDIR/run.sh $CURDIR/package/
|
||||
|
||||
# Detect architecture and copy appropriate libraries
|
||||
if [ -f "/lib64/ld-linux-x86-64.so.2" ]; then
|
||||
@@ -42,11 +43,13 @@ elif [ -f "/lib/ld-linux-aarch64.so.1" ]; then
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libdl.so.2 $CURDIR/package/lib/libdl.so.2
|
||||
cp -arfLv /lib/aarch64-linux-gnu/librt.so.1 $CURDIR/package/lib/librt.so.1
|
||||
cp -arfLv /lib/aarch64-linux-gnu/libpthread.so.0 $CURDIR/package/lib/libpthread.so.0
|
||||
elif [ $(uname -s) = "Darwin" ]; then
|
||||
echo "Detected Darwin"
|
||||
else
|
||||
echo "Error: Could not detect architecture"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Packaging completed successfully"
|
||||
echo "Packaging completed successfully"
|
||||
ls -liah $CURDIR/package/
|
||||
ls -liah $CURDIR/package/lib/
|
||||
ls -liah $CURDIR/package/lib/
|
||||
|
||||
@@ -1,14 +1,52 @@
|
||||
#!/bin/bash
|
||||
set -ex
|
||||
|
||||
# Get the absolute current dir where the script is located
|
||||
CURDIR=$(dirname "$(realpath $0)")
|
||||
|
||||
cd /
|
||||
|
||||
echo "CPU info:"
|
||||
if [ "$(uname)" != "Darwin" ]; then
|
||||
grep -e "model\sname" /proc/cpuinfo | head -1
|
||||
grep -e "flags" /proc/cpuinfo | head -1
|
||||
fi
|
||||
|
||||
LIBRARY="$CURDIR/libgowhisper-fallback.so"
|
||||
|
||||
if [ "$(uname)" != "Darwin" ]; then
|
||||
if grep -q -e "\savx\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX found OK"
|
||||
if [ -e $CURDIR/libgowhisper-avx.so ]; then
|
||||
LIBRARY="$CURDIR/libgowhisper-avx.so"
|
||||
fi
|
||||
fi
|
||||
|
||||
if grep -q -e "\savx2\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX2 found OK"
|
||||
if [ -e $CURDIR/libgowhisper-avx2.so ]; then
|
||||
LIBRARY="$CURDIR/libgowhisper-avx2.so"
|
||||
fi
|
||||
fi
|
||||
|
||||
# Check avx 512
|
||||
if grep -q -e "\savx512f\s" /proc/cpuinfo ; then
|
||||
echo "CPU: AVX512F found OK"
|
||||
if [ -e $CURDIR/libgowhisper-avx512.so ]; then
|
||||
LIBRARY="$CURDIR/libgowhisper-avx512.so"
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
export LD_LIBRARY_PATH=$CURDIR/lib:$LD_LIBRARY_PATH
|
||||
export WHISPER_LIBRARY=$LIBRARY
|
||||
|
||||
# If there is a lib/ld.so, use it
|
||||
if [ -f $CURDIR/lib/ld.so ]; then
|
||||
echo "Using lib/ld.so"
|
||||
echo "Using library: $LIBRARY"
|
||||
exec $CURDIR/lib/ld.so $CURDIR/whisper "$@"
|
||||
fi
|
||||
|
||||
echo "Using library: $LIBRARY"
|
||||
exec $CURDIR/whisper "$@"
|
||||
@@ -1,105 +0,0 @@
|
||||
package main
|
||||
|
||||
// This is a wrapper to statisfy the GRPC service interface
|
||||
// It is meant to be used by the main executable that is the server for the specific backend type (falcon, gpt3, etc)
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper"
|
||||
"github.com/go-audio/wav"
|
||||
"github.com/mudler/LocalAI/pkg/grpc/base"
|
||||
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
|
||||
"github.com/mudler/LocalAI/pkg/utils"
|
||||
)
|
||||
|
||||
type Whisper struct {
|
||||
base.SingleThread
|
||||
whisper whisper.Model
|
||||
}
|
||||
|
||||
func (sd *Whisper) Load(opts *pb.ModelOptions) error {
|
||||
// Note: the Model here is a path to a directory containing the model files
|
||||
w, err := whisper.New(opts.ModelFile)
|
||||
sd.whisper = w
|
||||
return err
|
||||
}
|
||||
|
||||
func (sd *Whisper) AudioTranscription(opts *pb.TranscriptRequest) (pb.TranscriptResult, error) {
|
||||
|
||||
dir, err := os.MkdirTemp("", "whisper")
|
||||
if err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
}
|
||||
defer os.RemoveAll(dir)
|
||||
|
||||
convertedPath := filepath.Join(dir, "converted.wav")
|
||||
|
||||
if err := utils.AudioToWav(opts.Dst, convertedPath); err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
}
|
||||
|
||||
// Open samples
|
||||
fh, err := os.Open(convertedPath)
|
||||
if err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
}
|
||||
defer fh.Close()
|
||||
|
||||
// Read samples
|
||||
d := wav.NewDecoder(fh)
|
||||
buf, err := d.FullPCMBuffer()
|
||||
if err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
}
|
||||
|
||||
data := buf.AsFloat32Buffer().Data
|
||||
|
||||
// Process samples
|
||||
context, err := sd.whisper.NewContext()
|
||||
if err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
|
||||
}
|
||||
|
||||
context.SetThreads(uint(opts.Threads))
|
||||
|
||||
if opts.Language != "" {
|
||||
context.SetLanguage(opts.Language)
|
||||
} else {
|
||||
context.SetLanguage("auto")
|
||||
}
|
||||
|
||||
if opts.Translate {
|
||||
context.SetTranslate(true)
|
||||
}
|
||||
|
||||
if err := context.Process(data, nil, nil, nil); err != nil {
|
||||
return pb.TranscriptResult{}, err
|
||||
}
|
||||
|
||||
segments := []*pb.TranscriptSegment{}
|
||||
text := ""
|
||||
for {
|
||||
s, err := context.NextSegment()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
|
||||
var tokens []int32
|
||||
for _, t := range s.Tokens {
|
||||
tokens = append(tokens, int32(t.Id))
|
||||
}
|
||||
|
||||
segment := &pb.TranscriptSegment{Id: int32(s.Num), Text: s.Text, Start: int64(s.Start), End: int64(s.End), Tokens: tokens}
|
||||
segments = append(segments, segment)
|
||||
|
||||
text += s.Text
|
||||
}
|
||||
|
||||
return pb.TranscriptResult{
|
||||
Segments: segments,
|
||||
Text: text,
|
||||
}, nil
|
||||
|
||||
}
|
||||
@@ -45,6 +45,7 @@
|
||||
default: "cpu-whisper"
|
||||
nvidia: "cuda12-whisper"
|
||||
intel: "intel-sycl-f16-whisper"
|
||||
metal: "metal-whisper"
|
||||
amd: "rocm-whisper"
|
||||
vulkan: "vulkan-whisper"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-whisper"
|
||||
@@ -71,7 +72,7 @@
|
||||
# amd: "rocm-stablediffusion-ggml"
|
||||
vulkan: "vulkan-stablediffusion-ggml"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-stablediffusion-ggml"
|
||||
# metal: "metal-stablediffusion-ggml"
|
||||
metal: "metal-stablediffusion-ggml"
|
||||
# darwin-x86: "darwin-x86-stablediffusion-ggml"
|
||||
- &rfdetr
|
||||
name: "rfdetr"
|
||||
@@ -147,7 +148,7 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-mlx-vlm"
|
||||
icon: https://avatars.githubusercontent.com/u/102832242?s=200&v=4
|
||||
urls:
|
||||
- https://github.com/ml-explore/mlx-vlm
|
||||
- https://github.com/Blaizzy/mlx-vlm
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-metal-darwin-arm64-mlx-vlm
|
||||
license: MIT
|
||||
@@ -159,6 +160,23 @@
|
||||
- vision-language
|
||||
- LLM
|
||||
- MLX
|
||||
- &mlx-audio
|
||||
name: "mlx-audio"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-mlx-audio"
|
||||
icon: https://avatars.githubusercontent.com/u/102832242?s=200&v=4
|
||||
urls:
|
||||
- https://github.com/Blaizzy/mlx-audio
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-metal-darwin-arm64-mlx-audio
|
||||
license: MIT
|
||||
description: |
|
||||
Run Audio Models with MLX
|
||||
tags:
|
||||
- audio-to-text
|
||||
- audio-generation
|
||||
- text-to-audio
|
||||
- LLM
|
||||
- MLX
|
||||
- &rerankers
|
||||
name: "rerankers"
|
||||
alias: "rerankers"
|
||||
@@ -252,6 +270,7 @@
|
||||
nvidia: "cuda12-kokoro"
|
||||
intel: "intel-kokoro"
|
||||
amd: "rocm-kokoro"
|
||||
nvidia-l4t: "nvidia-l4t-kokoro"
|
||||
- &coqui
|
||||
urls:
|
||||
- https://github.com/idiap/coqui-ai-TTS
|
||||
@@ -332,6 +351,9 @@
|
||||
alias: "chatterbox"
|
||||
capabilities:
|
||||
nvidia: "cuda12-chatterbox"
|
||||
metal: "metal-chatterbox"
|
||||
default: "cpu-chatterbox"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-chatterbox"
|
||||
- &piper
|
||||
name: "piper"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-piper"
|
||||
@@ -405,6 +427,68 @@
|
||||
- text-to-speech
|
||||
- TTS
|
||||
license: apache-2.0
|
||||
- &neutts
|
||||
name: "neutts"
|
||||
urls:
|
||||
- https://github.com/neuphonic/neutts-air
|
||||
description: |
|
||||
NeuTTS Air is the world’s first super-realistic, on-device, TTS speech language model with instant voice cloning. Built off a 0.5B LLM backbone, NeuTTS Air brings natural-sounding speech, real-time performance, built-in security and speaker cloning to your local device - unlocking a new category of embedded voice agents, assistants, toys, and compliance-safe apps.
|
||||
tags:
|
||||
- text-to-speech
|
||||
- TTS
|
||||
license: apache-2.0
|
||||
capabilities:
|
||||
default: "cpu-neutts"
|
||||
nvidia: "cuda12-neutts"
|
||||
amd: "rocm-neutts"
|
||||
nvidia-l4t: "nvidia-l4t-neutts"
|
||||
- !!merge <<: *neutts
|
||||
name: "neutts-development"
|
||||
capabilities:
|
||||
default: "cpu-neutts-development"
|
||||
nvidia: "cuda12-neutts-development"
|
||||
amd: "rocm-neutts-development"
|
||||
nvidia-l4t: "nvidia-l4t-neutts-development"
|
||||
- !!merge <<: *neutts
|
||||
name: "cpu-neutts"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-neutts"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-neutts
|
||||
- !!merge <<: *neutts
|
||||
name: "cuda12-neutts"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-neutts"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-cuda-12-neutts
|
||||
- !!merge <<: *neutts
|
||||
name: "rocm-neutts"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-rocm-hipblas-neutts"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-rocm-hipblas-neutts
|
||||
- !!merge <<: *neutts
|
||||
name: "nvidia-l4t-neutts"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-nvidia-l4t-arm64-neutts"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-nvidia-l4t-arm64-neutts
|
||||
- !!merge <<: *neutts
|
||||
name: "cpu-neutts-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-neutts"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-neutts
|
||||
- !!merge <<: *neutts
|
||||
name: "cuda12-neutts-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-neutts"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-cuda-12-neutts
|
||||
- !!merge <<: *neutts
|
||||
name: "rocm-neutts-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-rocm-hipblas-neutts"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-rocm-hipblas-neutts
|
||||
- !!merge <<: *neutts
|
||||
name: "nvidia-l4t-neutts-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-arm64-neutts"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-nvidia-l4t-arm64-neutts
|
||||
- !!merge <<: *mlx
|
||||
name: "mlx-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-mlx"
|
||||
@@ -415,6 +499,11 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-mlx-vlm"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-mlx-vlm
|
||||
- !!merge <<: *mlx-audio
|
||||
name: "mlx-audio-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-mlx-audio"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-mlx-audio
|
||||
- !!merge <<: *kitten-tts
|
||||
name: "kitten-tts-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-kitten-tts"
|
||||
@@ -557,6 +646,16 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-whisper"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-whisper
|
||||
- !!merge <<: *whispercpp
|
||||
name: "metal-whisper"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-whisper"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-metal-darwin-arm64-whisper
|
||||
- !!merge <<: *whispercpp
|
||||
name: "metal-whisper-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-whisper"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-whisper
|
||||
- !!merge <<: *whispercpp
|
||||
name: "cpu-whisper-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-whisper"
|
||||
@@ -643,6 +742,16 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-stablediffusion-ggml"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-stablediffusion-ggml
|
||||
- !!merge <<: *stablediffusionggml
|
||||
name: "metal-stablediffusion-ggml"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-stablediffusion-ggml"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-metal-darwin-arm64-stablediffusion-ggml
|
||||
- !!merge <<: *stablediffusionggml
|
||||
name: "metal-stablediffusion-ggml-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-stablediffusion-ggml"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-stablediffusion-ggml
|
||||
- !!merge <<: *stablediffusionggml
|
||||
name: "vulkan-stablediffusion-ggml"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-vulkan-stablediffusion-ggml"
|
||||
@@ -1004,6 +1113,7 @@
|
||||
nvidia: "cuda12-kokoro-development"
|
||||
intel: "intel-kokoro-development"
|
||||
amd: "rocm-kokoro-development"
|
||||
nvidia-l4t: "nvidia-l4t-kokoro-development"
|
||||
- !!merge <<: *kokoro
|
||||
name: "cuda11-kokoro-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-11-kokoro"
|
||||
@@ -1029,6 +1139,16 @@
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-intel-kokoro"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-intel-kokoro
|
||||
- !!merge <<: *kokoro
|
||||
name: "nvidia-l4t-kokoro"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-l4t-kokoro"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-l4t-kokoro
|
||||
- !!merge <<: *kokoro
|
||||
name: "nvidia-l4t-kokoro-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-l4t-kokoro"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-l4t-kokoro
|
||||
- !!merge <<: *kokoro
|
||||
name: "cuda11-kokoro"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-11-kokoro"
|
||||
@@ -1180,6 +1300,39 @@
|
||||
name: "chatterbox-development"
|
||||
capabilities:
|
||||
nvidia: "cuda12-chatterbox-development"
|
||||
metal: "metal-chatterbox-development"
|
||||
default: "cpu-chatterbox-development"
|
||||
nvidia-l4t: "nvidia-l4t-arm64-chatterbox"
|
||||
- !!merge <<: *chatterbox
|
||||
name: "cpu-chatterbox"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-chatterbox"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-cpu-chatterbox
|
||||
- !!merge <<: *chatterbox
|
||||
name: "cpu-chatterbox-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-cpu-chatterbox"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-cpu-chatterbox
|
||||
- !!merge <<: *chatterbox
|
||||
name: "nvidia-l4t-arm64-chatterbox"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-l4t-arm64-chatterbox"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-gpu-nvidia-l4t-arm64-chatterbox
|
||||
- !!merge <<: *chatterbox
|
||||
name: "nvidia-l4t-arm64-chatterbox-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-l4t-arm64-chatterbox"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-gpu-nvidia-l4t-arm64-chatterbox
|
||||
- !!merge <<: *chatterbox
|
||||
name: "metal-chatterbox"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-chatterbox"
|
||||
mirrors:
|
||||
- localai/localai-backends:latest-metal-darwin-arm64-chatterbox
|
||||
- !!merge <<: *chatterbox
|
||||
name: "metal-chatterbox-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-chatterbox"
|
||||
mirrors:
|
||||
- localai/localai-backends:master-metal-darwin-arm64-chatterbox
|
||||
- !!merge <<: *chatterbox
|
||||
name: "cuda12-chatterbox-development"
|
||||
uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-chatterbox"
|
||||
|
||||
@@ -1,38 +1,190 @@
|
||||
# Common commands about conda environment
|
||||
# Python Backends for LocalAI
|
||||
|
||||
## Create a new empty conda environment
|
||||
This directory contains Python-based AI backends for LocalAI, providing support for various AI models and hardware acceleration targets.
|
||||
|
||||
```
|
||||
conda create --name <env-name> python=<your version> -y
|
||||
## Overview
|
||||
|
||||
conda create --name autogptq python=3.11 -y
|
||||
The Python backends use a unified build system based on `libbackend.sh` that provides:
|
||||
- **Automatic virtual environment management** with support for both `uv` and `pip`
|
||||
- **Hardware-specific dependency installation** (CPU, CUDA, Intel, MLX, etc.)
|
||||
- **Portable Python support** for standalone deployments
|
||||
- **Consistent backend execution** across different environments
|
||||
|
||||
## Available Backends
|
||||
|
||||
### Core AI Models
|
||||
- **transformers** - Hugging Face Transformers framework (PyTorch-based)
|
||||
- **vllm** - High-performance LLM inference engine
|
||||
- **mlx** - Apple Silicon optimized ML framework
|
||||
- **exllama2** - ExLlama2 quantized models
|
||||
|
||||
### Audio & Speech
|
||||
- **bark** - Text-to-speech synthesis
|
||||
- **coqui** - Coqui TTS models
|
||||
- **faster-whisper** - Fast Whisper speech recognition
|
||||
- **kitten-tts** - Lightweight TTS
|
||||
- **mlx-audio** - Apple Silicon audio processing
|
||||
- **chatterbox** - TTS model
|
||||
- **kokoro** - TTS models
|
||||
|
||||
### Computer Vision
|
||||
- **diffusers** - Stable Diffusion and image generation
|
||||
- **mlx-vlm** - Vision-language models for Apple Silicon
|
||||
- **rfdetr** - Object detection models
|
||||
|
||||
### Specialized
|
||||
|
||||
- **rerankers** - Text reranking models
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Prerequisites
|
||||
- Python 3.10+ (default: 3.10.18)
|
||||
- `uv` package manager (recommended) or `pip`
|
||||
- Appropriate hardware drivers for your target (CUDA, Intel, etc.)
|
||||
|
||||
### Installation
|
||||
|
||||
Each backend can be installed individually:
|
||||
|
||||
```bash
|
||||
# Navigate to a specific backend
|
||||
cd backend/python/transformers
|
||||
|
||||
# Install dependencies
|
||||
make transformers
|
||||
# or
|
||||
bash install.sh
|
||||
|
||||
# Run the backend
|
||||
make run
|
||||
# or
|
||||
bash run.sh
|
||||
```
|
||||
|
||||
## To activate the environment
|
||||
### Using the Unified Build System
|
||||
|
||||
As of conda 4.4
|
||||
```
|
||||
conda activate autogptq
|
||||
The `libbackend.sh` script provides consistent commands across all backends:
|
||||
|
||||
```bash
|
||||
# Source the library in your backend script
|
||||
source $(dirname $0)/../common/libbackend.sh
|
||||
|
||||
# Install requirements (automatically handles hardware detection)
|
||||
installRequirements
|
||||
|
||||
# Start the backend server
|
||||
startBackend $@
|
||||
|
||||
# Run tests
|
||||
runUnittests
|
||||
```
|
||||
|
||||
The conda version older than 4.4
|
||||
## Hardware Targets
|
||||
|
||||
```
|
||||
source activate autogptq
|
||||
The build system automatically detects and configures for different hardware:
|
||||
|
||||
- **CPU** - Standard CPU-only builds
|
||||
- **CUDA** - NVIDIA GPU acceleration (supports CUDA 11/12)
|
||||
- **Intel** - Intel XPU/GPU optimization
|
||||
- **MLX** - Apple Silicon (M1/M2/M3) optimization
|
||||
- **HIP** - AMD GPU acceleration
|
||||
|
||||
### Target-Specific Requirements
|
||||
|
||||
Backends can specify hardware-specific dependencies:
|
||||
- `requirements.txt` - Base requirements
|
||||
- `requirements-cpu.txt` - CPU-specific packages
|
||||
- `requirements-cublas11.txt` - CUDA 11 packages
|
||||
- `requirements-cublas12.txt` - CUDA 12 packages
|
||||
- `requirements-intel.txt` - Intel-optimized packages
|
||||
- `requirements-mps.txt` - Apple Silicon packages
|
||||
|
||||
## Configuration Options
|
||||
|
||||
### Environment Variables
|
||||
|
||||
- `PYTHON_VERSION` - Python version (default: 3.10)
|
||||
- `PYTHON_PATCH` - Python patch version (default: 18)
|
||||
- `BUILD_TYPE` - Force specific build target
|
||||
- `USE_PIP` - Use pip instead of uv (default: false)
|
||||
- `PORTABLE_PYTHON` - Enable portable Python builds
|
||||
- `LIMIT_TARGETS` - Restrict backend to specific targets
|
||||
|
||||
### Example: CUDA 12 Only Backend
|
||||
|
||||
```bash
|
||||
# In your backend script
|
||||
LIMIT_TARGETS="cublas12"
|
||||
source $(dirname $0)/../common/libbackend.sh
|
||||
```
|
||||
|
||||
## Install the packages to your environment
|
||||
### Example: Intel-Optimized Backend
|
||||
|
||||
Sometimes you need to install the packages from the conda-forge channel
|
||||
|
||||
By using `conda`
|
||||
```
|
||||
conda install <your-package-name>
|
||||
|
||||
conda install -c conda-forge <your package-name>
|
||||
```bash
|
||||
# In your backend script
|
||||
LIMIT_TARGETS="intel"
|
||||
source $(dirname $0)/../common/libbackend.sh
|
||||
```
|
||||
|
||||
Or by using `pip`
|
||||
## Development
|
||||
|
||||
### Adding a New Backend
|
||||
|
||||
1. Create a new directory in `backend/python/`
|
||||
2. Copy the template structure from `common/template/`
|
||||
3. Implement your `backend.py` with the required gRPC interface
|
||||
4. Add appropriate requirements files for your target hardware
|
||||
5. Use `libbackend.sh` for consistent build and execution
|
||||
|
||||
### Testing
|
||||
|
||||
```bash
|
||||
# Run backend tests
|
||||
make test
|
||||
# or
|
||||
bash test.sh
|
||||
```
|
||||
pip install <your-package-name>
|
||||
|
||||
### Building
|
||||
|
||||
```bash
|
||||
# Install dependencies
|
||||
make <backend-name>
|
||||
|
||||
# Clean build artifacts
|
||||
make clean
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
Each backend follows a consistent structure:
|
||||
```
|
||||
backend-name/
|
||||
├── backend.py # Main backend implementation
|
||||
├── requirements.txt # Base dependencies
|
||||
├── requirements-*.txt # Hardware-specific dependencies
|
||||
├── install.sh # Installation script
|
||||
├── run.sh # Execution script
|
||||
├── test.sh # Test script
|
||||
├── Makefile # Build targets
|
||||
└── test.py # Unit tests
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
1. **Missing dependencies**: Ensure all requirements files are properly configured
|
||||
2. **Hardware detection**: Check that `BUILD_TYPE` matches your system
|
||||
3. **Python version**: Verify Python 3.10+ is available
|
||||
4. **Virtual environment**: Use `ensureVenv` to create/activate environments
|
||||
|
||||
## Contributing
|
||||
|
||||
When adding new backends or modifying existing ones:
|
||||
1. Follow the established directory structure
|
||||
2. Use `libbackend.sh` for consistent behavior
|
||||
3. Include appropriate requirements files for all target hardware
|
||||
4. Add comprehensive tests
|
||||
5. Update this README if adding new backend types
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
bark==0.1.5
|
||||
grpcio==1.74.0
|
||||
grpcio==1.76.0
|
||||
protobuf
|
||||
certifi
|
||||
@@ -1,6 +1,6 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
This is an extra gRPC server of LocalAI for Bark TTS
|
||||
This is an extra gRPC server of LocalAI for Chatterbox TTS
|
||||
"""
|
||||
from concurrent import futures
|
||||
import time
|
||||
@@ -14,15 +14,98 @@ import backend_pb2_grpc
|
||||
import torch
|
||||
import torchaudio as ta
|
||||
from chatterbox.tts import ChatterboxTTS
|
||||
|
||||
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
|
||||
import grpc
|
||||
import tempfile
|
||||
|
||||
def is_float(s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
def is_int(s):
|
||||
"""Check if a string can be converted to int."""
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
def split_text_at_word_boundary(text, max_length=250):
|
||||
"""
|
||||
Split text at word boundaries without truncating words.
|
||||
Returns a list of text chunks.
|
||||
"""
|
||||
if not text or len(text) <= max_length:
|
||||
return [text]
|
||||
|
||||
chunks = []
|
||||
words = text.split()
|
||||
current_chunk = ""
|
||||
|
||||
for word in words:
|
||||
# Check if adding this word would exceed the limit
|
||||
if len(current_chunk) + len(word) + 1 <= max_length:
|
||||
if current_chunk:
|
||||
current_chunk += " " + word
|
||||
else:
|
||||
current_chunk = word
|
||||
else:
|
||||
# If current chunk is not empty, add it to chunks
|
||||
if current_chunk:
|
||||
chunks.append(current_chunk)
|
||||
current_chunk = word
|
||||
else:
|
||||
# If a single word is longer than max_length, we have to include it anyway
|
||||
chunks.append(word)
|
||||
current_chunk = ""
|
||||
|
||||
# Add the last chunk if it's not empty
|
||||
if current_chunk:
|
||||
chunks.append(current_chunk)
|
||||
|
||||
return chunks
|
||||
|
||||
def merge_audio_files(audio_files, output_path, sample_rate):
|
||||
"""
|
||||
Merge multiple audio files into a single audio file.
|
||||
"""
|
||||
if not audio_files:
|
||||
return
|
||||
|
||||
if len(audio_files) == 1:
|
||||
# If only one file, just copy it
|
||||
import shutil
|
||||
shutil.copy2(audio_files[0], output_path)
|
||||
return
|
||||
|
||||
# Load all audio files
|
||||
waveforms = []
|
||||
for audio_file in audio_files:
|
||||
waveform, sr = ta.load(audio_file)
|
||||
if sr != sample_rate:
|
||||
# Resample if necessary
|
||||
resampler = ta.transforms.Resample(sr, sample_rate)
|
||||
waveform = resampler(waveform)
|
||||
waveforms.append(waveform)
|
||||
|
||||
# Concatenate all waveforms
|
||||
merged_waveform = torch.cat(waveforms, dim=1)
|
||||
|
||||
# Save the merged audio
|
||||
ta.save(output_path, merged_waveform, sample_rate)
|
||||
|
||||
# Clean up temporary files
|
||||
for audio_file in audio_files:
|
||||
if os.path.exists(audio_file):
|
||||
os.remove(audio_file)
|
||||
|
||||
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
|
||||
# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
|
||||
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
||||
COQUI_LANGUAGE = os.environ.get('COQUI_LANGUAGE', None)
|
||||
|
||||
# Implement the BackendServicer class with the service methods
|
||||
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
@@ -47,6 +130,28 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
if not torch.cuda.is_available() and request.CUDA:
|
||||
return backend_pb2.Result(success=False, message="CUDA is not available")
|
||||
|
||||
|
||||
options = request.Options
|
||||
|
||||
# empty dict
|
||||
self.options = {}
|
||||
|
||||
# The options are a list of strings in this form optname:optvalue
|
||||
# We are storing all the options in a dict so we can use it later when
|
||||
# generating the images
|
||||
for opt in options:
|
||||
if ":" not in opt:
|
||||
continue
|
||||
key, value = opt.split(":")
|
||||
# if value is a number, convert it to the appropriate type
|
||||
if is_float(value):
|
||||
value = float(value)
|
||||
elif is_int(value):
|
||||
value = int(value)
|
||||
elif value.lower() in ["true", "false"]:
|
||||
value = value.lower() == "true"
|
||||
self.options[key] = value
|
||||
|
||||
self.AudioPath = None
|
||||
|
||||
if os.path.isabs(request.AudioPath):
|
||||
@@ -56,10 +161,14 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
modelFileBase = os.path.dirname(request.ModelFile)
|
||||
# modify LoraAdapter to be relative to modelFileBase
|
||||
self.AudioPath = os.path.join(modelFileBase, request.AudioPath)
|
||||
|
||||
try:
|
||||
print("Preparing models, please wait", file=sys.stderr)
|
||||
self.model = ChatterboxTTS.from_pretrained(device=device)
|
||||
if "multilingual" in self.options:
|
||||
# remove key from options
|
||||
del self.options["multilingual"]
|
||||
self.model = ChatterboxMultilingualTTS.from_pretrained(device=device)
|
||||
else:
|
||||
self.model = ChatterboxTTS.from_pretrained(device=device)
|
||||
except Exception as err:
|
||||
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
|
||||
# Implement your logic here for the LoadModel service
|
||||
@@ -68,14 +177,43 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
|
||||
def TTS(self, request, context):
|
||||
try:
|
||||
# Generate audio using ChatterboxTTS
|
||||
kwargs = {}
|
||||
|
||||
if "language" in self.options:
|
||||
kwargs["language_id"] = self.options["language"]
|
||||
if self.AudioPath is not None:
|
||||
wav = self.model.generate(request.text, audio_prompt_path=self.AudioPath)
|
||||
kwargs["audio_prompt_path"] = self.AudioPath
|
||||
|
||||
# add options to kwargs
|
||||
kwargs.update(self.options)
|
||||
|
||||
# Check if text exceeds 250 characters
|
||||
# (chatterbox does not support long text)
|
||||
# https://github.com/resemble-ai/chatterbox/issues/60
|
||||
# https://github.com/resemble-ai/chatterbox/issues/110
|
||||
if len(request.text) > 250:
|
||||
# Split text at word boundaries
|
||||
text_chunks = split_text_at_word_boundary(request.text, max_length=250)
|
||||
print(f"Splitting text into chunks of 250 characters: {len(text_chunks)}", file=sys.stderr)
|
||||
# Generate audio for each chunk
|
||||
temp_audio_files = []
|
||||
for i, chunk in enumerate(text_chunks):
|
||||
# Generate audio for this chunk
|
||||
wav = self.model.generate(chunk, **kwargs)
|
||||
|
||||
# Create temporary file for this chunk
|
||||
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
|
||||
temp_file.close()
|
||||
ta.save(temp_file.name, wav, self.model.sr)
|
||||
temp_audio_files.append(temp_file.name)
|
||||
|
||||
# Merge all audio files
|
||||
merge_audio_files(temp_audio_files, request.dst, self.model.sr)
|
||||
else:
|
||||
wav = self.model.generate(request.text)
|
||||
|
||||
# Save the generated audio
|
||||
ta.save(request.dst, wav, self.model.sr)
|
||||
# Generate audio using ChatterboxTTS for short text
|
||||
wav = self.model.generate(request.text, **kwargs)
|
||||
# Save the generated audio
|
||||
ta.save(request.dst, wav, self.model.sr)
|
||||
|
||||
except Exception as err:
|
||||
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
|
||||
|
||||
@@ -15,5 +15,6 @@ fi
|
||||
if [ "x${BUILD_PROFILE}" == "xintel" ]; then
|
||||
EXTRA_PIP_INSTALL_FLAGS+=" --upgrade --index-strategy=unsafe-first-match"
|
||||
fi
|
||||
EXTRA_PIP_INSTALL_FLAGS+=" --no-build-isolation"
|
||||
|
||||
installRequirements
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu
|
||||
accelerate
|
||||
torch==2.6.0
|
||||
torchaudio==2.6.0
|
||||
transformers==4.46.3
|
||||
chatterbox-tts
|
||||
torch
|
||||
torchaudio
|
||||
transformers
|
||||
# https://github.com/mudler/LocalAI/pull/6240#issuecomment-3329518289
|
||||
chatterbox-tts@git+https://git@github.com/mudler/chatterbox.git@faster
|
||||
#chatterbox-tts==0.1.4
|
||||
@@ -2,5 +2,6 @@
|
||||
torch==2.6.0+cu118
|
||||
torchaudio==2.6.0+cu118
|
||||
transformers==4.46.3
|
||||
chatterbox-tts
|
||||
# https://github.com/mudler/LocalAI/pull/6240#issuecomment-3329518289
|
||||
chatterbox-tts@git+https://git@github.com/mudler/chatterbox.git@faster
|
||||
accelerate
|
||||
@@ -1,5 +1,6 @@
|
||||
torch==2.6.0
|
||||
torchaudio==2.6.0
|
||||
transformers==4.46.3
|
||||
chatterbox-tts
|
||||
torch
|
||||
torchaudio
|
||||
transformers
|
||||
# https://github.com/mudler/LocalAI/pull/6240#issuecomment-3329518289
|
||||
chatterbox-tts@git+https://git@github.com/mudler/chatterbox.git@faster
|
||||
accelerate
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/rocm6.0
|
||||
torch==2.6.0+rocm6.1
|
||||
torchaudio==2.6.0+rocm6.1
|
||||
transformers==4.46.3
|
||||
chatterbox-tts
|
||||
transformers
|
||||
# https://github.com/mudler/LocalAI/pull/6240#issuecomment-3329518289
|
||||
chatterbox-tts@git+https://git@github.com/mudler/chatterbox.git@faster
|
||||
accelerate
|
||||
|
||||
@@ -2,10 +2,10 @@
|
||||
intel-extension-for-pytorch==2.3.110+xpu
|
||||
torch==2.3.1+cxx11.abi
|
||||
torchaudio==2.3.1+cxx11.abi
|
||||
transformers==4.46.3
|
||||
chatterbox-tts
|
||||
transformers
|
||||
# https://github.com/mudler/LocalAI/pull/6240#issuecomment-3329518289
|
||||
chatterbox-tts@git+https://git@github.com/mudler/chatterbox.git@faster
|
||||
accelerate
|
||||
oneccl_bind_pt==2.3.100+xpu
|
||||
optimum[openvino]
|
||||
setuptools
|
||||
accelerate
|
||||
setuptools
|
||||
6
backend/python/chatterbox/requirements-l4t.txt
Normal file
6
backend/python/chatterbox/requirements-l4t.txt
Normal file
@@ -0,0 +1,6 @@
|
||||
--extra-index-url https://pypi.jetson-ai-lab.io/jp6/cu126/
|
||||
torch
|
||||
torchaudio
|
||||
transformers
|
||||
chatterbox-tts@git+https://git@github.com/mudler/chatterbox.git@faster
|
||||
accelerate
|
||||
@@ -2,4 +2,5 @@ grpcio==1.71.0
|
||||
protobuf
|
||||
certifi
|
||||
packaging
|
||||
setuptools
|
||||
setuptools
|
||||
poetry
|
||||
@@ -286,7 +286,8 @@ _makeVenvPortable() {
|
||||
function ensureVenv() {
|
||||
local interpreter=""
|
||||
|
||||
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
|
||||
if [ "x${PORTABLE_PYTHON}" == "xtrue" ] || [ -e "$(_portable_python)" ]; then
|
||||
echo "Using portable Python"
|
||||
ensurePortablePython
|
||||
interpreter="$(_portable_python)"
|
||||
else
|
||||
@@ -384,6 +385,11 @@ function installRequirements() {
|
||||
requirementFiles+=("${EDIR}/requirements-${BUILD_PROFILE}-after.txt")
|
||||
fi
|
||||
|
||||
# This is needed to build wheels that e.g. depends on Python.h
|
||||
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
|
||||
export C_INCLUDE_PATH="${C_INCLUDE_PATH:-}:$(_portable_dir)/include/python${PYTHON_VERSION}"
|
||||
fi
|
||||
|
||||
for reqFile in ${requirementFiles[@]}; do
|
||||
if [ -f "${reqFile}" ]; then
|
||||
echo "starting requirements install for ${reqFile}"
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
grpcio==1.74.0
|
||||
grpcio==1.76.0
|
||||
protobuf
|
||||
grpcio-tools
|
||||
@@ -1,4 +1,4 @@
|
||||
grpcio==1.74.0
|
||||
grpcio==1.76.0
|
||||
protobuf
|
||||
certifi
|
||||
packaging==24.1
|
||||
@@ -18,7 +18,7 @@ import backend_pb2_grpc
|
||||
import grpc
|
||||
|
||||
from diffusers import SanaPipeline, StableDiffusion3Pipeline, StableDiffusionXLPipeline, StableDiffusionDepth2ImgPipeline, DPMSolverMultistepScheduler, StableDiffusionPipeline, DiffusionPipeline, \
|
||||
EulerAncestralDiscreteScheduler, FluxPipeline, FluxTransformer2DModel, QwenImageEditPipeline
|
||||
EulerAncestralDiscreteScheduler, FluxPipeline, FluxTransformer2DModel, QwenImageEditPipeline, AutoencoderKLWan, WanPipeline, WanImageToVideoPipeline
|
||||
from diffusers import StableDiffusionImg2ImgPipeline, AutoPipelineForText2Image, ControlNetModel, StableVideoDiffusionPipeline, Lumina2Text2ImgPipeline
|
||||
from diffusers.pipelines.stable_diffusion import safety_checker
|
||||
from diffusers.utils import load_image, export_to_video
|
||||
@@ -66,19 +66,21 @@ from diffusers.schedulers import (
|
||||
)
|
||||
|
||||
def is_float(s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
def is_int(s):
|
||||
"""Check if a string can be converted to int."""
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
# The scheduler list mapping was taken from here: https://github.com/neggles/animatediff-cli/blob/6f336f5f4b5e38e85d7f06f1744ef42d0a45f2a7/src/animatediff/schedulers.py#L39
|
||||
# Credits to https://github.com/neggles
|
||||
# See https://github.com/huggingface/diffusers/issues/4167 for more details on sched mapping from A1111
|
||||
@@ -187,6 +189,8 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
value = float(value)
|
||||
elif is_int(value):
|
||||
value = int(value)
|
||||
elif value.lower() in ["true", "false"]:
|
||||
value = value.lower() == "true"
|
||||
self.options[key] = value
|
||||
|
||||
# From options, extract if present "torch_dtype" and set it to the appropriate type
|
||||
@@ -334,6 +338,32 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
torch_dtype=torch.bfloat16)
|
||||
self.pipe.vae.to(torch.bfloat16)
|
||||
self.pipe.text_encoder.to(torch.bfloat16)
|
||||
elif request.PipelineType == "WanPipeline":
|
||||
# WAN2.2 pipeline requires special VAE handling
|
||||
vae = AutoencoderKLWan.from_pretrained(
|
||||
request.Model,
|
||||
subfolder="vae",
|
||||
torch_dtype=torch.float32
|
||||
)
|
||||
self.pipe = WanPipeline.from_pretrained(
|
||||
request.Model,
|
||||
vae=vae,
|
||||
torch_dtype=torchType
|
||||
)
|
||||
self.txt2vid = True # WAN2.2 is a text-to-video pipeline
|
||||
elif request.PipelineType == "WanImageToVideoPipeline":
|
||||
# WAN2.2 image-to-video pipeline
|
||||
vae = AutoencoderKLWan.from_pretrained(
|
||||
request.Model,
|
||||
subfolder="vae",
|
||||
torch_dtype=torch.float32
|
||||
)
|
||||
self.pipe = WanImageToVideoPipeline.from_pretrained(
|
||||
request.Model,
|
||||
vae=vae,
|
||||
torch_dtype=torchType
|
||||
)
|
||||
self.img2vid = True # WAN2.2 image-to-video pipeline
|
||||
|
||||
if CLIPSKIP and request.CLIPSkip != 0:
|
||||
self.clip_skip = request.CLIPSkip
|
||||
@@ -475,11 +505,24 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
"num_inference_steps": steps,
|
||||
}
|
||||
|
||||
if request.src != "" and not self.controlnet and not self.img2vid:
|
||||
image = Image.open(request.src)
|
||||
# Handle image source: prioritize RefImages over request.src
|
||||
image_src = None
|
||||
if hasattr(request, 'ref_images') and request.ref_images and len(request.ref_images) > 0:
|
||||
# Use the first reference image if available
|
||||
image_src = request.ref_images[0]
|
||||
print(f"Using reference image: {image_src}", file=sys.stderr)
|
||||
elif request.src != "":
|
||||
# Fall back to request.src if no ref_images
|
||||
image_src = request.src
|
||||
print(f"Using source image: {image_src}", file=sys.stderr)
|
||||
else:
|
||||
print("No image source provided", file=sys.stderr)
|
||||
|
||||
if image_src and not self.controlnet and not self.img2vid:
|
||||
image = Image.open(image_src)
|
||||
options["image"] = image
|
||||
elif self.controlnet and request.src:
|
||||
pose_image = load_image(request.src)
|
||||
elif self.controlnet and image_src:
|
||||
pose_image = load_image(image_src)
|
||||
options["image"] = pose_image
|
||||
|
||||
if CLIPSKIP and self.clip_skip != 0:
|
||||
@@ -521,7 +564,11 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
|
||||
if self.img2vid:
|
||||
# Load the conditioning image
|
||||
image = load_image(request.src)
|
||||
if image_src:
|
||||
image = load_image(image_src)
|
||||
else:
|
||||
# Fallback to request.src for img2vid if no ref_images
|
||||
image = load_image(request.src)
|
||||
image = image.resize((1024, 576))
|
||||
|
||||
generator = torch.manual_seed(request.seed)
|
||||
@@ -558,6 +605,96 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
|
||||
return backend_pb2.Result(message="Media generated", success=True)
|
||||
|
||||
def GenerateVideo(self, request, context):
|
||||
try:
|
||||
prompt = request.prompt
|
||||
if not prompt:
|
||||
return backend_pb2.Result(success=False, message="No prompt provided for video generation")
|
||||
|
||||
# Set default values from request or use defaults
|
||||
num_frames = request.num_frames if request.num_frames > 0 else 81
|
||||
fps = request.fps if request.fps > 0 else 16
|
||||
cfg_scale = request.cfg_scale if request.cfg_scale > 0 else 4.0
|
||||
num_inference_steps = request.step if request.step > 0 else 40
|
||||
|
||||
# Prepare generation parameters
|
||||
kwargs = {
|
||||
"prompt": prompt,
|
||||
"negative_prompt": request.negative_prompt if request.negative_prompt else "",
|
||||
"height": request.height if request.height > 0 else 720,
|
||||
"width": request.width if request.width > 0 else 1280,
|
||||
"num_frames": num_frames,
|
||||
"guidance_scale": cfg_scale,
|
||||
"num_inference_steps": num_inference_steps,
|
||||
}
|
||||
|
||||
# Add custom options from self.options (including guidance_scale_2 if specified)
|
||||
kwargs.update(self.options)
|
||||
|
||||
# Set seed if provided
|
||||
if request.seed > 0:
|
||||
kwargs["generator"] = torch.Generator(device=self.device).manual_seed(request.seed)
|
||||
|
||||
# Handle start and end images for video generation
|
||||
if request.start_image:
|
||||
kwargs["start_image"] = load_image(request.start_image)
|
||||
if request.end_image:
|
||||
kwargs["end_image"] = load_image(request.end_image)
|
||||
|
||||
print(f"Generating video with {kwargs=}", file=sys.stderr)
|
||||
|
||||
# Generate video frames based on pipeline type
|
||||
if self.PipelineType == "WanPipeline":
|
||||
# WAN2.2 text-to-video generation
|
||||
output = self.pipe(**kwargs)
|
||||
frames = output.frames[0] # WAN2.2 returns frames in this format
|
||||
elif self.PipelineType == "WanImageToVideoPipeline":
|
||||
# WAN2.2 image-to-video generation
|
||||
if request.start_image:
|
||||
# Load and resize the input image according to WAN2.2 requirements
|
||||
image = load_image(request.start_image)
|
||||
# Use request dimensions or defaults, but respect WAN2.2 constraints
|
||||
request_height = request.height if request.height > 0 else 480
|
||||
request_width = request.width if request.width > 0 else 832
|
||||
max_area = request_height * request_width
|
||||
aspect_ratio = image.height / image.width
|
||||
mod_value = self.pipe.vae_scale_factor_spatial * self.pipe.transformer.config.patch_size[1]
|
||||
height = round((max_area * aspect_ratio) ** 0.5 / mod_value) * mod_value
|
||||
width = round((max_area / aspect_ratio) ** 0.5 / mod_value) * mod_value
|
||||
image = image.resize((width, height))
|
||||
kwargs["image"] = image
|
||||
kwargs["height"] = height
|
||||
kwargs["width"] = width
|
||||
|
||||
output = self.pipe(**kwargs)
|
||||
frames = output.frames[0]
|
||||
elif self.img2vid:
|
||||
# Generic image-to-video generation
|
||||
if request.start_image:
|
||||
image = load_image(request.start_image)
|
||||
image = image.resize((request.width if request.width > 0 else 1024,
|
||||
request.height if request.height > 0 else 576))
|
||||
kwargs["image"] = image
|
||||
|
||||
output = self.pipe(**kwargs)
|
||||
frames = output.frames[0]
|
||||
elif self.txt2vid:
|
||||
# Generic text-to-video generation
|
||||
output = self.pipe(**kwargs)
|
||||
frames = output.frames[0]
|
||||
else:
|
||||
return backend_pb2.Result(success=False, message=f"Pipeline {self.PipelineType} does not support video generation")
|
||||
|
||||
# Export video
|
||||
export_to_video(frames, request.dst, fps=fps)
|
||||
|
||||
return backend_pb2.Result(message="Video generated successfully", success=True)
|
||||
|
||||
except Exception as err:
|
||||
print(f"Error generating video: {err}", file=sys.stderr)
|
||||
traceback.print_exc()
|
||||
return backend_pb2.Result(success=False, message=f"Error generating video: {err}")
|
||||
|
||||
|
||||
def serve(address):
|
||||
server = grpc.server(futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
|
||||
|
||||
@@ -8,4 +8,5 @@ compel
|
||||
peft
|
||||
sentencepiece
|
||||
torch==2.7.1
|
||||
optimum-quanto
|
||||
optimum-quanto
|
||||
ftfy
|
||||
@@ -1,11 +1,12 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/cu118
|
||||
torch==2.7.1+cu118
|
||||
torchvision==0.22.1+cu118
|
||||
git+https://github.com/huggingface/diffusers
|
||||
opencv-python
|
||||
transformers
|
||||
torchvision==0.22.1
|
||||
accelerate
|
||||
compel
|
||||
peft
|
||||
sentencepiece
|
||||
optimum-quanto
|
||||
torch==2.7.1
|
||||
optimum-quanto
|
||||
ftfy
|
||||
@@ -1,10 +1,12 @@
|
||||
torch==2.7.1
|
||||
torchvision==0.22.1
|
||||
--extra-index-url https://download.pytorch.org/whl/cu121
|
||||
git+https://github.com/huggingface/diffusers
|
||||
opencv-python
|
||||
transformers
|
||||
torchvision
|
||||
accelerate
|
||||
compel
|
||||
peft
|
||||
sentencepiece
|
||||
optimum-quanto
|
||||
torch
|
||||
ftfy
|
||||
optimum-quanto
|
||||
|
||||
@@ -8,4 +8,5 @@ accelerate
|
||||
compel
|
||||
peft
|
||||
sentencepiece
|
||||
optimum-quanto
|
||||
optimum-quanto
|
||||
ftfy
|
||||
@@ -12,4 +12,5 @@ accelerate
|
||||
compel
|
||||
peft
|
||||
sentencepiece
|
||||
optimum-quanto
|
||||
optimum-quanto
|
||||
ftfy
|
||||
@@ -8,4 +8,5 @@ peft
|
||||
optimum-quanto
|
||||
numpy<2
|
||||
sentencepiece
|
||||
torchvision
|
||||
torchvision
|
||||
ftfy
|
||||
@@ -7,4 +7,5 @@ accelerate
|
||||
compel
|
||||
peft
|
||||
sentencepiece
|
||||
optimum-quanto
|
||||
optimum-quanto
|
||||
ftfy
|
||||
@@ -1,5 +1,5 @@
|
||||
setuptools
|
||||
grpcio==1.74.0
|
||||
grpcio==1.76.0
|
||||
pillow
|
||||
protobuf
|
||||
certifi
|
||||
|
||||
@@ -31,7 +31,7 @@ class TestBackendServicer(unittest.TestCase):
|
||||
"""
|
||||
This method tests if the server starts up successfully
|
||||
"""
|
||||
time.sleep(10)
|
||||
time.sleep(20)
|
||||
try:
|
||||
self.setUp()
|
||||
with grpc.insecure_channel("localhost:50051") as channel:
|
||||
@@ -48,7 +48,7 @@ class TestBackendServicer(unittest.TestCase):
|
||||
"""
|
||||
This method tests if the model is loaded successfully
|
||||
"""
|
||||
time.sleep(10)
|
||||
time.sleep(20)
|
||||
try:
|
||||
self.setUp()
|
||||
with grpc.insecure_channel("localhost:50051") as channel:
|
||||
@@ -66,7 +66,7 @@ class TestBackendServicer(unittest.TestCase):
|
||||
"""
|
||||
This method tests if the backend can generate images
|
||||
"""
|
||||
time.sleep(10)
|
||||
time.sleep(20)
|
||||
try:
|
||||
self.setUp()
|
||||
with grpc.insecure_channel("localhost:50051") as channel:
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
grpcio==1.74.0
|
||||
grpcio==1.76.0
|
||||
protobuf
|
||||
certifi
|
||||
wheel
|
||||
|
||||
@@ -64,15 +64,15 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
# Generate audio using Kokoro pipeline
|
||||
generator = self.pipeline(request.text, voice=voice)
|
||||
|
||||
# Get the first (and typically only) audio segment
|
||||
speechs = []
|
||||
# Get all the audio segment
|
||||
for i, (gs, ps, audio) in enumerate(generator):
|
||||
# Save audio to the destination file
|
||||
sf.write(request.dst, audio, 24000)
|
||||
speechs.append(audio)
|
||||
print(f"Generated audio segment {i}: gs={gs}, ps={ps}", file=sys.stderr)
|
||||
# For now, we only process the first segment
|
||||
# If you need to handle multiple segments, you might want to modify this
|
||||
break
|
||||
|
||||
# Merges the audio segments and writes them to the destination
|
||||
speech = torch.cat(speechs, dim=0)
|
||||
sf.write(request.dst, speech, 24000)
|
||||
|
||||
except Exception as err:
|
||||
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
|
||||
|
||||
|
||||
7
backend/python/kokoro/requirements-l4t.txt
Normal file
7
backend/python/kokoro/requirements-l4t.txt
Normal file
@@ -0,0 +1,7 @@
|
||||
--extra-index-url https://pypi.jetson-ai-lab.io/jp6/cu126/
|
||||
torch
|
||||
torchaudio
|
||||
transformers
|
||||
accelerate
|
||||
kokoro
|
||||
soundfile
|
||||
23
backend/python/mlx-audio/Makefile
Normal file
23
backend/python/mlx-audio/Makefile
Normal file
@@ -0,0 +1,23 @@
|
||||
.PHONY: mlx-audio
|
||||
mlx-audio:
|
||||
bash install.sh
|
||||
|
||||
.PHONY: run
|
||||
run: mlx-audio
|
||||
@echo "Running mlx-audio..."
|
||||
bash run.sh
|
||||
@echo "mlx run."
|
||||
|
||||
.PHONY: test
|
||||
test: mlx-audio
|
||||
@echo "Testing mlx-audio..."
|
||||
bash test.sh
|
||||
@echo "mlx tested."
|
||||
|
||||
.PHONY: protogen-clean
|
||||
protogen-clean:
|
||||
$(RM) backend_pb2_grpc.py backend_pb2.py
|
||||
|
||||
.PHONY: clean
|
||||
clean: protogen-clean
|
||||
rm -rf venv __pycache__
|
||||
465
backend/python/mlx-audio/backend.py
Normal file
465
backend/python/mlx-audio/backend.py
Normal file
@@ -0,0 +1,465 @@
|
||||
#!/usr/bin/env python3
|
||||
import asyncio
|
||||
from concurrent import futures
|
||||
import argparse
|
||||
import signal
|
||||
import sys
|
||||
import os
|
||||
import shutil
|
||||
import glob
|
||||
from typing import List
|
||||
import time
|
||||
import tempfile
|
||||
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
|
||||
import grpc
|
||||
from mlx_audio.tts.utils import load_model
|
||||
import soundfile as sf
|
||||
import numpy as np
|
||||
import uuid
|
||||
|
||||
def is_float(s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
def is_int(s):
|
||||
"""Check if a string can be converted to int."""
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
|
||||
# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
|
||||
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
||||
|
||||
# Implement the BackendServicer class with the service methods
|
||||
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
"""
|
||||
A gRPC servicer that implements the Backend service defined in backend.proto.
|
||||
This backend provides TTS (Text-to-Speech) functionality using MLX-Audio.
|
||||
"""
|
||||
|
||||
def Health(self, request, context):
|
||||
"""
|
||||
Returns a health check message.
|
||||
|
||||
Args:
|
||||
request: The health check request.
|
||||
context: The gRPC context.
|
||||
|
||||
Returns:
|
||||
backend_pb2.Reply: The health check reply.
|
||||
"""
|
||||
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
|
||||
|
||||
async def LoadModel(self, request, context):
|
||||
"""
|
||||
Loads a TTS model using MLX-Audio.
|
||||
|
||||
Args:
|
||||
request: The load model request.
|
||||
context: The gRPC context.
|
||||
|
||||
Returns:
|
||||
backend_pb2.Result: The load model result.
|
||||
"""
|
||||
try:
|
||||
print(f"Loading MLX-Audio TTS model: {request.Model}", file=sys.stderr)
|
||||
print(f"Request: {request}", file=sys.stderr)
|
||||
|
||||
# Parse options like in the kokoro backend
|
||||
options = request.Options
|
||||
self.options = {}
|
||||
|
||||
# The options are a list of strings in this form optname:optvalue
|
||||
# We store all the options in a dict for later use
|
||||
for opt in options:
|
||||
if ":" not in opt:
|
||||
continue
|
||||
key, value = opt.split(":", 1) # Split only on first colon to handle values with colons
|
||||
|
||||
# Convert numeric values to appropriate types
|
||||
if is_float(value):
|
||||
value = float(value)
|
||||
elif is_int(value):
|
||||
value = int(value)
|
||||
elif value.lower() in ["true", "false"]:
|
||||
value = value.lower() == "true"
|
||||
|
||||
self.options[key] = value
|
||||
|
||||
print(f"Options: {self.options}", file=sys.stderr)
|
||||
|
||||
# Load the model using MLX-Audio's load_model function
|
||||
try:
|
||||
self.tts_model = load_model(request.Model)
|
||||
self.model_path = request.Model
|
||||
print(f"TTS model loaded successfully from {request.Model}", file=sys.stderr)
|
||||
except Exception as model_err:
|
||||
print(f"Error loading TTS model: {model_err}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message=f"Failed to load model: {model_err}")
|
||||
|
||||
except Exception as err:
|
||||
print(f"Error loading MLX-Audio TTS model {err=}, {type(err)=}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message=f"Error loading MLX-Audio TTS model: {err}")
|
||||
|
||||
print("MLX-Audio TTS model loaded successfully", file=sys.stderr)
|
||||
return backend_pb2.Result(message="MLX-Audio TTS model loaded successfully", success=True)
|
||||
|
||||
def TTS(self, request, context):
|
||||
"""
|
||||
Generates TTS audio from text using MLX-Audio.
|
||||
|
||||
Args:
|
||||
request: A TTSRequest object containing text, model, destination, voice, and language.
|
||||
context: A grpc.ServicerContext object that provides information about the RPC.
|
||||
|
||||
Returns:
|
||||
A Result object indicating success or failure.
|
||||
"""
|
||||
try:
|
||||
# Check if model is loaded
|
||||
if not hasattr(self, 'tts_model') or self.tts_model is None:
|
||||
return backend_pb2.Result(success=False, message="TTS model not loaded. Please call LoadModel first.")
|
||||
|
||||
print(f"Generating TTS with MLX-Audio - text: {request.text[:50]}..., voice: {request.voice}, language: {request.language}", file=sys.stderr)
|
||||
|
||||
# Handle speed parameter based on model type
|
||||
speed_value = self._handle_speed_parameter(request, self.model_path)
|
||||
|
||||
# Map language names to codes if needed
|
||||
lang_code = self._map_language_code(request.language, request.voice)
|
||||
|
||||
# Prepare generation parameters
|
||||
gen_params = {
|
||||
"text": request.text,
|
||||
"speed": speed_value,
|
||||
"verbose": False,
|
||||
}
|
||||
|
||||
# Add model-specific parameters
|
||||
if request.voice and request.voice.strip():
|
||||
gen_params["voice"] = request.voice
|
||||
|
||||
# Check if model supports language codes (primarily Kokoro)
|
||||
if "kokoro" in self.model_path.lower():
|
||||
gen_params["lang_code"] = lang_code
|
||||
|
||||
# Add pitch and gender for Spark models
|
||||
if "spark" in self.model_path.lower():
|
||||
gen_params["pitch"] = 1.0 # Default to moderate
|
||||
gen_params["gender"] = "female" # Default to female
|
||||
|
||||
print(f"Generation parameters: {gen_params}", file=sys.stderr)
|
||||
|
||||
# Generate audio using the loaded model
|
||||
try:
|
||||
results = self.tts_model.generate(**gen_params)
|
||||
except Exception as gen_err:
|
||||
print(f"Error during TTS generation: {gen_err}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message=f"TTS generation failed: {gen_err}")
|
||||
|
||||
# Process the generated audio segments
|
||||
audio_arrays = []
|
||||
for segment in results:
|
||||
audio_arrays.append(segment.audio)
|
||||
|
||||
# If no segments, return error
|
||||
if not audio_arrays:
|
||||
print("No audio segments generated", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message="No audio generated")
|
||||
|
||||
# Concatenate all segments
|
||||
cat_audio = np.concatenate(audio_arrays, axis=0)
|
||||
|
||||
# Generate output filename and path
|
||||
if request.dst:
|
||||
output_path = request.dst
|
||||
else:
|
||||
unique_id = str(uuid.uuid4())
|
||||
filename = f"tts_{unique_id}.wav"
|
||||
output_path = filename
|
||||
|
||||
# Write the audio as a WAV
|
||||
try:
|
||||
sf.write(output_path, cat_audio, 24000)
|
||||
print(f"Successfully wrote audio file to {output_path}", file=sys.stderr)
|
||||
|
||||
# Verify the file exists and has content
|
||||
if not os.path.exists(output_path):
|
||||
print(f"File was not created at {output_path}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message="Failed to create audio file")
|
||||
|
||||
file_size = os.path.getsize(output_path)
|
||||
if file_size == 0:
|
||||
print("File was created but is empty", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message="Generated audio file is empty")
|
||||
|
||||
print(f"Audio file size: {file_size} bytes", file=sys.stderr)
|
||||
|
||||
except Exception as write_err:
|
||||
print(f"Error writing audio file: {write_err}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message=f"Failed to save audio: {write_err}")
|
||||
|
||||
return backend_pb2.Result(success=True, message=f"TTS audio generated successfully: {output_path}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in MLX-Audio TTS: {e}", file=sys.stderr)
|
||||
return backend_pb2.Result(success=False, message=f"TTS generation failed: {str(e)}")
|
||||
|
||||
async def Predict(self, request, context):
|
||||
"""
|
||||
Generates TTS audio based on the given prompt using MLX-Audio TTS.
|
||||
This is a fallback method for compatibility with the Predict endpoint.
|
||||
|
||||
Args:
|
||||
request: The predict request.
|
||||
context: The gRPC context.
|
||||
|
||||
Returns:
|
||||
backend_pb2.Reply: The predict result.
|
||||
"""
|
||||
try:
|
||||
# Check if model is loaded
|
||||
if not hasattr(self, 'tts_model') or self.tts_model is None:
|
||||
context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
|
||||
context.set_details("TTS model not loaded. Please call LoadModel first.")
|
||||
return backend_pb2.Reply(message=bytes("", encoding='utf-8'))
|
||||
|
||||
# For TTS, we expect the prompt to contain the text to synthesize
|
||||
if not request.Prompt:
|
||||
context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
|
||||
context.set_details("Prompt is required for TTS generation")
|
||||
return backend_pb2.Reply(message=bytes("", encoding='utf-8'))
|
||||
|
||||
# Handle speed parameter based on model type
|
||||
speed_value = self._handle_speed_parameter(request, self.model_path)
|
||||
|
||||
# Map language names to codes if needed
|
||||
lang_code = self._map_language_code(None, None) # Use defaults for Predict
|
||||
|
||||
# Prepare generation parameters
|
||||
gen_params = {
|
||||
"text": request.Prompt,
|
||||
"speed": speed_value,
|
||||
"verbose": False,
|
||||
}
|
||||
|
||||
# Add model-specific parameters
|
||||
if hasattr(self, 'options') and 'voice' in self.options:
|
||||
gen_params["voice"] = self.options['voice']
|
||||
|
||||
# Check if model supports language codes (primarily Kokoro)
|
||||
if "kokoro" in self.model_path.lower():
|
||||
gen_params["lang_code"] = lang_code
|
||||
|
||||
print(f"Generating TTS with MLX-Audio - text: {request.Prompt[:50]}..., params: {gen_params}", file=sys.stderr)
|
||||
|
||||
# Generate audio using the loaded model
|
||||
try:
|
||||
results = self.tts_model.generate(**gen_params)
|
||||
except Exception as gen_err:
|
||||
print(f"Error during TTS generation: {gen_err}", file=sys.stderr)
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(f"TTS generation failed: {gen_err}")
|
||||
return backend_pb2.Reply(message=bytes("", encoding='utf-8'))
|
||||
|
||||
# Process the generated audio segments
|
||||
audio_arrays = []
|
||||
for segment in results:
|
||||
audio_arrays.append(segment.audio)
|
||||
|
||||
# If no segments, return error
|
||||
if not audio_arrays:
|
||||
print("No audio segments generated", file=sys.stderr)
|
||||
return backend_pb2.Reply(message=bytes("No audio generated", encoding='utf-8'))
|
||||
|
||||
# Concatenate all segments
|
||||
cat_audio = np.concatenate(audio_arrays, axis=0)
|
||||
duration = len(cat_audio) / 24000 # Assuming 24kHz sample rate
|
||||
|
||||
# Return success message with audio information
|
||||
response = f"TTS audio generated successfully. Duration: {duration:.2f}s, Sample rate: 24000Hz"
|
||||
return backend_pb2.Reply(message=bytes(response, encoding='utf-8'))
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in MLX-Audio TTS Predict: {e}", file=sys.stderr)
|
||||
context.set_code(grpc.StatusCode.INTERNAL)
|
||||
context.set_details(f"TTS generation failed: {str(e)}")
|
||||
return backend_pb2.Reply(message=bytes("", encoding='utf-8'))
|
||||
|
||||
def _handle_speed_parameter(self, request, model_path):
|
||||
"""
|
||||
Handle speed parameter based on model type.
|
||||
|
||||
Args:
|
||||
request: The TTSRequest object.
|
||||
model_path: The model path to determine model type.
|
||||
|
||||
Returns:
|
||||
float: The processed speed value.
|
||||
"""
|
||||
# Get speed from options if available
|
||||
speed = 1.0
|
||||
if hasattr(self, 'options') and 'speed' in self.options:
|
||||
speed = self.options['speed']
|
||||
|
||||
# Handle speed parameter based on model type
|
||||
if "spark" in model_path.lower():
|
||||
# Spark actually expects float values that map to speed descriptions
|
||||
speed_map = {
|
||||
"very_low": 0.0,
|
||||
"low": 0.5,
|
||||
"moderate": 1.0,
|
||||
"high": 1.5,
|
||||
"very_high": 2.0,
|
||||
}
|
||||
if isinstance(speed, str) and speed in speed_map:
|
||||
speed_value = speed_map[speed]
|
||||
else:
|
||||
# Try to use as float, default to 1.0 (moderate) if invalid
|
||||
try:
|
||||
speed_value = float(speed)
|
||||
if speed_value not in [0.0, 0.5, 1.0, 1.5, 2.0]:
|
||||
speed_value = 1.0 # Default to moderate
|
||||
except:
|
||||
speed_value = 1.0 # Default to moderate
|
||||
else:
|
||||
# Other models use float speed values
|
||||
try:
|
||||
speed_value = float(speed)
|
||||
if speed_value < 0.5 or speed_value > 2.0:
|
||||
speed_value = 1.0 # Default to 1.0 if out of range
|
||||
except ValueError:
|
||||
speed_value = 1.0 # Default to 1.0 if invalid
|
||||
|
||||
return speed_value
|
||||
|
||||
def _map_language_code(self, language, voice):
|
||||
"""
|
||||
Map language names to codes if needed.
|
||||
|
||||
Args:
|
||||
language: The language parameter from the request.
|
||||
voice: The voice parameter from the request.
|
||||
|
||||
Returns:
|
||||
str: The language code.
|
||||
"""
|
||||
if not language:
|
||||
# Default to voice[0] if not found
|
||||
return voice[0] if voice else "a"
|
||||
|
||||
# Map language names to codes if needed
|
||||
language_map = {
|
||||
"american_english": "a",
|
||||
"british_english": "b",
|
||||
"spanish": "e",
|
||||
"french": "f",
|
||||
"hindi": "h",
|
||||
"italian": "i",
|
||||
"portuguese": "p",
|
||||
"japanese": "j",
|
||||
"mandarin_chinese": "z",
|
||||
# Also accept direct language codes
|
||||
"a": "a", "b": "b", "e": "e", "f": "f", "h": "h", "i": "i", "p": "p", "j": "j", "z": "z",
|
||||
}
|
||||
|
||||
return language_map.get(language.lower(), language)
|
||||
|
||||
def _build_generation_params(self, request, default_speed=1.0):
|
||||
"""
|
||||
Build generation parameters from request attributes and options for MLX-Audio TTS.
|
||||
|
||||
Args:
|
||||
request: The gRPC request.
|
||||
default_speed: Default speed if not specified.
|
||||
|
||||
Returns:
|
||||
dict: Generation parameters for MLX-Audio
|
||||
"""
|
||||
# Initialize generation parameters for MLX-Audio TTS
|
||||
generation_params = {
|
||||
'speed': default_speed,
|
||||
'voice': 'af_heart', # Default voice
|
||||
'lang_code': 'a', # Default language code
|
||||
}
|
||||
|
||||
# Extract parameters from request attributes
|
||||
if hasattr(request, 'Temperature') and request.Temperature > 0:
|
||||
# Temperature could be mapped to speed variation
|
||||
generation_params['speed'] = 1.0 + (request.Temperature - 0.5) * 0.5
|
||||
|
||||
# Override with options if available
|
||||
if hasattr(self, 'options'):
|
||||
# Speed from options
|
||||
if 'speed' in self.options:
|
||||
generation_params['speed'] = self.options['speed']
|
||||
|
||||
# Voice from options
|
||||
if 'voice' in self.options:
|
||||
generation_params['voice'] = self.options['voice']
|
||||
|
||||
# Language code from options
|
||||
if 'lang_code' in self.options:
|
||||
generation_params['lang_code'] = self.options['lang_code']
|
||||
|
||||
# Model-specific parameters
|
||||
param_option_mapping = {
|
||||
'temp': 'speed',
|
||||
'temperature': 'speed',
|
||||
'top_p': 'speed', # Map top_p to speed variation
|
||||
}
|
||||
|
||||
for option_key, param_key in param_option_mapping.items():
|
||||
if option_key in self.options:
|
||||
if param_key == 'speed':
|
||||
# Ensure speed is within reasonable bounds
|
||||
speed_val = float(self.options[option_key])
|
||||
if 0.5 <= speed_val <= 2.0:
|
||||
generation_params[param_key] = speed_val
|
||||
|
||||
return generation_params
|
||||
|
||||
async def serve(address):
|
||||
# Start asyncio gRPC server
|
||||
server = grpc.aio.server(migration_thread_pool=futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
|
||||
options=[
|
||||
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
|
||||
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
|
||||
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
|
||||
])
|
||||
# Add the servicer to the server
|
||||
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
|
||||
# Bind the server to the address
|
||||
server.add_insecure_port(address)
|
||||
|
||||
# Gracefully shutdown the server on SIGTERM or SIGINT
|
||||
loop = asyncio.get_event_loop()
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
loop.add_signal_handler(
|
||||
sig, lambda: asyncio.ensure_future(server.stop(5))
|
||||
)
|
||||
|
||||
# Start the server
|
||||
await server.start()
|
||||
print("MLX-Audio TTS Server started. Listening on: " + address, file=sys.stderr)
|
||||
# Wait for the server to be terminated
|
||||
await server.wait_for_termination()
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Run the MLX-Audio TTS gRPC server.")
|
||||
parser.add_argument(
|
||||
"--addr", default="localhost:50051", help="The address to bind the server to."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
asyncio.run(serve(args.addr))
|
||||
14
backend/python/mlx-audio/install.sh
Executable file
14
backend/python/mlx-audio/install.sh
Executable file
@@ -0,0 +1,14 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
USE_PIP=true
|
||||
|
||||
backend_dir=$(dirname $0)
|
||||
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
installRequirements
|
||||
1
backend/python/mlx-audio/requirements-mps.txt
Normal file
1
backend/python/mlx-audio/requirements-mps.txt
Normal file
@@ -0,0 +1 @@
|
||||
git+https://github.com/Blaizzy/mlx-audio
|
||||
7
backend/python/mlx-audio/requirements.txt
Normal file
7
backend/python/mlx-audio/requirements.txt
Normal file
@@ -0,0 +1,7 @@
|
||||
grpcio==1.71.0
|
||||
protobuf
|
||||
certifi
|
||||
setuptools
|
||||
mlx-audio
|
||||
soundfile
|
||||
numpy
|
||||
11
backend/python/mlx-audio/run.sh
Executable file
11
backend/python/mlx-audio/run.sh
Executable file
@@ -0,0 +1,11 @@
|
||||
#!/bin/bash
|
||||
|
||||
backend_dir=$(dirname $0)
|
||||
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
startBackend $@
|
||||
142
backend/python/mlx-audio/test.py
Normal file
142
backend/python/mlx-audio/test.py
Normal file
@@ -0,0 +1,142 @@
|
||||
import unittest
|
||||
import subprocess
|
||||
import time
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
|
||||
import grpc
|
||||
|
||||
import unittest
|
||||
import subprocess
|
||||
import time
|
||||
import grpc
|
||||
import backend_pb2_grpc
|
||||
import backend_pb2
|
||||
|
||||
class TestBackendServicer(unittest.TestCase):
|
||||
"""
|
||||
TestBackendServicer is the class that tests the gRPC service.
|
||||
|
||||
This class contains methods to test the startup and shutdown of the gRPC service.
|
||||
"""
|
||||
def setUp(self):
|
||||
self.service = subprocess.Popen(["python", "backend.py", "--addr", "localhost:50051"])
|
||||
time.sleep(10)
|
||||
|
||||
def tearDown(self) -> None:
|
||||
self.service.terminate()
|
||||
self.service.wait()
|
||||
|
||||
def test_server_startup(self):
|
||||
try:
|
||||
self.setUp()
|
||||
with grpc.insecure_channel("localhost:50051") as channel:
|
||||
stub = backend_pb2_grpc.BackendStub(channel)
|
||||
response = stub.Health(backend_pb2.HealthMessage())
|
||||
self.assertEqual(response.message, b'OK')
|
||||
except Exception as err:
|
||||
print(err)
|
||||
self.fail("Server failed to start")
|
||||
finally:
|
||||
self.tearDown()
|
||||
def test_load_model(self):
|
||||
"""
|
||||
This method tests if the TTS model is loaded successfully
|
||||
"""
|
||||
try:
|
||||
self.setUp()
|
||||
with grpc.insecure_channel("localhost:50051") as channel:
|
||||
stub = backend_pb2_grpc.BackendStub(channel)
|
||||
response = stub.LoadModel(backend_pb2.ModelOptions(Model="mlx-community/Kokoro-82M-4bit"))
|
||||
self.assertTrue(response.success)
|
||||
self.assertEqual(response.message, "MLX-Audio TTS model loaded successfully")
|
||||
except Exception as err:
|
||||
print(err)
|
||||
self.fail("LoadModel service failed")
|
||||
finally:
|
||||
self.tearDown()
|
||||
|
||||
def test_tts_generation(self):
|
||||
"""
|
||||
This method tests if TTS audio is generated successfully
|
||||
"""
|
||||
try:
|
||||
self.setUp()
|
||||
with grpc.insecure_channel("localhost:50051") as channel:
|
||||
stub = backend_pb2_grpc.BackendStub(channel)
|
||||
response = stub.LoadModel(backend_pb2.ModelOptions(Model="mlx-community/Kokoro-82M-4bit"))
|
||||
self.assertTrue(response.success)
|
||||
|
||||
# Test TTS generation
|
||||
tts_req = backend_pb2.TTSRequest(
|
||||
text="Hello, this is a test of the MLX-Audio TTS system.",
|
||||
model="mlx-community/Kokoro-82M-4bit",
|
||||
voice="af_heart",
|
||||
language="a"
|
||||
)
|
||||
tts_resp = stub.TTS(tts_req)
|
||||
self.assertTrue(tts_resp.success)
|
||||
self.assertIn("TTS audio generated successfully", tts_resp.message)
|
||||
except Exception as err:
|
||||
print(err)
|
||||
self.fail("TTS service failed")
|
||||
finally:
|
||||
self.tearDown()
|
||||
|
||||
def test_tts_with_options(self):
|
||||
"""
|
||||
This method tests if TTS works with various options and parameters
|
||||
"""
|
||||
try:
|
||||
self.setUp()
|
||||
with grpc.insecure_channel("localhost:50051") as channel:
|
||||
stub = backend_pb2_grpc.BackendStub(channel)
|
||||
response = stub.LoadModel(backend_pb2.ModelOptions(
|
||||
Model="mlx-community/Kokoro-82M-4bit",
|
||||
Options=["voice:af_soft", "speed:1.2", "lang_code:b"]
|
||||
))
|
||||
self.assertTrue(response.success)
|
||||
|
||||
# Test TTS generation with different voice and language
|
||||
tts_req = backend_pb2.TTSRequest(
|
||||
text="Hello, this is a test with British English accent.",
|
||||
model="mlx-community/Kokoro-82M-4bit",
|
||||
voice="af_soft",
|
||||
language="b"
|
||||
)
|
||||
tts_resp = stub.TTS(tts_req)
|
||||
self.assertTrue(tts_resp.success)
|
||||
self.assertIn("TTS audio generated successfully", tts_resp.message)
|
||||
except Exception as err:
|
||||
print(err)
|
||||
self.fail("TTS with options service failed")
|
||||
finally:
|
||||
self.tearDown()
|
||||
|
||||
|
||||
def test_tts_multilingual(self):
|
||||
"""
|
||||
This method tests if TTS works with different languages
|
||||
"""
|
||||
try:
|
||||
self.setUp()
|
||||
with grpc.insecure_channel("localhost:50051") as channel:
|
||||
stub = backend_pb2_grpc.BackendStub(channel)
|
||||
response = stub.LoadModel(backend_pb2.ModelOptions(Model="mlx-community/Kokoro-82M-4bit"))
|
||||
self.assertTrue(response.success)
|
||||
|
||||
# Test Spanish TTS
|
||||
tts_req = backend_pb2.TTSRequest(
|
||||
text="Hola, esto es una prueba del sistema TTS MLX-Audio.",
|
||||
model="mlx-community/Kokoro-82M-4bit",
|
||||
voice="af_heart",
|
||||
language="e"
|
||||
)
|
||||
tts_resp = stub.TTS(tts_req)
|
||||
self.assertTrue(tts_resp.success)
|
||||
self.assertIn("TTS audio generated successfully", tts_resp.message)
|
||||
except Exception as err:
|
||||
print(err)
|
||||
self.fail("Multilingual TTS service failed")
|
||||
finally:
|
||||
self.tearDown()
|
||||
12
backend/python/mlx-audio/test.sh
Executable file
12
backend/python/mlx-audio/test.sh
Executable file
@@ -0,0 +1,12 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
backend_dir=$(dirname $0)
|
||||
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
runUnittests
|
||||
@@ -21,6 +21,21 @@ import io
|
||||
from PIL import Image
|
||||
import tempfile
|
||||
|
||||
def is_float(s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
def is_int(s):
|
||||
"""Check if a string can be converted to int."""
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
|
||||
# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
|
||||
@@ -32,22 +47,6 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
A gRPC servicer that implements the Backend service defined in backend.proto.
|
||||
"""
|
||||
|
||||
def _is_float(self, s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
def _is_int(self, s):
|
||||
"""Check if a string can be converted to int."""
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
def Health(self, request, context):
|
||||
"""
|
||||
Returns a health check message.
|
||||
@@ -87,10 +86,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
continue
|
||||
key, value = opt.split(":", 1) # Split only on first colon to handle values with colons
|
||||
|
||||
# Convert numeric values to appropriate types
|
||||
if self._is_float(value):
|
||||
if is_float(value):
|
||||
value = float(value)
|
||||
elif self._is_int(value):
|
||||
elif is_int(value):
|
||||
value = int(value)
|
||||
elif value.lower() in ["true", "false"]:
|
||||
value = value.lower() == "true"
|
||||
|
||||
@@ -24,28 +24,27 @@ _ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
|
||||
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
||||
|
||||
def is_float(s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
def is_int(s):
|
||||
"""Check if a string can be converted to int."""
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
# Implement the BackendServicer class with the service methods
|
||||
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
"""
|
||||
A gRPC servicer that implements the Backend service defined in backend.proto.
|
||||
"""
|
||||
|
||||
def _is_float(self, s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
def _is_int(self, s):
|
||||
"""Check if a string can be converted to int."""
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
def Health(self, request, context):
|
||||
"""
|
||||
Returns a health check message.
|
||||
@@ -86,9 +85,9 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
key, value = opt.split(":", 1) # Split only on first colon to handle values with colons
|
||||
|
||||
# Convert numeric values to appropriate types
|
||||
if self._is_float(value):
|
||||
if is_float(value):
|
||||
value = float(value)
|
||||
elif self._is_int(value):
|
||||
elif is_int(value):
|
||||
value = int(value)
|
||||
elif value.lower() in ["true", "false"]:
|
||||
value = value.lower() == "true"
|
||||
|
||||
23
backend/python/neutts/Makefile
Normal file
23
backend/python/neutts/Makefile
Normal file
@@ -0,0 +1,23 @@
|
||||
.PHONY: neutts
|
||||
neutts:
|
||||
bash install.sh
|
||||
|
||||
.PHONY: run
|
||||
run: neutts
|
||||
@echo "Running neutts..."
|
||||
bash run.sh
|
||||
@echo "neutts run."
|
||||
|
||||
.PHONY: test
|
||||
test: neutts
|
||||
@echo "Testing neutts..."
|
||||
bash test.sh
|
||||
@echo "neutts tested."
|
||||
|
||||
.PHONY: protogen-clean
|
||||
protogen-clean:
|
||||
$(RM) backend_pb2_grpc.py backend_pb2.py
|
||||
|
||||
.PHONY: clean
|
||||
clean: protogen-clean
|
||||
rm -rf venv __pycache__
|
||||
162
backend/python/neutts/backend.py
Normal file
162
backend/python/neutts/backend.py
Normal file
@@ -0,0 +1,162 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
This is an extra gRPC server of LocalAI for NeuTTSAir
|
||||
"""
|
||||
from concurrent import futures
|
||||
import time
|
||||
import argparse
|
||||
import signal
|
||||
import sys
|
||||
import os
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
import torch
|
||||
from neuttsair.neutts import NeuTTSAir
|
||||
import soundfile as sf
|
||||
|
||||
import grpc
|
||||
|
||||
def is_float(s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
float(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
def is_int(s):
|
||||
"""Check if a string can be converted to int."""
|
||||
try:
|
||||
int(s)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
||||
|
||||
# If MAX_WORKERS are specified in the environment use it, otherwise default to 1
|
||||
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
||||
|
||||
# Implement the BackendServicer class with the service methods
|
||||
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
"""
|
||||
BackendServicer is the class that implements the gRPC service
|
||||
"""
|
||||
def Health(self, request, context):
|
||||
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
|
||||
def LoadModel(self, request, context):
|
||||
|
||||
# Get device
|
||||
# device = "cuda" if request.CUDA else "cpu"
|
||||
if torch.cuda.is_available():
|
||||
print("CUDA is available", file=sys.stderr)
|
||||
device = "cuda"
|
||||
else:
|
||||
print("CUDA is not available", file=sys.stderr)
|
||||
device = "cpu"
|
||||
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
if mps_available:
|
||||
device = "mps"
|
||||
if not torch.cuda.is_available() and request.CUDA:
|
||||
return backend_pb2.Result(success=False, message="CUDA is not available")
|
||||
|
||||
|
||||
options = request.Options
|
||||
|
||||
# empty dict
|
||||
self.options = {}
|
||||
self.ref_text = None
|
||||
|
||||
# The options are a list of strings in this form optname:optvalue
|
||||
# We are storing all the options in a dict so we can use it later when
|
||||
# generating the images
|
||||
for opt in options:
|
||||
if ":" not in opt:
|
||||
continue
|
||||
key, value = opt.split(":")
|
||||
# if value is a number, convert it to the appropriate type
|
||||
if is_float(value):
|
||||
value = float(value)
|
||||
elif is_int(value):
|
||||
value = int(value)
|
||||
elif value.lower() in ["true", "false"]:
|
||||
value = value.lower() == "true"
|
||||
self.options[key] = value
|
||||
|
||||
codec_repo = "neuphonic/neucodec"
|
||||
if "codec_repo" in self.options:
|
||||
codec_repo = self.options["codec_repo"]
|
||||
del self.options["codec_repo"]
|
||||
if "ref_text" in self.options:
|
||||
self.ref_text = self.options["ref_text"]
|
||||
del self.options["ref_text"]
|
||||
|
||||
self.AudioPath = None
|
||||
|
||||
if os.path.isabs(request.AudioPath):
|
||||
self.AudioPath = request.AudioPath
|
||||
elif request.AudioPath and request.ModelFile != "" and not os.path.isabs(request.AudioPath):
|
||||
# get base path of modelFile
|
||||
modelFileBase = os.path.dirname(request.ModelFile)
|
||||
# modify LoraAdapter to be relative to modelFileBase
|
||||
self.AudioPath = os.path.join(modelFileBase, request.AudioPath)
|
||||
try:
|
||||
print("Preparing models, please wait", file=sys.stderr)
|
||||
self.model = NeuTTSAir(backbone_repo=request.Model, backbone_device=device, codec_repo=codec_repo, codec_device=device)
|
||||
except Exception as err:
|
||||
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
|
||||
# Implement your logic here for the LoadModel service
|
||||
# Replace this with your desired response
|
||||
return backend_pb2.Result(message="Model loaded successfully", success=True)
|
||||
|
||||
def TTS(self, request, context):
|
||||
try:
|
||||
kwargs = {}
|
||||
|
||||
# add options to kwargs
|
||||
kwargs.update(self.options)
|
||||
|
||||
ref_codes = self.model.encode_reference(self.AudioPath)
|
||||
|
||||
wav = self.model.infer(request.text, ref_codes, self.ref_text)
|
||||
|
||||
sf.write(request.dst, wav, 24000)
|
||||
except Exception as err:
|
||||
return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
|
||||
return backend_pb2.Result(success=True)
|
||||
|
||||
def serve(address):
|
||||
server = grpc.server(futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
|
||||
options=[
|
||||
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
|
||||
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
|
||||
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
|
||||
])
|
||||
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
|
||||
server.add_insecure_port(address)
|
||||
server.start()
|
||||
print("Server started. Listening on: " + address, file=sys.stderr)
|
||||
|
||||
# Define the signal handler function
|
||||
def signal_handler(sig, frame):
|
||||
print("Received termination signal. Shutting down...")
|
||||
server.stop(0)
|
||||
sys.exit(0)
|
||||
|
||||
# Set the signal handlers for SIGINT and SIGTERM
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
|
||||
try:
|
||||
while True:
|
||||
time.sleep(_ONE_DAY_IN_SECONDS)
|
||||
except KeyboardInterrupt:
|
||||
server.stop(0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Run the gRPC server.")
|
||||
parser.add_argument(
|
||||
"--addr", default="localhost:50051", help="The address to bind the server to."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
serve(args.addr)
|
||||
33
backend/python/neutts/install.sh
Executable file
33
backend/python/neutts/install.sh
Executable file
@@ -0,0 +1,33 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
backend_dir=$(dirname $0)
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
# This is here because the Intel pip index is broken and returns 200 status codes for every package name, it just doesn't return any package links.
|
||||
# This makes uv think that the package exists in the Intel pip index, and by default it stops looking at other pip indexes once it finds a match.
|
||||
# We need uv to continue falling through to the pypi default index to find optimum[openvino] in the pypi index
|
||||
# the --upgrade actually allows us to *downgrade* torch to the version provided in the Intel pip index
|
||||
if [ "x${BUILD_PROFILE}" == "xintel" ]; then
|
||||
EXTRA_PIP_INSTALL_FLAGS+=" --upgrade --index-strategy=unsafe-first-match"
|
||||
fi
|
||||
|
||||
if [ "x${BUILD_TYPE}" == "xcublas" ] || [ "x${BUILD_TYPE}" == "xl4t" ]; then
|
||||
export CMAKE_ARGS="-DGGML_CUDA=on"
|
||||
fi
|
||||
|
||||
if [ "x${BUILD_TYPE}" == "xhipblas" ]; then
|
||||
export CMAKE_ARGS="-DGGML_HIPBLAS=on"
|
||||
fi
|
||||
|
||||
EXTRA_PIP_INSTALL_FLAGS+=" --no-build-isolation"
|
||||
|
||||
git clone https://github.com/neuphonic/neutts-air neutts-air
|
||||
|
||||
cp -rfv neutts-air/neuttsair ./
|
||||
|
||||
installRequirements
|
||||
2
backend/python/neutts/requirements-after.txt
Normal file
2
backend/python/neutts/requirements-after.txt
Normal file
@@ -0,0 +1,2 @@
|
||||
datasets==4.1.1
|
||||
torchtune==0.6.1
|
||||
10
backend/python/neutts/requirements-cpu.txt
Normal file
10
backend/python/neutts/requirements-cpu.txt
Normal file
@@ -0,0 +1,10 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu
|
||||
accelerate
|
||||
torch==2.8.0
|
||||
transformers==4.56.1
|
||||
librosa==0.11.0
|
||||
neucodec>=0.0.4
|
||||
phonemizer==3.3.0
|
||||
soundfile==0.13.1
|
||||
resemble-perth==1.0.1
|
||||
llama-cpp-python
|
||||
8
backend/python/neutts/requirements-cublas12.txt
Normal file
8
backend/python/neutts/requirements-cublas12.txt
Normal file
@@ -0,0 +1,8 @@
|
||||
librosa==0.11.0
|
||||
neucodec>=0.0.4
|
||||
phonemizer==3.3.0
|
||||
soundfile==0.13.1
|
||||
torch==2.8.0
|
||||
transformers==4.56.1
|
||||
resemble-perth==1.0.1
|
||||
accelerate
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user