ComfyUI 源代码阅读:程序启动

近期接触 AIGC 方面的内容,项目中用到 ComfyUI 这个组件,这个系列文章主要是阅读 ComfyUI 的源代码,尝试分析主要的执行流程。

启动参数

ComfyUI 支持的启动命令比较多,在启动前会先解析启动参数,具体的启动参数和说明整理如下:

image.png

程序启动流程

comfyui-startup.png

主要的启动流程就如上图所示,我们分成不同的模块来具体的分析:

执行启动前的脚本

ComfyUI 启动时会扫描 models 文件夹和 custom_nodes 文件夹,将文件名和路径、文件扩展类型的映射存储在全局变量 folder_names_and_paths 中。

启动前从这个变量中查询到 custom_nodes 的节点路径,找到自定义节点中的前置启动脚本:prestartup_script.py。通过动态加载执行这个脚本文件,并且保存执行信息。

配置 cuda

启动流程的第二步,会设置三个cuda相关的配置:

  1. CUDA_VISIBLE_DEVICES:为当前的实例绑定一个指定的 cuda 设备,通常是一个数字。
  2. CUBLAS_WORKSPACE_CONFIG:用于配置NVIDIA的cuBLAS库。cuBLAS库是一个GPU-accelerated的基础线性代数库,用于在NVIDIA的CUDA平台上进行高性能的矩阵运算。CUBLAS_WORKSPACE_CONFIG的值用于控制cuBLAS库的工作空间配置。它的值是一个由两部分组成的字符串,格式为"minimumSize:maximumSize",其中minimumSizemaximumSize都是以字节为单位的整数。这里会默认设置成 4096字节,即4KB。
  3. PYTORCH_CUDA_ALLOC_CONF:设置 cudaMallocAsync,启用异步分配显存的能力,提高GPU内存分配的效率。

创建 http 服务和任务队列

ComfyUI 中封装了一个类用来表示整个http服务,包括两种类型的参数:

  1. web 相关:MIME、中间件(权限、cors等)、最大上传大小、前端文件、路由链接。
  2. 任务执行相关:任务队列、事件循环、异步消息队列,任务数、最后运行的节点id、客户端的id。

创建的过程就是创建一个 PromptServer 的类,初始化上面的两种类型的参数。任务队列是另一个类 PromptQueue,主要包含的都是任务相关的参数。

从外部文件加载模型路径

ComfyUI 是支持复用 webUI 的模型的,实现的方式即通过一个配置文件:extra_model_paths.yaml 来实现,如果在该配置文件中定义了 webUI 模型的路径,ComfyUI 在启动时会将这些路径添加到前面提到的全局变量 folder_names_and_paths 中。

#config for a1111 ui  
#all you have to do is change the base_path to where yours is installed  
a111:  
    base_path: path/to/stable-diffusion-webui/  
  
    checkpoints: models/Stable-diffusion  
    configs: models/Stable-diffusion  
    vae: models/VAE  
    loras: |  
         models/Lora  
         models/LyCORIS  
    upscale_models: |  
                  models/ESRGAN  
                  models/RealESRGAN  
                  models/SwinIR  
    embeddings: embeddings  
    hypernetworks: models/hypernetworks  
    controlnet: models/ControlNet

初始化自定义节点

节点的来源有三种:

  1. comfyui 自带的基础节点,实现在 nodes.py 文件中。
  2. comfy_extras 文件夹中以 nodes_xxx 开头的文件加载的节点,这部分属于 comfyui 支持一些新特性添加的默认实现,例如 flux,sd3等
  3. 外部的节点,从 custom_nodes 文件夹加载,主要是开源的一些节点,例如 comfyui-manager

初始化自定义节点就是将定义在 comfy_extras 文件夹和 custom_nodes 文件夹的节点通过动态加载的方式加载到内存中,并将它们添加到 NODE_CLASS_MAPPINGS 表和NODE_DISPLAY_NAME_MAPPINGS 表中。

启动worker消费任务

这一步会创建并启动一个新的守护线程,这个线程将执行prompt_worker函数,函数的参数是qserver。 worker 中具体的消费逻辑,我们将在下一部分的内容中详细分析。

启动服务

ComfyUI 的 HTTP 服务是使用 aiohttp 库实现的,它是基于 Python 的 asyncio 库,可以用于处理并发的 HTTP 请求。

HTTP 服务启动时会启动两个协程,一个是启动 HTTP 服务,另一个是用于监听异步消息队列,并发送消息,实现 websocket 服务中的消息交互能力。

小结

那么到这里 ComfyUI 主要启动流程我们已经分析完毕了,从启动的代码可以发现,ComfyUI 的代码非常简洁,大量运用Python动态加载库来实现代码的动态加载,实现自定义节点的插件体系,方便扩展。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值