geekai/api/service/function/func_mj.go

65 lines
1.7 KiB
Go

package function
import (
"chatplus/service"
"chatplus/utils"
"fmt"
)
// AI 绘画函数
type FuncMidJourney struct {
name string
service *service.MjService
}
func NewMidJourneyFunc(mjService *service.MjService) FuncMidJourney {
return FuncMidJourney{
name: "MidJourney AI 绘画",
service: mjService}
}
func (f FuncMidJourney) Invoke(params map[string]interface{}) (string, error) {
logger.Infof("MJ 绘画参数:%+v", params)
prompt := utils.InterfaceToString(params["prompt"])
if !utils.IsEmptyValue(params["--ar"]) {
prompt = fmt.Sprintf("%s --ar %s", prompt, params["--ar"])
delete(params, "--ar")
}
if !utils.IsEmptyValue(params["--s"]) {
prompt = fmt.Sprintf("%s --s %s", prompt, params["--s"])
delete(params, "--s")
}
if !utils.IsEmptyValue(params["--seed"]) {
prompt = fmt.Sprintf("%s --seed %s", prompt, params["--seed"])
delete(params, "--seed")
}
if !utils.IsEmptyValue(params["--no"]) {
prompt = fmt.Sprintf("%s --no %s", prompt, params["--no"])
delete(params, "--no")
}
if !utils.IsEmptyValue(params["--niji"]) {
prompt = fmt.Sprintf("%s --niji %s", prompt, params["--niji"])
delete(params, "--niji")
} else {
prompt = prompt + " --v 5.2"
}
f.service.PushTask(service.MjTask{
Id: utils.InterfaceToString(params["session_id"]),
Src: service.TaskSrcChat,
Prompt: prompt,
UserId: utils.IntValue(utils.InterfaceToString(params["user_id"]), 0),
RoleId: utils.IntValue(utils.InterfaceToString(params["role_id"]), 0),
Icon: utils.InterfaceToString(params["icon"]),
ChatId: utils.InterfaceToString(params["chat_id"]),
})
return prompt, nil
}
func (f FuncMidJourney) Name() string {
return f.name
}
var _ Function = &FuncMidJourney{}