Procházet zdrojové kódy

fix 返回预计生成时间

zwhui před 2 roky
rodič
revize
e6922bc739

+ 40 - 0
midjourney/src/main/java/com/yhlxj/dao/model/dto/MjVersionStatus.java

@@ -0,0 +1,40 @@
+package com.yhlxj.dao.model.dto;
+
+import lombok.Data;
+
+@Data
+public class MjVersionStatus {
+    //{
+    //    "name": "jobs.time_to_start.fast.10min.job_type_kdpt_diffusion_anime",
+    //        "date": "2024-07-11T03:24:01.784Z",
+    //        "value": "N/A"
+    //}
+
+    private String name;
+
+    private String date;
+
+    private String value;
+    /**
+     * 转换为毫秒
+     */
+    private Long millisecond;
+
+
+    public void setValue(String value) {
+        this.value = value;
+        // 转换为秒并设置 second
+        if (value != null && !value.equals("N/A")) {
+            try {
+                double minutes = Double.parseDouble(value);
+                this.millisecond = Math.round(minutes * 60 * 1000);
+            } catch (NumberFormatException e) {
+                this.millisecond = null;
+                System.err.println("Invalid value format: " + value);
+            }
+        } else {
+            this.millisecond = null;
+        }
+    }
+
+}

+ 63 - 25
midjourney/src/main/java/com/yhlxj/service/midjourney/impl/MidjourneyServiceImpl.java

@@ -22,6 +22,7 @@ import com.yhlxj.dao.mapper.midjourney.MidjourneyUserConversationMapper;
 import com.yhlxj.dao.mapper.midjourney.MidjourneyUserMapper;
 import com.yhlxj.dao.model.dto.BlendDimensions;
 import com.yhlxj.dao.model.dto.MessageButton;
+import com.yhlxj.dao.model.dto.MjVersionStatus;
 import com.yhlxj.dao.model.dto.SubmitUploadDTO;
 import com.yhlxj.dao.model.entity.*;
 import com.yhlxj.dao.model.enums.WsMessageTypeEnum;
@@ -107,6 +108,9 @@ public class MidjourneyServiceImpl implements MidjourneyService {
 
     private final MidjourneyApiChannelMapper midjourneyApiChannelMapper;
 
+    public static final Map<String, MjVersionStatus> MJ_VERSION_STATUS_MAP = new ConcurrentHashMap<>();
+
+
     @Override
     public MidjourneyUser getUser(String userToken) {
         MidjourneyUser midjourneyUser = midjourneyUserMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUser.class)
@@ -153,7 +157,18 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         }
         saveConversation(user.getId(), mode,result,"IMAGINE", StringUtils.EMPTY, botType,result.getApiChannelId());
         Map<String, Object> properties = result.getProperties();
-        properties.put("messageContent","**"+prompt+"** - <@mj>" + (mode == 1 ? "(fast)" : "(relaxed)"));
+        String modeStr;
+        String mjVersionStatusKey;
+        if (mode == 1){
+            modeStr = "(fast)";
+            mjVersionStatusKey = "v6-fast-model";
+        }else {
+            modeStr = "(relaxed)";
+            mjVersionStatusKey = "v6-relax-model";
+        }
+        properties.put("messageContent","**"+prompt+"** - <@mj>" + modeStr);
+        MjVersionStatus mjVersionStatus = MJ_VERSION_STATUS_MAP.get(mjVersionStatusKey);
+        properties.put("estimateTime", mjVersionStatus == null ? null : System.currentTimeMillis() + mjVersionStatus.getMillisecond());
         return result;
     }
 
@@ -224,15 +239,15 @@ public class MidjourneyServiceImpl implements MidjourneyService {
             redisService.hincr(RedisService.key.MIDJOURNEY_ACCOUNT.getName(), result.getInstanceId().toString(), 1.0);
         }
         saveConversation(user.getId(), mode,result,"BLEND", StringUtils.EMPTY, botType, result.getApiChannelId());
+        Map<String, Object> properties = result.getProperties();
+        MjVersionStatus mjVersionStatus = MJ_VERSION_STATUS_MAP.get(mode == 1 ? "v6-fast-model" : "v6-relax-model");
+        properties.put("estimateTime", mjVersionStatus == null ? null : System.currentTimeMillis() + mjVersionStatus.getMillisecond());
         return result;
     }
 
     @Override
     public SubmitResult submitModal(MidjourneyUser user,Integer mode, Long taskId, String prompt, String maskBase64) throws Exception {
         MidjourneyUserConversation midjourneyUserConversation = conversationMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUserConversation.class).eq(MidjourneyUserConversation::getTaskId, taskId).last("limit 1"));
