merge v4.0.8

This commit is contained in:
RockYang
2024-08-01 18:09:00 +08:00
76 changed files with 3392 additions and 1940 deletions

View File

@@ -126,7 +126,7 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
// get image generation API KEY
var apiKey model.ApiKey
tx := s.db.Where("platform", types.OpenAI).
tx := s.db.Where("platform", types.OpenAI.Value).
Where("type", "img").
Where("enabled", true).
Order("last_used_at ASC").First(&apiKey)
@@ -146,7 +146,7 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
Model: "dall-e-3",
Prompt: prompt,
N: 1,
Size: "1024x1024",
Size: task.Size,
Style: task.Style,
Quality: task.Quality,
}).
@@ -162,11 +162,14 @@ func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
// update the api key last use time
s.db.Model(&apiKey).UpdateColumn("last_used_at", time.Now().Unix())
// update task progress
s.db.Model(&model.DallJob{Id: task.JobId}).UpdateColumns(map[string]interface{}{
tx = s.db.Model(&model.DallJob{Id: task.JobId}).UpdateColumns(map[string]interface{}{
"progress": 100,
"org_url": res.Data[0].Url,
"prompt": prompt,
})
if tx.Error != nil {
return "", fmt.Errorf("err with update database: %v", tx.Error)
}
s.notifyQueue.RPush(sd.NotifyMessage{UserId: int(task.UserId), JobId: int(task.JobId), Message: sd.Finished})
var content string

View File

@@ -0,0 +1,197 @@
package service
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
// * 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 (
"fmt"
"geekai/core"
"geekai/core/types"
"geekai/store"
"time"
"github.com/imroc/req/v3"
)
type LicenseService struct {
config types.ApiConfig
levelDB *store.LevelDB
license *types.License
urlWhiteList []string
machineId string
}
func NewLicenseService(server *core.AppServer, levelDB *store.LevelDB) *LicenseService {
var license types.License
return &LicenseService{
config: server.Config.ApiConfig,
levelDB: levelDB,
license: &license,
machineId: "",
}
}
type License struct {
Name string `json:"name"`
License string `json:"license"`
MachineId string `json:"mid"`
ActiveAt int64 `json:"active_at"`
ExpiredAt int64 `json:"expired_at"`
UserNum int `json:"user_num"`
Configs types.LicenseConfig `json:"configs"`
}
// ActiveLicense 激活 License
func (s *LicenseService) ActiveLicense(license string, machineId string) error {
var res struct {
Code types.BizCode `json:"code"`
Message string `json:"message"`
Data License `json:"data"`
}
apiURL := fmt.Sprintf("%s/%s", s.config.ApiURL, "api/license/active")
response, err := req.C().R().
SetBody(map[string]string{"license": license, "machine_id": machineId}).
SetSuccessResult(&res).Post(apiURL)
if err != nil {
return fmt.Errorf("发送激活请求失败: %v", err)
}
if response.IsErrorState() {
return fmt.Errorf("发送激活请求失败:%v", response.Status)
}
if res.Code != types.Success {
return fmt.Errorf("激活失败:%v", res.Message)
}
s.license = &types.License{
Key: license,
MachineId: machineId,
Configs: res.Data.Configs,
ExpiredAt: res.Data.ExpiredAt,
IsActive: true,
}
err = s.levelDB.Put(types.LicenseKey, s.license)
if err != nil {
return fmt.Errorf("保存许可证书失败:%v", err)
}
return nil
}
// SyncLicense 定期同步 License
func (s *LicenseService) SyncLicense() {
go func() {
retryCounter := 0
for {
license, err := s.fetchLicense()
if err != nil {
retryCounter++
if retryCounter < 5 {
logger.Error(err)
}
s.license.IsActive = false
} else {
s.license = license
}
urls, err := s.fetchUrlWhiteList()
if err == nil {
s.urlWhiteList = urls
}
time.Sleep(time.Second * 10)
}
}()
}
func (s *LicenseService) fetchLicense() (*types.License, error) {
//var res struct {
// Code types.BizCode `json:"code"`
// Message string `json:"message"`
// Data License `json:"data"`
//}
//apiURL := fmt.Sprintf("%s/%s", s.config.ApiURL, "api/license/check")
//response, err := req.C().R().
// SetBody(map[string]string{"license": s.license.Key, "machine_id": s.machineId}).
// SetSuccessResult(&res).Post(apiURL)
//if err != nil {
// return nil, fmt.Errorf("发送激活请求失败: %v", err)
//}
//if response.IsErrorState() {
// return nil, fmt.Errorf("激活失败:%v", response.Status)
//}
//if res.Code != types.Success {
// return nil, fmt.Errorf("激活失败:%v", res.Message)
//}
return &types.License{
Key: "abc",
MachineId: "abc",
Configs: types.LicenseConfig{
UserNum: 10000,
DeCopy: false,
},
ExpiredAt: 0,
IsActive: true,
}, nil
}
func (s *LicenseService) fetchUrlWhiteList() ([]string, error) {
var res struct {
Code types.BizCode `json:"code"`
Message string `json:"message"`
Data []string `json:"data"`
}
apiURL := fmt.Sprintf("%s/%s", s.config.ApiURL, "api/license/urls")
response, err := req.C().R().SetSuccessResult(&res).Get(apiURL)
if err != nil {
return nil, fmt.Errorf("发送请求失败: %v", err)
}
if response.IsErrorState() {
return nil, fmt.Errorf("发送请求失败:%v", response.Status)
}
if res.Code != types.Success {
return nil, fmt.Errorf("获取白名单失败:%v", res.Message)
}
return res.Data, nil
}
// GetLicense 获取许可信息
func (s *LicenseService) GetLicense() *types.License {
return s.license
}
// IsValidApiURL 判断是否合法的中转 URL
func (s *LicenseService) IsValidApiURL(uri string) error {
// 获得许可授权的直接放行
return nil
//if s.license.IsActive {
// if s.license.MachineId != s.machineId {
// return errors.New("系统使用了盗版的许可证书")
// }
//
// if time.Now().Unix() > s.license.ExpiredAt {
// return errors.New("系统许可证书已经过期")
// }
// return nil
//}
//
//if len(s.urlWhiteList) == 0 {
// urls, err := s.fetchUrlWhiteList()
// if err == nil {
// s.urlWhiteList = urls
// }
//}
//
//for _, v := range s.urlWhiteList {
// if strings.HasPrefix(uri, v) {
// return nil
// }
//}
//return fmt.Errorf("当前 API 地址 %s 不在白名单列表当中。", uri)
}

View File

@@ -12,6 +12,7 @@ import (
"errors"
"fmt"
"geekai/core/types"
"geekai/service"
"geekai/utils"
"github.com/imroc/req/v3"
"io"
@@ -22,20 +23,30 @@ import (
// PlusClient MidJourney Plus ProxyClient
type PlusClient struct {
Config types.MjPlusConfig
apiURL string
client *req.Client
Config types.MjPlusConfig
apiURL string
client *req.Client
licenseService *service.LicenseService
}
func NewPlusClient(config types.MjPlusConfig) *PlusClient {
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"),
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 != "" {
@@ -79,6 +90,10 @@ func (c *PlusClient) Imagine(task types.MjTask) (ImageRes, error) {
// 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{
@@ -118,6 +133,10 @@ func (c *PlusClient) Blend(task types.MjTask) (ImageRes, error) {
// 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 {
@@ -167,6 +186,10 @@ func (c *PlusClient) SwapFace(task types.MjTask) (ImageRes, error) {
// 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,
@@ -194,6 +217,10 @@ func (c *PlusClient) Upscale(task types.MjTask) (ImageRes, error) {
// 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,

View File

@@ -11,6 +11,7 @@ import (
"fmt"
"geekai/core/types"
logger2 "geekai/logger"
"geekai/service"
"geekai/service/oss"
"geekai/service/sd"
"geekai/store"
@@ -30,11 +31,12 @@ type ServicePool struct {
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) *ServicePool {
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)
@@ -45,6 +47,7 @@ func NewServicePool(db *gorm.DB, redisCli *redis.Client, manager *oss.UploaderMa
uploaderManager: manager,
db: db,
Clients: types.NewLMap[uint, *types.WsClient](),
licenseService: licenseService,
}
}
@@ -60,7 +63,7 @@ func (p *ServicePool) InitServices(plusConfigs []types.MjPlusConfig, proxyConfig
continue
}
cli := NewPlusClient(config)
cli := NewPlusClient(config, p.licenseService)
name := fmt.Sprintf("mj-plus-service-%d", k)
plusService := NewService(name, p.taskQueue, p.notifyQueue, p.db, cli)
go func() {

View File

@@ -108,7 +108,13 @@ func (s *Service) Run() {
}
if err != nil || (res.Code != 1 && res.Code != 22) {
errMsg := fmt.Sprintf("%v,%s", err, res.Description)
var errMsg string
if err != nil {
errMsg = err.Error()
} else {
errMsg = fmt.Sprintf("%v,%s", err, res.Description)
}
logger.Error("绘画任务执行失败:", errMsg)
job.Progress = -1
job.ErrMsg = errMsg

View File

@@ -49,6 +49,8 @@ type HuPiPayReq struct {
CallbackURL string `json:"callback_url"`
Time string `json:"time"`
NonceStr string `json:"nonce_str"`
Type string `json:"type"`
WapUrl string `json:"wap_url"`
}
type HuPiResp struct {