feat: midjourney image upscale function is ready

api/Makefile

@@ -12,7 +12,7 @@ linux:
 .PHONY: linux
 
 darwin:
-	CGO_ENABLED=1 GOOS=darwin GOARCH=amd64 go build -o bin/$(NAME)-amd64-darwin main.go
+	CGO_ENABLED=0 GOOS=darwin GOARCH=amd64 go build -o bin/$(NAME)-amd64-darwin main.go
 .PHONY: darwin
 
 clean:

api/core/app_server.go

@@ -151,7 +151,7 @@ func corsMiddleware() gin.HandlerFunc {
 			c.Header("Access-Control-Allow-Origin", origin)
 			c.Header("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE, UPDATE")
 			//允许跨域设置可以返回其他子段，可以自定义字段
-			c.Header("Access-Control-Allow-Headers", "Authorization, Content-Length, Content-Type, ChatGPT-TOKEN, ADMIN-SESSION-TOKEN")
+			c.Header("Access-Control-Allow-Headers", "Authorization, Content-Length, Content-Type")
 			// 允许浏览器（客户端）可以解析的头部 （重要）
 			c.Header("Access-Control-Expose-Headers", "Content-Length, Access-Control-Allow-Origin, Access-Control-Allow-Headers")
 			//设置缓存时间

api/core/types/client.go

@@ -6,18 +6,14 @@ import (
 	"sync"
 )
 
-var ErrConClosed = errors.New("connection closed")
+var ErrConClosed = errors.New("connection Closed")
-
-type Client interface {
-	Close()
-}
 
 // WsClient websocket client
 type WsClient struct {
 	Conn   *websocket.Conn
 	lock   sync.Mutex
 	mt     int
-	closed bool
+	Closed bool
 }
 
 func NewWsClient(conn *websocket.Conn) *WsClient {
@@ -25,7 +21,7 @@ func NewWsClient(conn *websocket.Conn) *WsClient {
 		Conn:   conn,
 		lock:   sync.Mutex{},
 		mt:     2, // fixed bug for 'Invalid UTF-8 in text frame'
-		closed: false,
+		Closed: false,
 	}
 }
 
@@ -33,7 +29,7 @@ func (wc *WsClient) Send(message []byte) error {
 	wc.lock.Lock()
 	defer wc.lock.Unlock()
 
-	if wc.closed {
+	if wc.Closed {
 		return ErrConClosed
 	}
 
@@ -41,7 +37,7 @@ func (wc *WsClient) Send(message []byte) error {
 }
 
 func (wc *WsClient) Receive() (int, []byte, error) {
-	if wc.closed {
+	if wc.Closed {
 		return 0, nil, ErrConClosed
 	}
 
@@ -52,10 +48,10 @@ func (wc *WsClient) Close() {
 	wc.lock.Lock()
 	defer wc.lock.Unlock()
 
-	if wc.closed {
+	if wc.Closed {
 		return
 	}
 
 	_ = wc.Conn.Close()
-	wc.closed = true
+	wc.Closed = true
 }

api/handler/mj_handler.go

@@ -3,12 +3,14 @@ package handler
 import (
 	"chatplus/core"
 	"chatplus/core/types"
+	"chatplus/service/function"
 	"chatplus/store"
 	"chatplus/store/model"
 	"chatplus/utils"
 	"chatplus/utils/resp"
 	"github.com/gin-gonic/gin"
 	"gorm.io/gorm"
+	"time"
 )
 
 type TaskStatus string
@@ -34,10 +36,11 @@ type MidJourneyHandler struct {
 	BaseHandler
 	leveldb *store.LevelDB
 	db      *gorm.DB
+	mjFunc  function.FuncMidJourney
 }
 
-func NewMidJourneyHandler(app *core.AppServer, leveldb *store.LevelDB, db *gorm.DB) *MidJourneyHandler {
+func NewMidJourneyHandler(app *core.AppServer, leveldb *store.LevelDB, db *gorm.DB, functions map[string]function.Function) *MidJourneyHandler {
-	h := MidJourneyHandler{leveldb: leveldb, db: db}
+	h := MidJourneyHandler{leveldb: leveldb, db: db, mjFunc: functions[types.FuncMidJourney].(function.FuncMidJourney)}
 	h.App = app
 	return &h
 }
@@ -50,34 +53,43 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
 	}
 
 	var data struct {
-		Type      string     `json:"type"`
+		MessageId   string     `json:"message_id"`
-		MessageId string     `json:"message_id"`
+		ReferenceId string     `json:"reference_id"`
-		Image     Image      `json:"image"`
+		Image       Image      `json:"image"`
-		Content   string     `json:"content"`
+		Content     string     `json:"content"`
-		Prompt    string     `json:"prompt"`
+		Prompt      string     `json:"prompt"`
-		Status    TaskStatus `json:"status"`
+		Status      TaskStatus `json:"status"`
-		Key       string     `json:"key"`
+		Key         string     `json:"key"`
 	}
 	if err := c.ShouldBindJSON(&data); err != nil || data.Prompt == "" {
 		resp.ERROR(c, types.InvalidArgs)
 		return
 	}
 
-	key := utils.Sha256(data.Prompt)
+	logger.Infof("收到 MidJourney 回调请求：%+v", data)
-	data.Key = key
+
+	// the job is saved
+	var job model.MidJourneyJob
+	res := h.db.Where("message_id = ?", data.MessageId).First(&job)
+	if res.Error == nil {
+		resp.SUCCESS(c)
+		return
+	}
+
+	data.Key = utils.Sha256(data.Prompt)
 	//logger.Info(data.Prompt, ",", key)
 	if data.Status == Finished {
 		var task types.MjTask
-		err := h.leveldb.Get(types.TaskStorePrefix+key, &task)
+		err := h.leveldb.Get(types.TaskStorePrefix+data.Key, &task)
 		if err != nil {
 			logger.Error("error with get MidJourney task: ", err)
-			resp.ERROR(c)
+			resp.ERROR(c, err.Error())
 			return
 		}
 
 		// TODO: 是否需要把图片下载到本地服务器？
 
-		historyUserMsg := model.HistoryMessage{
+		message := model.HistoryMessage{
 			UserId:     task.UserId,
 			ChatId:     task.ChatId,
 			RoleId:     task.RoleId,
@@ -87,18 +99,30 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
 			Tokens:     0,
 			UseContext: false,
 		}
-		res := h.db.Save(&historyUserMsg)
+		res := h.db.Create(&message)
 		if res.Error != nil {
-			logger.Error("error with save MidJourney message: ", res.Error)
+			logger.Error("error with save chat history message: ", res.Error)
 		}
 
-		// delete task from leveldb
+		// save the job
-		_ = h.leveldb.Delete(types.TaskStorePrefix + key)
+		job.UserId = task.UserId
+		job.ChatId = task.ChatId
+		job.MessageId = data.MessageId
+		job.Content = data.Content
+		job.Prompt = data.Prompt
+		job.Image = utils.JsonEncode(data.Image)
+		job.Hash = data.Image.Hash
+		job.CreatedAt = time.Now()
+		res = h.db.Create(&job)
+		if res.Error != nil {
+			logger.Error("error with save MidJourney Job: ", res.Error)
+		}
 	}
 
 	// 推送消息到客户端
-	wsClient := h.App.MjTaskClients.Get(key)
+	wsClient := h.App.MjTaskClients.Get(data.Key)
 	if wsClient == nil { // 客户端断线，则丢弃
+		logger.Errorf("Client is offline: %+v", data)
 		resp.SUCCESS(c, "Client is offline")
 		return
 	}
@@ -107,9 +131,48 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
 		utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsMjImg, Content: data})
 		utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsEnd})
 		// delete client
-		h.App.MjTaskClients.Delete(key)
+		h.App.MjTaskClients.Delete(data.Key)
 	} else {
 		utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsMjImg, Content: data})
 	}
 	resp.SUCCESS(c, "SUCCESS")
 }
