通过list建立websocket

routers/images.go

@@ -19,9 +19,48 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/gorilla/mux"
+	"github.com/gorilla/websocket"
 )
 
+var images_websocket_manager = models.NewWebSocketManager()
+
 func ImagesGet(w http.ResponseWriter, r *http.Request) {
+
+	// websocket 推理图像
+	if r.Header.Get("Upgrade") == "websocket" {
+		upgrader := websocket.Upgrader{}
+		conn, err := upgrader.Upgrade(w, r, nil)
+		if err != nil {
+			log.Println(err)
+			return
+		}
+		defer conn.Close()
+
+		task := r.URL.Query().Get("task")
+		if task == "" {
+			w.WriteHeader(http.StatusBadRequest)
+			w.Write([]byte("task 参数不能为空"))
+			return
+		}
+
+		wsid := images_websocket_manager.AddConnection(conn)
+		defer images_websocket_manager.RemoveConnection(wsid)
+
+		for {
+			_, msg, err := conn.ReadMessage()
+			if err != nil {
+				log.Println(err)
+				return
+			}
+			log.Println(string(msg))
+			if string(msg) == "close" {
+				break
+			}
+		}
+		return
+
+	}
+
 	var listview models.ListView
 	listview.Page = utils.ParamInt(r.URL.Query().Get("page"), 1)
 	listview.PageSize = utils.ParamInt(r.URL.Query().Get("pageSize"), 10)
@@ -67,8 +106,10 @@ func ImagesPost(w http.ResponseWriter, r *http.Request) {
 				return
 			}
 
-			// 直接创建一组图片
+			// TODO: 创建任务获得任务编号, 多张图时期望可以流式推理
 			task := uuid.New().String()
+
+			// 直接创建一组图片
 			var image_list []models.Image
 			for i := 0; i < template.Number; i++ {
 				var image models.Image