Extensions/ComfyUiRagCustomNodes
ComfyUI Extension

ComfyUiRagCustomNodes

This package contains custom nodes for ComfyUI: Image generation using Stable Diffusion and LLM, Uploading and searching images via ChromaDB and MinIO (S3), Integration with CLIP, Ollama, and other tools

By Santat2023·Created 12 months ago·Updated 29 days ago· 0
Santat2023/ComfyUiRagNodes
Nodes2
On cloudLocal install
CategoryMyNodes
Stars0
Updated29 days ago

Nodes (2)

DB Load Node

Your reference image, but found by vector search instead of a folder

MyNodes
LLM Node

One rough idea in, an enhanced prompt plus two search queries out

MyNodes
Readme

ComfyUiRagNodes

Кастомные ноды для ComfyUI, реализующие прототип локальной RAG-системы для работы иллюстратора.

Проект является частью дипломной работы и предназначен для исследования применения Retrieval-Augmented Generation (RAG) совместно с локально запускаемыми генеративными моделями.

Основная идея системы — по текстовому запросу пользователя автоматически обработать промпт, сформировать специализированные запросы для поиска референсов и получить изображения из локальной коллекции через Qdrant и MinIO/S3.

Возможности

Система позволяет:

  • принимать исходный промпт пользователя непосредственно в ComfyUI;
  • обрабатывать промпт с помощью LLM;
  • получать итоговый улучшенный промпт для генерации;
  • выделять из исходного запроса описание позы;
  • выделять описание стиля;
  • использовать эти описания для семантического поиска референсов;
  • искать изображения в локальной базе Qdrant;
  • загружать найденные изображения из MinIO/S3;
  • передавать найденные изображения дальше в граф ComfyUI;
  • использовать отдельные референсы для ControlNet и IP-Adapter.

Общая схема:

                         Исходный промпт
                                │
                                ▼
                         ┌─────────────┐
                         │   LLM Node  │
                         └──────┬──────┘
                                │
                 ┌──────────────┼──────────────┐
                 │              │              │
                 ▼              ▼              ▼
          Итоговый промпт   Pose query     Style query
                 │              │              │
                 │              ▼              ▼
                 │         ┌─────────┐    ┌─────────┐
                 │         │ Qdrant  │    │ Qdrant  │
                 │         └────┬────┘    └────┬────┘
                 │              │              │
                 │              ▼              ▼
                 │        Pose image      Style image
                 │              │              │
                 │              ▼              ▼
                 │         ControlNet      IP-Adapter
                 │              │              │
                 └──────────────┴──────────────┘
                                │
                                ▼
                         Генеративная модель
                                │
                                ▼
                         Итоговое изображение

Архитектура

Проект разделён на несколько уровней:

ComfyUI
   │
   ├── nodes/
   │     ├── llm_node.py
   │     └── db_load_node.py
   │
   ├── services/
   │     ├── llm_service.py
   │     └── retrieval_service.py
   │
   ├── repositories/
   │     ├── qdrant_repo.py
   │     └── s3_repo.py
   │
   ├── models/
   │     └── clip_model.py
   │
   ├── config.example.json
   └── __init__.py

nodes/

Содержит непосредственно кастомные ноды ComfyUI.

llm_node.py

Нода обработки пользовательского промпта.

Она принимает исходный текст и CLIP ComfyUI и возвращает результаты, которые могут использоваться дальше в графе.

Основная задача:

Исходный промпт
      │
      ▼
     LLM
      │
      ├──► итоговый промпт
      ├──► описание позы
      └──► описание стиля

db_load_node.py

Нода поиска и загрузки изображения из базы референсов.

Она использует RetrievalService, выполняющий поиск по Qdrant и получение файла из MinIO/S3.

services/

Содержит основную бизнес-логику приложения.

llm_service.py

Отвечает за работу с LLM.

Сервис используется нодой LLM_Node и скрывает детали взаимодействия с конкретным LLM-провайдером.

