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

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;
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.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;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.UUID;

@Service
public class CrtGenerateServiceImpl implements CrtGenerateService {

    @Autowired
    private ComfyUIService comfyUIService;
    @Autowired
    private ComfyUISocketService comfyUISocketService;
    @Autowired
    private CrtDramaProjectStoryboardDao crtDramaProjectStoryboardDao;

    /**
     * [ComfyUI] 查询任务队列
     */
    @Override
    public Map<String, Object> getQueue() {
        // [ComfyUI] 执行任务
        Mono<CFQueue> cfQueueMono = comfyUIService.getQueue();
        CFQueue response = cfQueueMono.block();
        System.out.println("结果: " + response);

        Map<String, Object> resp = new LinkedHashMap<>();
        resp.put("response", response);
        return resp;
    }

    /**
     * 短剧创作-生成图片
     */
    @Override
    public Map<String, Object> generateImage(CrtDramaProjectStoryboard crtDramaProjectStoryboard) {

        Long drama_project_storyboard_id = crtDramaProjectStoryboard.getDrama_project_storyboard_id();

        CrtDramaProjectStoryboard detail = crtDramaProjectStoryboardDao.selectById(drama_project_storyboard_id);
        if (detail == null) throw new CustException("分镜不存在");

        // -- 前端生成的UUID ---------------------------------------------
        String client_id = Convert.toStr(UUID.randomUUID());

        // -- [ComfyUI] 创建 WebSocket 监听连接 ---------------------------
        comfyUISocketService.connectToSse(client_id, 8001).subscribe();

        // -- [ComfyUI] 执行任务 -----------------------------------------

        // [Demo-基础生图]
        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\": \"保存图像\" }" +
                    "}," +
                "}";

        JSONObject prompt_object = JSONUtil.parseObj(prompt);

        System.out.println("prompt_object: " + prompt_object);

        // [ComfyUI] 执行任务
        Mono<CFPromptResponse> cfPromptResponseMono = comfyUIService.prompt(client_id, prompt_object);
        CFPromptResponse response = cfPromptResponseMono.block();
        response.setClient_id(client_id);
        System.out.println("结果: " + response);

        Map<String, Object> resp = new LinkedHashMap<>();
        resp.put("drama_project_storyboard_id", drama_project_storyboard_id);
        resp.put("response", response);
        return resp;
    }

}