Dev comfyui

db/crt_generate_image.sql

@@ -13,6 +13,7 @@ CREATE TABLE `crt_generate_image` (
     `prompt_id` VARCHAR(255) COMMENT '任务ID',
     `url_origin` VARCHAR(2000) NOT NULL COMMENT '原图',
     `url` VARCHAR(2000) COMMENT '转存图',
+    `object_key` VARCHAR(500) COMMENT 'ObjectKey',
     `target` INT COMMENT '上传目标 (-1:本地, 1:腾讯云, 2:阿里云, 3.抖音云)',
     `create_time` DATETIME DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
     `update_time` DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',

src/main/java/com/backendsys/modules/crt/entity/CrtGenerateImage.java

@@ -24,6 +24,7 @@ public class CrtGenerateImage {
     @TableField(exist = false)
     private String url_thumb;                   // 转存图-缩略图
     private Integer target;                     // 上传目标 (-1:本地, 1:腾讯云, 2:阿里云, 3.抖音云)
+    private String object_key;
 
     @JsonAdapter(LocalDateTimeAdapter.class)
     private LocalDateTime create_time;

src/main/java/com/backendsys/modules/crt/service/impl/CrtGenerateServiceImpl.java

@@ -156,9 +156,14 @@ public class CrtGenerateServiceImpl implements CrtGenerateService {
         comfyuiText2Image.setBatch_size(param_batch_size);
         comfyuiText2Image.setPrompt_flux_guidance(param_prompt_flux_guidance);
         comfyuiText2Image.setSampler(param_sampler);
+
         comfyuiText2Image.setStep(param_step);
         comfyuiText2Image.setWidth(width);
         comfyuiText2Image.setHeight(height);
+//        comfyuiText2Image.setStep(5);
+//        comfyuiText2Image.setWidth(400);
+//        comfyuiText2Image.setHeight(200);
+
         comfyuiText2Image.setModel_name(model_name);
         comfyuiText2Image.setPrompt_text(prompt_text);
         comfyuiText2Image.setSeed(param_seed);

src/main/java/com/backendsys/modules/sdk/comfyui/service/impl/ComfyuiSocketServiceImpl.java

@@ -9,6 +9,7 @@ import com.backendsys.modules.crt.dao.CrtGenerateImageDao;
 import com.backendsys.modules.crt.entity.CrtGenerateImage;
 import com.backendsys.modules.sdk.comfyui.enums.TaskStatusEnums;
 import com.backendsys.modules.sdk.comfyui.service.ComfyuiSocketService;
+import com.backendsys.modules.sdk.comfyui.utils.ComfyUtil;
 import com.backendsys.modules.sse.entity.SseResponse;
 import com.backendsys.modules.sse.entity.SseResponseEnum;
 import com.backendsys.modules.sse.utils.SseUtil;
@@ -38,6 +39,9 @@ public class ComfyuiSocketServiceImpl implements ComfyuiSocketService {
 
     @Autowired
     private SseUtil sseUtil;
+
+    @Autowired
+    private ComfyUtil comfyUtil;
     @Autowired
     private SysFileService sysFileService;
     @Autowired
@@ -142,32 +146,50 @@ public class ComfyuiSocketServiceImpl implements ComfyuiSocketService {
 
                             JSONObject data = JSONUtil.parseObj(message);
 
-                            if (is_save) {
-                                // == [任务执行完成] =======================================================
-                                // { "type": "executed", .. }
-                                String type = Convert.toStr(data.get("type"));
-                                if (TaskStatusEnums.EXECUTED.getValue().equals(type)) {
-
-                                    JSONObject dataChildren = JSONUtil.parseObj(data.get("data"));
-                                    JSONObject output = JSONUtil.parseObj(dataChildren.get("output"));
-
-                                    // -- [生成图片 -> 转存图片 -> 新增记录] ------------------------------------------------------
-                                    // 由于图片地址不是公开的，需要加 Token 访问，因此不能公开返回原始图片地址，比如：
-                                    // http://43.128.1.201:8000/api/view?filename=fenjing_1_00012_.png&token=$2b$12$.MR4qGaFetN1FPQzbfyIrehsyjnPJ12xAZhR/l7KZpLkUPQTCG4gy
-                                    // 因此，直接在 ComfyUI 的物理目录上构建 nginx 静态资源访问目录
-                                    Object imagesObj = output.get("images");
-                                    if (imagesObj != null) {
-                                        String prompt_id = Convert.toStr(dataChildren.get("prompt_id"));
-                                        Long drama_project_storyboard_id = Convert.toLong(params.get("drama_project_storyboard_id"));
-                                        List<String> images_path = imagesToStoryboard(user_id, prompt_id, port, imagesObj, drama_project_storyboard_id);
-                                        output.put("images_path", images_path);
-                                        dataChildren.put("output", output);
-                                        data.put("data", dataChildren);
+
+                            // == [任务执行完成] =======================================================
+                            // { "type": "executed", .. }
+                            String type = Convert.toStr(data.get("type"));
+                            if (TaskStatusEnums.EXECUTED.getValue().equals(type)) {
+
+                                JSONObject dataChildren = JSONUtil.parseObj(data.get("data"));
+                                JSONObject output = JSONUtil.parseObj(dataChildren.get("output"));
+
+                                // -- [生成图片 -> 转存图片 -> 新增记录] ------------------------------------------------------
+                                // 由于图片地址不是公开的，需要加 Token 访问，因此不能公开返回原始图片地址，比如：
+                                // http://43.128.1.201:8000/api/view?filename=fenjing_1_00012_.png&token=$2b$12$.MR4qGaFetN1FPQzbfyIrehsyjnPJ12xAZhR/l7KZpLkUPQTCG4gy
+                                // 因此，直接在 ComfyUI 的物理目录上构建 nginx 静态资源访问目录
+                                Object imagesObj = output.get("images");
+                                if (imagesObj != null) {
+                                    String prompt_id = Convert.toStr(dataChildren.get("prompt_id"));
+                                    Long drama_project_storyboard_id = Convert.toLong(params.get("drama_project_storyboard_id"));
+
+                                    // 临时图片相对路径
+                                    JSONArray images = JSONUtil.parseArray(imagesObj);
+                                    // 返回完整的临时图片路径
+                                    List<String> images_temp_path_full = imagesToFullpath(images, port);
+                                    System.out.println("images_temp_path_full = " + images_temp_path_full);
+                                    output.put("images_temp_path", images_temp_path_full);
+
+                                    if (is_save) {
+                                        // 创建一个 CompletableFuture 来执行异步任务
+                                        CompletableFuture.runAsync(() -> {
+                                            // 返回转存后的图片路径
+                                            List<String> images_path = imagesToStoryboard(user_id, prompt_id, images_temp_path_full, drama_project_storyboard_id);
+                                            output.put("images_path", images_path);
+                                        });
                                     }
 
+                                    dataChildren.put("output", output);
+                                    data.put("data", dataChildren);
+
+                                    // [db] 新增任务记录
+                                    comfyUtil.asyncUpdateTaskRecord(prompt_id, JSONUtil.toJsonStr(dataChildren));
+
                                 }
-                                // ========================================================================
+
                             }
+                            // ========================================================================
 
                             sseUtil.send(user_id, new SseResponse(SseResponseEnum.COMFYUI, data).toJsonStr());
                         })
@@ -218,41 +240,51 @@ public class ComfyuiSocketServiceImpl implements ComfyuiSocketService {
         return Flux.fromIterable(sessions.keySet());
     }
 
-    // 转存图片
-    private List<String> imagesToStoryboard(Long user_id, String prompt_id, Integer port, Object imagesObj, Long drama_project_storyboard_id) {
-        JSONArray images = JSONUtil.parseArray(imagesObj);
+    // 返回临时图片完整路径
+    private List<String> imagesToFullpath(JSONArray images, Integer port) {
         List<String> images_path = new ArrayList<>();
         if (images.size() > 0) {
             for (int i = 0; i < images.size(); i++) {
                 JSONObject image = images.getJSONObject(i);
                 String filename = image.getStr("filename");
-                String filepath = "http://o.daogu.ai/" + port + "/" + filename;     // 生成图片临时目录
+                String filepath = "http://o.daogu.ai/" + port + "/" + filename;
+                images_path.add(filepath);
+            }
+        }
+        return images_path;
+    }
+
+    // 转存图片
+    private List<String> imagesToStoryboard(Long user_id, String prompt_id, List<String> images_temp_path_full, Long drama_project_storyboard_id) {
+
+        List<String> images_path = new ArrayList<>();
+        if (images_temp_path_full.size() > 0) {
+
+            images_temp_path_full.stream().forEach(temp_image -> {
 
                 // -- [图片转存储存桶] -------------------------------------
-                SysFileResult result = sysFileService.urlToUploadFile(filepath, user_id);
+                SysFileResult result = sysFileService.urlToUploadFile(temp_image, user_id);
                 System.out.println("urlToUploadFile (result) = " + result);
 
+                String image = result.getUrl();
+                images_path.add(image);
+
                 // -- [记录到生成图片记录表] --------------------------------
-                // 创建一个 CompletableFuture 来执行异步任务
-                CompletableFuture.runAsync(() -> {
-
-                    if (drama_project_storyboard_id != null) {
-                        CrtGenerateImage entity = new CrtGenerateImage();
-                        entity.setUser_id(user_id);
-                        entity.setDrama_project_storyboard_id(drama_project_storyboard_id);
-                        entity.setPrompt_id(prompt_id);
-                        entity.setUrl_origin(filepath);
-                        entity.setUrl(result.getUrl());
-                        entity.setTarget(result.getTarget());
-                        crtGenerateImageDao.insert(entity);
-                    }
-
-                });
+                if (drama_project_storyboard_id != null) {
+                    CrtGenerateImage entity = new CrtGenerateImage();
+                    entity.setUser_id(user_id);
+                    entity.setDrama_project_storyboard_id(drama_project_storyboard_id);
+                    entity.setPrompt_id(prompt_id);
+                    entity.setUrl_origin(temp_image);
+                    entity.setUrl(image);
+                    entity.setTarget(result.getTarget());
+                    entity.setObject_key(result.getKey());
+                    crtGenerateImageDao.insert(entity);
+                }
                 // ------------------------------------------------------
 
-                images_path.add(filepath);
+            });
 
-            }
         }
         return images_path;
     }

src/main/java/com/backendsys/modules/sdk/comfyui/service/impl/ComfyuiText2ImageServiceImpl.java

@@ -7,17 +7,21 @@ import com.backendsys.modules.sdk.comfyui.entity.ComfyuiResponse;
 import com.backendsys.modules.sdk.comfyui.entity.ComfyuiText2Image;
 import com.backendsys.modules.sdk.comfyui.service.ComfyuiService;
 import com.backendsys.modules.sdk.comfyui.service.ComfyuiText2ImageService;
+import com.backendsys.modules.sdk.comfyui.utils.ComfyUtil;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
 import reactor.core.publisher.Mono;
 
+import java.util.concurrent.CompletableFuture;
+
 @Service
 public class ComfyuiText2ImageServiceImpl implements ComfyuiText2ImageService {
 
     @Autowired
     private ComfyuiService comfyUIService;
+
     @Autowired
-    private ComfyuiTaskDao comfyuiTaskDao;
+    private ComfyUtil comfyUtil;
 
     /**
      * [ComfyUI] 文生图 (7.4生图.json)
@@ -337,6 +341,10 @@ public class ComfyuiText2ImageServiceImpl implements ComfyuiText2ImageService {
         Mono<ComfyuiResponse> cfPromptResponseMono = comfyUIService.prompt(client_id, prompt_object);
         ComfyuiResponse response = cfPromptResponseMono.block();
         response.setClient_id(client_id);
+
+        // [db] 新增任务记录
+        comfyUtil.asyncInsertTaskRecord(response.getPrompt_id(), prompt);
+
         return response;
     }
 
@@ -421,6 +429,10 @@ public class ComfyuiText2ImageServiceImpl implements ComfyuiText2ImageService {
         Mono<ComfyuiResponse> cfPromptResponseMono = comfyUIService.prompt(client_id, prompt_object);
         ComfyuiResponse response = cfPromptResponseMono.block();
         response.setClient_id(client_id);
+
+        // [db] 新增任务记录
+        comfyUtil.asyncInsertTaskRecord(response.getPrompt_id(), prompt);
+
         return response;
     }
 

src/main/java/com/backendsys/modules/sdk/comfyui/utils/ComfyUtil.java

@@ -1,8 +1,61 @@
 package com.backendsys.modules.sdk.comfyui.utils;
 
+import com.backendsys.modules.common.config.security.utils.SecurityUtil;
+import com.backendsys.modules.sdk.comfyui.dao.ComfyuiTaskDao;
+import com.backendsys.modules.sdk.comfyui.entity.ComfyuiTask;
+import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Component;
 
+import java.util.concurrent.CompletableFuture;
+
 @Component
 public class ComfyUtil {
 
+    @Autowired
+    private ComfyuiTaskDao comfyuiTaskDao;
+
+    /**
+     * [新增] 生成任务的记录 (异步)
+     */
+    public void asyncInsertTaskRecord(String prompt_id, String prompt) {
+        Long user_id = SecurityUtil.getUserId();
+        CompletableFuture.runAsync(() -> {
+            // [db] 新增任务记录
+            try {
+                System.out.println("-- 新增任务记录 --");
+                ComfyuiTask comfyuiTask = new ComfyuiTask();
+                comfyuiTask.setUser_id(user_id);
+                comfyuiTask.setPrompt_id(prompt_id);
+                comfyuiTask.setGenerate_request(prompt);
+                comfyuiTaskDao.insert(comfyuiTask);
+            } catch (Exception e) {
+                System.out.println(e.getMessage());
+            }
+        });
+    }
+    /**
+     * [更新] 生成任务的记录 (异步)
+     */
+    public void asyncUpdateTaskRecord(String prompt_id, String response) {
+        CompletableFuture.runAsync(() -> {
+            // [db] 新增任务记录
+            try {
+                System.out.println("-- 更新任务记录 --");
+                LambdaQueryWrapper<ComfyuiTask> wrapper = new LambdaQueryWrapper<>();
+                wrapper.eq(ComfyuiTask::getPrompt_id, prompt_id);
+
+                ComfyuiTask comfyuiTask = new ComfyuiTask();
+
+
+                // 目前 [进行中、失败] 的任务状态 还没有做记录？
+                comfyuiTask.setStatus(2);
+                comfyuiTask.setGenerate_response(response);
+                comfyuiTaskDao.update(comfyuiTask, wrapper);
+            } catch (Exception e) {
+                System.out.println(e.getMessage());
+            }
+        });
+    }
+
 }

src/main/java/com/backendsys/modules/sse/entity/SseResponseEnum.java

@@ -7,7 +7,7 @@ public enum SseResponseEnum {
     LOGOUT("logout", "退出登录"),
     NOTICE("notice", "通知"),
     UPLOAD("upload", "上传"),
-    COMFYUI("comfyui", "ComfyUI"),
+    COMFYUI("comfyui", "Comfyui"),
     OLLAMA("ollama", "Ollama"),
     DEEPSEEK("deepseek", "Deepseek"),
     HUNYUAN("hunyuan", "Hunyuan");