You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

668 rivejä
23 KiB

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