geekai/api/handler/sd_handler.go

package handler

import (
	"chatplus/core"
	"chatplus/core/types"
	"chatplus/service"
	"chatplus/service/oss"
	"chatplus/store/model"
	"chatplus/store/vo"
	"chatplus/utils"
	"chatplus/utils/resp"
	"encoding/base64"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis/v8"
	"github.com/gorilla/websocket"
	"gorm.io/gorm"
	"net/http"
	"strings"
	"sync"
	"time"
)

type SdJobHandler struct {
	BaseHandler
	redis           *redis.Client
	db              *gorm.DB
	mjService       *service.MjService
	uploaderManager *oss.UploaderManager
	lock            sync.Mutex
	clients         *types.LMap[string, *types.WsClient]
}

func NewSdJobHandler(
	app *core.AppServer,
	client *redis.Client,
	db *gorm.DB,
	manager *oss.UploaderManager,
	mjService *service.MjService) *MidJourneyHandler {
	h := MidJourneyHandler{
		redis:           client,
		db:              db,
		uploaderManager: manager,
		lock:            sync.Mutex{},
		mjService:       mjService,
		clients:         types.NewLMap[string, *types.WsClient](),
	}
	h.App = app
	return &h
}

// Client WebSocket 客户端，用于通知任务状态变更
func (h *SdJobHandler) Client(c *gin.Context) {
	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		logger.Error(err)
		return
	}

	sessionId := c.Query("session_id")
	client := types.NewWsClient(ws)
	// 删除旧的连接
	h.clients.Delete(sessionId)
	h.clients.Put(sessionId, client)
	logger.Infof("New websocket connected, IP: %s", c.ClientIP())
}

type sdNotifyData struct {
	TaskId    string
	ImageName string
	ImageData string
	Progress  int
	Seed      string
	Success   bool
	Message   string
}

func (h *SdJobHandler) Notify(c *gin.Context) {
	token := c.GetHeader("Authorization")
	if token != h.App.Config.ExtConfig.Token {
		resp.NotAuth(c)
		return
	}
	var data sdNotifyData
	if err := c.ShouldBindJSON(&data); err != nil || data.TaskId == "" {
		resp.ERROR(c, types.InvalidArgs)
		return
	}
	logger.Debugf("收到 MidJourney 回调请求：%+v", data)

	h.lock.Lock()
	defer h.lock.Unlock()

	err, finished := h.notifyHandler(c, data)
	if err != nil {
		resp.ERROR(c, err.Error())
		return
	}

	// 解除任务锁定
	if finished && (data.Progress == 100) {
		h.redis.Del(c, service.MjRunningJobKey)
	}
	resp.SUCCESS(c)

}

func (h *SdJobHandler) notifyHandler(c *gin.Context, data sdNotifyData) (error, bool) {
	taskString, err := h.redis.Get(c, service.MjRunningJobKey).Result()
	if err != nil { // 过期任务，丢弃
		logger.Warn("任务已过期：", err)
		return nil, true
	}

	var task types.SdTask
	err = utils.JsonDecode(taskString, &task)
	if err != nil { // 非标准任务，丢弃
		logger.Warn("任务解析失败：", err)
		return nil, false
	}

	var job model.SdJob
	res := h.db.Where("id = ?", task.Id).First(&job)
	if res.Error != nil {
		logger.Warn("非法任务：", res.Error)
		return nil, false
	}
	job.Params = utils.JsonEncode(task.Params)
	job.ReferenceId = data.ImageData
	job.Progress = data.Progress
	job.Prompt = data.Prompt
	job.Hash = data.Image.Hash

	// 任务完成，将最终的图片下载下来
	if data.Progress == 100 {
		imgURL, err := h.uploaderManager.GetUploadHandler().PutImg(data.Image.URL)
		if err != nil {
			logger.Error("error with download img: ", err.Error())
			return err, false
		}
		job.ImgURL = imgURL
	} else {
		// 临时图片直接保存，访问的时候使用代理进行转发
		job.ImgURL = data.Image.URL
	}
	res = h.db.Updates(&job)
	if res.Error != nil {
		logger.Error("error with update job: ", res.Error)
		return res.Error, false
	}

	var jobVo vo.MidJourneyJob
	err := utils.CopyObject(job, &jobVo)
	if err == nil {
		if data.Progress < 100 {
			image, err := utils.DownloadImage(jobVo.ImgURL, h.App.Config.ProxyURL)
			if err == nil {
				jobVo.ImgURL = "data:image/png;base64," + base64.StdEncoding.EncodeToString(image)
			}
		}

		// 推送任务到前端
		client := h.clients.Get(task.SessionId)
		if client != nil {
			utils.ReplyChunkMessage(client, jobVo)
		}
	}

	// 更新用户剩余绘图次数
	if data.Progress == 100 {
		h.db.Model(&model.User{}).Where("id = ?", task.UserId).UpdateColumn("img_calls", gorm.Expr("img_calls - ?", 1))
	}

	return nil, true
}

func (h *SdJobHandler) checkLimits(c *gin.Context) bool {
	user, err := utils.GetLoginUser(c, h.db)
	if err != nil {
		resp.NotAuth(c)
		return false
	}

	if user.ImgCalls <= 0 {
		resp.ERROR(c, "您的绘图次数不足，请联系管理员充值！")
		return false
	}

	return true

}

// Image 创建一个绘画任务
func (h *SdJobHandler) Image(c *gin.Context) {
	var data struct {
		SessionId string  `json:"session_id"`
		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
	}
	if !h.checkLimits(c) {
		return
	}

	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.Raw {
		prompt += " --style raw"
	}
	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:        int(job.Id),
		SessionId: data.SessionId,
		Src:       service.TaskSrcImg,
		Type:      service.Image,
		Prompt:    prompt,
		UserId:    userId,
	})

	var jobVo vo.MidJourneyJob
	err := utils.CopyObject(job, &jobVo)
	if err == nil {
		// 推送任务到前端
		client := h.clients.Get(data.SessionId)
		if client != nil {
			utils.ReplyChunkMessage(client, jobVo)
		}
	}
	resp.SUCCESS(c)
}

// JobList 获取 MJ 任务列表
func (h *SdJobHandler) JobList(c *gin.Context) {
	status := h.GetInt(c, "status", 0)
	var items []model.MidJourneyJob
	var res *gorm.DB
	userId, _ := c.Get(types.LoginUserID)
	if status == 1 {
		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).Order("id ASC").Find(&items)
	}
	if res.Error != nil {
		resp.ERROR(c, types.NoData)
		return
	}

	var jobs = make([]vo.MidJourneyJob, 0)
	for _, item := range items {
		var job vo.MidJourneyJob
		err := utils.CopyObject(item, &job)
		if err != nil {
			continue
		}
		if item.Progress < 100 {
			// 30 分钟还没完成的任务直接删除
			if time.Now().Sub(item.CreatedAt) > time.Minute*30 {
				h.db.Delete(&item)
				continue
			}
			if 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)
}