mirror of
https://github.com/yangjian102621/geekai.git
synced 2025-11-11 03:33:48 +08:00
refactor midjourney service, use api key in database
This commit is contained in:
RockYang
@@ -7,15 +7,28 @@ package mj
|
||||
// * @Author yangjian102621@163.com
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
|
||||
import "geekai/core/types"
|
||||
import (
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"geekai/core/types"
|
||||
logger2 "geekai/logger"
|
||||
"geekai/service"
|
||||
"geekai/store/model"
|
||||
"geekai/utils"
|
||||
"github.com/imroc/req/v3"
|
||||
"gorm.io/gorm"
|
||||
"io"
|
||||
"time"
|
||||
|
||||
type Client interface {
|
||||
Imagine(task types.MjTask) (ImageRes, error)
|
||||
Blend(task types.MjTask) (ImageRes, error)
|
||||
SwapFace(task types.MjTask) (ImageRes, error)
|
||||
Upscale(task types.MjTask) (ImageRes, error)
|
||||
Variation(task types.MjTask) (ImageRes, error)
|
||||
QueryTask(taskId string) (QueryRes, error)
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// Client MidJourney client
|
||||
type Client struct {
|
||||
client *req.Client
|
||||
licenseService *service.LicenseService
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
type ImageReq struct {
|
||||
@@ -33,7 +46,8 @@ type ImageRes struct {
|
||||
Description string `json:"description"`
|
||||
Properties struct {
|
||||
} `json:"properties"`
|
||||
Result string `json:"result"`
|
||||
Result string `json:"result"`
|
||||
Channel string `json:"channel,omitempty"`
|
||||
}
|
||||
|
||||
type ErrRes struct {
|
||||
@@ -66,3 +80,184 @@ type QueryRes struct {
|
||||
Status string `json:"status"`
|
||||
SubmitTime int `json:"submitTime"`
|
||||
}
|
||||
|
||||
var logger = logger2.GetLogger()
|
||||
|
||||
func NewClient(licenseService *service.LicenseService, db *gorm.DB) *Client {
|
||||
return &Client{
|
||||
client: req.C().SetTimeout(time.Minute).SetUserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36"),
|
||||
licenseService: licenseService,
|
||||
db: db,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) Imagine(task types.MjTask) (ImageRes, error) {
|
||||
apiPath := fmt.Sprintf("mj-%s/mj/submit/imagine", task.Mode)
|
||||
prompt := fmt.Sprintf("%s %s", task.Prompt, task.Params)
|
||||
if task.NegPrompt != "" {
|
||||
prompt += fmt.Sprintf(" --no %s", task.NegPrompt)
|
||||
}
|
||||
body := ImageReq{
|
||||
BotType: "MID_JOURNEY",
|
||||
Prompt: prompt,
|
||||
Base64Array: make([]string, 0),
|
||||
}
|
||||
// 生成图片 Base64 编码
|
||||
if len(task.ImgArr) > 0 {
|
||||
imageData, err := utils.DownloadImage(task.ImgArr[0], "")
|
||||
if err != nil {
|
||||
logger.Error("error with download image: ", err)
|
||||
} else {
|
||||
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
|
||||
}
|
||||
|
||||
}
|
||||
return c.doRequest(body, apiPath, task.ChannelId)
|
||||
}
|
||||
|
||||
// Blend 融图
|
||||
func (c *Client) Blend(task types.MjTask) (ImageRes, error) {
|
||||
apiPath := fmt.Sprintf("mj-%s/mj/submit/blend", task.Mode)
|
||||
body := ImageReq{
|
||||
BotType: "MID_JOURNEY",
|
||||
Dimensions: "SQUARE",
|
||||
Base64Array: make([]string, 0),
|
||||
}
|
||||
// 生成图片 Base64 编码
|
||||
if len(task.ImgArr) > 0 {
|
||||
for _, imgURL := range task.ImgArr {
|
||||
imageData, err := utils.DownloadImage(imgURL, "")
|
||||
if err != nil {
|
||||
logger.Error("error with download image: ", err)
|
||||
} else {
|
||||
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
|
||||
}
|
||||
}
|
||||
}
|
||||
return c.doRequest(body, apiPath, task.ChannelId)
|
||||
}
|
||||
|
||||
// SwapFace 换脸
|
||||
func (c *Client) SwapFace(task types.MjTask) (ImageRes, error) {
|
||||
apiPath := fmt.Sprintf("mj-%s/mj/insight-face/swap", task.Mode)
|
||||
// 生成图片 Base64 编码
|
||||
if len(task.ImgArr) != 2 {
|
||||
return ImageRes{}, errors.New("参数错误，必须上传2张图片")
|
||||
}
|
||||
var sourceBase64 string
|
||||
var targetBase64 string
|
||||
imageData, err := utils.DownloadImage(task.ImgArr[0], "")
|
||||
if err != nil {
|
||||
logger.Error("error with download image: ", err)
|
||||
} else {
|
||||
sourceBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(imageData)
|
||||
}
|
||||
imageData, err = utils.DownloadImage(task.ImgArr[1], "")
|
||||
if err != nil {
|
||||
logger.Error("error with download image: ", err)
|
||||
} else {
|
||||
targetBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(imageData)
|
||||
}
|
||||
|
||||
body := gin.H{
|
||||
"sourceBase64": sourceBase64,
|
||||
"targetBase64": targetBase64,
|
||||
"accountFilter": gin.H{
|
||||
"instanceId": "",
|
||||
},
|
||||
"state": "",
|
||||
}
|
||||
return c.doRequest(body, apiPath, task.ChannelId)
|
||||
}
|
||||
|
||||
// Upscale 放大指定的图片
|
||||
func (c *Client) Upscale(task types.MjTask) (ImageRes, error) {
|
||||
body := map[string]string{
|
||||
"customId": fmt.Sprintf("MJ::JOB::upsample::%d::%s", task.Index, task.MessageHash),
|
||||
"taskId": task.MessageId,
|
||||
}
|
||||
apiPath := fmt.Sprintf("mj-%s/mj/submit/action", task.Mode)
|
||||
return c.doRequest(body, apiPath, task.ChannelId)
|
||||
}
|
||||
|
||||
// Variation 以指定的图片的视角进行变换再创作，注意需要在对应的频道中关闭 Remix 变换，否则 Variation 指令将不会生效
|
||||
func (c *Client) Variation(task types.MjTask) (ImageRes, error) {
|
||||
body := map[string]string{
|
||||
"customId": fmt.Sprintf("MJ::JOB::variation::%d::%s", task.Index, task.MessageHash),
|
||||
"taskId": task.MessageId,
|
||||
}
|
||||
apiPath := fmt.Sprintf("mj-%s/mj/submit/action", task.Mode)
|
||||
|
||||
return c.doRequest(body, apiPath, task.ChannelId)
|
||||
}
|
||||
|
||||
func (c *Client) doRequest(body interface{}, apiPath string, channel string) (ImageRes, error) {
|
||||
var res ImageRes
|
||||
var errRes ErrRes
|
||||
session := c.db.Session(&gorm.Session{}).Where("type", "mj").Where("enabled", true)
|
||||
if channel != "" {
|
||||
session = session.Where("api_url", channel)
|
||||
}
|
||||
|
||||
var apiKey model.ApiKey
|
||||
err := session.Order("last_used_at ASC").First(&apiKey).Error
|
||||
if err != nil {
|
||||
return ImageRes{}, fmt.Errorf("no available MidJourney api key: %v", err)
|
||||
}
|
||||
|
||||
if err = c.licenseService.IsValidApiURL(apiKey.ApiURL); err != nil {
|
||||
return ImageRes{}, err
|
||||
}
|
||||
|
||||
apiURL := fmt.Sprintf("%s/%s", apiKey.ApiURL, apiPath)
|
||||
logger.Info("API URL: ", apiURL)
|
||||
r, err := req.C().R().
|
||||
SetHeader("Authorization", "Bearer "+apiKey.Value).
|
||||
SetBody(body).
|
||||
SetSuccessResult(&res).
|
||||
SetErrorResult(&errRes).
|
||||
Post(apiURL)
|
||||
if err != nil {
|
||||
errMsg := err.Error()
|
||||
if r != nil {
|
||||
errStr, _ := io.ReadAll(r.Body)
|
||||
logger.Error("请求 API 出错：", string(errStr))
|
||||
errMsg = errMsg + " " + string(errStr)
|
||||
}
|
||||
return ImageRes{}, fmt.Errorf("请求 API 出错：%v", errMsg)
|
||||
}
|
||||
|
||||
if r.IsErrorState() {
|
||||
return ImageRes{}, fmt.Errorf("API 返回错误：%s", errRes.Error.Message)
|
||||
}
|
||||
|
||||
// update the api key last used time
|
||||
if err = c.db.Model(&apiKey).Update("last_used_at", time.Now().Unix()).Error; err != nil {
|
||||
logger.Error("update api key last used time error: ", err)
|
||||
}
|
||||
res.Channel = apiKey.ApiURL
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (c *Client) QueryTask(taskId string, channel string) (QueryRes, error) {
|
||||
var apiKey model.ApiKey
|
||||
err := c.db.Where("type", "mj").Where("enabled", true).Where("api_url", channel).First(&apiKey).Error
|
||||
if err != nil {
|
||||
return QueryRes{}, fmt.Errorf("no available MidJourney api key: %v", err)
|
||||
}
|
||||
apiURL := fmt.Sprintf("%s/mj/task/%s/fetch", apiKey.ApiURL, taskId)
|
||||
var res QueryRes
|
||||
r, err := c.client.R().SetHeader("Authorization", "Bearer "+apiKey.Value).
|
||||
SetSuccessResult(&res).
|
||||
Get(apiURL)
|
||||
|
||||
if err != nil {
|
||||
return QueryRes{}, err
|
||||
}
|
||||
|
||||
if r.IsErrorState() {
|
||||
return QueryRes{}, errors.New("error status:" + r.Status)
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
@@ -1,211 +0,0 @@
|
||||
package mj
|
||||
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
// * Copyright 2023 The Geek-AI Authors. All rights reserved.
|
||||
// * Use of this source code is governed by a Apache-2.0 license
|
||||
// * that can be found in the LICENSE file.
|
||||
// * @Author yangjian102621@163.com
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"geekai/core/types"
|
||||
"geekai/service"
|
||||
"geekai/utils"
|
||||
"github.com/imroc/req/v3"
|
||||
"io"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// PlusClient MidJourney Plus ProxyClient
|
||||
type PlusClient struct {
|
||||
Config types.MjPlusConfig
|
||||
apiURL string
|
||||
client *req.Client
|
||||
licenseService *service.LicenseService
|
||||
}
|
||||
|
||||
func NewPlusClient(config types.MjPlusConfig, licenseService *service.LicenseService) *PlusClient {
|
||||
return &PlusClient{
|
||||
Config: config,
|
||||
apiURL: config.ApiURL,
|
||||
client: req.C().SetTimeout(time.Minute).SetUserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36"),
|
||||
licenseService: licenseService,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *PlusClient) preCheck() error {
|
||||
return c.licenseService.IsValidApiURL(c.Config.ApiURL)
|
||||
}
|
||||
|
||||
func (c *PlusClient) Imagine(task types.MjTask) (ImageRes, error) {
|
||||
if err := c.preCheck(); err != nil {
|
||||
return ImageRes{}, err
|
||||
}
|
||||
|
||||
apiURL := fmt.Sprintf("%s/mj-%s/mj/submit/imagine", c.apiURL, c.Config.Mode)
|
||||
prompt := fmt.Sprintf("%s %s", task.Prompt, task.Params)
|
||||
if task.NegPrompt != "" {
|
||||
prompt += fmt.Sprintf(" --no %s", task.NegPrompt)
|
||||
}
|
||||
body := ImageReq{
|
||||
BotType: "MID_JOURNEY",
|
||||
Prompt: prompt,
|
||||
Base64Array: make([]string, 0),
|
||||
}
|
||||
// 生成图片 Base64 编码
|
||||
if len(task.ImgArr) > 0 {
|
||||
imageData, err := utils.DownloadImage(task.ImgArr[0], "")
|
||||
if err != nil {
|
||||
logger.Error("error with download image: ", err)
|
||||
} else {
|
||||
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
|
||||
}
|
||||
|
||||
}
|
||||
return c.doRequest(body, apiURL)
|
||||
}
|
||||
|
||||
// Blend 融图
|
||||
func (c *PlusClient) Blend(task types.MjTask) (ImageRes, error) {
|
||||
if err := c.preCheck(); err != nil {
|
||||
return ImageRes{}, err
|
||||
}
|
||||
|
||||
apiURL := fmt.Sprintf("%s/mj-%s/mj/submit/blend", c.apiURL, c.Config.Mode)
|
||||
logger.Info("API URL: ", apiURL)
|
||||
body := ImageReq{
|
||||
BotType: "MID_JOURNEY",
|
||||
Dimensions: "SQUARE",
|
||||
Base64Array: make([]string, 0),
|
||||
}
|
||||
// 生成图片 Base64 编码
|
||||
if len(task.ImgArr) > 0 {
|
||||
for _, imgURL := range task.ImgArr {
|
||||
imageData, err := utils.DownloadImage(imgURL, "")
|
||||
if err != nil {
|
||||
logger.Error("error with download image: ", err)
|
||||
} else {
|
||||
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
|
||||
}
|
||||
}
|
||||
}
|
||||
return c.doRequest(body, apiURL)
|
||||
}
|
||||
|
||||
// SwapFace 换脸
|
||||
func (c *PlusClient) SwapFace(task types.MjTask) (ImageRes, error) {
|
||||
if err := c.preCheck(); err != nil {
|
||||
return ImageRes{}, err
|
||||
}
|
||||
|
||||
apiURL := fmt.Sprintf("%s/mj-%s/mj/insight-face/swap", c.apiURL, c.Config.Mode)
|
||||
// 生成图片 Base64 编码
|
||||
if len(task.ImgArr) != 2 {
|
||||
return ImageRes{}, errors.New("参数错误，必须上传2张图片")
|
||||
}
|
||||
var sourceBase64 string
|
||||
var targetBase64 string
|
||||
imageData, err := utils.DownloadImage(task.ImgArr[0], "")
|
||||
if err != nil {
|
||||
logger.Error("error with download image: ", err)
|
||||
} else {
|
||||
sourceBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(imageData)
|
||||
}
|
||||
imageData, err = utils.DownloadImage(task.ImgArr[1], "")
|
||||
if err != nil {
|
||||
logger.Error("error with download image: ", err)
|
||||
} else {
|
||||
targetBase64 = "data:image/png;base64," + base64.StdEncoding.EncodeToString(imageData)
|
||||
}
|
||||
|
||||
body := gin.H{
|
||||
"sourceBase64": sourceBase64,
|
||||
"targetBase64": targetBase64,
|
||||
"accountFilter": gin.H{
|
||||
"instanceId": "",
|
||||
},
|
||||
"state": "",
|
||||
}
|
||||
return c.doRequest(body, apiURL)
|
||||
}
|
||||
|
||||
// Upscale 放大指定的图片
|
||||
func (c *PlusClient) Upscale(task types.MjTask) (ImageRes, error) {
|
||||
if err := c.preCheck(); err != nil {
|
||||
return ImageRes{}, err
|
||||
}
|
||||
|
||||
body := map[string]string{
|
||||
"customId": fmt.Sprintf("MJ::JOB::upsample::%d::%s", task.Index, task.MessageHash),
|
||||
"taskId": task.MessageId,
|
||||
}
|
||||
apiURL := fmt.Sprintf("%s/mj-%s/mj/submit/action", c.apiURL, c.Config.Mode)
|
||||
return c.doRequest(body, apiURL)
|
||||
}
|
||||
|
||||
// Variation 以指定的图片的视角进行变换再创作，注意需要在对应的频道中关闭 Remix 变换，否则 Variation 指令将不会生效
|
||||
func (c *PlusClient) Variation(task types.MjTask) (ImageRes, error) {
|
||||
if err := c.preCheck(); err != nil {
|
||||
return ImageRes{}, err
|
||||
}
|
||||
|
||||
body := map[string]string{
|
||||
"customId": fmt.Sprintf("MJ::JOB::variation::%d::%s", task.Index, task.MessageHash),
|
||||
"taskId": task.MessageId,
|
||||
}
|
||||
apiURL := fmt.Sprintf("%s/mj-%s/mj/submit/action", c.apiURL, c.Config.Mode)
|
||||
|
||||
return c.doRequest(body, apiURL)
|
||||
}
|
||||
|
||||
func (c *PlusClient) doRequest(body interface{}, apiURL string) (ImageRes, error) {
|
||||
var res ImageRes
|
||||
var errRes ErrRes
|
||||
logger.Info("API URL: ", apiURL)
|
||||
r, err := req.C().R().
|
||||
SetHeader("Authorization", "Bearer "+c.Config.ApiKey).
|
||||
SetBody(body).
|
||||
SetSuccessResult(&res).
|
||||
SetErrorResult(&errRes).
|
||||
Post(apiURL)
|
||||
if err != nil {
|
||||
errMsg := err.Error()
|
||||
if r != nil {
|
||||
errStr, _ := io.ReadAll(r.Body)
|
||||
logger.Error("请求 API 出错：", string(errStr))
|
||||
errMsg = errMsg + " " + string(errStr)
|
||||
}
|
||||
return ImageRes{}, fmt.Errorf("请求 API 出错：%v", errMsg)
|
||||
}
|
||||
|
||||
if r.IsErrorState() {
|
||||
return ImageRes{}, fmt.Errorf("API 返回错误：%s", errRes.Error.Message)
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (c *PlusClient) QueryTask(taskId string) (QueryRes, error) {
|
||||
apiURL := fmt.Sprintf("%s/mj/task/%s/fetch", c.apiURL, taskId)
|
||||
var res QueryRes
|
||||
r, err := c.client.R().SetHeader("Authorization", "Bearer "+c.Config.ApiKey).
|
||||
SetSuccessResult(&res).
|
||||
Get(apiURL)
|
||||
|
||||
if err != nil {
|
||||
return QueryRes{}, err
|
||||
}
|
||||
|
||||
if r.IsErrorState() {
|
||||
return QueryRes{}, errors.New("error status:" + r.Status)
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
var _ Client = &PlusClient{}
|
||||
@@ -1,215 +0,0 @@
|
||||
package mj
|
||||
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
// * Copyright 2023 The Geek-AI Authors. All rights reserved.
|
||||
// * Use of this source code is governed by a Apache-2.0 license
|
||||
// * that can be found in the LICENSE file.
|
||||
// * @Author yangjian102621@163.com
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
|
||||
import (
|
||||
"geekai/core/types"
|
||||
logger2 "geekai/logger"
|
||||
"geekai/service"
|
||||
"geekai/service/oss"
|
||||
"geekai/service/sd"
|
||||
"geekai/store"
|
||||
"geekai/store/model"
|
||||
"geekai/utils"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ServicePool Mj service pool
|
||||
type ServicePool struct {
|
||||
services []*Service
|
||||
taskQueue *store.RedisQueue
|
||||
notifyQueue *store.RedisQueue
|
||||
db *gorm.DB
|
||||
uploaderManager *oss.UploaderManager
|
||||
Clients *types.LMap[uint, *types.WsClient] // UserId => Client
|
||||
licenseService *service.LicenseService
|
||||
}
|
||||
|
||||
var logger = logger2.GetLogger()
|
||||
|
||||
func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderManager, licenseService *service.LicenseService) *ServicePool {
|
||||
services := make([]*Service, 0)
|
||||
taskQueue := store.NewRedisQueue("MidJourney_Task_Queue", redisCli)
|
||||
notifyQueue := store.NewRedisQueue("MidJourney_Notify_Queue", redisCli)
|
||||
return &ServicePool{
|
||||
taskQueue: taskQueue,
|
||||
notifyQueue: notifyQueue,
|
||||
services: services,
|
||||
uploaderManager: manager,
|
||||
db: db,
|
||||
Clients: types.NewLMap[uint, *types.WsClient](),
|
||||
licenseService: licenseService,
|
||||
}
|
||||
}
|
||||
|
||||
func (p *ServicePool) InitServices(plusConfigs []types.MjPlusConfig, proxyConfigs []types.MjProxyConfig) {
|
||||
// stop old service
|
||||
for _, s := range p.services {
|
||||
s.Stop()
|
||||
}
|
||||
p.services = make([]*Service, 0)
|
||||
|
||||
for _, config := range plusConfigs {
|
||||
if config.Enabled == false {
|
||||
continue
|
||||
}
|
||||
|
||||
cli := NewPlusClient(config, p.licenseService)
|
||||
name := utils.Md5(config.ApiURL)
|
||||
plusService := NewService(name, p.taskQueue, p.notifyQueue, p.db, cli)
|
||||
go func() {
|
||||
plusService.Run()
|
||||
}()
|
||||
p.services = append(p.services, plusService)
|
||||
}
|
||||
|
||||
// for mid-journey proxy
|
||||
for _, config := range proxyConfigs {
|
||||
if config.Enabled == false {
|
||||
continue
|
||||
}
|
||||
cli := NewProxyClient(config)
|
||||
name := utils.Md5(config.ApiURL)
|
||||
proxyService := NewService(name, p.taskQueue, p.notifyQueue, p.db, cli)
|
||||
go func() {
|
||||
proxyService.Run()
|
||||
}()
|
||||
p.services = append(p.services, proxyService)
|
||||
}
|
||||
}
|
||||
|
||||
func (p *ServicePool) CheckTaskNotify() {
|
||||
go func() {
|
||||
for {
|
||||
var message sd.NotifyMessage
|
||||
err := p.notifyQueue.LPop(&message)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
cli := p.Clients.Get(uint(message.UserId))
|
||||
if cli == nil {
|
||||
continue
|
||||
}
|
||||
err = cli.Send([]byte(message.Message))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (p *ServicePool) DownloadImages() {
|
||||
go func() {
|
||||
var items []model.MidJourneyJob
|
||||
for {
|
||||
res := p.db.Where("img_url = ? AND progress = ?", "", 100).Find(&items)
|
||||
if res.Error != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// download images
|
||||
for _, v := range items {
|
||||
if v.OrgURL == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Infof("try to download image: %s", v.OrgURL)
|
||||
mjService := p.getService(v.ChannelId)
|
||||
if mjService == nil {
|
||||
logger.Errorf("Invalid task: %+v", v)
|
||||
continue
|
||||
}
|
||||
|
||||
task, _ := mjService.Client.QueryTask(v.TaskId)
|
||||
if len(task.Buttons) > 0 {
|
||||
v.Hash = GetImageHash(task.Buttons[0].CustomId)
|
||||
}
|
||||
// 如果是返回的是 discord 图片地址，则使用代理下载
|
||||
proxy := false
|
||||
if strings.HasPrefix(v.OrgURL, "https://cdn.discordapp.com") {
|
||||
proxy = true
|
||||
}
|
||||
imgURL, err := p.uploaderManager.GetUploadHandler().PutUrlFile(v.OrgURL, proxy)
|
||||
|
||||
if err != nil {
|
||||
logger.Errorf("error with download image %s, %v", v.OrgURL, err)
|
||||
continue
|
||||
} else {
|
||||
logger.Infof("download image %s successfully.", v.OrgURL)
|
||||
}
|
||||
|
||||
v.ImgURL = imgURL
|
||||
p.db.Updates(&v)
|
||||
|
||||
cli := p.Clients.Get(uint(v.UserId))
|
||||
if cli == nil {
|
||||
continue
|
||||
}
|
||||
err = cli.Send([]byte(sd.Finished))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
time.Sleep(time.Second * 5)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// PushTask push a new mj task in to task queue
|
||||
func (p *ServicePool) PushTask(task types.MjTask) {
|
||||
logger.Debugf("add a new MidJourney task to the task list: %+v", task)
|
||||
p.taskQueue.RPush(task)
|
||||
}
|
||||
|
||||
// HasAvailableService check if it has available mj service in pool
|
||||
func (p *ServicePool) HasAvailableService() bool {
|
||||
return len(p.services) > 0
|
||||
}
|
||||
|
||||
// SyncTaskProgress 异步拉取任务
|
||||
func (p *ServicePool) SyncTaskProgress() {
|
||||
go func() {
|
||||
var jobs []model.MidJourneyJob
|
||||
for {
|
||||
res := p.db.Where("progress < ?", 100).Find(&jobs)
|
||||
if res.Error != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, job := range jobs {
|
||||
// 5 分钟还没完成的任务标记为失败
|
||||
if time.Now().Sub(job.CreatedAt) > time.Minute*5 {
|
||||
job.Progress = 101
|
||||
job.ErrMsg = "任务超时"
|
||||
p.db.Updates(&job)
|
||||
continue
|
||||
}
|
||||
|
||||
if servicePlus := p.getService(job.ChannelId); servicePlus != nil {
|
||||
_ = servicePlus.Notify(job)
|
||||
}
|
||||
}
|
||||
|
||||
time.Sleep(time.Second * 10)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (p *ServicePool) getService(name string) *Service {
|
||||
for _, s := range p.services {
|
||||
if s.Name == name {
|
||||
return s
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -1,185 +0,0 @@
|
||||
package mj
|
||||
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
// * Copyright 2023 The Geek-AI Authors. All rights reserved.
|
||||
// * Use of this source code is governed by a Apache-2.0 license
|
||||
// * that can be found in the LICENSE file.
|
||||
// * @Author yangjian102621@163.com
|
||||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"geekai/core/types"
|
||||
"geekai/utils"
|
||||
"github.com/imroc/req/v3"
|
||||
"io"
|
||||
)
|
||||
|
||||
// ProxyClient MidJourney Proxy Client
|
||||
type ProxyClient struct {
|
||||
Config types.MjProxyConfig
|
||||
apiURL string
|
||||
}
|
||||
|
||||
func NewProxyClient(config types.MjProxyConfig) *ProxyClient {
|
||||
return &ProxyClient{Config: config, apiURL: config.ApiURL}
|
||||
}
|
||||
|
||||
func (c *ProxyClient) Imagine(task types.MjTask) (ImageRes, error) {
|
||||
apiURL := fmt.Sprintf("%s/mj/submit/imagine", c.apiURL)
|
||||
prompt := fmt.Sprintf("%s %s", task.Prompt, task.Params)
|
||||
if task.NegPrompt != "" {
|
||||
prompt += fmt.Sprintf(" --no %s", task.NegPrompt)
|
||||
}
|
||||
body := ImageReq{
|
||||
Prompt: prompt,
|
||||
Base64Array: make([]string, 0),
|
||||
}
|
||||
// 生成图片 Base64 编码
|
||||
if len(task.ImgArr) > 0 {
|
||||
imageData, err := utils.DownloadImage(task.ImgArr[0], "")
|
||||
if err != nil {
|
||||
logger.Error("error with download image: ", err)
|
||||
} else {
|
||||
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
|
||||
}
|
||||
|
||||
}
|
||||
logger.Info("API URL: ", apiURL)
|
||||
var res ImageRes
|
||||
var errRes ErrRes
|
||||
r, err := req.C().R().
|
||||
SetHeader("mj-api-secret", c.Config.ApiKey).
|
||||
SetBody(body).
|
||||
SetSuccessResult(&res).
|
||||
SetErrorResult(&errRes).
|
||||
Post(apiURL)
|
||||
if err != nil {
|
||||
return ImageRes{}, fmt.Errorf("请求 API %s 出错：%v", apiURL, err)
|
||||
}
|
||||
|
||||
if r.IsErrorState() {
|
||||
errStr, _ := io.ReadAll(r.Body)
|
||||
return ImageRes{}, fmt.Errorf("API 返回错误：%s，%v", errRes.Error.Message, string(errStr))
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// Blend 融图
|
||||
func (c *ProxyClient) Blend(task types.MjTask) (ImageRes, error) {
|
||||
apiURL := fmt.Sprintf("%s/mj/submit/blend", c.apiURL)
|
||||
body := ImageReq{
|
||||
Dimensions: "SQUARE",
|
||||
Base64Array: make([]string, 0),
|
||||
}
|
||||
// 生成图片 Base64 编码
|
||||
if len(task.ImgArr) > 0 {
|
||||
for _, imgURL := range task.ImgArr {
|
||||
imageData, err := utils.DownloadImage(imgURL, "")
|
||||
if err != nil {
|
||||
logger.Error("error with download image: ", err)
|
||||
} else {
|
||||
body.Base64Array = append(body.Base64Array, "data:image/png;base64,"+base64.StdEncoding.EncodeToString(imageData))
|
||||
}
|
||||
}
|
||||
}
|
||||
var res ImageRes
|
||||
var errRes ErrRes
|
||||
r, err := req.C().R().
|
||||
SetHeader("mj-api-secret", c.Config.ApiKey).
|
||||
SetBody(body).
|
||||
SetSuccessResult(&res).
|
||||
SetErrorResult(&errRes).
|
||||
Post(apiURL)
|
||||
if err != nil {
|
||||
return ImageRes{}, fmt.Errorf("请求 API %s 出错：%v", apiURL, err)
|
||||
}
|
||||
|
||||
if r.IsErrorState() {
|
||||
return ImageRes{}, fmt.Errorf("API 返回错误：%s", errRes.Error.Message)
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// SwapFace 换脸
|
||||
func (c *ProxyClient) SwapFace(_ types.MjTask) (ImageRes, error) {
|
||||
return ImageRes{}, errors.New("MidJourney-Proxy暂未实现该功能，请使用 MidJourney-Plus")
|
||||
}
|
||||
|
||||
// Upscale 放大指定的图片
|
||||
func (c *ProxyClient) Upscale(task types.MjTask) (ImageRes, error) {
|
||||
body := map[string]interface{}{
|
||||
"action": "UPSCALE",
|
||||
"index": task.Index,
|
||||
"taskId": task.MessageId,
|
||||
}
|
||||
apiURL := fmt.Sprintf("%s/mj/submit/change", c.apiURL)
|
||||
var res ImageRes
|
||||
var errRes ErrRes
|
||||
r, err := req.C().R().
|
||||
SetHeader("mj-api-secret", c.Config.ApiKey).
|
||||
SetBody(body).
|
||||
SetSuccessResult(&res).
|
||||
SetErrorResult(&errRes).
|
||||
Post(apiURL)
|
||||
if err != nil {
|
||||
return ImageRes{}, fmt.Errorf("请求 API 出错：%v", err)
|
||||
}
|
||||
|
||||
if r.IsErrorState() {
|
||||
return ImageRes{}, fmt.Errorf("API 返回错误：%s", errRes.Error.Message)
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// Variation 以指定的图片的视角进行变换再创作，注意需要在对应的频道中关闭 Remix 变换，否则 Variation 指令将不会生效
|
||||
func (c *ProxyClient) Variation(task types.MjTask) (ImageRes, error) {
|
||||
body := map[string]interface{}{
|
||||
"action": "VARIATION",
|
||||
"index": task.Index,
|
||||
"taskId": task.MessageId,
|
||||
}
|
||||
apiURL := fmt.Sprintf("%s/mj/submit/change", c.apiURL)
|
||||
var res ImageRes
|
||||
var errRes ErrRes
|
||||
r, err := req.C().R().
|
||||
SetHeader("mj-api-secret", c.Config.ApiKey).
|
||||
SetBody(body).
|
||||
SetSuccessResult(&res).
|
||||
SetErrorResult(&errRes).
|
||||
Post(apiURL)
|
||||
if err != nil {
|
||||
return ImageRes{}, fmt.Errorf("请求 API 出错：%v", err)
|
||||
}
|
||||
|
||||
if r.IsErrorState() {
|
||||
return ImageRes{}, fmt.Errorf("API 返回错误：%s", errRes.Error.Message)
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (c *ProxyClient) QueryTask(taskId string) (QueryRes, error) {
|
||||
apiURL := fmt.Sprintf("%s/mj/task/%s/fetch", c.apiURL, taskId)
|
||||
var res QueryRes
|
||||
r, err := req.C().R().SetHeader("mj-api-secret", c.Config.ApiKey).
|
||||
SetSuccessResult(&res).
|
||||
Get(apiURL)
|
||||
|
||||
if err != nil {
|
||||
return QueryRes{}, err
|
||||
}
|
||||
|
||||
if r.IsErrorState() {
|
||||
return QueryRes{}, errors.New("error status:" + r.Status)
|
||||
}
|
||||
|
||||
return res, nil
|
||||
}
|
||||
|
||||
var _ Client = &ProxyClient{}
|
||||
@@ -11,10 +11,11 @@ import (
|
||||
"fmt"
|
||||
"geekai/core/types"
|
||||
"geekai/service"
|
||||
"geekai/service/sd"
|
||||
"geekai/service/oss"
|
||||
"geekai/store"
|
||||
"geekai/store/model"
|
||||
"geekai/utils"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -23,127 +24,112 @@ import (
|
||||
|
||||
// Service MJ 绘画服务
|
||||
type Service struct {
|
||||
Name string // service Name
|
||||
Client Client // MJ Client
|
||||
taskQueue *store.RedisQueue
|
||||
notifyQueue *store.RedisQueue
|
||||
db *gorm.DB
|
||||
running bool
|
||||
retryCount map[uint]int
|
||||
client *Client // MJ Client
|
||||
taskQueue *store.RedisQueue
|
||||
notifyQueue *store.RedisQueue
|
||||
db *gorm.DB
|
||||
Clients *types.LMap[uint, *types.WsClient] // UserId => Client
|
||||
uploaderManager *oss.UploaderManager
|
||||
}
|
||||
|
||||
func NewService(name string, taskQueue *store.RedisQueue, notifyQueue *store.RedisQueue, db *gorm.DB, cli Client) *Service {
|
||||
func NewService(redisCli *redis.Client, db *gorm.DB, client *Client, manager *oss.UploaderManager) *Service {
|
||||
return &Service{
|
||||
Name: name,
|
||||
db: db,
|
||||
taskQueue: taskQueue,
|
||||
notifyQueue: notifyQueue,
|
||||
Client: cli,
|
||||
running: true,
|
||||
retryCount: make(map[uint]int),
|
||||
db: db,
|
||||
taskQueue: store.NewRedisQueue("MidJourney_Task_Queue", redisCli),
|
||||
notifyQueue: store.NewRedisQueue("MidJourney_Notify_Queue", redisCli),
|
||||
client: client,
|
||||
Clients: types.NewLMap[uint, *types.WsClient](),
|
||||
uploaderManager: manager,
|
||||
}
|
||||
}
|
||||
|
||||
const failedProgress = 101
|
||||
|
||||
func (s *Service) Run() {
|
||||
logger.Infof("Starting MidJourney job consumer for %s", s.Name)
|
||||
for s.running {
|
||||
var task types.MjTask
|
||||
err := s.taskQueue.LPop(&task)
|
||||
if err != nil {
|
||||
logger.Errorf("taking task with error: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// 如果配置了多个中转平台的 API KEY
|
||||
// U,V 操作必须和 Image 操作属于同一个平台，否则找不到关联任务，需重新放回任务列表
|
||||
if task.ChannelId != "" && task.ChannelId != s.Name {
|
||||
if s.retryCount[task.Id] > 5 {
|
||||
s.db.Model(model.MidJourneyJob{Id: task.Id}).Delete(&model.MidJourneyJob{})
|
||||
logger.Info("Starting MidJourney job consumer for service")
|
||||
go func() {
|
||||
for {
|
||||
var task types.MjTask
|
||||
err := s.taskQueue.LPop(&task)
|
||||
if err != nil {
|
||||
logger.Errorf("taking task with error: %v", err)
|
||||
continue
|
||||
}
|
||||
logger.Debugf("handle other service task, name: %s, channel_id: %s, drop it.", s.Name, task.ChannelId)
|
||||
s.taskQueue.RPush(task)
|
||||
s.retryCount[task.Id]++
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
// translate prompt
|
||||
if utils.HasChinese(task.Prompt) {
|
||||
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.Prompt), "gpt-4o-mini")
|
||||
if err == nil {
|
||||
task.Prompt = content
|
||||
} else {
|
||||
logger.Warnf("error with translate prompt: %v", err)
|
||||
// translate prompt
|
||||
if utils.HasChinese(task.Prompt) {
|
||||
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.Prompt), "gpt-4o-mini")
|
||||
if err == nil {
|
||||
task.Prompt = content
|
||||
} else {
|
||||
logger.Warnf("error with translate prompt: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
// translate negative prompt
|
||||
if task.NegPrompt != "" && utils.HasChinese(task.NegPrompt) {
|
||||
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.RewritePromptTemplate, task.NegPrompt), "gpt-4o-mini")
|
||||
if err == nil {
|
||||
task.NegPrompt = content
|
||||
} else {
|
||||
logger.Warnf("error with translate prompt: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
var job model.MidJourneyJob
|
||||
tx := s.db.Where("id = ?", task.Id).First(&job)
|
||||
if tx.Error != nil {
|
||||
logger.Error("任务不存在，任务ID：", task.TaskId)
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Infof("%s handle a new MidJourney task: %+v", s.Name, task)
|
||||
var res ImageRes
|
||||
switch task.Type {
|
||||
case types.TaskImage:
|
||||
res, err = s.Client.Imagine(task)
|
||||
break
|
||||
case types.TaskUpscale:
|
||||
res, err = s.Client.Upscale(task)
|
||||
break
|
||||
case types.TaskVariation:
|
||||
res, err = s.Client.Variation(task)
|
||||
break
|
||||
case types.TaskBlend:
|
||||
res, err = s.Client.Blend(task)
|
||||
break
|
||||
case types.TaskSwapFace:
|
||||
res, err = s.Client.SwapFace(task)
|
||||
break
|
||||
}
|
||||
|
||||
if err != nil || (res.Code != 1 && res.Code != 22) {
|
||||
var errMsg string
|
||||
if err != nil {
|
||||
errMsg = err.Error()
|
||||
} else {
|
||||
errMsg = fmt.Sprintf("%v,%s", err, res.Description)
|
||||
// translate negative prompt
|
||||
if task.NegPrompt != "" && utils.HasChinese(task.NegPrompt) {
|
||||
content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, task.NegPrompt), "gpt-4o-mini")
|
||||
if err == nil {
|
||||
task.NegPrompt = content
|
||||
} else {
|
||||
logger.Warnf("error with translate prompt: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
logger.Error("绘画任务执行失败：", errMsg)
|
||||
job.Progress = failedProgress
|
||||
job.ErrMsg = errMsg
|
||||
// update the task progress
|
||||
// use fast mode as default
|
||||
if task.Mode == "" {
|
||||
task.Mode = "fast"
|
||||
}
|
||||
|
||||
var job model.MidJourneyJob
|
||||
tx := s.db.Where("id = ?", task.Id).First(&job)
|
||||
if tx.Error != nil {
|
||||
logger.Error("任务不存在，任务ID：", task.TaskId)
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Infof("handle a new MidJourney task: %+v", task)
|
||||
var res ImageRes
|
||||
switch task.Type {
|
||||
case types.TaskImage:
|
||||
res, err = s.client.Imagine(task)
|
||||
break
|
||||
case types.TaskUpscale:
|
||||
res, err = s.client.Upscale(task)
|
||||
break
|
||||
case types.TaskVariation:
|
||||
res, err = s.client.Variation(task)
|
||||
break
|
||||
case types.TaskBlend:
|
||||
res, err = s.client.Blend(task)
|
||||
break
|
||||
case types.TaskSwapFace:
|
||||
res, err = s.client.SwapFace(task)
|
||||
break
|
||||
}
|
||||
|
||||
if err != nil || (res.Code != 1 && res.Code != 22) {
|
||||
var errMsg string
|
||||
if err != nil {
|
||||
errMsg = err.Error()
|
||||
} else {
|
||||
errMsg = fmt.Sprintf("%v,%s", err, res.Description)
|
||||
}
|
||||
|
||||
logger.Error("绘画任务执行失败：", errMsg)
|
||||
job.Progress = service.FailTaskProgress
|
||||
job.ErrMsg = errMsg
|
||||
// update the task progress
|
||||
s.db.Updates(&job)
|
||||
// 任务失败，通知前端
|
||||
s.notifyQueue.RPush(service.NotifyMessage{UserId: task.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
|
||||
continue
|
||||
}
|
||||
logger.Infof("任务提交成功：%+v", res)
|
||||
// 更新任务 ID/频道
|
||||
job.TaskId = res.Result
|
||||
job.MessageId = res.Result
|
||||
job.ChannelId = res.Channel
|
||||
s.db.Updates(&job)
|
||||
// 任务失败，通知前端
|
||||
s.notifyQueue.RPush(sd.NotifyMessage{UserId: task.UserId, JobId: int(job.Id), Message: sd.Failed})
|
||||
continue
|
||||
}
|
||||
logger.Infof("任务提交成功：%+v", res)
|
||||
// 更新任务 ID/频道
|
||||
job.TaskId = res.Result
|
||||
job.MessageId = res.Result
|
||||
job.ChannelId = s.Name
|
||||
s.db.Updates(&job)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) Stop() {
|
||||
s.running = false
|
||||
}()
|
||||
}
|
||||
|
||||
type CBReq struct {
|
||||
@@ -164,46 +150,6 @@ type CBReq struct {
|
||||
} `json:"properties"`
|
||||
}
|
||||
|
||||
func (s *Service) Notify(job model.MidJourneyJob) error {
|
||||
task, err := s.Client.QueryTask(job.TaskId)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 任务执行失败了
|
||||
if task.FailReason != "" {
|
||||
s.db.Model(&model.MidJourneyJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
|
||||
"progress": failedProgress,
|
||||
"err_msg": task.FailReason,
|
||||
})
|
||||
s.notifyQueue.RPush(sd.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: sd.Failed})
|
||||
return fmt.Errorf("task failed: %v", task.FailReason)
|
||||
}
|
||||
|
||||
if len(task.Buttons) > 0 {
|
||||
job.Hash = GetImageHash(task.Buttons[0].CustomId)
|
||||
}
|
||||
oldProgress := job.Progress
|
||||
job.Progress = utils.IntValue(strings.Replace(task.Progress, "%", "", 1), 0)
|
||||
job.Prompt = task.PromptEn
|
||||
if task.ImageUrl != "" {
|
||||
job.OrgURL = task.ImageUrl
|
||||
}
|
||||
tx := s.db.Updates(&job)
|
||||
if tx.Error != nil {
|
||||
return fmt.Errorf("error with update database: %v", tx.Error)
|
||||
}
|
||||
// 通知前端更新任务进度
|
||||
if oldProgress != job.Progress {
|
||||
message := sd.Running
|
||||
if job.Progress == 100 {
|
||||
message = sd.Finished
|
||||
}
|
||||
s.notifyQueue.RPush(sd.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: message})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func GetImageHash(action string) string {
|
||||
split := strings.Split(action, "::")
|
||||
if len(split) > 5 {
|
||||
@@ -211,3 +157,143 @@ func GetImageHash(action string) string {
|
||||
}
|
||||
return split[len(split)-1]
|
||||
}
|
||||
|
||||
func (s *Service) CheckTaskNotify() {
|
||||
go func() {
|
||||
for {
|
||||
var message service.NotifyMessage
|
||||
err := s.notifyQueue.LPop(&message)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
cli := s.Clients.Get(uint(message.UserId))
|
||||
if cli == nil {
|
||||
continue
|
||||
}
|
||||
err = cli.Send([]byte(message.Message))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *Service) DownloadImages() {
|
||||
go func() {
|
||||
var items []model.MidJourneyJob
|
||||
for {
|
||||
res := s.db.Where("img_url = ? AND progress = ?", "", 100).Find(&items)
|
||||
if res.Error != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// download images
|
||||
for _, v := range items {
|
||||
if v.OrgURL == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
logger.Infof("try to download image: %s", v.OrgURL)
|
||||
// 如果是返回的是 discord 图片地址，则使用代理下载
|
||||
proxy := false
|
||||
if strings.HasPrefix(v.OrgURL, "https://cdn.discordapp.com") {
|
||||
proxy = true
|
||||
}
|
||||
imgURL, err := s.uploaderManager.GetUploadHandler().PutUrlFile(v.OrgURL, proxy)
|
||||
|
||||
if err != nil {
|
||||
logger.Errorf("error with download image %s, %v", v.OrgURL, err)
|
||||
continue
|
||||
} else {
|
||||
logger.Infof("download image %s successfully.", v.OrgURL)
|
||||
}
|
||||
|
||||
v.ImgURL = imgURL
|
||||
s.db.Updates(&v)
|
||||
|
||||
cli := s.Clients.Get(uint(v.UserId))
|
||||
if cli == nil {
|
||||
continue
|
||||
}
|
||||
err = cli.Send([]byte(service.TaskStatusFinished))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
time.Sleep(time.Second * 5)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// PushTask push a new mj task in to task queue
|
||||
func (s *Service) PushTask(task types.MjTask) {
|
||||
logger.Debugf("add a new MidJourney task to the task list: %+v", task)
|
||||
s.taskQueue.RPush(task)
|
||||
}
|
||||
|
||||
// SyncTaskProgress 异步拉取任务
|
||||
func (s *Service) SyncTaskProgress() {
|
||||
go func() {
|
||||
var jobs []model.MidJourneyJob
|
||||
for {
|
||||
res := s.db.Where("progress < ?", 100).Where("channel_id <> ?", "").Find(&jobs)
|
||||
if res.Error != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, job := range jobs {
|
||||
// 10 分钟还没完成的任务标记为失败
|
||||
if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
|
||||
job.Progress = service.FailTaskProgress
|
||||
job.ErrMsg = "任务超时"
|
||||
s.db.Updates(&job)
|
||||
continue
|
||||
}
|
||||
|
||||
task, err := s.client.QueryTask(job.TaskId, job.ChannelId)
|
||||
if err != nil {
|
||||
logger.Errorf("error with query task: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// 任务执行失败了
|
||||
if task.FailReason != "" {
|
||||
s.db.Model(&model.MidJourneyJob{Id: job.Id}).UpdateColumns(map[string]interface{}{
|
||||
"progress": service.FailTaskProgress,
|
||||
"err_msg": task.FailReason,
|
||||
})
|
||||
logger.Errorf("task failed: %v", task.FailReason)
|
||||
s.notifyQueue.RPush(service.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: service.TaskStatusFailed})
|
||||
continue
|
||||
}
|
||||
|
||||
if len(task.Buttons) > 0 {
|
||||
job.Hash = GetImageHash(task.Buttons[0].CustomId)
|
||||
}
|
||||
oldProgress := job.Progress
|
||||
job.Progress = utils.IntValue(strings.Replace(task.Progress, "%", "", 1), 0)
|
||||
job.Prompt = task.PromptEn
|
||||
if task.ImageUrl != "" {
|
||||
job.OrgURL = task.ImageUrl
|
||||
}
|
||||
err = s.db.Updates(&job).Error
|
||||
if err != nil {
|
||||
logger.Errorf("error with update database: %v", err)
|
||||
continue
|
||||
}
|
||||
|
||||
// 通知前端更新任务进度
|
||||
if oldProgress != job.Progress {
|
||||
message := service.TaskStatusRunning
|
||||
if job.Progress == 100 {
|
||||
message = service.TaskStatusFinished
|
||||
}
|
||||
s.notifyQueue.RPush(service.NotifyMessage{UserId: job.UserId, JobId: int(job.Id), Message: message})
|
||||
}
|
||||
}
|
||||
|
||||
time.Sleep(time.Second * 5)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user