mirror of
https://github.com/yangjian102621/geekai.git
synced 2026-09-16 18:27:13 +00:00
RockYang
d4fd38ab7e
- 同步 Plus v4.3.1 功能源并移除商业 License 闭环 - 更新开源镜像命名、Docker 部署版本和 geekai 数据库配置 - 补充前端 ESLint 检查配置并修复存量解析与模板问题 - 保留 JWT、管理员权限、API Key 和 OAuth 等正常鉴权机制
318 lines
8.6 KiB
Go
318 lines
8.6 KiB
Go
package image
|
||
|
||
// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
|
||
// * 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/types"
|
||
"geekai/log"
|
||
"geekai/service"
|
||
"geekai/service/oss"
|
||
"geekai/store"
|
||
"geekai/store/model"
|
||
"geekai/utils"
|
||
"io"
|
||
"time"
|
||
|
||
"github.com/go-redis/redis/v8"
|
||
|
||
"github.com/imroc/req/v3"
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
var logger = log.GetLogger()
|
||
|
||
// Image Generation Service
|
||
|
||
type Service struct {
|
||
httpClient *req.Client
|
||
db *gorm.DB
|
||
uploadManager *oss.UploaderManager
|
||
taskQueue *store.RedisQueue
|
||
userService *service.UserService
|
||
}
|
||
|
||
func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, userService *service.UserService) *Service {
|
||
return &Service{
|
||
httpClient: req.C().SetTimeout(time.Minute * 3),
|
||
db: db,
|
||
taskQueue: store.NewRedisQueue("Image_Task_Queue", redisCli),
|
||
uploadManager: manager,
|
||
userService: userService,
|
||
}
|
||
}
|
||
|
||
// PushTask push a new image task in to task queue
|
||
func (s *Service) PushTask(task types.ImageTask) {
|
||
logger.Infof("add a new Image generation task to the task list: %+v", task)
|
||
if err := s.taskQueue.RPush(task); err != nil {
|
||
logger.Errorf("push image task to queue failed: %v", err)
|
||
}
|
||
}
|
||
|
||
func (s *Service) Run() {
|
||
// 将数据库中未提交的任务加载到队列
|
||
var jobs []model.ImageJob
|
||
s.db.Where("task_id", "").Where("status IN ?", []string{model.ImageStatusPending, model.ImageStatusInProgress}).Find(&jobs)
|
||
for _, v := range jobs {
|
||
var task types.ImageTask
|
||
err := utils.JsonDecode(v.Params, &task)
|
||
if err != nil {
|
||
logger.Errorf("decode task params with error: %v", err)
|
||
continue
|
||
}
|
||
task.Id = v.Id
|
||
s.PushTask(task)
|
||
}
|
||
|
||
logger.Info("Starting Image generation job consumer...")
|
||
go func() {
|
||
for {
|
||
var task types.ImageTask
|
||
err := s.taskQueue.LPop(&task)
|
||
if err != nil {
|
||
logger.Errorf("taking task with error: %v", err)
|
||
continue
|
||
}
|
||
logger.Infof("handle a new Image generation task: %+v", task)
|
||
go func() {
|
||
s.db.Model(&model.ImageJob{Id: task.Id}).UpdateColumns(map[string]any{
|
||
"status": model.ImageStatusInProgress,
|
||
"progress": 1,
|
||
"err_msg": "",
|
||
})
|
||
_, err = s.Image(task, false)
|
||
if err != nil {
|
||
logger.Errorf("error with image task: %v", err)
|
||
s.db.Model(&model.ImageJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
|
||
"status": model.ImageStatusFailed,
|
||
"progress": 0,
|
||
"err_msg": err.Error(),
|
||
})
|
||
}
|
||
}()
|
||
}
|
||
}()
|
||
}
|
||
|
||
type imgReq struct {
|
||
Model string `json:"model"`
|
||
Image []string `json:"image,omitempty"`
|
||
Prompt string `json:"prompt"`
|
||
AspectRatio string `json:"aspect_ratio,omitempty"`
|
||
Size string `json:"size,omitempty"`
|
||
ResponseFormat string `json:"response_format,omitempty"`
|
||
}
|
||
|
||
type imgRes struct {
|
||
Created int64 `json:"created"`
|
||
Data []struct {
|
||
RevisedPrompt string `json:"revised_prompt,omitempty"`
|
||
Url string `json:"url,omitempty"`
|
||
B64Json string `json:"b64_json,omitempty"`
|
||
} `json:"data"`
|
||
}
|
||
|
||
type ErrRes struct {
|
||
Error struct {
|
||
Code any `json:"code"`
|
||
Message string `json:"message"`
|
||
Param any `json:"param"`
|
||
Type string `json:"type"`
|
||
} `json:"error"`
|
||
}
|
||
|
||
func (s *Service) Image(task types.ImageTask, sync bool) (string, error) {
|
||
logger.Debugf("绘画参数：%+v", task)
|
||
|
||
var chatModel model.ChatModel
|
||
if task.ModelId > 0 {
|
||
s.db.Where("id", task.ModelId).First(&chatModel)
|
||
} else {
|
||
s.db.Where("value", task.ModelValue).First(&chatModel)
|
||
}
|
||
|
||
// get image generation API KEY
|
||
var apiKey model.ApiKey
|
||
session := s.db.Where("enabled", true)
|
||
if chatModel.KeyId > 0 {
|
||
session = session.Where("id = ?", chatModel.KeyId)
|
||
} else {
|
||
session = session.Where("type = ?", "image")
|
||
}
|
||
err := session.Order("last_used_at ASC").First(&apiKey).Error
|
||
if err != nil {
|
||
return "", fmt.Errorf("no available Image Generation api key: %v", err)
|
||
}
|
||
|
||
var res imgRes
|
||
var errRes ErrRes
|
||
if len(apiKey.ProxyURL) > 5 {
|
||
s.httpClient.SetProxyURL(apiKey.ProxyURL).R()
|
||
}
|
||
apiURL := fmt.Sprintf("%s/v1/images/generations", apiKey.ApiURL)
|
||
reqBody := imgReq{
|
||
Model: chatModel.Value,
|
||
Prompt: task.Prompt,
|
||
AspectRatio: task.AspectRatio,
|
||
Size: task.Size,
|
||
ResponseFormat: "url",
|
||
}
|
||
|
||
// 图片编辑
|
||
if len(task.Image) > 0 {
|
||
reqBody.Image = task.Image
|
||
}
|
||
|
||
logger.Infof("Channel:%s, API KEY:%s, BODY: %+v", apiURL, apiKey.Value, reqBody)
|
||
r, err := s.httpClient.R().SetHeader("Body-Type", "application/json").
|
||
SetHeader("Authorization", "Bearer "+apiKey.Value).
|
||
SetBody(reqBody).
|
||
SetErrorResult(&errRes).
|
||
SetSuccessResult(&res).
|
||
Post(apiURL)
|
||
if err != nil {
|
||
logger.Errorf("error with send request: %v", err)
|
||
return "", fmt.Errorf("error with send request: %v", err)
|
||
}
|
||
|
||
if r.IsErrorState() {
|
||
logger.Errorf("error with send request, status: %s, %+v", r.Status, errRes.Error)
|
||
return "", fmt.Errorf("error with send request, status: %s, %+v", r.Status, errRes.Error)
|
||
}
|
||
|
||
if len(res.Data) == 0 && r.Body != nil {
|
||
body, _ := io.ReadAll(r.Body)
|
||
return "", fmt.Errorf("%s", string(body))
|
||
}
|
||
|
||
// update the api key last use time
|
||
s.db.Model(&apiKey).UpdateColumn("last_used_at", time.Now().Unix())
|
||
var imgURL string
|
||
var data = map[string]any{
|
||
"status": model.ImageStatusDownloading,
|
||
"progress": 100,
|
||
"prompt": task.Prompt,
|
||
"task_id": utils.RandString(12),
|
||
}
|
||
// 如果返回的是base64，则需要上传到oss
|
||
if res.Data[0].B64Json != "" {
|
||
imgURL, err = s.uploadManager.GetUploadHandler().PutBase64(res.Data[0].B64Json)
|
||
if err != nil {
|
||
return "", fmt.Errorf("error with upload image: %v", err)
|
||
}
|
||
logger.Infof("upload image to oss: %s", imgURL)
|
||
data["img_url"] = imgURL
|
||
} else {
|
||
imgURL = res.Data[0].Url
|
||
}
|
||
data["org_url"] = imgURL
|
||
// update task progress
|
||
err = s.db.Model(&model.ImageJob{Id: task.Id}).UpdateColumns(data).Error
|
||
if err != nil {
|
||
return "", fmt.Errorf("err with update database: %v", err)
|
||
}
|
||
|
||
var content string
|
||
if sync {
|
||
content = fmt.Sprintf("```\n%s\n```\n下面是我为你创作的图片：\n\n\n", task.Prompt, imgURL)
|
||
}
|
||
|
||
return content, nil
|
||
}
|
||
|
||
func (s *Service) CheckTaskStatus() {
|
||
go func() {
|
||
logger.Info("Running Image generation task status checking ...")
|
||
for {
|
||
// 检查未完成任务进度
|
||
var jobs []model.ImageJob
|
||
s.db.Where("status IN ?", []string{model.ImageStatusPending, model.ImageStatusInProgress, model.ImageStatusDownloading}).Find(&jobs)
|
||
for _, job := range jobs {
|
||
// 超时的任务标记为失败
|
||
if time.Since(job.CreatedAt) > time.Minute*10 {
|
||
job.Status = model.ImageStatusFailed
|
||
job.Progress = 0
|
||
job.ErrMsg = "任务超时"
|
||
s.db.Updates(&job)
|
||
}
|
||
}
|
||
|
||
// 找出失败的任务，并恢复其扣减积分
|
||
s.db.Where("status", model.ImageStatusFailed).Where("power > ?", 0).Find(&jobs)
|
||
for _, job := range jobs {
|
||
var task types.ImageTask
|
||
err := utils.JsonDecode(job.Params, &task)
|
||
if err != nil {
|
||
continue
|
||
}
|
||
err = s.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{
|
||
Type: types.PowerRefund,
|
||
Model: task.ModelName,
|
||
Remark: fmt.Sprintf("任务失败，退回积分。任务ID：%d，Err: %s", job.Id, job.ErrMsg),
|
||
})
|
||
if err != nil {
|
||
continue
|
||
}
|
||
// 更新任务状态
|
||
s.db.Model(&job).UpdateColumn("power", 0)
|
||
}
|
||
time.Sleep(time.Second * 10)
|
||
}
|
||
}()
|
||
}
|
||
|
||
func (s *Service) DownloadImages() {
|
||
go func() {
|
||
var items []model.ImageJob
|
||
for {
|
||
res := s.db.Where("img_url = ? AND status = ?", "", model.ImageStatusDownloading).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)
|
||
imgURL, err := s.downloadImage(v.Id, v.OrgURL)
|
||
if err != nil {
|
||
logger.Error("error with download image: %s, error: %v", imgURL, err)
|
||
continue
|
||
} else {
|
||
logger.Infof("download image %s successfully.", v.OrgURL)
|
||
}
|
||
|
||
}
|
||
|
||
time.Sleep(time.Second * 5)
|
||
}
|
||
}()
|
||
}
|
||
|
||
func (s *Service) downloadImage(jobId uint, orgURL string) (string, error) {
|
||
// sava image
|
||
imgURL, err := s.uploadManager.GetUploadHandler().PutUrlFile(orgURL, ".png", false)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
|
||
// update img_url
|
||
res := s.db.Model(&model.ImageJob{Id: jobId}).UpdateColumns(map[string]any{
|
||
"img_url": imgURL,
|
||
"status": model.ImageStatusSuccess,
|
||
})
|
||
if res.Error != nil {
|
||
return "", err
|
||
}
|
||
return imgURL, nil
|
||
}
|