生成图像(默认参数)

routers/images.go

@@ -62,8 +62,8 @@ func ImagesGet(w http.ResponseWriter, r *http.Request) {
 		log.Println("任务编号:", task, "任务数量:", len(image_list))
 
 		// 加入连接池
-		wsid := images_websocket_manager.AddConnection(conn)
-		defer images_websocket_manager.RemoveConnection(wsid)
+		images_websocket_manager.AddConnection(conn, task)
+		defer images_websocket_manager.RemoveConnection(conn)
 
 		for {
 			_, msg, err := conn.ReadMessage()
@@ -111,6 +111,7 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
 				Scheduler         string  `json:"scheduler"`           // 调度器
 				Seed              string  `json:"seed"`                // 随机种子(单张图生成时使用)
 				Number            int     `json:"number"`              // 生成数量
+				ModelID           int     `json:"model_id"`            // 模型ID
 			}{}
 			body, err := ioutil.ReadAll(r.Body)
 			if err != nil {
@@ -136,12 +137,26 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
 			if template.GuidanceScale > 20 {
 				template.GuidanceScale = 20
 			}
+			if template.Scheduler == "" {
+				template.Scheduler = "DDIM"
+			}
+			if template.ModelID <= 0 {
+				w.WriteHeader(http.StatusBadRequest)
+				w.Write([]byte("model_id 参数不能为空"))
+				return
+			}
 
-			// TODO: 创建任务获得任务编号, 多张图时期望可以流式推理
-			task := uuid.New().String()
+			// 从数据库中读取模型信息
+			var model models.Model = models.Model{ID: template.ModelID}
+			if err := configs.ORMDB().First(&model).Error; err != nil {
+				w.WriteHeader(http.StatusBadRequest)
+				w.Write([]byte("模型不存在"))
+				return
+			}
 
 			// 直接创建一组图片
 			var image_list []models.Image
+			var task string = uuid.New().String()
 			for i := 0; i < template.Number; i++ {
 				var image models.Image
 				image.UserID = account.ID
@@ -157,6 +172,10 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
 				image_list = append(image_list, image)
 			}
 
+			go model.Inference(image_list, func() {
+				images_websocket_manager.NotifyTaskChange(task, image_list)
+			})
+
 			// 存储图片信息到数据库
 			if err := configs.ORMDB().Create(&image_list).Error; err != nil {
 				log.Println(err)