Extensions/comfyui-tag-rag
ComfyUI Extension

comfyui-tag-rag

ComfyUI custom node package for tag-based prompt RAG using CSV datasets and local LLM generation.

By FrostySDXL·Created 2 months ago·Updated 2 months ago· 0
FrostySDXL/comfyui-tag-rag
Nodes—
On cloudLocal install
Stars0
Updated2 months ago
Readme

comfyui-tag-rag

ComfyUI custom nodes for tag-based prompt RAG using CSV tag datasets, lexical/optional FAISS retrieval, and local llama.cpp generation.

Status

Experimental / Pre-Alpha. The current usable path is:

CSV tags -> TagRAG Index -> lexical/semantic retrieval -> TagRAG Prompt/Chat/Image -> native llama-cpp-python generation by default or alternate llama-server generation

Implemented now:

  • CSV loading with encoding fallback.
  • Normalized TagRecord / TagCatalog data model.
  • Lexical retrieval over tags, aliases, categories, descriptions, exclusions, and weights.
  • Basic TagRAG Index node that returns a serialized JSON index_ref.
  • Basic TagRAG Prompt node that reloads the CSV from index_ref, retrieves tags, builds a deterministic one-shot prompt, and calls native llama-cpp-python by default or a local llama-server endpoint when selected.
  • Native llama-cpp-python embeddings for semantic/hybrid retrieval when an embedding-capable GGUF is provided, plus alternate llama-server embeddings through /v1/embeddings or legacy /embedding.
  • Optional FAISS semantic indexes persisted beside the CSV under .tag_rag_indexes/<index_name>/.
  • Semantic and hybrid retrieval modes when faiss-cpu and an embedding endpoint are available.
  • Basic TagRAG Chat node with bounded in-memory conversation history and RAG retrieval per turn.
  • TagRAG Image node that accepts existing tag strings as the primary path, can optionally run built-in WD14-style ONNX tagging when image extras/model files are manually installed, retrieves related CSV tags, and optionally refines through llama-server.
  • Native llama-cpp-python in-process GGUF generation backend for Prompt, Chat, and Image refinement is the normal/default path; llama-server remains supported.

Future:

  • Advanced vision-language image conversation/collaboration workflows.

Goal

Help users craft better ComfyUI image-generation prompts by combining:

  • A personal CSV tag dataset.
  • Retrieval over that dataset.
  • Local llama.cpp-based text generation.
  • Optional chat and tag-to-prompt/image-to-tag workflows.

Quick Start: Lexical Retrieval + Prompt Generation

  1. Install the package into your ComfyUI custom_nodes/ directory.

  2. Install core requirements:

    pip install -r requirements.txt

  3. Provide an existing local generation GGUF through model_path for the default native backend. No model is downloaded automatically.

    Or start/point to an existing local llama.cpp server if you prefer generation_backend=llama_server:

    http://127.0.0.1:8080/v1/chat/completions

  4. Add TagRAG Index in ComfyUI.

    • Set csv_path to your CSV file.
    • Set tag_column to the column containing tag text.
    • Optional columns: aliases, category, description, source, weight.
    • The node returns a JSON index_ref string.

  5. Wire index_ref into TagRAG Prompt.

  6. Enter an idea, model_path, top_k, max_tokens, and temperature.

  7. Leave generation_backend=llama_cpp_python for the default path, or switch to llama_server and use base_url/model.

  8. Use the positive_prompt output in your ComfyUI generation workflow.

No FAISS or embedding model/server is required for the lexical retrieval path.

Optional: Semantic or Hybrid Retrieval

  1. Normal ComfyUI installs use requirements.txt, which now includes FAISS and native/image dependencies. If you maintain split installs, semantic retrieval needs:

    pip install -r requirements-rag.txt

  2. Choose an embedding backend in TagRAG Index.

    Native default for semantic/hybrid:

    embedding_backend: llama_cpp_python
embedding_model_path: path to an embedding-capable GGUF

    Alternate llama-server backend:

    OpenAI-compatible: http://127.0.0.1:8080/v1/embeddings