+
+// Upscale send upscale command to MidJourney Bot
+func (h *MidJourneyHandler) Upscale(c *gin.Context) {
+	var data struct {
+		Index       int32  `json:"index"`
+		MessageId   string `json:"message_id"`
+		MessageHash string `json:"message_hash"`
+		SessionId   string `json:"session_id"`
+		Key         string `json:"key"`
+	}
+
+	if err := c.ShouldBindJSON(&data); err != nil ||
+		data.SessionId == "" ||
+		data.Key == "" {
+		resp.ERROR(c, types.InvalidArgs)
+		return
+	}
+	wsClient := h.App.ChatClients.Get(data.SessionId)
+	if wsClient == nil {
+		resp.ERROR(c, "No Websocket client online")
+		return
+	}
+
+	err := h.mjFunc.Upscale(function.MjUpscaleReq{
+		Index:       data.Index,
+		MessageId:   data.MessageId,
+		MessageHash: data.MessageHash,
+	})
+	if err != nil {
+		resp.ERROR(c, err.Error())
+		return
+	}
+
+	utils.ReplyMessage(wsClient, "已推送放大图片任务到 MidJourney 机器人，请耐心等待任务执行...")
+	if h.App.MjTaskClients.Get(data.Key) == nil {
+		h.App.MjTaskClients.Put(data.Key, wsClient)
+	}
+	resp.SUCCESS(c)
+}

