From c1888a8062fc7a000457714d8426584d42d0f142 Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@users.noreply.github.com>
Date: Sat, 30 Dec 2023 12:55:18 -0500
Subject: [PATCH] feat(preload): prepare models in galleries (#1515)

Previously if applying models from the gallery API, we didn't actually
allowed remote URLs as models as nothing was actually downloading the
models referenced in the configuration file. Now we call Preload after
we have all the models loaded in memory.
---
 api/api.go             | 24 ++++++++++++------------
 api/config/config.go   |  1 +
 api/localai/gallery.go |  6 ++++++
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/api/api.go b/api/api.go
index e3409effe17a..b89f47e1ba0e 100644
--- a/api/api.go
+++ b/api/api.go
@@ -51,6 +51,18 @@ func Startup(opts ...options.AppOption) (*options.Option, *config.ConfigLoader,
 		log.Error().Msgf("error downloading models: %s", err.Error())
 	}
 
+	if options.PreloadJSONModels != "" {
+		if err := localai.ApplyGalleryFromString(options.Loader.ModelPath, options.PreloadJSONModels, cl, options.Galleries); err != nil {
+			return nil, nil, err
+		}
+	}
+
+	if options.PreloadModelsFromPath != "" {
+		if err := localai.ApplyGalleryFromFile(options.Loader.ModelPath, options.PreloadModelsFromPath, cl, options.Galleries); err != nil {
+			return nil, nil, err
+		}
+	}
+
 	if options.Debug {
 		for _, v := range cl.ListConfigs() {
 			cfg, _ := cl.GetConfig(v)
@@ -67,18 +79,6 @@ func Startup(opts ...options.AppOption) (*options.Option, *config.ConfigLoader,
 		}
 	}
 
-	if options.PreloadJSONModels != "" {
-		if err := localai.ApplyGalleryFromString(options.Loader.ModelPath, options.PreloadJSONModels, cl, options.Galleries); err != nil {
-			return nil, nil, err
-		}
-	}
-
-	if options.PreloadModelsFromPath != "" {
-		if err := localai.ApplyGalleryFromFile(options.Loader.ModelPath, options.PreloadModelsFromPath, cl, options.Galleries); err != nil {
-			return nil, nil, err
-		}
-	}
-
 	// turn off any process that was started by GRPC if the context is canceled
 	go func() {
 		<-options.Context.Done()
diff --git a/api/config/config.go b/api/config/config.go
index 4c28f49b33a5..350d2c6585f8 100644
--- a/api/config/config.go
+++ b/api/config/config.go
@@ -267,6 +267,7 @@ func (cm *ConfigLoader) ListConfigs() []string {
 	return res
 }
 
+// Preload prepare models if they are not local but url or huggingface repositories
 func (cm *ConfigLoader) Preload(modelPath string) error {
 	cm.Lock()
 	defer cm.Unlock()
diff --git a/api/localai/gallery.go b/api/localai/gallery.go
index a1cab6dc3edd..a2ad5bd1ac46 100644
--- a/api/localai/gallery.go
+++ b/api/localai/gallery.go
@@ -130,6 +130,12 @@ func (g *galleryApplier) Start(c context.Context, cm *config.ConfigLoader) {
 					continue
 				}
 
+				err = cm.Preload(g.modelPath)
+				if err != nil {
+					updateError(err)
+					continue
+				}
+
 				g.updateStatus(op.id, &galleryOpStatus{Processed: true, Message: "completed", Progress: 100})
 			}
 		}