opt: add sessionId for mj task

api/handler/mj_handler.go

@@ -13,7 +13,9 @@ import (
 	"fmt"
 	"github.com/gin-gonic/gin"
 	"github.com/go-redis/redis/v8"
+	"github.com/gorilla/websocket"
 	"gorm.io/gorm"
+	"net/http"
 	"strings"
 	"sync"
 	"time"
@@ -43,6 +45,7 @@ type MidJourneyHandler struct {
 	mjService       *service.MjService
 	uploaderManager *oss.UploaderManager
 	lock            sync.Mutex
+	clients         *types.LMap[string, *types.WsClient]
 }
 
 func NewMidJourneyHandler(
@@ -57,6 +60,7 @@ func NewMidJourneyHandler(
 		uploaderManager: manager,
 		lock:            sync.Mutex{},
 		mjService:       mjService,
+		clients:         types.NewLMap[string, *types.WsClient](),
 	}
 	h.App = app
 	return &h
@@ -72,6 +76,23 @@ type notifyData struct {
 	Progress    int        `json:"progress"`
 }
 
+// Client WebSocket 客户端，用于通知任务状态变更
+func (h *MidJourneyHandler) Client(c *gin.Context) {
+	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		logger.Error(err)
+		return
+	}
+
+	sessionId := c.Query("session_id")
+	client := types.NewWsClient(ws)
+	// 关闭旧的连接
+	if h.clients.Has(sessionId) {
+		h.clients.Get(sessionId).Close()
+	}
+	h.clients.Put(sessionId, client)
+}
+
 func (h *MidJourneyHandler) Notify(c *gin.Context) {
 	token := c.GetHeader("Authorization")
 	if token != h.App.Config.ExtConfig.Token {
@@ -154,8 +175,23 @@ func (h *MidJourneyHandler) notifyHandler(c *gin.Context, data notifyData) (erro
 			return res.Error, false
 		}
 
+		var jobVo vo.MidJourneyJob
+		err := utils.CopyObject(job, &jobVo)
+		if err == nil {
+			image, err := utils.DownloadImage(jobVo.ImgURL, h.App.Config.ProxyURL)
+			if err == nil {
+				jobVo.ImgURL = "data:image/png;base64," + base64.StdEncoding.EncodeToString(image)
+			}
+
+			// 推送任务到前端
+			client := h.clients.Get(task.SessionId)
+			if client != nil {
+				utils.ReplyChunkMessage(client, jobVo)
+			}
+		}
+
 	} else if task.Src == service.TaskSrcChat { // 聊天任务
-		wsClient := h.App.MjTaskClients.Get(task.Id)
+		wsClient := h.App.MjTaskClients.Get(task.SessionId)
 		if data.Status == Finished {
 			if wsClient != nil && data.ReferenceId != "" {
 				content := fmt.Sprintf("**%s** 任务执行成功，正在从 MidJourney 服务器下载图片，请稍后...", data.Prompt)
@@ -216,7 +252,7 @@ func (h *MidJourneyHandler) notifyHandler(c *gin.Context, data notifyData) (erro
 			utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsMjImg, Content: data})
 			utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsEnd})
 			// 本次绘画完毕，移除客户端
-			h.App.MjTaskClients.Delete(task.Id)
+			h.App.MjTaskClients.Delete(task.SessionId)
 		} else {
 			// 使用代理临时转发图片
 			if data.Image.URL != "" {
@@ -235,15 +271,16 @@ func (h *MidJourneyHandler) notifyHandler(c *gin.Context, data notifyData) (erro
 // Image 创建一个绘画任务
 func (h *MidJourneyHandler) Image(c *gin.Context) {
 	var data struct {
-		Prompt  string  `json:"prompt"`
-		Rate    string  `json:"rate"`
-		Model   string  `json:"model"`
-		Chaos   int     `json:"chaos"`
-		Raw     bool    `json:"raw"`
-		Seed    int64   `json:"seed"`
-		Stylize int     `json:"stylize"`
-		Img     string  `json:"img"`
-		Weight  float32 `json:"weight"`
+		SessionId string  `json:"session_id"`
+		Prompt    string  `json:"prompt"`
+		Rate      string  `json:"rate"`
+		Model     string  `json:"model"`
+		Chaos     int     `json:"chaos"`
+		Raw       bool    `json:"raw"`
+		Seed      int64   `json:"seed"`
+		Stylize   int     `json:"stylize"`
+		Img       string  `json:"img"`
+		Weight    float32 `json:"weight"`
 	}
 	if err := c.ShouldBindJSON(&data); err != nil {
 		resp.ERROR(c, types.InvalidArgs)
@@ -268,6 +305,9 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
 			prompt += fmt.Sprintf(" --iw %f", data.Weight)
 		}
 	}
+	if data.Raw {
+		prompt += " --style raw"
+	}
 	if data.Model != "" && !strings.Contains(prompt, "--v") && !strings.Contains(prompt, "--niji") {
 		prompt += data.Model
 	}
@@ -287,12 +327,23 @@ func (h *MidJourneyHandler) Image(c *gin.Context) {
 	}
 
 	h.mjService.PushTask(service.MjTask{
-		Id:     fmt.Sprintf("%d", job.Id),
-		Src:    service.TaskSrcImg,
-		Type:   service.Image,
-		Prompt: prompt,
-		UserId: userId,
+		Id:        int(job.Id),
+		SessionId: data.SessionId,
+		Src:       service.TaskSrcImg,
+		Type:      service.Image,
+		Prompt:    prompt,
+		UserId:    userId,
 	})
+
+	var jobVo vo.MidJourneyJob
+	err := utils.CopyObject(job, &jobVo)
+	if err == nil {
+		// 推送任务到前端
+		client := h.clients.Get(data.SessionId)
+		if client != nil {
+			utils.ReplyChunkMessage(client, jobVo)
+		}
+	}
 	resp.SUCCESS(c)
 }
 
@@ -317,7 +368,7 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
 	}
 
 	idValue, _ := c.Get(types.LoginUserID)
-	jobId := data.SessionId
+	jobId := 0
 	userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
 	src := service.TaskSrc(data.Src)
 	if src == service.TaskSrcImg {
@@ -330,14 +381,25 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
 			CreatedAt: time.Now(),
 		}
 		if res := h.db.Create(&job); res.Error == nil {
-			jobId = fmt.Sprintf("%d", job.Id)
+			jobId = int(job.Id)
 		} else {
 			resp.ERROR(c, "添加任务失败："+res.Error.Error())
 			return
 		}
+
+		var jobVo vo.MidJourneyJob
+		err := utils.CopyObject(job, &jobVo)
+		if err == nil {
+			// 推送任务到前端
+			client := h.clients.Get(data.SessionId)
+			if client != nil {
+				utils.ReplyChunkMessage(client, jobVo)
+			}
+		}
 	}
 	h.mjService.PushTask(service.MjTask{
 		Id:          jobId,
+		SessionId:   data.SessionId,
 		Src:         src,
 		Type:        service.Upscale,
 		Prompt:      data.Prompt,
@@ -358,6 +420,7 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
 			h.App.MjTaskClients.Put(data.SessionId, wsClient)
 		}
 	}
+
 	resp.SUCCESS(c)
 }
 