В зависимости от конфигурации может использоваться локальная Ollama или внешний OpenAI-compatible API.

Основная задача:

Промпт пользователя
        │
        ▼
    LLMService
        │
        ├──► итоговый промпт
        ├──► pose query
        └──► style query

retrieval_service.py

Координирует процесс поиска изображения.

Основной процесс:

Текстовый запрос
       │
       ▼
      CLIP
       │
       ▼
   embedding
       │
       ▼
     Qdrant
       │
       ▼
  metadata / key
       │
       ▼
    MinIO/S3
       │
       ▼
  bytes изображения

repositories/

Содержит интерфейсы взаимодействия с внешними системами хранения.

qdrant_repo.py

Отвечает за работу с Qdrant.

Используется для:

  • получения списка коллекций;
  • поиска ближайших векторов;
  • получения информации о найденных объектах.

Qdrant используется как векторная база данных изображений.

s3_repo.py

Отвечает за получение файлов из MinIO/S3.

Qdrant хранит вектор и метаданные, а само изображение хранится в объектном хранилище.

Таким образом:

Qdrant
  ├── embedding
  └── metadata
         │
         └──► ссылка / ключ изображения
                    │
                    ▼
                 MinIO
                    │
                    ▼
              image bytes

models/

clip_model.py

Содержит работу с CLIP.

CLIP используется для преобразования текстового запроса в embedding, совместимый с векторным представлением изображений в Qdrant.

Пример:

"женщина стоит, одна рука на бедре"
              │
              ▼
             CLIP
              │
              ▼
        text embedding
              │
              ▼
            Qdrant
              │
              ▼
     наиболее подходящий
       референс позы

Конфигурация

В репозитории присутствует файл:

config.example.json

Он является примером конфигурации.

Перед использованием рекомендуется создать собственный:

config.json

и указать необходимые параметры.

Пример

{
  "llm": {
    "provider": "Ollama",

    "ollama": {
      "base_url": "http://localhost:11434",
      "model": "qwen2:7b"
    },

    "openai": {
      "base_url": "https://api.openai.com/v1",
      "api_key": "",
      "model": ""
    },

    "yandex": {
      "base_url": "https://ai.api.cloud.yandex.net/v1",
      "api_key": "",
      "model": ""
    }
  }
}

Не следует помещать реальные API-ключи в Git-репозиторий.

Ollama

Для локальной работы с LLM рекомендуется использовать Ollama.

После установки Ollama необходимо загрузить используемую модель.

Например:

ollama pull qwen2:7b

После этого Ollama должна быть доступна по адресу:

http://localhost:11434

Конкретная модель задаётся в config.json.

Преимущество такого подхода для проекта — возможность выполнять LLM-обработку локально без обязательного использования платных внешних API.

Qdrant

Qdrant используется как векторная база данных для коллекций изображений.

Типичный локальный адрес:

http://localhost:6333

Коллекции создаются внешним приложением ImageManager.

Векторный поиск выполняется следующим образом:

Запрос пользователя
        │
        ▼
       CLIP
        │
        ▼
  Text embedding
        │
        ▼
      Qdrant
        │
        ▼
  Similarity Search
        │
        ▼
  найденный объект

MinIO

MinIO используется как локальное S3-совместимое хранилище изображений.

Типичная архитектура:

Qdrant
 │
 ├── vector
 └── metadata
       │
       └── image key
              │
              ▼
            MinIO
              │
              ▼
        исходное изображение

Такое разделение позволяет не хранить бинарные изображения непосредственно в векторной базе.

Подготовка коллекций

Для создания и управления коллекциями используется отдельное приложение:

ImageManager

Репозиторий:

https://github.com/Santat2023/ImageManager

ImageManager предназначен для:

  • загрузки изображений;
  • формирования коллекций;
  • генерации описаний;
  • создания CLIP-векторов;
  • сохранения изображений в MinIO/S3;
  • сохранения векторов и метаданных в Qdrant.

