Parcourir la source

调整生图语法细节

tsurumure il y a 1 mois
Parent
commit
0dd04b90cc

+ 3 - 3
db/comfyui_task.sql

@@ -10,8 +10,8 @@ CREATE TABLE `comfyui_task` (
     `id` BIGINT AUTO_INCREMENT COMMENT 'ID',
     `user_id` BIGINT NOT NULL COMMENT '用户ID',
     `prompt_id` VARCHAR(255) NOT NULL COMMENT '任务ID',
-#     `status` TINYINT DEFAULT '-1' COMMENT '任务状态 (-1:未开始, 1:进行中, 2:成功, 3:失败)',
-    `status` TINYINT DEFAULT '-1' COMMENT '任务状态 (-1:未完成, 1:已完成)',
+    `type` VARCHAR(255) NOT NULL COMMENT '任务类型 (Text2Image, ..)',
+    `status` TINYINT DEFAULT '-1' COMMENT '任务状态 (-1:未开始, 1:进行中, 2:成功, 3:失败)',
     `msg` VARCHAR(2000) COMMENT '任务状态信息,当任务失败时展示失败原因(如触发平台的内容风控等)',
     `generate_request` TEXT COMMENT '任务请求原始参数 (JSONString)',
     `generate_response` TEXT COMMENT '任务生成原始结果 (JSONString)',
@@ -19,7 +19,7 @@ CREATE TABLE `comfyui_task` (
     `update_time` DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
     INDEX `idx_user_id` (`user_id`),
     INDEX `idx_tprompt_id` (`prompt_id`)
-) ENGINE=INNODB AUTO_INCREMENT=1 DEFAULT CHARSET=utf8mb4 COMMENT='短剧创作-生图任务表';
+) ENGINE=INNODB AUTO_INCREMENT=1 DEFAULT CHARSET=utf8mb4 COMMENT='任务表';
 
 /*
     [可灵] 生视频结果

+ 3 - 2
src/main/java/com/backendsys/modules/sdk/comfyui/entity/ComfyuiTask.java

@@ -6,14 +6,15 @@ import com.baomidou.mybatisplus.annotation.TableName;
 import lombok.Data;
 
 @Data
-@TableName("comfyui_task")
+@TableName("crt_task")
 public class ComfyuiTask {
 
     @TableId(type = IdType.AUTO)
     private Long id;
     private Long user_id;                 // 用户ID
     private String prompt_id;             // 任务ID
-    private Integer status;               // 任务状态 (-1:未完成, 1:已完成)
+    private String type;                  // 任务类型 (Text2Image, ..)
+    private Integer status;               // 任务状态 (-1:未开始, 1:进行中, 2:成功, 3:失败)
     private String msg;                   // 任务状态信息,当任务失败时展示失败原因(如触发平台的内容风控等)
     private String generate_request;      // 任务请求原始参数 (JSONString)
     private String generate_response;     // 任务生成原始结果 (JSONString)

+ 1 - 0
src/main/java/com/backendsys/modules/sdk/comfyui/enums/TaskStatusEnums.java

@@ -21,4 +21,5 @@ public enum TaskStatusEnums {
     public String getValue() {
         return this.value;
     }
+
 }

+ 21 - 0
src/main/java/com/backendsys/modules/sdk/comfyui/enums/TaskTypeEnums.java

@@ -0,0 +1,21 @@
+package com.backendsys.modules.sdk.comfyui.enums;
+
+public enum TaskTypeEnums {
+    TEXT_2_IMAGE("Text2Image", "文生图"),
+    TEXT_2_VIDEO("Text2Video", "文生视频")
+    ;
+    private final String value;
+    private final String label;
+
+    TaskTypeEnums(String value, String label) {
+        this.value = value;
+        this.label = label;
+    }
+    public String getLabel() {
+        return this.label;
+    }
+    public String getValue() {
+        return this.value;
+    }
+
+}

+ 3 - 1
src/main/java/com/backendsys/modules/sdk/comfyui/service/impl/ComfyuiSocketServiceImpl.java

@@ -187,7 +187,9 @@ public class ComfyuiSocketServiceImpl implements ComfyuiSocketService {
                                     data.put("data", dataChildren);
 
                                     // [DB] 新增任务记录
-                                    comfyUtil.asyncUpdateTaskRecord(prompt_id, JSONUtil.toJsonStr(dataChildren), 1);
+                                    CompletableFuture.runAsync(() -> {
+                                        comfyUtil.updateComfyuiTask(prompt_id, JSONUtil.toJsonStr(dataChildren), 2);
+                                    });
 
                                 }
 

+ 7 - 6
src/main/java/com/backendsys/modules/sdk/comfyui/service/impl/ComfyuiText2ImageServiceImpl.java

@@ -5,6 +5,7 @@ import cn.hutool.json.JSONUtil;
 import com.backendsys.modules.sdk.comfyui.dao.ComfyuiTaskDao;
 import com.backendsys.modules.sdk.comfyui.entity.ComfyuiResponse;
 import com.backendsys.modules.sdk.comfyui.entity.ComfyuiText2Image;
+import com.backendsys.modules.sdk.comfyui.enums.TaskTypeEnums;
 import com.backendsys.modules.sdk.comfyui.service.ComfyuiService;
 import com.backendsys.modules.sdk.comfyui.service.ComfyuiText2ImageService;
 import com.backendsys.modules.sdk.comfyui.utils.ComfyUtil;
@@ -337,14 +338,14 @@ public class ComfyuiText2ImageServiceImpl implements ComfyuiText2ImageService {
         ;
         JSONObject prompt_object = JSONUtil.parseObj(prompt);
 
+        // [DB] 新增任务记录
+        comfyUtil.insertComfyuiTask(prompt, TaskTypeEnums.TEXT_2_IMAGE.getValue());
+
         // [ComfyUI] 执行任务
         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;
     }
 
@@ -425,14 +426,14 @@ public class ComfyuiText2ImageServiceImpl implements ComfyuiText2ImageService {
         ;
         JSONObject prompt_object = JSONUtil.parseObj(prompt);
 
+        // [DB] 新增任务记录
+        comfyUtil.insertComfyuiTask(prompt, TaskTypeEnums.TEXT_2_IMAGE.getValue());
+
         // [ComfyUI] 执行任务
         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;
     }
 

+ 30 - 34
src/main/java/com/backendsys/modules/sdk/comfyui/utils/ComfyUtil.java

@@ -7,8 +7,6 @@ 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 {
 
@@ -16,43 +14,41 @@ public class ComfyUtil {
     private ComfyuiTaskDao comfyuiTaskDao;
 
     /**
-     * [新增] 生成任务的记录 (异步)
+     * [新增] 生成任务的记录
      */
