模型切换

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 {
+		// TODO: 上传模型到推理机
-		//	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
-		//}
 
 		// 记录到推理机
 		server.Models = append(server.Models, model.ID)
@@ -121,22 +95,42 @@ func (model *Model) Inference(image_list []Image, callback func(Image)) {
 		configs.ORMDB().Take(&server)
 	}
 
-	// 排队等待切换模型(先检查是否已经切换完成)
+	// 检查推理机是否已经加载了模型
-	for {
+	if server.ModelID != model.ID {
-		var form = struct {
+		log.Println("推理机未加载模型, 开始排队加载模型")
-			Components []struct {
+
-				ID    int    `json:"id"`
+		// 执行切换模型(推理机需要先处理完当前的任务才能加载新的模型)
-				Type  string `json:"type"`
+		if err := goreq.Post(fmt.Sprintf("http://%s:%d/sdapi/v1/options", server.IP, server.Port)).SetJsonBody(map[string]interface{}{
-				Props struct {
+			"sd_model_checkpoint":      model.ModelCheckpoint,
-					Value string `json:"value"`
+			"CLIP_stop_at_last_layers": 2,
+		}).Do().Error(); err != nil {
+			log.Println("切换模型失败:", err)
+			return
 		}
-			} `json:"components"`
+
+		var form = struct {
+			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 {
+		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