В текущей версии ImageManager также может использовать дополнительные ручные описания изображений.

Например:

images/
├── image01.jpg
├── image01.txt
├── image02.jpg
└── image03.jpg

Если рядом с изображением существует .txt с таким же именем, его содержимое может использоваться как дополнительное описание изображения.

Это особенно полезно для коллекций поз и композиции, поскольку автоматическое captioning не всегда достаточно точно описывает положение тела, направление движения и композицию кадра.

Поиск референсов

В проекте используется семантический поиск.

Например, пользователь задаёт:

Девушка-воин стоит на вершине скалы на закате,
в драматической позе, кинематографический стиль.

LLM может сформировать специализированные запросы:

Pose:

female warrior standing on a rock,
dynamic heroic pose, full body,
one leg forward, arms positioned dramatically

Style:

cinematic fantasy illustration,
dramatic sunset lighting,
high contrast, detailed digital art

Далее эти запросы независимо отправляются в поиск.

Pose query
    │
    ▼
 Qdrant
    │
    ▼
Pose reference

и:

Style query
    │
    ▼
 Qdrant
    │
    ▼
Style reference

Полученные изображения могут использоваться независимо:

Pose reference ───► ControlNet

Style reference ──► IP-Adapter

Пример графа ComfyUI

Концептуально граф может выглядеть следующим образом:

                   ┌───────────────┐
                   │  User Prompt  │
                   └───────┬───────┘
                           │
                           ▼
                   ┌───────────────┐
                   │    LLM Node   │
                   └───────┬───────┘
                           │
             ┌─────────────┼─────────────┐
             │             │             │
             ▼             ▼             ▼
        Conditioning   Pose Query    Style Query
             │             │             │
             │             ▼             ▼
             │        ┌─────────┐   ┌─────────┐
             │        │ DB Load │   │ DB Load │
             │        │  Pose   │   │  Style  │
             │        └────┬────┘   └────┬────┘
             │             │             │
             │             ▼             ▼
             │        Pose Image    Style Image
             │             │             │
             │             ▼             ▼
             │         ControlNet    IP-Adapter
             │             │             │
             └─────────────┴──────┬──────┘
                                  │
                                  ▼
                         Stable Diffusion
                                  │
                                  ▼
                         Generated Image

Установка

1. Клонирование

Репозиторий необходимо разместить в каталоге custom_nodes ComfyUI:

cd ComfyUI/custom_nodes

git clone https://github.com/Santat2023/ComfyUiRagNodes.git

В результате должна получиться структура:

ComfyUI/
└── custom_nodes/
    └── ComfyUiRagNodes/
        ├── models/
        ├── nodes/
        ├── repositories/
        ├── services/
        ├── config.example.json
        ├── __init__.py
        └── README.md

2. Зависимости

Необходимы компоненты, используемые проектом:

  • Python;
  • PyTorch;
  • CLIP;
  • Qdrant Client;
  • boto3 / S3-compatible client;
  • Pillow;
  • NumPy;
  • Ollama или совместимый LLM API;
  • ComfyUI.

Конкретные версии зависят от версии ComfyUI и окружения.

Запуск

  1. Запустить Qdrant.
  2. Запустить MinIO.
  3. Запустить Ollama, если используется локальная LLM.
  4. Убедиться, что коллекции изображений созданы через ImageManager.
  5. Запустить ComfyUI.
  6. Перейти в интерфейс ComfyUI.
  7. Найти ноды категории:
MyNodes

Устранение проблем

Ноды не появились в ComfyUI

Проверьте структуру:

ComfyUI/
└── custom_nodes/
    └── ComfyUiRagCustomNodes/
        ├── __init__.py
        ├── nodes/
        ├── services/
        ├── repositories/
        └── models/

После изменения Python-кода необходимо перезапустить ComfyUI.

Qdrant недоступен

Проверьте, что Qdrant запущен и доступен по адресу:

http://localhost:6333

Также проверьте наличие необходимых коллекций.