api/handler/user_handler.go

@@ -173,29 +173,8 @@ func (h *UserHandler) Login(c *gin.Context) {
 		LoginIp:      c.ClientIP(),
 		LoginAddress: utils.Ip2Region(h.searcher, c.ClientIP()),
 	})
-	var chatConfig types.ChatConfig
-	err = utils.JsonDecode(user.ChatConfig, &chatConfig)
-	if err != nil {
-		resp.ERROR(c, err.Error())
-		return
-	}
 
-	resp.SUCCESS(c, gin.H{
+	resp.SUCCESS(c, sessionId)
-		"session_id":     sessionId,
-		"id":             user.Id,
-		"nickname":       user.Nickname,
-		"avatar":         user.Avatar,
-		"username":       user.Username,
-		"tokens":         user.Tokens,
-		"calls":          user.Calls,
-		"expired_time":   user.ExpiredTime,
-		"api_key":        chatConfig.ApiKey,
-		"model":          chatConfig.Model,
-		"temperature":    chatConfig.Temperature,
-		"max_tokens":     chatConfig.MaxTokens,
-		"enable_context": chatConfig.EnableContext,
-		"enable_history": chatConfig.EnableHistory,
-	})
 }
 
 // Logout 注 销

api/main.go

@@ -178,6 +178,7 @@ func main() {
 		}),
 		fx.Invoke(func(s *core.AppServer, h *handler.MidJourneyHandler) {
 			s.Engine.POST("/api/mj/notify", h.Notify)
+			s.Engine.POST("/api/mj/upscale", h.Upscale)
 		}),
 
 		// 管理后台控制器

api/service/function/mid_journey.go

@@ -29,7 +29,7 @@ func (f FuncMidJourney) Invoke(params map[string]interface{}) (string, error) {
 		return "", errors.New("无效的 API Token")
 	}
 
