From e9637f85a71aad693df9662bf0643299ee42c4ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=A1=9C=E8=8F=AF?= <huan0016@gmail.com>
Date: Sun, 2 Jul 2023 06:13:47 +0800
Subject: [PATCH] =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E5=88=87=E6=8D=A2?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 models/Model.go | 74 +++++++++++++++++++++++--------------------------
 1 file changed, 34 insertions(+), 40 deletions(-)

diff --git a/models/Model.go b/models/Model.go
index cd90383..5d07c6e 100644
--- a/models/Model.go
+++ b/models/Model.go
@@ -81,33 +81,7 @@ func (model *Model) Inference(image_list []Image, callback func(Image)) {
 		data, _ := json.MarshalIndent(server, "", "    ")
 		fmt.Println(string(data))
 
-		//var form = struct {
-		//	Components []struct {
-		//		ID    int    `json:"id"`
-		//		Type  string `json:"type"`
-		//		Props struct {
-		//			Value string `json:"value"`
-		//		}
-		//	} `json:"components"`
-		//}{}
-		//// 检查当前是否为目标模型, 不是则执行切换模型 http://106.15.192.42:7860/config
-		//if err := goreq.Get(fmt.Sprintf("http://%s:%d/config", server.IP, server.Port)).Do().BindJSON(&form); err != nil {
-		//	log.Println("获取推理机配置失败:", err)
-		//	return
-		//}
-		//var isSet = false
-		//for _, component := range form.Components {
-		//	if component.Type == "dropdown" && component.ID == 1514 && component.Props.Value == model.Name {
-		//		log.Println("当前推理机已经部署了目标模型")
-		//		isSet = true
-		//		break
-		//	}
-		//}
-		//if !isSet {
-		//	log.Println("当前推理机未部署目标模型, 开始部署目标模型")
-		//	// 没有切换模型接口
-		//	return
-		//}
+		// TODO: 上传模型到推理机
 
 		// 记录到推理机
 		server.Models = append(server.Models, model.ID)
@@ -121,20 +95,40 @@ func (model *Model) Inference(image_list []Image, callback func(Image)) {
 		configs.ORMDB().Take(&server)
 	}
 
-	// 排队等待切换模型(先检查是否已经切换完成)
-	for {
+	// 检查推理机是否已经加载了模型
+	if server.ModelID != model.ID {
+		log.Println("推理机未加载模型, 开始排队加载模型")
+
+		// 执行切换模型(推理机需要先处理完当前的任务才能加载新的模型)
+		if err := goreq.Post(fmt.Sprintf("http://%s:%d/sdapi/v1/options", server.IP, server.Port)).SetJsonBody(map[string]interface{}{
+			"sd_model_checkpoint":      model.ModelCheckpoint,
+			"CLIP_stop_at_last_layers": 2,
+		}).Do().Error(); err != nil {
+			log.Println("切换模型失败:", err)
+			return
+		}
+
 		var form = struct {
-			Components []struct {
-				ID    int    `json:"id"`
-				Type  string `json:"type"`
-				Props struct {
-					Value string `json:"value"`
-				}
-			} `json:"components"`
+			SdCheckpointHash  string `json:"sd_checkpoint_hash"`
+			SdModelCheckpoint string `json:"sd_model_checkpoint"`
 		}{}
-		// 检查当前是否为目标模型, 不是则执行切换模型 http://
-		if err := goreq.Get(fmt.Sprintf("http://%s:%d/config", server.IP, server.Port)).Do().BindJSON(&form); err != nil {
-			log.Println("获取推理机配置失败:", err)
+
+		for {
+			if err := goreq.Get(fmt.Sprintf("http://%s:%d/sdapi/v1/options", server.IP, server.Port)).Do().BindJSON(&form); err != nil {
+				log.Println("获取推理机配置失败:", err)
+				return
+			}
+			if form.SdModelCheckpoint == model.ModelCheckpoint {
+				log.Println("模型切换完成:", form.SdModelCheckpoint)
+				break
+			}
+			time.Sleep(time.Second)
+		}
+
+		// 更新推理机模型ID
+		server.ModelID = model.ID
+		if err := configs.ORMDB().Save(&server).Error; err != nil {
+			log.Println("更新推理机模型ID失败:", err)
 			return
 		}
 	}
@@ -177,7 +171,7 @@ func (model *Model) Inference(image_list []Image, callback func(Image)) {
 		image_list[index].Progress = 100
 		callback(image_list[index])
 	}
-	log.Println("模型未部署到推理機, 取消推理模型")
+	log.Println("推理完成:", model.ID, model.Name)
 }
 
 // 将base64编码的图片保存到本地webp