Legacy llama.cpp:  http://127.0.0.1:8080/embedding

  3. In TagRAG Index, set retrieval_mode to semantic or hybrid and enable rebuild when creating or refreshing the FAISS index.

  4. Wire the returned index_ref into TagRAG Prompt; leave retrieval_mode as from_index or override it explicitly.

If FAISS or the selected embedding backend is not available, lexical mode continues to work and semantic/hybrid modes return a clear status string.

Embedding model notes:

  • retrieval_mode=lexical requires no embedding model.
  • retrieval_mode=semantic or hybrid requires an embedding model/backend.
  • Native embeddings use embedding_model_path and require an embedding-capable GGUF.
  • Generation and embedding GGUFs can be different; do not assume a generation GGUF is a good embedding model.

CSV Format

Minimum CSV:

tag
cat ears
blue hair
space background

Recommended CSV:

tag,aliases,category,description,source,weight
cat ears,nekomimi|kitty ears,feature,animal ears,personal,1.25
blue hair,azure hair,hair,blue colored hair,personal,1.0

Notes:

  • Alias separators: | or ;.
  • Encoding fallback order: utf-8-sig, utf-8, latin-1.
  • Tags and aliases are normalized for matching by lowercasing, trimming, and treating underscores as spaces.

Node Reference

| Node | Status | Purpose | |------|--------|---------| | TagRAG Index | Functional lexical + optional semantic path | Load CSV data, optionally build/reuse a persisted FAISS semantic index, test lexical retrieval, emit JSON index_ref. | | TagRAG Prompt | Functional one-shot path | Retrieve tags with lexical/semantic/hybrid mode, build one-shot messages, call native llama-cpp-python by default or alternate llama-server, output prompt text and debug status. | | TagRAG Chat | Functional lexical + optional semantic path | Bounded multi-turn prompt refinement with RAG retrieval per turn, resettable in-memory history, and native/default or llama-server generation. | | TagRAG Image | Functional tag-string + optional WD14 path | Accept existing Danbooru-style tags from manual input or an external tagger node, optionally tag IMAGE tensors with a manually installed WD14-style ONNX model, retrieve related tags, and optionally refine through native/default or llama-server generation. |

All TagRAG nodes are marked as ComfyUI output nodes so they can be queued standalone during manual testing and string-output workflows. User-supplied text inputs that are needed for useful work, such as csv_path, idea, message, and index_ref, are exposed as optional widgets with safe defaults so ComfyUI can queue the nodes without wires connected. Missing or blank values return status strings instead of Python TypeError crashes; for example, a blank csv_path returns an explicit ERROR: csv_path is required status.

Chat history is stored in memory on the node instance, keyed by conversation_id. It resets when ComfyUI restarts and may reset if ComfyUI recreates the node instance; it is not persisted to disk.

TagRAG Image Mode

TagRAG Image remains tag-string first and composable with external ComfyUI tagger nodes:

  1. Wire a TagRAG Index index_ref into TagRAG Image.
  2. Provide comma-separated Danbooru-style tags through the tags STRING input.
  3. Leave refine disabled to return the original tags plus retrieved related CSV tags without any generation call.
  4. Enable refine to build image-prompt messages from prompts/image_prompt_system.txt and call the selected generation backend for a refined prompt.

Built-in tagging is opt-in through image_tagging_backend:

  • existing_tags (default): use the tags STRING input. This path does not require onnxruntime, Pillow, or model files.
  • built_in_wd14: use the optional in-node WD14-style tagger when an IMAGE is provided and dependencies/model files are present.

If both tags and image are provided, the tags STRING takes precedence so external/manual taggers stay deterministic and composable.

Optional built-in WD14 setup:

  1. Install image extras manually, outside node execution:

    pip install -r requirements-image.txt

  2. Place model files manually. Preferred layout:

    ComfyUI/models/tag_rag/wd14/<model_name>/
