瀏覽代碼

Dev generateImage

tsurumure 2 月之前
父節點
當前提交
680c961c5d

+ 5 - 0
src/main/java/com/backendsys/modules/crt/entity/CrtDramaProjectStoryboard.java

@@ -9,6 +9,7 @@ import com.baomidou.mybatisplus.annotation.TableId;
 import com.baomidou.mybatisplus.annotation.TableName;
 import com.google.gson.annotations.JsonAdapter;
 import jakarta.validation.constraints.Max;
+import jakarta.validation.constraints.NotEmpty;
 import jakarta.validation.constraints.NotNull;
 import jakarta.validation.constraints.Size;
 import lombok.Data;
@@ -35,6 +36,10 @@ public class CrtDramaProjectStoryboard {
     @NotNull(message = "分镜ID不能为空", groups = { Update.class, GenerateImage.class })
     private Long drama_project_storyboard_id;
 
+    @TableField(exist = false)
+    @NotEmpty(message = "生成任务ID不能为空", groups = { GenerateImage.class })
+    private String client_id;
+
     private Long user_id;
 
     @NotNull(message = "项目ID不能为空", groups = { Create.class, StoryboardDetail.class, Clear.class, Delete.class  })

+ 27 - 2
src/main/java/com/backendsys/modules/crt/service/impl/CrtGenerateServiceImpl.java

@@ -1,5 +1,8 @@
 package com.backendsys.modules.crt.service.impl;
 
+import cn.hutool.core.convert.Convert;
+import cn.hutool.json.JSONObject;
+import cn.hutool.json.JSONUtil;
 import com.backendsys.exception.CustException;
 import com.backendsys.modules.common.utils.Result;
 import com.backendsys.modules.crt.dao.CrtDramaProjectStoryboardDao;
@@ -7,11 +10,14 @@ import com.backendsys.modules.crt.entity.CrtDramaProjectStoryboard;
 import com.backendsys.modules.crt.service.CrtGenerateService;
 import com.backendsys.modules.sdk.comfyui.entity.CFPromptResponse;
 import com.backendsys.modules.sdk.comfyui.service.ComfyUIService;
+import com.backendsys.modules.sdk.comfyui.service.ComfyUISocketService;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
 import reactor.core.publisher.Mono;
 
+import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.UUID;
 
 @Service
 public class CrtGenerateServiceImpl implements CrtGenerateService {
@@ -19,6 +25,9 @@ public class CrtGenerateServiceImpl implements CrtGenerateService {
     @Autowired
     private ComfyUIService comfyUIService;
 
+    @Autowired
+    private ComfyUISocketService comfyUISocketService;
+
     @Autowired
     private CrtDramaProjectStoryboardDao crtDramaProjectStoryboardDao;
 
@@ -30,13 +39,29 @@ public class CrtGenerateServiceImpl implements CrtGenerateService {
         CrtDramaProjectStoryboard detail = crtDramaProjectStoryboardDao.selectById(drama_project_storyboard_id);
         if (detail == null) throw new CustException("分镜不存在");
 
+        // -- 前端生成的UUID ---------------------------------------------
+//        String client_id = Convert.toStr(UUID.randomUUID());
+        String client_id = crtDramaProjectStoryboard.getClient_id();
+
+        // -- [ComfyUI] 创建 WebSocket 监听连接 ---------------------------
+        comfyUISocketService.connect(client_id, "ws://43.128.1.201:8007/ws").subscribe();
+
+        // -- [ComfyUI] 执行任务 -----------------------------------------
+//        String prompt = "{{\"3\":{\"inputs\":{\"seed\":449753344472378,\"steps\":20,\"cfg\":8,\"sampler_name\":\"euler\",\"scheduler\":\"normal\",\"denoise\":1,\"model\":[\"4\",0],\"positive\":[\"6\",0],\"negative\":[\"7\",0],\"latent_image\":[\"5\",0]},\"class_type\":\"KSampler\",\"_meta\":{\"title\":\"K采样器\"}},\"4\":{\"inputs\":{\"ckpt_name\":\"v1-5-pruned-emaonly-fp16.safetensors\"},\"class_type\":\"CheckpointLoaderSimple\",\"_meta\":{\"title\":\"Checkpoint加载器(简易)\"}},\"5\":{\"inputs\":{\"width\":512,\"height\":512,\"batch_size\":1},\"class_type\":\"EmptyLatentImage\",\"_meta\":{\"title\":\"空Latent图像\"}},\"6\":{\"inputs\":{\"text\":\"beautiful scenery nature glass bottle landscape, , purple galaxy bottle,\",\"speak_and_recognation\":{\"__value__\":[false,true]},\"clip\":[\"4\",1]},\"class_type\":\"CLIPTextEncode\",\"_meta\":{\"title\":\"CLIP文本编码\"}},\"7\":{\"inputs\":{\"text\":\"text, watermark\",\"speak_and_recognation\":{\"__value__\":[false,true]},\"clip\":[\"4\",1]},\"class_type\":\"CLIPTextEncode\",\"_meta\":{\"title\":\"CLIP文本编码\"}},\"8\":{\"inputs\":{\"samples\":[\"3\",0],\"vae\":[\"4\",2]},\"class_type\":\"VAEDecode\",\"_meta\":{\"title\":\"VAE解码\"}},\"9\":{\"inputs\":{\"filename_prefix\":\"ComfyUI\",\"images\":[\"8\",0]},\"class_type\":\"SaveImage\",\"_meta\":{\"title\":\"保存图像\"}}}}";
         String prompt = "{}";
+        JSONObject prompt_object = JSONUtil.parseObj(prompt);
+
+        System.out.println("prompt_object: " + prompt_object);
 
-        Mono<CFPromptResponse> cfPromptResponse = comfyUIService.prompt(prompt);
+        Mono<CFPromptResponse> cfPromptResponse = comfyUIService.prompt(client_id, prompt_object);
         CFPromptResponse response = cfPromptResponse.block();
         System.out.println("结果: " + response);
 
-        return Map.of("drama_project_storyboard_id", drama_project_storyboard_id);
+        Map<String, Object> resp = new LinkedHashMap<>();
+        resp.put("drama_project_storyboard_id", drama_project_storyboard_id);
+        resp.put("client_id", client_id);
+        resp.put("response", response);
+        return resp;
     }
 
 }

+ 12 - 0
src/main/java/com/backendsys/modules/sdk/comfyui/entity/CFPromptRequest.java

@@ -0,0 +1,12 @@
+package com.backendsys.modules.sdk.comfyui.entity;
+
+import cn.hutool.json.JSONObject;
+import lombok.Data;
+
+@Data
+public class CFPromptRequest {
+
+    private String client_id;
+    private JSONObject prompt;
+
+}

+ 4 - 3
src/main/java/com/backendsys/modules/sdk/comfyui/entity/CFPromptResponse.java

@@ -5,8 +5,9 @@ import lombok.Data;
 @Data
 public class CFPromptResponse {
 
-    private String prompt_id;
-    private Integer number;
-    private Object node_errors;
+    private String prompt_id;       // 任务ID
+    private Integer number;         // 当前任务序号,可用于后续获取需要等待任务数的计算
+    private Object node_errors;     // 错误信息
+    private Object error;
 
 }

+ 2 - 1
src/main/java/com/backendsys/modules/sdk/comfyui/service/ComfyUIService.java

@@ -1,11 +1,12 @@
 package com.backendsys.modules.sdk.comfyui.service;
 
+import cn.hutool.json.JSONObject;
 import com.backendsys.modules.sdk.comfyui.entity.CFPromptResponse;
 import reactor.core.publisher.Mono;
 
 public interface ComfyUIService {
 
     // [ComfyUI] 执行任务
-    Mono<CFPromptResponse> prompt(String prompt);
+    Mono<CFPromptResponse> prompt(String client_id, JSONObject prompt);
 
 }

+ 38 - 16
src/main/java/com/backendsys/modules/sdk/comfyui/service/impl/ComfyUIServiceImpl.java

@@ -1,10 +1,13 @@
 package com.backendsys.modules.sdk.comfyui.service.impl;
 
 import cn.hutool.core.convert.Convert;
+import cn.hutool.json.JSONObject;
 import com.backendsys.modules.common.Filter.WebClientFilter;
+import com.backendsys.modules.sdk.comfyui.entity.CFPromptRequest;
 import com.backendsys.modules.sdk.comfyui.entity.CFPromptResponse;
 import com.backendsys.modules.sdk.comfyui.service.ComfyUIService;
 import org.springframework.beans.factory.annotation.Value;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
 import org.springframework.stereotype.Service;
 import org.springframework.util.LinkedMultiValueMap;
@@ -14,6 +17,7 @@ import org.springframework.web.util.UriComponentsBuilder;
 import reactor.core.publisher.Mono;
 
 import java.util.UUID;
+import java.util.function.Consumer;
 
 @Service
 public class ComfyUIServiceImpl implements ComfyUIService {
@@ -32,32 +36,50 @@ public class ComfyUIServiceImpl implements ComfyUIService {
     }
 
     /**
-     * [ComfyUI] 执行任务
+     * 获取 Token 拼接好的 HttpHeaders
      */
-    @Override
-    public Mono<CFPromptResponse> prompt(String prompt) {
+    public Consumer<HttpHeaders> getHttpHeaders() {
+        Consumer<HttpHeaders> httpHeaders = (headers) -> {
+            headers.add("Authorization", "Bearer " + COMFYUI_TOKEN);
+        };
+        return httpHeaders;
+    }
 
-        String client_id = Convert.toStr(UUID.randomUUID());
+    public CFPromptResponse mapErrorResponse(CFPromptResponse response) {
+        return response;
+    }
 
-        MultiValueMap<String, String> params = new LinkedMultiValueMap<>();
-        params.add("client_id", client_id);          // Unix时间戳、单位ms
-        params.add("prompt", prompt);
+    /**
+     * [ComfyUI] 执行任务
+     */
+    @Override
+    public Mono<CFPromptResponse> prompt(String client_id, JSONObject prompt) {
 
-        System.out.println("params = " + params);
+        CFPromptRequest bodyValue = new CFPromptRequest();
+        bodyValue.setClient_id(client_id);
+        bodyValue.setPrompt(prompt);
 
-        String url = "/prompt?token=" + COMFYUI_TOKEN;
+//        String url = "http://43.128.1.201:8007/prompt?token=" + COMFYUI_TOKEN;
+        String url = "http://43.128.1.201:8007/prompt";
         String uri = UriComponentsBuilder.fromUriString(url).toUriString();
+
         WebClient webClient = getWebClient();
-        return webClient.get()
+        return webClient.post()
                 .uri(uri)
+                .headers(getHttpHeaders())
                 .accept(MediaType.APPLICATION_JSON)
+                .bodyValue(bodyValue)
                 .exchangeToMono(response -> {
-                    return response.bodyToMono(CFPromptResponse.class);
-//                    if (response.statusCode().is2xxSuccessful()) {
-//                        return response.bodyToMono(CFPromptResponse.class); // 成功响应
-//                    } else {
-//                        return response.bodyToMono(CFPromptResponse.class).map(e -> KLingUtil.mapErrorResponse(e));
-//                    }
+                    if (response.statusCode().is2xxSuccessful()) {
+                        return response.bodyToMono(CFPromptResponse.class);
+                    } else {
+                        return response.bodyToMono(CFPromptResponse.class).map(e -> mapErrorResponse(e));
+                    }
+                }).onErrorResume(e -> {
+                    // 捕获所有异常(包括上面抛出的 RuntimeException)
+                    CFPromptResponse response = new CFPromptResponse();
+                    response.setNode_errors(e.getMessage());
+                    return Mono.just(response);
                 });
 
     }