feat: mj advance drawing page function is ready, use better task scheduling argorithm

api/core/types/function.go

@@ -83,31 +83,7 @@ var InnerFunctions = []Function{
 			Properties: map[string]Property{
 				"prompt": {
 					Type:        "string",
-					Description: "绘画内容描述，提示词，如果该参数中有中文的话，则需要翻译成英文",
+					Description: "提示词，如果该参数中有中文的话，则需要翻译成英文。提示词中的参数作为提示的一部分，不要删除",
-				},
-				"--ar": {
-					Type:        "string",
-					Description: "图片长宽比，默认值 16:9",
-				},
-				"--niji": {
-					Type:        "string",
-					Description: "动漫模型版本，默认值空",
-				},
-				"--s": {
-					Type:        "string",
-					Description: "风格，stylize",
-				},
-				"--seed": {
-					Type:        "string",
-					Description: "随机种子",
-				},
-				"--no": {
-					Type:        "string",
-					Description: "负面提示词，指定不要什么元素或者风格,如果该参数中有中文的话，则需要翻译成英文",
-				},
-				"--v": {
-					Type:        "string",
-					Description: "模型版本，默认值: 5.2",
 				},
 			},
 			Required: []string{},

api/handler/mj_handler.go

@@ -14,7 +14,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/go-redis/redis/v8"
 	"gorm.io/gorm"
-	"net/http"
+	"strings"
 	"sync"
 	"time"
 )