├── model.onnx
└── WD14 tag metadata CSV

    If ComfyUI folder_paths is unavailable, status messages describe the safe fallback path used for discovery.

  3. Select image_tagging_backend=built_in_wd14, set tagger_model or tagger_model_path, thresholds, exclude_tags, and replace_underscores as needed.

The node does not install packages, start services, or silently download model files. Missing dependencies or model files return actionable status text. Explicit model download/setup is deferred.

Out of scope for this image path: vision-language chat, “talk about this image,” pose/composition transfer, Florence2, ColPali, and llama.cpp vision/mmproj workflows.

llama.cpp Server Preset

llama-server remains a supported alternate generation and embedding backend. It avoids duplicated native model loads when a user already runs a local llama.cpp service.

Current generation target:

base_url: http://127.0.0.1:8080
chat endpoint: /v1/chat/completions
streaming: disabled
response content path: choices[0].message.content

Optional embedding target:

embedding endpoint: http://127.0.0.1:8080/v1/embeddings
retrieval mode: semantic or hybrid
index backend: FAISS

The project does not start, stop, or manage llama-server processes. It only calls a user-controlled local server when embedding_backend=llama_server or generation_backend=llama_server is selected.

Native llama-cpp-python Backend

Prompt, Chat, and Image refinement default to generation_backend=llama_cpp_python for in-process GGUF generation.

Setup notes:

  • llama-cpp-python is listed in normal install requirements because ComfyUI custom-node installs commonly install node requirements. Capability checks remain because wheel availability is environment-specific and installs can fail.
  • Provide an existing local GGUF file through model_path.
  • No GGUF models are downloaded, discovered, started, stopped, or managed by the nodes.
  • The native generation backend is text-only. Vision/mmproj, image chat/collaboration, streaming, and tool calling remain out of scope.
  • Native semantic/hybrid embeddings are supported through a separate embedding_model_path; use an embedding-capable GGUF.

Initial native options exposed in nodes:

  • model_path
  • chat_format
  • n_ctx
  • n_gpu_layers
  • max_tokens
  • temperature
  • top_p
  • top_k_sampling
  • repeat_penalty
  • seed

Model instances are cached by model_path, chat_format, n_ctx, and n_gpu_layers to avoid reloading the same GGUF on every call. A small backend cleanup helper exists in tag_rag.llama_cpp_backend, but nodes do not aggressively call llama.cpp backend-free or torch cleanup.

Embedding model instances are cached separately by embedding_model_path, n_ctx, and n_gpu_layers.

Requirements

| File | Purpose | |------|---------| | requirements.txt | Normal ComfyUI install: numpy, requests, faiss-cpu, onnxruntime, Pillow, llama-cpp-python. | | requirements-rag.txt | Split-install semantic retrieval reference: FAISS. | | requirements-image.txt | Split-install image tagging reference: onnxruntime, Pillow. |

All modules remain capability-guarded and should import safely even if FAISS/image/llama-cpp-python installs fail or are removed from a customized environment.

Project Structure

| Path | Role | |------|------| | __init__.py | ComfyUI package entrypoint and node registration. Loaded by ComfyUI when this directory is git-cloned into custom_nodes/. | | comfyui_tag_rag/__init__.py | Wheel-install entrypoint that re-exports NODE_CLASS_MAPPINGS plus WEB_DIRECTORY so pip install comfyui-tag-rag (which drops the package into site-packages/comfyui_tag_rag/) is a drop-in replacement for the git-clone install. | | nodes/ | ComfyUI node classes and node-facing error handling. | | tag_rag/ | Core library for CSV loading, tag catalog, lexical/semantic search, prompt building, embeddings, llama-server/native generation clients, and persistence helpers. | | prompts/ | Plain-text prompt templates loaded by tag_rag.prompt_builder for git-clone installs. | | tag_rag/prompt_templates/ | Packaged fallback prompt templates for wheel installs where the prompts/ directory is not adjacent to the package. | | examples/ | Explicit workflow JSON placeholders plus graph-shape documentation for one-shot, chat, and image modes. | | tests/ | Pytest suite for imports, catalog/search behavior, prompt construction, llama-server mocking, and end-to-end one-shot flow. | | scripts/verify.py | Import and compile verifier for environments where pytest may not be available. | | web/js/tag_rag.js | Intentionally minimal placeholder loaded through WEB_DIRECTORY; no frontend features are implemented yet. |

