mjproxy_handler.go 22 KB

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