feat(oci): support OCI images and Ollama models (#2628)

* Support specifying oci:// and ollama:// for model URLs

Fixes: https://github.com/mudler/LocalAI/issues/2527
Fixes: https://github.com/mudler/LocalAI/issues/1028

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Lower watcher warnings

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Allow to install ollama models from CLI

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fixup tests

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Do not keep file ownership

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Skip test on darwin

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto 2024-06-22 08:17:41 +02:00 committed by GitHub
parent e265a618d9
commit f569237a50
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 638 additions and 97 deletions

View file

@ -11,19 +11,24 @@ import (
"strconv"
"strings"
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
"github.com/go-skynet/LocalAI/pkg/oci"
"github.com/go-skynet/LocalAI/pkg/utils"
"github.com/rs/zerolog/log"
)
const (
HuggingFacePrefix = "huggingface://"
OCIPrefix = "oci://"
OllamaPrefix = "ollama://"
HTTPPrefix = "http://"
HTTPSPrefix = "https://"
GithubURI = "github:"
GithubURI2 = "github://"
)
func GetURI(url string, basePath string, f func(url string, i []byte) error) error {
func DownloadAndUnmarshal(url string, basePath string, f func(url string, i []byte) error) error {
url = ConvertURL(url)
if strings.HasPrefix(url, "file://") {
@ -76,9 +81,15 @@ func LooksLikeURL(s string) bool {
strings.HasPrefix(s, HTTPSPrefix) ||
strings.HasPrefix(s, HuggingFacePrefix) ||
strings.HasPrefix(s, GithubURI) ||
strings.HasPrefix(s, OllamaPrefix) ||
strings.HasPrefix(s, OCIPrefix) ||
strings.HasPrefix(s, GithubURI2)
}
func LooksLikeOCI(s string) bool {
return strings.HasPrefix(s, OCIPrefix) || strings.HasPrefix(s, OllamaPrefix)
}
func ConvertURL(s string) string {
switch {
case strings.HasPrefix(s, GithubURI2):
@ -149,6 +160,32 @@ func removePartialFile(tmpFilePath string) error {
func DownloadFile(url string, filePath, sha string, fileN, total int, downloadStatus func(string, string, string, float64)) error {
url = ConvertURL(url)
if LooksLikeOCI(url) {
progressStatus := func(desc ocispec.Descriptor) io.Writer {
return &progressWriter{
fileName: filePath,
total: desc.Size,
hash: sha256.New(),
fileNo: fileN,
totalFiles: total,
downloadStatus: downloadStatus,
}
}
if strings.HasPrefix(url, OllamaPrefix) {
url = strings.TrimPrefix(url, OllamaPrefix)
return oci.OllamaFetchModel(url, filePath, progressStatus)
}
url = strings.TrimPrefix(url, OCIPrefix)
img, err := oci.GetImage(url, "", nil, nil)
if err != nil {
return fmt.Errorf("failed to get image %q: %v", url, err)
}
return oci.ExtractOCIImage(img, filepath.Dir(filePath))
}
// Check if the file already exists
_, err := os.Stat(filePath)
if err == nil {