@@ -24,7 +24,6 @@ type TaskStatus string
 const (
 	Stopped  = TaskStatus("Stopped")
 	Finished = TaskStatus("Finished")
-	Running  = TaskStatus("Running")
 )
 
 type Image struct {
@@ -117,18 +116,25 @@ func (h *MidJourneyHandler) notifyHandler(c *gin.Context, data notifyData) (erro
 		return nil, false
 	}
 
+	var job model.MidJourneyJob
+	res := h.db.Where("message_id = ?", data.MessageId).First(&job)
+	if res.Error == nil && data.Status == Finished {
+		logger.Warn("重复消息：", data.MessageId)
+		return nil, false
+	}
+
 	if task.Src == service.TaskSrcImg { // 绘画任务
-		logger.Error(err)
 		var job model.MidJourneyJob
-		res := h.db.First(&job, task.Id)
+		res := h.db.Where("id = ?", task.Id).First(&job)
 		if res.Error != nil {
-			logger.Warn("非法任务：", err)
+			logger.Warn("非法任务：", res.Error)
 			return nil, false
 		}
 		job.MessageId = data.MessageId
 		job.ReferenceId = data.ReferenceId
 		job.Progress = data.Progress
 		job.Prompt = data.Prompt
+		job.Hash = data.Image.Hash
 
 		// 任务完成，将最终的图片下载下来
 		if data.Progress == 100 {
@@ -144,18 +150,11 @@ func (h *MidJourneyHandler) notifyHandler(c *gin.Context, data notifyData) (erro
 		}
 		res = h.db.Updates(&job)
 		if res.Error != nil {
-			logger.Error("error with update job: ", err.Error())
+			logger.Error("error with update job: ", res.Error)
 			return res.Error, false
 		}
 
 	} else if task.Src == service.TaskSrcChat { // 聊天任务
-		var job model.MidJourneyJob
-		res := h.db.Where("message_id = ?", data.MessageId).First(&job)
-		if res.Error == nil {
-			logger.Warn("重复消息：", data.MessageId)
-			return nil, false
-		}
-
 		wsClient := h.App.MjTaskClients.Get(task.Id)
 		if data.Status == Finished {
 			if wsClient != nil && data.ReferenceId != "" {
@@ -233,15 +232,68 @@ func (h *MidJourneyHandler) notifyHandler(c *gin.Context, data notifyData) (erro
 	return nil, true
 }
 
-// Proxy 通过代理访问 discord f服务器图片
+// Image 创建一个绘画任务
-func (h *MidJourneyHandler) Proxy(c *gin.Context) {
+func (h *MidJourneyHandler) Image(c *gin.Context) {
-	imgURL := c.Query("url")
+	var data struct {
-	imageData, err := utils.DownloadImage(imgURL, h.App.Config.ProxyURL)
+		Prompt  string  `json:"prompt"`
-	if err != nil {
+		Rate    string  `json:"rate"`
-		c.String(http.StatusOK, err.Error())
+		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)
 		return
 	}
-	c.String(http.StatusOK, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
+	var prompt = data.Prompt
+	if data.Rate != "" && !strings.Contains(prompt, "--ar") {
+		prompt += " --ar " + data.Rate
+	}
+	if data.Seed > 0 && !strings.Contains(prompt, "--seed") {
+		prompt += fmt.Sprintf(" --seed %d", data.Seed)
+	}
+	if data.Stylize > 0 && !strings.Contains(prompt, "--s") && !strings.Contains(prompt, "--stylize") {
+		prompt += fmt.Sprintf(" --s %d", data.Stylize)
+	}
+	if data.Chaos > 0 && !strings.Contains(prompt, "--c") && !strings.Contains(prompt, "--chaos") {
+		prompt += fmt.Sprintf(" --c %d", data.Chaos)
+	}
+	if data.Img != "" {
+		prompt = fmt.Sprintf("%s %s", data.Img, prompt)
+		if data.Weight > 0 {
+			prompt += fmt.Sprintf(" --iw %f", data.Weight)
+		}
+	}
+	if data.Model != "" && !strings.Contains(prompt, "--v") && !strings.Contains(prompt, "--niji") {
+		prompt += data.Model
+	}
+
+	idValue, _ := c.Get(types.LoginUserID)
+	userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
+	job := model.MidJourneyJob{
+		Type:      service.Image.String(),
+		UserId:    userId,
+		Progress:  0,
+		Prompt:    prompt,
+		CreatedAt: time.Now(),
+	}
+	if res := h.db.Create(&job); res.Error != nil {
+		resp.ERROR(c, "添加任务失败："+res.Error.Error())
+		return
+	}
+
+	h.mjService.PushTask(service.MjTask{
+		Id:     fmt.Sprintf("%d", job.Id),
+		Src:    service.TaskSrcImg,
+		Type:   service.Image,
+		Prompt: prompt,
+		UserId: userId,
+	})
+	resp.SUCCESS(c)
 }
 
 type reqVo struct {
@@ -264,13 +316,32 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
 		return
 	}
 
-	userId, _ := c.Get(types.LoginUserID)
+	idValue, _ := c.Get(types.LoginUserID)
+	jobId := data.SessionId
+	userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
+	src := service.TaskSrc(data.Src)
+	if src == service.TaskSrcImg {
+		job := model.MidJourneyJob{
+			Type:      service.Upscale.String(),
+			UserId:    userId,
+			Hash:      data.MessageHash,
+			Progress:  0,
+			Prompt:    data.Prompt,
+			CreatedAt: time.Now(),
+		}
+		if res := h.db.Create(&job); res.Error == nil {
+			jobId = fmt.Sprintf("%d", job.Id)
+		} else {
+			resp.ERROR(c, "添加任务失败："+res.Error.Error())
+			return
+		}
+	}
 	h.mjService.PushTask(service.MjTask{
-		Id:          data.SessionId,
+		Id:          jobId,
-		Src:         service.TaskSrc(data.Src),
+		Src:         src,
 		Type:        service.Upscale,
 		Prompt:      data.Prompt,
-		UserId:      utils.IntValue(utils.InterfaceToString(userId), 0),
+		UserId:      userId,
 		RoleId:      data.RoleId,
 		Icon:        data.Icon,
 		ChatId:      data.ChatId,
@@ -298,13 +369,33 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
 		return
 	}
 
-	userId, _ := c.Get(types.LoginUserID)
+	idValue, _ := c.Get(types.LoginUserID)
+	jobId := data.SessionId
+	userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
+	src := service.TaskSrc(data.Src)
+	if src == service.TaskSrcImg {
+		job := model.MidJourneyJob{
+			Type:      service.Variation.String(),
+			UserId:    userId,
+			ImgURL:    "",
+			Hash:      data.MessageHash,
+			Progress:  0,
+			Prompt:    data.Prompt,
+			CreatedAt: time.Now(),
+		}
+		if res := h.db.Create(&job); res.Error == nil {
+			jobId = fmt.Sprintf("%d", job.Id)
+		} else {
+			resp.ERROR(c, "添加任务失败："+res.Error.Error())
+			return
+		}
+	}
 	h.mjService.PushTask(service.MjTask{
-		Id:          data.SessionId,
+		Id:          jobId,
-		Src:         service.TaskSrc(data.Src),
+		Src:         src,
 		Type:        service.Variation,
 		Prompt:      data.Prompt,
-		UserId:      utils.IntValue(utils.InterfaceToString(userId), 0),
+		UserId:      userId,
 		RoleId:      data.RoleId,
 		Icon:        data.Icon,
 		ChatId:      data.ChatId,
@@ -332,9 +423,9 @@ func (h *MidJourneyHandler) JobList(c *gin.Context) {
 	var res *gorm.DB
 	userId, _ := c.Get(types.LoginUserID)
 	if status == 1 {
-		res = h.db.Where("user_id = ? AND progress = 100", userId).Find(&items)
+		res = h.db.Where("user_id = ? AND progress = 100", userId).Order("id DESC").Find(&items)
 	} else {
-		res = h.db.Where("user_id = ? AND progress < 100", userId).Find(&items)
+		res = h.db.Where("user_id = ? AND progress < 100", userId).Order("id ASC").Find(&items)
 	}
 	if res.Error != nil {
 		resp.ERROR(c, types.NoData)
@@ -348,6 +439,12 @@ func (h *MidJourneyHandler) JobList(c *gin.Context) {
 		if err != nil {
 			continue
 		}
+		if item.Progress < 100 && item.ImgURL != "" { // 正在运行中任务使用代理访问图片
+			image, err := utils.DownloadImage(item.ImgURL, h.App.Config.ProxyURL)
+			if err == nil {
+				job.ImgURL = "data:image/png;base64," + base64.StdEncoding.EncodeToString(image)
+			}
+		}
 		jobs = append(jobs, job)
 	}
 	resp.SUCCESS(c, jobs)

api/main.go

@@ -191,10 +191,10 @@ func main() {
 		fx.Invoke(func(s *core.AppServer, h *handler.MidJourneyHandler) {
 			group := s.Engine.Group("/api/mj/")
 			group.POST("notify", h.Notify)
+			group.POST("image", h.Image)
 			group.POST("upscale", h.Upscale)
 			group.POST("variation", h.Variation)
 			group.GET("jobs", h.JobList)
-			group.GET("proxy", h.Proxy)
 		}),
 
 		// 管理后台控制器

api/service/function/func_mj.go

@@ -3,7 +3,6 @@ package function
 import (
 	"chatplus/service"
 	"chatplus/utils"
-	"fmt"
 )
 
 // AI 绘画函数
@@ -22,29 +21,6 @@ func NewMidJourneyFunc(mjService *service.MjService) FuncMidJourney {
 func (f FuncMidJourney) Invoke(params map[string]interface{}) (string, error) {
 	logger.Infof("MJ 绘画参数：%+v", params)
 	prompt := utils.InterfaceToString(params["prompt"])
-	if !utils.IsEmptyValue(params["ar"]) {
-		prompt = fmt.Sprintf("%s --ar %s", prompt, params["ar"])
-		delete(params, "ar")
-	}
-	if !utils.IsEmptyValue(params["s"]) {
-		prompt = fmt.Sprintf("%s --s %s", prompt, params["s"])
-		delete(params, "s")
-	}
-	if !utils.IsEmptyValue(params["seed"]) {
-		prompt = fmt.Sprintf("%s --seed %s", prompt, params["seed"])
-		delete(params, "seed")
-	}
-	if !utils.IsEmptyValue(params["no"]) {
-		prompt = fmt.Sprintf("%s --no %s", prompt, params["no"])
-		delete(params, "no")
-	}
-	if !utils.IsEmptyValue(params["niji"]) {
-		prompt = fmt.Sprintf("%s --niji %s", prompt, params["niji"])
-		delete(params, "niji")
-	} else {
-		prompt = prompt + " --v 5.2"
-	}
-
 	f.service.PushTask(service.MjTask{
 		Id:     utils.InterfaceToString(params["session_id"]),
 		Src:    service.TaskSrcChat,

api/service/mj_service.go

@@ -4,12 +4,14 @@ import (
 	"chatplus/core/types"
 	logger2 "chatplus/logger"
 	"chatplus/store"
+	"chatplus/store/model"
 	"chatplus/utils"
 	"context"
 	"errors"
 	"fmt"
 	"github.com/go-redis/redis/v8"
 	"github.com/imroc/req/v3"
+	"gorm.io/gorm"
 	"time"
 )
 
@@ -58,12 +60,14 @@ type MjService struct {
 	client    *req.Client
 	taskQueue *store.RedisQueue
 	redis     *redis.Client
+	db        *gorm.DB
 }
 
-func NewMjService(appConfig *types.AppConfig, client *redis.Client) *MjService {
+func NewMjService(appConfig *types.AppConfig, client *redis.Client, db *gorm.DB) *MjService {
 	return &MjService{
 		config:    appConfig.ExtConfig,
 		redis:     client,
+		db:        db,
 		taskQueue: store.NewRedisQueue("midjourney_task_queue", client),
 		client:    req.C().SetTimeout(30 * time.Second)}
 }
@@ -104,9 +108,11 @@ func (s *MjService) Run() {
 		}
 		if err != nil {
 			logger.Error("绘画任务执行失败：", err)
-			//if task.RetryCount > 5 {
+			if task.RetryCount > 5 {
-			//	continue
+				// 取消并删除任务
-			//}
+				s.db.Where("id = ?", task.Id).Delete(&model.MidJourneyJob{})
+				continue
+			}
 			task.RetryCount += 1
 			s.taskQueue.RPush(task)
 			// TODO: 执行失败通知聊天客户端

database/update-v3.1.3.sql

@@ -7,4 +7,6 @@ ALTER TABLE `chatgpt_mj_jobs` ADD `hash` VARCHAR(100) NULL DEFAULT NULL COMMENT
 ALTER TABLE `chatgpt_mj_jobs` ADD `img_url` VARCHAR(255) NULL DEFAULT NULL COMMENT '图片URL' AFTER `prompt`;
 
 -- 2023-09-15
-ALTER TABLE `chatgpt_mj_jobs` ADD `type` VARCHAR(20) NULL DEFAULT 'image' COMMENT '任务类别' AFTER `user_id`;
+ALTER TABLE `chatgpt_mj_jobs` ADD `type` VARCHAR(20) NULL DEFAULT 'image' COMMENT '任务类别' AFTER `user_id`;
+ALTER TABLE `chatgpt_mj_jobs` DROP INDEX `message_id`;
+ALTER TABLE `chatgpt_mj_jobs` ADD INDEX(`message_id`);

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

@@ -234,7 +234,6 @@
               display flex
               justify-content center
               align-items center
-              background-color: rgba(0, 0, 0, 0.5)
 
               span {
                 font-size 20px
@@ -249,6 +248,8 @@
       .finish-job-list {
         .job-item {
           margin-bottom 20px
+          width 100%
+          height 100%
 
           .opt {
             .opt-line {
@@ -289,12 +290,24 @@
         height 100%
         max-height 240px
 
+        img {
+          height 240px
+        }
+
+        .el-image-viewer__wrapper {
+          img {
+            width auto
+            height auto
+          }
+        }
+
         .image-slot {
           display flex
           flex-flow column
           justify-content center
           align-items center
           height 100%
+          min-height 200px
           color #ffffff
 
           .iconfont {

web/src/views/ImageMj.vue

@@ -60,7 +60,7 @@
               <el-form-item label="创意度">
                 <template #default>
                   <div class="form-item-inner">
-                    <el-input v-model="params.chaos" size="small"/>
+                    <el-input v-model.number="params.chaos" size="small"/>
                     <el-tooltip
                         effect="light"
                         content="参数用法：--chaos 或--c，取值范围: 0-100 <br/> 取值越高结果越发散，反之则稳定收敛<br /> 默认值0最为精准稳定"
@@ -80,7 +80,7 @@
               <el-form-item label="风格化">
                 <template #default>
                   <div class="form-item-inner">
-                    <el-input v-model="params.stylize" size="small"/>
+                    <el-input v-model.number="params.stylize" size="small"/>
                     <el-tooltip
                         effect="light"
                         content="风格化：--stylize 或 --s，范围 1-1000，默认值100 <br/>高取值会产生非常艺术化但与提示关联性较低的图像"
@@ -100,7 +100,7 @@
               <el-form-item label="随机种子">
                 <template #default>
                   <div class="form-item-inner">
-                    <el-input v-model="params.seed" size="small"/>
+                    <el-input v-model.number="params.seed" size="small"/>
                     <el-tooltip
                         effect="light"
                         content="随机种子：--seed，默认值0表示随机产生 <br/>使用相同的种子参数和描述将产生相似的图像"
@@ -179,7 +179,7 @@
               <el-form-item label="图像权重">
                 <template #default>
                   <div class="form-item-inner">
-                    <el-slider v-model="params.weight" :max="1" :step="0.01"
+                    <el-slider v-model.number="params.weight" :max="1" :step="0.01"
                                style="width: 180px;--el-slider-main-bg-color:#47fff1"/>
                     <el-tooltip
                         effect="light"
@@ -201,6 +201,7 @@
                   v-model="params.prompt"
                   :autosize="{ minRows: 4, maxRows: 6 }"
                   type="textarea"
+                  ref="promptRef"
                   placeholder="这里输入你的咒语，例如：A chinese girl walking in the middle of a cobblestone street"
               />
             </div>
@@ -222,7 +223,7 @@
                       placement="top-start"
                       :title="getTaskType(scope.item.type)"
                       :width="240"
-                      trigger="hover"
+                      trigger="click"
                   >
                     <template #reference>
                       <div v-if="scope.item.progress > 0" class="job-item-inner">
@@ -321,19 +322,19 @@
                   <div class="opt">
                     <div class="opt-line">
                       <ul>
-                        <li><a @click="upscale(1)">U1</a></li>
+                        <li><a @click="upscale(1,scope.item)">U1</a></li>
-                        <li><a @click="upscale(2)">U2</a></li>
+                        <li><a @click="upscale(2,scope.item)">U2</a></li>
-                        <li><a @click="upscale(3)">U3</a></li>
+                        <li><a @click="upscale(3,scope.item)">U3</a></li>
-                        <li><a @click="upscale(4)">U4</a></li>
+                        <li><a @click="upscale(4,scope.item)">U4</a></li>
                       </ul>
                     </div>
 
                     <div class="opt-line">
                       <ul>
-                        <li><a @click="variation(1)">V1</a></li>
+                        <li><a @click="variation(1,scope.item)">V1</a></li>
-                        <li><a @click="variation(2)">V2</a></li>
+                        <li><a @click="variation(2,scope.item)">V2</a></li>
-                        <li><a @click="variation(3)">V3</a></li>
+                        <li><a @click="variation(3,scope.item)">V3</a></li>
-                        <li><a @click="variation(4)">V4</a></li>
+                        <li><a @click="variation(4,scope.item)">V4</a></li>
                       </ul>
                     </div>
                   </div>
@@ -358,20 +359,28 @@ import {httpGet, httpPost} from "@/utils/http";
 import {ElMessage} from "element-plus";
 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";
 
-const listBoxHeight = window.innerHeight - 40
+const listBoxHeight = ref(window.innerHeight - 40)
-const mjBoxHeight = window.innerHeight - 150
+const mjBoxHeight = ref(window.innerHeight - 150)
+
+window.onresize = () => {
+  listBoxHeight.value = window.innerHeight - 40
+  mjBoxHeight.value = window.innerHeight - 150
+}
 const rates = [
   {css: "horizontal", value: "16:9", text: "横图"},
   {css: "square", value: "1:1", text: "方图"},
   {css: "vertical", value: "9:16", text: "竖图"},
 ]
 const models = [
-  {text: "标准模型", value: "--v 5.2", img: "/images/mj-normal.png"},
+  {text: "标准模型", value: " --v 5.2", img: "/images/mj-normal.png"},
-  {text: "动漫模型", value: "--niji 5", img: "/images/mj-niji.png"},
+  {text: "动漫模型", value: " --niji 5", img: "/images/mj-niji.png"},
 ]
 const params = ref({
-  rate: rates[0].value,
+  rate: rates[1].value,
   model: models[0].value,
   chaos: 0,
   stylize: 100,
@@ -385,14 +394,18 @@ const params = ref({
 const runningJobs = ref([])
 const finishedJobs = ref([])
 const previewImgList = ref([])
-
+const router = useRouter()
 onMounted(() => {
-  fetchFinishedJobs()
+  checkSession().then(() => {
-  fetchRunningJobs()
+    fetchFinishedJobs()
+    fetchRunningJobs()
+  }).catch(() => {
+    router.push('/login')
+  });
 
   const clipboard = new Clipboard('.copy-prompt');
   clipboard.on('success', () => {
-    ElMessage.success('复制成功！');
+    ElMessage.success({message: "复制成功！", duration: 500});
   })
 
   clipboard.on('error', () => {
@@ -403,9 +416,13 @@ 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)
   })
@@ -414,6 +431,9 @@ const fetchFinishedJobs = () => {
 const fetchRunningJobs = () => {
   httpGet("/api/mj/jobs?status=0").then(res => {
     runningJobs.value = res.data
+    setTimeout(() => {
+      fetchRunningJobs()
+    }, 1000)
   }).catch(e => {
     ElMessage.error("获取任务失败：" + e.message)
   })
@@ -463,9 +483,44 @@ const getTaskType = (type) => {
 }
 
 // 创建绘图任务
+const promptRef = ref(null)
 const generate = () => {
-
+  if (params.value.prompt === '') {
+    promptRef.value.focus()
+    return ElMessage.error("请输入绘画提示词！")
+  }
+  httpPost("/api/mj/image", params.value).then(() => {
+    ElMessage.success("绘画任务推送成功，请耐心等待任务执行...")
+  }).catch(e => {
+    ElMessage.error("任务推送失败：" + e.message)
+  })
 }
+
+// 图片放大任务
+const upscale = (index, item) => {
+  send('/api/mj/upscale', index, item)
+}
+
+// 图片变换任务
+const variation = (index, item) => {
+  send('/api/mj/variation', index, item)
+}
+
+const send = (url, index, item) => {
+  httpPost(url, {
+    index: index,
+    src: "img",
+    message_id: item.message_id,
+    message_hash: item.hash,
+    session_id: getSessionId(),
+    prompt: item.prompt,
+  }).then(() => {
+    ElMessage.success("任务推送成功，请耐心等待任务执行...")
+  }).catch(e => {
+    ElMessage.error("任务推送失败：" + e.message)
+  })
+}
+
 </script>
 
 <style lang="stylus">