ComfyUI：节点式AI图像工作流

"""ComfyUI节点式工作流核心逻辑模拟"""from typing import Any  class Node:    """节点基类，模拟ComfyUI中所有节点的抽象父类"""     CATEGORY = "default"     def __init__(self):        self._id = None        self._name = None     @classmethod    def INPUT_TYPES(cls):        """定义节点输入参数类型，返回包含required和optional的字典"""        return {"required": {}}     @classmethod    def RETURN_TYPES(cls):        """定义节点返回值类型元组"""        return ()     @classmethod    def FUNCTION(cls):        """指定节点执行时调用的方法名"""        return "execute"     def execute(self, **kwargs):        raise NotImplementedError  class CheckpointLoaderSimple(Node):    """加载SD模型检查点，返回模型、CLIP和VAE三个组件"""     CATEGORY = "loaders"     @classmethod    def INPUT_TYPES(cls):        return {"required": {"ckpt_name": ("STRING",)}}     @classmethod    def RETURN_TYPES(cls):        return ("MODEL", "CLIP", "VAE")     @classmethod    def FUNCTION(cls):        return "load_checkpoint"     def load_checkpoint(self, ckpt_name):        return (f"model:{ckpt_name}", f"clip:{ckpt_name}", f"vae:{ckpt_name}")  class KSampler(Node):    """K采样器节点，使用指定算法在潜在空间中进行去噪采样"""     CATEGORY = "sampling"     @classmethod    def INPUT_TYPES(cls):        return {"required": {            "model": ("MODEL",),            "seed": ("INT",),            "steps": ("INT",),            "cfg": ("FLOAT",),            "sampler_name": ("STRING",),            "positive": ("CONDITIONING",),            "negative": ("CONDITIONING",),            "latent_image": ("LATENT",),        }}     @classmethod    def RETURN_TYPES(cls):        return ("LATENT",)     @classmethod    def FUNCTION(cls):        return "sample"     def sample(self, model, seed, steps, cfg, sampler_name,               positive, negative, latent_image):        return ({"samples": f"latent_{seed}_{steps}steps"},)  class VAEDecode(Node):    """VAE解码器，将潜在空间数据解码为像素空间图像"""     CATEGORY = "latent"     @classmethod    def INPUT_TYPES(cls):        return {"required": {            "samples": ("LATENT",),            "vae": ("VAE",),        }}     @classmethod    def RETURN_TYPES(cls):        return ("IMAGE",)     @classmethod    def FUNCTION(cls):        return "decode"     def decode(self, samples, vae):        return ("decoded_image",)  class PreviewImage(Node):    """预览图像节点，将图像输出到用户界面"""     CATEGORY = "image"     @classmethod    def INPUT_TYPES(cls):        return {"required": {"images": ("IMAGE",)}}     @classmethod    def RETURN_TYPES(cls):        return ()     def execute(self, images):        return ()  class ExecutionEngine:    """执行引擎，负责构建DAG图并按拓扑序执行节点"""     def __init__(self):        self.nodes = {}        self.connections = []     def add_node(self, node_id, node_instance):        node_instance._id = node_id        self.nodes[node_id] = node_instance     def add_connection(self, from_id, from_slot, to_id, to_slot):        self.connections.append({            "from": from_id, "from_slot": from_slot,            "to": to_id, "to_slot": to_slot        })     def get_execution_order(self):        in_degree = {nid: 0 for nid in self.nodes}        adj = {nid: [] for nid in self.nodes}        for conn in self.connections:            adj[conn["from"]].append(conn["to"])            in_degree[conn["to"]] += 1        queue = [nid for nid in self.nodes if in_degree[nid] == 0]        order = []        while queue:            nid = queue.pop(0)            order.append(nid)            for nxt in adj[nid]:                in_degree[nxt] -= 1                if in_degree[nxt] == 0:                    queue.append(nxt)        return order     def _get_input_names(self, node):        """获取节点输入参数名列表（按槽位顺序）"""        names = []        input_types = node.INPUT_TYPES()        for name in input_types.get("required", {}):            names.append(name)        return names     def _get_defaults(self, node):        """为未连接的输入参数生成默认值"""        defaults = {}        input_types = node.INPUT_TYPES()        for name, spec in input_types.get("required", {}).items():            type_name = spec[0] if isinstance(spec, tuple) else spec            if type_name == "STRING":                defaults[name] = "default_model"            elif type_name == "INT":                defaults[name] = 42            elif type_name == "FLOAT":                defaults[name] = 7.5            elif type_name == "CONDITIONING":                defaults[name] = "empty_cond"            elif type_name == "LATENT":                defaults[name] = {"samples": "empty_latent"}        return defaults     def execute(self):        order = self.get_execution_order()        outputs = {}        for nid in order:            node = self.nodes[nid]            inputs = self._get_defaults(node)            param_names = self._get_input_names(node)            for conn in self.connections:                if conn["to"] == nid:                    src_out = outputs[conn["from"]]                    slot_name = param_names[conn["to_slot"]]                    inputs[slot_name] = src_out[conn["from_slot"]]            func_name = node.FUNCTION()            method = getattr(node, func_name)            result = method(**inputs)            outputs[nid] = result        return outputs  if __name__ == "__main__":    print("=== ComfyUI节点式工作流模拟 ===\n")     engine = ExecutionEngine()     loader = CheckpointLoaderSimple()    engine.add_node("1", loader)     sampler = KSampler()    engine.add_node("2", sampler)     decoder = VAEDecode()    engine.add_node("3", decoder)     preview = PreviewImage()    engine.add_node("4", preview)     engine.add_connection("1", 0, "2", 0)    engine.add_connection("1", 2, "3", 1)    engine.add_connection("2", 0, "3", 0)    engine.add_connection("3", 0, "4", 0)     print("执行顺序:", engine.get_execution_order())    results = engine.execute()    print("节点1输出:", results["1"])    print("节点2输出:", results["2"])    print("节点3输出:", results["3"])     print("\n=== 自定义节点: 超分辨率 ===\n")     class UpscaleNode(Node):        """自定义超分辨率节点，将图像分辨率提升2倍"""        CATEGORY = "custom"         @classmethod        def INPUT_TYPES(cls):            return {"required": {"image": ("IMAGE",)}}         @classmethod        def RETURN_TYPES(cls):            return ("IMAGE",)         def execute(self, image):            return (f"upscaled_2x({image})",)     engine2 = ExecutionEngine()    engine2.add_node("1", CheckpointLoaderSimple())    engine2.add_node("2", KSampler())    engine2.add_node("3", VAEDecode())    engine2.add_node("4", UpscaleNode())     engine2.add_connection("1", 0, "2", 0)    engine2.add_connection("1", 2, "3", 1)    engine2.add_connection("2", 0, "3", 0)    engine2.add_connection("3", 0, "4", 0)     results2 = engine2.execute()    print("超分输出:", results2["4"])    print("执行路径: 1 -> 2 -> 3 -> 4")