From 43cf62e5f8fe649e2768d4ad57430435ddab316a Mon Sep 17 00:00:00 2001
From: mudler <mudler@localai.io>
Date: Mon, 26 Jun 2023 00:40:54 +0200
Subject: [PATCH] refactor: allow to preload from gallery, cleanup dup code

---
 api/api.go     |  4 ++--
 api/gallery.go | 33 ++++++++++++++-------------------
 2 files changed, 16 insertions(+), 21 deletions(-)

diff --git a/api/api.go b/api/api.go
index 527258b8..05aec9b6 100644
--- a/api/api.go
+++ b/api/api.go
@@ -80,13 +80,13 @@ func App(opts ...AppOption) (*fiber.App, error) {
 	app.Use(recover.New())
 
 	if options.preloadJSONModels != "" {
-		if err := ApplyGalleryFromString(options.loader.ModelPath, options.preloadJSONModels, cm); err != nil {
+		if err := ApplyGalleryFromString(options.loader.ModelPath, options.preloadJSONModels, cm, options.galleries); err != nil {
 			return nil, err
 		}
 	}
 
 	if options.preloadModelsFromPath != "" {
-		if err := ApplyGalleryFromFile(options.loader.ModelPath, options.preloadModelsFromPath, cm); err != nil {
+		if err := ApplyGalleryFromFile(options.loader.ModelPath, options.preloadModelsFromPath, cm, options.galleries); err != nil {
 			return nil, err
 		}
 	}
diff --git a/api/gallery.go b/api/gallery.go
index 5260a239..1c0cec91 100644
--- a/api/gallery.go
+++ b/api/gallery.go
@@ -143,40 +143,35 @@ func displayDownload(fileName string, current string, total string, percentage f
 	}
 }
 
-func ApplyGalleryFromFile(modelPath, s string, cm *ConfigMerger) error {
+type galleryModel struct {
+	gallery.GalleryModel
+	ID string `json:"id"`
+}
+
+func ApplyGalleryFromFile(modelPath, s string, cm *ConfigMerger, galleries []gallery.Gallery) error {
 	dat, err := os.ReadFile(s)
 	if err != nil {
 		return err
 	}
-	var requests []gallery.GalleryModel
-	err = json.Unmarshal(dat, &requests)
-	if err != nil {
-		return err
-	}
-
-	for _, r := range requests {
-		if err := prepareModel(modelPath, r, cm, displayDownload); err != nil {
-			return err
-		}
-	}
-
-	return nil
+	return ApplyGalleryFromString(modelPath, string(dat), cm, galleries)
 }
 
-func ApplyGalleryFromString(modelPath, s string, cm *ConfigMerger) error {
-	var requests []gallery.GalleryModel
+func ApplyGalleryFromString(modelPath, s string, cm *ConfigMerger, galleries []gallery.Gallery) error {
+	var requests []galleryModel
 	err := json.Unmarshal([]byte(s), &requests)
 	if err != nil {
 		return err
 	}
 
 	for _, r := range requests {
-		if err := prepareModel(modelPath, r, cm, displayDownload); err != nil {
-			return err
+		if r.ID == "" {
+			err = prepareModel(modelPath, r.GalleryModel, cm, displayDownload)
+		} else {
+			err = gallery.InstallModelFromGallery(galleries, r.ID, modelPath, r.GalleryModel, displayDownload)
 		}
 	}
 
-	return nil
+	return err
 }
 
 func getOpStatus(g *galleryApplier) func(c *fiber.Ctx) error {