Install entrypoints

The package supports both common ComfyUI install patterns:

  1. Git-clone install (typical ComfyUI workflow): clone the repository into ComfyUI's custom_nodes/ directory. ComfyUI loads the root __init__.py, which uses ComfyUI-relative imports to expose the four TagRAG* node classes. The repo's prompts/ and web/js/ are used directly.
  2. Wheel install (pip install comfyui-tag-rag): setuptools builds a wheel that ships a comfyui_tag_rag/ distribution package plus the same nodes/ and tag_rag/ source. comfyui_tag_rag/__init__.py re-exports the same mappings and WEB_DIRECTORY so ComfyUI's loader picks them up the same way. Prompt templates ship as tag_rag/prompt_templates/*.txt (declared in [tool.setuptools.package-data]) so PromptBuilder can still resolve them through importlib.resources.

Module-scoped guidance exists in:

  • nodes/README.md
  • tag_rag/README.md
  • prompts/README.md
  • examples/README.md
  • tests/README.md

Development Verification

Run these from the repo root after code or documentation changes that reference behavior:

python -m compileall . -q
python scripts/verify.py
python -m pytest tests/ -v
git status --short

Troubleshooting

| Symptom | Likely cause | Resolution | |---------|--------------|------------| | ERROR: llama-server request failed from Prompt/Chat/Image generation | llama-server is not reachable at base_url or is not exposing /v1/chat/completions. | Confirm the server is running, use the base URL only (for example http://127.0.0.1:8080), verify the selected model is loaded server-side, or switch to native generation with model_path. | | Lexical retrieval works but semantic/hybrid returns FAISS unavailable | faiss-cpu is not installed. | Install requirements-rag.txt, then rebuild the index with retrieval_mode=semantic or hybrid. | | Semantic/hybrid returns missing semantic_index_dir | The index_ref came from a lexical-only build or lacks persisted semantic index metadata. | Re-run TagRAG Index with semantic/hybrid mode and rebuild=true, then wire the new JSON index_ref. | | Image tagging reports missing image dependencies | Optional image extras are absent. | Install requirements-image.txt; the default existing_tags path continues to work without these extras. | | Built-in WD14 reports missing model files | Manual WD14 model files are not in the selected model directory. | Place the ONNX model and its tag metadata CSV in the expected model directory, or set tagger_model_path explicitly. No files are downloaded automatically. | | Native generation reports llama-cpp-python unavailable or invalid GGUF path | Native backend is not installed or model_path does not point to a readable GGUF. | Install llama-cpp-python for your platform/CUDA stack and provide an explicit local GGUF path, or switch to generation_backend=llama_server. | | Semantic/hybrid native embeddings report missing embedding_model_path | embedding_backend=llama_cpp_python needs a separate embedding-capable GGUF path. | Set embedding_model_path, or switch embedding_backend=llama_server and configure the embedding endpoint. | | No custom UI changes appear | web/js/tag_rag.js is currently a placeholder. | This is expected; core nodes work through standard ComfyUI widgets. |

Design Boundaries

In scope for the early alpha:

  • Local-first ComfyUI prompt assistance.
  • CSV tag datasets.
  • Lexical retrieval by default.
  • Native llama-cpp-python in-process GGUF generation by default.
  • llama-server generation over HTTP as an alternate.
  • Optional FAISS semantic retrieval and optional built-in WD14-style image tagging.

Out of scope:

  • Cloud LLM providers.
  • Agent frameworks.
  • LoRA or recipe management.
  • Silent automatic model downloads.
  • Automatic llama-server process management.
  • Automatic repair of failed FAISS, image tagging, or native llama-cpp-python installations.
  • llama.cpp vision/mmproj support or image collaboration workflows.

License

MIT