@@ -370,7 +433,7 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
 	}
 
 	idValue, _ := c.Get(types.LoginUserID)
-	jobId := data.SessionId
+	jobId := 0
 	userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
 	src := service.TaskSrc(data.Src)
 	if src == service.TaskSrcImg {
@@ -384,14 +447,25 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
 			CreatedAt: time.Now(),
 		}
 		if res := h.db.Create(&job); res.Error == nil {
-			jobId = fmt.Sprintf("%d", job.Id)
+			jobId = int(job.Id)
 		} else {
 			resp.ERROR(c, "添加任务失败："+res.Error.Error())
 			return
 		}
+
+		var jobVo vo.MidJourneyJob
+		err := utils.CopyObject(job, &jobVo)
+		if err == nil {
+			// 推送任务到前端
+			client := h.clients.Get(data.SessionId)
+			if client != nil {
+				utils.ReplyChunkMessage(client, jobVo)
+			}
+		}
 	}
 	h.mjService.PushTask(service.MjTask{
 		Id:          jobId,
+		SessionId:   data.SessionId,
 		Src:         src,
 		Type:        service.Variation,
 		Prompt:      data.Prompt,

api/service/function/func_mj.go

@@ -22,14 +22,14 @@ func (f FuncMidJourney) Invoke(params map[string]interface{}) (string, error) {
 	logger.Infof("MJ 绘画参数：%+v", params)
 	prompt := utils.InterfaceToString(params["prompt"])
 	f.service.PushTask(service.MjTask{
-		Id:     utils.InterfaceToString(params["session_id"]),
-		Src:    service.TaskSrcChat,
-		Type:   service.Image,
-		Prompt: prompt,
-		UserId: utils.IntValue(utils.InterfaceToString(params["user_id"]), 0),
-		RoleId: utils.IntValue(utils.InterfaceToString(params["role_id"]), 0),
-		Icon:   utils.InterfaceToString(params["icon"]),
-		ChatId: utils.InterfaceToString(params["chat_id"]),
+		SessionId: utils.InterfaceToString(params["session_id"]),
+		Src:       service.TaskSrcChat,
+		Type:      service.Image,
+		Prompt:    prompt,
+		UserId:    utils.IntValue(utils.InterfaceToString(params["user_id"]), 0),
+		RoleId:    utils.IntValue(utils.InterfaceToString(params["role_id"]), 0),
+		Icon:      utils.InterfaceToString(params["icon"]),
+		ChatId:    utils.InterfaceToString(params["chat_id"]),
 	})
 	return prompt, nil
 }

api/service/mj_service.go

@@ -41,7 +41,8 @@ const (
 )
 
 type MjTask struct {
-	Id          string   `json:"id"`
+	Id          int      `json:"id"`
+	SessionId   string   `json:"session_id"`
 	Src         TaskSrc  `json:"src"`
 	Type        TaskType `json:"type"`
 	UserId      int      `json:"user_id"`

api/utils/net.go

@@ -12,7 +12,7 @@ import (
 var logger = logger2.GetLogger()
 
 // ReplyChunkMessage 回复客户片段端消息
-func ReplyChunkMessage(client *types.WsClient, message types.WsMessage) {
+func ReplyChunkMessage(client *types.WsClient, message interface{}) {
 	msg, err := json.Marshal(message)
 	if err != nil {
 		logger.Errorf("Error for decoding json data: %v", err.Error())

web/src/assets/css/image-mj.css

@@ -185,31 +185,15 @@
   display: flex;
   justify-content: center;
   align-items: center;
-  background-color: rgba(0,0,0,0.5);
 }
 .page-mj .inner .task-list-box .running-job-list .job-item .job-item-inner .progress span {
   font-size: 20px;
   color: #fff;
 }
-.page-mj .inner .task-list-box .running-job-list .job-item .el-image {
+.page-mj .inner .task-list-box .finish-job-list .job-item {
   width: 100%;
   height: 100%;
 }
-.page-mj .inner .task-list-box .running-job-list .job-item .el-image .image-slot {
-  display: flex;
-  flex-flow: column;
-  justify-content: center;
-  align-items: center;
-  height: 100%;
-  color: #fff;
-}
-.page-mj .inner .task-list-box .running-job-list .job-item .el-image .image-slot .iconfont {
-  font-size: 50px;
-  margin-bottom: 10px;
-}
-.page-mj .inner .task-list-box .finish-job-list .job-item {
-  margin-bottom: 20px;
-}
 .page-mj .inner .task-list-box .finish-job-list .job-item .opt .opt-line {
   margin: 6px 0;
 }
@@ -233,6 +217,37 @@
 .page-mj .inner .task-list-box .finish-job-list .job-item .opt .opt-line ul li a:hover {
   background-color: #6d6f78;
 }
+.page-mj .inner .task-list-box .el-image {
+  width: 100%;
+  height: 100%;
+  max-height: 240px;
+}
+.page-mj .inner .task-list-box .el-image img {
+  height: 240px;
+}
+.page-mj .inner .task-list-box .el-image .el-image-viewer__wrapper img {
+  width: auto;
+  height: auto;
+}
+.page-mj .inner .task-list-box .el-image .image-slot {
+  display: flex;
+  flex-flow: column;
+  justify-content: center;
+  align-items: center;
+  height: 100%;
+  min-height: 200px;
+  color: #fff;
+}
+.page-mj .inner .task-list-box .el-image .image-slot .iconfont {
+  font-size: 50px;
+  margin-bottom: 10px;
+}
+.page-mj .inner .task-list-box .el-image.upscale {
+  max-height: 304px;
+}
+.page-mj .inner .task-list-box .el-image.upscale img {
+  height: 304px;
+}
 .mj-list-item-prompt .el-icon {
   margin-left: 10px;
   cursor: pointer;

web/src/assets/css/image-mj.styl

@@ -247,7 +247,6 @@
 
       .finish-job-list {
         .job-item {
-          margin-bottom 20px
           width 100%
           height 100%
 
@@ -316,6 +315,14 @@
           }
         }
       }
+
+      .el-image.upscale {
+        max-height 304px
+
+        img {
+          height 304px
+        }
+      }
     }
   }
 

web/src/utils/http.js

@@ -1,7 +1,7 @@
 import axios from 'axios'
 import {getAdminToken, getSessionId, getUserToken} from "@/store/session";
 
-axios.defaults.timeout = 10000
+axios.defaults.timeout = 30000
 axios.defaults.baseURL = process.env.VUE_APP_API_HOST
 axios.defaults.withCredentials = true;
 axios.defaults.headers.post['Content-Type'] = 'application/json'

web/src/views/ImageMj.vue

@@ -285,10 +285,10 @@
                       placement="top-start"
                       title="提示词"
                       :width="240"
-                      trigger="hover"
+                      trigger="click"
                   >
                     <template #reference>
-                      <el-image :src="scope.item.img_url"
+                      <el-image :src="scope.item.img_url" :class="scope.item.type === 'upscale'?'upscale':''"
                                 :zoom-rate="1.2"
                                 :preview-src-list="previewImgList"
                                 fit="cover"
@@ -319,7 +319,7 @@
                     </template>
                   </el-popover>
 
-                  <div class="opt">
+                  <div class="opt" v-if="scope.item.type !== 'upscale'">
                     <div class="opt-line">
                       <ul>
                         <li><a @click="upscale(1,scope.item)">U1</a></li>
@@ -352,7 +352,7 @@
 </template>
 
 <script setup>
-import {onMounted, ref} from "vue"
+import {nextTick, onMounted, ref} from "vue"
 import {DeleteFilled, DocumentCopy, InfoFilled, Picture, Plus} from "@element-plus/icons-vue";
 import Compressor from "compressorjs";
 import {httpGet, httpPost} from "@/utils/http";
@@ -361,7 +361,9 @@ import ItemList from "@/components/ItemList.vue";
 import Clipboard from "clipboard";
 import {checkSession} from "@/action/session";
 import {useRouter} from "vue-router";
-import {getSessionId} from "@/store/session";
+import {getSessionId, getUserToken} from "@/store/session";
+import {randString} from "@/utils/libs";
+import hl from "highlight.js";
 
 const listBoxHeight = ref(window.innerHeight - 40)
 const mjBoxHeight = ref(window.innerHeight - 150)
@@ -395,10 +397,85 @@ const runningJobs = ref([])
 const finishedJobs = ref([])
 const previewImgList = ref([])
 const router = useRouter()
+
+const socket = ref(null)
+
+const connect = () => {
+  let host = process.env.VUE_APP_WS_HOST
+  if (host === '') {
+    if (location.protocol === 'https:') {
+      host = 'wss://' + location.host;
+    } else {
+      host = 'ws://' + location.host;
+    }
+  }
+  const _socket = new WebSocket(host + `/api/mj/client?session_id=${getSessionId()}`);
+  _socket.addEventListener('open', () => {
+    socket.value = _socket;
+  });
+
+  _socket.addEventListener('message', event => {
+    if (event.data instanceof Blob) {
+      const reader = new FileReader();
+      reader.readAsText(event.data, "UTF-8");
+      reader.onload = () => {
+        const data = JSON.parse(String(reader.result));
+        console.log(data)
+        let isNew = false
+        if (data.progress === 100) {
+          for (let i = 0; i < finishedJobs.value.length; i++) {
+            if (finishedJobs.value[i].id === data.id) {
+              isNew = false
+              break
+            }
+          }
+          if (isNew) {
+            finishedJobs.value.unshift(data)
+          }
+        } else {
+          for (let i = 0; i < runningJobs.value; i++) {
+            if (runningJobs.value[i].id === data.id) {
+              isNew = false
+              runningJobs.value[i] = data
+              break
+            }
+          }
+          if (isNew) {
+            runningJobs.value.push(data)
+          }
+        }
+      }
+    }
+  });
+
+  _socket.addEventListener('close', () => {
+    connect()
+  });
+}
+
 onMounted(() => {
   checkSession().then(() => {
-    fetchFinishedJobs()
-    fetchRunningJobs()
+    // 获取运行中的任务
+    httpGet("/api/mj/jobs?status=1").then(res => {
+      if (finishedJobs.value.length !== res.data.length) {
+        finishedJobs.value = res.data
+      }
+      previewImgList.value = []
+      for (let index in finishedJobs.value) {
+        previewImgList.value.push(finishedJobs.value[index]["img_url"])
+      }
+    }).catch(e => {
+      ElMessage.error("获取任务失败：" + e.message)
+    })
+
+    // 获取运行中的任务
+    httpGet("/api/mj/jobs?status=0").then(res => {
+      if (runningJobs.value.length !== res.data.length) {
+        runningJobs.value = res.data
+      }
+    }).catch(e => {
+      ElMessage.error("获取任务失败：" + e.message)
+    })
   }).catch(() => {
     router.push('/login')
   });
@@ -413,32 +490,6 @@ onMounted(() => {
   })
 })
 
-const fetchFinishedJobs = () => {
-  httpGet("/api/mj/jobs?status=1").then(res => {
-    finishedJobs.value = res.data
-    previewImgList.value = []
-    for (let index in finishedJobs.value) {
-      previewImgList.value.push(finishedJobs.value[index]["img_url"])
-    }
-    setTimeout(() => {
-      fetchFinishedJobs()
-    }, 2000)
-  }).catch(e => {
-    ElMessage.error("获取任务失败：" + e.message)
-  })
-}
-
-const fetchRunningJobs = () => {
-  httpGet("/api/mj/jobs?status=0").then(res => {
-    runningJobs.value = res.data
-    setTimeout(() => {
-      fetchRunningJobs()
-    }, 1000)
-  }).catch(e => {
-    ElMessage.error("获取任务失败：" + e.message)
-  })
-}
-
 // 切换图片比例
 const changeRate = (item) => {
   params.value.rate = item.value
@@ -489,6 +540,10 @@ const generate = () => {
     promptRef.value.focus()
     return ElMessage.error("请输入绘画提示词！")
   }
+  if (params.value.model.indexOf("niji") !== -1 && params.value.raw) {
+    return ElMessage.error("动漫模型不允许启用原始模式")
+  }
+  params.value.session_id = getSessionId()
   httpPost("/api/mj/image", params.value).then(() => {
     ElMessage.success("绘画任务推送成功，请耐心等待任务执行...")
   }).catch(e => {