ソースを参照

Dev comfy generate

tsurumure 1 ヶ月 前
コミット
f29a68881d

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

@@ -13,6 +13,7 @@ import com.backendsys.modules.sdk.comfyui.entity.CFQueue;
 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.beans.factory.annotation.Value;
 import org.springframework.stereotype.Service;
 import reactor.core.publisher.Mono;
 
@@ -25,10 +26,8 @@ public class CrtGenerateServiceImpl implements CrtGenerateService {
 
     @Autowired
     private ComfyUIService comfyUIService;
-
     @Autowired
     private ComfyUISocketService comfyUISocketService;
-
     @Autowired
     private CrtDramaProjectStoryboardDao crtDramaProjectStoryboardDao;
 
@@ -62,7 +61,7 @@ public class CrtGenerateServiceImpl implements CrtGenerateService {
         String client_id = Convert.toStr(UUID.randomUUID());
 
         // -- [ComfyUI] 创建 WebSocket 监听连接 ---------------------------
-        comfyUISocketService.connectToSse(client_id, "ws://43.128.1.201:8001/ws").subscribe();
+        comfyUISocketService.connectToSse(client_id, 8001).subscribe();
 
         // -- [ComfyUI] 执行任务 -----------------------------------------
 

+ 1 - 1
src/main/java/com/backendsys/modules/sdk/comfyui/controller/ComfyUIDemoController.java

@@ -33,7 +33,7 @@ public class ComfyUIDemoController {
     @Anonymous
     @PostMapping("/api/comfyui/ws/connect")
     public String connect(String clientId) {
-        comfyUISocketService.connect(clientId, "ws://43.128.1.201:8001/ws").subscribe();
+        comfyUISocketService.connect(clientId, 8001).subscribe();
         return "Connection initiated for: " + clientId;
     }
 

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

@@ -5,10 +5,10 @@ import reactor.core.publisher.Mono;
 public interface ComfyUISocketService {
 
     // [ComfyUI] 创建 WebSocket 监听连接
-    Mono<Void> connect(String clientId, String wsUrl);
+    Mono<Void> connect(String clientId, Integer port);
 
     // [ComfyUI] 创建 WebSocket 监听连接 (转发到 SSE)
-    Mono<Void> connectToSse(String clientId, String wsUrl);
+    Mono<Void> connectToSse(String clientId, Integer port);
 
     // [ComfyUI] 断开 WebSocket 监听连接
     Mono<Void> disconnect(String clientId);

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

@@ -24,15 +24,15 @@ import java.util.function.Consumer;
 @Service
 public class ComfyUIServiceImpl implements ComfyUIService {
 
-    @Value("${comfyui.url}")
-    private String COMFYUI_URL;
+    @Value("${comfyui.host}")
+    private String COMFYUI_HOST;
     @Value("${comfyui.token}")
     private String COMFYUI_TOKEN;
 
     private WebClient webClient;
     public WebClient getWebClient() {
         if (webClient == null) {
-//            webClient = WebClient.builder().baseUrl(COMFYUI_URL + ":8001").filter(WebClientFilter.logFilter).build();
+//            webClient = WebClient.builder().baseUrl(COMFYUI_HOST + ":8001").filter(WebClientFilter.logFilter).build();
             webClient = WebClient.builder().filter(WebClientFilter.logFilter).build();
         }
         return webClient;
@@ -53,7 +53,7 @@ public class ComfyUIServiceImpl implements ComfyUIService {
      * [ComfyUI] 查询任务队列
      */
     public Mono<CFQueue> getQueue() {
-        String url = COMFYUI_URL + "/queue";
+        String url = "http://" + COMFYUI_HOST + ":8001/queue";
         WebClient webClient = getWebClient();
         return webClient.get()
                 .uri(url)
@@ -72,7 +72,7 @@ public class ComfyUIServiceImpl implements ComfyUIService {
         bodyValue.setClient_id(client_id);
         bodyValue.setPrompt(prompt);
 
-        String url = COMFYUI_URL + ":8001/prompt";
+        String url = "http://" + COMFYUI_HOST + ":8001/prompt";
 
         WebClient webClient = getWebClient();
         return webClient.post()

+ 29 - 5
src/main/java/com/backendsys/modules/sdk/comfyui/service/impl/ComfyUISocketServiceImpl.java

@@ -31,6 +31,8 @@ public class ComfyUISocketServiceImpl implements ComfyUISocketService {
     @Autowired
     private SseUtil sseUtil;
 
+    @Value("${comfyui.host}")
+    private String COMFYUI_HOST;
     @Value("${comfyui.token}")
     private String COMFYUI_TOKEN;
 
@@ -55,7 +57,9 @@ public class ComfyUISocketServiceImpl implements ComfyUISocketService {
      * [ComfyUI] 创建 WebSocket 监听连接
      */
     @Override
-    public Mono<Void> connect(String clientId, String wsUrl) {
+    public Mono<Void> connect(String clientId, Integer port) {
+
+        String wsUrl =  "ws://" + COMFYUI_HOST + ":" + port + "/ws";
         return Mono.defer(() -> {
             if (sessions.containsKey(clientId)) {
                 return Mono.error(new IllegalStateException("Connection already exists for client: " + clientId));
@@ -88,10 +92,11 @@ public class ComfyUISocketServiceImpl implements ComfyUISocketService {
      * [ComfyUI] 创建 WebSocket 监听连接 (转发到 SSE)
      */
     @Override
-    public Mono<Void> connectToSse(String clientId, String wsUrl) {
+    public Mono<Void> connectToSse(String clientId, Integer port) {
 
         Long user_id = SecurityUtil.getUserId();
 
+        String wsUrl =  "ws://" + COMFYUI_HOST + ":" + port + "/ws";
         return Mono.defer(() -> {
             if (sessions.containsKey(clientId)) {
                 return Mono.error(new IllegalStateException("Connection already exists for client: " + clientId));
@@ -117,11 +122,30 @@ public class ComfyUISocketServiceImpl implements ComfyUISocketService {
                                 messageObj = JSONUtil.parseObj(message);
                                 // 记录生成数据
                                 JSONObject output = JSONUtil.parseObj(messageObj.get("output"));
-                                JSONArray images = JSONUtil.parseArray(output.get("images"));
-                                // [{"filename": "ComfyUI_00122_.png", "subfolder": "", "type": "output"}]
-                                // http://43.128.1.201:8001/api/view?filename=ComfyUI_00118_.png
+
+                                Object imagesObj = output.get("images");
+                                if (imagesObj != null) {
+                                    JSONArray images = JSONUtil.parseArray(imagesObj);
+                                    // [{"filename": "ComfyUI_00122_.png", "subfolder": "", "type": "output"}]
+                                    // http://43.128.1.201:8001/api/view?filename=ComfyUI_00118_.png
+
+                                    if (images.size() > 0) {
+                                        for (int i = 0; i < images.size(); i++) {
+
+                                            JSONObject image = images.getJSONObject(i);
+//                                    JSONObject image = images.getJSONObject(i);
+//                                    String filename = image.getStr("filename");
+//                                    String subfolder = image.getStr("subfolder");
+//                                    String type = image.getStr("type");
+//                                    System.out.println("filename: " + filename + ", subfolder: " + subfolder + ", type: " + type);
+//                                    sseUtil.send(user_id, new SseResponse(SseResponseEnum.COMFYUI, "filename: " + filename + ", subfolder: " + subfolder + ", type: " + type).toJsonStr());
+                                        }
+                                    }
+                                }
+
                             }
 
+
                             sseUtil.send(user_id, new SseResponse(SseResponseEnum.COMFYUI, messageObj).toJsonStr());
                         })
                         .doOnError(e -> {

+ 1 - 1
src/main/resources/application-dev.yml

@@ -191,5 +191,5 @@ klingai:
   token-duration-time: 10000
 
 comfyui:
-  url: http://127.0.0.1
+  host: 127.0.0.1
   token: $2b$12$.MR4qGaFetN1FPQzbfyIrehsyjnPJ12xAZhR/l7KZpLkUPQTCG4gy

+ 1 - 1
src/main/resources/application-local.yml

@@ -204,5 +204,5 @@ klingai:
   token-duration-time: 10000
 
 comfyui:
-  url: http://43.128.1.201
+  host: 43.128.1.201
   token: $2b$12$.MR4qGaFetN1FPQzbfyIrehsyjnPJ12xAZhR/l7KZpLkUPQTCG4gy

+ 1 - 1
src/main/resources/application-prod.yml

@@ -192,5 +192,5 @@ klingai:
   token-duration-time: 10000
 
 comfyui:
-  url: http://127.0.0.1
+  host: 127.0.0.1
   token: $2b$12$.MR4qGaFetN1FPQzbfyIrehsyjnPJ12xAZhR/l7KZpLkUPQTCG4gy