chenbiao 2 gadi atpakaļ
vecāks
revīzija
fac074ca06

+ 108 - 47
midjourney/src/main/java/com/yhlxj/service/midjourney/impl/MidjourneyServiceImpl.java

@@ -176,7 +176,16 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         }
         properties.put("messageContent","**" + prompt + "** - <@mj>" + modeStr);
         MjVersionStatus mjVersionStatus = MJ_VERSION_STATUS_MAP.get(mjVersionStatusKey);
-        properties.put("estimateTime", mjVersionStatus == null ? null : System.currentTimeMillis() + mjVersionStatus.getMillisecond());
+        int count = 1;
+        long second = 1000 * 60L;
+        if (mjVersionStatus != null) {
+            second = Math.max(mjVersionStatus.getMillisecond(), second);
+        }
+
+        if (result.getCode() == 22) {
+            count = extractNumberOrDefault(result.getDescription());
+        }
+        properties.put("estimateTime", System.currentTimeMillis() + (second * count));
         return result;
     }
 
@@ -249,7 +258,16 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         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());
+        int count = 1;
+        long second = 1000 * 60L;
+        if (mjVersionStatus != null) {
+            second = Math.max(mjVersionStatus.getMillisecond(), second);
+        }
+
+        if (result.getCode() == 22) {
+            count = extractNumberOrDefault(result.getDescription());
+        }
+        properties.put("estimateTime", System.currentTimeMillis() + (second * count));
         return result;
     }
 
@@ -282,7 +300,16 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         }
         properties.put("messageContent","**"+prompt+"** - <@mj>" + modeStr);
         MjVersionStatus mjVersionStatus = MJ_VERSION_STATUS_MAP.get(mjVersionStatusKey);
-        properties.put("estimateTime", mjVersionStatus == null ? null : System.currentTimeMillis() + mjVersionStatus.getMillisecond());
+        int count = 1;
+        long second = 1000 * 60L;
+        if (mjVersionStatus != null) {
+            second = Math.max(mjVersionStatus.getMillisecond(), second);
+        }
+
+        if (result.getCode() == 22) {
+            count = extractNumberOrDefault(result.getDescription());
+        }
+        properties.put("estimateTime", System.currentTimeMillis() + (second * count));
         properties.put("finalPrompt",midjourneyUserConversation == null ? prompt : midjourneyUserConversation.getPrompt());
         return result;
     }
@@ -339,7 +366,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
     //            JSONObject jsonObject = JSONUtil.parseObj(submitResult.getResult());
     //            throw BusinessRuntimeException.getInstance("prompt包含敏感词:"+jsonObject.getStr("bannedWord"));
     //        }
-    //        throw BusinessRuntimeException.getInstance("队列已满,请稍后试");
+    //        throw BusinessRuntimeException.getInstance("队列已满,请稍后试");
     //    }
     //    if (StringUtils.isNotBlank(accountWithMinUsage) && mode == 2) {
     //        submitResult.setInstanceId(Long.valueOf(accountWithMinUsage));
@@ -447,7 +474,16 @@ public class MidjourneyServiceImpl implements MidjourneyService {
             }
         }
         MjVersionStatus mjVersionStatus = MJ_VERSION_STATUS_MAP.get(mjVersionStatusKey);
-        properties.put("estimateTime", mjVersionStatus == null ? null : System.currentTimeMillis() + mjVersionStatus.getMillisecond());
+        int count = 1;
+        long second = 1000 * 60L;
+        if (mjVersionStatus != null) {
+            second = Math.max(mjVersionStatus.getMillisecond(), second);
+        }
+
+        if (result.getCode() == 22) {
+            count = extractNumberOrDefault(result.getDescription());
+        }
+        properties.put("estimateTime", System.currentTimeMillis() + (second * count));
         properties.put("finalPrompt",conversation.getPrompt());
         MidjourneyUserConversation midjourneyUserConversation = saveConversation(user.getId(), mode, result, "ACTION", StringUtils.EMPTY, botType, apiChannelId);
         properties.put("id",midjourneyUserConversation.getId());
@@ -498,7 +534,9 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         }
         if (StringUtils.isNotBlank(accountWithMinUsage) && mode == 2) {
             submitResult.setInstanceId(Long.valueOf(accountWithMinUsage));
-            userSubmitLimit(user,action,param.get("customId"));
+            synchronized (user.getId()){
+                userSubmitLimit(user,action,param.get("customId"));
+            }
         }
         return submitResult;
     }
@@ -532,7 +570,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
                                 throw BusinessRuntimeException.getInstance("该任务所属的源Mj账号已被官网封禁或暂时风控。请稍后再试,或重新生成图片");
                             }else {
                                 log.error("服务商:{}, action:{}, error message:{}", midjourneyApiChannel.getName(), action, submitResult.getDescription());
-                                throw BusinessRuntimeException.getInstance("该任务源Mj账号队列已满,请稍后试");
+                                throw BusinessRuntimeException.getInstance("该任务源Mj账号队列已满,请稍后试");
                             }
                         }
                     }
@@ -545,9 +583,14 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         }else if (code == 24) {
             Map<String, Object> properties = submitResult.getProperties();
             throw BusinessRuntimeException.getInstance("prompt包含敏感词:"+properties.get("bannedWord"));
-        }else {
+        }
+        else if (code == 3) {
+            log.error("action:" + action + " error message:" + submitResult.getDescription());
+            throw BusinessRuntimeException.getInstance("该任务所属的源Mj账号已被官网封禁或暂时风控,请重新生成图片");
+        }
+        else {
             log.error("action:" + action + " error message:" + submitResult.getDescription());
-            throw BusinessRuntimeException.getInstance("该任务源Mj账号队列已满,请稍后尝试");
+            throw BusinessRuntimeException.getInstance("网络出现异常");
         }
     }
 