-	//logger.Infof("MJ 绘画参数：%+v", params)
+	logger.Infof("MJ 绘画参数：%+v", params)
 	prompt := utils.InterfaceToString(params["prompt"])
 	if !utils.IsEmptyValue(params["ar"]) {
 		prompt = fmt.Sprintf("%s --ar %s", prompt, params["ar"])
@@ -60,6 +60,31 @@ func (f FuncMidJourney) Invoke(params map[string]interface{}) (string, error) {
 	return prompt, nil
 }
 
+type MjUpscaleReq struct {
+	Index       int32  `json:"index"`
+	MessageId   string `json:"message_id"`
+	MessageHash string `json:"message_hash"`
+}
+
+func (f FuncMidJourney) Upscale(upReq MjUpscaleReq) error {
+	url := fmt.Sprintf("%s/api/mj/upscale", f.config.ApiURL)
+	var res types.BizVo
+	r, err := f.client.R().
+		SetHeader("Authorization", f.config.Token).
+		SetHeader("Content-Type", "application/json").
+		SetBody(upReq).
+		SetSuccessResult(&res).Post(url)
+	if err != nil || r.IsErrorState() {
+		return fmt.Errorf("%v%v", r.String(), err)
+	}
+
+	if res.Code != types.Success {
+		return errors.New(res.Message)
+	}
+	
+	return nil
+}
+
 func (f FuncMidJourney) Name() string {
 	return f.name
 }

api/store/model/mj_job.go

@@ -0,0 +1,19 @@
+package model
+
+import "time"
+
+type MidJourneyJob struct {
+	Id        uint `gorm:"primarykey;column:id"`
+	UserId    uint
+	ChatId    string
+	MessageId string
+	Hash      string
+	Content   string
+	Prompt    string
+	Image     string
+	CreatedAt time.Time
+}
+
+func (MidJourneyJob) TableName() string {
+	return "chatgpt_mj_jobs"
+}

database/update-v3.0.7.sql

@@ -0,0 +1,34 @@
+CREATE TABLE `chatgpt_mj_jobs` (
+                                   `id` int NOT NULL,
+                                   `user_id` int NOT NULL COMMENT '用户 ID',
+                                   `chat_id` char(40) NOT NULL COMMENT '聊天会话 ID',
+                                   `message_id` char(40) NOT NULL COMMENT '消息 ID',
+                                   `hash` char(40) NOT NULL COMMENT '图片哈希',
+                                   `content` varchar(2000) NOT NULL COMMENT '消息内容',
+                                   `prompt` varchar(2000) NOT NULL COMMENT '会话提示词',
+                                   `image` text NOT NULL COMMENT '图片信息 json',
+                                   `created_at` datetime NOT NULL
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci COMMENT='MidJourney 任务表';
+
+--
+-- 转储表的索引
+--
+
+--
+-- 表的索引 `chatgpt_mj_jobs`
+--
+ALTER TABLE `chatgpt_mj_jobs`
+    ADD PRIMARY KEY (`id`),
+  ADD UNIQUE KEY `message_id` (`message_id`),
+  ADD UNIQUE KEY `hash` (`hash`);
+
+--
+-- 在导出的表使用AUTO_INCREMENT
+--
+
+--
+-- 使用表AUTO_INCREMENT `chatgpt_mj_jobs`
+--
+ALTER TABLE `chatgpt_mj_jobs`
+    MODIFY `id` int NOT NULL AUTO_INCREMENT;
+COMMIT;

web/src/components/ChatMidJourney.vue

@@ -1,5 +1,5 @@
 <template>
-  <div class="chat-line chat-line-mj">
+  <div class="chat-line chat-line-mj" v-loading="loading">
     <div class="chat-line-inner">
       <div class="chat-icon">
         <img :src="icon" alt="User"/>
@@ -30,7 +30,7 @@
           </div>
         </div>
 
-        <div class="opt" v-if="data.image?.hash !== ''">
+        <div class="opt" v-if="!data['reference_id'] &&data.image?.hash !== ''">
           <div class="opt-line">
             <ul>
               <li><a @click="upscale(1)">U1</a></li>
@@ -64,6 +64,8 @@
 import {ref, watch} from "vue";
 import {Clock} from "@element-plus/icons-vue";
 import {ElMessage} from "element-plus";
+import {httpPost} from "@/utils/http";
+import {getSessionId} from "@/store/session";
 
 const props = defineProps({
   content: Object,
@@ -72,10 +74,10 @@ const props = defineProps({
 });
 
 const data = ref(props.content)
-console.log(data.value)
 const tokens = ref(0)
 const cacheKey = "img_placeholder_height"
 const item = localStorage.getItem(cacheKey);
+const loading = ref(false)
 const height = ref(0)
 if (item) {
   height.value = parseInt(item)
@@ -88,9 +90,23 @@ if (data.value["image"]?.width > 0) {
 watch(() => props.content, (newVal) => {
   data.value = newVal;
 });
-
+const emits = defineEmits(['disable-input', 'disable-input']);
 const upscale = (index) => {
-  ElMessage.warning("当前版本暂未实现 Variation 功能！")
+  loading.value = true
+  emits('disable-input')
+  httpPost("/api/mj/upscale", {
+    index: index,
+    message_id: data.value?.["message_id"],
+    message_hash: data.value?.["image"]?.hash,
+    session_id: getSessionId(),
+    key: data.value?.["key"]
+  }).then(() => {
+    ElMessage.success("任务推送成功，请耐心等待任务执行...")
+    loading.value = false
+  }).catch(e => {
+    ElMessage.error("任务推送失败：" + e.message)
+    emits('disable-input')
+  })
 }
 
 const variation = (index) => {

web/src/store/session.js

@@ -4,26 +4,16 @@
  * storage handler
  */
 
-const SessionUserKey = 'LOGIN_USER';
+const SessionUserKey = 'SESSION_ID';
 
 export function getSessionId() {
-    const user = getLoginUser();
+    return sessionStorage.getItem(SessionUserKey)
-    return user ? user['session_id'] : '';
 }
 
 export function removeLoginUser() {
     sessionStorage.removeItem(SessionUserKey)
 }
 
-export function getLoginUser() {
+export function setSessionId(sessionId) {
-    const value = sessionStorage.getItem(SessionUserKey);
+    sessionStorage.setItem(SessionUserKey, sessionId)
-    if (value) {
-        return JSON.parse(value);
-    } else {
-        return null;
-    }
-}
-
-export function setLoginUser(user) {
-    sessionStorage.setItem(SessionUserKey, JSON.stringify(user))
 }

web/src/utils/http.js

@@ -1,5 +1,4 @@
 import axios from 'axios'
-import {getSessionId} from "@/store/session";
 
 axios.defaults.timeout = 10000
 axios.defaults.baseURL = process.env.VUE_APP_API_HOST
@@ -10,7 +9,7 @@ axios.defaults.headers.post['Content-Type'] = 'application/json'
 axios.interceptors.request.use(
     config => {
         // set token
-        config.headers['ChatGPT-TOKEN'] = getSessionId();
+        // config.headers['ChatGPT-TOKEN'] = getSessionId();
         return config
     }, error => {
         return Promise.reject(error)

web/src/views/ChatPlus.vue

@@ -169,6 +169,8 @@
                   <chat-mid-journey v-else-if="item.type==='mj'"
                                     :content="item.content"
                                     :icon="item.icon"
+                                    @disable-input="disableInput(true)"
+                                    @enable-input="enableInput"
                                     :created-at="dateFormat(item['created_at'])"/>
                 </div>
               </div><!-- end chat box -->
@@ -517,7 +519,7 @@ const connect = function (chat_id, role_id) {
   _socket.addEventListener('open', () => {
     chatData.value = []; // 初始化聊天数据
     previousText.value = '';
-    canSend.value = true;
+    enableInput()
     activelyClose.value = false;
 
     if (isNewChat) { // 加载打招呼信息
@@ -548,9 +550,7 @@ const connect = function (chat_id, role_id) {
             content: ""
           });
         } else if (data.type === "mj") {
-          canSend.value = false;
+          disableInput(true)
-          showReGenerate.value = false;
-          showStopGenerate.value = true;
           const content = data.content;
           const md = require('markdown-it')({breaks: true});
           content.content = md.render(content.content)
@@ -568,16 +568,14 @@ const connect = function (chat_id, role_id) {
           if (flag === false) {
             chatData.value.push({
               type: "mj",
-              id: content.key,
+              id: content["message_id"],
               icon: "/images/avatar/mid_journey.png",
               content: content
             });
           }
 
           if (content.status === "Finished") {
-            canSend.value = true;
+            enableInput()
-            showReGenerate.value = true;
-            showStopGenerate.value = false;
           }
         } else if (data.type === 'end') { // 消息接收完毕
           // 追加当前会话到会话列表
@@ -588,18 +586,12 @@ const connect = function (chat_id, role_id) {
             activeChat.value = newChatItem.value;
             newChatItem.value = null; // 只追加一次
           }
-          const reply = chatData.value[chatData.value.length - 1]
-          if (reply.content.indexOf("绘画提示词：") === -1) {
-            return
-          }
 
-          canSend.value = true;
+          enableInput()
-          showReGenerate.value = true;
-          showStopGenerate.value = false;
           lineBuffer.value = ''; // 清空缓冲
           
-
           // 获取 token
+          const reply = chatData.value[chatData.value.length - 1]
           httpGet(`/api/chat/tokens?text=${reply.orgContent}&model=${model.value}`).then(res => {
             reply['created_at'] = new Date().getTime();
             reply['tokens'] = res.data;
@@ -639,7 +631,7 @@ const connect = function (chat_id, role_id) {
       return;
     }
     // 停止发送消息
-    canSend.value = true;
+    disableInput(true)
     socket.value = null;
     loading.value = true;
     checkSession().then(() => {
@@ -657,6 +649,18 @@ const connect = function (chat_id, role_id) {
   socket.value = _socket;
 }
 
+const disableInput = (force) => {
+  canSend.value = false;
+  showReGenerate.value = false;
+  showStopGenerate.value = !force;
+}
+
+const enableInput = () => {
+  canSend.value = true;
+  showReGenerate.value = previousText.value !== "";
+  showStopGenerate.value = false;
+}
+
 // 登录输入框输入事件处理
 const inputKeyDown = function (e) {
   if (e.keyCode === 13) {
@@ -700,9 +704,7 @@ const sendMessage = function () {
   })
 
   showHello.value = false
-  canSend.value = false;
+  disableInput(false)
-  showStopGenerate.value = true;
-  showReGenerate.value = false;
   socket.value.send(prompt.value);
   previousText.value = prompt.value;
   prompt.value = '';
@@ -794,18 +796,13 @@ const loadChatHistory = function (chatId) {
 const stopGenerate = function () {
   showStopGenerate.value = false;
   httpGet("/api/chat/stop?session_id=" + getSessionId()).then(() => {
-    canSend.value = true;
+    enableInput()
-    if (previousText.value !== '') {
-      showReGenerate.value = true;
-    }
   })
 }
 
 // 重新生成
 const reGenerate = function () {
-  canSend.value = false;
+  disableInput(false)
-  showStopGenerate.value = true;
-  showReGenerate.value = false;
   const text = '重新生成上述问题的答案：' + previousText.value;
   // 追加消息
   chatData.value.push({

web/src/views/Login.vue

@@ -52,7 +52,7 @@ import {onMounted, ref} from "vue";
 import {Lock, UserFilled} from "@element-plus/icons-vue";
 import {httpPost} from "@/utils/http";
 import {ElMessage} from "element-plus";
-import {setLoginUser} from "@/store/session";
+import {setSessionId} from "@/store/session";
 import {useRouter} from "vue-router";
 import FooterBar from "@/components/FooterBar.vue";
 import {isMobile} from "@/utils/libs";
@@ -79,7 +79,7 @@ const login = function () {
   }
 
   httpPost('/api/user/login', {username: username.value.trim(), password: password.value.trim()}).then((res) => {
-    setLoginUser(res.data)
+    setSessionId(res.data)
     if (isMobile()) {
       router.push('/mobile')
     } else {

web/src/views/admin/Login.vue

@@ -44,7 +44,7 @@ import {onMounted, ref} from "vue";
 import {Lock, UserFilled} from "@element-plus/icons-vue";
 import {httpPost} from "@/utils/http";
 import {ElMessage} from "element-plus";
-import {setLoginUser} from "@/store/session";
+import {setSession} from "@/store/session";
 import {useRouter} from "vue-router";
 import FooterBar from "@/components/FooterBar.vue";
 
@@ -70,7 +70,6 @@ const login = function () {
   }
 
   httpPost('/api/admin/login', {username: username.value.trim(), password: password.value.trim()}).then((res) => {
-    setLoginUser(res.data)
     router.push("/admin")
   }).catch((e) => {
     ElMessage.error('登录失败，' + e.message)