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 11 months ago·Updated 9 days ago· 0
Santat2023/ComfyUiRagNodes
Nodes2
On cloudLocal install
CategoryMyNodes
Stars0
Updated9 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-поиска может улучшить управляемость процесса генерации иллюстраций по сравнению с генерацией только по текстовому промпту.