4.5 KiB
CLAUDE.md
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
Project Overview
ComfyUI custom node that generates Stable Diffusion prompts using Ollama. Users provide a brief description and select a style preset; the node outputs a detailed, optimized prompt. Supports LoRA-enhanced models with automatic discovery and prioritization.
Current version: 1.4.1
Architecture
__init__.py # ComfyUI entry point, exports NODE_CLASS_MAPPINGS
nodes/
prompt_generator_node.py # PromptGeneratorNode class with all logic
config/
templates.yaml # Jinja2 style templates (editable by users)
Modelfile.limbicnation # Ollama Modelfile for LoRA-enhanced model
Single node design: All functionality is in PromptGeneratorNode. The node:
- Loads style templates from YAML (or uses hardcoded defaults)
- Renders the selected template with Jinja2 (or regex fallback)
- Calls Ollama via streaming API with per-chunk timeouts and ComfyUI ProgressBar integration
- Strips reasoning/thinking from output unless
include_reasoning=True
Graceful degradation: Optional imports (yaml, jinja2, ollama) have fallbacks. The node works with subprocess calls even if Python packages aren't installed.
Streaming & timeouts (v1.1.6): _generate_with_streaming() uses a background thread to enforce per-chunk (30s) and total timeouts, avoiding the 120s blocking timeout issue. Cold starts get a 1.3x timeout multiplier.
Development
Prerequisites:
- Ollama running locally with a model pulled (e.g.
qwen3:8b) - ComfyUI installation for integration testing
Install dependencies:
pip install -r requirements.txt
Lint with ruff (already configured in project):
ruff check .
ruff format .
Test the node manually: Restart ComfyUI, add the node from text/generation category, verify prompt generation works.
Key Patterns
ComfyUI node structure:
INPUT_TYPES(): Class method returning dict withrequiredandoptionalinputsRETURN_TYPES,RETURN_NAMES: Output type definitionsFUNCTION: Name of the method to call (generate)CATEGORY: Where node appears in ComfyUI menu (text/generation)
Template system: Templates in config/templates.yaml use Jinja2 syntax. Variables: {{ description }}, {{ emphasis }}, {{ mood }}. Conditionals: {% if emphasis %}...{% endif %}.
Output cleaning: extract_final_prompt() strips Qwen3's "Thinking..." blocks and markdown formatting from responses.
Model discovery: The node auto-discovers Ollama models and prioritizes those containing lora or limbicnation keywords. All available models appear in the dropdown.
Publishing
Registry: registry.comfy.org - Node ID: comfyui-prompt-generator
Publish a new version:
# 1. Update version in pyproject.toml
# 2. Commit changes
git add -A && git commit -m "chore: bump version to X.Y.Z"
# 3. Tag and push (triggers GitHub Actions workflow)
git tag vX.Y.Z
git push origin main && git push origin vX.Y.Z
# Or publish manually:
comfy node publish --confirm
CI/CD: .github/workflows/publish.yml auto-publishes on version tags (v*.*.*).
Adding New Styles
- Add entry to
config/templates.yamlfollowing existing format - Add style key to
INPUT_TYPES()style combo list inprompt_generator_node.py - Optionally add to
DEFAULT_STYLESdict for fallback when YAML unavailable
LoRA-Enhanced Models
- Training: QLoRA on Qwen3-4B-Instruct-2507 using the Limbicnation/Images-Diffusion-Prompt-Style dataset (750 prompts)
- Quantization: Merged LoRA and converted to Q8_0 GGUF for Ollama
- Integration:
PromptGeneratorNodeauto-discovers and prioritizes models containingloraorlimbicnationkeywords
Creating a LoRA-Enhanced Model
- Fine-tune a LoRA on the dataset above
- Export as
.safetensors(non-quantized recommended) - Create the Ollama model:
# Edit config/Modelfile.limbicnation with your adapter path ollama create qwen3-limbicnation -f config/Modelfile.limbicnation - Restart ComfyUI — the new model will appear in the dropdown
Modelfile Template
See config/Modelfile.limbicnation for a pre-configured template with:
- Limbicnation system prompt
- Optimal temperature/top_p settings
- ADAPTER placeholder for your LoRA