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: midjourney image upscale function is ready

Browse Source
This commit is contained in:
RockYang 2023-08-15 15:28:40 +08:00
parent 6c76086916
commit 827acdd3f9
16 changed files with 227 additions and 109 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api/Makefile
View File

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

2
api/core/app_server.go
View File

@ -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")
//设置缓存时间

18
api/core/types/client.go
View File

@ -6,18 +6,14 @@ import (
"sync"
)
var ErrConClosed = errors.New("connection closed")
type Client interface {
Close()
}
var ErrConClosed = errors.New("connection Closed")
// 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
}

91
api/handler/mj_handler.go
View File

@ -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 {
h := MidJourneyHandler{leveldb: leveldb, db: db}
func NewMidJourneyHandler(app *core.AppServer, leveldb *store.LevelDB, db *gorm.DB, functions map[string]function.Function) *MidJourneyHandler {
h := MidJourneyHandler{leveldb: leveldb, db: db, mjFunc: functions[types.FuncMidJourney].(function.FuncMidJourney)}
h.App = app
return &h
}
@ -50,8 +53,8 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
}
var data struct {
Type string `json:"type"`
MessageId string `json:"message_id"`
ReferenceId string `json:"reference_id"`
Image Image `json:"image"`
Content string `json:"content"`
Prompt string `json:"prompt"`
@ -63,21 +66,30 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
return
}
key := utils.Sha256(data.Prompt)
data.Key = key
logger.Infof("收到 MidJourney 回调请求：%+v", data)
// 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
_ = h.leveldb.Delete(types.TaskStorePrefix + key)
// save the job
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)
}

23
api/handler/user_handler.go
View File

@ -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{
"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,
})
resp.SUCCESS(c, sessionId)
}
// Logout 注 销

1
api/main.go
View File

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

27
api/service/function/mid_journey.go
View File

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

19
api/store/model/mj_job.go Normal file
View File

@ -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"
}

0
database/update.sql → database/update-v3.0.6.sql
View File

34
database/update-v3.0.7.sql Normal file
View File

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

26
web/src/components/ChatMidJourney.vue
View File

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

18
web/src/store/session.js
View File

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

3
web/src/utils/http.js
View File

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

51
web/src/views/ChatPlus.vue
View File

@ -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;
showReGenerate.value = false;
showStopGenerate.value = true;
disableInput(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;
showReGenerate.value = true;
showStopGenerate.value = false;
enableInput()
}
} 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;
showReGenerate.value = true;
showStopGenerate.value = false;
enableInput()
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;
showStopGenerate.value = true;
showReGenerate.value = false;
disableInput(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;
if (previousText.value !== '') {
showReGenerate.value = true;
}
enableInput()
})
}
// 重新生成
const reGenerate = function () {
canSend.value = false;
showStopGenerate.value = true;
showReGenerate.value = false;
disableInput(false)
const text = '重新生成上述问题的答案：' + previousText.value;
// 追加消息
chatData.value.push({

4
web/src/views/Login.vue
View File

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

3
web/src/views/admin/Login.vue
View File

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