relay-mj.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645
  1. package relay
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "fmt"
  7. "io"
  8. "log"
  9. "net/http"
  10. "one-api/common"
  11. "one-api/constant"
  12. "one-api/dto"
  13. "one-api/model"
  14. relaycommon "one-api/relay/common"
  15. relayconstant "one-api/relay/constant"
  16. "one-api/service"
  17. "one-api/setting"
  18. "strconv"
  19. "strings"
  20. "time"
  21. "github.com/gin-gonic/gin"
  22. )
  23. func RelayMidjourneyImage(c *gin.Context) {
  24. taskId := c.Param("id")
  25. midjourneyTask := model.GetByOnlyMJId(taskId)
  26. if midjourneyTask == nil {
  27. c.JSON(400, gin.H{
  28. "error": "midjourney_task_not_found",
  29. })
  30. return
  31. }
  32. resp, err := http.Get(midjourneyTask.ImageUrl)
  33. if err != nil {
  34. c.JSON(http.StatusInternalServerError, gin.H{
  35. "error": "http_get_image_failed",
  36. })
  37. return
  38. }
  39. defer resp.Body.Close()
  40. if resp.StatusCode != http.StatusOK {
  41. responseBody, _ := io.ReadAll(resp.Body)
  42. c.JSON(resp.StatusCode, gin.H{
  43. "error": string(responseBody),
  44. })
  45. return
  46. }
  47. // 从Content-Type头获取MIME类型
  48. contentType := resp.Header.Get("Content-Type")
  49. if contentType == "" {
  50. // 如果无法确定内容类型,则默认为jpeg
  51. contentType = "image/jpeg"
  52. }
  53. // 设置响应的内容类型
  54. c.Writer.Header().Set("Content-Type", contentType)
  55. // 将图片流式传输到响应体
  56. _, err = io.Copy(c.Writer, resp.Body)
  57. if err != nil {
  58. log.Println("Failed to stream image:", err)
  59. }
  60. return
  61. }
  62. func RelayMidjourneyNotify(c *gin.Context) *dto.MidjourneyResponse {
  63. var midjRequest dto.MidjourneyDto
  64. err := common.UnmarshalBodyReusable(c, &midjRequest)
  65. if err != nil {
  66. return &dto.MidjourneyResponse{
  67. Code: 4,
  68. Description: "bind_request_body_failed",
  69. Properties: nil,
  70. Result: "",
  71. }
  72. }
  73. midjourneyTask := model.GetByOnlyMJId(midjRequest.MjId)
  74. if midjourneyTask == nil {
  75. return &dto.MidjourneyResponse{
  76. Code: 4,
  77. Description: "midjourney_task_not_found",
  78. Properties: nil,
  79. Result: "",
  80. }
  81. }
  82. midjourneyTask.Progress = midjRequest.Progress
  83. midjourneyTask.PromptEn = midjRequest.PromptEn
  84. midjourneyTask.State = midjRequest.State
  85. midjourneyTask.SubmitTime = midjRequest.SubmitTime
  86. midjourneyTask.StartTime = midjRequest.StartTime
  87. midjourneyTask.FinishTime = midjRequest.FinishTime
  88. midjourneyTask.ImageUrl = midjRequest.ImageUrl
  89. midjourneyTask.Status = midjRequest.Status
  90. midjourneyTask.FailReason = midjRequest.FailReason
  91. err = midjourneyTask.Update()
  92. if err != nil {
  93. return &dto.MidjourneyResponse{
  94. Code: 4,
  95. Description: "update_midjourney_task_failed",
  96. }
  97. }
  98. return nil
  99. }
  100. func coverMidjourneyTaskDto(c *gin.Context, originTask *model.Midjourney) (midjourneyTask dto.MidjourneyDto) {
  101. midjourneyTask.MjId = originTask.MjId
  102. midjourneyTask.Progress = originTask.Progress
  103. midjourneyTask.PromptEn = originTask.PromptEn
  104. midjourneyTask.State = originTask.State
  105. midjourneyTask.SubmitTime = originTask.SubmitTime
  106. midjourneyTask.StartTime = originTask.StartTime
  107. midjourneyTask.FinishTime = originTask.FinishTime
  108. midjourneyTask.ImageUrl = ""
  109. if originTask.ImageUrl != "" && setting.MjForwardUrlEnabled {
  110. midjourneyTask.ImageUrl = setting.ServerAddress + "/mj/image/" + originTask.MjId
  111. if originTask.Status != "SUCCESS" {
  112. midjourneyTask.ImageUrl += "?rand=" + strconv.FormatInt(time.Now().UnixNano(), 10)
  113. }
  114. } else {
  115. midjourneyTask.ImageUrl = originTask.ImageUrl
  116. }
  117. midjourneyTask.Status = originTask.Status
  118. midjourneyTask.FailReason = originTask.FailReason
  119. midjourneyTask.Action = originTask.Action
  120. midjourneyTask.Description = originTask.Description
  121. midjourneyTask.Prompt = originTask.Prompt
  122. if originTask.Buttons != "" {
  123. var buttons []dto.ActionButton
  124. err := json.Unmarshal([]byte(originTask.Buttons), &buttons)
  125. if err == nil {
  126. midjourneyTask.Buttons = buttons
  127. }
  128. }
  129. if originTask.Properties != "" {
  130. var properties dto.Properties
  131. err := json.Unmarshal([]byte(originTask.Properties), &properties)
  132. if err == nil {
  133. midjourneyTask.Properties = &properties
  134. }
  135. }
  136. return
  137. }
  138. func RelaySwapFace(c *gin.Context) *dto.MidjourneyResponse {
  139. startTime := time.Now().UnixNano() / int64(time.Millisecond)
  140. tokenId := c.GetInt("token_id")
  141. userId := c.GetInt("id")
  142. group := c.GetString("group")
  143. channelId := c.GetInt("channel_id")
  144. relayInfo := relaycommon.GenRelayInfo(c)
  145. var swapFaceRequest dto.SwapFaceRequest
  146. err := common.UnmarshalBodyReusable(c, &swapFaceRequest)
  147. if err != nil {
  148. return service.MidjourneyErrorWrapper(constant.MjRequestError, "bind_request_body_failed")
  149. }
  150. if swapFaceRequest.SourceBase64 == "" || swapFaceRequest.TargetBase64 == "" {
  151. return service.MidjourneyErrorWrapper(constant.MjRequestError, "sour_base64_and_target_base64_is_required")
  152. }
  153. modelName := service.CoverActionToModelName(constant.MjActionSwapFace)
  154. modelPrice, success := common.GetModelPrice(modelName, true)
  155. // 如果没有配置价格,则使用默认价格
  156. if !success {
  157. defaultPrice, ok := common.GetDefaultModelRatioMap()[modelName]
  158. if !ok {
  159. modelPrice = 0.1
  160. } else {
  161. modelPrice = defaultPrice
  162. }
  163. }
  164. groupRatio := setting.GetGroupRatio(group)
  165. ratio := modelPrice * groupRatio
  166. userQuota, err := model.GetUserQuota(userId, false)
  167. if err != nil {
  168. return &dto.MidjourneyResponse{
  169. Code: 4,
  170. Description: err.Error(),
  171. }
  172. }
  173. quota := int(ratio * common.QuotaPerUnit)
  174. if userQuota-quota < 0 {
  175. return &dto.MidjourneyResponse{
  176. Code: 4,
  177. Description: "quota_not_enough",
  178. }
  179. }
  180. requestURL := getMjRequestPath(c.Request.URL.String())
  181. baseURL := c.GetString("base_url")
  182. fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)
  183. mjResp, _, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL)
  184. if err != nil {
  185. return &mjResp.Response
  186. }
  187. defer func(ctx context.Context) {
  188. if mjResp.StatusCode == 200 && mjResp.Response.Code == 1 {
  189. err := service.PostConsumeQuota(relayInfo, quota, 0, true)
  190. if err != nil {
  191. common.SysError("error consuming token remain quota: " + err.Error())
  192. }
  193. //err = model.CacheUpdateUserQuota(userId)
  194. if err != nil {
  195. common.SysError("error update user quota cache: " + err.Error())
  196. }
  197. if quota != 0 {
  198. tokenName := c.GetString("token_name")
  199. logContent := fmt.Sprintf("模型固定价格 %.2f,分组倍率 %.2f,操作 %s", modelPrice, groupRatio, constant.MjActionSwapFace)
  200. other := make(map[string]interface{})
  201. other["model_price"] = modelPrice
  202. other["group_ratio"] = groupRatio
  203. model.RecordConsumeLog(ctx, userId, channelId, 0, 0, modelName, tokenName,
  204. quota, logContent, tokenId, userQuota, 0, false, group, other)
  205. model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
  206. channelId := c.GetInt("channel_id")
  207. model.UpdateChannelUsedQuota(channelId, quota)
  208. }
  209. }
  210. }(c.Request.Context())
  211. midjResponse := &mjResp.Response
  212. midjourneyTask := &model.Midjourney{
  213. UserId: userId,
  214. Code: midjResponse.Code,
  215. Action: constant.MjActionSwapFace,
  216. MjId: midjResponse.Result,
  217. Prompt: "InsightFace",
  218. PromptEn: "",
  219. Description: midjResponse.Description,
  220. State: "",
  221. SubmitTime: startTime,
  222. StartTime: time.Now().UnixNano() / int64(time.Millisecond),
  223. FinishTime: 0,
  224. ImageUrl: "",
  225. Status: "",
  226. Progress: "0%",
  227. FailReason: "",
  228. ChannelId: c.GetInt("channel_id"),
  229. Quota: quota,
  230. }
  231. err = midjourneyTask.Insert()
  232. if err != nil {
  233. return service.MidjourneyErrorWrapper(constant.MjRequestError, "insert_midjourney_task_failed")
  234. }
  235. c.Writer.WriteHeader(mjResp.StatusCode)
  236. respBody, err := json.Marshal(midjResponse)
  237. if err != nil {
  238. return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
  239. }
  240. _, err = io.Copy(c.Writer, bytes.NewBuffer(respBody))
  241. if err != nil {
  242. return service.MidjourneyErrorWrapper(constant.MjRequestError, "copy_response_body_failed")
  243. }
  244. return nil
  245. }
  246. func RelayMidjourneyTaskImageSeed(c *gin.Context) *dto.MidjourneyResponse {
  247. taskId := c.Param("id")
  248. userId := c.GetInt("id")
  249. originTask := model.GetByMJId(userId, taskId)
  250. if originTask == nil {
  251. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_no_found")
  252. }
  253. channel, err := model.GetChannelById(originTask.ChannelId, true)
  254. if err != nil {
  255. return service.MidjourneyErrorWrapper(constant.MjRequestError, "get_channel_info_failed")
  256. }
  257. if channel.Status != common.ChannelStatusEnabled {
  258. return service.MidjourneyErrorWrapper(constant.MjRequestError, "该任务所属渠道已被禁用")
  259. }
  260. c.Set("channel_id", originTask.ChannelId)
  261. c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key))
  262. requestURL := getMjRequestPath(c.Request.URL.String())
  263. fullRequestURL := fmt.Sprintf("%s%s", channel.GetBaseURL(), requestURL)
  264. midjResponseWithStatus, _, err := service.DoMidjourneyHttpRequest(c, time.Second*30, fullRequestURL)
  265. if err != nil {
  266. return &midjResponseWithStatus.Response
  267. }
  268. midjResponse := &midjResponseWithStatus.Response
  269. c.Writer.WriteHeader(midjResponseWithStatus.StatusCode)
  270. respBody, err := json.Marshal(midjResponse)
  271. if err != nil {
  272. return service.MidjourneyErrorWrapper(constant.MjRequestError, "unmarshal_response_body_failed")
  273. }
  274. _, err = io.Copy(c.Writer, bytes.NewBuffer(respBody))
  275. if err != nil {
  276. return service.MidjourneyErrorWrapper(constant.MjRequestError, "copy_response_body_failed")
  277. }
  278. return nil
  279. }
  280. func RelayMidjourneyTask(c *gin.Context, relayMode int) *dto.MidjourneyResponse {
  281. userId := c.GetInt("id")
  282. var err error
  283. var respBody []byte
  284. switch relayMode {
  285. case relayconstant.RelayModeMidjourneyTaskFetch:
  286. taskId := c.Param("id")
  287. originTask := model.GetByMJId(userId, taskId)
  288. if originTask == nil {
  289. return &dto.MidjourneyResponse{
  290. Code: 4,
  291. Description: "task_no_found",
  292. }
  293. }
  294. midjourneyTask := coverMidjourneyTaskDto(c, originTask)
  295. respBody, err = json.Marshal(midjourneyTask)
  296. if err != nil {
  297. return &dto.MidjourneyResponse{
  298. Code: 4,
  299. Description: "unmarshal_response_body_failed",
  300. }
  301. }
  302. case relayconstant.RelayModeMidjourneyTaskFetchByCondition:
  303. var condition = struct {
  304. IDs []string `json:"ids"`
  305. }{}
  306. err = c.BindJSON(&condition)
  307. if err != nil {
  308. return &dto.MidjourneyResponse{
  309. Code: 4,
  310. Description: "do_request_failed",
  311. }
  312. }
  313. var tasks []dto.MidjourneyDto
  314. if len(condition.IDs) != 0 {
  315. originTasks := model.GetByMJIds(userId, condition.IDs)
  316. for _, originTask := range originTasks {
  317. midjourneyTask := coverMidjourneyTaskDto(c, originTask)
  318. tasks = append(tasks, midjourneyTask)
  319. }
  320. }
  321. if tasks == nil {
  322. tasks = make([]dto.MidjourneyDto, 0)
  323. }
  324. respBody, err = json.Marshal(tasks)
  325. if err != nil {
  326. return &dto.MidjourneyResponse{
  327. Code: 4,
  328. Description: "unmarshal_response_body_failed",
  329. }
  330. }
  331. }
  332. c.Writer.Header().Set("Content-Type", "application/json")
  333. _, err = io.Copy(c.Writer, bytes.NewBuffer(respBody))
  334. if err != nil {
  335. return &dto.MidjourneyResponse{
  336. Code: 4,
  337. Description: "copy_response_body_failed",
  338. }
  339. }
  340. return nil
  341. }
  342. func RelayMidjourneySubmit(c *gin.Context, relayMode int) *dto.MidjourneyResponse {
  343. tokenId := c.GetInt("token_id")
  344. //channelType := c.GetInt("channel")
  345. userId := c.GetInt("id")
  346. group := c.GetString("group")
  347. channelId := c.GetInt("channel_id")
  348. relayInfo := relaycommon.GenRelayInfo(c)
  349. consumeQuota := true
  350. var midjRequest dto.MidjourneyRequest
  351. err := common.UnmarshalBodyReusable(c, &midjRequest)
  352. if err != nil {
  353. return service.MidjourneyErrorWrapper(constant.MjRequestError, "bind_request_body_failed")
  354. }
  355. if relayMode == relayconstant.RelayModeMidjourneyAction { // midjourney plus,需要从customId中获取任务信息
  356. mjErr := service.CoverPlusActionToNormalAction(&midjRequest)
  357. if mjErr != nil {
  358. return mjErr
  359. }
  360. relayMode = relayconstant.RelayModeMidjourneyChange
  361. }
  362. if relayMode == relayconstant.RelayModeMidjourneyImagine { //绘画任务,此类任务可重复
  363. if midjRequest.Prompt == "" {
  364. return service.MidjourneyErrorWrapper(constant.MjRequestError, "prompt_is_required")
  365. }
  366. midjRequest.Action = constant.MjActionImagine
  367. } else if relayMode == relayconstant.RelayModeMidjourneyDescribe { //按图生文任务,此类任务可重复
  368. midjRequest.Action = constant.MjActionDescribe
  369. } else if relayMode == relayconstant.RelayModeMidjourneyShorten { //缩短任务,此类任务可重复,plus only
  370. midjRequest.Action = constant.MjActionShorten
  371. } else if relayMode == relayconstant.RelayModeMidjourneyBlend { //绘画任务,此类任务可重复
  372. midjRequest.Action = constant.MjActionBlend
  373. } else if relayMode == relayconstant.RelayModeMidjourneyUpload { //绘画任务,此类任务可重复
  374. midjRequest.Action = constant.MjActionUpload
  375. } else if midjRequest.TaskId != "" { //放大、变换任务,此类任务,如果重复且已有结果,远端api会直接返回最终结果
  376. mjId := ""
  377. if relayMode == relayconstant.RelayModeMidjourneyChange {
  378. if midjRequest.TaskId == "" {
  379. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_id_is_required")
  380. } else if midjRequest.Action == "" {
  381. return service.MidjourneyErrorWrapper(constant.MjRequestError, "action_is_required")
  382. } else if midjRequest.Index == 0 {
  383. return service.MidjourneyErrorWrapper(constant.MjRequestError, "index_is_required")
  384. }
  385. //action = midjRequest.Action
  386. mjId = midjRequest.TaskId
  387. } else if relayMode == relayconstant.RelayModeMidjourneySimpleChange {
  388. if midjRequest.Content == "" {
  389. return service.MidjourneyErrorWrapper(constant.MjRequestError, "content_is_required")
  390. }
  391. params := service.ConvertSimpleChangeParams(midjRequest.Content)
  392. if params == nil {
  393. return service.MidjourneyErrorWrapper(constant.MjRequestError, "content_parse_failed")
  394. }
  395. mjId = params.TaskId
  396. midjRequest.Action = params.Action
  397. } else if relayMode == relayconstant.RelayModeMidjourneyModal {
  398. //if midjRequest.MaskBase64 == "" {
  399. // return service.MidjourneyErrorWrapper(constant.MjRequestError, "mask_base64_is_required")
  400. //}
  401. mjId = midjRequest.TaskId
  402. midjRequest.Action = constant.MjActionModal
  403. }
  404. originTask := model.GetByMJId(userId, mjId)
  405. if originTask == nil {
  406. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_not_found")
  407. } else { //原任务的Status=SUCCESS,则可以做放大UPSCALE、变换VARIATION等动作,此时必须使用原来的请求地址才能正确处理
  408. if setting.MjActionCheckSuccessEnabled {
  409. if originTask.Status != "SUCCESS" && relayMode != relayconstant.RelayModeMidjourneyModal {
  410. return service.MidjourneyErrorWrapper(constant.MjRequestError, "task_status_not_success")
  411. }
  412. }
  413. channel, err := model.GetChannelById(originTask.ChannelId, true)
  414. if err != nil {
  415. return service.MidjourneyErrorWrapper(constant.MjRequestError, "get_channel_info_failed")
  416. }
  417. if channel.Status != common.ChannelStatusEnabled {
  418. return service.MidjourneyErrorWrapper(constant.MjRequestError, "该任务所属渠道已被禁用")
  419. }
  420. c.Set("base_url", channel.GetBaseURL())
  421. c.Set("channel_id", originTask.ChannelId)
  422. c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key))
  423. log.Printf("检测到此操作为放大、变换、重绘,获取原channel信息: %s,%s", strconv.Itoa(originTask.ChannelId), channel.GetBaseURL())
  424. }
  425. midjRequest.Prompt = originTask.Prompt
  426. //if channelType == common.ChannelTypeMidjourneyPlus {
  427. // // plus
  428. //} else {
  429. // // 普通版渠道
  430. //
  431. //}
  432. }
  433. if midjRequest.Action == constant.MjActionInPaint || midjRequest.Action == constant.MjActionCustomZoom {
  434. consumeQuota = false
  435. }
  436. //baseURL := common.ChannelBaseURLs[channelType]
  437. requestURL := getMjRequestPath(c.Request.URL.String())
  438. baseURL := c.GetString("base_url")
  439. //midjRequest.NotifyHook = "http://127.0.0.1:3000/mj/notify"
  440. fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)
  441. modelName := service.CoverActionToModelName(midjRequest.Action)
  442. modelPrice, success := common.GetModelPrice(modelName, true)
  443. // 如果没有配置价格,则使用默认价格
  444. if !success {
  445. defaultPrice, ok := common.GetDefaultModelRatioMap()[modelName]
  446. if !ok {
  447. modelPrice = 0.1
  448. } else {
  449. modelPrice = defaultPrice
  450. }
  451. }
  452. groupRatio := setting.GetGroupRatio(group)
  453. ratio := modelPrice * groupRatio
  454. userQuota, err := model.GetUserQuota(userId, false)
  455. if err != nil {
  456. return &dto.MidjourneyResponse{
  457. Code: 4,
  458. Description: err.Error(),
  459. }
  460. }
  461. quota := int(ratio * common.QuotaPerUnit)
  462. if consumeQuota && userQuota-quota < 0 {
  463. return &dto.MidjourneyResponse{
  464. Code: 4,
  465. Description: "quota_not_enough",
  466. }
  467. }
  468. midjResponseWithStatus, responseBody, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL)
  469. if err != nil {
  470. return &midjResponseWithStatus.Response
  471. }
  472. midjResponse := &midjResponseWithStatus.Response
  473. defer func(ctx context.Context) {
  474. if consumeQuota && midjResponseWithStatus.StatusCode == 200 {
  475. err := service.PostConsumeQuota(relayInfo, quota, 0, true)
  476. if err != nil {
  477. common.SysError("error consuming token remain quota: " + err.Error())
  478. }
  479. if quota != 0 {
  480. tokenName := c.GetString("token_name")
  481. logContent := fmt.Sprintf("模型固定价格 %.2f,分组倍率 %.2f,操作 %s,ID %s", modelPrice, groupRatio, midjRequest.Action, midjResponse.Result)
  482. other := make(map[string]interface{})
  483. other["model_price"] = modelPrice
  484. other["group_ratio"] = groupRatio
  485. model.RecordConsumeLog(ctx, userId, channelId, 0, 0, modelName, tokenName,
  486. quota, logContent, tokenId, userQuota, 0, false, group, other)
  487. model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
  488. channelId := c.GetInt("channel_id")
  489. model.UpdateChannelUsedQuota(channelId, quota)
  490. }
  491. }
  492. }(c.Request.Context())
  493. // 文档:https://github.com/novicezk/midjourney-proxy/blob/main/docs/api.md
  494. //1-提交成功
  495. // 21-任务已存在(处理中或者有结果了) {"code":21,"description":"任务已存在","result":"0741798445574458","properties":{"status":"SUCCESS","imageUrl":"https://xxxx"}}
  496. // 22-排队中 {"code":22,"description":"排队中,前面还有1个任务","result":"0741798445574458","properties":{"numberOfQueues":1,"discordInstanceId":"1118138338562560102"}}
  497. // 23-队列已满,请稍后再试 {"code":23,"description":"队列已满,请稍后尝试","result":"14001929738841620","properties":{"discordInstanceId":"1118138338562560102"}}
  498. // 24-prompt包含敏感词 {"code":24,"description":"可能包含敏感词","properties":{"promptEn":"nude body","bannedWord":"nude"}}
  499. // other: 提交错误,description为错误描述
  500. midjourneyTask := &model.Midjourney{
  501. UserId: userId,
  502. Code: midjResponse.Code,
  503. Action: midjRequest.Action,
  504. MjId: midjResponse.Result,
  505. Prompt: midjRequest.Prompt,
  506. PromptEn: "",
  507. Description: midjResponse.Description,
  508. State: "",
  509. SubmitTime: time.Now().UnixNano() / int64(time.Millisecond),
  510. StartTime: 0,
  511. FinishTime: 0,
  512. ImageUrl: "",
  513. Status: "",
  514. Progress: "0%",
  515. FailReason: "",
  516. ChannelId: c.GetInt("channel_id"),
  517. Quota: quota,
  518. }
  519. if midjResponse.Code == 3 {
  520. //无实例账号自动禁用渠道(No available account instance)
  521. channel, err := model.GetChannelById(midjourneyTask.ChannelId, true)
  522. if err != nil {
  523. common.SysError("get_channel_null: " + err.Error())
  524. }
  525. if channel.GetAutoBan() && common.AutomaticDisableChannelEnabled {
  526. model.UpdateChannelStatusById(midjourneyTask.ChannelId, 2, "No available account instance")
  527. }
  528. }
  529. if midjResponse.Code != 1 && midjResponse.Code != 21 && midjResponse.Code != 22 {
  530. //非1-提交成功,21-任务已存在和22-排队中,则记录错误原因
  531. midjourneyTask.FailReason = midjResponse.Description
  532. consumeQuota = false
  533. }
  534. if midjResponse.Code == 21 { //21-任务已存在(处理中或者有结果了)
  535. // 将 properties 转换为一个 map
  536. properties, ok := midjResponse.Properties.(map[string]interface{})
  537. if ok {
  538. imageUrl, ok1 := properties["imageUrl"].(string)
  539. status, ok2 := properties["status"].(string)
  540. if ok1 && ok2 {
  541. midjourneyTask.ImageUrl = imageUrl
  542. midjourneyTask.Status = status
  543. if status == "SUCCESS" {
  544. midjourneyTask.Progress = "100%"
  545. midjourneyTask.StartTime = time.Now().UnixNano() / int64(time.Millisecond)
  546. midjourneyTask.FinishTime = time.Now().UnixNano() / int64(time.Millisecond)
  547. midjResponse.Code = 1
  548. }
  549. }
  550. }
  551. //修改返回值
  552. if midjRequest.Action != constant.MjActionInPaint && midjRequest.Action != constant.MjActionCustomZoom {
  553. newBody := strings.Replace(string(responseBody), `"code":21`, `"code":1`, -1)
  554. responseBody = []byte(newBody)
  555. }
  556. }
  557. if midjResponse.Code == 1 && midjRequest.Action == "UPLOAD" {
  558. midjourneyTask.Progress = "100%"
  559. midjourneyTask.Status = "SUCCESS"
  560. }
  561. err = midjourneyTask.Insert()
  562. if err != nil {
  563. return &dto.MidjourneyResponse{
  564. Code: 4,
  565. Description: "insert_midjourney_task_failed",
  566. }
  567. }
  568. if midjResponse.Code == 22 { //22-排队中,说明任务已存在
  569. //修改返回值
  570. newBody := strings.Replace(string(responseBody), `"code":22`, `"code":1`, -1)
  571. responseBody = []byte(newBody)
  572. }
  573. //resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
  574. bodyReader := io.NopCloser(bytes.NewBuffer(responseBody))
  575. //for k, v := range resp.Header {
  576. // c.Writer.Header().Set(k, v[0])
  577. //}
  578. c.Writer.WriteHeader(midjResponseWithStatus.StatusCode)
  579. _, err = io.Copy(c.Writer, bodyReader)
  580. if err != nil {
  581. return &dto.MidjourneyResponse{
  582. Code: 4,
  583. Description: "copy_response_body_failed",
  584. }
  585. }
  586. err = bodyReader.Close()
  587. if err != nil {
  588. return &dto.MidjourneyResponse{
  589. Code: 4,
  590. Description: "close_response_body_failed",
  591. }
  592. }
  593. return nil
  594. }
  595. type taskChangeParams struct {
  596. ID string
  597. Action string
  598. Index int
  599. }
  600. func getMjRequestPath(path string) string {
  601. requestURL := path
  602. if strings.Contains(requestURL, "/mj-") {
  603. urls := strings.Split(requestURL, "/mj/")
  604. if len(urls) < 2 {
  605. return requestURL
  606. }
  607. requestURL = "/mj/" + urls[1]
  608. }
  609. return requestURL
  610. }