@@ -924,48 +967,51 @@ public class MidjourneyServiceImpl implements MidjourneyService {
             log.info("获取到的图片 url:{}", conversation.getImageUrl());
             ImgSize imgSize = TASK_IMG_SIZE_MAP.get(conversation.getTaskId());
             URL url = new URL(conversation.getImageUrl());
-            url = new URL(url.getProtocol(), "cdn.mj.liuliangbang.vip", url.getPort(), url.getFile());
-            String filePath = url.getFile();
-            // 注册 WebP 插件
-            ImageIO.scanForPlugins();
-
-            // 从 URL 中读取图像,特别处理 WebP 格式
-            BufferedImage image = null;
-            if (filePath.contains("webp")) {
-                Iterator<ImageReader> readers = ImageIO.getImageReadersByMIMEType("image/webp");
-                if (readers.hasNext()) {
-                    ImageReader reader = readers.next();
-                    try {
-                        ImageInputStream iis = ImageIO.createImageInputStream(url.openStream());
-                        reader.setInput(iis, true);
-                        image = reader.read(0);
-                    }catch (Exception e){
-                        log.error("下载图获取尺寸错误,{}", StringUtil.getErrorText(e));
+            try {
+                String filePath = url.getFile();
+                // 注册 WebP 插件
+                ImageIO.scanForPlugins();
+
+                // 从 URL 中读取图像,特别处理 WebP 格式
+                BufferedImage image = null;
+                if (filePath.contains("webp")) {
+                    Iterator<ImageReader> readers = ImageIO.getImageReadersByMIMEType("image/webp");
+                    if (readers.hasNext()) {
+                        ImageReader reader = readers.next();
+                        try {
+                            ImageInputStream iis = ImageIO.createImageInputStream(url.openStream());
+                            reader.setInput(iis, true);
+                            image = reader.read(0);
+                        }catch (Exception e){
+                            log.error("下载图获取尺寸错误,{}", StringUtil.getErrorText(e));
+                        }
                     }
+                } else {
+                    image = ImageIO.read(url);
                 }
-            } else {
-                image = ImageIO.read(url);
-            }
 
-            if (image == null) {
-                throw new IOException("Failed to read image from URL: " + conversation.getImageUrl());
-            }
+                if (image == null) {
+                    throw new IOException("Failed to read image from URL: " + conversation.getImageUrl());
+                }
 
-            if (imgSize == null) {
-                try {
-                    //BufferedImage image = ImageIO.read(new URL(url.getProtocol(), "cdn.mj.liuliangbang.vip", url.getPort(), url.getFile()));
-                    // 获取图像的宽度和高度
-                    int width = image.getWidth();
-                    int height = image.getHeight();
-                    log.info("首次获取到的图片 尺寸:{},{}", width, height);
-                    // 缩小宽度和高度的一半
-                    imgSize = new ImgSize();
-                    imgSize.setHeight(height);
-                    imgSize.setWidth(width);
-                    TASK_IMG_SIZE_MAP.put(conversation.getTaskId(), imgSize);
-                }catch (Exception e){
-                    log.error("下载图获取尺寸错误,{}", StringUtil.getErrorText(e));
+                if (imgSize == null) {
+                    try {
+                        //BufferedImage image = ImageIO.read(new URL(url.getProtocol(), "cdn.mj.liuliangbang.vip", url.getPort(), url.getFile()));
+                        // 获取图像的宽度和高度
+                        int width = image.getWidth();
+                        int height = image.getHeight();
+                        log.info("首次获取到的图片 尺寸:{},{}", width, height);
+                        // 缩小宽度和高度的一半
+                        imgSize = new ImgSize();
+                        imgSize.setHeight(height);
+                        imgSize.setWidth(width);
+                        TASK_IMG_SIZE_MAP.put(conversation.getTaskId(), imgSize);
+                    }catch (Exception e){
+                        log.error("下载图获取尺寸错误,{}", StringUtil.getErrorText(e));
+                    }
                 }
+            }catch (Exception e){
+                log.error("下载图获取尺寸错误,{}", StringUtil.getErrorText(e));
             }
 
             String scaling = "";
@@ -979,7 +1025,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
             log.info("任务执行回调 imgUrl:{}", conversation.getImageUrl());
             // 使用 Hutool 提取各个部分并构建新 URL
             String newUrl = URLUtil.toURI(new URL(url.getProtocol(), "cdn.mj.liuliangbang.vip", url.getPort(), url.getFile() + scaling)).toString();
-            //
+
             conversation.setImageUrl(newUrl);
         }
         redisService.set(RedisService.key.MIDJOURNEY_CONVERSATION.getName() + conversation.getTaskId(), conversation,RedisService.key.MIDJOURNEY_CONVERSATION.getTimeout());
@@ -1148,4 +1194,19 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         }
     }
 
+    public static int extractNumberOrDefault(String input) {
+       try {
+           Pattern pattern = Pattern.compile("\\d+");
+           Matcher matcher = pattern.matcher(input);
+
+           if (matcher.find()) {
+               return Integer.parseInt(matcher.group());
+           } else {
+               return 1;
+           }
+       }catch (Exception e){
+           return 1;
+       }
+    }
+
 }