zwhui hace 1 año
padre
commit
c3f8076421

+ 6 - 0
midjourney/src/main/java/com/yhlxj/dao/model/entity/MidjourneyUserQueue.java

@@ -1,14 +1,18 @@
 package com.yhlxj.dao.model.entity;
 
 import com.yhlxj.dao.model.BaseEntity;
+import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.Data;
+import lombok.NoArgsConstructor;
 
 /**
  * @author zwhui
  * @date 2025/6/6 13:42
  */
 @Data
+@NoArgsConstructor
+@AllArgsConstructor
 @Builder
 public class MidjourneyUserQueue extends BaseEntity {
     private Long userId;
@@ -20,4 +24,6 @@ public class MidjourneyUserQueue extends BaseEntity {
     private String prompt;
 
     private String base64Array;
+
+    private String fileExt;
 }

+ 1 - 0
midjourney/src/main/java/com/yhlxj/redis/RedisService.java

@@ -814,6 +814,7 @@ public class RedisService {
         MIDJOURNEY_RELEASE_LIMIT("midjourney:release:limit:","midjourney解除限制次数", 30L),
         MIDJOURNEY_QUEUE_LIMIT("midjourney:queue:limit:","midjourney慢速队列", 60 * 60 * 48L),
         MIDJOURNEY_QUEUE_DETAIL_LIMIT("midjourney:queue:detail:limit:","midjourney慢速队列详情", 30 * 60L),
+        MIDJOURNEY_QUEUE_DETAIL_IMAGINE_LIMIT("midjourney:queue:detail:imagine:limit:","midjourney慢速队列图片", 30 * 60L),
         ;
 
         private String name;

+ 4 - 1
midjourney/src/main/java/com/yhlxj/service/midjourney/MidjourneyService.java

@@ -12,6 +12,7 @@ import com.yhlxj.dao.model.entity.MidjourneyUserQueue;
 import com.yhlxj.dao.model.response.SubmitResult;
 import com.yhlxj.dao.model.views.MidjourneyUserQueueView;
 
+import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
@@ -45,6 +46,8 @@ public interface MidjourneyService {
 
     MidjourneyUser getUser(String userToken) throws ExecutionException, InterruptedException;
 
+    MidjourneyUser getUserSimple(String userToken);
+
     void syscConversation() throws Exception;
 
     MidjourneyUserConversation uploadFile(SubmitUploadDTO uploadDTO,MidjourneyUser user) throws Exception;
@@ -61,7 +64,7 @@ public interface MidjourneyService {
 
     void releaseLimit(String userToken);
 
-    MidjourneyUserQueueView queueLimit(Long userId, SubmitImagineDTO submitImagineDTO);
+    MidjourneyUserQueueView queueLimit(Long userId, SubmitImagineDTO submitImagineDTO) throws IOException;
 
     List<MidjourneyUserQueue> getTaskPosition(Long userId);
 }

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

@@ -3,6 +3,8 @@ package com.yhlxj.service.midjourney.impl;
 import cn.hutool.core.bean.BeanUtil;
 import cn.hutool.core.collection.CollectionUtil;
 import cn.hutool.core.date.DateUtil;
+import cn.hutool.core.io.FileUtil;
+import cn.hutool.core.io.file.FileNameUtil;
 import cn.hutool.core.map.MapUtil;
 import cn.hutool.core.util.NumberUtil;
 import cn.hutool.http.HttpRequest;
@@ -30,6 +32,7 @@ import com.yhlxj.service.midjourney.MidjourneyAccountService;
 import com.yhlxj.service.midjourney.MidjourneyService;
 import com.yhlxj.service.midjourney.MidjourneyUserSettingsService;
 import com.yhlxj.util.COSUtil;
+import io.undertow.util.FileUtils;
 import lombok.Data;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -49,8 +52,10 @@ import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URL;
+import java.nio.file.DirectoryStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.time.LocalTime;
 import java.util.*;
 import java.util.concurrent.*;
@@ -119,7 +124,13 @@ public class MidjourneyServiceImpl implements MidjourneyService {
     private static final int MAX_PROCESSING_TIME = 60;
 
 
-
+    public MidjourneyUser getUserSimple(String userToken){
+        MidjourneyUser midjourneyUser = getMidjourneyUser(userToken);
+        if (midjourneyUser.getIsBlack()){
+            throw new BusinessRuntimeException("10002","账号异常,请联系客服");
+        }
+        return midjourneyUser;
+    }
     @Override
     public MidjourneyUser getUser(String userToken) throws ExecutionException, InterruptedException {
         MidjourneyUser midjourneyUser = getMidjourneyUser(userToken);
@@ -1373,17 +1384,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
 
 
     @Override
-    public MidjourneyUserQueueView queueLimit(Long userId, SubmitImagineDTO submitImagineDTO) {
-        Long queueId = sequence.nextId();
-
-        RedisService.key key = RedisService.key.MIDJOURNEY_QUEUE_DETAIL_LIMIT;
-        String queueKey = key.getName() + userId;
-        MidjourneyUserQueue queue = MidjourneyUserQueue.builder()
-                .queueId(queueId).userId(userId).prompt(submitImagineDTO.getPrompt())
-                .base64Array(JSONUtil.toJsonStr(submitImagineDTO.getBase64Array())).build();
-
-        redisService.lSet(queueKey, JSONUtil.toJsonStr(queue),key.getTimeout());
-
+    public MidjourneyUserQueueView queueLimit(Long userId, SubmitImagineDTO submitImagineDTO) throws IOException {
         // 计算当前时段需要的额外等待时间(分钟)
         int additionalWaitMinutes = calculateAdditionalWaitTime();
         // 清理过期的假任务,保持队列干净
@@ -1422,7 +1423,7 @@ public class MidjourneyServiceImpl implements MidjourneyService {
                 }
             }
         }
-
+        Long queueId = sequence.nextId();
         // 真实任务的时间戳设置为"当前时间+额外等待时间"
         long processingTime = MIN_PROCESSING_TIME +
                 (long)(Math.random() * (MAX_PROCESSING_TIME - MIN_PROCESSING_TIME));
@@ -1437,6 +1438,13 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         MidjourneyUserQueueView view = new MidjourneyUserQueueView();
         view.setQueueId(queueId);
         view.setPosition(position);
+
+        RedisService.key key = RedisService.key.MIDJOURNEY_QUEUE_DETAIL_LIMIT;
+        String queueKey = key.getName() + userId;
+        String fileExt = base64ToPic(queueId, submitImagineDTO.getBase64Array());
+        MidjourneyUserQueue queue = MidjourneyUserQueue.builder().position(position)
+                .queueId(queueId).userId(userId).prompt(submitImagineDTO.getPrompt()).fileExt(fileExt).build();
+        redisService.lSet(queueKey, JSONUtil.toJsonStr(queue),key.getTimeout());
         return view;
     }
 
@@ -1451,8 +1459,13 @@ public class MidjourneyServiceImpl implements MidjourneyService {
         List<MidjourneyUserQueue> list = redisService.lGet(RedisService.key.MIDJOURNEY_QUEUE_DETAIL_LIMIT.getName() + userId, 0L, -1L)
                 .stream().map(item -> JSONUtil.toBean((String) item, MidjourneyUserQueue.class)).collect(Collectors.toList());
         list.forEach(item -> {
-            Long rank = redisService.zRank(RedisService.key.MIDJOURNEY_QUEUE_LIMIT.getName() + userId, item.getQueueId());
-            item.setPosition((rank != null) ? rank + 1 : 1);
+            try {
+                item.setBase64Array(encodeImageToBase64("/mnt/picture/"+item.getQueueId()+"."+item.getFileExt()));
+                Long rank = redisService.zRank(RedisService.key.MIDJOURNEY_QUEUE_LIMIT.getName() + userId, item.getQueueId());
+                item.setPosition((rank != null) ? rank + 1 : 1);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
         });
         return list;
     }
@@ -1504,4 +1517,33 @@ public class MidjourneyServiceImpl implements MidjourneyService {
     }
 
 
+    /**
+     * base64图片存本地
+     */
+    public static String base64ToPic(Long queueId,List<String> base64Array) throws IOException {
+        for (String base64 : base64Array) {
+            String fileExt = "jpeg";
+            if (base64.contains(";")){
+                fileExt = base64.split(";")[0].split("/")[1];
+                base64 = base64.split(",")[1];
+            }
+            byte[] body = Base64.getDecoder().decode(base64);
+            File file = new File("/mnt/picture", queueId + "." + fileExt);
+            try (FileImageOutputStream imageOutput = new FileImageOutputStream(file)) {
+                imageOutput.write(body, 0, body.length);
+            }
+            return fileExt;
+        }
+        return null;
+    }
+    public static String encodeImageToBase64(String imagePathStr) throws IOException {
+        Path imagePath = Paths.get(imagePathStr);
+        if (!Files.exists(imagePath)) {
+            throw new IOException("文件未找到: " + imagePathStr);
+        }
+        // 1. 读取文件所有字节
+        byte[] imageBytes = Files.readAllBytes(imagePath);
+        return Base64.getEncoder().encodeToString(imageBytes);
+    }
+
 }

+ 32 - 10
midjourney/src/main/java/com/yhlxj/web/mirror/MidjourneyController.java

@@ -1,5 +1,6 @@
 package com.yhlxj.web.mirror;
 
+import cn.hutool.core.io.FileUtil;
 import cn.hutool.core.lang.Validator;
 import cn.hutool.core.util.ObjectUtil;
 import cn.hutool.core.util.StrUtil;
@@ -48,6 +49,10 @@ import javax.servlet.http.HttpServletResponse;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.DirectoryStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.List;
@@ -97,6 +102,16 @@ public class MidjourneyController {
        return midjourneyService.getUser(userToken);
     }
 
+    public MidjourneyUser getUserSimple() throws IOException {
+        String userToken = request.getHeader("user-token");
+        String ip = IpUtils.getIpAddr(request);
+        log.info("用户请求 ip:{}, userToken:{}", ip, userToken);
+        if (StringUtils.isBlank(userToken)) {
+            response.sendRedirect("https://nf.video");
+        }
+        return midjourneyService.getUserSimple(userToken);
+    }
+
     @GetMapping("/midjourneyMirrorWithToken/{userToken}")
     public void midjourneyMirrorWithToken(@PathVariable String userToken) throws IOException {
         Cookie cookie = new Cookie("userToken", userToken);
@@ -230,6 +245,7 @@ public class MidjourneyController {
         Integer mode = submitImagineDTO.getMode();
         if (submitImagineDTO.getQueueId() != null){
             redisService.lRemoveLeft(RedisService.key.MIDJOURNEY_QUEUE_DETAIL_LIMIT.getName() + user.getId());
+            delFile(submitImagineDTO.getQueueId());
         }
         Map<String, Object> checkResult = checkUserLimit(user, mode);
         Long num = (Long) checkResult.get("num");
@@ -253,7 +269,13 @@ public class MidjourneyController {
         //syncUser(user.getId(), user.getMode(), num);
         return GatewayResponse.SUCCESS.newBuilder().toResult(conversation);
     }
-
+    public void delFile(Long queueId) throws IOException {
+        try (DirectoryStream<Path> stream = Files.newDirectoryStream(Paths.get("/mnt/picture"), queueId +".*")) {
+            for (Path filePath : stream) {
+                FileUtil.del(filePath);
+            }
+        }
+    }
     /**
      * 提交Describe任务
      */
@@ -405,7 +427,7 @@ public class MidjourneyController {
     @NoSubmit
     public Result<SubmitResult> seed(@Validated @RequestBody SubmitSeedDTO seedDTO) throws Exception {
         log.info("任务id:{},执行动作:{}",seedDTO.getTaskId(), "seed");
-        MidjourneyUser user = getUser();
+        MidjourneyUser user = getUserSimple();
         try {
             return GatewayResponse.SUCCESS.newBuilder().toResult(midjourneyService.submitSeed(user, seedDTO.getTaskId(), seedDTO.getBotType()));
         } catch (Exception e) {
@@ -420,7 +442,7 @@ public class MidjourneyController {
      */
     @GetMapping("/conversation/list")
     public Result<SearchResult<MidjourneyUserConversation>> conversationList() throws Exception {
-        MidjourneyUser user = getUser();
+        MidjourneyUser user = getUserSimple();
         log.info("查询会话列表 userId:{}",user.getId());
         MapBuilder builder = MapUtils.flatBuilder(request.getParameterMap());
         return GatewayResponse.SUCCESS.newBuilder().toResult(beanSearcher.search(MidjourneyUserConversation.class,
@@ -435,7 +457,7 @@ public class MidjourneyController {
      */
     @GetMapping("/conversation/listByIds")
     public Result<List<MidjourneyUserConversation>> conversationListByIds(Integer mode,@RequestParam(value = "ids" ) List<String> ids) throws Exception {
-        MidjourneyUser user = getUser();
+        MidjourneyUser user = getUserSimple();
         log.info("根据ids查询会话 userId:{}",user.getId());
         return GatewayResponse.SUCCESS.newBuilder().toResult(midjourneyService.listConversationByIds(mode,ids));
     }
@@ -446,7 +468,7 @@ public class MidjourneyController {
     @PostMapping("/conversation/{id}/cancel")
     @NoSubmit
     public Result<Void> conversationCancel(@PathVariable("id") Long id) throws Exception {
-        MidjourneyUser user = getUser();
+        MidjourneyUser user = getUserSimple();
         log.info("取消任务 id:{}",id);
         midjourneyService.cancelConversation(user,id);
         return GatewayResponse.SUCCESS.newBuilder().toResult();
@@ -520,7 +542,7 @@ public class MidjourneyController {
      */
     @GetMapping("/get/painting/record")
     public Result<SearchResult<MidjourneyUserPaintingDayRecordView>> getPaintingRecord(String prompt, Boolean bookmark) throws Exception {
-        MidjourneyUser user = getUser();
+        MidjourneyUser user = getUserSimple();
         log.info("用户绘画记录 userId:{},prompt:{},bookmark:{}", user.getId(), prompt, bookmark);
         String condition = String.format(" and mc.user_id = %s", user.getId());
         if (StrUtil.isNotBlank(prompt)) {
@@ -552,7 +574,7 @@ public class MidjourneyController {
      */
     @GetMapping("/get/settings")
     public Result<MidjourneyUserSettings> getSettings() throws Exception {
-        MidjourneyUser user = getUser();
+        MidjourneyUser user = getUserSimple();
         return GatewayResponse.SUCCESS.newBuilder().toResult(midjourneyUserSettingsService.getById(user.getSettings().getId()));
     }
 
@@ -561,7 +583,7 @@ public class MidjourneyController {
      */
     @PostMapping("/put/settings")
     public Result<MidjourneyUserSettings> upSettings(@RequestBody MidjourneyUserSettings settings) throws Exception {
-        MidjourneyUser user = getUser();
+        MidjourneyUser user = getUserSimple();
         if (!user.getId().equals(settings.getUserId())) {
             throw BusinessRuntimeException.getInstance("参数错误");
         }
@@ -677,7 +699,7 @@ public class MidjourneyController {
     @PostMapping("/queue")
     public Result<MidjourneyUserQueueView> saveQueue(@Validated @RequestBody SubmitImagineDTO submitImagineDTO) throws Exception {
         log.info("提交Imagine任务,提示:{}",submitImagineDTO.getPrompt());
-        MidjourneyUser user = getUser();
+        MidjourneyUser user = getUserSimple();
         Integer mode = submitImagineDTO.getMode();
         if (mode != 2){
             throw BusinessRuntimeException.getInstance("网络异常!");
@@ -697,7 +719,7 @@ public class MidjourneyController {
      */
     @GetMapping("/getTaskPosition")
     public Result<List<MidjourneyUserQueue>> getTaskPosition() throws Exception{
-        MidjourneyUser user = getUser();
+        MidjourneyUser user = getUserSimple();
         return GatewayResponse.SUCCESS.newBuilder().toResult(midjourneyService.getTaskPosition(user.getId()));
     }
 }