ComfyUI AI系统深度实践：从工作流编排到自定义节点开发

import json
import requests
import websocket
import uuid
import time

class ComfyUIRuntime:
    def __init__(self, server_address="127.0.0.1:8188"):
        self.base_url = f"http://{server_address}"
        self.ws_url = f"ws://{server_address}/ws"
    
    def load_workflow(self, json_path):
        with open(json_path, 'r') as f:
            self.workflow = json.load(f)
        return self
    
    def set_prompt(self, text, node_id="6"):  # 默认 CLIP Text Encode 节点为 6
        self.workflow[node_id]["inputs"]["text"] = text
        return self
    
    def set_seed(self, seed, node_id="3"):   # 默认 KSampler 节点为 3
        self.workflow[node_id]["inputs"]["seed"] = seed
        return self
    
    def queue_prompt(self):
        """提交任务并返回 prompt_id"""
        payload = {"prompt": self.workflow}
        resp = requests.post(f"{self.base_url}/prompt", json=payload)
        return resp.json()["prompt_id"]
    
    def wait_for_image(self, prompt_id, timeout=120):
        """通过 WebSocket 监听执行进度，并获取输出图像路径"""
        ws = websocket.WebSocket()
        ws.connect(self.ws_url)
        images = []
        start = time.time()
        while time.time() - start < timeout:
            out = ws.recv()
            if isinstance(out, str):
                data = json.loads(out)
                if data["type"] == "executing" and data["data"]["prompt_id"] == prompt_id:
                    if data["data"]["node"] is None:  # 执行完毕
                        break
            else:
                # 二进制数据（图像预览）
                pass
            time.sleep(0.1)
        ws.close()
        # 实际生产环境从 output 文件夹读取图片或直接解析 Base64
        return f"http://{self.server_address}/view?filename=output.png"

# custom_nodes/prompt_enhancer.py
import comfy
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

class PromptEnhancer:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "raw_prompt": ("STRING", {"multiline": True, "default": "一只猫"}),
                "creativity": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.0, "step": 0.05})
            }
        }
    
    RETURN_TYPES = ("STRING",)
    FUNCTION = "enhance"
    CATEGORY = "custom/llm"

    def __init__(self):
        self.model = None
        self.tokenizer = None
    
    def load_model(self):
        if self.model is None:
            self.tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
            self.model = AutoModelForCausalLM.from_pretrained(
                "Qwen/Qwen2.5-1.5B-Instruct",
                torch_dtype=torch.float16,
                device_map="cuda"
            )
        return self.model, self.tokenizer
    
    def enhance(self, raw_prompt, creativity):
        model, tokenizer = self.load_model()
        sys_msg = f"将以下简短描述扩展为Stable Diffusion正向提示词，添细节、光效、风格，扩展系数{creativity}："
        messages = [{"role": "system", "content": sys_msg}, {"role": "user", "content": raw_prompt}]
        text = tokenizer.apply_chat_template(messages, tokenize=False)
        inputs = tokenizer(text, return_tensors="pt").to("cuda")
        outputs = model.generate(**inputs, max_new_tokens=150, temperature=creativity)
        enhanced = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # 清洗掉系统指令部分，仅保留扩展后的提示词
        return (enhanced.split("assistant")[-1].strip(),)

# 节点注册
NODE_CLASS_MAPPINGS = {"PromptEnhancer": PromptEnhancer}
NODE_DISPLAY_NAME_MAPPINGS = {"PromptEnhancer": "LLM 提示词增强器"}

from fastapi import FastAPI, BackgroundTasks
from rq import Queue
from redis import Redis
from pydantic import BaseModel
import uuid

app = FastAPI()
redis_conn = Redis(host='localhost', port=6379)
task_queue = Queue('comfy_tasks', connection=redis_conn)

class GenRequest(BaseModel):
    prompt: str
    width: int = 512
    height: int = 512
    callback_url: str = None

def run_comfy_task(prompt, width, height, callback_url):
    """后台工作函数，由 RQ Worker 执行"""
    runtime = ComfyUIRuntime()
    runtime.load_workflow("base_workflow.json") \
           .set_prompt(prompt) \
           .set_seed(int(time.time()))
    pid = runtime.queue_prompt()
    image_url = runtime.wait_for_image(pid)
    
    if callback_url:
        requests.post(callback_url, json={"status": "done", "image_url": image_url, "prompt": prompt})
    return image_url

@app.post("/generate")
async def generate(request: GenRequest, background_tasks: BackgroundTasks):
    task_id = str(uuid.uuid4())
    # 将任务推入 Redis 队列，由独立 Worker 进程消费
    job = task_queue.enqueue(run_comfy_task, request.prompt, request.width, request.height, request.callback_url)
    return {"task_id": job.id, "status": "queued", "eta": "约15秒"}

# 优化后的图像获取（直接捕获 WebSocket 二进制）
def wait_and_capture(ws, prompt_id):
    while True:
        msg = ws.recv()
        if isinstance(msg, bytes):
            # 解析 ComfyUI 二进制头，提取图片数据
            import struct
            header = msg[:8]
            # 实际逻辑略，直接保存为 BytesIO
            from PIL import Image
            import io
            img = Image.open(io.BytesIO(msg[8:]))
            return img