Skip to content

phsl / new-api

Go to a project
Commit 6ebf1b4f by CaIon
Options

fix mj bug

parent 833c4779
Showing with 139 additions and 14 deletions
controller/midjourney.go
...@@ -2,14 +2,17 @@ package controller ...@@ -2,14 +2,17 @@ package controller
import ( import (
"bytes" "bytes"
"context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"io"
"log" "log"
"net/http" "net/http"
"one-api/common" "one-api/common"
"one-api/model" "one-api/model"
"strconv" "strconv"
"strings"
"time" "time"
) )
...@@ -25,7 +28,9 @@ func UpdateMidjourneyTask() { ...@@ -25,7 +28,9 @@ func UpdateMidjourneyTask() {
time.Sleep(time.Duration(15) * time.Second) time.Sleep(time.Duration(15) * time.Second)
tasks := model.GetAllUnFinishTasks() tasks := model.GetAllUnFinishTasks()
if len(tasks) != 0 { if len(tasks) != 0 {
log.Printf("检测到未完成的任务数有: %v", len(tasks))
for _, task := range tasks { for _, task := range tasks {
log.Printf("未完成的任务信息: %v", task)
midjourneyChannel, err := model.GetChannelById(task.ChannelId, true) midjourneyChannel, err := model.GetChannelById(task.ChannelId, true)
if err != nil { if err != nil {
log.Printf("UpdateMidjourneyTask: %v", err) log.Printf("UpdateMidjourneyTask: %v", err)
...@@ -39,6 +44,7 @@ func UpdateMidjourneyTask() { ...@@ -39,6 +44,7 @@ func UpdateMidjourneyTask() {
continue continue
} }
requestUrl := fmt.Sprintf("%s/mj/task/%s/fetch", *midjourneyChannel.BaseURL, task.MjId) requestUrl := fmt.Sprintf("%s/mj/task/%s/fetch", *midjourneyChannel.BaseURL, task.MjId)
log.Printf("requestUrl: %s", requestUrl)
req, err := http.NewRequest("GET", requestUrl, bytes.NewBuffer([]byte(""))) req, err := http.NewRequest("GET", requestUrl, bytes.NewBuffer([]byte("")))
if err != nil { if err != nil {
...@@ -46,7 +52,16 @@ func UpdateMidjourneyTask() { ...@@ -46,7 +52,16 @@ func UpdateMidjourneyTask() {
continue continue
} }
// 设置超时时间
timeout := time.Second * 5
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
// 使用带有超时的 context 创建新的请求
req = req.WithContext(ctx)
req.Header.Set("Content-Type", "application/json") req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer midjourney-proxy")
req.Header.Set("mj-api-secret", midjourneyChannel.Key) req.Header.Set("mj-api-secret", midjourneyChannel.Key)
resp, err := httpClient.Do(req) resp, err := httpClient.Do(req)
if err != nil { if err != nil {
...@@ -54,12 +69,38 @@ func UpdateMidjourneyTask() { ...@@ -54,12 +69,38 @@ func UpdateMidjourneyTask() {
continue continue
} }
defer resp.Body.Close() defer resp.Body.Close()
responseBody, err := io.ReadAll(resp.Body)
log.Printf("responseBody: %s", string(responseBody))
var responseItem Midjourney var responseItem Midjourney
err = json.NewDecoder(resp.Body).Decode(&responseItem) // err = json.NewDecoder(resp.Body).Decode(&responseItem)
err = json.Unmarshal(responseBody, &responseItem)
if err != nil { if err != nil {
log.Printf("UpdateMidjourneyTask error: %v", err) if strings.Contains(err.Error(), "cannot unmarshal number into Go struct field Midjourney.status of type string") {
var responseWithoutStatus MidjourneyWithoutStatus
var responseStatus MidjourneyStatus
err1 := json.Unmarshal(responseBody, &responseWithoutStatus)
err2 := json.Unmarshal(responseBody, &responseStatus)
if err1 == nil && err2 == nil {
jsonData, err3 := json.Marshal(responseWithoutStatus)
if err3 != nil {
log.Fatalf("UpdateMidjourneyTask error1: %v", err3)
continue continue
} }
err4 := json.Unmarshal(jsonData, &responseStatus)
if err4 != nil {
log.Fatalf("UpdateMidjourneyTask error2: %v", err4)
continue
}
responseItem.Status = strconv.Itoa(responseStatus.Status)
} else {
log.Printf("UpdateMidjourneyTask error3: %v", err)
continue
}
} else {
log.Printf("UpdateMidjourneyTask error4: %v", err)
continue
}
}
task.Code = 1 task.Code = 1
task.Progress = responseItem.Progress task.Progress = responseItem.Progress
task.PromptEn = responseItem.PromptEn task.PromptEn = responseItem.PromptEn
...@@ -94,7 +135,7 @@ func UpdateMidjourneyTask() { ...@@ -94,7 +135,7 @@ func UpdateMidjourneyTask() {
err = task.Update() err = task.Update()
if err != nil { if err != nil {
log.Printf("UpdateMidjourneyTask error: %v", err) log.Printf("UpdateMidjourneyTask error5: %v", err)
} }
log.Printf("UpdateMidjourneyTask success: %v", task) log.Printf("UpdateMidjourneyTask success: %v", task)
} }
......
controller/relay-mj.go
...@@ -12,6 +12,7 @@ import ( ...@@ -12,6 +12,7 @@ import (
"one-api/model" "one-api/model"
"strconv" "strconv"
"strings" "strings"
"time"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
) )
...@@ -32,6 +33,28 @@ type Midjourney struct { ...@@ -32,6 +33,28 @@ type Midjourney struct {
FailReason string `json:"failReason"` FailReason string `json:"failReason"`
} }
type MidjourneyStatus struct {
Status int `json:"status"`
}
type MidjourneyWithoutStatus struct {
Id int `json:"id"`
Code int `json:"code"`
UserId int `json:"user_id" gorm:"index"`
Action string `json:"action"`
MjId string `json:"mj_id" gorm:"index"`
Prompt string `json:"prompt"`
PromptEn string `json:"prompt_en"`
Description string `json:"description"`
State string `json:"state"`
SubmitTime int64 `json:"submit_time"`
StartTime int64 `json:"start_time"`
FinishTime int64 `json:"finish_time"`
ImageUrl string `json:"image_url"`
Progress string `json:"progress"`
FailReason string `json:"fail_reason"`
ChannelId int `json:"channel_id"`
}
func RelayMidjourneyImage(c *gin.Context) { func RelayMidjourneyImage(c *gin.Context) {
taskId := c.Param("id") taskId := c.Param("id")
midjourneyTask := model.GetByMJId(taskId) midjourneyTask := model.GetByMJId(taskId)
...@@ -115,7 +138,13 @@ func relayMidjourneyTask(c *gin.Context, relayMode int) *MidjourneyResponse { ...@@ -115,7 +138,13 @@ func relayMidjourneyTask(c *gin.Context, relayMode int) *MidjourneyResponse {
midjourneyTask.SubmitTime = originTask.SubmitTime midjourneyTask.SubmitTime = originTask.SubmitTime
midjourneyTask.StartTime = originTask.StartTime midjourneyTask.StartTime = originTask.StartTime
midjourneyTask.FinishTime = originTask.FinishTime midjourneyTask.FinishTime = originTask.FinishTime
midjourneyTask.ImageUrl = ""
if originTask.ImageUrl != "" {
midjourneyTask.ImageUrl = common.ServerAddress + "/mj/image/" + originTask.MjId midjourneyTask.ImageUrl = common.ServerAddress + "/mj/image/" + originTask.MjId
if originTask.Status != "SUCCESS" {
midjourneyTask.ImageUrl += "?rand=" + strconv.FormatInt(time.Now().UnixNano(), 10)
}
}
midjourneyTask.Status = originTask.Status midjourneyTask.Status = originTask.Status
midjourneyTask.FailReason = originTask.FailReason midjourneyTask.FailReason = originTask.FailReason
midjourneyTask.Action = originTask.Action midjourneyTask.Action = originTask.Action
...@@ -157,7 +186,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse { ...@@ -157,7 +186,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
} }
} }
} }
if relayMode == RelayModeMidjourneyImagine { if relayMode == RelayModeMidjourneyImagine { //绘画任务，此类任务可重复
if midjRequest.Prompt == "" { if midjRequest.Prompt == "" {
return &MidjourneyResponse{ return &MidjourneyResponse{
Code: 4, Code: 4,
...@@ -165,7 +194,11 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse { ...@@ -165,7 +194,11 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
} }
} }
midjRequest.Action = "IMAGINE" midjRequest.Action = "IMAGINE"
} else if midjRequest.TaskId != "" { } else if relayMode == RelayModeMidjourneyDescribe { //按图生文任务，此类任务可重复
midjRequest.Action = "DESCRIBE"
} else if relayMode == RelayModeMidjourneyBlend { //绘画任务，此类任务可重复
midjRequest.Action = "BLEND"
} else if midjRequest.TaskId != "" { //放大、变换任务，此类任务，如果重复且已有结果，远端api会直接返回最终结果
originTask := model.GetByMJId(midjRequest.TaskId) originTask := model.GetByMJId(midjRequest.TaskId)
if originTask == nil { if originTask == nil {
return &MidjourneyResponse{ return &MidjourneyResponse{
...@@ -183,7 +216,17 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse { ...@@ -183,7 +216,17 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
Code: 4, Code: 4,
Description: "task_status_is_not_success", Description: "task_status_is_not_success",
} }
} else { //原任务的Status=SUCCESS，则可以做放大UPSCALE、变换VARIATION等动作，此时必须使用原来的请求地址才能正确处理
channel, err := model.GetChannelById(originTask.ChannelId, false)
if err != nil {
return &MidjourneyResponse{
Code: 4,
Description: "channel_not_found",
}
}
c.Set("base_url", channel.GetBaseURL())
c.Set("channel_id", originTask.ChannelId)
log.Printf("检测到此操作为放大、变换，获取原channel信息: %s,%s", strconv.Itoa(originTask.ChannelId), channel.GetBaseURL())
} }
midjRequest.Prompt = originTask.Prompt midjRequest.Prompt = originTask.Prompt
} else if relayMode == RelayModeMidjourneyChange { } else if relayMode == RelayModeMidjourneyChange {
...@@ -234,6 +277,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse { ...@@ -234,6 +277,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
//midjRequest.NotifyHook = "http://127.0.0.1:3000/mj/notify" //midjRequest.NotifyHook = "http://127.0.0.1:3000/mj/notify"
fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL) fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)
log.Printf("fullRequestURL: %s", fullRequestURL)
var requestBody io.Reader var requestBody io.Reader
if isModelMapped { if isModelMapped {
...@@ -275,14 +319,15 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse { ...@@ -275,14 +319,15 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
Description: "create_request_failed", Description: "create_request_failed",
} }
} }
//req.HeaderBar.Set("Authorization", c.Request.HeaderBar.Get("Authorization")) //req.Header.Set("Authorization", c.Request.Header.Get("Authorization"))
req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type")) req.Header.Set("Content-Type", c.Request.Header.Get("Content-Type"))
req.Header.Set("Accept", c.Request.Header.Get("Accept")) req.Header.Set("Accept", c.Request.Header.Get("Accept"))
//mjToken := "" //mjToken := ""
//if c.Request.HeaderBar.Get("Authorization") != "" { //if c.Request.Header.Get("Authorization") != "" {
// mjToken = strings.Split(c.Request.HeaderBar.Get("Authorization"), " ")[1] // mjToken = strings.Split(c.Request.Header.Get("Authorization"), " ")[1]
//} //}
req.Header.Set("Authorization", "Bearer midjourney-proxy")
req.Header.Set("mj-api-secret", strings.Split(c.Request.Header.Get("Authorization"), " ")[1]) req.Header.Set("mj-api-secret", strings.Split(c.Request.Header.Get("Authorization"), " ")[1])
// print request header // print request header
log.Printf("request header: %s", req.Header) log.Printf("request header: %s", req.Header)
...@@ -367,10 +412,14 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse { ...@@ -367,10 +412,14 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
Description: "unmarshal_response_body_failed", Description: "unmarshal_response_body_failed",
} }
} }
if midjResponse.Code == 24 || midjResponse.Code == 21 || midjResponse.Code == 4 {
consumeQuota = false
}
// 文档：https://github.com/novicezk/midjourney-proxy/blob/main/docs/api.md
//1-提交成功
// 21-任务已存在（处理中或者有结果了） {"code":21,"description":"任务已存在","result":"0741798445574458","properties":{"status":"SUCCESS","imageUrl":"https://xxxx"}}
// 22-排队中 {"code":22,"description":"排队中，前面还有1个任务","result":"0741798445574458","properties":{"numberOfQueues":1,"discordInstanceId":"1118138338562560102"}}
// 23-队列已满，请稍后再试 {"code":23,"description":"队列已满，请稍后尝试","result":"14001929738841620","properties":{"discordInstanceId":"1118138338562560102"}}
// 24-prompt包含敏感词 {"code":24,"description":"可能包含敏感词","properties":{"promptEn":"nude body","bannedWord":"nude"}}
// other: 提交错误，description为错误描述
midjourneyTask := &model.Midjourney{ midjourneyTask := &model.Midjourney{
UserId: userId, UserId: userId,
Code: midjResponse.Code, Code: midjResponse.Code,
...@@ -380,7 +429,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse { ...@@ -380,7 +429,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
PromptEn: "", PromptEn: "",
Description: midjResponse.Description, Description: midjResponse.Description,
State: "", State: "",
SubmitTime: 0, SubmitTime: time.Now().UnixNano() / int64(time.Millisecond),
StartTime: 0, StartTime: 0,
FinishTime: 0, FinishTime: 0,
ImageUrl: "", ImageUrl: "",
...@@ -389,9 +438,35 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse { ...@@ -389,9 +438,35 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
FailReason: "", FailReason: "",
ChannelId: c.GetInt("channel_id"), ChannelId: c.GetInt("channel_id"),
} }
if midjResponse.Code == 4 || midjResponse.Code == 24 {
if midjResponse.Code != 1 && midjResponse.Code != 21 && midjResponse.Code != 22 {
//非1-提交成功,21-任务已存在和22-排队中，则记录错误原因
midjourneyTask.FailReason = midjResponse.Description midjourneyTask.FailReason = midjResponse.Description
consumeQuota = false
} }
if midjResponse.Code == 21 { //21-任务已存在（处理中或者有结果了）
// 将 properties 转换为一个 map
properties, ok := midjResponse.Properties.(map[string]interface{})
if ok {
imageUrl, ok1 := properties["imageUrl"].(string)
status, ok2 := properties["status"].(string)
if ok1 && ok2 {
midjourneyTask.ImageUrl = imageUrl
midjourneyTask.Status = status
if status == "SUCCESS" {
midjourneyTask.Progress = "100%"
midjourneyTask.StartTime = time.Now().UnixNano() / int64(time.Millisecond)
midjourneyTask.FinishTime = time.Now().UnixNano() / int64(time.Millisecond)
midjResponse.Code = 1
}
}
}
//修改返回值
newBody := strings.Replace(string(responseBody), `"code":21`, `"code":1`, -1)
responseBody = []byte(newBody)
}
err = midjourneyTask.Insert() err = midjourneyTask.Insert()
if err != nil { if err != nil {
return &MidjourneyResponse{ return &MidjourneyResponse{
...@@ -399,6 +474,13 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse { ...@@ -399,6 +474,13 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
Description: "insert_midjourney_task_failed", Description: "insert_midjourney_task_failed",
} }
} }
if midjResponse.Code == 22 { //22-排队中，说明任务已存在
//修改返回值
newBody := strings.Replace(string(responseBody), `"code":22`, `"code":1`, -1)
responseBody = []byte(newBody)
}
resp.Body = io.NopCloser(bytes.NewBuffer(responseBody)) resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
for k, v := range resp.Header { for k, v := range resp.Header {
......
controller/relay.go
...@@ -26,6 +26,8 @@ const ( ...@@ -26,6 +26,8 @@ const (
RelayModeImagesGenerations RelayModeImagesGenerations
RelayModeEdits RelayModeEdits
RelayModeMidjourneyImagine RelayModeMidjourneyImagine
RelayModeMidjourneyDescribe
RelayModeMidjourneyBlend
RelayModeMidjourneyChange RelayModeMidjourneyChange
RelayModeMidjourneyNotify RelayModeMidjourneyNotify
RelayModeMidjourneyTaskFetch RelayModeMidjourneyTaskFetch
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment