GTE-Base-ZH模型Docker部署与ComfyUI可视化流程搭建

最近在折腾一些文本相关的AI应用，发现很多场景都需要把文字转换成计算机能理解的“向量”，也就是做文本嵌入。试了几个开源模型，感觉GTE-Base-ZH在中文处理上效果挺扎实，就想把它用起来。但光有个模型服务还不够，如果能和现在流行的ComfyUI可视化工作流结合起来，比如把一段描述文字变成向量，再用这个向量去影响图片生成，那玩法就多了。

今天就来聊聊怎么用Docker Compose把GTE-Base-ZH模型服务稳稳当当地跑起来，再把它“嫁接”到ComfyUI里，变成一个可以拖拽的节点。整个过程不算复杂，但有些细节需要注意，我会把踩过的坑和解决方法都写清楚。

1. 准备工作与环境概览

在开始动手之前，我们先理清一下要做什么，以及需要准备哪些东西。你不是在部署一个孤零零的模型，而是在搭建一个小型的服务生态。

首先，你需要有一台带GPU的Linux服务器，这是保证向量生成速度的基础。内存建议16GB以上，硬盘空间留出20GB左右。操作系统我用的Ubuntu 22.04，其他发行版也大同小异。

核心的组件有三个：

1. GTE-Base-ZH模型服务：这是主角，负责接收中文文本，吐出对应的向量。
2. ComfyUI：可视化工作流编辑器，我们将把GTE服务集成进去。
3. Docker与Docker Compose：用来管理和编排上面这些服务，让它们能一起愉快地工作。

整个方案的思路是，用Docker Compose定义一个配置文件，把GTE模型服务和其他可能需要的辅助服务（比如数据库）一次性启动。然后，我们为ComfyUI编写一个自定义节点，这个节点会去调用我们部署好的GTE服务，完成文本到向量的转换。最后，在ComfyUI的画布上，你就可以像搭积木一样，把文本向量节点和Stable Diffusion的图片生成节点连起来，创造出更可控的生成效果。

2. 使用Docker Compose部署GTE-Base-ZH服务

Docker Compose的好处是能把多个服务的关系定义清楚，一键启停，非常方便。我们首先来部署最核心的GTE模型服务。

2.1 编写Docker Compose配置文件

在你的项目根目录下，创建一个名为 docker-compose.yml 的文件。这个文件就像乐高说明书，告诉Docker如何组装我们的服务。

version: '3.8'

services:
  gte-service:
    image: your-registry/gte-base-zh:latest # 假设你已构建好镜像
    container_name: gte-embedding-service
    restart: unless-stopped
    ports:
      - "8000:8000" # 将容器内的8000端口映射到主机
    environment:
      - MODEL_NAME=GTE-Base-ZH
      - DEVICE=cuda # 如果使用GPU
    volumes:
      - ./model_cache:/app/model_cache # 挂载缓存目录，避免重复下载模型
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu] # 声明需要GPU资源
    networks:
      - ai-network

  # 可选：可以在这里添加一个向量数据库，如Qdrant，用于存储和检索生成的向量
  # vector-db:
  #   image: qdrant/qdrant:latest
  #   container_name: qdrant-db
  #   restart: unless-stopped
  #   ports:
  #     - "6333:6333"
  #   volumes:
  #     - ./qdrant_storage:/qdrant/storage
  #   networks:
  #     - ai-network

networks:
  ai-network:
    driver: bridge

几点关键说明：

• image：你需要先构建一个包含GTE-Base-ZH模型和推理代码的Docker镜像，并推送到镜像仓库。或者，可以直接使用一些社区提供的镜像。
• ports：我们把服务暴露在主机的8000端口，这样ComfyUI或者其他应用都能访问。
• volumes：挂载一个本地目录到容器的模型缓存路径，这样下次重启容器时，就不用重新下载巨大的模型文件了。
• deploy.resources：这部分配置是告诉Docker Compose我们需要GPU。确保你的Docker已正确配置NVIDIA Container Toolkit。
• networks：创建一个独立的网络，让多个服务在内部能通过服务名互相访问，比如ComfyUI容器可以通过 http://gte-service:8000 访问GTE服务。

2.2 准备模型服务代码与Dockerfile

光有编排文件还不够，我们需要GTE模型服务本身的代码和Docker镜像。通常，这个服务会提供一个HTTP API，接收文本，返回向量。

创建一个 gte_service 目录，里面至少包含：

• app/main.py：使用FastAPI或Flask编写的API服务代码。
• requirements.txt：Python依赖列表。
• Dockerfile：构建镜像的指令。

一个极简的 Dockerfile 示例：

FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

# 假设你的启动命令是运行main.py
CMD ["python", "app/main.py"]

对应的 app/main.py 核心部分可能长这样（使用FlagEmbedding库）：

from fastapi import FastAPI
from pydantic import BaseModel
from FlagEmbedding import FlagModel
import torch

app = FastAPI()

# 加载模型，这里假设模型已提前下载到指定路径
model = FlagModel('BAAI/bge-large-zh-v1.5', 
                  query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
                  use_fp16=True) # 使用半精度加快推理

class TextRequest(BaseModel):
    texts: list[str]
    normalize_embeddings: bool = True

@app.post("/embed")
async def create_embeddings(request: TextRequest):
    """接收文本列表，返回向量列表"""
    with torch.no_grad():
        embeddings = model.encode(request.texts, 
                                  normalize_embeddings=request.normalize_embeddings)
    # 将numpy数组转换为列表
    embeddings_list = embeddings.tolist()
    return {"embeddings": embeddings_list}

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

2.3 启动服务与验证

编写好所有文件后，在包含 docker-compose.yml 的目录下，执行一条命令即可启动所有服务：

docker-compose up -d

-d 参数表示在后台运行。用下面的命令查看服务状态和日志：

docker-compose ps
docker-compose logs -f gte-service # 查看GTE服务日志

服务启动后，打开浏览器或使用 curl 测试一下API是否正常：

curl -X POST http://localhost:8000/embed \
  -H "Content-Type: application/json" \
  -d '{"texts": ["今天天气真好", "人工智能是未来科技的核心"], "normalize_embeddings": true}'

如果返回一串长长的数字向量，恭喜你，GTE服务部署成功了！

3. 开发ComfyUI自定义节点调用向量服务

现在，我们的GTE服务已经在8000端口待命了。接下来，我们要让ComfyUI认识它。这需要通过开发一个自定义节点来实现。

3.1 创建ComfyUI节点文件

在ComfyUI的 custom_nodes 目录下，新建一个文件夹，比如叫 comfyui_gte_node。在里面创建 __init__.py 和 nodes.py 文件。

nodes.py 是这个自定义节点的核心，它定义了节点的输入、输出和功能逻辑。

import comfy.model_management
import comfy.utils
import torch
import numpy as np
import requests
import json
from typing import List

# 一个简单的文本嵌入节点
class GTE_TextEmbedding:
    """
    调用远程GTE服务，将中文文本转换为嵌入向量。
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {"multiline": True, "default": "请输入中文文本"}),
                "api_url": ("STRING", {"default": "http://localhost:8000/embed"}),
                "normalize": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ("EMBEDS",) # 定义输出类型为EMBEDS
    RETURN_NAMES = ("embeddings",)
    FUNCTION = "encode"
    CATEGORY = "embedding" # 在ComfyUI的节点菜单中分类

    def encode(self, text, api_url, normalize):
        # 准备请求数据
        payload = {
            "texts": [text], # 目前处理单条文本，可扩展为多条
            "normalize_embeddings": normalize
        }
        
        try:
            response = requests.post(api_url, json=payload, timeout=30)
            response.raise_for_status() # 检查HTTP错误
            result = response.json()
            embeddings_list = result.get("embeddings", [])
            
            if not embeddings_list:
                raise ValueError("API返回的嵌入向量为空")
                
            # 将列表转换为PyTorch张量，ComfyUI内部处理多用张量
            embeddings_np = np.array(embeddings_list[0], dtype=np.float32)
            embeddings_tensor = torch.from_numpy(embeddings_np).unsqueeze(0) # 增加批次维度
            
            return (embeddings_tensor,)
            
        except requests.exceptions.RequestException as e:
            comfy.utils.log_message(f"调用GTE API失败: {e}")
            # 返回一个零向量作为降级处理，避免工作流完全中断
            dummy_embed = torch.zeros((1, 1024)) # 假设维度是1024
            return (dummy_embed,)
        except (KeyError, ValueError, json.JSONDecodeError) as e:
            comfy.utils.log_message(f"解析API响应失败: {e}")
            dummy_embed = torch.zeros((1, 1024))
            return (dummy_embed,)

# 一个更复杂的节点示例：将文本嵌入与条件生成结合
class TextToConditioningWithGTE:
    """
    结合GTE文本嵌入和CLIP文本编码，生成用于图像生成的混合条件。
    这是一个示意性节点，展示如何将向量用于下游任务。
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {"multiline": True}),
                "gte_api_url": ("STRING", {"default": "http://localhost:8000/embed"}),
                "clip": ("CLIP",), # 连接一个CLIP文本编码器节点
            },
        }

    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION = "encode"
    CATEGORY = "conditioning"

    def encode(self, text, gte_api_url, clip):
        # 1. 获取GTE文本嵌入
        gte_payload = {"texts": [text], "normalize_embeddings": True}
        try:
            gte_resp = requests.post(gte_api_url, json=gte_payload, timeout=15)
            gte_embedding = torch.from_numpy(np.array(gte_resp.json()["embeddings"][0]))
        except Exception as e:
            print(f"GTE嵌入失败，使用备用方案: {e}")
            gte_embedding = torch.randn(1024) # 备用随机向量

        # 2. 获取CLIP文本编码 (这里简化处理，实际需适配ComfyUI的CLIP接口)
        # 假设clip对象有encode方法
        clip_encoding = clip.encode(text)
        
        # 3. 这里可以进行向量融合操作，例如拼接或加权平均
        # combined_conditioning = some_fusion_function(gte_embedding, clip_encoding)
        # 此处仅为示例，直接返回CLIP编码
        # 你需要根据实际模型输入要求设计融合逻辑
        
        return (clip_encoding, ) # 实际应返回融合后的conditioning

# 告诉ComfyUI有哪些节点
NODE_CLASS_MAPPINGS = {
    "GTE_TextEmbedding": GTE_TextEmbedding,
    "TextToConditioningWithGTE": TextToConditioningWithGTE,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "GTE_TextEmbedding": "GTE文本嵌入",
    "TextToConditioningWithGTE": "文本转条件 (GTE增强)",
}

3.2 安装依赖与注册节点

在自定义节点目录下创建 requirements.txt，列出必要的Python包，比如 requests。

为了让ComfyUI加载这个节点，通常需要在 __init__.py 中导入 nodes.py 并暴露映射关系。更常见的做法是，ComfyUI会自动扫描 custom_nodes 目录下的子文件夹。确保你的节点文件结构符合要求后，重启ComfyUI。

重启后，在ComfyUI的节点搜索框中输入“GTE”或“嵌入”，应该就能找到你新添加的节点了。

4. 在ComfyUI中构建可视化工作流

节点准备好之后，就可以开始搭积木了。这里我设计两个简单的工作流示例，展示如何将文本向量用于图像生成。

4.1 基础工作流：文本嵌入展示

这个工作流很简单，就是为了验证节点是否正常工作。

1. 在画布上，右键搜索并添加 GTE文本嵌入 节点。
2. 在节点的 text 输入框里，写上你想测试的中文句子，比如“一只在星空下奔跑的狐狸”。
3. 确保 api_url 指向你运行GTE服务的地址（如果ComfyUI和GTE服务不在同一台机器，需要改成IP地址）。
4. 连接一个 预览图像 或 调试输出 节点到 GTE文本嵌入 节点的输出端。由于输出是向量数据，预览节点可能无法直接显示，你可以连接一个 保存数据 节点，将向量保存为 .npy 或 .pt 文件，然后用Python脚本加载查看。或者，ComfyUI社区有一些专门用于可视化高维数据的节点。

这个流程主要目的是测试连通性，确保点击“执行”后，节点能成功调用远程API并返回数据，不报错。

4.2 进阶工作流：结合文生图与向量控制

这个工作流更有趣一些，我们尝试用文本向量来影响图像生成。思路是：将GTE生成的向量，通过某种方式“注入”到Stable Diffusion的生成过程中。一种常见的方法是使用 “自定义条件”。

请注意：Stable Diffusion的标准流程通常使用CLIP的文本编码作为条件。直接注入另一个模型的向量，需要模型本身支持多模态条件，或者你需要使用支持额外条件输入的LoRA、Textual Inversion等适配器。以下是一个概念性流程：

1. 文本输入：添加一个 文本输入 节点，写入详细的中文描述。
2. 生成文本嵌入：添加你的 GTE文本嵌入 节点，接收上一步的文本，生成向量。
3. CLIP文本编码：同时，添加标准的 CLIP文本编码器 节点，也接收同样的文本描述，生成CLIP向量。
4. 向量融合/处理：添加一个自定义的 向量处理 节点（这需要你额外开发）。这个节点的功能可能是：
   • 拼接：将GTE向量和CLIP向量拼接成一个更长的向量。
   • 加权求和：将两个向量按权重相加。
   • 作为附加输入：将GTE向量作为 KSampler 节点的 positive 或 negative 条件的一个额外部分。有些高级采样器或自定义节点支持传入额外的潜在表示。
5. 连接至采样器：将处理后的条件向量，连接到 KSampler 节点的 positive 输入。
6. 执行生成：配置好模型、VAE、采样步数等参数，点击执行。

这个工作流的关键在于第4步的“向量处理”节点。你需要明确GTE向量如何参与到扩散模型的生成过程中。这可能涉及到：

• 使用一个预先训练好的、能理解GTE向量空间的 “适配层” （一个小神经网络），将GTE向量映射到CLIP的潜在空间。
• 利用 ComfyUI的“条件组合”节点，尝试将GTE向量作为一种额外的条件与CLIP条件进行组合。

这属于比较进阶的用法，可能需要一些实验和调试。一个更简单直接的落地思路是：用GTE向量进行图像检索或排序。例如，你可以用GTE向量为一批生成的图片计算相似度，筛选出与文本语义最匹配的几张图，这在构建图像生成流水线时非常实用。

5. 总结

走完这一趟，你应该已经成功地把GTE-Base-ZH模型封装成了独立的Docker服务，并且为ComfyUI打造了一个可以调用的自定义节点。从部署到集成，整个过程涉及了服务编排、API调用和可视化编程，算是把AI应用开发的几个关键环节串了起来。

部署方面，Docker Compose让服务管理变得清晰简单，特别是未来如果你想加入向量数据库、缓存服务等其他组件，只需要在YAML文件里继续添加就行。节点开发部分，虽然看起来代码不少，但核心逻辑就是发送HTTP请求和处理返回结果，模式是固定的，以后集成其他模型服务也可以照葫芦画瓢。

最有意思的还是ComfyUI工作流的构建。把文本向量化只是第一步，如何让这个向量在图像生成中发挥作用，才是开放性的挑战。我上面提到的向量融合方法只是一个起点，你可以探索更多方式，比如用向量来控制生成风格、构图比例，或者作为搜索条件来管理你的图库。

在实际使用中，你可能会遇到网络延迟、向量维度不匹配、服务稳定性等问题。对于延迟，可以考虑在ComfyUI节点里加入简单的缓存机制；对于稳定性，可以在Docker Compose里配置健康检查和重启策略。多试试，多调调，这套组合的灵活性很高，能玩出很多花样。

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。