-        if (midjourneyUserConversation == null) {
-            throw BusinessRuntimeException.getInstance("关联任务不存在或已失效");
-        }
         Map<String, Object> param = MapUtil.builder(new HashMap<String,Object>())
                 .put("taskId", taskId.toString())
                 .put("state", user.getId().toString())
@@ -243,14 +258,24 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         if (StringUtils.isNotBlank(maskBase64)) {
             param.put("maskBase64", maskBase64);
         }
-        SubmitResult result = submit(user,mode,"modal", null, param,midjourneyUserConversation.getApiChannelId());
-        if (mode == 2) {
+        Long apiChannelId = midjourneyUserConversation == null ? null : midjourneyUserConversation.getApiChannelId();
+        SubmitResult result = submit(user,mode,"modal", null, param, apiChannelId);
+        saveConversation(user.getId(),  mode,result,"MODAL", StringUtils.EMPTY, StringUtils.EMPTY,apiChannelId);
+        Map<String, Object> properties = result.getProperties();
+        String modeStr;
+        String mjVersionStatusKey;
+        if (mode == 1){
+            modeStr = "(fast)";
+            mjVersionStatusKey = "v6-fast-model";
+        }else {
+            modeStr = "(relaxed)";
+            mjVersionStatusKey = "v6-relax-model";
             redisService.hincr(RedisService.key.MIDJOURNEY_ACCOUNT.getName(), result.getInstanceId().toString(), 1.0);
         }
-        saveConversation(user.getId(),  mode,result,"MODAL", StringUtils.EMPTY, StringUtils.EMPTY,midjourneyUserConversation.getApiChannelId());
-        Map<String, Object> properties = result.getProperties();
-        properties.put("finalPrompt",midjourneyUserConversation.getPrompt());
-        properties.put("messageContent","**"+prompt+"** - <@mj>" + (mode == 1 ? "(fast)" : "(relaxed)"));
+        properties.put("messageContent","**"+prompt+"** - <@mj>" + modeStr);
+        MjVersionStatus mjVersionStatus = MJ_VERSION_STATUS_MAP.get(mjVersionStatusKey);
+        properties.put("estimateTime", mjVersionStatus == null ? null : System.currentTimeMillis() + mjVersionStatus.getMillisecond());
+        properties.put("finalPrompt",midjourneyUserConversation == null ? prompt : midjourneyUserConversation.getPrompt());
         return result;
     }
 
