xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-09-17 08:46:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
RockYang 2023-09-17 18:03:45 +08:00
parent 26944f9e39
commit 5034a20345
8 changed files with 234 additions and 109 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
api/core/types/function.go
View File

@ -83,31 +83,7 @@ var InnerFunctions = []Function{
Properties: map[string]Property{
"prompt": {
Type: "string",
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",
Description: "提示词，如果该参数中有中文的话，则需要翻译成英文。提示词中的参数作为提示的一部分，不要删除",
},
},
Required: []string{},

157
api/handler/mj_handler.go
View File

@ -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服务器图片
func (h *MidJourneyHandler) Proxy(c *gin.Context) {
imgURL := c.Query("url")
imageData, err := utils.DownloadImage(imgURL, h.App.Config.ProxyURL)
if err != nil {
c.String(http.StatusOK, err.Error())
// 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"`
}
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,
Src: service.TaskSrc(data.Src),
Id: jobId,
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,
Src: service.TaskSrc(data.Src),
Id: jobId,
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)

2
api/main.go
View File

@ -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)
}),
// 管理后台控制器

24
api/service/function/func_mj.go
View File

@ -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,

14
api/service/mj_service.go
View File

@ -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 {
// continue
//}
if task.RetryCount > 5 {
// 取消并删除任务
s.db.Where("id = ?", task.Id).Delete(&model.MidJourneyJob{})
continue
}
task.RetryCount += 1
s.taskQueue.RPush(task)
// TODO: 执行失败通知聊天客户端

4
database/update-v3.1.3.sql
View File

@ -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`);

15
web/src/assets/css/image-mj.styl
View File

@ -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 {

101
web/src/views/ImageMj.vue
View File

@ -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(2)">U2</a></li>
<li><a @click="upscale(3)">U3</a></li>
<li><a @click="upscale(4)">U4</a></li>
<li><a @click="upscale(1,scope.item)">U1</a></li>
<li><a @click="upscale(2,scope.item)">U2</a></li>
<li><a @click="upscale(3,scope.item)">U3</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(2)">V2</a></li>
<li><a @click="variation(3)">V3</a></li>
<li><a @click="variation(4)">V4</a></li>
<li><a @click="variation(1,scope.item)">V1</a></li>
<li><a @click="variation(2,scope.item)">V2</a></li>
<li><a @click="variation(3,scope.item)">V3</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 mjBoxHeight = window.innerHeight - 150
const listBoxHeight = ref(window.innerHeight - 40)
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: "--niji 5", img: "/images/mj-niji.png"},
{text: "标准模型", value: " --v 5.2", img: "/images/mj-normal.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()
fetchRunningJobs()
checkSession().then(() => {
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">