-    public void asyncInsertTaskRecord(String prompt_id, String prompt) {
+    public void insertComfyuiTask(String prompt, String type) {
         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());
-            }
-        });
+        try {
+            System.out.println("-- 新增任务记录 --");
+            ComfyuiTask comfyuiTask = new ComfyuiTask();
+            comfyuiTask.setUser_id(user_id);
+            comfyuiTask.setType(type);
+            comfyuiTask.setGenerate_request(prompt);
+
+            comfyuiTaskDao.insert(comfyuiTask);
+        } catch (Exception e) {
+            System.out.println(e.getMessage());
+        }
     }
     /**
-     * [更新] 生成任务的记录 (异步)
+     * [更新] 生成任务的记录
+     * - status: 任务状态 (-1:未开始, 1:进行中, 2:成功, 3:失败)
      */
-    public void asyncUpdateTaskRecord(String prompt_id, String response, Integer status) {
-        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(status);
-                comfyuiTask.setGenerate_response(response);
-                comfyuiTaskDao.update(comfyuiTask, wrapper);
-            } catch (Exception e) {
-                System.out.println(e.getMessage());
-            }
-        });
+    public void updateComfyuiTask(String prompt_id, String response, Integer status) {
+        try {
+            System.out.println("-- 更新任务记录 --");
+            LambdaQueryWrapper<ComfyuiTask> wrapper = new LambdaQueryWrapper<>();
+            wrapper.eq(ComfyuiTask::getPrompt_id, prompt_id);
+
+            ComfyuiTask comfyuiTask = new ComfyuiTask();
+            comfyuiTask.setStatus(status);
+            comfyuiTask.setPrompt_id(prompt_id);
+            comfyuiTask.setGenerate_response(response);
+
+            comfyuiTaskDao.update(comfyuiTask, wrapper);
+        } catch (Exception e) {
+            System.out.println(e.getMessage());
+        }
     }
 
 }