@@ -388,9 +413,6 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         if (result.getCode() == 21) {
             if (mode == 2 && modalFlag.get()){
                 SubmitResult modal = submitModal(user, mode,Long.valueOf(result.getResult()), conversation.getPrompt(), null);
-                Map<String, Object> properties = modal.getProperties();
-                properties.put("finalPrompt",conversation.getPrompt());
-                properties.put("messageContent", "**"+conversation.getPrompt()+"** - <@mj>" + " (relaxed)");
                 log.info("action-modal:{}", modal);
                 return modal;
             }else {
@@ -404,8 +426,19 @@ public class MidjourneyServiceImpl implements MidjourneyService {
             }
         }
         Map<String, Object> properties = result.getProperties();
+        String modeStr;
+        String mjVersionStatusKey;
+        if (mode == 1){
+            modeStr = "(fast)";
+            mjVersionStatusKey = "v6-fast-upscaler";
+        }else {
+            modeStr = "(relaxed)";
+            mjVersionStatusKey = "v6-relax-upscaler";
+        }
+        properties.put("messageContent","**"+conversation.getPrompt()+"** - <@mj>" + modeStr);
+        MjVersionStatus mjVersionStatus = MJ_VERSION_STATUS_MAP.get(mjVersionStatusKey);
+        properties.put("estimateTime", mjVersionStatus == null ? null : System.currentTimeMillis() + mjVersionStatus.getMillisecond());
         properties.put("finalPrompt",conversation.getPrompt());
-        properties.put("messageContent","**"+conversation.getPrompt()+"** - <@mj>" + (mode == 1 ? " (fast)" : " (relaxed)"));
         MidjourneyUserConversation midjourneyUserConversation = saveConversation(user.getId(), mode, result, "ACTION", StringUtils.EMPTY, botType, apiChannelId);
         properties.put("id",midjourneyUserConversation.getId());
         return result;
@@ -874,7 +907,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         synchronized (key.intern()) {
             String currentProgressStr = redisService.getStr(key);
             int currentProgress = StringUtils.isBlank(currentProgressStr) ? 0 : Integer.parseInt(currentProgressStr);
-            if (conversation.getProgressNum() > currentProgress) {
+            if (conversation.getProgressNum() > currentProgress || "FAILURE".equals(conversation.getStatus())) {
                 redisService.set(key, String.valueOf(conversation.getProgressNum()), RedisService.key.MIDJOURNEY_PROGRESS.getTimeout());
                 sendProgressToClient(conversation);
             }
@@ -966,16 +999,21 @@ public class MidjourneyServiceImpl implements MidjourneyService {
 
     @Override
     public void syncMidjourneyStatus() {
-
+        String url = "https://storage.googleapis.com/midjourney-status/status.json";
+        String res = HttpUtil.get(url);
+        JSONArray metrics = new JSONObject(res).getJSONArray("metrics");
+        List<MjVersionStatus> list = metrics.toList(MjVersionStatus.class);
+        for (MjVersionStatus mjVersionStatus : list) {
+            if (mjVersionStatus.getName().contains("jobs.v6")) {
+                // 分类数据并组合键
+                String[] parts = mjVersionStatus.getName().split("\\.");
+                String mainCategory = parts[1];
+                String subCategory = parts[2] + "-" + parts[3];
+                String key = mainCategory + "-" + subCategory;
+                MJ_VERSION_STATUS_MAP
+                        .putIfAbsent(key, mjVersionStatus);
+            }
+        }
     }
 
-    public static void main(String[] args) throws Exception {
-        String body = HttpUtil.get("https://storage.googleapis.com/midjourney-status/status.json");
-        System.out.println(body);
-        JSONObject jsonObject = JSONUtil.parseObj(body);
-        System.out.println(jsonObject);
-        JSONArray metrics1 = jsonObject.getJSONArray("metrics");
-        Map<String, Object> metrics = Jsons.toMap(metrics1);
-        System.out.println(metrics);
-    }
 }

+ 1 - 62
midjourney/src/main/java/com/yhlxj/service/task/Scheduler.java

@@ -22,8 +22,6 @@ import java.util.concurrent.ConcurrentHashMap;
 @RequiredArgsConstructor
 public class Scheduler {
 
-    public static final Map<String, MjVersionStatus> MJ_VERSION_STATUS_MAP = new ConcurrentHashMap<>();
-
     private final MidjourneyService midjourneyService;
 
     private final MidjourneyAccountService midjourneyAccountService;
@@ -38,69 +36,10 @@ public class Scheduler {
         midjourneyAccountService.syncAccountByMJPlus();
     }
 
-    @Scheduled(cron = "0 0/2 * * * ?")
-    public void getMjStatus() throws Exception {
-        String url = "https://storage.googleapis.com/midjourney-status/status.json";
-        String res = HttpUtil.get(url);
-        JSONObject json = new JSONObject(res);
-        JSONArray metrics = json.getJSONArray("metrics");
-        List<MjVersionStatus> list = metrics.toList(MjVersionStatus.class);
-        for (MjVersionStatus mjVersionStatus : list) {
-            if (mjVersionStatus.getName().contains("jobs.v6")) {
-                // 分类数据并组合键
-                String[] parts = mjVersionStatus.getName().split("\\.");
-                String mainCategory = parts[1];
-                String subCategory = parts[2] + "-" + parts[3];
-                String key = mainCategory + "-" + subCategory;
-                MJ_VERSION_STATUS_MAP
-                        .putIfAbsent(key, mjVersionStatus);
-
-            }
-
-        }
-    }
-
 
-    @Scheduled(cron = "0 0/5 * * * ?")
+    @Scheduled(cron = "0 0/2 * * * ?")
     public void syncMidjourneyStatus() throws Exception {
         midjourneyService.syncMidjourneyStatus();
     }
 
-}
-
-@Data
-class MjVersionStatus {
-    //{
-    //    "name": "jobs.time_to_start.fast.10min.job_type_kdpt_diffusion_anime",
-    //        "date": "2024-07-11T03:24:01.784Z",
-    //        "value": "N/A"
-    //}
-
-    private String name;
-
-    private String date;
-
-    private String value;
-    /**
-     * 转换为秒
-     */
-    private Long second;
-
-
-    public void setValue(String value) {
-        this.value = value;
-        // 转换为秒并设置 second
-        if (value != null && !value.equals("N/A")) {
-            try {
-                double minutes = Double.parseDouble(value);
-                this.second = Math.round(minutes * 60);
-            } catch (NumberFormatException e) {
-                this.second = null;
-                System.err.println("Invalid value format: " + value);
-            }
-        } else {
-            this.second = null;
-        }
-    }
-
 }