ComfyUI 图片转提示词扩展使用教程

本文档介绍如何在 ComfyUI 环境中安装和使用图片转提示词扩展，该扩展基于视觉语言模型实现图像到文本的转换功能。

1. 项目目录结构

扩展插件的目录组织方式如下：

ComfyUI_Image2Prompt/
├── custom_nodes/
│   └── Image2Prompt/
│       ├── __init__.py
│       ├── core.py
│       ├── settings.yml
│       └── model_weights/
│           ├── vision_encoder/
│           └── text_decoder/
├── docs/
│   └── tutorial.md
└── dependencies.txt

各目录和文件的作用说明：

• custom_nodes/：存放 ComfyUI 自定义节点的核心代码。
  • Image2Prompt/：插件主目录。
    • __init__.py：节点注册与初始化模块。
    • core.py：核心处理逻辑，包含图像编码和提示词生成。
    • settings.yml：运行时配置参数。
    • model_weights/：模型权重存储目录。
      • vision_encoder/：视觉编码器权重。
      • text_decoder/：文本解码器权重。
• docs/：项目文档目录。
• dependencies.txt：Python 依赖包列表。

2. 核心启动模块

core.py 是插件的核心模块，负责模型加载和图像处理。以下是模块的主要结构：

import yaml
import torch
from pathlib import Path

class Image2PromptEngine:
    def __init__(self, config_path):
        self.config = self._load_config(config_path)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = None
    
    def _load_config(self, path):
        with open(path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
    
    def initialize_model(self):
        model_name = self.config.get('model', {}).get('name', 'vision_model')
        model_path = Path(self.config['model']['weight_path'])
        print(f"正在加载模型: {model_name}")
        # 模型加载逻辑
        self.model = self._load_model_weights(model_path)
    
    def process_image(self, image_path):
        image_tensor = self._preprocess_image(image_path)
        prompt = self.model.generate(image_tensor)
        return prompt

def initialize_extension():
    engine = Image2PromptEngine("settings.yml")
    engine.initialize_model()
    return engine

if __name__ == "__main__":
    extension = initialize_extension()
    print("插件初始化完成")

核心类 Image2PromptEngine 提供了以下功能：

• __init__：初始化配置和计算设备。
• _load_config：读取 YAML 配置文件。
• initialize_model：加载预训练模型权重。
• process_image：处理输入图像并生成对应的提示词文本。

3. 配置文件说明

settings.yml 用于配置插件的运行参数。以下是配置文件的结构示例：

extension:
  name: Image2Prompt
  release: "1.0.0"
  model:
    name: vision_transformer
    weight_path: "./model_weights/vision_encoder/"
    device: "auto"
  generation:
    max_tokens: 150
    temperature: 0.7
    top_p: 0.9
  options:
    auto_load_weights: true
    cache_enabled: true

配置项详细说明：

• extension：插件基本信息。
  • name：插件名称。
  • release：版本号。
  • model：模型相关配置。
    • name：使用的模型名称。
    • weight_path：模型权重文件路径。
    • device：计算设备选择，可设为 auto、cuda 或 cpu。
  • generation：文本生成参数。
    • max_tokens：生成文本的最大token数量。
    • temperature：采样温度，控制随机性。
    • top_p：核采样参数。
  • options：可选功能开关。
    • auto_load_weights：是否在启动时自动加载模型权重。
    • cache_enabled：是否启用推理结果缓存。

  通过修改配置文件，用户可以灵活调整插件的行为，包括选择不同的模型、调整生成参数以及控制可选功能的开关。

  4. 使用示例

  在 ComfyUI 工作流中集成该插件的基本步骤：

  1. 将插件目录复制到 ComfyUI 的 custom_nodes 文件夹中。
  2. 启动 ComfyUI，插件会自动加载。
  3. 在节点菜单中找到「图像转提示词」节点。
  4. 连接图像输入节点，配置相应的参数。
  5. 运行工作流获取生成的提示词。