浏览代码

fix SUBMITTED状态返回预估时间

zwhui 2 年之前
父节点
当前提交
9c06c0eabe
共有 1 个文件被更改,包括 15 次插入1 次删除
  1. 15 1
      midjourney/src/main/java/com/yhlxj/service/midjourney/impl/MidjourneyServiceImpl.java

+ 15 - 1
midjourney/src/main/java/com/yhlxj/service/midjourney/impl/MidjourneyServiceImpl.java

@@ -880,8 +880,22 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         String content = jsonObject.getStr("messageContent");
         if(StringUtils.isNotBlank(content)){
             jsonObject.set("messageContent","**"+conversation.getPrompt()+"** - <@mj>");
-            conversation.setProperties(jsonObject.toString());
         }
+        if ("SUBMITTED".equals(conversation.getStatus())){
+            MidjourneyUserConversation midjourneyUserConversation = conversationMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUserConversation.class).eq(MidjourneyUserConversation::getTaskId, conversation.getTaskId()).last("limit 1"));
+            String mjVersionStatusKey;
+            if (midjourneyUserConversation != null){
+                if (midjourneyUserConversation.getMode() == 1){
+                    mjVersionStatusKey = "v6-fast-";
+                }else {
+                    mjVersionStatusKey = "v6-relax-";
+                }
+                mjVersionStatusKey += (List.of("UPSCALE","VARIATION").contains(midjourneyUserConversation.getAction()) ? "upscaler" : "model");
+                MjVersionStatus mjVersionStatus = MJ_VERSION_STATUS_MAP.get(mjVersionStatusKey);
+                jsonObject.set("estimateTime", mjVersionStatus == null ? null : System.currentTimeMillis() + mjVersionStatus.getMillisecond());
+            }
+        }
+        conversation.setProperties(jsonObject.toString());
         if (StringUtil.isNotBlank(conversation.getImageUrl())) {
             log.info("获取到的图片 url:{}", conversation.getImageUrl());
             ImgSize imgSize = TASK_IMG_SIZE_MAP.get(conversation.getTaskId());