MinIO недоступен

Проверьте:

  • запущен ли MinIO;
  • корректность endpoint;
  • access key;
  • secret key;
  • название bucket;
  • наличие изображения, указанного в metadata Qdrant.

Ollama недоступна

Проверьте:

http://localhost:11434

и наличие модели:

ollama list

При необходимости:

ollama pull qwen2:7b

Структура репозитория

ComfyUiRagNodes/
│
├── models/
│   └── clip_model.py
│
├── nodes/
│   ├── db_load_node.py
│   └── llm_node.py
│
├── repositories/
│   ├── qdrant_repo.py
│   └── s3_repo.py
│
├── services/
│   ├── llm_service.py
│   └── retrieval_service.py
│
├── config.example.json
├── __init__.py
├── .gitignore
└── README.md

Принцип разделения ответственности

В проекте используется разделение компонентов по ответственности.

Ноды

Отвечают за интеграцию с ComfyUI:

  • входные параметры;
  • выходные параметры;
  • типы данных ComfyUI;
  • соединение компонентов в workflow.

Сервисы

Отвечают за основную логику:

  • работу с LLM;
  • обработку промптов;
  • организацию процесса поиска.

Репозитории

Отвечают за инфраструктуру:

  • Qdrant;
  • MinIO/S3.

Модели

Отвечают за ML-компоненты:

  • CLIP;
  • преобразование текста в embedding.

Такое разделение позволяет заменить отдельный компонент без необходимости переписывать остальные части системы.

Исследовательская часть

Репозиторий является программной частью дипломного проекта, посвящённого применению RAG для поддержки работы иллюстратора с использованием локально запускаемых генеративных моделей.

Исследуемый подход отличается от простого поиска похожих изображений тем, что исходный запрос пользователя сначала анализируется LLM и разделяется на функционально различные компоненты:

Исходный запрос
       │
       ▼
      LLM
       │
       ├──► генеративный prompt
       │
       ├──► описание позы
       │
       └──► описание стиля
                    │
          ┌─────────┴─────────┐
          ▼                   ▼
     поиск позы          поиск стиля
          │                   │
          ▼                   ▼
     ControlNet           IP-Adapter

Таким образом, RAG используется не только для получения общего контекста, а для формирования специализированных визуальных условий генерации.

Связанные проекты

ImageManager

Приложение для формирования и управления локальными коллекциями референсов:

https://github.com/Santat2023/ImageManager

ComfyUI

https://github.com/Comfy-Org/ComfyUI

Qdrant

https://qdrant.tech/

MinIO

https://min.io/

Ollama

https://ollama.com/

Назначение проекта

Проект представляет собой proof of concept локальной системы RAG для генерации изображений.

Целевая архитектура:

┌─────────────────────────────────────────────────────┐
│                    Иллюстратор                      │
└─────────────────────────┬───────────────────────────┘
                          │
                          ▼
                  Текстовый промпт
                          │
                          ▼
                 ┌─────────────────┐
                 │       LLM       │
                 └────────┬────────┘
                          │
             ┌────────────┼────────────┐
             │            │            │
             ▼            ▼            ▼
          Prompt         Pose         Style
             │            │            │
             │            └─────┬──────┘
             │                  │
             │             ┌────▼────┐
             │             │ Qdrant  │
             │             └────┬────┘
             │                  │
             │             ┌────▼────┐
             │             │ MinIO   │
             │             └────┬────┘
             │                  │
             │        ┌─────────┴─────────┐
             │        │                   │
             │        ▼                   ▼
             │    ControlNet          IP-Adapter
             │        │                   │
             └────────┴─────────┬─────────┘
                                │
                                ▼
                       Генеративная модель
                                │
                                ▼
                       Итоговая иллюстрация

Основная цель proof of concept — экспериментально проверить, насколько использование локальной базы визуальных референсов и специализированного RAG-поиска может улучшить управляемость процесса генерации иллюстраций по сравнению с генерацией только по текстовому промпту.