From 081e7a54fd232b293ccbc665c226721d19b66be9 Mon Sep 17 00:00:00 2001 From: Daniel Florian Date: Fri, 21 Nov 2025 15:40:17 -0700 Subject: [PATCH] Initial commit: ComfyUI-WorkflowGenerator --- .github/workflows/ci.yml | 37 ++ .github/workflows/node-diff.yml | 28 + .pre-commit-config.yaml | 10 + README.md | 251 ++++++++ __init__.py | 45 ++ catalog/.gitkeep | 4 + cataloging/catalog.py | 366 +++++++++++ config/config.py | 336 ++++++++++ generators/node_validator.py | 751 ++++++++++++++++++++++ generators/workflow_builder.py | 529 +++++++++++++++ generators/workflow_generator.py | 245 +++++++ nodes/node_validator_node.py | 344 ++++++++++ nodes/pipeline_node.py | 461 +++++++++++++ nodes/update_catalog_node.py | 135 ++++ nodes/workflow_builder_node.py | 178 +++++ nodes/workflow_generator_node.py | 209 ++++++ pyproject.toml | 114 ++++ requirements.txt | 5 + tests/test_nodes.py | 52 ++ utils/file_utils.py | 88 +++ utils/model_loaders/device_utils.py | 256 ++++++++ utils/model_loaders/gguf_loader.py | 249 +++++++ utils/model_loaders/huggingface_loader.py | 220 +++++++ utils/model_manager.py | 443 +++++++++++++ utils/node_utils.py | 141 ++++ utils/process_utils.py | 163 +++++ utils/workflow_layout.py | 149 +++++ wiki/Configuration.md | 67 ++ wiki/Home.md | 30 + wiki/Installation.md | 141 ++++ wiki/Instruction-Prompt-Usage.md | 76 +++ wiki/Node-Reference.md | 131 ++++ wiki/README.md | 78 +++ wiki/Troubleshooting.md | 110 ++++ workflows/WG_3-Chained-Nodes.json | 319 +++++++++ workflows/WG_Pipeline.json | 143 ++++ workflows/WG_UpdateNodeCatalog.json | 113 ++++ 37 files changed, 7017 insertions(+) create mode 100644 .github/workflows/ci.yml create mode 100644 .github/workflows/node-diff.yml create mode 100644 .pre-commit-config.yaml create mode 100644 README.md create mode 100644 __init__.py create mode 100644 catalog/.gitkeep create mode 100644 cataloging/catalog.py create mode 100644 config/config.py create mode 100644 generators/node_validator.py create mode 100644 generators/workflow_builder.py create mode 100644 generators/workflow_generator.py create mode 100644 nodes/node_validator_node.py create mode 100644 nodes/pipeline_node.py create mode 100644 nodes/update_catalog_node.py create mode 100644 nodes/workflow_builder_node.py create mode 100644 nodes/workflow_generator_node.py create mode 100644 pyproject.toml create mode 100644 requirements.txt create mode 100644 tests/test_nodes.py create mode 100644 utils/file_utils.py create mode 100644 utils/model_loaders/device_utils.py create mode 100644 utils/model_loaders/gguf_loader.py create mode 100644 utils/model_loaders/huggingface_loader.py create mode 100644 utils/model_manager.py create mode 100644 utils/node_utils.py create mode 100644 utils/process_utils.py create mode 100644 utils/workflow_layout.py create mode 100644 wiki/Configuration.md create mode 100644 wiki/Home.md create mode 100644 wiki/Installation.md create mode 100644 wiki/Instruction-Prompt-Usage.md create mode 100644 wiki/Node-Reference.md create mode 100644 wiki/README.md create mode 100644 wiki/Troubleshooting.md create mode 100644 workflows/WG_3-Chained-Nodes.json create mode 100644 workflows/WG_Pipeline.json create mode 100644 workflows/WG_UpdateNodeCatalog.json diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..a256ea0 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,37 @@ +name: CI + +on: + push: + branches: [ main, develop ] + pull_request: + branches: [ main, develop ] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10", "3.11", "3.12"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[dev]" + + - name: Run Ruff + run: ruff check . + + - name: Run Ruff format check + run: ruff format --check . + + - name: Run tests + run: pytest tests/ -v + diff --git a/.github/workflows/node-diff.yml b/.github/workflows/node-diff.yml new file mode 100644 index 0000000..b6389e0 --- /dev/null +++ b/.github/workflows/node-diff.yml @@ -0,0 +1,28 @@ +name: Node Diff + +on: + pull_request: + branches: [ main ] + +jobs: + node-diff: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Install node-diff + run: | + python -m pip install --upgrade pip + pip install node-diff + + - name: Run node-diff + run: | + node-diff --base HEAD~1 --head HEAD + diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..365adad --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,10 @@ +# Pre-commit hooks configuration +# See https://pre-commit.com for more information +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.6.9 + hooks: + - id: ruff + args: [--fix, --exit-non-zero-on-fix] + - id: ruff-format + diff --git a/README.md b/README.md new file mode 100644 index 0000000..f479532 --- /dev/null +++ b/README.md @@ -0,0 +1,251 @@ +# ComfyUI Workflow Generator + +Generate ComfyUI workflows from natural language descriptions using Large Language Models (LLMs). + +This custom node enables users to describe a workflow in natural language (e.g., *"Create a text-to-image workflow using SDXL"*) and automatically constructs the corresponding node graph. It leverages Large Language Models (LLMs) to bridge the gap between intent and execution. + +## Implementation Note + +This project is an independent implementation of the [ComfyGPT research](https://arxiv.org/abs/2503.17671), designed to bring that architecture directly into ComfyUI as a native node suite. + +My goal was to preserve the core functionality of the original research—specifically the multi-stage pipeline of generation, validation, and construction, while optimizing it for the local ComfyUI environment. This implementation focuses on modularity, allowing users to inspect and intervene at each stage of the generation process. + +**Current Limitations:** +The system's knowledge is bounded by its training data. While it excels at standard patterns and known nodes, it cannot inherently "know" about custom nodes released after its training cutoff without additional context. It acts as a powerful accelerator for workflow creation, but supervision is recommended. + +## Based on ComfyGPT Research + +**Title:** "ComfyGPT: A Self-Optimizing Multi-Agent System for Comprehensive ComfyUI Workflow Generation" +- **Original Repository:** https://github.com/comfygpt/comfygpt +- **Project Website:** https://comfygpt.github.io/ + +### How It Works + +This implementation uses a specialized LLM fine-tuned on workflow data to execute a three-stage pipeline: +1. **Generator:** Interprets the natural language prompt to generate a logical graph structure (JSON). +2. **Validator:** Verifies node names against the local installation or semantic embeddings to ensure compatibility. +3. **Builder:** Compiles the validated structure into an executable ComfyUI workflow format. + +--- + +### Model Sources + +The models used in this implementation are based on the original ComfyGPT research: + +- **WorkflowGenerator Model**: Original fine-tuned model from [xiatianzs/resources](https://huggingface.co/xiatianzs/resources/tree/main) (ComfyGPT research team) + - Base model: Qwen2.5-14B, fine-tuned for workflow generation + - Training context window: 8,192 tokens (cutoff length used during training - see [training config](https://github.com/comfygpt/comfygpt/blob/main/train/sft/qwen2.5_full_sft.yaml)) + - Model architecture context window: 131,072 tokens (128K) - Maximum supported by model architecture (see [model config](https://huggingface.co/xiatianzs/resources/blob/main/model/flow_agent/config.json)) + - Quantized to GGUF format (q8_0) for efficient inference + +- **Embedding Model**: [sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2](https://huggingface.co/sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2) (original SentenceTransformer model) + - Used for semantic search in NodeValidator + +- **NodeValidator Model**: Base [Qwen2.5-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-7B-Instruct) model (not fine-tuned) + - Context window: 32,768 tokens + - Used for LLM refinement (optional) + - Quantized to GGUF format (q8_0) for efficient inference + +**Model Repositories:** +- **Original Models**: [xiatianzs/resources](https://huggingface.co/xiatianzs/resources/tree/main) - Original fine-tuned models from ComfyGPT research team (HuggingFace format) +- **Pre-quantized GGUF Models**: [DanielPFlorian/comfyui-workflowgenerator-models](https://huggingface.co/DanielPFlorian/comfyui-workflowgenerator-models) - Ready-to-use quantized GGUF models for this implementation + + +## Installation + +### Prerequisites + +- ComfyUI installed and running +- Python 3.10 or higher (see note below) +- CUDA-capable GPU (recommended) or CPU +- Git installed (for cloning the repository) + +**Python Requirements:** +- **Standard ComfyUI Installation:** Python 3.10 or higher must be installed separately. ComfyUI requires Python to run, so if ComfyUI is working, Python is already installed. +- **Portable ComfyUI Installation:** Python is included (embedded in the `python_embeded` folder), so no separate Python installation is needed. The custom node will use ComfyUI's embedded Python. + +**Note:** If Git is not installed, download it from [git-scm.com](https://git-scm.com/). You can verify Git installation by running `git --version` in your terminal. + +### Model Recommendation + +**Model Format Comparison:** + +- **GGUF models**: Use significantly less VRAM with similar generation quality compared to HuggingFace models. Quantization (q8_0, q4_0) provides a good balance between quality and memory usage. + +- **HuggingFace models**: Use more VRAM but offer full precision. Both formats are fully supported and produce similar quality results. + +### Step 1: Clone the Repository + +Navigate to your ComfyUI `custom_nodes` directory and clone this repository: + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/danielpflorian/ComfyUI-WorkflowGenerator.git +``` + +### Step 2: Install Dependencies + +Install the required Python dependencies by opening your terminal inside the ComfyUI-WorkflowGenerator folder: + +```bash +cd ComfyUI/custom_nodes/ComfyUI-WorkflowGenerator +pip install -r requirements.txt +``` + +**For Portable ComfyUI Installations:** + +For portable installations, use the embedded Python from the portable ComfyUI root directory: + +```bash +cd +python_embeded\python.exe -s -m pip install -r ComfyUI\custom_nodes\ComfyUI-WorkflowGenerator\requirements.txt +``` + +### Step 3: Install llama-cpp-python (Required for GGUF Models) + +**Important:** `llama-cpp-python` is **required** for GGUF models. It must be installed separately based on your system configuration. + +**Note:** The quick install commands below may not work for all Windows/CUDA configurations. If they fail or if CUDA support isn't detected, **local compilation is often necessary**. See the [Wiki - llama-cpp-python Installation](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/wiki/Installation#llama-cpp-python-installation) for detailed instructions on compiling from source. + +**Quick install (try this first):** +- **CPU only**: `pip install llama-cpp-python` +- **CUDA (NVIDIA)**: `pip install llama-cpp-python[cuda]` +- **Metal (macOS)**: `pip install llama-cpp-python[metal]` + +**For Portable ComfyUI Installations:** + +For portable installations, use the embedded Python from the portable ComfyUI root directory: + +```bash +cd +python_embeded\python.exe -s -m pip install llama-cpp-python[cuda] +``` + +**Note:** If you plan to use HuggingFace models instead, you can skip this step. + +### Step 4: Copy Models + +Copy your GGUF models and tokenizers to `ComfyUI/models/LLM/`: + +``` +ComfyUI/models/LLM/ +├── workflow-generator-q8_0.gguf # WorkflowGenerator model +├── workflow-generator/ # WorkflowGenerator tokenizer +├── Qwen2.5-7B-Instruct-q8_0.gguf # NodeValidator model (optional) +├── Qwen2.5-7B-Instruct/ # NodeValidator tokenizer (optional) +└── paraphrase-multilingual-MiniLM-L12-v2/ # Embedding model +``` + +### Step 5: Restart ComfyUI + +Restart ComfyUI to load the custom nodes. The nodes will appear in the `WorkflowGenerator` category. + +## Quick Start + +The easiest way to get started is using the **Workflow Generator Pipeline** node, which processes your instruction through all three stages sequentially: + +1. **Install the custom node** (see Installation above) +2. **Install llama-cpp-python** (required for GGUF models - see [Wiki - llama-cpp-python Installation](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/wiki/Installation#llama-cpp-python-installation)) +3. **Run the "Update Node Catalog" node first** - This scans and catalogs all available ComfyUI nodes (native and custom) and is required before generating workflows +4. **Add the "Workflow Generator Pipeline" node** to your ComfyUI workflow +5. **Enter your instruction**, for example: `"Create a text-to-image workflow"` +6. **Configure model path** (see Model Format Comparison above for guidance) +7. **Execute the node** to generate a complete workflow +8. **Use the generated workflow** - it will appear as a workflow JSON that you can load or save + +**Expected Results:** +- The Pipeline node automatically generates the workflow diagram, validates node names, and converts it to executable ComfyUI workflow JSON +- You can save the workflow to a file for later use. By default it will be saved in comfyUI/output. + +**Note:** For detailed documentation on individual nodes and advanced usage, see the [Wiki](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/wiki). + +## Available Nodes + +This section provides a high-level overview of each node. + +#### [Workflow Generator Pipeline](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/wiki/Node-Reference#workflow-generator-pipeline) + +- **Purpose**: One-click solution for complete workflow generation +- **What it does**: Processes your instruction through all three stages sequentially: generates the workflow diagram, validates and corrects node names, then builds the final ComfyUI workflow JSON +- **Best for**: Quick workflow creation without intermediate steps +- **Key inputs**: Instruction, model path, configuration options +- **Key outputs**: Complete workflow JSON, optional file save + +For more control, you can use individual nodes to inspect and modify intermediate results: + +**Important:** Before using individual nodes or the pipeline, **run the "Update Node Catalog" node first** to scan and catalog all available ComfyUI nodes. This is required for proper node validation and workflow building. + +**Note:** WorkflowGenerator, NodeValidator, and WorkflowBuilder are designed to be used in sequence. NodeValidator is optional—you can connect WorkflowGenerator directly to WorkflowBuilder if you want to skip validation. + +#### 1. [WorkflowGenerator](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/wiki/Node-Reference#workflowgenerator-node) + +- **Purpose**: Generates workflow diagrams from natural language +- **Key inputs**: Instruction, model selection +- **Key outputs**: Workflow diagram (JSON string) +- **Usage**: First step in the pipeline + +#### 2. [NodeValidator](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/wiki/Node-Reference#nodevalidator-node) + +- **Purpose**: Validates and corrects node names in workflow diagrams +- **Key inputs**: Workflow diagram, optional instruction for context +- **Key outputs**: Refined diagram with corrected node names +- **Usage**: Optional second step (can be skipped) +- **Two modes**: + - Semantic search only (faster, deterministic) + - LLM refinement (slower, more accurate) + +#### 3. [WorkflowBuilder](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/wiki/Node-Reference#workflowbuilder-node) + +- **Purpose**: Converts workflow diagrams into executable ComfyUI workflow JSON +- **Key inputs**: Refined diagram (from NodeValidator) or raw diagram (from WorkflowGenerator) +- **Key outputs**: Workflow JSON, optional file save +- **Usage**: Final step in the pipeline + +#### [UpdateNodeCatalog](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/wiki/Node-Reference#updatenodecatalog-node) + +- **Purpose**: Scans and catalogs all available ComfyUI nodes (native and custom) +- **When to run**: After installing new custom nodes or updating ComfyUI +- **Key inputs**: Catalog directory path (optional) +- **Key outputs**: Updated catalog files + +## Troubleshooting + +**Quick fixes:** +- **Models not found**: Verify models are in `ComfyUI/models/LLM/` +- **OOM errors**: `auto_gpu_layers` is enabled by default to prevent this. If issues persist, verify your VRAM is sufficient for the model size. +- **Slow performance**: Use GGUF models (lower VRAM usage), disable LLM refinement +- **Invalid nodes**: Run **Update Node Catalog** node + +## Architectural Insights & Future Vision + +While this implementation successfully brings the ComfyGPT architecture to ComfyUI, the rapid evolution of the generative AI landscape suggests that the next generation of workflow generation tools will need to evolve beyond static fine-tuning. + +### The Problem with Static Models +The ComfyUI custom node ecosystem changes daily. New nodes and unforeseen architectures are released constantly. A model fine-tuned on a dataset from last month (like the current WorkflowGenerator) is inherently "frozen" in time. It cannot hallucinate the correct connections for a node it has never seen during training. + +### Ideas on Future Workflow Generators +To achieve true state-of-the-art performance that keeps pace with the community, future architectures should likely move toward: + +1. **Retrieval-Augmented Generation (RAG) for Nodes**: Instead of baking node knowledge into the model weights, an agent could query a dynamic, vector-embedded database that includes both the user's *currently installed* nodes and an updated database of repositories. This allows the system to discover and use completely new nodes or updated versions directly from the internet, enabling it to "read" the documentation of a node released today and use it immediately. + +2. **Input/Output Type Awareness**: The current approach uses semantic search to correct node names (e.g., matching "Load Image" to "Image Loader"). However, a robust system needs to understand **I/O Schema**. + * *Current*: "Does this node name look right?" + * *Future*: "Does the `LATENT` output of Node A actually fit into the `LATENT` input of Node B?" + * An agent needs to reason about data types (Float, Image, Conditioning, Model) to ensure generated workflows are not just linguistically plausible, but executably valid. + +3. **Small Graph-Reasoning Models (SLMs)**: We may not need massive 70B+ parameter models for this task. Specialized, smaller models trained specifically on Graph Theory and Directed Acyclic Graphs (DAGs) could offer superior reasoning capabilities for connecting logical blocks, while relying on the RAG system for the specific node vocabulary. + +4. **Case-Based Reasoning & Workflow Retrieval**: Thousands of high-quality workflows exist online, often accompanied by images and descriptions. A future system should not just generate from scratch but actively **retrieve and adapt existing workflows**. By indexing these community-created workflows as "knowledge patterns," the agent can identify how experts typically solve specific sub-problems (e.g. "How is ControlNet usually connected to Wan Video Generation Nodes?") and apply those proven sub-graphs to the current task. + +This project serves as a foundational step, implementing the current research (ComfyGPT). However, the future lies in agents that can actively explore, reason about, and validate the live ComfyUI environment they inhabit. + +## License + +GNU General Public License v3 + +## Acknowledgments + +- **ComfyUI**: Built on top of [ComfyUI](https://github.com/comfyanonymous/ComfyUI) +- **ComfyGPT Research Team**: This project is based on the [ComfyGPT research paper](https://arxiv.org/abs/2503.17671) and [original repository](https://github.com/comfygpt/comfygpt) by Oucheng Huang, Yuhang Ma, Zeng Zhao, Mingrui Wu, Jiayi Ji, Rongsheng Zhang, Zhipeng Hu, Xiaoshuai Sun, and Rongrong Ji +- **llama-cpp-python**: Uses [llama-cpp-python](https://github.com/abetlen/llama-cpp-python) for GGUF model support + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..44f2d60 --- /dev/null +++ b/__init__.py @@ -0,0 +1,45 @@ +""" +ComfyUI-WorkflowGenerator Custom Nodes for ComfyUI +Generates ComfyUI workflows from natural language descriptions using LLMs. +""" + +from comfy_api.latest import ComfyExtension, io +from typing_extensions import override + +from .nodes.node_validator_node import NodeValidatorNode +from .nodes.pipeline_node import WGPipelineNode +from .nodes.update_catalog_node import UpdateNodeCatalogNode +from .nodes.workflow_builder_node import WorkflowBuilderNode + +# Import node classes from individual files +from .nodes.workflow_generator_node import WorkflowGeneratorNode + +# Import model manager for folder registration +from .utils.model_manager import register_llm_folder_type + + +class WGExtension(ComfyExtension): + """ComfyUI WorkflowGenerator extension for ComfyUI V3 API.""" + + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + """ + Return list of ComfyUI-WorkflowGenerator node classes. + + Returns: + List of node classes to register + """ + return [ + UpdateNodeCatalogNode, + WorkflowGeneratorNode, + NodeValidatorNode, + WorkflowBuilderNode, + WGPipelineNode, + ] + + +def comfy_entrypoint(): + """Entry point for ComfyUI to load the extension.""" + # Register LLM folder type when extension loads + register_llm_folder_type() + return WGExtension() diff --git a/catalog/.gitkeep b/catalog/.gitkeep new file mode 100644 index 0000000..200ca6a --- /dev/null +++ b/catalog/.gitkeep @@ -0,0 +1,4 @@ +# This file ensures the catalog directory is tracked by Git +# The catalog files (node_list.json, node_info.json) are generated by UpdateNodeCatalog +# and are excluded from Git via .gitignore due to their large size (100MB+) + diff --git a/cataloging/catalog.py b/cataloging/catalog.py new file mode 100644 index 0000000..3cc8f99 --- /dev/null +++ b/cataloging/catalog.py @@ -0,0 +1,366 @@ +""" +Node cataloging system for scanning and cataloging ComfyUI nodes. +Supports both V1 and V3 API nodes. +""" + +import json +import logging +from pathlib import Path +from typing import Any + +from tinydb import TinyDB + +# Try to import ComfyUI nodes +try: + import nodes + + COMFYUI_NODES_AVAILABLE = True +except ImportError: + COMFYUI_NODES_AVAILABLE = False + logging.warning("ComfyUI nodes module not available") + +# Try to import V3 API +try: + from comfy_api.latest import io # noqa: F401 + + V3_API_AVAILABLE = True +except ImportError: + V3_API_AVAILABLE = False + logging.warning("ComfyUI V3 API not available") + + +def extract_v1_node_info(node_name: str, node_class: Any) -> dict | None: + """ + Extract information from a V1 API node. + + Args: + node_name: Node class name + node_class: Node class object + + Returns: + Dictionary with node information or None if extraction fails + """ + try: + input_types = {} + if hasattr(node_class, "INPUT_TYPES"): + input_types = node_class.INPUT_TYPES() + + return_types = [] + if hasattr(node_class, "RETURN_TYPES"): + return_types = node_class.RETURN_TYPES + + return_names = None + if hasattr(node_class, "RETURN_NAMES"): + return_names = node_class.RETURN_NAMES + + output_tooltips = None + if hasattr(node_class, "OUTPUT_TOOLTIPS"): + output_tooltips = node_class.OUTPUT_TOOLTIPS + + category = None + if hasattr(node_class, "CATEGORY"): + category = node_class.CATEGORY + + description = None + if hasattr(node_class, "DESCRIPTION"): + description = node_class.DESCRIPTION + + display_name = node_name + if COMFYUI_NODES_AVAILABLE: + try: + if hasattr(nodes, "NODE_DISPLAY_NAME_MAPPINGS") and node_name in nodes.NODE_DISPLAY_NAME_MAPPINGS: + display_name = nodes.NODE_DISPLAY_NAME_MAPPINGS[node_name] + except Exception: + pass + + is_output_node = False + if hasattr(node_class, "OUTPUT_NODE") and node_class.OUTPUT_NODE: + is_output_node = True + + is_deprecated = False + if getattr(node_class, "DEPRECATED", False): + is_deprecated = True + + is_experimental = False + if getattr(node_class, "EXPERIMENTAL", False): + is_experimental = True + + is_api_node = None + if hasattr(node_class, "API_NODE"): + is_api_node = node_class.API_NODE + + node_info = { + "node_type": node_name, + "display_name": display_name, + "inputs": input_types, + "outputs": return_types, + "output_names": return_names, + "output_tooltips": output_tooltips, + "category": category, + "description": description, + "is_output_node": is_output_node, + "is_deprecated": is_deprecated, + "is_experimental": is_experimental, + "is_api_node": is_api_node, + } + + return node_info + + except Exception as e: + logging.warning(f"Failed to extract V1 node info for '{node_name}': {e}") + return None + + +def extract_v3_node_info(node_name: str, node_class: Any) -> dict | None: + """ + Extract information from a V3 API node. + + Args: + node_name: Node class name + node_class: Node class object (io.ComfyNode) + + Returns: + Dictionary with node information or None if extraction fails + """ + try: + schema = node_class.GET_SCHEMA() + + input_types = {} + input_metadata = {} + if hasattr(schema, "inputs") and schema.inputs is not None: + inputs = schema.inputs + required = {} + optional = {} + hidden = {} + + for input_def in inputs: + input_name = getattr(input_def, "id", getattr(input_def, "name", None)) + if input_name is None: + continue + input_type = input_def.io_type + is_optional = getattr(input_def, "optional", False) + is_required = not is_optional + is_hidden = getattr(input_def, "hidden", False) + + input_data = [input_type] + if hasattr(input_def, "default") and input_def.default is not None: + input_data.append({"default": input_def.default}) + + input_metadata[input_name] = { + "display_name": getattr(input_def, "display_name", None), + "tooltip": getattr(input_def, "tooltip", None), + } + + if is_hidden: + hidden[input_name] = input_data + elif is_required: + required[input_name] = input_data + else: + optional[input_name] = input_data + + if required: + input_types["required"] = required + if optional: + input_types["optional"] = optional + if hidden: + input_types["hidden"] = hidden + + return_types = [] + return_names = [] + output_display_names = [] + output_tooltips = [] + if hasattr(schema, "outputs") and schema.outputs is not None: + for output_def in schema.outputs: + return_types.append(output_def.io_type) + if hasattr(output_def, "name"): + return_names.append(output_def.name) + else: + return_names.append(output_def.io_type) + + output_display_names.append(getattr(output_def, "display_name", None)) + output_tooltips.append(getattr(output_def, "tooltip", None)) + + category = None + if hasattr(schema, "category"): + category = schema.category + + description = None + if hasattr(schema, "description"): + description = schema.description + + display_name = getattr(schema, "display_name", None) + is_output_node = getattr(schema, "is_output_node", False) + is_deprecated = getattr(schema, "is_deprecated", False) + is_experimental = getattr(schema, "is_experimental", False) + is_api_node = getattr(schema, "is_api_node", False) + + node_info = { + "node_type": node_name, + "display_name": display_name, + "inputs": input_types, + "input_metadata": input_metadata if input_metadata else None, + "outputs": return_types, + "output_names": return_names if return_names else None, + "output_display_names": output_display_names if output_display_names else None, + "output_tooltips": output_tooltips if output_tooltips else None, + "category": category, + "description": description, + "is_output_node": is_output_node, + "is_deprecated": is_deprecated, + "is_experimental": is_experimental, + "is_api_node": is_api_node, + } + + return node_info + + except Exception as e: + logging.warning(f"Failed to extract V3 node info for '{node_name}': {e}") + return None + + +def scan_nodes(progress_callback=None) -> tuple[list[str], list[dict]]: + """ + Scan all available ComfyUI nodes (V1 and V3). + + Args: + progress_callback: Optional callback function(current, total) to report progress + + Returns: + Tuple of (node_list, node_info_list) + - node_list: List of node type names + - node_info_list: List of node information dictionaries + """ + if not COMFYUI_NODES_AVAILABLE: + logging.error("ComfyUI nodes module not available") + return [], [] + + node_list = [] + node_info_list = [] + + logging.info("Scanning ComfyUI nodes...") + + try: + node_mappings = nodes.NODE_CLASS_MAPPINGS + total_nodes = len(node_mappings) + logging.info(f"Found {total_nodes} nodes in NODE_CLASS_MAPPINGS") + + processed = 0 + v1_count = 0 + v3_count = 0 + error_count = 0 + + # Update progress every N nodes to avoid too many updates + progress_update_interval = max(1, total_nodes // 100) + + for idx, (node_name, node_class) in enumerate(node_mappings.items(), 1): + current_position = idx + try: + is_v3 = False + if V3_API_AVAILABLE: + try: + if hasattr(node_class, "GET_SCHEMA"): + schema = node_class.GET_SCHEMA() + if hasattr(schema, "node_id"): + is_v3 = True + except Exception: + pass + + if is_v3: + node_info = extract_v3_node_info(node_name, node_class) + if node_info: + v3_count += 1 + else: + node_info = extract_v1_node_info(node_name, node_class) + if node_info: + v1_count += 1 + else: + node_info = extract_v1_node_info(node_name, node_class) + if node_info: + v1_count += 1 + + if node_info: + node_list.append(node_name) + node_info_list.append(node_info) + processed += 1 + else: + error_count += 1 + logging.warning(f"Failed to extract info for node '{node_name}'") + + # Update progress periodically (every N nodes or on last node) + if progress_callback and (current_position % progress_update_interval == 0 or current_position == total_nodes): + progress_callback(current_position, total_nodes) + + except Exception as e: + error_count += 1 + logging.warning(f"Error processing node '{node_name}': {e}") + if progress_callback and (current_position % progress_update_interval == 0 or current_position == total_nodes): + progress_callback(current_position, total_nodes) + continue + + logging.info(f"Node scanning complete: {processed} nodes processed ({v1_count} V1, {v3_count} V3, {error_count} errors)") + + if progress_callback: + progress_callback(total_nodes, total_nodes) + + except Exception as e: + logging.error(f"Failed to scan nodes: {e}") + return [], [] + + return node_list, node_info_list + + +def update_node_catalog(node_list_path: str, node_info_db_path: str, force_update: bool = False, progress_callback=None) -> dict[str, Any]: + """ + Update the node catalog by scanning ComfyUI nodes. + + Args: + node_list_path: Path to save node list JSON + node_info_db_path: Path to save node info database JSON + force_update: If True, overwrite existing files + + Returns: + Dictionary with update statistics + """ + logging.info("Updating node catalog...") + + node_list_file = Path(node_list_path) + node_info_file = Path(node_info_db_path) + + if not force_update and node_list_file.exists() and node_info_file.exists(): + logging.info("Node catalog files already exist, skipping update") + return {"status": "skipped", "message": "Catalog files already exist"} + + node_list, node_info_list = scan_nodes(progress_callback=progress_callback) + + if not node_list: + logging.error("No nodes found during scan") + return {"status": "error", "message": "No nodes found during scan"} + + try: + node_list_file.parent.mkdir(parents=True, exist_ok=True) + with open(node_list_file, "w") as f: + json.dump(node_list, f, indent=2) + logging.info(f"Saved node list to {node_list_path} ({len(node_list)} nodes)") + except Exception as e: + logging.error(f"Failed to save node list: {e}") + return {"status": "error", "message": f"Failed to save node list: {e}"} + + try: + node_info_file.parent.mkdir(parents=True, exist_ok=True) + + db = TinyDB(node_info_file) + db.truncate() + db.insert_multiple(node_info_list) + db.close() + logging.info(f"Saved node info database to {node_info_db_path} ({len(node_info_list)} entries)") + except Exception as e: + logging.error(f"Failed to save node info database: {e}") + return {"status": "error", "message": f"Failed to save node info database: {e}"} + + return { + "status": "success", + "node_count": len(node_list), + "node_info_count": len(node_info_list), + "node_list_path": str(node_list_path), + "node_info_db_path": str(node_info_db_path), + } diff --git a/config/config.py b/config/config.py new file mode 100644 index 0000000..6184012 --- /dev/null +++ b/config/config.py @@ -0,0 +1,336 @@ +""" +Configuration management for ComfyUI-WorkflowGenerator nodes. +Builds configuration objects from node inputs and default values. +""" + +import os +from pathlib import Path + +from omegaconf import DictConfig, OmegaConf + +from ..utils.model_manager import resolve_model_path + +# Models are in ComfyUI/models/LLM/, catalog files are in custom node catalog/ +DEFAULT_CONFIG = { + "workflow_generator": { + "model_format": "gguf", + "model_path": "models/LLM/workflow-generator-q8_0.gguf", + "tokenizer_path": "models/LLM/workflow-generator", + "max_new_tokens": 8192, + "top_p": 0.7, + "temperature": 0.95, + "allow_primitive_nodes": False, + }, + "node_validator": { + "local": False, + "model_format": "gguf", + "model_path": "models/LLM/Qwen2.5-7B-Instruct-q8_0.gguf", + "tokenizer_path": "models/LLM/Qwen2.5-7B-Instruct", + "k": 5, + "max_new_tokens": 4096, + }, + "node_base": { + "catalog_directory": "catalog", + "embedding_model_path": "models/LLM/paraphrase-multilingual-MiniLM-L12-v2", + }, +} + + +def get_custom_node_path() -> Path: + """ + Get the path to the ComfyUI-WorkflowGenerator custom node directory. + + Returns: + Path to custom node directory + """ + current_file = Path(__file__).resolve() + + # config/config.py -> ComfyUI-WorkflowGenerator/ + if current_file.parent.name == "config": + return current_file.parent.parent + + return Path.cwd() + + +def resolve_config_path(path: str, base_path: Path | None = None) -> str: + """ + Resolve a configuration path relative to the custom node directory or ComfyUI root. + + Args: + path: Path string (can be relative or absolute) + - If starts with "custom_nodes/", resolves relative to ComfyUI root + - Otherwise, resolves relative to custom node directory + base_path: Base path for resolution (default: custom node directory) + + Returns: + Resolved absolute path + """ + if base_path is None: + base_path = get_custom_node_path() + + if os.path.isabs(path) and os.path.exists(path): + return path + + if path.startswith("custom_nodes/"): + from ..utils.model_manager import get_comfyui_base_path + + comfyui_base = get_comfyui_base_path() + if comfyui_base: + full_path = comfyui_base / path + if full_path.exists(): + return str(full_path) + return str(full_path.resolve()) + + full_path = base_path / path + if full_path.exists(): + return str(full_path) + + catalog_path = base_path / "catalog" / path + if catalog_path.exists(): + return str(catalog_path) + + return str(full_path.resolve()) + + +class ConfigBuilder: + """ + Builder class for creating configuration objects from node inputs. + """ + + def __init__(self, base_config: dict | None = None): + """ + Initialize ConfigBuilder. + + Args: + base_config: Base configuration dictionary (default: DEFAULT_CONFIG) + """ + if base_config is None: + base_config = DEFAULT_CONFIG.copy() + + self.base_config = base_config + self.custom_node_path = get_custom_node_path() + + def build_workflow_generator_config( + self, + model_format: str = "gguf", + model_path: str = "", + tokenizer_path: str = "", + max_new_tokens: int = 8192, + top_p: float = 0.7, + temperature: float = 0.95, + allow_primitive_nodes: bool = False, + **kwargs, + ) -> DictConfig: + """ + Build configuration for WorkflowGenerator. + + Args: + model_format: Model format ("gguf" or "huggingface") + model_path: Path to model file/directory + tokenizer_path: Path to tokenizer directory + max_new_tokens: Maximum tokens to generate + top_p: Top-p sampling parameter + temperature: Temperature for sampling + allow_primitive_nodes: Whether to include primitive nodes + **kwargs: Additional configuration overrides + + Returns: + OmegaConf DictConfig object + """ + config = self.base_config["workflow_generator"].copy() + + if model_format: + config["model_format"] = model_format + if model_path: + config["model_path"] = model_path + # Only set tokenizer_path if explicitly provided (not None) + # If None, remove it from config so auto-detection happens + if tokenizer_path is not None: + if tokenizer_path: # Non-empty string + config["tokenizer_path"] = tokenizer_path + else: + # Empty string means remove from config (use auto-detection) + config.pop("tokenizer_path", None) + else: + # None means remove from config (use auto-detection) + config.pop("tokenizer_path", None) + if max_new_tokens is not None: + config["max_new_tokens"] = max_new_tokens + if top_p is not None: + config["top_p"] = top_p + if temperature is not None: + config["temperature"] = temperature + if allow_primitive_nodes is not None: + config["allow_primitive_nodes"] = allow_primitive_nodes + + config.update(kwargs) + + # Model/tokenizer paths are in ComfyUI models/, not custom node dir + if "model_path" in config: + config["model_path"] = resolve_model_path(config["model_path"], model_type="llm") + if "tokenizer_path" in config: + tokenizer_path_val = config["tokenizer_path"] + if not os.path.isabs(tokenizer_path_val): + resolved = resolve_model_path(tokenizer_path_val, model_type="llm") + if os.path.exists(resolved) and os.path.isdir(resolved): + config["tokenizer_path"] = resolved + else: + from ..utils.model_manager import get_comfyui_base_path + + comfyui_base = get_comfyui_base_path() + if comfyui_base: + full_path = comfyui_base / tokenizer_path_val + if full_path.exists(): + config["tokenizer_path"] = str(full_path) + else: + config["tokenizer_path"] = str(full_path.resolve()) + else: + config["tokenizer_path"] = os.path.abspath(tokenizer_path_val) + else: + config["tokenizer_path"] = os.path.normpath(os.path.abspath(tokenizer_path_val)) + + return OmegaConf.create({"workflow_generator": config}) + + def build_node_validator_config( + self, + local: bool = False, + model_format: str = "gguf", + model_path: str = "", + tokenizer_path: str = "", + k: int = 5, + max_new_tokens: int = 4096, + **kwargs, + ) -> DictConfig: + """ + Build configuration for NodeValidator. + + Args: + local: Whether to use local model + model_format: Model format ("gguf" or "huggingface") + model_path: Path to model file/directory + tokenizer_path: Path to tokenizer directory + k: Number of similar nodes to consider + max_new_tokens: Maximum tokens to generate + **kwargs: Additional configuration overrides + + Returns: + OmegaConf DictConfig object + """ + config = self.base_config["node_validator"].copy() + + if local is not None: + config["local"] = local + if model_format: + config["model_format"] = model_format + if model_path: + config["model_path"] = model_path + # Only set tokenizer_path if explicitly provided (not None) + # If None, remove it from config so auto-detection happens + if tokenizer_path is not None: + if tokenizer_path: # Non-empty string + config["tokenizer_path"] = tokenizer_path + else: + # Empty string means remove from config (use auto-detection) + config.pop("tokenizer_path", None) + else: + # None means remove from config (use auto-detection) + config.pop("tokenizer_path", None) + if k is not None: + config["k"] = k + if max_new_tokens is not None: + config["max_new_tokens"] = max_new_tokens + + config.update(kwargs) + + # Model/tokenizer paths are in ComfyUI models/, not custom node dir + if "model_path" in config: + config["model_path"] = resolve_model_path(config["model_path"], model_type="llm") + if "tokenizer_path" in config: + tokenizer_path_val = config["tokenizer_path"] + if not os.path.isabs(tokenizer_path_val): + resolved = resolve_model_path(tokenizer_path_val, model_type="llm") + if os.path.exists(resolved) and os.path.isdir(resolved): + config["tokenizer_path"] = resolved + else: + from ..utils.model_manager import get_comfyui_base_path + + comfyui_base = get_comfyui_base_path() + if comfyui_base: + full_path = comfyui_base / tokenizer_path_val + if full_path.exists(): + config["tokenizer_path"] = str(full_path) + else: + config["tokenizer_path"] = str(full_path.resolve()) + else: + config["tokenizer_path"] = os.path.abspath(tokenizer_path_val) + else: + config["tokenizer_path"] = os.path.normpath(os.path.abspath(tokenizer_path_val)) + + return OmegaConf.create({"node_validator": config}) + + def build_node_base_config(self, catalog_directory: str = "", embedding_model_path: str | None = None, **kwargs) -> DictConfig: + """ + Build configuration for NodeBase. + + Args: + catalog_directory: Directory containing node_list.json and node_info.json + embedding_model_path: Path to embedding model (optional - only needed for semantic search) + **kwargs: Additional configuration overrides + + Returns: + OmegaConf DictConfig object + """ + config = self.base_config["node_base"].copy() + + if catalog_directory: + config["catalog_directory"] = catalog_directory + + if embedding_model_path is not None: + if embedding_model_path: + config["embedding_model_path"] = embedding_model_path + else: + config.pop("embedding_model_path", None) + + config.update(kwargs) + + if "catalog_directory" in config: + catalog_dir = resolve_config_path(config["catalog_directory"], self.custom_node_path) + config["node_list_path"] = os.path.join(catalog_dir, "node_list.json") + config["node_info_db_path"] = os.path.join(catalog_dir, "node_info.json") + config.pop("catalog_directory", None) + else: + config["node_list_path"] = resolve_config_path("catalog/node_list.json", self.custom_node_path) + config["node_info_db_path"] = resolve_config_path("catalog/node_info.json", self.custom_node_path) + + if "embedding_model_path" in config: + from ..utils.model_manager import resolve_model_path + + config["embedding_model_path"] = resolve_model_path(config["embedding_model_path"], model_type="llm") + + return OmegaConf.create({"node_base": config}) + + def build_full_config( + self, + workflow_generator_kwargs: dict | None = None, + node_validator_kwargs: dict | None = None, + node_base_kwargs: dict | None = None, + workflow_builder_kwargs: dict | None = None, + ) -> DictConfig: + """ + Build full configuration for all generators. + + Args: + workflow_generator_kwargs: WorkflowGenerator configuration overrides + node_validator_kwargs: NodeValidator configuration overrides + node_base_kwargs: NodeBase configuration overrides + + Returns: + OmegaConf DictConfig object with all generator configurations + """ + workflow_generator_cfg = self.build_workflow_generator_config(**(workflow_generator_kwargs or {})) + node_validator_cfg = self.build_node_validator_config(**(node_validator_kwargs or {})) + node_base_cfg = self.build_node_base_config(**(node_base_kwargs or {})) + + full_config = OmegaConf.merge(workflow_generator_cfg, node_validator_cfg, node_base_cfg) + + return full_config diff --git a/generators/node_validator.py b/generators/node_validator.py new file mode 100644 index 0000000..2700f88 --- /dev/null +++ b/generators/node_validator.py @@ -0,0 +1,751 @@ +""" +NodeValidator: Validates and corrects node names in workflow diagrams. +Enhanced with GPU layer management, error handling, and ComfyUI integration. +""" + +import json +import logging +import time +import traceback +from typing import Any + +import numpy as np +import torch +from sentence_transformers import SentenceTransformer +from sklearn.metrics.pairwise import cosine_similarity +from tinydb import Query, TinyDB + +try: + import nodes as comfy_nodes + + COMFY_NODES_AVAILABLE = True +except ImportError: + COMFY_NODES_AVAILABLE = False + logging.warning("ComfyUI nodes module not available, validation will use catalog only") + +from ..utils.model_loaders.device_utils import load_model +from ..utils.process_utils import fetch_node_input_names, fetch_node_output_names, format, json_format + + +class NODE_MISS_EXCEPTION(Exception): + """Exception raised when a node type is not found in the database.""" + + pass + + +class NodeBase: + """ + Base class for node information management and semantic search. + + Maintains a database of available ComfyUI nodes and provides + semantic search capabilities for finding similar nodes. + """ + + def __init__(self, cfg: Any): + """ + Initialize NodeBase with node database and optionally embedding model. + + Args: + cfg: Configuration object with node_base settings + """ + logging.info("Initializing NodeBase...") + + node_list_path = cfg.node_base.node_list_path + logging.debug(f"Loading node list from: {node_list_path}") + try: + with open(node_list_path) as f: + self.node_list = json.load(f) + logging.info(f"Loaded {len(self.node_list)} node types") + except Exception as e: + logging.error(f"Failed to load node list: {e}") + raise + + node_info_db_path = cfg.node_base.node_info_db_path + logging.debug(f"Loading node info database from: {node_info_db_path}") + try: + self.node_info_db = TinyDB(node_info_db_path) + node_info_count = len(self.node_info_db.all()) + logging.info(f"Loaded {node_info_count} node info entries") + except (TypeError, ValueError) as e: + # Handle case where file is in old format (plain JSON array instead of TinyDB format) + logging.warning("Node info database appears to be in old format, attempting to convert...") + try: + with open(node_info_db_path, encoding="utf-8") as f: + old_data = json.load(f) + + if isinstance(old_data, list): + db = TinyDB(node_info_db_path) + db.truncate() + db.insert_multiple(old_data) + db.close() + logging.info(f"Converted {len(old_data)} entries from old format to TinyDB format") + + self.node_info_db = TinyDB(node_info_db_path) + node_info_count = len(self.node_info_db.all()) + logging.info(f"Loaded {node_info_count} node info entries") + else: + raise ValueError(f"Unexpected data format in {node_info_db_path}") + except Exception as conv_e: + logging.error(f"Failed to convert node info database: {conv_e}") + logging.error("Please run 'UpdateNodeCatalog' node to regenerate the catalog files") + raise ValueError( + f"Node info database format is invalid. Please run 'UpdateNodeCatalog' node to regenerate it.\n" + f"Original error: {e}\n" + f"Conversion error: {conv_e}" + ) from conv_e + except Exception as e: + logging.error(f"Failed to load node info database: {e}") + raise + + self.cache_node_meta_info = {} + + # Load embedding model (optional - only needed for semantic search) + embedding_model_path = getattr(cfg.node_base, "embedding_model_path", None) + if embedding_model_path: + logging.debug(f"Loading embedding model: {embedding_model_path}") + try: + device = "cuda" if torch.cuda.is_available() else "cpu" + logging.debug(f"Using {device} for embedding model") + self.embedding_model = SentenceTransformer(embedding_model_path, device=device) + logging.debug("Computing embeddings for semantic search...") + self.list_embeddings = self.embedding_model.encode(self.node_list, batch_size=32, show_progress_bar=False) + logging.info(f"Embeddings computed for {len(self.list_embeddings)} nodes") + except Exception as e: + logging.error(f"Failed to load embedding model: {e}") + raise + else: + self.embedding_model = None + self.list_embeddings = None + logging.info("Embedding model not loaded (semantic search disabled)") + + def fetch_node_meta_info(self, node_type: str) -> dict | None: + """ + Fetch metadata information for a node type. + + Args: + node_type: Node type name + + Returns: + Node metadata dictionary or None if not found + + Raises: + NODE_MISS_EXCEPTION: If node type is not in the database + """ + if node_type not in self.node_list: + logging.warning(f"Node type '{node_type}' not in node list") + return None + + node_meta_info = self.cache_node_meta_info.get(node_type, None) + if node_meta_info is not None: + return node_meta_info + + query = Query() + node_meta_info_list = self.node_info_db.search(query.node_type == node_type) + + if len(node_meta_info_list) == 0: + logging.warning(f"Node type '{node_type}' not found in database") + raise NODE_MISS_EXCEPTION(node_type) + + node_meta_info = node_meta_info_list[0] + self.cache_node_meta_info[node_type] = node_meta_info + return node_meta_info + + def find_most_similar(self, node_type: str, k: int = 5) -> list[dict]: + """ + Find the k most similar nodes to the given node type using semantic search. + + Args: + node_type: Node type name to search for + k: Number of similar nodes to return + + Returns: + List of dictionaries containing node information + + Raises: + ValueError: If embedding model is not loaded + """ + if self.embedding_model is None or self.list_embeddings is None: + raise ValueError("Embedding model not loaded. Semantic search requires embedding_model_path to be provided.") + + query_embedding = self.embedding_model.encode([node_type]) + similarities = cosine_similarity(query_embedding, self.list_embeddings)[0] + top_k_indices = np.argsort(similarities)[-k:][::-1] + similar_node_list = [self.node_list[i] for i in top_k_indices] + + similar_node_infos = [] + for similar_node_type in similar_node_list: + try: + similar_node_infos.append(self.get_node_diagram_info(similar_node_type)) + except Exception as e: + logging.warning(f"Failed to get info for node '{similar_node_type}': {e}") + continue + + return similar_node_infos + + def not_in_sql(self, node_type: str) -> bool: + """ + Check if a node type exists in the node list. + + Args: + node_type: Node type name to check + + Returns: + True if node exists, False otherwise + """ + return node_type in self.node_list + + def is_valid_comfy_node(self, node_type: str) -> bool: + """ + Check if a node type is a valid ComfyUI node class name. + + This checks against ComfyUI's actual NODE_CLASS_MAPPINGS, not just the catalog. + The catalog may contain display names or aliases that aren't actual node class names. + + Args: + node_type: Node type name to check + + Returns: + True if node is a valid ComfyUI node class, False otherwise + """ + if not COMFY_NODES_AVAILABLE: + # Fallback to catalog check if ComfyUI nodes aren't available + return node_type in self.node_list + + # Check against ComfyUI's actual node registry + return node_type in comfy_nodes.NODE_CLASS_MAPPINGS + + def get_node_diagram_info(self, node_type: str) -> dict: + """ + Get diagram information for a node type with enhanced metadata. + + Args: + node_type: Node type name + + Returns: + Dictionary with node_name, input_names, output_names, and additional metadata + """ + meta_info = self.fetch_node_meta_info(node_type) + input_names = fetch_node_input_names(meta_info) + output_names = fetch_node_output_names(meta_info) + + info = {"node_name": node_type, "input_names": input_names, "output_names": output_names} + + if meta_info: + if "category" in meta_info: + info["category"] = meta_info["category"] + if "display_name" in meta_info: + info["display_name"] = meta_info["display_name"] + if "description" in meta_info: + info["description"] = meta_info["description"] + if "is_deprecated" in meta_info: + info["is_deprecated"] = meta_info["is_deprecated"] + if "is_experimental" in meta_info: + info["is_experimental"] = meta_info["is_experimental"] + if "is_output_node" in meta_info: + info["is_output_node"] = meta_info["is_output_node"] + if "is_api_node" in meta_info: + info["is_api_node"] = meta_info["is_api_node"] + if "input_metadata" in meta_info: + info["input_metadata"] = meta_info["input_metadata"] + if "output_display_names" in meta_info: + info["output_display_names"] = meta_info["output_display_names"] + if "output_tooltips" in meta_info: + info["output_tooltips"] = meta_info["output_tooltips"] + + return info + + +class NodeValidator: + """ + Validator responsible for validating and correcting node names in workflow diagrams. + + Uses semantic search and optionally a language model to fix invalid node names + by finding the most similar valid node names. + """ + + def __init__( + self, + node_base: NodeBase, + cfg: Any, + n_gpu_layers: int | str = "auto", + dtype: str | torch.dtype = "auto", + device_preference: str = "auto", + attn_implementation: str = "auto", + context_size: int = 4096, + temperature: float = 0.95, + top_p: float = 0.7, + use_mmap: bool = True, + use_mlock: bool = False, + n_batch: int = 512, + n_threads: int | None = None, + ): + """ + Initialize NodeValidator. + + Args: + node_base: NodeBase instance for node information + cfg: Configuration object with node_validator settings + n_gpu_layers: For GGUF: "auto", -1 for all, 0 for CPU, or specific count + dtype: For HuggingFace: "auto", "fp16", "bf16", "fp32", "fp8", or torch.dtype + device_preference: "auto", "cuda", or "cpu" + """ + logging.info("WorkflowGenerator NodeValidator: Initializing...") + + self.k = getattr(cfg.node_validator, "k", 5) + self.node_base = node_base + self.temperature = temperature + self.top_p = top_p + self.local = getattr(cfg.node_validator, "local", False) + + if self.local: + self.max_new_tokens = getattr(cfg.node_validator, "max_new_tokens", 4096) + model_format = getattr(cfg.node_validator, "model_format", "huggingface") + model_path = cfg.node_validator.model_path + tokenizer_path = getattr(cfg.node_validator, "tokenizer_path", None) + + logging.info(f"WorkflowGenerator NodeValidator: Loading local model: {model_format}, path: {model_path}") + + load_kwargs = {} + + if model_format.lower() == "huggingface": + load_kwargs = { + "dtype": dtype, + "device_preference": device_preference, + "attn_implementation": attn_implementation, + } + else: + # GGUF models use n_gpu_layers, context_size, and performance parameters + load_kwargs = { + "n_gpu_layers": n_gpu_layers, + "device_preference": device_preference, + "context_size": context_size, + "use_mmap": use_mmap, + "use_mlock": use_mlock, + "n_batch": n_batch, + "n_threads": n_threads, + } + + load_start = time.time() + try: + self.model_wrapper = load_model( + model_format=model_format, model_path=model_path, tokenizer_path=tokenizer_path, **load_kwargs + ) + load_time = time.time() - load_start + logging.info(f"WorkflowGenerator NodeValidator: Model loaded in {load_time:.2f}s") + except Exception as e: + logging.error(f"WorkflowGenerator NodeValidator: Failed to load model: {e}") + raise + + self.model = self.model_wrapper.model + self.tokenizer = self.model_wrapper.tokenizer + + try: + if torch.cuda.is_available() and model_format.lower() == "huggingface": + vram_used = torch.cuda.memory_allocated() / 1024**3 + logging.debug(f"GPU memory allocated: {vram_used:.2f} GB") + except Exception: + pass + else: + logging.info("WorkflowGenerator NodeValidator: Using non-local mode (semantic search only)") + + def refine_diagram(self, diagram: list, prompt: str, seed: int | None = None) -> tuple[list, dict]: + """ + Refine a workflow diagram by validating and correcting node names. + + Args: + diagram: List of edges representing the workflow diagram + prompt: Original user prompt/description + seed: Random seed for reproducible refinement (only used when local=True) + + Returns: + Tuple of (refined_diagram, debug_info) where: + - refined_diagram: Refined diagram with corrected node names + - debug_info: Dictionary containing: + - llm_prompts: List of prompts sent to LLM (one per invalid node) + - candidate_nodes: Dictionary mapping invalid node names to their candidate nodes + """ + if seed is not None: + max_seed = 2**32 - 1 + if seed < 0 or seed > max_seed: + seed = seed % (max_seed + 1) # Wrap around if out of range + + logging.info("Refining workflow diagram...") + + invalid_nodes = [] + swap_map = {} + debug_info = {"llm_prompts": [], "candidate_nodes": {}} + + logging.debug("Scanning diagram for invalid node names...") + + for edge in diagram: + src_type, src_name, dst_type, dst_name = edge + src_node = format(src_type) + dst_node = format(dst_type) + + is_in_catalog = self.node_base.not_in_sql(src_node) + is_valid_comfy = self.node_base.is_valid_comfy_node(src_node) + + if (not is_in_catalog or not is_valid_comfy) and src_node not in swap_map: + invalid_nodes.append(src_node) + if not is_in_catalog: + logging.debug(f"Source node '{src_node}' not found in node database") + elif not is_valid_comfy: + logging.debug(f"Source node '{src_node}' is not a valid ComfyUI node class (may be display name/alias)") + + try: + candidate_nodes = self.node_base.find_most_similar(src_node, self.k) + debug_info["candidate_nodes"][src_node] = candidate_nodes + except Exception as e: + logging.warning(f"Failed to get candidate nodes for '{src_node}': {e}") + debug_info["candidate_nodes"][src_node] = [] + + if self.local: + swap_name, llm_prompt = self.fix_node_name_local(prompt, diagram, src_node, seed=seed, return_prompt=True) + if llm_prompt: + debug_info["llm_prompts"].append({"invalid_node": src_node, "prompt": llm_prompt}) + else: + swap_name = self.fix_node_name(prompt, diagram, src_node) + if swap_name is not None: + swap_map[src_node] = swap_name + logging.info(f"Fixed node name: '{src_node}' -> '{swap_name}'") + + is_in_catalog = self.node_base.not_in_sql(dst_node) + is_valid_comfy = self.node_base.is_valid_comfy_node(dst_node) + + if (not is_in_catalog or not is_valid_comfy) and dst_node not in swap_map: + invalid_nodes.append(dst_node) + if not is_in_catalog: + logging.debug(f"Destination node '{dst_node}' not found in node database") + elif not is_valid_comfy: + logging.debug(f"Destination node '{dst_node}' is not a valid ComfyUI node class (may be display name/alias)") + + try: + candidate_nodes = self.node_base.find_most_similar(dst_node, self.k) + debug_info["candidate_nodes"][dst_node] = candidate_nodes + except Exception as e: + logging.warning(f"Failed to get candidate nodes for '{dst_node}': {e}") + debug_info["candidate_nodes"][dst_node] = [] + + if self.local: + swap_name, llm_prompt = self.fix_node_name_local(prompt, diagram, dst_node, seed=seed, return_prompt=True) + if llm_prompt: + debug_info["llm_prompts"].append({"invalid_node": dst_node, "prompt": llm_prompt}) + else: + swap_name = self.fix_node_name(prompt, diagram, dst_node) + if swap_name is not None: + swap_map[dst_node] = swap_name + logging.info(f"Fixed node name: '{dst_node}' -> '{swap_name}'") + + logging.debug("Validating connections (input/output names)...") + connection_issues = [] + + for i, edge in enumerate(diagram): + src_type, src_name, dst_type, dst_name = edge + src_node = format(src_type) + dst_node = format(dst_type) + + if src_node in swap_map: + src_node = swap_map[src_node] + edge[0] = src_node + if dst_node in swap_map: + dst_node = swap_map[dst_node] + edge[2] = dst_node + + try: + src_meta = self.node_base.fetch_node_meta_info(src_node) + if src_meta: + src_output_names = fetch_node_output_names(src_meta) + if src_name not in src_output_names: + connection_issues.append( + {"edge_index": i, "node": src_node, "type": "output", "name": src_name, "available": src_output_names} + ) + logging.debug( + f"Source node '{src_node}' has no output '{src_name}'. Available: {src_output_names[:5]}..." + ) + except Exception as e: + logging.warning(f"Failed to validate source node '{src_node}' output: {e}") + + try: + dst_meta = self.node_base.fetch_node_meta_info(dst_node) + if dst_meta: + dst_input_names = fetch_node_input_names(dst_meta) + if dst_name not in dst_input_names: + connection_issues.append( + {"edge_index": i, "node": dst_node, "type": "input", "name": dst_name, "available": dst_input_names} + ) + logging.debug( + f"Destination node '{dst_node}' has no input '{dst_name}'. Available: {dst_input_names[:5]}..." + ) + + if dst_node not in swap_map: + try: + candidate_nodes = self.node_base.find_most_similar(dst_node, self.k) + debug_info["candidate_nodes"][dst_node] = candidate_nodes + + for candidate in candidate_nodes: + candidate_meta = self.node_base.fetch_node_meta_info(candidate["node_name"]) + if candidate_meta: + candidate_inputs = fetch_node_input_names(candidate_meta) + if dst_name in candidate_inputs: + logging.debug(f"Candidate node '{candidate['node_name']}' has input '{dst_name}'") + if self.local: + connection_issue_info = {"type": "input", "name": dst_name, "available": dst_input_names} + swap_name, llm_prompt = self.fix_node_name_local( + prompt, + diagram, + dst_node, + seed=seed, + return_prompt=True, + preferred_candidate=candidate["node_name"], + connection_issue=connection_issue_info, + ) + if llm_prompt: + debug_info["llm_prompts"].append( + { + "invalid_node": dst_node, + "invalid_connection": f"input '{dst_name}'", + "prompt": llm_prompt, + } + ) + if swap_name and swap_name != dst_node: + swap_map[dst_node] = swap_name + logging.info(f"Fixed node name: '{dst_node}' -> '{swap_name}' (has required input '{dst_name}')") + break + else: + swap_map[dst_node] = candidate["node_name"] + logging.info( + f"Fixed node name: '{dst_node}' -> '{candidate['node_name']}' (has required input '{dst_name}')" + ) + break + except Exception as e: + logging.warning(f"Failed to find candidate nodes for '{dst_node}': {e}") + except Exception as e: + logging.warning(f"Failed to validate destination node '{dst_node}' input: {e}") + + if len(swap_map) == 0 and len(connection_issues) == 0: + logging.info("Diagram refinement complete: no corrections needed") + else: + if len(swap_map) > 0: + logging.debug(f"Applying {len(swap_map)} node name corrections...") + for i, edge in enumerate(diagram): + src_type, src_name, dst_type, dst_name = edge + src_node = format(src_type) + dst_node = format(dst_type) + + if src_node in swap_map: + swap_name = swap_map[src_node] + edge[0] = edge[0].replace(src_node, swap_name) + + if dst_node in swap_map: + swap_name = swap_map[dst_node] + edge[2] = edge[2].replace(dst_node, swap_name) + + diagram[i] = edge + + if len(connection_issues) > 0: + logging.warning(f"Found {len(connection_issues)} connection issues") + + logging.info(f"Diagram refinement complete: {len(swap_map)} node corrections, {len(connection_issues)} connection issues") + + return diagram, debug_info + + def fix_node_name(self, desc: str, diagram: list, error_name: str) -> str | None: + """ + Fix node name using semantic search only (non-local mode). + + Args: + desc: User description + diagram: Workflow diagram + error_name: Invalid node name to fix + + Returns: + Corrected node name or None if no suitable match found + """ + logging.info(f"Fixing node name '{error_name}' using semantic search") + + try: + candidate_nodes = self.node_base.find_most_similar(error_name, self.k) + if candidate_nodes: + return candidate_nodes[0]["node_name"] + except Exception as e: + logging.warning(f"Failed to find similar node for '{error_name}': {e}") + + return None + + def fix_node_name_local( + self, + desc: str, + diagram: list, + error_name: str, + seed: int | None = None, + return_prompt: bool = False, + preferred_candidate: str | None = None, + connection_issue: dict | None = None, + ) -> str | None | tuple[str | None, str | None]: + """ + Fix node name using local language model. + + Args: + desc: User description + diagram: Workflow diagram + error_name: Invalid node name to fix + seed: Random seed for reproducible generation (optional) + return_prompt: If True, return tuple of (result, full_prompt_text), else just result + preferred_candidate: Specific candidate to prioritize (e.g., one that has required input/output) + connection_issue: Dictionary describing the connection problem, e.g.: + {"type": "input", "name": "text_list", "available": ["text_a", "text_b", ...]} + + Returns: + If return_prompt=False: Corrected node name or None if no suitable match found + If return_prompt=True: Tuple of (corrected_node_name, full_prompt_text) where full_prompt_text is the complete prompt sent to LLM + """ + if not self.local: + result = self.fix_node_name(desc, diagram, error_name) + if return_prompt: + return result, None + return result + + logging.info(f"Fixing node name '{error_name}' using local model") + + try: + candidate_nodes = self.node_base.find_most_similar(error_name, self.k) + candidate_node_names = [e["node_name"] for e in candidate_nodes] + + if preferred_candidate and preferred_candidate in candidate_node_names: + preferred_nodes = [] + other_nodes = [] + for c in candidate_nodes: + if c["node_name"] == preferred_candidate: + c_copy = c.copy() + c_copy["_recommended"] = True + if connection_issue: + c_copy["_reason"] = f"Has required {connection_issue['type']}: '{connection_issue['name']}'" + preferred_nodes.append(c_copy) + else: + other_nodes.append(c) + candidate_nodes = preferred_nodes + other_nodes + candidate_node_names = [e["node_name"] for e in candidate_nodes] + + if connection_issue: + issue_type = connection_issue.get("type", "input") + missing_name = connection_issue.get("name", "unknown") + available_names = connection_issue.get("available", []) + available_str = ", ".join([f"'{n}'" for n in available_names[:10]]) + if len(available_names) > 10: + available_str += ", ..." + + prompt = ( + f"I would like you to act as an expert in ComfyUI platform. " + f"I will provide a workflow description and a logical diagram in JSON format representing the ComfyUI workflow. " + f"The logical diagram is a list of edges [edge_1, edge_2, ...], where each edge is " + f"[source_node_name, output_name, destination_node_name, input_name], representing a connection from source to destination.\n\n" + f"**Workflow Description:** {desc if desc else 'Not provided'}\n\n" + f"**Logical Diagram:** {str(diagram)}\n\n" + f"**PROBLEM DETECTED:**\n" + f"The workflow uses node '{error_name}', but there is a CONNECTION ERROR:\n" + f"- The workflow tries to connect to {issue_type} '{missing_name}'\n" + f"- However, node '{error_name}' does NOT have this {issue_type}\n" + f"- Available {issue_type}s on '{error_name}': {available_str}\n\n" + f"**YOUR TASK:**\n" + f"Select a replacement node from the candidates below that:\n" + f"1. Has the required {issue_type} '{missing_name}' (CRITICAL)\n" + f"2. Is semantically similar to '{error_name}'\n" + f"3. Fits the workflow's purpose\n\n" + f"**Candidate Nodes:** {str(candidate_nodes)}\n\n" + f"**IMPORTANT NOTES:**\n" + f"- Look for candidates with '_recommended': True - these have the required {issue_type}\n" + f"- Check each candidate's '{issue_type}_names' list carefully\n" + f"- The replacement MUST have '{missing_name}' in its {issue_type}s\n\n" + f"**Response Format:**\n" + f"Return ONLY the chosen node name in pure JSON format:\n" + f"'''json{{\"candidate_node_name\":\"...\"}}'''" + ) + else: + prompt = ( + f"I would like you to act as an expert in ComfyUI platform. " + f"I will provide a workflow description and a logical diagram in JSON format representing the ComfyUI workflow. " + f"The logical diagram is a list of edges [edge_1, edge_2, ...], where each edge is " + f"[source_node_name, output_name, destination_node_name, input_name], representing a connection from source to destination.\n\n" + f"**Workflow Description:** {desc if desc else 'Not provided'}\n\n" + f"**Logical Diagram:** {str(diagram)}\n\n" + f"**PROBLEM DETECTED:**\n" + f"The node name '{error_name}' is not valid in ComfyUI.\n\n" + f"**YOUR TASK:**\n" + f"Select the most suitable replacement node from the candidates below that:\n" + f"1. Is semantically similar to '{error_name}'\n" + f"2. Fits the workflow's purpose based on the description and diagram\n" + f"3. Has compatible inputs/outputs for the connections in the diagram\n\n" + f"**Candidate Nodes:** {str(candidate_nodes)}\n\n" + f"**Response Format:**\n" + f"Return ONLY the chosen node name in pure JSON format:\n" + f"'''json{{\"candidate_node_name\":\"...\"}}'''" + ) + + messages = [ + {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."}, + {"role": "user", "content": prompt}, + ] + + # Use HuggingFace tokenizer even for GGUF models (embedded tokenizer doesn't expose apply_chat_template) + text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + + full_prompt_text = text if return_prompt else None + + model_inputs = self.tokenizer([text], return_tensors="pt") + if hasattr(self.model_wrapper, "device") and isinstance(self.model_wrapper.device, torch.device): + model_inputs = model_inputs.to(self.model_wrapper.device) + + generate_kwargs = {"max_new_tokens": self.max_new_tokens, "temperature": self.temperature, "top_p": self.top_p} + if seed is not None: + max_seed = 2**32 - 1 + if seed < 0 or seed > max_seed: + seed = seed % (max_seed + 1) # Wrap around if out of range + generate_kwargs["seed"] = seed + + generated_ids = self.model_wrapper.generate(model_inputs.input_ids, **generate_kwargs) + + generated_ids = [output_ids[len(input_ids) :] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)] + + result = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] + + try: + result = json_format(result) + if isinstance(result, dict): + result = result.get("candidate_node_name", None) + except Exception as e: + logging.warning(f"Failed to parse JSON response: {e}") + traceback.print_exc() + if return_prompt: + return None, full_prompt_text + return None + + if result not in candidate_node_names: + logging.warning(f"Model returned '{result}' which is not in candidate list") + if return_prompt: + return None, full_prompt_text + return None + + if return_prompt: + return result, full_prompt_text + return result + + except Exception as e: + logging.error(f"Failed to fix node name using local model: {e}") + traceback.print_exc() + if return_prompt: + return None, None + return None + + def cleanup(self): + """Clean up model resources.""" + if hasattr(self, "model_wrapper"): + del self.model_wrapper + if hasattr(self, "model"): + del self.model + if hasattr(self, "tokenizer"): + del self.tokenizer + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + logging.info("WorkflowGenerator NodeValidator: Resources cleaned up") diff --git a/generators/workflow_builder.py b/generators/workflow_builder.py new file mode 100644 index 0000000..c1c37ef --- /dev/null +++ b/generators/workflow_builder.py @@ -0,0 +1,529 @@ +""" +WorkflowBuilder: Converts workflow diagrams to ComfyUI workflow JSON. +Enhanced with error handling and ComfyUI integration. +""" + +import logging +from typing import Any + +from ..utils.process_utils import format + +NO_LINK_TYPES = ["STRING", "FLOAT", "INT", "BOOLEAN"] + + +class NAME_EXCEPTION(Exception): + """Exception raised when a node name is invalid.""" + + pass + + +class REQUIR_MISS_EXCEPTION(Exception): + """Exception raised when required inputs are missing.""" + + pass + + +class WorkflowBuilder: + """ + Builder responsible for converting workflow diagrams to ComfyUI workflow JSON. + + Takes a diagram (list of edges) and converts it into a complete ComfyUI + workflow JSON structure with nodes, links, and metadata. + """ + + def __init__(self, node_base: Any): + """ + Initialize WorkflowBuilder. + + Args: + node_base: NodeBase instance for node information + """ + logging.info("WorkflowGenerator WorkflowBuilder: Initializing...") + self.node_base = node_base + + def parse_diagram_to_workflow(self, diagram: list, version: float = 0.4, type_to_pos: dict | None = None) -> dict: + """ + Parse a workflow diagram into ComfyUI workflow JSON. + + Args: + diagram: List of edges representing the workflow diagram + version: Workflow version (default: 0.4) + type_to_pos: Optional dictionary mapping node types to positions + + Returns: + Complete ComfyUI workflow dictionary + """ + logging.info("Parsing diagram to workflow...") + logging.debug("Processing diagram edges...") + + count = 1 + type_to_id = {} + links = [] + link_count = 1 + inputs_dict = {} + outputs_dict = {} + catch_candidate_inputs = {} + processed_edges = 0 + + for i, (src_node, src_name, dst_node, dst_name) in enumerate(diagram): + processed_edges += 1 + + if src_node not in type_to_id: + type_to_id[src_node] = count + count += 1 + if dst_node not in type_to_id: + type_to_id[dst_node] = count + count += 1 + + src_id = type_to_id[src_node] + dst_id = type_to_id[dst_node] + src_type = format(src_node) + dst_type = format(dst_node) + + if src_id not in outputs_dict: + try: + outputs_dict[src_id] = self.fetch_node_outputs(src_type) + except Exception as e: + logging.error(f"Failed to fetch outputs for node '{src_type}': {e}") + raise + + if dst_id not in inputs_dict: + try: + inputs_dict[dst_id], catch_candidate_inputs[dst_type] = self.fetch_node_inputs(dst_type) + except Exception as e: + logging.error(f"Failed to fetch inputs for node '{dst_type}': {e}") + raise + + outputs = outputs_dict[src_id] + inputs = inputs_dict[dst_id] + candidate_inputs = catch_candidate_inputs[dst_type] + + try: + outputs, inputs, link, link_count = decode_edge( + outputs, + inputs, + candidate_inputs, + link_count, + src_id, + dst_id, + src_name, + dst_name, + dst_type, + ) + + # Only add link if decode_edge returned one (None means edge was skipped gracefully) + if link is not None: + links.append(link) + link_count += 1 + else: + # Edge was skipped - don't increment link_count to avoid gaps + logging.debug(f"Edge {i} skipped (no valid connection found)") + + except NAME_EXCEPTION as e: + logging.warning(f"Failed to decode edge {i}: {e}. Skipping edge.") + logging.debug(f" Source node: {src_type}, Destination node: {dst_type}") + logging.debug(f" Output name: {src_name}, Input name: {dst_name}") + continue + except Exception as e: + logging.error(f"Failed to decode edge {i}: {e}") + logging.error(f" Source node: {src_type}, Destination node: {dst_type}") + logging.error(f" Output name: {src_name}, Input name: {dst_name}") + raise + + outputs_dict[src_id] = outputs + inputs_dict[dst_id] = inputs + + logging.debug(f"Processed {processed_edges} edges, created {len(type_to_id)} unique nodes") + logging.debug("Building workflow structure...") + + for id, inputs in inputs_dict.items(): + tmp_input = [] + for input in inputs: + if len(input) > 0: + tmp_input.append(input) + inputs_dict[id] = tmp_input + + nodes = [] + for node_type, node_id in type_to_id.items(): + inputs = [] + if node_id in inputs_dict: + inputs = inputs_dict[node_id] + + outputs = [] + if node_id in outputs_dict: + outputs = outputs_dict[node_id] + + node = { + "id": node_id, + "type": format(node_type), + "inputs": inputs, + "outputs": outputs, + } + + if type_to_pos is not None and node_type in type_to_pos: + node["pos"] = type_to_pos[node_type]["pos"] + node["size"] = type_to_pos[node_type]["size"] + + nodes.append(node) + + if type_to_pos is None: + try: + from ..utils.workflow_layout import WorkflowLayout + + logging.info("WorkflowGenerator WorkflowBuilder: Applying auto-layout to generated workflow") + WorkflowLayout.process(nodes, links) + except Exception as e: + logging.warning(f"WorkflowGenerator WorkflowBuilder: Failed to apply auto-layout: {e}") + + decode_payload = { + "last_link_id": link_count - 1, + "nodes": nodes, + "links": links, + "version": version, + } + + decode_payload["extra"] = { + "extra": { + "ds": { + "scale": 0.952486143572623, + "offset": [163.72358533046707, 22.018287636666063], + } + } + } + + logging.info(f"Workflow structure created: {len(nodes)} nodes, {len(links)} links") + + return decode_payload + + def fetch_node_inputs(self, node_type: str) -> tuple[list, list]: + """ + Fetch input information for a node type. + + Args: + node_type: Node type name + + Returns: + Tuple of (inputs list, candidate_inputs list) + """ + try: + meta_info = self.node_base.fetch_node_meta_info(node_type) + except Exception as e: + logging.error(f"Failed to fetch meta info for '{node_type}': {e}") + raise + + if meta_info is None: + logging.warning(f"Node type '{node_type}' not found in catalog, returning empty inputs") + return [], [] + + input_types = meta_info.get("inputs", {}) + if input_types is None: + input_types = {} + + inputs = [] + else_inputs = [] + + for key, item in input_types.items(): + if key != "hidden": + for name, val in item.items(): + val_type = val[0] + + # Skip list types (COMBO) + if isinstance(val_type, list): + continue + + # Non-linkable types (primitives) + if val_type not in NO_LINK_TYPES: + input = {"name": name, "type": val_type} + inputs.append(input) + elif len(val) > 1: + # Optional inputs with default values + val_info = val[1] + if "defaultInput" in val_info: + input = {"name": name, "type": val_type} + else_inputs.append(input) + + inputs.extend(else_inputs) + candidate_inputs = get_node_candidate_inputs(input_types, inputs) + + return inputs, candidate_inputs + + def fetch_node_outputs(self, node_type: str) -> list[dict]: + """ + Fetch output information for a node type. + + Args: + node_type: Node type name + + Returns: + List of output dictionaries + """ + try: + node_meta_info = self.node_base.fetch_node_meta_info(node_type) + except Exception as e: + logging.error(f"Failed to fetch meta info for '{node_type}': {e}") + raise + + if node_meta_info is None: + logging.warning(f"Node type '{node_type}' not found in catalog, returning empty outputs") + return [] + + output_types = node_meta_info.get("outputs", []) + if output_types is None: + output_types = [] + + output_names = node_meta_info.get("output_names", None) + outputs = [] + + for i, output_type in enumerate(output_types): + output = {"type": output_type, "slot_index": i, "links": []} + + if output_names is not None and i < len(output_names): + output["name"] = output_names[i] + else: + output["name"] = output_type + + outputs.append(output) + + return outputs + + +def get_output_info(outputs: list[dict]) -> dict[str, int]: + """ + Extract output name to index mapping from outputs list. + + Args: + outputs: List of output dictionaries + + Returns: + Dictionary mapping output names to indices + """ + has_name = False + error_tag = False + type_dict = [] + output_names = {} + + for index, val in enumerate(outputs): + if "name" in val: + has_name = True + if val["type"] in type_dict: + error_tag = True + type_dict.append(val["type"]) + output_names[val["name"]] = index + + if not has_name and error_tag: + raise ValueError("Outputs format error: duplicate types without names") + + return output_names + + +def get_node_candidate_inputs(input_types: dict, inputs: list[dict]) -> list[dict]: + """ + Get candidate inputs that can be added dynamically. + + Args: + input_types: Dictionary of input type definitions + inputs: List of existing inputs + + Returns: + List of candidate input dictionaries + """ + candidate_inputs = [] + inputs_names = [input["name"] for input in inputs] + + for key, item in input_types.items(): + if key != "hidden": + for name, val in item.items(): + val_type = val[0] + + if name in inputs_names: + continue + + if val_type not in NO_LINK_TYPES and not isinstance(val, list): + continue + + input = {"name": name, "type": val_type} + candidate_inputs.append(input) + + return candidate_inputs + + +def decode_edge( + outputs: list[dict], + inputs: list[dict], + candidate_inputs: list[dict], + link_count: int, + src_id: int, + dst_id: int, + output_name: str, + input_name: str, + dst_node_type: str = None, +) -> tuple[list[dict], list[dict], list, int]: + """ + Decode an edge and create a link between source and destination nodes. + + Args: + outputs: List of output dictionaries from source node + inputs: List of input dictionaries from destination node + candidate_inputs: List of candidate inputs that can be added + link_count: Current link count + src_id: Source node ID + dst_id: Destination node ID + output_name: Name of the output port + input_name: Name of the input port + + Returns: + Tuple of (updated outputs, updated inputs, link, link_count) + """ + output_names = get_output_info(outputs) + + if output_name not in output_names: + raise NAME_EXCEPTION(f"Output name '{output_name}' not found") + + src_port = output_names[output_name] + output = outputs[src_port] + + links = output.get("links", []) + links.append(link_count) + output["links"] = links + outputs[src_port] = output + + input = None + dst_port = None + + for i, _input in enumerate(inputs): + if _input.get("name") == input_name: + input = _input + input["link"] = link_count + dst_port = i + break + + # If not found, try candidate inputs + if input is None: + for _input in candidate_inputs: + if _input.get("name") == input_name: + input = { + "name": input_name, + "type": _input["type"], + "link": link_count, + "widget": {"name": input_name}, + } + dst_port = len(inputs) + inputs.append(input) + break + + # If still not found, try fuzzy matching on input names + if input is None: + all_input_names = [inp.get("name") for inp in inputs + candidate_inputs if inp.get("name")] + + best_match = None + best_score = 0 + + input_name_normalized = input_name.lower() + for suffix in ["_list", "_array", "s", "_items", "_values"]: + if input_name_normalized.endswith(suffix): + input_name_base = input_name_normalized[: -len(suffix)] + break + else: + input_name_base = input_name_normalized + + for available_name in all_input_names: + available_normalized = available_name.lower() + + if available_normalized == input_name_normalized: + best_match = available_name + best_score = 1.0 + break + + if available_normalized == input_name_base or input_name_base == available_normalized: + best_match = available_name + best_score = 0.9 + continue + + if input_name_normalized in available_normalized or available_normalized in input_name_normalized: + length_score = min(len(input_name_normalized), len(available_normalized)) / max( + len(input_name_normalized), len(available_normalized) + ) + if available_normalized.startswith(input_name_base) or input_name_base.startswith(available_normalized): + length_score = max(length_score, 0.7) + if length_score > best_score: + best_match = available_name + best_score = length_score + + if best_match and best_score > 0.5: + logging.warning( + f"Input name '{input_name}' not found for node, using similar name '{best_match}' (similarity: {best_score:.2f})" + ) + for _input in inputs + candidate_inputs: + if _input.get("name") == best_match: + input = { + "name": best_match, + "type": _input["type"], + "link": link_count, + } + if best_match not in [inp.get("name") for inp in inputs]: + input["widget"] = {"name": best_match} + dst_port = len(inputs) + inputs.append(input) + else: + for i, inp in enumerate(inputs): + if inp.get("name") == best_match: + inp["link"] = link_count + dst_port = i + input = inp + break + break + else: + type_matched_inputs = [inp for inp in inputs + candidate_inputs if inp.get("type") == output.get("type") and inp.get("name")] + + if type_matched_inputs: + best_type_match = None + for matched_input in type_matched_inputs: + matched_name = matched_input.get("name", "").lower() + if input_name_base in matched_name or matched_name.startswith(input_name_base): + best_type_match = matched_input + break + + if best_type_match is None: + best_type_match = type_matched_inputs[0] + + logging.warning( + f"Input name '{input_name}' not found for node '{dst_node_type}', " + f"using type-compatible input '{best_type_match.get('name')}' " + f"(type: {output.get('type')})" + ) + + matched_name = best_type_match.get("name") + if matched_name not in [inp.get("name") for inp in inputs]: + input = { + "name": matched_name, + "type": best_type_match.get("type"), + "link": link_count, + "widget": {"name": matched_name}, + } + dst_port = len(inputs) + inputs.append(input) + else: + for i, inp in enumerate(inputs): + if inp.get("name") == matched_name: + inp["link"] = link_count + dst_port = i + input = inp + break + else: + suggestions = ", ".join(all_input_names[:5]) + node_info = f" for node '{dst_node_type}'" if dst_node_type else "" + logging.warning( + f"Cannot connect edge: Input name '{input_name}' not found{node_info}. " + f"Output type: {output.get('type')}. " + f"Available inputs: {suggestions}" + ("..." if len(all_input_names) > 5 else "") + ". " + "Skipping this edge." + ) + return outputs, inputs, None, link_count + + if input["type"] != output["type"]: + logging.warning(f"Type mismatch for edge: output type '{output['type']}' != input type '{input['type']}'. Skipping this edge.") + return outputs, inputs, None, link_count + + link = [link_count, src_id, src_port, dst_id, dst_port, input["type"]] + + return outputs, inputs, link, link_count diff --git a/generators/workflow_generator.py b/generators/workflow_generator.py new file mode 100644 index 0000000..c453ed3 --- /dev/null +++ b/generators/workflow_generator.py @@ -0,0 +1,245 @@ +""" +WorkflowGenerator: Generates workflow diagrams from natural language descriptions. +Enhanced with GPU layer management, error handling, and ComfyUI integration. +""" + +import json +import logging +import time +from typing import Any + +import torch + +from ..utils.model_loaders.device_utils import load_model +from ..utils.process_utils import del_digram_primitive + + +class WorkflowGenerator: + """ + Generator responsible for creating workflow diagrams from natural language. + + Uses a language model to convert user descriptions into structured workflow + diagrams represented as JSON edge lists. + """ + + def __init__( + self, + cfg: Any, + n_gpu_layers: int | str = "auto", + dtype: str | torch.dtype = "auto", + device_preference: str = "auto", + attn_implementation: str = "auto", + context_size: int = 4096, + use_mmap: bool = True, + use_mlock: bool = False, + n_batch: int = 512, + n_threads: int | None = None, + ): + """ + Initialize WorkflowGenerator with model configuration. + + Args: + cfg: Configuration object with workflow_generator settings + n_gpu_layers: For GGUF: "auto", -1 for all, 0 for CPU, or specific count + dtype: For HuggingFace: "auto", "fp16", "bf16", "fp32", "fp8", or torch.dtype + device_preference: "auto", "cuda", or "cpu" + attn_implementation: For HuggingFace: "auto", "flash_attention_2", or "sdpa" + """ + logging.info("WorkflowGenerator: Initializing...") + + # Get model format from config (default to huggingface for backward compatibility) + model_format = getattr(cfg.workflow_generator, "model_format", "huggingface") + model_path = cfg.workflow_generator.model_path + tokenizer_path = getattr(cfg.workflow_generator, "tokenizer_path", None) + + logging.info(f"Model format: {model_format}, path: {model_path}") + + load_kwargs = {} + + if model_format.lower() == "huggingface": + # HuggingFace models use dtype, device_map, and attention implementation + load_kwargs = { + "dtype": dtype, + "device_preference": device_preference, + "attn_implementation": attn_implementation, + } + logging.debug("Model format: HuggingFace") + else: + # GGUF models use n_gpu_layers, context_size, and performance parameters + load_kwargs = { + "n_gpu_layers": n_gpu_layers, + "device_preference": device_preference, + "context_size": context_size, + "use_mmap": use_mmap, + "use_mlock": use_mlock, + "n_batch": n_batch, + "n_threads": n_threads, + } + logging.debug("Model format: GGUF (already quantized)") + + load_start = time.time() + try: + self.model_wrapper = load_model(model_format=model_format, model_path=model_path, tokenizer_path=tokenizer_path, **load_kwargs) + load_time = time.time() - load_start + logging.info(f"Model loaded in {load_time:.2f}s") + except Exception as e: + logging.error(f"Failed to load model: {e}") + raise + + self.model = self.model_wrapper.model + self.tokenizer = self.model_wrapper.tokenizer + + try: + if torch.cuda.is_available() and model_format.lower() == "huggingface": + vram_used = torch.cuda.memory_allocated() / 1024**3 + logging.debug(f"GPU memory allocated: {vram_used:.2f} GB") + except Exception: + pass + + self.top_p = getattr(cfg.workflow_generator, "top_p", 0.7) + self.max_new_tokens = getattr(cfg.workflow_generator, "max_new_tokens", 8192) + self.temperature = getattr(cfg.workflow_generator, "temperature", 0.95) + self.allow_primitive_nodes = getattr(cfg.workflow_generator, "allow_primitive_nodes", False) + # Template matches FlowDataset training format: instruction + "description: [description]" + # Model (Qwen2.5-14B, not instruct) was fine-tuned on this exact format + self.template = "Based on the description I provided, generate a JSON example of the required ComfyUi workflow. description:" + + logging.info("WorkflowGenerator: Initialized successfully") + + def generate(self, prompt: str, seed: int | None = None) -> list: + """ + Generate workflow diagram from natural language prompt. + + Args: + prompt: Natural language description of the desired workflow + seed: Random seed for reproducible generation (optional) + + Returns: + List of edges representing the workflow diagram + """ + # Format prompt to match FlowDataset training format: Model (Qwen2.5-14B, not instruct) was fine-tuned on this exact format + connection_guidance = ( + " Ensure all connections are valid: every input_name must exist on the destination node, " + "and every output_name must exist on the source node." + ) + + if "Based on the description I provided" not in prompt: + formatted_prompt = self.template + " " + prompt + connection_guidance + else: + if "description:" in prompt: + desc_idx = prompt.find("description:") + if desc_idx != -1: + desc_prefix = prompt[: desc_idx + len("description:")] + desc_content = prompt[desc_idx + len("description:") :].strip() + formatted_prompt = desc_prefix + " " + desc_content + connection_guidance + else: + formatted_prompt = prompt + connection_guidance + else: + formatted_prompt = prompt + " description: " + connection_guidance + + prompt = formatted_prompt + + messages = [{"role": "user", "content": prompt}] + + try: + # Use HuggingFace tokenizer even for GGUF models (embedded tokenizer doesn't expose apply_chat_template) + text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + except Exception as e: + logging.error(f"Failed to apply chat template: {e}") + text = prompt + + logging.debug(f"Tokenizing input (max_new_tokens={self.max_new_tokens}, temperature={self.temperature}, top_p={self.top_p})...") + + try: + model_inputs = self.tokenizer([text], return_tensors="pt") + if hasattr(self.model_wrapper, "device") and isinstance(self.model_wrapper.device, torch.device): + model_inputs = model_inputs.to(self.model_wrapper.device) + input_tokens = model_inputs.input_ids.shape[1] + logging.debug(f"Input tokens: {input_tokens}") + except Exception as e: + logging.error(f"Failed to tokenize input: {e}") + raise + + logging.info("Generating workflow diagram...") + + try: + generate_kwargs = { + "max_new_tokens": self.max_new_tokens, + "temperature": self.temperature, + "do_sample": True, + "top_p": self.top_p, + } + if seed is not None: + max_seed = 2**32 - 1 + if seed < 0 or seed > max_seed: + seed = seed % (max_seed + 1) + generate_kwargs["seed"] = seed + + generated_ids = self.model_wrapper.generate(model_inputs.input_ids, **generate_kwargs) + + generated_ids = [output_ids[len(input_ids) :] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)] + output_tokens = len(generated_ids[0]) + logging.debug(f"Generated {output_tokens} tokens") + except Exception as e: + logging.error(f"Failed to generate tokens: {e}") + raise + + try: + response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] + except Exception as e: + logging.error(f"Failed to decode tokens: {e}") + raise + + logging.debug("Post-processing diagram...") + try: + diagram = self.post_process(response) + logging.info(f"Workflow diagram generated successfully: {len(diagram)} edges") + return diagram + except Exception as e: + logging.error(f"Failed to post-process diagram: {e}") + raise + + def post_process(self, response: str | list) -> list: + """ + Post-process the model response into a valid diagram format. + + Args: + response: Raw model response (string or list) + + Returns: + List of edges representing the workflow diagram + """ + diagram = response + + if isinstance(diagram, str): + try: + diagram = json.loads(diagram) + except json.JSONDecodeError as e: + logging.error(f"Failed to parse JSON response: {e}") + logging.error(f"Response content: {response[:500]}") + raise + + if not self.allow_primitive_nodes: + try: + diagram = del_digram_primitive(diagram) + except Exception as e: + logging.warning(f"Failed to remove primitive nodes: {e}") + + if not isinstance(diagram, list): + raise ValueError(f"Invalid diagram format: expected list, got {type(diagram)}") + + return diagram + + def cleanup(self): + """Clean up model resources.""" + if hasattr(self, "model_wrapper"): + del self.model_wrapper + if hasattr(self, "model"): + del self.model + if hasattr(self, "tokenizer"): + del self.tokenizer + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + logging.info("WorkflowGenerator: Resources cleaned up") diff --git a/nodes/node_validator_node.py b/nodes/node_validator_node.py new file mode 100644 index 0000000..885c538 --- /dev/null +++ b/nodes/node_validator_node.py @@ -0,0 +1,344 @@ +""" +NodeValidatorNode: Validates and corrects node names in workflow diagrams. +Step 2 of the ComfyUI-WorkflowGenerator pipeline (optional but recommended). +""" + +import json +import logging +import random +from typing import Any + +import numpy as np +import torch +from comfy_api.latest import io +from omegaconf import DictConfig + +from ..config.config import ConfigBuilder +from ..generators.node_validator import NodeBase, NodeValidator +from ..utils.model_manager import ( + detect_model_format, + get_embedding_models, + get_model_full_path, + get_refine_agent_models, +) +from ..utils.node_utils import check_llm_lazy_status, get_n_threads_options, parse_n_threads + +# Generate options list at module import time (safe fallback if detection fails) +_N_THREADS_OPTIONS = get_n_threads_options() + + +class NodeValidatorNode(io.ComfyNode): + """ + NodeValidator: Validates and corrects node names in workflow diagrams. + Uses semantic search and optionally an LLM to fix invalid node references. + """ + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="NodeValidator", + category="WorkflowGenerator", + display_name="2. NodeValidator (Optional)", + description=( + "Step 2 (Optional) - Validates and corrects node names.\n" + "Uses semantic search and optionally an LLM to fix invalid node references.\n" + "Recommended to ensure workflow generation succeeds if Step 1 produces errors." + ), + inputs=[ + io.String.Input( + "instruction", + multiline=True, + default="", + tooltip="Optional context for LLM refinement. Connect 'instruction' output from Step 1 or add your own instructions to guide refinement.", + ), + io.String.Input("workflow_edges", force_input=True, tooltip="Input diagram from Step 1 (WorkflowGenerator)."), + io.Boolean.Input( + "use_llm_refinement", + default=False, + tooltip="Enable LLM-based correction for better accuracy (slower). False uses only semantic search.", + ), + io.Combo.Input( + "refine_model_path", + options=get_refine_agent_models(), # Called dynamically in define_schema() - updates on frontend reload + default="Qwen2.5-7B-Instruct-q8_0.gguf", + lazy=True, + tooltip="LLM model for refinement (if enabled).", + ), + io.Combo.Input( + "embedding_model_path", + options=get_embedding_models(), # Called dynamically in define_schema() - updates on frontend reload + default="paraphrase-multilingual-MiniLM-L12-v2", + tooltip="Embedding model for semantic search.", + ), + io.String.Input("catalog_directory", default="catalog", tooltip="Directory containing node catalog files."), + io.Combo.Input( + "dtype", + options=["auto", "fp16", "bf16", "fp32", "fp8"], + default="auto", + lazy=True, + tooltip="Data type for refinement model (HuggingFace).", + ), + io.Combo.Input( + "device_preference", + options=["auto", "cuda", "cpu"], + default="auto", + lazy=True, + tooltip="Device preference for refinement model.", + ), + io.Combo.Input( + "attn_implementation", + options=["auto", "flash_attention_2", "sdpa"], + default="auto", + lazy=True, + tooltip="Attention implementation (HuggingFace).", + ), + io.Boolean.Input("auto_gpu_layers", default=True, lazy=True, tooltip="Auto-calculate GPU layers based on available VRAM (prevents OOM errors) (GGUF)."), + io.Int.Input("n_gpu_layers", default=-1, min=-1, max=1000, lazy=True, tooltip="Number of GPU layers (GGUF). -1 for all."), + io.Int.Input("max_new_tokens", default=4096, min=1, max=16384, lazy=True, tooltip="Max tokens for refinement."), + io.Int.Input("context_size", default=4096, min=512, max=32768, lazy=True, tooltip="Context window size (GGUF)."), + io.Boolean.Input("use_mmap", default=True, lazy=True, tooltip="Use memory-mapped loading (GGUF)."), + io.Boolean.Input("use_mlock", default=False, lazy=True, tooltip="Lock memory to prevent swapping (GGUF)."), + io.Int.Input("n_batch", default=512, min=32, max=2048, lazy=True, tooltip="Batch size for prompt processing (GGUF)."), + io.Combo.Input( + "n_threads", options=_N_THREADS_OPTIONS, default="Auto", lazy=True, tooltip="CPU threads for inference (GGUF)." + ), + io.Float.Input("temperature", default=0.95, min=0.0, max=2.0, step=0.01, lazy=True, tooltip="Sampling temperature."), + io.Int.Input("top_k", default=5, min=1, max=20, tooltip="Number of similar nodes to consider."), + io.Float.Input("top_p", default=0.7, min=0.0, max=1.0, step=0.01, lazy=True, tooltip="Top-p sampling."), + io.Int.Input( + "seed", + default=0, + min=0, + max=0xFFFFFFFF, # 2**32 - 1 (4294967295) - maximum valid seed value + control_after_generate=True, + lazy=True, + tooltip="Random seed.", + ), + ], + outputs=[ + io.String.Output(display_name="workflow_edges (refined)", tooltip="Refined workflow diagram with corrected names."), + io.String.Output(display_name="llm_prompts (debug)", tooltip="Prompts sent to LLM (debug)."), + io.String.Output(display_name="candidate_nodes (debug)", tooltip="Candidate nodes found (debug)."), + ], + ) + + @classmethod + def validate_inputs(cls, instruction: str = "", catalog_directory: str = "", embedding_model_path: str = "", **kwargs) -> bool | str: + """ + Validate constant inputs before execution. + + Note: Only constant inputs (widget values) are available here. + Connected inputs (like workflow_edges) are not available and should be validated in execute(). + + Returns: + True if valid, or error message string if invalid + """ + # workflow_edges is force_input=True (connection-only), so it's not available here + # instruction can be constant or connected - only validate if it's a constant + # Other validation happens in execute() when upstream nodes have run + return True + + @classmethod + def check_lazy_status( + cls, + use_llm_refinement: bool = False, + refine_model_path: str = None, + max_new_tokens: int = None, + device_preference: str = None, + n_gpu_layers: int | None = None, + dtype: str = None, + attn_implementation: str = None, + context_size: int | None = None, + temperature: float = None, + top_p: float = None, + seed: int = None, + **kwargs, + ) -> list[str]: + """ + Determine which lazy inputs need to be evaluated. + + Returns: + List of input names that need to be evaluated. + LLM-related inputs are only needed if use_llm_refinement is True. + """ + return check_llm_lazy_status( + use_llm_refinement=use_llm_refinement, + refine_model_path=refine_model_path, + max_new_tokens=max_new_tokens, + device_preference=device_preference, + n_gpu_layers=n_gpu_layers, + dtype=dtype, + attn_implementation=attn_implementation, + context_size=context_size, + temperature=temperature, + top_p=top_p, + seed=seed, + **kwargs, + ) + + @classmethod + def fingerprint_inputs(cls, workflow_edges: str = "", seed: int = 0, use_llm_refinement: bool = False, **kwargs) -> Any: + """ + Fingerprint inputs to determine if node should re-execute. + + Returns: + Tuple of relevant inputs (no hashing needed - ComfyUI compares with !=) + """ + # Include seed only if using LLM refinement (LLM is stochastic) + # Semantic search is deterministic, so seed doesn't matter when use_llm_refinement=False + if use_llm_refinement: + return (workflow_edges, seed) + else: + return (workflow_edges,) + + @classmethod + def execute( + cls, + workflow_edges: str, + instruction: str = "", + catalog_directory: str = "catalog", + embedding_model_path: str = "", + use_llm_refinement: bool = False, + refine_model_path: str | None = None, + top_k: int = 5, + max_new_tokens: int | None = None, + device_preference: str | None = None, + auto_gpu_layers: bool | None = True, + n_gpu_layers: int | None = None, + dtype: str | None = None, + attn_implementation: str | None = None, + context_size: int | None = None, + temperature: float | None = None, + top_p: float | None = None, + use_mmap: bool | None = None, + use_mlock: bool | None = None, + n_batch: int | None = None, + n_threads: str | None = None, + seed: int | None = None, + ) -> io.NodeOutput: + """Execute NodeValidator to validate and correct node names.""" + try: + logging.info("WorkflowGenerator NodeValidator: Starting diagram refinement") + + if not workflow_edges or not workflow_edges.strip(): + raise ValueError( + "workflow_edges input is required - connect the 'workflow_edges' output from Step 1 (WorkflowGenerator). " + "If Step 1 failed, check its error message (e.g., missing llama-cpp-python)." + ) + + if use_llm_refinement and seed is not None: + max_seed = 2**32 - 1 + if seed < 0 or seed > max_seed: + seed = seed % (max_seed + 1) + + random.seed(seed) + np.random.seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.manual_seed(seed) + + embedding_model_full_path = get_model_full_path(embedding_model_path) if embedding_model_path else None + + refine_model_full_path = None + refine_model_format = None + if use_llm_refinement: + if refine_model_path is None: + refine_model_path = "models/LLM/Qwen2.5-7B-Instruct-q8_0.gguf" + refine_model_full_path = get_model_full_path(refine_model_path) + refine_model_format = detect_model_format(refine_model_path) + + if max_new_tokens is None: + max_new_tokens = 4096 + if device_preference is None: + device_preference = "auto" + if auto_gpu_layers is None: + auto_gpu_layers = False + if n_gpu_layers is None: + n_gpu_layers = -1 + if dtype is None: + dtype = "auto" + if attn_implementation is None: + attn_implementation = "auto" + if temperature is None: + temperature = 0.95 + if top_p is None: + top_p = 0.7 + if seed is None: + seed = 0 + if context_size is None: + context_size = 4096 + if use_mmap is None: + use_mmap = True + if use_mlock is None: + use_mlock = False + if n_batch is None: + n_batch = 512 + if n_threads is None: + n_threads = "Auto" + n_threads_parsed = parse_n_threads(n_threads) + + if auto_gpu_layers: + n_gpu_layers = "auto" + + try: + diagram_obj = json.loads(workflow_edges) + except json.JSONDecodeError as e: + raise ValueError( + f"workflow_edges contains invalid JSON: {e}\n" + "The workflow_edges must be a valid JSON string (typically a list of edges). " + "If you're connecting from Step 1 (WorkflowGenerator), ensure Step 1 executed successfully. " + "If entering manually, ensure the JSON is properly formatted." + ) from e + + config_builder = ConfigBuilder() + node_base_cfg = config_builder.build_node_base_config( + catalog_directory=catalog_directory, embedding_model_path=embedding_model_full_path + ) + + refine_cfg = config_builder.build_node_validator_config( + local=use_llm_refinement, + model_format=refine_model_format if use_llm_refinement else "gguf", + model_path=refine_model_full_path if use_llm_refinement else "", + tokenizer_path=None, # Auto-detected in model loader + k=top_k, + max_new_tokens=max_new_tokens if use_llm_refinement else 4096, + ) + + cfg = DictConfig({"node_base": node_base_cfg.node_base, "node_validator": refine_cfg.node_validator}) + + node_base = NodeBase(cfg) + + node_validator = NodeValidator( + node_base, + cfg, + n_gpu_layers=n_gpu_layers if use_llm_refinement else -1, + dtype=dtype if use_llm_refinement else "auto", + device_preference=device_preference if use_llm_refinement else "auto", + attn_implementation=attn_implementation if use_llm_refinement else "auto", + context_size=context_size if use_llm_refinement else 4096, + temperature=temperature if use_llm_refinement else 0.95, + top_p=top_p if use_llm_refinement else 0.7, + use_mmap=use_mmap if use_llm_refinement else True, + use_mlock=use_mlock if use_llm_refinement else False, + n_batch=n_batch if use_llm_refinement else 512, + n_threads=n_threads_parsed if use_llm_refinement else None, + ) + + try: + refine_seed = seed if use_llm_refinement else None + refine_prompt = instruction if instruction else "" + refined_diagram, debug_info = node_validator.refine_diagram(diagram_obj, prompt=refine_prompt, seed=refine_seed) + + refined_diagram_json = json.dumps(refined_diagram, indent=2) + llm_prompts_json = json.dumps(debug_info.get("llm_prompts", []), indent=2) + candidate_nodes_json = json.dumps(debug_info.get("candidate_nodes", {}), indent=2) + + logging.info("WorkflowGenerator NodeValidator: Diagram refined successfully") + return io.NodeOutput(refined_diagram_json, llm_prompts_json, candidate_nodes_json) + + finally: + node_validator.cleanup() + + except Exception as e: + logging.error(f"WorkflowGenerator NodeValidator error: {e}", exc_info=True) + raise diff --git a/nodes/pipeline_node.py b/nodes/pipeline_node.py new file mode 100644 index 0000000..c3ea646 --- /dev/null +++ b/nodes/pipeline_node.py @@ -0,0 +1,461 @@ +""" +WGPipelineNode: Complete ComfyUI-WorkflowGenerator workflow generation pipeline. +Executes WorkflowGenerator → NodeValidator → WorkflowBuilder in sequence. +""" + +import hashlib +import json +import logging +import random +from typing import Any + +import numpy as np +import torch +from comfy_api.latest import io +from omegaconf import DictConfig + +from ..config.config import ConfigBuilder +from ..generators.node_validator import NodeBase, NodeValidator +from ..generators.workflow_builder import WorkflowBuilder +from ..generators.workflow_generator import WorkflowGenerator +from ..utils.file_utils import save_workflow_json +from ..utils.model_manager import ( + detect_model_format, + get_embedding_models, + get_model_full_path, + get_refine_agent_models, + get_workflow_generator_models, +) +from ..utils.node_utils import check_llm_lazy_status + +try: + import folder_paths + + COMFYUI_AVAILABLE = True +except ImportError: + COMFYUI_AVAILABLE = False + logging.warning("ComfyUI folder_paths not available") + +# Internal cache for built workflows in pipeline (keyed by refined_diagram + catalog_directory hash) +_pipeline_workflow_cache: dict[str, str] = {} + + +class WGPipelineNode(io.ComfyNode): + """ + Pipeline: Complete ComfyUI-WorkflowGenerator workflow generation pipeline. + Executes WorkflowGenerator → NodeValidator → WorkflowBuilder in sequence. + """ + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="WG_Pipeline", + category="WorkflowGenerator", + display_name="Workflow Generator Pipeline", + description=( + "Complete Pipeline - Executes steps 1-3 sequentially:\n" + "1. WorkflowGenerator -> Creates initial diagram\n" + "2. NodeValidator -> Validates/corrects node names\n" + "3. Build and Save -> Converts to JSON\n" + "Convenience node for the full process." + ), + inputs=[ + # WorkflowGenerator inputs (Step 1) - in same order as WorkflowGeneratorNode + io.String.Input("instruction", multiline=True, tooltip="Description of the desired ComfyUI workflow."), + io.Combo.Input( + "model_path", + options=get_workflow_generator_models(), # Called dynamically in define_schema() - updates on frontend reload + default="workflow-generator-q8_0.gguf", + tooltip="Model file (GGUF) or directory (HuggingFace).", + ), + io.Combo.Input( + "dtype", options=["auto", "fp16", "bf16", "fp32", "fp8"], default="auto", tooltip="Data type (HuggingFace only)." + ), + io.Combo.Input( + "device_preference", options=["auto", "cuda", "cpu"], default="auto", tooltip="Device preference (cuda/cpu)." + ), + io.Combo.Input( + "attn_implementation", + options=["auto", "flash_attention_2", "sdpa"], + default="auto", + tooltip="Attention implementation (HuggingFace).", + ), + io.Boolean.Input("auto_gpu_layers", default=True, tooltip="Auto-calculate GPU layers based on available VRAM (prevents OOM errors) (GGUF)."), + io.Int.Input("n_gpu_layers", default=-1, min=-1, max=1000, tooltip="Number of GPU layers (GGUF). -1 for all."), + io.Int.Input("max_new_tokens", default=8192, min=1, max=32768, tooltip="Max tokens to generate."), + io.Int.Input("context_size", default=4096, min=512, max=131072, tooltip="Context window size (GGUF)."), + io.Float.Input("temperature", default=0.95, min=0.0, max=2.0, step=0.01, tooltip="Sampling temperature."), + io.Float.Input("top_p", default=0.7, min=0.0, max=1.0, step=0.01, tooltip="Top-p sampling."), + io.Boolean.Input( + "allow_primitive_nodes", + default=False, + tooltip="Include primitive nodes in diagram. False removes them for cleaner workflows.", + ), + # NodeValidator inputs (Step 2) - in same order as NodeValidatorNode + io.Boolean.Input("refine_enabled", default=True, tooltip="Enable Step 2 (NodeValidator)."), + io.Boolean.Input( + "use_llm_refinement", default=False, tooltip="Enable LLM-based node correction (slower but more accurate)." + ), + io.Combo.Input( + "refine_model_path", + options=get_refine_agent_models(), # Called dynamically in define_schema() - updates on frontend reload + default="Qwen2.5-7B-Instruct-q8_0.gguf", + lazy=True, + tooltip="LLM model for refinement.", + ), + io.Combo.Input( + "embedding_model_path", + options=get_embedding_models(), # Called dynamically in define_schema() - updates on frontend reload + default="paraphrase-multilingual-MiniLM-L12-v2", + tooltip="Embedding model for semantic search.", + ), + io.String.Input("catalog_directory", default="catalog", tooltip="Directory containing node catalog files."), + io.Combo.Input( + "refine_dtype", + options=["auto", "fp16", "bf16", "fp32", "fp8"], + default="auto", + lazy=True, + tooltip="Data type for refinement model (HuggingFace).", + ), + io.Combo.Input( + "refine_device_preference", + options=["auto", "cuda", "cpu"], + default="auto", + lazy=True, + tooltip="Device preference for refinement model.", + ), + io.Combo.Input( + "refine_attn_implementation", + options=["auto", "flash_attention_2", "sdpa"], + default="auto", + lazy=True, + tooltip="Attention implementation for refinement model.", + ), + io.Boolean.Input( + "refine_auto_gpu_layers", default=False, lazy=True, tooltip="Auto-calculate GPU layers for refinement model (GGUF)." + ), + io.Int.Input( + "refine_n_gpu_layers", + default=-1, + min=-1, + max=1000, + lazy=True, + tooltip="GPU layers for refinement model (GGUF). -1 for all.", + ), + io.Int.Input("refine_max_new_tokens", default=4096, min=1, max=16384, lazy=True, tooltip="Max tokens for refinement."), + io.Int.Input( + "refine_context_size", default=4096, min=512, max=32768, lazy=True, tooltip="Context window size for refinement model (GGUF)." + ), + io.Float.Input( + "refine_temperature", default=0.95, min=0.0, max=2.0, step=0.01, lazy=True, tooltip="Temperature for refinement." + ), + io.Int.Input("top_k", default=5, min=1, max=20, tooltip="Number of similar nodes to consider."), + io.Float.Input("refine_top_p", default=0.7, min=0.0, max=1.0, step=0.01, lazy=True, tooltip="Top-p for refinement."), + # WorkflowBuilder inputs (Step 3) - in same order as WorkflowBuilderNode + io.Boolean.Input("save_workflow", default=True, tooltip="Save workflow to file."), + io.String.Input("filename_prefix", default="generated_workflow", tooltip="Filename prefix or path for saved workflow."), + # Seed for reproducibility (shared by WorkflowGenerator and NodeValidator) + io.Int.Input( + "seed", + default=0, + min=0, + max=0xFFFFFFFF, # 2**32 - 1 (4294967295) - maximum valid seed value + control_after_generate=True, + tooltip="Random seed.", + ), + ], + outputs=[ + io.String.Output(display_name="workflow_edges", tooltip="Initial workflow diagram (JSON)."), + io.String.Output(display_name="workflow_edges (refined)", tooltip="Refined workflow diagram (JSON)."), + io.String.Output(display_name="workflow_json", tooltip="Final ComfyUI workflow JSON."), + io.String.Output(display_name="file_path", tooltip="Path to saved file."), + ], + is_output_node=True, # Mark as output node since WorkflowBuilder saves files + ) + + @classmethod + def validate_inputs(cls, instruction: str = "", **kwargs) -> bool | str: + """ + Validate constant inputs before execution. + + Note: Only constant inputs (widget values) are available here. + Connected inputs are not available and should be validated in execute(). + + Returns: + True if valid, or error message string if invalid + + Note: Validation moved to execute() to avoid ComfyUI's per-input error formatting. + This prevents the "Custom validation failed for node: X -" prefix and duplicate errors + that appear on every input field. The execute() method will raise a ValueError with + a clear message if instruction is missing. + """ + return True + + @classmethod + def check_lazy_status( + cls, + refine_enabled: bool = True, + use_llm_refinement: bool = False, + refine_model_path: str = None, + refine_context_size: int | None = None, + refine_max_new_tokens: int | None = None, + refine_temperature: float | None = None, + refine_top_p: float | None = None, + refine_auto_gpu_layers: bool | None = None, + refine_n_gpu_layers: int | None = None, + refine_dtype: str | None = None, + refine_device_preference: str | None = None, + refine_attn_implementation: str | None = None, + **kwargs, + ) -> list[str]: + """ + Determine which lazy inputs need to be evaluated. + + Returns: + List of input names that need to be evaluated. + Refine model inputs are only needed if refine_enabled=True AND use_llm_refinement=True. + """ + if not (refine_enabled and use_llm_refinement): + return [] + + return check_llm_lazy_status( + use_llm_refinement=True, + refine_model_path=refine_model_path, + refine_context_size=refine_context_size, + refine_max_new_tokens=refine_max_new_tokens, + refine_temperature=refine_temperature, + refine_top_p=refine_top_p, + refine_auto_gpu_layers=refine_auto_gpu_layers, + refine_n_gpu_layers=refine_n_gpu_layers, + refine_dtype=refine_dtype, + refine_device_preference=refine_device_preference, + refine_attn_implementation=refine_attn_implementation, + **kwargs, + ) + + @classmethod + def fingerprint_inputs(cls, instruction: str = "", seed: int = 0, **kwargs) -> Any: + """ + Fingerprint inputs to determine if node should re-execute. + + Returns: + Tuple of relevant inputs (no hashing needed - ComfyUI compares with !=) + """ + # Pipeline contains WorkflowGenerator (stochastic) and potentially NodeValidator with LLM (stochastic) + # Same instruction + same seed = same output + return (instruction, seed) + + @classmethod + def execute( + cls, + instruction: str, + model_path: str = "workflow-generator-q8_0.gguf", + max_new_tokens: int = 8192, + top_p: float = 0.7, + temperature: float = 0.95, + allow_primitive_nodes: bool = False, + refine_enabled: bool = True, + use_llm_refinement: bool = False, + refine_model_path: str | None = None, + top_k: int = 5, + refine_max_new_tokens: int | None = None, + refine_temperature: float | None = None, + refine_top_p: float | None = None, + refine_auto_gpu_layers: bool | None = None, + refine_n_gpu_layers: int | None = None, + refine_dtype: str | None = None, + refine_device_preference: str | None = None, + refine_attn_implementation: str | None = None, + refine_context_size: int | None = None, + catalog_directory: str = "catalog", + embedding_model_path: str = "paraphrase-multilingual-MiniLM-L12-v2", + device_preference: str = "auto", + auto_gpu_layers: bool = True, + n_gpu_layers: int = -1, + dtype: str = "auto", + attn_implementation: str = "auto", + context_size: int = 4096, + save_workflow: bool = True, + filename_prefix: str = "generated_workflow", + seed: int = 0, + ) -> io.NodeOutput: + """Execute complete ComfyUI-WorkflowGenerator pipeline.""" + try: + logging.info("WorkflowGenerator Pipeline: Starting complete workflow generation") + + if not instruction or not instruction.strip(): + raise ValueError("instruction input is required - enter a description of the workflow you want to generate") + + max_seed = 2**32 - 1 + if seed < 0 or seed > max_seed: + seed = seed % (max_seed + 1) + + random.seed(seed) + np.random.seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.manual_seed(seed) + + model_full_path = get_model_full_path(model_path) + model_format = detect_model_format(model_path) + + embedding_model_full_path = get_model_full_path(embedding_model_path) + + workflow_generator_n_gpu_layers = n_gpu_layers + if auto_gpu_layers: + workflow_generator_n_gpu_layers = "auto" + + if refine_enabled and use_llm_refinement: + if refine_model_path is None: + refine_model_path = "Qwen2.5-7B-Instruct-q8_0.gguf" + if refine_context_size is None: + refine_context_size = 4096 + if refine_max_new_tokens is None: + refine_max_new_tokens = 4096 + if refine_temperature is None: + refine_temperature = 0.95 + if refine_top_p is None: + refine_top_p = 0.7 + if refine_auto_gpu_layers is None: + refine_auto_gpu_layers = False + if refine_n_gpu_layers is None: + refine_n_gpu_layers = -1 + if refine_dtype is None: + refine_dtype = "auto" + if refine_device_preference is None: + refine_device_preference = "auto" + if refine_attn_implementation is None: + refine_attn_implementation = "auto" + + refine_model_full_path = get_model_full_path(refine_model_path) + refine_model_format = detect_model_format(refine_model_path) + + refine_n_gpu_layers_value = refine_n_gpu_layers + if refine_auto_gpu_layers: + refine_n_gpu_layers_value = "auto" + else: + refine_model_full_path = None + refine_model_format = "gguf" + refine_n_gpu_layers_value = -1 + refine_max_new_tokens = 4096 + refine_temperature = 0.95 + refine_top_p = 0.7 + refine_dtype = "auto" + refine_device_preference = "auto" + refine_attn_implementation = "auto" + + config_builder = ConfigBuilder() + + workflow_generator_cfg = config_builder.build_workflow_generator_config( + model_format=model_format, + model_path=model_full_path, + tokenizer_path=None, # Auto-detected in model loader + max_new_tokens=max_new_tokens, + top_p=top_p, + temperature=temperature, + allow_primitive_nodes=allow_primitive_nodes, + ) + + node_base_cfg = config_builder.build_node_base_config( + catalog_directory=catalog_directory, embedding_model_path=embedding_model_full_path + ) + + node_validator_cfg = config_builder.build_node_validator_config( + local=use_llm_refinement if refine_enabled else False, + model_format=refine_model_format if (refine_enabled and use_llm_refinement) else "gguf", + model_path=refine_model_full_path if (refine_enabled and use_llm_refinement) else "", + tokenizer_path=None, # Auto-detected in model loader + k=top_k, + max_new_tokens=refine_max_new_tokens, + ) + + cfg = DictConfig( + { + "workflow_generator": workflow_generator_cfg.workflow_generator, + "node_validator": node_validator_cfg.node_validator, + "node_base": node_base_cfg.node_base, + } + ) + + logging.info("Pipeline Stage 1: WorkflowGenerator - Generating diagram") + workflow_generator = WorkflowGenerator( + cfg, + n_gpu_layers=workflow_generator_n_gpu_layers, + dtype=dtype, + device_preference=device_preference, + attn_implementation=attn_implementation, + context_size=context_size, + ) + + try: + diagram = workflow_generator.generate(instruction, seed=seed) + diagram_json = json.dumps(diagram, indent=2) + finally: + workflow_generator.cleanup() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + refined_diagram = diagram + refined_diagram_json = "" + + if refine_enabled: + logging.info("Pipeline Stage 2: NodeValidator - Refining diagram") + node_base = NodeBase(cfg) + node_validator = NodeValidator( + node_base, + cfg, + n_gpu_layers=refine_n_gpu_layers_value, + dtype=refine_dtype, + device_preference=refine_device_preference, + attn_implementation=refine_attn_implementation, + context_size=refine_context_size, + temperature=refine_temperature, + top_p=refine_top_p, + ) + + try: + refine_seed = seed if use_llm_refinement else None + refine_prompt = instruction if use_llm_refinement else "" + refined_diagram, _ = node_validator.refine_diagram(diagram, prompt=refine_prompt, seed=refine_seed) + refined_diagram_json = json.dumps(refined_diagram, indent=2) + finally: + node_validator.cleanup() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + else: + refined_diagram_json = diagram_json + + refined_diagram_json_str = json.dumps(refined_diagram, sort_keys=True) + cache_key_data = f"{refined_diagram_json_str}|{catalog_directory}" + cache_key = hashlib.sha256(cache_key_data.encode("utf-8")).hexdigest() + + workflow_json = _pipeline_workflow_cache.get(cache_key) + + if workflow_json is not None: + logging.info("Pipeline Stage 3: WorkflowBuilder - Using cached workflow JSON") + else: + logging.info("Pipeline Stage 3: WorkflowBuilder - Converting to workflow") + node_base = NodeBase(cfg) + workflow_builder = WorkflowBuilder(node_base) + + workflow = workflow_builder.parse_diagram_to_workflow(refined_diagram) + workflow_json = json.dumps(workflow, indent=2) + + _pipeline_workflow_cache[cache_key] = workflow_json + logging.info("Pipeline Stage 3: WorkflowBuilder - Workflow conversion complete (cached)") + + file_path = "" + if save_workflow and COMFYUI_AVAILABLE: + try: + output_dir = folder_paths.get_output_directory() + file_path = save_workflow_json(workflow_json, filename_prefix, output_dir) + logging.info(f"Pipeline: Workflow saved to {file_path}") + except Exception as e: + logging.warning(f"Failed to save workflow: {e}") + + logging.info("WorkflowGenerator Pipeline: Complete workflow generation finished") + return io.NodeOutput(diagram_json, refined_diagram_json, workflow_json, file_path) + + except Exception as e: + logging.error(f"WorkflowGenerator Pipeline error: {e}", exc_info=True) + raise diff --git a/nodes/update_catalog_node.py b/nodes/update_catalog_node.py new file mode 100644 index 0000000..0fbdae9 --- /dev/null +++ b/nodes/update_catalog_node.py @@ -0,0 +1,135 @@ +""" +UpdateNodeCatalogNode: Scans ComfyUI installation and updates node catalog files. +Generates node_info.json and node_list.json for use by ComfyUI-WorkflowGenerator. +""" + +import json +import logging +import os +from typing import Any + +from comfy_api.latest import io + +try: + from comfy.utils import ProgressBar + from comfy_execution.utils import get_executing_context + + PROGRESS_AVAILABLE = True +except ImportError: + PROGRESS_AVAILABLE = False + logging.warning("ComfyUI progress utilities not available") + +from ..cataloging.catalog import update_node_catalog +from ..config.config import get_custom_node_path, resolve_config_path + + +class UpdateNodeCatalogNode(io.ComfyNode): + """ + UpdateNodeCatalog: Scans ComfyUI installation and updates node catalog files. + Generates node_info.json and node_list.json for use by ComfyUI-WorkflowGenerator agents. + """ + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="WG_UpdateNodeCatalog", + category="WorkflowGenerator", + display_name="UpdateNodeCatalog", + description=( + "Scans ComfyUI for nodes and updates catalog files (node_list.json, node_info.json).\n" + "Run this after installing new custom nodes to make them available to the workflow validator." + ), + inputs=[ + io.String.Input("catalog_directory", default="catalog", tooltip="Directory to save catalog files."), + io.Boolean.Input("force_update", default=False, tooltip="Force update even if files exist."), + ], + outputs=[ + io.String.Output(display_name="status", tooltip="Update status (JSON)."), + io.Int.Output(display_name="node_count", tooltip="Number of nodes found."), + ], + is_output_node=True, # Mark as output node so it can run standalone + ) + + @classmethod + def validate_inputs(cls, catalog_directory: str = "", **kwargs) -> bool | str: + """ + Validate inputs before execution. + + Returns: + True if valid, or error message string if invalid + """ + # Paths are validated during execution, no pre-validation needed + return True + + @classmethod + def fingerprint_inputs(cls, force_update: bool = False, **kwargs) -> Any: + """ + Fingerprint inputs to determine if node should re-execute. + + Returns: + Hashable value that changes when node should re-execute. + For catalog updates, we want it to run when force_update changes. + """ + # Return force_update value - node will re-execute when it changes + # This allows manual triggering of catalog updates + return force_update + + @classmethod + def execute( + cls, + catalog_directory: str = "catalog", + force_update: bool = False, + ) -> io.NodeOutput: + """Execute node catalog update.""" + try: + logging.info("WorkflowGenerator UpdateNodeCatalog: Starting node catalog update") + + node_id = None + if PROGRESS_AVAILABLE: + try: + context = get_executing_context() + if context: + node_id = context.node_id + except Exception: + pass + + custom_node_path = get_custom_node_path() + catalog_dir = resolve_config_path(catalog_directory, custom_node_path) + + node_list_path = os.path.join(catalog_dir, "node_list.json") + node_info_db_path = os.path.join(catalog_dir, "node_info.json") + + progress_bar = None + if PROGRESS_AVAILABLE and node_id: + progress_bar = ProgressBar(100, node_id=node_id) + + def progress_callback(current, total): + if progress_bar: + scan_progress = int((current / total) * 80) if total > 0 else 0 + progress_bar.update_absolute(scan_progress, total=100) + logging.debug(f"Progress: {current}/{total} nodes scanned ({scan_progress}%)") + + result = update_node_catalog( + node_list_path=node_list_path, + node_info_db_path=node_info_db_path, + force_update=force_update, + progress_callback=progress_callback, + ) + + if progress_bar: + progress_bar.update_absolute(90, total=100) + + if progress_bar: + progress_bar.update_absolute(100, total=100) + + status_json = json.dumps(result, indent=2) + node_count = result.get("node_count", 0) + + logging.info(f"WorkflowGenerator UpdateNodeCatalog: Catalog update complete - {node_count} nodes") + return io.NodeOutput(status_json, node_count) + + except Exception as e: + logging.error(f"WorkflowGenerator UpdateNodeCatalog error: {e}", exc_info=True) + error_result = {"status": "error", "message": str(e), "node_count": 0} + status_json = json.dumps(error_result, indent=2) + return io.NodeOutput(status_json, 0) diff --git a/nodes/workflow_builder_node.py b/nodes/workflow_builder_node.py new file mode 100644 index 0000000..427acc5 --- /dev/null +++ b/nodes/workflow_builder_node.py @@ -0,0 +1,178 @@ +""" +WorkflowBuilderNode: Converts workflow diagrams to ComfyUI workflow JSON. +Step 3 of the ComfyUI-WorkflowGenerator pipeline. +""" + +import hashlib +import json +import logging +from typing import Any + +from comfy_api.latest import io +from omegaconf import DictConfig + +# Import utilities +from ..config.config import ConfigBuilder +from ..generators.node_validator import NodeBase + +# Import generator classes +from ..generators.workflow_builder import WorkflowBuilder +from ..utils.file_utils import save_workflow_json + +# Try to import ComfyUI utilities +try: + import folder_paths + + COMFYUI_AVAILABLE = True +except ImportError: + COMFYUI_AVAILABLE = False + logging.warning("ComfyUI folder_paths not available") + +# Internal cache for built workflows (keyed by workflow_edges + catalog_directory hash) +_workflow_cache: dict[str, str] = {} + + +class WorkflowBuilderNode(io.ComfyNode): + """ + WorkflowBuilder: Converts workflow diagrams to ComfyUI workflow JSON. + Can save the workflow to a file in ComfyUI's output directory. + """ + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="WorkflowBuilder", + category="WorkflowGenerator", + display_name="3. Build and Save Workflow", + description=("Step 3 - Converts workflow diagram to ComfyUI JSON.\nCan save the workflow to a file for loading."), + inputs=[ + io.String.Input("workflow_edges", force_input=True, tooltip="Input diagram from Step 1 or Step 2."), + io.String.Input("catalog_directory", default="catalog", tooltip="Directory containing node catalog files."), + io.Boolean.Input("save_workflow", default=True, tooltip="Save workflow to file."), + io.String.Input("filename_prefix", default="generated_workflow", tooltip="Filename prefix or path for saved workflow."), + ], + outputs=[ + io.String.Output(display_name="workflow_json", tooltip="Complete ComfyUI workflow JSON."), + io.String.Output(display_name="file_path", tooltip="Path to saved file."), + ], + is_output_node=True, # Mark as output node since it can save files + ) + + @classmethod + def validate_inputs(cls, catalog_directory: str = "", **kwargs) -> bool | str: + """ + Validate constant inputs before execution. + + Note: Only constant inputs (widget values) are available here. + Connected inputs (like workflow_edges) are not available and should be validated in execute(). + + Returns: + True if valid, or error message string if invalid + """ + # workflow_edges is force_input=True (connection-only), so it's not available here + # Validation of connected inputs happens in execute() when upstream nodes have run + return True + + @classmethod + def fingerprint_inputs( + cls, workflow_edges: str = "", catalog_directory: str = "catalog", save_workflow: bool = False, filename_prefix: str = "", **kwargs + ) -> Any: + """ + Fingerprint inputs to determine if node should re-execute. + + Caching strategy: + - Only includes inputs that affect the workflow build (workflow_edges, catalog_directory) + - save_workflow and filename_prefix are excluded because they only affect saving, not building + - Internal caching in execute() will reuse built workflow JSON when only save settings change + - When rerunning with the same seed, workflow_edges will be identical (from cached upstream nodes), + so the workflow build is cached and only the save step runs if needed + + Returns: + Tuple of relevant inputs (no hashing needed - ComfyUI compares with !=) + """ + # Only include inputs that affect the workflow build itself + # save_workflow and filename_prefix don't affect the build, only the saving step + # Internal cache in execute() handles reusing built workflow when only save settings change + return (workflow_edges, catalog_directory) + + @classmethod + def execute( + cls, + workflow_edges: str, + catalog_directory: str = "catalog", + save_workflow: bool = True, + filename_prefix: str = "generated_workflow", + ) -> io.NodeOutput: + """ + Execute WorkflowBuilder to convert diagram to workflow. + + Uses internal caching to avoid rebuilding workflows when only save settings change. + The workflow JSON is cached based on workflow_edges and catalog_directory. + """ + try: + # Create cache key from workflow_edges and catalog_directory + cache_key_data = f"{workflow_edges}|{catalog_directory}" + cache_key = hashlib.sha256(cache_key_data.encode("utf-8")).hexdigest() + + # Check if we have a cached workflow JSON + workflow_json = _workflow_cache.get(cache_key) + + if workflow_json is not None: + logging.info("WorkflowGenerator WorkflowBuilder: Using cached workflow JSON") + else: + logging.info("WorkflowGenerator WorkflowBuilder: Starting workflow conversion") + + # Validate workflow_edges is provided + if not workflow_edges or not workflow_edges.strip(): + raise ValueError( + "workflow_edges input is required - connect the 'workflow_edges' output from Step 1 (WorkflowGenerator) " + "or 'workflow_edges (refined)' output from Step 2 (NodeValidator). " + "If upstream nodes failed, check their error messages (e.g., missing llama-cpp-python)." + ) + + # Parse workflow_edges + try: + diagram_obj = json.loads(workflow_edges) + except json.JSONDecodeError as e: + raise ValueError(f"Invalid workflow_edges JSON: {e}") from e + + # Build configuration + # Note: catalog_directory is resolved by build_node_base_config using resolve_config_path + # embedding_model_path is not needed for WorkflowBuilder (only Step 2 (NodeValidator) uses semantic search) + config_builder = ConfigBuilder() + node_base_cfg = config_builder.build_node_base_config( + catalog_directory=catalog_directory, + embedding_model_path=None, # Not needed for WorkflowBuilder + ) + + # Merge configs + cfg = DictConfig({"node_base": node_base_cfg.node_base}) + + # Initialize NodeBase and WorkflowBuilder + node_base = NodeBase(cfg) + workflow_builder = WorkflowBuilder(node_base) + + # Convert diagram to workflow + workflow = workflow_builder.parse_diagram_to_workflow(diagram_obj) + + # Convert to JSON string and cache it + workflow_json = json.dumps(workflow, indent=2) + _workflow_cache[cache_key] = workflow_json + + logging.info("WorkflowGenerator WorkflowBuilder: Workflow conversion complete (cached)") + + # Save workflow if requested (this step always runs, even with cached workflow) + file_path = "" + if save_workflow and COMFYUI_AVAILABLE: + try: + output_dir = folder_paths.get_output_directory() + file_path = save_workflow_json(workflow_json, filename_prefix, output_dir) + logging.info(f"WorkflowGenerator WorkflowBuilder: Workflow saved to {file_path}") + except Exception as e: + logging.warning(f"Failed to save workflow: {e}") + + return io.NodeOutput(workflow_json, file_path) + + except Exception as e: + logging.error(f"WorkflowGenerator WorkflowBuilder error: {e}", exc_info=True) + raise diff --git a/nodes/workflow_generator_node.py b/nodes/workflow_generator_node.py new file mode 100644 index 0000000..5ad4ab3 --- /dev/null +++ b/nodes/workflow_generator_node.py @@ -0,0 +1,209 @@ +""" +WorkflowGeneratorNode: Generates workflow diagrams from natural language descriptions. +Step 1 of the ComfyUI-WorkflowGenerator pipeline. +""" + +import json +import logging +import random +from typing import Any + +import numpy as np +import torch +from comfy_api.latest import io + +from ..config.config import ConfigBuilder +from ..generators.workflow_generator import WorkflowGenerator +from ..utils.model_manager import ( + detect_model_format, + get_model_full_path, + get_workflow_generator_models, +) +from ..utils.node_utils import get_n_threads_options, parse_n_threads + +# Generate options list at module import time (safe fallback if detection fails) +_N_THREADS_OPTIONS = get_n_threads_options() + + +class WorkflowGeneratorNode(io.ComfyNode): + """ + WorkflowGenerator: Generates workflow diagrams from natural language descriptions. + Uses a language model to convert user instructions into structured workflow diagrams. + """ + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="WorkflowGenerator", + category="WorkflowGenerator", + display_name="1. WorkflowGenerator", + description=( + "Step 1 - Generates a workflow diagram from natural language instructions.\n" + "This is the first stage of the pipeline, converting text to a structured plan." + ), + inputs=[ + io.String.Input("instruction", multiline=True, tooltip="Description of the desired ComfyUI workflow."), + io.Combo.Input( + "model_path", + options=get_workflow_generator_models(), # Called dynamically in define_schema() - updates on frontend reload + default="workflow-generator-q8_0.gguf", + tooltip="Model file (GGUF) or directory (HuggingFace).", + ), + io.Combo.Input( + "dtype", options=["auto", "fp16", "bf16", "fp32", "fp8"], default="auto", tooltip="Data type (HuggingFace only)." + ), + io.Combo.Input( + "device_preference", options=["auto", "cuda", "cpu"], default="auto", tooltip="Device preference (cuda/cpu)." + ), + io.Combo.Input( + "attn_implementation", + options=["auto", "flash_attention_2", "sdpa"], + default="auto", + tooltip="Attention implementation (HuggingFace).", + ), + io.Boolean.Input("auto_gpu_layers", default=True, tooltip="Auto-calculate GPU layers based on available VRAM (prevents OOM errors) (GGUF)."), + io.Int.Input("n_gpu_layers", default=-1, min=-1, max=1000, tooltip="Number of GPU layers (GGUF). -1 for all."), + io.Int.Input("max_new_tokens", default=8192, min=1, max=32768, tooltip="Max tokens to generate."), + io.Int.Input("context_size", default=4096, min=512, max=131072, tooltip="Context window size (GGUF)."), + io.Boolean.Input("use_mmap", default=True, lazy=True, tooltip="Use memory-mapped loading (GGUF)."), + io.Boolean.Input("use_mlock", default=False, lazy=True, tooltip="Lock memory to prevent swapping (GGUF)."), + io.Int.Input("n_batch", default=512, min=32, max=2048, lazy=True, tooltip="Batch size for prompt processing (GGUF)."), + io.Combo.Input( + "n_threads", options=_N_THREADS_OPTIONS, default="Auto", lazy=True, tooltip="CPU threads for inference (GGUF)." + ), + io.Float.Input("temperature", default=0.95, min=0.0, max=2.0, step=0.01, tooltip="Sampling temperature."), + io.Float.Input("top_p", default=0.7, min=0.0, max=1.0, step=0.01, tooltip="Top-p sampling."), + io.Boolean.Input( + "allow_primitive_nodes", + default=False, + tooltip="Include primitive nodes in diagram. False (recommended) removes them for cleaner workflows.", + ), + io.Int.Input( + "seed", + default=0, + min=0, + max=0xFFFFFFFF, # 2**32 - 1 (4294967295) - maximum valid seed value + control_after_generate=True, + tooltip="Random seed.", + ), + ], + outputs=[ + io.String.Output(display_name="workflow_edges", tooltip="Generated workflow diagram (JSON)."), + io.String.Output(display_name="instruction", tooltip="Original instruction text (pass to Step 2)."), + ], + ) + + @classmethod + def validate_inputs(cls, instruction: str = "", **kwargs) -> bool | str: + """ + Validate constant inputs before execution. + + Note: Only constant inputs (widget values) are available here. + Connected inputs are not available and should be validated in execute(). + + Returns: + True if valid, or error message string if invalid + """ + # Validation moved to execute() to avoid ComfyUI's per-input error formatting + # This prevents the "Custom validation failed for node: X -" prefix and duplicate errors + return True + + @classmethod + def fingerprint_inputs(cls, instruction: str = "", seed: int = 0, **kwargs) -> Any: + """ + Fingerprint inputs to determine if node should re-execute. + + Returns: + Tuple of relevant inputs (no hashing needed - ComfyUI compares with !=) + """ + # Return simple tuple - ComfyUI handles comparison + return (instruction, seed) + + @classmethod + def execute( + cls, + instruction: str, + model_path: str = "workflow-generator-q8_0.gguf", + max_new_tokens: int = 8192, + top_p: float = 0.7, + temperature: float = 0.95, + allow_primitive_nodes: bool = False, + device_preference: str = "auto", + auto_gpu_layers: bool = True, + n_gpu_layers: int = -1, + dtype: str = "auto", + attn_implementation: str = "auto", + context_size: int = 4096, + use_mmap: bool = True, + use_mlock: bool = False, + n_batch: int = 512, + n_threads: str = "Auto", + seed: int = 0, + ) -> io.NodeOutput: + """Execute WorkflowGenerator to generate workflow diagram.""" + if not instruction or not instruction.strip(): + raise ValueError("instruction input is required - enter a description of the workflow you want to generate") + + try: + logging.info("WorkflowGenerator: Starting workflow diagram generation") + + model_full_path = get_model_full_path(model_path) + model_format = detect_model_format(model_path) + + # Handle auto_gpu_layers: if True, calculate optimal layers; otherwise use manual value + if auto_gpu_layers: + # Auto-calculate will be done in model loader + n_gpu_layers = "auto" # Special value for model loader to trigger auto-calculation + # Otherwise, n_gpu_layers is an integer: -1 (all), 0 (CPU only), or 1+ (specific count) + # If value exceeds model layers, model loader will handle it gracefully + + # Build configuration (tokenizer path will be auto-detected) + config_builder = ConfigBuilder() + cfg = config_builder.build_workflow_generator_config( + model_format=model_format, + model_path=model_full_path, + tokenizer_path=None, # Auto-detected in model loader + max_new_tokens=max_new_tokens, + top_p=top_p, + temperature=temperature, + allow_primitive_nodes=allow_primitive_nodes, + ) + + n_threads_parsed = parse_n_threads(n_threads) + + workflow_generator = WorkflowGenerator( + cfg, + n_gpu_layers=n_gpu_layers, + dtype=dtype, + device_preference=device_preference, + attn_implementation=attn_implementation, + context_size=context_size, + use_mmap=use_mmap, + use_mlock=use_mlock, + n_batch=n_batch, + n_threads=n_threads_parsed, + ) + + try: + max_seed = 2**32 - 1 + if seed < 0 or seed > max_seed: + seed = seed % (max_seed + 1) + + random.seed(seed) + np.random.seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + torch.manual_seed(seed) + + diagram = workflow_generator.generate(instruction, seed=seed) + diagram_json = json.dumps(diagram, indent=2) + + logging.info("WorkflowGenerator: Diagram generated successfully") + return io.NodeOutput(diagram_json, instruction) + + finally: + workflow_generator.cleanup() + + except Exception as e: + logging.error(f"WorkflowGenerator error: {e}", exc_info=True) + raise diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..86636f7 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,114 @@ +[build-system] +requires = ["setuptools>=70.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "ComfyUI-WorkflowGenerator" +version = "0.0.1" +description = "ComfyUI-WorkflowGenerator custom nodes for generating ComfyUI workflows from natural language" +readme = "README.md" +license = {text = "GNU General Public License v3"} +requires-python = ">=3.10" +classifiers = [] +dependencies = [ + "sentence-transformers>=2.2.1", + "tinydb>=4.8.2", + "scikit-learn>=1.3.0", + "omegaconf>=2.3.0", +] + +[project.optional-dependencies] +dev = [ + "bump-my-version", + "coverage", + "mypy", + "pre-commit", + "pytest", + "ruff", +] + +[project.urls] +Repository = "https://github.com/danielpflorian/ComfyUI-WorkflowGenerator" +BugTracker = "https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/issues" +Documentation = "https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/wiki" + +[tool.comfy] +PublisherId = "danielpflorian" +DisplayName = "ComfyUI-WorkflowGenerator" +Icon = "" +Tags = ["workflow", "llm", "ai", "generation"] +Repository = "https://github.com/danielpflorian/ComfyUI-WorkflowGenerator" + +includes = [] + +[tool.setuptools.package-data] +"*" = ["*.*"] + +[tool.pytest.ini_options] +minversion = "8.0" +testpaths = [ + "tests", +] + +[tool.mypy] +files = "." +strict = true +warn_unreachable = true +warn_no_return = true + +[[tool.mypy.overrides]] +module = "tests.*" +allow_untyped_defs = true +disable_error_code = "attr-defined" + +[tool.ruff] +line-length = 140 +src = ["nodes", "generators", "utils", "config", "cataloging", "tests"] +target-version = "py310" +exclude = [ + ".git", + "__pycache__", + ".venv", + "venv", + "build", + "dist", +] + +[tool.ruff.lint] +select = [ + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort + "N", # pep8-naming + "UP", # pyupgrade + "B", # flake8-bugbear + "C4", # flake8-comprehensions + "S", # flake8-bandit (security) + "T20", # flake8-print + "SIM", # flake8-simplify +] +ignore = [ + "E501", # line too long (handled by formatter) + "S101", # use of assert (needed for tests) + "T201", # print statements (useful for debugging) + "N801", # class name should use CapWords (allow existing exception names) + "N818", # exception name should be named with an Error suffix (allow existing exception names) + "S110", # try-except-pass detected (intentional silent exceptions) + "N999", # invalid module name (false positive for directory names with hyphens) + "N806", # variable in function should be lowercase (false positive for class name assignments) +] + +[tool.ruff.lint.isort] +known-first-party = ["nodes", "generators", "utils", "config", "cataloging"] + +[tool.ruff.lint.flake8-quotes] +inline-quotes = "double" +multiline-quotes = "double" + +[tool.ruff.format] +quote-style = "double" +indent-style = "space" +skip-magic-trailing-comma = false +line-ending = "auto" + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..22d9f7c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +sentence-transformers>=2.2.1 +tinydb>=4.8.2 +scikit-learn>=1.3.0 +omegaconf>=2.3.0 + diff --git a/tests/test_nodes.py b/tests/test_nodes.py new file mode 100644 index 0000000..f41bcba --- /dev/null +++ b/tests/test_nodes.py @@ -0,0 +1,52 @@ +"""Tests for ComfyUI-WorkflowGenerator custom nodes.""" + +from nodes.node_validator_node import NodeValidatorNode +from nodes.pipeline_node import WGPipelineNode +from nodes.update_catalog_node import UpdateNodeCatalogNode +from nodes.workflow_builder_node import WorkflowBuilderNode +from nodes.workflow_generator_node import WorkflowGeneratorNode + + +def test_workflow_generator_schema(): + """Test that WorkflowGenerator node has correct schema.""" + schema = WorkflowGeneratorNode.define_schema() + assert schema.node_id == "WorkflowGenerator" + assert schema.category == "WorkflowGenerator" + assert "instruction" in [inp.id for inp in schema.inputs] + assert "workflow_edges" in [out.display_name for out in schema.outputs] + + +def test_node_validator_schema(): + """Test that NodeValidator node has correct schema.""" + schema = NodeValidatorNode.define_schema() + assert schema.node_id == "NodeValidator" + assert schema.category == "WorkflowGenerator" + assert "workflow_edges" in [inp.id for inp in schema.inputs] + assert "workflow_edges (refined)" in [out.display_name for out in schema.outputs] + + +def test_workflow_builder_schema(): + """Test that WorkflowBuilder node has correct schema.""" + schema = WorkflowBuilderNode.define_schema() + assert schema.node_id == "WorkflowBuilder" + assert schema.category == "WorkflowGenerator" + assert "workflow_edges" in [inp.id for inp in schema.inputs] + assert "workflow_json" in [out.display_name for out in schema.outputs] + + +def test_pipeline_schema(): + """Test that Pipeline node has correct schema.""" + schema = WGPipelineNode.define_schema() + assert schema.node_id == "WG_Pipeline" + assert schema.category == "WorkflowGenerator" + assert "instruction" in [inp.id for inp in schema.inputs] + assert len(schema.outputs) == 4 # workflow_edges, workflow_edges (refined), workflow_json, file_path + + +def test_update_catalog_schema(): + """Test that UpdateNodeCatalog node has correct schema.""" + schema = UpdateNodeCatalogNode.define_schema() + assert schema.node_id == "WG_UpdateNodeCatalog" + assert schema.category == "WorkflowGenerator" + assert "catalog_directory" in [inp.id for inp in schema.inputs] + assert "node_count" in [out.display_name for out in schema.outputs] diff --git a/utils/file_utils.py b/utils/file_utils.py new file mode 100644 index 0000000..95c84dd --- /dev/null +++ b/utils/file_utils.py @@ -0,0 +1,88 @@ +import logging +import os +import re + + +def save_workflow_json(workflow_json: str, filename_prefix: str, output_dir: str) -> str: + """ + Save workflow JSON to a file with smart incrementing counter. + + Args: + workflow_json: The workflow JSON string to save. + filename_prefix: The filename prefix or path. + output_dir: The base output directory. + + Returns: + The full path to the saved file. + """ + try: + original_input = filename_prefix.strip() + ends_with_slash = original_input.endswith("/") or original_input.endswith("\\") + + normalized_input = original_input.replace("\\", "/").lstrip("/") + normalized_input = normalized_input.rstrip("/") + + if "/" in normalized_input: + parts = normalized_input.split("/") + parts = [p for p in parts if p] + if len(parts) > 1: + subdirectory = "/".join(parts[:-1]) + base_filename = parts[-1] + else: + subdirectory = parts[0] if parts else "" + base_filename = "generated_workflow" + else: + if ends_with_slash or not normalized_input: + subdirectory = normalized_input if normalized_input else "" + base_filename = "generated_workflow" + else: + subdirectory = "" + base_filename = normalized_input + + if base_filename.endswith(".workflow.json"): + base_filename = base_filename[:-15] + elif base_filename.endswith(".json"): + base_filename = base_filename[:-5] + + if not base_filename: + base_filename = "generated_workflow" + + full_output_folder = os.path.join(output_dir, subdirectory) if subdirectory else output_dir + + if not os.path.exists(full_output_folder): + try: + os.makedirs(full_output_folder, exist_ok=True) + except Exception as e: + logging.warning(f"Failed to create directory {full_output_folder}: {e}") + + counter = 1 + if os.path.exists(full_output_folder): + file_pattern = re.compile(rf"^{re.escape(base_filename)}_(\d+)\.json$") + + max_counter = 0 + try: + for file in os.listdir(full_output_folder): + match = file_pattern.match(file) + if match: + try: + cnt = int(match.group(1)) + if cnt > max_counter: + max_counter = cnt + except ValueError: + pass + except Exception as e: + logging.warning(f"Failed to scan output directory for counter: {e}") + + if max_counter > 0: + counter = max_counter + 1 + + file_path = os.path.join(full_output_folder, f"{base_filename}_{counter:03}.json") + + with open(file_path, "w") as f: + f.write(workflow_json) + + return file_path + + except Exception as e: + logging.warning(f"Failed to save workflow: {e}") + raise e diff --git a/utils/model_loaders/device_utils.py b/utils/model_loaders/device_utils.py new file mode 100644 index 0000000..07edf27 --- /dev/null +++ b/utils/model_loaders/device_utils.py @@ -0,0 +1,256 @@ +""" +Device, dtype, and VRAM utility functions for model loading. +Shared utilities for both GGUF and HuggingFace model loaders. +""" + +import logging + +import torch + +# Try to import ComfyUI model management for device detection +try: + from comfy import model_management + + COMFYUI_AVAILABLE = True +except ImportError: + COMFYUI_AVAILABLE = False + logging.warning("ComfyUI model_management not available, using basic device detection") + + +def calculate_optimal_gpu_layers( + model_size_gb: float, available_vram_gb: float, reserve_gb: float = 2.0, context_size: int = 4096 +) -> int: + """ + Calculate optimal number of GPU layers based on available VRAM. + + Args: + model_size_gb: Total model size in GB + available_vram_gb: Available VRAM in GB + reserve_gb: VRAM to reserve for other operations (default: 2.0 GB) + context_size: Context window size (used to estimate KV cache size) + + Returns: + Number of layers to load on GPU (-1 for all, 0 for CPU only, or specific count) + """ + if available_vram_gb <= 0: + return 0 # CPU only + + # 1. Reserve VRAM for system + KV cache + # Rough KV cache estimate: ~1-2GB for 4096 context, scales linearly with context + # Formula: 2 * 2 * layers * hidden_size * context_len (conservative estimate) + # Using a simpler heuristic: 0.5GB per 4096 tokens of context + kv_cache_gb = (context_size / 4096.0) * 0.5 + effective_reserve = reserve_gb + kv_cache_gb + + usable_vram = available_vram_gb - effective_reserve + if usable_vram <= 0: + return 0 + + if model_size_gb <= usable_vram: + return -1 # Fits fully + + # 2. Estimate per-layer size dynamically + # Assume ~40-80 layers for 7B-14B models. + # If we don't know layer count, estimating based on file size is better than fixed 100MB + # Qwen2.5-14B has 48 layers. Qwen2.5-7B has 28 layers. + # Average layers approx 40 for safe estimation. + estimated_layers = 40 + layer_size_gb = model_size_gb / estimated_layers + + # 3. Calculate layers that fit + max_layers = int(usable_vram / layer_size_gb) + return max(0, max_layers) + + +def get_available_vram() -> float: + """ + Get available VRAM in GB using ComfyUI's memory management if available. + + Returns: + Available VRAM in GB + """ + if COMFYUI_AVAILABLE: + try: + dev = model_management.get_torch_device() + free_memory = model_management.get_free_memory(dev) + return free_memory / (1024**3) + except Exception as e: + logging.warning(f"Failed to get VRAM from ComfyUI: {e}") + + if torch.cuda.is_available(): + try: + free_memory, _ = torch.cuda.mem_get_info(0) + return free_memory / (1024**3) + except Exception as e: + logging.warning(f"Failed to get VRAM from torch: {e}") + + return 0.0 + + +def get_device_for_model(device_preference: str = "auto") -> str: + """ + Get device string for model loading. + + Args: + device_preference: "auto", "cuda", or "cpu" + + Returns: + Device string ("cuda" or "cpu") + """ + if device_preference == "cpu": + return "cpu" + + if device_preference == "cuda" or device_preference == "auto": + if COMFYUI_AVAILABLE: + try: + dev = model_management.get_torch_device() + if hasattr(dev, "type"): + if dev.type == "cuda": + return "cuda" + elif dev.type == "cpu": + return "cpu" + except Exception: + pass + + # Fallback to torch detection + if torch.cuda.is_available(): + return "cuda" + + return "cpu" + + +def get_dtype_for_model(dtype_preference: str = "auto", device: str = "cuda") -> torch.dtype: + """ + Get dtype for HuggingFace model loading. + + Args: + dtype_preference: "auto", "fp16", "bf16", "fp32", or "fp8" + device: Device string ("cuda" or "cpu") + + Returns: + torch.dtype + """ + if dtype_preference == "fp32": + return torch.float32 + + if dtype_preference == "fp16": + return torch.float16 + + if dtype_preference == "bf16": + if device == "cuda" and torch.cuda.is_bf16_supported(): + return torch.bfloat16 + return torch.float16 + + if dtype_preference == "fp8": + # FP8 is CUDA-only and requires specific hardware support + if device == "cpu": + logging.debug("FP8 is not supported on CPU, falling back to fp16") + return torch.float16 + + # FP8 requires CUDA and specific PyTorch builds with FP8 support + if not torch.cuda.is_available(): + logging.debug("FP8 requires CUDA, but CUDA is not available, falling back to fp16") + return torch.float16 + + if COMFYUI_AVAILABLE: + try: + float8_types = model_management.get_supported_float8_types() + if float8_types: + selected_type = float8_types[0] + # Verify FP8 storage is actually available by testing set_default_dtype + # This is what transformers will do internally, so we catch the error early + try: + original_dtype = torch.get_default_dtype() + torch.set_default_dtype(selected_type) + torch.set_default_dtype(original_dtype) + type_name = str(selected_type).replace("torch.", "") + logging.debug(f"Using FP8: {type_name}") + return selected_type + except (TypeError, RuntimeError) as e: + # FP8 dtype exists but storage is not available + logging.debug(f"FP8 storage not available: {e}, falling back to fp16") + return torch.float16 + except Exception as e: + logging.debug(f"Failed to get FP8 types: {e}, falling back to fp16") + logging.debug("FP8 not available, falling back to fp16") + return torch.float16 + + if COMFYUI_AVAILABLE and device == "cuda": + try: + model_management.get_torch_device() + if torch.cuda.is_bf16_supported(): + return torch.bfloat16 + return torch.float16 + except Exception: + pass + + # Default: fp16 for CUDA, fp32 for CPU + if device == "cuda": + return torch.float16 + return torch.float32 + + +def load_model( + model_format: str, + model_path: str, + tokenizer_path: str | None = None, + n_gpu_layers: int | str = -1, + dtype: str | torch.dtype = "auto", + device_preference: str = "auto", + max_memory: dict | None = None, + attn_implementation: str = "auto", + context_size: int = 4096, + **kwargs, +): + """ + Unified model loader that supports both GGUF and HuggingFace formats. + + Args: + model_format: "gguf" or "huggingface" + model_path: Path to model file (GGUF) or directory (HuggingFace) + tokenizer_path: Path to tokenizer directory (auto-detected if None) + n_gpu_layers: For GGUF: "auto" (calculate optimal), -1 for all, 0 for CPU only, or 1+ for specific count. + If value exceeds model's total layers, all layers will be used (no error). + dtype: For HuggingFace: "auto", "fp16", "bf16", "fp32", "fp8", or torch.dtype + device_preference: "auto", "cuda", or "cpu" + max_memory: For HuggingFace: Dict for device_map memory limits + attn_implementation: For HuggingFace: "auto", "flash_attention_2", or "sdpa" + context_size: For GGUF: Context window size + **kwargs: Additional arguments passed to model wrapper initialization + + Returns: + Model wrapper instance (GGUFModelWrapper or HuggingFaceModelWrapper) + """ + import os + + from ..model_manager import auto_detect_tokenizer_path + + if model_format.lower() == "gguf": + from .gguf_loader import GGUFModelWrapper + + if tokenizer_path is None: + # Ensure model_path is absolute and normalized before auto-detection + if not os.path.isabs(model_path): + model_path = os.path.abspath(model_path) + model_path = os.path.normpath(model_path) + tokenizer_path = auto_detect_tokenizer_path(model_path, model_format="gguf") + logging.info(f"Auto-detected tokenizer path for GGUF model: {tokenizer_path}") + return GGUFModelWrapper( + model_path, tokenizer_path, n_gpu_layers=n_gpu_layers, device_preference=device_preference, context_size=context_size, **kwargs + ) + + elif model_format.lower() == "huggingface": + from .huggingface_loader import HuggingFaceModelWrapper + + return HuggingFaceModelWrapper( + model_path, + tokenizer_path, + dtype=dtype, + device_preference=device_preference, + max_memory=max_memory, + attn_implementation=attn_implementation, + **kwargs, + ) + + else: + raise ValueError(f"Unsupported model_format: {model_format}. Supported formats: 'gguf', 'huggingface'") diff --git a/utils/model_loaders/gguf_loader.py b/utils/model_loaders/gguf_loader.py new file mode 100644 index 0000000..f7dd152 --- /dev/null +++ b/utils/model_loaders/gguf_loader.py @@ -0,0 +1,249 @@ +""" +GGUF model loader using llama-cpp-python with GPU layer management. +""" + +import logging +import os +import time + +import numpy as np +import torch +from transformers import AutoTokenizer + +from ..model_manager import auto_detect_tokenizer_path +from .device_utils import ( + calculate_optimal_gpu_layers, + get_available_vram, + get_device_for_model, +) + + +class GGUFModelWrapper: + """Wrapper for GGUF models using llama-cpp-python with GPU layer management.""" + + def __init__( + self, + model_path: str, + tokenizer_path: str | None = None, + n_gpu_layers: int | str = -1, + device_preference: str = "auto", + **kwargs, + ): + """ + Initialize GGUF model wrapper. + + Args: + model_path: Path to GGUF model file + tokenizer_path: Path to HuggingFace tokenizer directory (auto-detected if None) + n_gpu_layers: Number of GPU layers (-1 for all, 0 for CPU only, or 1+ for specific count). + If value exceeds model's total layers, all layers will be used (no error). + device_preference: "auto", "cuda", or "cpu" + **kwargs: Additional arguments for llama-cpp-python Llama initialization (including context_size) + """ + if tokenizer_path is None: + if not os.path.isabs(model_path): + model_path = os.path.abspath(model_path) + model_path = os.path.normpath(model_path) + tokenizer_path = auto_detect_tokenizer_path(model_path, model_format="gguf") + logging.debug(f"Auto-detected tokenizer path: {tokenizer_path}") + try: + from llama_cpp import Llama + except ImportError: + raise ImportError( + "llama-cpp-python is not installed. " + "This package must be installed manually:\n" + " • CPU only: pip install llama-cpp-python\n" + " • CUDA support: pip install llama-cpp-python[cuda]\n" + "Note: CUDA support may require compilation (CMake, Visual Studio on Windows, or build tools on Linux/Mac). " + "See README.md for detailed installation instructions." + ) from None + + logging.info(f"Loading GGUF model: {os.path.basename(model_path)}") + load_start = time.time() + + device = get_device_for_model(device_preference) + cuda_available = device == "cuda" and torch.cuda.is_available() + + logging.debug(f"Device preference: {device_preference}, selected: {device}") + if cuda_available: + logging.debug(f"GPU: {torch.cuda.get_device_name(0)}") + + context_size = kwargs.pop("context_size", 4096) + + # If user explicitly requested CPU, force n_gpu_layers=0 regardless of other settings + if device_preference == "cpu": + n_gpu_layers = 0 + logging.debug("Device preference is CPU, forcing CPU-only execution") + + if isinstance(n_gpu_layers, str) and n_gpu_layers.lower() == "auto": + if cuda_available: + available_vram = get_available_vram() + logging.debug(f"Auto GPU layers: Available VRAM: {available_vram:.2f} GB") + + try: + model_size_gb = os.path.getsize(model_path) / (1024**3) + logging.debug(f"Auto GPU layers: Model size: {model_size_gb:.2f} GB") + except Exception: + model_size_gb = 4.0 + logging.debug(f"Auto GPU layers: Could not determine model size, using estimate: {model_size_gb:.2f} GB") + + optimal_layers = calculate_optimal_gpu_layers( + model_size_gb=model_size_gb, available_vram_gb=available_vram, context_size=context_size + ) + + if optimal_layers == -1: + logging.debug("Auto GPU layers: Using all GPU layers (model fits in VRAM)") + n_gpu_layers = -1 + elif optimal_layers == 0: + logging.info("Auto GPU layers: Not enough VRAM, using CPU only") + n_gpu_layers = 0 + else: + logging.debug(f"Auto GPU layers: Using {optimal_layers} layers on GPU") + n_gpu_layers = optimal_layers + else: + logging.debug("Auto GPU layers: CUDA not available, using CPU only") + n_gpu_layers = 0 + else: + n_gpu_layers = int(n_gpu_layers) + if n_gpu_layers == -1: + logging.debug("Using all GPU layers") + elif n_gpu_layers == 0: + logging.debug("Using CPU only") + else: + logging.debug(f"Using {n_gpu_layers} layers on GPU") + + use_mmap = kwargs.pop("use_mmap", True) + use_mlock = kwargs.pop("use_mlock", False) + n_batch = kwargs.pop("n_batch", 512) + n_threads = kwargs.pop("n_threads", None) + + llama_kwargs = { + "n_ctx": context_size, + "verbose": False, + "use_mmap": use_mmap, + "use_mlock": use_mlock, + "n_batch": n_batch, + } + + if n_threads is not None: + llama_kwargs["n_threads"] = n_threads + + llama_kwargs["n_gpu_layers"] = n_gpu_layers + llama_kwargs.update(kwargs) + + # Check if llama-cpp-python was built with CUDA support + try: + import llama_cpp + + if hasattr(llama_cpp, "__version__"): + logging.debug(f"llama-cpp-python version: {llama_cpp.__version__}") + except Exception: + pass + + self.model = Llama(model_path=model_path, **llama_kwargs) + load_time = time.time() - load_start + + if cuda_available and n_gpu_layers > 0: + try: + if torch.cuda.is_available(): + gpu_memory = torch.cuda.memory_allocated() / 1024**3 + logging.debug(f"GPU memory allocated: {gpu_memory:.2f} GB") + if gpu_memory > 0.1: + logging.debug("Model appears to be using GPU") + else: + logging.debug("Low GPU memory usage - model may be running on CPU") + except Exception: + pass + + logging.info(f"GGUF model loaded in {load_time:.2f}s") + + # Use HuggingFace tokenizer even though GGUF models can have embedded tokenizers + # Reason: Our code relies on apply_chat_template() which is HuggingFace-specific + logging.debug(f"Loading tokenizer from: {tokenizer_path}") + tokenizer_start = time.time() + tokenizer_path_abs = os.path.abspath(tokenizer_path) if not os.path.isabs(tokenizer_path) else tokenizer_path + + try: + from huggingface_hub.utils import HFValidationError + except ImportError: + HFValidationError = Exception + + if os.path.exists(tokenizer_path_abs) and os.path.isdir(tokenizer_path_abs): + try: + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path_abs, trust_remote_code=True, local_files_only=True) + except (HFValidationError, Exception) as e: + # If local_files_only fails, try without it (might be a cached HuggingFace model) + if isinstance(e, HFValidationError) or "Repo id must use alphanumeric" in str(e): + # Transformers is still validating the path - use absolute path with explicit local_files_only + logging.warning(f"Tokenizer path validation error, retrying with explicit local path handling: {e}") + else: + logging.warning(f"Failed to load tokenizer with local_files_only=True: {e}") + logging.warning("Trying without local_files_only...") + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path_abs, trust_remote_code=True) + else: + try: + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) + except (HFValidationError, Exception) as e: + if (isinstance(e, HFValidationError) or "Repo id must use alphanumeric" in str(e)) and ( + os.path.sep in tokenizer_path or (os.name == "nt" and len(tokenizer_path) > 1 and tokenizer_path[1] == ":") + ): + raise ValueError( + f"Tokenizer path does not exist or is invalid: {tokenizer_path_abs}\n" + f"Please ensure the tokenizer directory exists at this location.\n" + f"Expected location: {os.path.dirname(model_path)}/workflow-generator/ (or similar based on model name)" + ) from e + raise + tokenizer_time = time.time() - tokenizer_start + logging.debug(f"Tokenizer loaded in {tokenizer_time:.2f}s") + + self.device = device + + def generate( + self, input_ids, max_new_tokens: int = 512, temperature: float = 0.7, top_p: float = 0.9, do_sample: bool = True, **kwargs + ): + """ + Generate text using GGUF model. + + Args: + input_ids: Token IDs from HuggingFace tokenizer (tensor or list) + max_new_tokens: Maximum number of tokens to generate + temperature: Sampling temperature + top_p: Nucleus sampling parameter + do_sample: Whether to use sampling + **kwargs: Additional generation parameters + + Returns: + Generated token IDs tensor (full sequence: input + generated) + """ + if torch.is_tensor(input_ids): + input_ids_list = input_ids.cpu().tolist() + elif isinstance(input_ids, np.ndarray): + input_ids_list = input_ids.tolist() + else: + input_ids_list = input_ids + + # Handle batch dimension (llama-cpp-python expects single sequence) + if isinstance(input_ids_list[0], list): + input_ids_list = input_ids_list[0] + + output = self.model(input_ids_list, max_tokens=max_new_tokens, temperature=temperature, top_p=top_p, stop=[], **kwargs) + + if isinstance(output, dict) and "choices" in output: + generated_text = output["choices"][0]["text"] + elif hasattr(output, "choices"): + generated_text = output.choices[0].text + else: + generated_text = str(output) + + generated_token_ids = self.tokenizer.encode(generated_text, add_special_tokens=False) + full_token_ids = input_ids_list + generated_token_ids + full_tensor = torch.tensor([full_token_ids], dtype=torch.long) + + return full_tensor + + def __del__(self): + """Cleanup method for VRAM management.""" + if hasattr(self, "model"): + del self.model + if hasattr(self, "tokenizer"): + del self.tokenizer diff --git a/utils/model_loaders/huggingface_loader.py b/utils/model_loaders/huggingface_loader.py new file mode 100644 index 0000000..5a9f785 --- /dev/null +++ b/utils/model_loaders/huggingface_loader.py @@ -0,0 +1,220 @@ +""" +HuggingFace model loader using transformers with GPU layer management. +""" + +import logging +import os +import time + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +from ..model_manager import auto_detect_tokenizer_path +from .device_utils import ( + get_device_for_model, + get_dtype_for_model, +) + + +class HuggingFaceModelWrapper: + """Wrapper for HuggingFace models using transformers with GPU layer management.""" + + def __init__( + self, + model_path: str, + tokenizer_path: str | None = None, + dtype: str | torch.dtype = "auto", + device_preference: str = "auto", + max_memory: dict | None = None, + attn_implementation: str = "auto", + **kwargs, + ): + """ + Initialize HuggingFace model wrapper. + + Args: + model_path: Path to HuggingFace model directory + tokenizer_path: Optional separate tokenizer path (defaults to model_path) + dtype: "auto", "fp16", "bf16", "fp32", "fp8", or torch.dtype + device_preference: "auto", "cuda", or "cpu" + max_memory: Dict for device_map memory limits (e.g., {0: "10GB", "cpu": "30GB"}) + attn_implementation: "auto", "flash_attention_2", or "sdpa" + **kwargs: Additional arguments for AutoModelForCausalLM.from_pretrained + """ + logging.info(f"Loading HuggingFace model: {os.path.basename(model_path)}") + load_start = time.time() + + device = get_device_for_model(device_preference) + logging.debug(f"Device preference: {device_preference}, selected: {device}") + + if isinstance(dtype, str): + if dtype == "auto" and device_preference == "auto": + torch_dtype = get_dtype_for_model("fp16", device) + else: + torch_dtype = get_dtype_for_model(dtype, device) + else: + torch_dtype = dtype + + logging.debug(f"Using dtype: {torch_dtype}") + + if attn_implementation == "auto": + try: + import flash_attn # noqa: F401 + + if device == "cuda" and torch.cuda.is_available(): + attn_impl = "flash_attention_2" + logging.debug("Auto-detected attention: flash_attention_2") + else: + attn_impl = "sdpa" + logging.debug("Auto-detected attention: sdpa") + except ImportError: + attn_impl = "sdpa" + logging.debug("Auto-detected attention: sdpa") + else: + attn_impl = attn_implementation + if attn_impl == "flash_attention_2" and (device != "cuda" or not torch.cuda.is_available()): + logging.warning("flash_attention_2 requires NVIDIA CUDA, but CUDA is not available. Falling back to sdpa.") + attn_impl = "sdpa" + logging.debug(f"Using attention: {attn_impl}") + + load_kwargs = { + "torch_dtype": torch_dtype, + "trust_remote_code": True, + "attn_implementation": attn_impl, + } + + if device_preference == "cpu": + load_kwargs["device_map"] = {"": "cpu"} + logging.debug("Using CPU only") + elif device == "cuda": + load_kwargs["device_map"] = "auto" + if max_memory is not None: + load_kwargs["max_memory"] = max_memory + logging.debug("Using CUDA with device_map='auto' and max_memory limits") + else: + logging.debug("Using CUDA with device_map='auto' (automatic GPU/CPU split)") + else: + load_kwargs["device_map"] = {"": "cpu"} + logging.debug("Using CPU only") + + quantization_enabled = False + if "quantization_config" in kwargs: + quantization_enabled = True + logging.debug("8-bit quantization enabled") + + load_kwargs.update(kwargs) + + if quantization_enabled: + logging.info("Quantizing model to 8-bit (this may take a few minutes for large models)...") + + if tokenizer_path is None: + if not os.path.isabs(model_path): + model_path = os.path.abspath(model_path) + model_path = os.path.normpath(model_path) + if model_path.endswith(".gguf"): + raise ValueError( + f"Invalid model path for HuggingFace format: {model_path}\nHuggingFace models must be directories, not .gguf files.\n" + ) + tokenizer_path = auto_detect_tokenizer_path(model_path, model_format="huggingface") + logging.debug(f"Auto-detected tokenizer path: {tokenizer_path}") + + tokenizer_load_path = tokenizer_path if tokenizer_path else model_path + + self.model = AutoModelForCausalLM.from_pretrained(model_path, **load_kwargs) + + load_time = time.time() - load_start + logging.info(f"HuggingFace model loaded in {load_time:.2f}s") + + logging.debug(f"Loading tokenizer from: {tokenizer_load_path}") + tokenizer_start = time.time() + tokenizer_path_abs = os.path.abspath(tokenizer_load_path) if not os.path.isabs(tokenizer_load_path) else tokenizer_load_path + + try: + from huggingface_hub.utils import HFValidationError + except ImportError: + HFValidationError = Exception + + if os.path.exists(tokenizer_path_abs) and os.path.isdir(tokenizer_path_abs): + try: + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path_abs, trust_remote_code=True, local_files_only=True) + except (HFValidationError, Exception) as e: + if isinstance(e, HFValidationError) or "Repo id must use alphanumeric" in str(e): + logging.warning(f"Tokenizer path validation error, retrying with explicit local path handling: {e}") + else: + logging.warning(f"Failed to load tokenizer with local_files_only=True: {e}") + logging.warning("Trying without local_files_only...") + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path_abs, trust_remote_code=True) + else: + try: + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_load_path, trust_remote_code=True) + except (HFValidationError, Exception) as e: + if (isinstance(e, HFValidationError) or "Repo id must use alphanumeric" in str(e)) and ( + os.path.sep in tokenizer_load_path or + (os.name == "nt" and len(tokenizer_load_path) > 1 and tokenizer_load_path[1] == ":") + ): + raise ValueError( + f"Tokenizer path does not exist: {tokenizer_path_abs}\n" + f"Please ensure the tokenizer directory exists at this location." + ) from e + raise + tokenizer_time = time.time() - tokenizer_start + logging.debug(f"Tokenizer loaded in {tokenizer_time:.2f}s") + + if hasattr(self.model, "device"): + self.device = self.model.device + elif hasattr(self.model, "hf_device_map"): + self.device = torch.device(device) if device == "cuda" else torch.device("cpu") + else: + self.device = next(self.model.parameters()).device + + try: + if torch.cuda.is_available() and device == "cuda": + vram_used = torch.cuda.memory_allocated() / 1024**3 + logging.debug(f"GPU memory allocated: {vram_used:.2f} GB") + except Exception: + pass + + def generate( + self, input_ids, max_new_tokens: int = 512, temperature: float = 0.7, top_p: float = 0.9, do_sample: bool = True, **kwargs + ): + """ + Generate text using HuggingFace model. + + Args: + input_ids: Token IDs tensor + max_new_tokens: Maximum number of tokens to generate + temperature: Sampling temperature + top_p: Nucleus sampling parameter + do_sample: Whether to use sampling + **kwargs: Additional generation parameters (seed will be extracted and used separately) + + Returns: + Generated token IDs tensor + """ + seed = kwargs.pop("seed", None) + + if seed is not None: + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + if torch.is_tensor(input_ids): + if isinstance(self.device, dict): + device = list(self.device.values())[0] if self.device else torch.device("cpu") + else: + device = self.device + input_ids = input_ids.to(device) + + with torch.no_grad(): + generated_ids = self.model.generate( + input_ids, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, do_sample=do_sample, **kwargs + ) + + return generated_ids + + def __del__(self): + """Cleanup method for VRAM management.""" + if hasattr(self, "model"): + del self.model + if hasattr(self, "tokenizer"): + del self.tokenizer diff --git a/utils/model_manager.py b/utils/model_manager.py new file mode 100644 index 0000000..414c397 --- /dev/null +++ b/utils/model_manager.py @@ -0,0 +1,443 @@ +""" +Model manager for resolving model paths and detecting devices. +Integrates with ComfyUI's model management system. +""" + +import logging +import os +from pathlib import Path + +# Try to import ComfyUI folder_paths for model path resolution +try: + import folder_paths + + COMFYUI_AVAILABLE = True +except ImportError: + COMFYUI_AVAILABLE = False + logging.warning("ComfyUI folder_paths not available, using basic path resolution") + +# Try to import ComfyUI model management for device detection +try: + from comfy import model_management + + COMFYUI_MODEL_MGMT = True +except ImportError: + COMFYUI_MODEL_MGMT = False + + +def get_comfyui_base_path() -> Path | None: + """ + Get the base path of the ComfyUI installation. + + Returns: + Path to ComfyUI base directory or None if not found + """ + # Common locations: + # 1. Current working directory (if running from ComfyUI root) + # 2. Parent of custom_nodes directory + # 3. Environment variable + current_path = Path.cwd() + + if (current_path / "comfy").exists() or (current_path / "comfyUI").exists(): + return current_path + + for parent in current_path.parents: + if (parent / "comfy").exists() or (parent / "comfyUI").exists(): + return parent + + comfyui_path = os.environ.get("COMFYUI_PATH") + if comfyui_path: + return Path(comfyui_path) + + return None + + +def resolve_model_path(model_path: str, model_type: str = "llm", fallback_to_absolute: bool = True) -> str: + """ + Resolve a model path relative to ComfyUI's model directories. + + Args: + model_path: Model path (can be relative or absolute) + model_type: Type of model ("llm", "gguf", etc.) + fallback_to_absolute: If True, return absolute path if not found in ComfyUI dirs + + Returns: + Resolved absolute path to the model + """ + if os.path.isabs(model_path) and os.path.exists(model_path): + return model_path + + comfyui_base = get_comfyui_base_path() + if comfyui_base: + for base_dir in ["models/LLM", f"models/{model_type}"]: + # Handle subdirectory paths (e.g., "GGUF/llama-joycaption-beta-one-hf-llava.IQ4_XS.gguf") + full_path = comfyui_base / base_dir / model_path + if full_path.exists(): + return str(full_path) + + # Also try with just the filename (for backward compatibility) + filename = os.path.basename(model_path) + full_path = comfyui_base / base_dir / filename + if full_path.exists(): + return str(full_path) + + full_path = comfyui_base / model_path + if full_path.exists(): + return str(full_path) + + if COMFYUI_AVAILABLE: + try: + if model_path.endswith(".gguf"): + model_dirs = folder_paths.get_folder_paths("unet_gguf") + if not model_dirs: + model_dirs = folder_paths.get_folder_paths("clip_gguf") + if not model_dirs: + model_dirs = folder_paths.get_folder_paths("llm") + + for model_dir in model_dirs: + full_path = os.path.join(model_dir, model_path) + if os.path.exists(full_path): + return full_path + + filename = os.path.basename(model_path) + full_path = os.path.join(model_dir, filename) + if os.path.exists(full_path): + return full_path + + model_dirs = folder_paths.get_folder_paths("llm") + for model_dir in model_dirs: + full_path = os.path.join(model_dir, model_path) + if os.path.exists(full_path): + return full_path + + filename = os.path.basename(model_path) + full_path = os.path.join(model_dir, filename) + if os.path.exists(full_path): + return full_path + except Exception as e: + logging.warning(f"Failed to resolve path using folder_paths: {e}") + + comfyui_base = get_comfyui_base_path() + if comfyui_base: + llm_dir = comfyui_base / "models" / "LLM" + if llm_dir.exists(): + full_path = llm_dir / model_path + if full_path.exists(): + return str(full_path) + + filename = os.path.basename(model_path) + full_path = llm_dir / filename + if full_path.exists(): + return str(full_path) + + if fallback_to_absolute: + abs_path = os.path.abspath(model_path) + if os.path.exists(abs_path): + return abs_path + # Return the path anyway (let the caller handle missing file) + return abs_path + + return model_path + + +def auto_detect_tokenizer_path(model_path: str, model_format: str = "gguf") -> str: + """ + Automatically detect tokenizer path based on model path. + + For GGUF models: Tokenizer is typically in a directory with the same name + as the model (without .gguf extension) in the same parent directory. + Example: models/LLM/workflow-generator-q8_0.gguf -> models/LLM/workflow-generator/ + + For HuggingFace models: Tokenizer is in the same directory as the model. + Example: models/LLM/Qwen2.5-7B-Instruct/ -> models/LLM/Qwen2.5-7B-Instruct/ + + Args: + model_path: Resolved absolute path to the model file or directory + model_format: "gguf" or "huggingface" + + Returns: + Resolved absolute path to the tokenizer directory + """ + import os + + if not os.path.isabs(model_path): + model_path = os.path.abspath(model_path) + model_path = os.path.normpath(model_path) + model_path_obj = Path(model_path).resolve() + + if model_format.lower() == "gguf": + # For GGUF: tokenizer is in a directory with the same name (without .gguf) + # Example: workflow-generator-q8_0.gguf -> workflow-generator/ + model_name = model_path_obj.stem + # Remove quantization suffix if present (e.g., "workflow-generator-q8_0" -> "workflow-generator") + # Common suffixes: -q8_0, -q4_0, -q4_1, -q5_0, -q5_1, -q6_0, -f16, -f32 + import re + + model_name_clean = re.sub(r"-[qf]\d+[_\d]*$", "", model_name) + + parent_dir = model_path_obj.parent + tokenizer_dir = parent_dir / model_name_clean + + if tokenizer_dir.exists() and tokenizer_dir.is_dir(): + tokenizer_files = ["tokenizer.json", "tokenizer_config.json", "vocab.json"] + if any((tokenizer_dir / f).exists() for f in tokenizer_files): + return str(tokenizer_dir) + + # Fallback: try with full model name (including quantization suffix) + tokenizer_dir = parent_dir / model_name + if tokenizer_dir.exists() and tokenizer_dir.is_dir(): + tokenizer_files = ["tokenizer.json", "tokenizer_config.json", "vocab.json"] + if any((tokenizer_dir / f).exists() for f in tokenizer_files): + return str(tokenizer_dir) + + # If not found, return the expected path anyway (will be handled by error handling) + return str(tokenizer_dir) + + else: + if model_path_obj.is_dir(): + return str(model_path_obj) + else: + # Model path is a file, tokenizer should be in the same directory + # This shouldn't happen for HuggingFace, but handle it anyway + return str(model_path_obj.parent) + + +def get_device_info() -> tuple[str, bool]: + """ + Get device information using ComfyUI's model management. + + Returns: + Tuple of (device_string, cuda_available) + """ + if COMFYUI_MODEL_MGMT: + try: + dev = model_management.get_torch_device() + if hasattr(dev, "type"): + if dev.type == "cuda": + return "cuda", True + elif dev.type == "cpu": + return "cpu", False + elif dev.type == "mps": + return "mps", False + except Exception as e: + logging.warning(f"Failed to get device from ComfyUI: {e}") + + try: + import torch + + if torch.cuda.is_available(): + return "cuda", True + except ImportError: + pass + + return "cpu", False + + +def register_llm_folder_type(): + """ + Register "llm" folder type in ComfyUI's folder_paths. + This allows us to use folder_paths.get_filename_list("llm") and folder_paths.get_full_path("llm", ...). + """ + if not COMFYUI_AVAILABLE: + logging.warning("ComfyUI folder_paths not available, cannot register LLM folder type") + return + + try: + if "llm" in folder_paths.folder_names_and_paths: + return + + comfyui_base = get_comfyui_base_path() + if comfyui_base: + llm_dir = comfyui_base / "models" / "LLM" + if llm_dir.exists(): + # folder_names_and_paths structure: {key: (list_of_paths, set_of_extensions)} + # For LLM, we want to accept both .gguf files and directories (no extension filter) + folder_paths.folder_names_and_paths["llm"] = ([str(llm_dir)], set()) + logging.info(f"Registered LLM folder type: {llm_dir}") + else: + logging.warning(f"LLM directory not found: {llm_dir}") + else: + try: + # Look for existing model directories to infer the models path + if "checkpoints" in folder_paths.folder_names_and_paths: + checkpoints_paths, _ = folder_paths.folder_names_and_paths["checkpoints"] + if checkpoints_paths: + # models/LLM should be sibling to models/checkpoints + models_dir = Path(checkpoints_paths[0]).parent + llm_dir = models_dir / "LLM" + if llm_dir.exists(): + folder_paths.folder_names_and_paths["llm"] = ([str(llm_dir)], set()) + logging.info(f"Registered LLM folder type: {llm_dir}") + return + except Exception as e: + logging.warning(f"Failed to infer LLM directory from existing folders: {e}") + + logging.warning("Could not determine ComfyUI base path, LLM folder type not registered") + except Exception as e: + logging.error(f"Failed to register LLM folder type: {e}") + + +def get_llm_model_list() -> list[str]: + """ + Get all models from ComfyUI/models/LLM/ (both .gguf files and directories). + Recursively scans subdirectories for .gguf files. + Returns only model names (not full paths). + + Returns: + List of model names (e.g., ["workflow-generator-q8_0.gguf", "Qwen2.5-7B-Instruct"]) + """ + model_names = [] + + def scan_directory(directory: Path, prefix: str = ""): + """ + Recursively scan directory for .gguf files and HuggingFace model directories. + + Args: + directory: Directory to scan + prefix: Prefix to add to model names (for subdirectories) + """ + if not directory.exists(): + return + + try: + for item in directory.iterdir(): + if item.is_file() and item.suffix == ".gguf": + if prefix: + model_names.append(f"{prefix}{item.name}") + else: + model_names.append(item.name) + elif item.is_dir(): + has_gguf = False + try: + has_gguf = any(f.suffix == ".gguf" for f in item.iterdir() if f.is_file()) + except (PermissionError, OSError) as e: + logging.debug(f"Cannot read directory {item}: {e}") + continue + + if has_gguf: + # Directory contains .gguf files - scan recursively and add files with prefix + sub_prefix = f"{item.name}/" if prefix else f"{item.name}/" + scan_directory(item, sub_prefix) + else: + # Directory is likely a HuggingFace model - add directory name + if prefix: + model_names.append(f"{prefix}{item.name}") + else: + model_names.append(item.name) + except (PermissionError, OSError) as e: + logging.warning(f"Cannot scan directory {directory}: {e}") + + if not COMFYUI_AVAILABLE: + comfyui_base = get_comfyui_base_path() + if comfyui_base: + llm_dir = comfyui_base / "models" / "LLM" + scan_directory(llm_dir) + return sorted(model_names) + + try: + if "llm" in folder_paths.folder_names_and_paths: + llm_paths, _ = folder_paths.folder_names_and_paths["llm"] + for llm_path in llm_paths: + llm_dir = Path(llm_path) + scan_directory(llm_dir) + else: + comfyui_base = get_comfyui_base_path() + if comfyui_base: + llm_dir = comfyui_base / "models" / "LLM" + scan_directory(llm_dir) + except Exception as e: + logging.warning(f"Failed to get LLM model list: {e}") + + return sorted(model_names) + + +def get_workflow_generator_models() -> list[str]: + """ + Filter models suitable for workflow generator. + Includes both GGUF files and HuggingFace directories. + + Returns: + List of model names suitable for workflow generator + """ + all_models = get_llm_model_list() + # For now, return all models (both GGUF and HuggingFace work for workflow generator) + # In the future, we could filter by naming patterns if needed + return all_models + + +def get_refine_agent_models() -> list[str]: + """ + Filter models suitable for refine_agent (node validator). + Includes both GGUF files and HuggingFace directories. + + Returns: + List of model names suitable for refine_agent + """ + all_models = get_llm_model_list() + # For now, return all models (both GGUF and HuggingFace work for refine_agent) + # In the future, we could filter by naming patterns if needed + return all_models + + +def get_embedding_models() -> list[str]: + """ + Filter embedding models (SentenceTransformer/HuggingFace directories only). + Excludes .gguf files as embedding models are typically HuggingFace/SentenceTransformer. + + Returns: + List of embedding model directory names + """ + all_models = get_llm_model_list() + embedding_models = [m for m in all_models if not m.endswith(".gguf")] + return embedding_models + + +def detect_model_format(model_name: str) -> str: + """ + Auto-detect model format from model name. + + Args: + model_name: Model name (e.g., "workflow-generator-q8_0.gguf" or "Qwen2.5-7B-Instruct") + + Returns: + "gguf" if model ends with .gguf, "huggingface" if it's a directory + """ + if model_name.endswith(".gguf"): + return "gguf" + else: + return "huggingface" + + +def get_model_full_path(model_name: str) -> str: + """ + Get full path to a model using ComfyUI's folder_paths. + + Args: + model_name: Model name (e.g., "workflow-generator-q8_0.gguf" or "Qwen2.5-7B-Instruct") + + Returns: + Full absolute path to the model file or directory + """ + from pathlib import Path + + if not COMFYUI_AVAILABLE: + return resolve_model_path(model_name, model_type="llm") + + try: + if "llm" in folder_paths.folder_names_and_paths: + # Check for exact match first to avoid fuzzy matching issues + llm_paths, _ = folder_paths.folder_names_and_paths["llm"] + for llm_path in llm_paths: + exact_path = Path(llm_path) / model_name + if exact_path.exists() and ( + (not model_name.endswith(".gguf") and exact_path.is_dir()) or + (model_name.endswith(".gguf") and exact_path.is_file()) + ): + return str(exact_path) + + return folder_paths.get_full_path("llm", model_name) + else: + return resolve_model_path(model_name, model_type="llm") + except Exception as e: + logging.warning(f"Failed to get full path using folder_paths: {e}, falling back to resolve_model_path") + return resolve_model_path(model_name, model_type="llm") diff --git a/utils/node_utils.py b/utils/node_utils.py new file mode 100644 index 0000000..a41aeda --- /dev/null +++ b/utils/node_utils.py @@ -0,0 +1,141 @@ +import logging + +try: + import psutil +except ImportError: + psutil = None + + +def get_n_threads_options() -> list[str]: + """ + Generate options for n_threads Combo input. + Returns: List of strings ["All", "Auto"] + integer options up to CPU count + """ + try: + if psutil is not None: + cpu_count = psutil.cpu_count(logical=True) + else: + import multiprocessing + + cpu_count = multiprocessing.cpu_count() + + options = ["All", "Auto"] + options.extend([str(i) for i in range(1, cpu_count + 1)]) + return options + except (OSError, RuntimeError, AttributeError, Exception) as e: + logging.warning(f"Could not detect CPU count, using fallback options: {e}") + return ["All", "Auto"] + [str(i) for i in range(1, 33)] + + +def parse_n_threads(value: str) -> int | None: + """ + Parse n_threads Combo value to integer or None. + + Args: + value: String from Combo input ("All", "Auto", or integer string) + + Returns: + int for "All" or specific count, None for "Auto" + """ + if value == "All": + try: + if psutil is not None: + return psutil.cpu_count(logical=True) + else: + import multiprocessing + + return multiprocessing.cpu_count() + except Exception: + return 32 + elif value == "Auto": + return None + else: + try: + return int(value) + except (ValueError, TypeError): + return None + + +def check_llm_lazy_status( + use_llm_refinement: bool = False, + refine_model_path: str | None = None, + max_new_tokens: int | None = None, + device_preference: str | None = None, + n_gpu_layers: int | None = None, + dtype: str | None = None, + attn_implementation: str | None = None, + context_size: int | None = None, + temperature: float | None = None, + top_p: float | None = None, + seed: int | None = None, + # Additional parameters that might be present + refine_max_new_tokens: int | None = None, + refine_device_preference: str | None = None, + refine_n_gpu_layers: int | None = None, + refine_dtype: str | None = None, + refine_attn_implementation: str | None = None, + refine_context_size: int | None = None, + refine_temperature: float | None = None, + refine_top_p: float | None = None, + refine_auto_gpu_layers: bool | None = None, + auto_gpu_layers: bool | None = None, + **kwargs, +) -> list[str]: + """ + Determine which lazy inputs need to be evaluated for LLM refinement. + Handles both standard parameter names (NodeValidator) and prefixed names (Pipeline). + + Returns: + List of input names that need to be evaluated. + """ + needed = [] + + if not use_llm_refinement: + return needed + + def check_add(val, name): + if val is None: + needed.append(name) + + # Detect if we're in pipeline context (any prefixed parameter is not None) + is_pipeline_context = any( + [ + refine_max_new_tokens is not None, + refine_device_preference is not None, + refine_n_gpu_layers is not None, + refine_dtype is not None, + refine_attn_implementation is not None, + refine_context_size is not None, + refine_temperature is not None, + refine_top_p is not None, + refine_auto_gpu_layers is not None, + ] + ) + + if is_pipeline_context: + # Pipeline node context - check prefixed names + check_add(refine_model_path, "refine_model_path") + check_add(refine_max_new_tokens, "refine_max_new_tokens") + check_add(refine_device_preference, "refine_device_preference") + check_add(refine_n_gpu_layers, "refine_n_gpu_layers") + check_add(refine_dtype, "refine_dtype") + check_add(refine_attn_implementation, "refine_attn_implementation") + check_add(refine_context_size, "refine_context_size") + check_add(refine_temperature, "refine_temperature") + check_add(refine_top_p, "refine_top_p") + check_add(refine_auto_gpu_layers, "refine_auto_gpu_layers") + else: + # NodeValidator node context - check standard names + check_add(refine_model_path, "refine_model_path") + check_add(max_new_tokens, "max_new_tokens") + check_add(device_preference, "device_preference") + check_add(n_gpu_layers, "n_gpu_layers") + check_add(dtype, "dtype") + check_add(attn_implementation, "attn_implementation") + check_add(context_size, "context_size") + check_add(temperature, "temperature") + check_add(top_p, "top_p") + check_add(seed, "seed") + check_add(auto_gpu_layers, "auto_gpu_layers") + + return needed diff --git a/utils/process_utils.py b/utils/process_utils.py new file mode 100644 index 0000000..22f2143 --- /dev/null +++ b/utils/process_utils.py @@ -0,0 +1,163 @@ +""" +Utility functions for processing ComfyUI-WorkflowGenerator workflows and diagrams. +""" + +import json +import logging +from typing import Any + + +def format(node_type: str) -> str: + """ + Format node type by removing the last underscore and everything after it. + Also handles node names with parentheses (e.g., "Node Name (author)"). + + Args: + node_type: Node type string (e.g., "CheckpointLoaderSimple_1" or "SDXL Empty Latent Image (rgthree)") + + Returns: + Formatted node type (e.g., "CheckpointLoaderSimple" or "SDXL Empty Latent Image (rgthree)") + + Note: + - If node type contains underscore, removes suffix (e.g., "NodeName_0" -> "NodeName") + - If node type contains parentheses, keeps as-is (e.g., "Node Name (author)" -> "Node Name (author)") + - If node type has neither, returns as-is + """ + last_index = node_type.rfind("_") + if last_index != -1: + suffix = node_type[last_index + 1 :] + if suffix.isdigit() or (len(suffix) <= 2 and suffix.isalnum()): + result = node_type[:last_index] + return result + return node_type + + +def del_digram_primitive(digram: list[list[str]]) -> list[list[str]]: + """ + Remove primitive nodes from diagram. + + Args: + digram: List of edges in the diagram, where each edge is [src_type, src_name, dst_type, dst_name] + + Returns: + Filtered diagram without primitive nodes + """ + result_digram = [] + for edge in digram: + src_type, src_name, dst_type, dst_name = edge + src_node = format(src_type) + dst_node = format(dst_type) + if src_node == "PrimitiveNode" or dst_node == "PrimitiveNode": + continue + result_digram.append(edge) + return result_digram + + +def json_format(encode: str | Any) -> Any: + """ + Extract and parse JSON from encoded string. + + Args: + encode: String containing JSON (possibly wrapped in markdown code blocks) or already parsed JSON + + Returns: + Parsed JSON object + + Raises: + ValueError: If JSON format is invalid or cannot be parsed + """ + if not isinstance(encode, str): + return encode + + if "'''json" in encode: + start = encode.find("'''json") + len("'''json") + end = encode.find("'''", start) + if start == -1 or end == -1: + raise ValueError("Invalid JSON format: missing closing '''") + encode = encode[start:end].strip().strip('"') + + if "```json" in encode: + start = encode.find("```json") + len("```json") + end = encode.find("```", start) + if start == -1 or end == -1: + raise ValueError("Invalid JSON format: missing closing ```") + encode = encode[start:end].strip().strip('"') + + encode = encode.replace("'", '"') + encode = encode.replace("\n", "") + + obj_start = encode.find("{") + obj_end = encode.rfind("}") + + arr_start = encode.find("[") + arr_end = encode.rfind("]") + + # Determine which format to use (prefer object if both exist) + if obj_start != -1 and obj_end != -1 and obj_end > obj_start: + count_open = encode.count("{") + count_close = encode.count("}") + if count_open != count_close: + logging.debug(f"Invalid JSON format: object bracket mismatch. Content: {encode[:200]}") + raise ValueError("Invalid JSON format: object bracket mismatch") + encode = encode[obj_start : obj_end + 1] + elif arr_start != -1 and arr_end != -1 and arr_end > arr_start: + count_open = encode.count("[") + count_close = encode.count("]") + if count_open != count_close: + logging.debug(f"Invalid JSON format: array bracket mismatch. Content: {encode[:200]}") + raise ValueError("Invalid JSON format: array bracket mismatch") + encode = encode[arr_start : arr_end + 1] + else: + # Try to parse as-is (might already be valid JSON) + pass + + try: + result = json.loads(encode) + return result + except json.JSONDecodeError as e: + logging.debug(f"JSON decode error: {e}") + logging.debug(f"Content: {encode[:500]}") + raise + + +def fetch_node_input_names(meta_info: dict[str, Any]) -> list[str]: + """ + Extract input names from node meta info. + + Args: + meta_info: Node metadata dictionary containing 'inputs' key + + Returns: + List of input names + """ + input_names = [] + if "inputs" not in meta_info or meta_info["inputs"] is None: + return input_names + + inputs = meta_info["inputs"] + for key, item in inputs.items(): + if key != "hidden" and isinstance(item, dict): + for name in item: + input_names.append(name) + + return input_names + + +def fetch_node_output_names(meta_info: dict[str, Any]) -> list[str]: + """ + Extract output names from node meta info. + + Args: + meta_info: Node metadata dictionary containing 'output_names' or 'outputs' key + + Returns: + List of output names + """ + if "output_names" in meta_info and meta_info["output_names"] is not None: + return meta_info["output_names"] + + # Fallback: use output types as names + if "outputs" in meta_info and meta_info["outputs"] is not None: + return [str(output) for output in meta_info["outputs"]] + + return [] diff --git a/utils/workflow_layout.py b/utils/workflow_layout.py new file mode 100644 index 0000000..bb6f55a --- /dev/null +++ b/utils/workflow_layout.py @@ -0,0 +1,149 @@ +from collections import defaultdict, deque + + +class WorkflowLayout: + """ + Layout engine for ComfyUI workflows. + Positions nodes in a logical left-to-right flow (DAG layout). + """ + + # Standard ComfyUI node dimensions and spacing + NODE_WIDTH = 210 + NODE_HEIGHT_BASE = 60 + SLOT_HEIGHT = 20 + + # Layout spacing + RANK_SEP = 450 # Horizontal spacing between layers + NODE_SEP = 150 # Vertical spacing between nodes + + def __init__(self, nodes: list[dict], links: list[list]): + """ + Initialize layout engine. + + Args: + nodes: List of node dictionaries + links: List of link arrays [id, src_id, src_slot, dst_id, dst_slot, type] + """ + self.nodes = nodes + self.links = links + self.node_map = {node["id"]: node for node in nodes} + + def apply_layout(self): + """ + Calculate and apply positions to all nodes. + Modifies the 'pos' attribute of nodes in-place. + """ + if not self.nodes: + return + + adj = defaultdict(list) + in_degree = defaultdict(int) + parents = defaultdict(list) + + for node in self.nodes: + in_degree[node["id"]] = 0 + + # Link format: [id, src_id, src_slot, dst_id, dst_slot, type] + for link in self.links: + if not link: + continue + + if isinstance(link, (list, tuple)) and len(link) >= 5: + src_id = link[1] + dst_id = link[3] + + if src_id in self.node_map and dst_id in self.node_map: + adj[src_id].append(dst_id) + parents[dst_id].append(src_id) + in_degree[dst_id] += 1 + + # Assign ranks (levels) using topological sort approach + ranks = {} + queue = deque() + + for node_id in self.node_map: + if in_degree[node_id] == 0: + queue.append((node_id, 0)) + ranks[node_id] = 0 + + # We use a relaxed approach to handle cycles if any (though unlikely in valid workflows) + processed_count = 0 + while queue: + u, rank = queue.popleft() + processed_count += 1 + + for v in adj[u]: + in_degree[v] -= 1 + current_rank = ranks.get(v, 0) + new_rank = max(current_rank, rank + 1) + ranks[v] = new_rank + + if in_degree[v] == 0: + queue.append((v, new_rank)) + + # Handle any nodes remaining (cycles or disconnected components) + for node_id in self.node_map: + if node_id not in ranks: + max_parent_rank = -1 + for p in parents[node_id]: + if p in ranks: + max_parent_rank = max(max_parent_rank, ranks[p]) + ranks[node_id] = max_parent_rank + 1 + + layers = defaultdict(list) + for node_id, rank in ranks.items(): + layers[rank].append(node_id) + + # Sort within layers to minimize crossings (heuristic) - sort by average parent Y position + sorted_ranks = sorted(layers.keys()) + + positions = {nid: [0, 0] for nid in self.node_map} + + for r in sorted_ranks: + layer_nodes = layers[r] + + if r > 0: + + def get_avg_parent_y(nid): + ps = parents[nid] + if not ps: + return 0 + known_parents = [p for p in ps if p in positions and positions[p][1] != 0] + if not known_parents: + return 0 + return sum(positions[p][1] for p in known_parents) / len(known_parents) + + layer_nodes.sort(key=get_avg_parent_y) + + current_y = 0 + for node_id in layer_nodes: + node = self.node_map[node_id] + + # Estimate node height based on input/output count + input_count = len(node.get("inputs", [])) + output_count = len(node.get("outputs", [])) + height = self.NODE_HEIGHT_BASE + max(input_count, output_count) * self.SLOT_HEIGHT + + x = r * self.RANK_SEP + y = current_y + + positions[node_id] = [x, y] + node["pos"] = [x, y] + + current_y += height + self.NODE_SEP + + # Center the layer vertically around Y=0 + layer_height = current_y - self.NODE_SEP + offset_y = -layer_height / 2 + + for node_id in layer_nodes: + new_y = int(self.node_map[node_id]["pos"][1] + offset_y) + new_x = int(self.node_map[node_id]["pos"][0]) + self.node_map[node_id]["pos"] = [new_x, new_y] + positions[node_id] = [new_x, new_y] + + @staticmethod + def process(nodes: list[dict], links: list[list]): + """Static convenience method""" + layout = WorkflowLayout(nodes, links) + layout.apply_layout() diff --git a/wiki/Configuration.md b/wiki/Configuration.md new file mode 100644 index 0000000..7570bc3 --- /dev/null +++ b/wiki/Configuration.md @@ -0,0 +1,67 @@ +# Configuration Guide + +Complete configuration options, performance tuning, and model-specific settings. + +## Model Paths + +- **Relative paths**: `models/LLM/workflow-generator-q8_0.gguf` (relative to ComfyUI root) +- **Absolute paths**: `/path/to/model.gguf` +- System automatically resolves paths relative to ComfyUI's model directories + +## GPU Settings + +- **GPU Layers**: `"auto"` (default & recommended), `"all"`, `"none"`, or number (0-50) + - `"auto"`: Automatically determines optimal GPU layer allocation based on available VRAM to prevent OOM. + - `"all"`: Use all available GPU layers (fastest, but may cause OOM on lower VRAM). + - `"none"`: CPU only. + - Number: Specific count of GPU layers to use. +- **Device Preference**: `"auto"` (detects best), `"cuda"`, or `"cpu"` + +## Model-Specific Settings + +### GGUF Models + +- Handle quantization internally (no dtype settings needed) +- Use `n_gpu_layers` to control GPU/CPU split +- Use `context_size` to set context window (default: 4096) +- Performance parameters: + - `use_mmap` (bool): Memory-mapped file access (default: True) + - `use_mlock` (bool): Lock memory pages (default: False) + - `n_batch` (int): Batch size for processing (default: 512) + - `n_threads` (int, optional): CPU thread count (auto-detected if not set) + +### HuggingFace Models + +- Support dtype options: `"auto"`, `"fp16"`, `"bf16"`, `"fp32"`, `"fp8"`, or `torch.dtype` +- Attention implementations: + - `"auto"`: Automatically selects best available + - `"flash_attention_2"`: Flash Attention 2 (requires flash-attn package) + - `"sdpa"`: Scaled Dot Product Attention (built into PyTorch 2.0+) +- Device map: Automatically handles device placement + +## Performance Tuning + +### For Lower VRAM Usage + +- Use GGUF models (typically 4-8GB for q8_0 quantization vs 12-16GB+ for HuggingFace) +- Set `n_gpu_layers="all"` if VRAM allows +- Disable LLM refinement (`use_llm_refinement=False`) +- Use smaller context windows if possible + +### For Lower VRAM Usage + +- Reduce `n_gpu_layers` (try "auto" or lower number) +- Use smaller quantization (q4_0 instead of q8_0) +- Set `device_preference="cpu"` for some operations +- Use HuggingFace models with CPU offloading + +### For Better Accuracy + +- Enable LLM refinement (`use_llm_refinement=True`) +- Use larger models (q8_0 instead of q4_0) +- Increase `top_k` in NodeValidator for more candidate nodes + +--- + +[← Back to Home](Home) | [← Node Reference](Node-Reference) | [Next: Troubleshooting →](Troubleshooting) + diff --git a/wiki/Home.md b/wiki/Home.md new file mode 100644 index 0000000..efb7579 --- /dev/null +++ b/wiki/Home.md @@ -0,0 +1,30 @@ +# ComfyUI-WorkflowGenerator Wiki + +Welcome to the ComfyUI-WorkflowGenerator documentation wiki. This wiki contains detailed guides, references, and troubleshooting information. + +## Quick Links + +- [Installation Guide](Installation) - Complete setup instructions including llama-cpp-python installation +- [Node Reference](Node-Reference) - Detailed documentation for each node +- [Configuration Guide](Configuration) - Model paths, GPU settings, and performance tuning +- [Troubleshooting](Troubleshooting) - Common issues and solutions +- [Instruction Prompt Usage](Instruction-Prompt-Usage) - How instructions are used across pipeline steps + +## Overview + +ComfyUI-WorkflowGenerator is a custom node suite that generates ComfyUI workflows from natural language descriptions. It uses a three-stage pipeline: + +1. **WorkflowGenerator** - Converts natural language to workflow diagrams +2. **NodeValidator** (optional) - Validates and corrects node names +3. **WorkflowBuilder** - Converts diagrams to executable ComfyUI workflow JSON + +## Getting Started + +1. Follow the [Installation Guide](Installation) to set up the custom node +2. Read the [Node Reference](Node-Reference) to understand each component +3. Check [Troubleshooting](Troubleshooting) if you encounter issues + +## Main Repository + +For the main project repository, installation instructions, and quick start guide, see the [GitHub README](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator). + diff --git a/wiki/Installation.md b/wiki/Installation.md new file mode 100644 index 0000000..7304959 --- /dev/null +++ b/wiki/Installation.md @@ -0,0 +1,141 @@ +# Installation Guide + +This guide covers detailed installation instructions, including the `llama-cpp-python` setup required for GGUF models. + +## Prerequisites + +- ComfyUI installed and running +- Python 3.10 or higher +- CUDA-capable GPU (recommended) or CPU +- Git installed + +## Basic Installation + +Follow the main [README](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator) for basic installation steps (cloning, dependencies, etc.). + +## llama-cpp-python Installation + +**Important:** `llama-cpp-python` is **required** for GGUF models (recommended). It must be installed separately based on your system configuration. It is not included in `requirements.txt` because it requires different installation methods for CPU, CUDA, and Metal. + +For detailed `llama-cpp-python` installation instructions, see [INSTALL_LLAMACPP.md](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/blob/main/INSTALL_LLAMACPP.md) in the main repository. + +### Quick Install Options + +#### CPU Only (No GPU acceleration) +```bash +pip install llama-cpp-python +``` + +#### CUDA Support (NVIDIA GPUs) - Recommended + +**IMPORTANT NOTE FOR WINDOWS USERS:** +Installing `llama-cpp-python` with CUDA support on Windows is notoriously difficult because pre-built binaries (wheels) often don't match your specific combination of Python version, CUDA version, and OS. + +**If the quick install command below fails or falls back to CPU mode, you MUST compile from source.** This is normal and expected. + +**Try this first (Quick Install):** +```bash +pip install llama-cpp-python[cuda] +``` + +**If that fails (Compile from Source):** +This is the most reliable method for Windows. It takes about 20-30 minutes to compile. + +1. **Install Prerequisites:** + - [Visual Studio 2022 Community](https://visualstudio.microsoft.com/vs/community/) (Select "Desktop development with C++") + - [CUDA Toolkit 12.x](https://developer.nvidia.com/cuda-downloads) (Match your driver version) + - [CMake](https://cmake.org/download/) (Add to PATH during installation) + +2. **Compile and Install:** + ```powershell + # PowerShell + $env:CMAKE_ARGS="-DGGML_CUDA=on -DCMAKE_CUDA_ARCHITECTURES=native" + $env:FORCE_CMAKE=1 + pip install llama-cpp-python --no-cache-dir --verbose --force-reinstall + ``` + + ```batch + :: Command Prompt (cmd) + set CMAKE_ARGS=-DGGML_CUDA=on -DCMAKE_CUDA_ARCHITECTURES=native + set FORCE_CMAKE=1 + pip install llama-cpp-python --no-cache-dir --verbose --force-reinstall + ``` + +For a complete step-by-step guide with troubleshooting, please see the [Detailed Installation Guide (INSTALL_LLAMACPP.md)](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/blob/main/INSTALL_LLAMACPP.md). + +#### Metal Support (macOS with Apple Silicon) +```bash +pip install llama-cpp-python[metal] +``` + +### CUDA/PyTorch Compatibility + +`llama-cpp-python` with CUDA support is compatible with: +- **PyTorch versions:** 2.0+ (tested with 2.1.0, 2.2.0, 2.3.0, 2.4.0+) +- **CUDA versions:** 11.8, 12.1, 12.4+ (pre-built wheels available for common versions) +- **Python versions:** 3.10, 3.11, 3.12 + +**Note:** Pre-built wheels are available for common CUDA/Python combinations. If your specific combination isn't available, you'll need to compile from source. + +### Quick Check: Will I Need to Compile from Source? + +To check if a pre-built wheel is available for your setup: + +```bash +# Check your Python, PyTorch, and CUDA versions +python -c "import sys, torch; print(f'Python: {sys.version_info.major}.{sys.version_info.minor}'); print(f'PyTorch: {torch.__version__}'); print(f'CUDA: {torch.version.cuda if torch.cuda.is_available() else \"N/A\"}')" + +# Try installing (this will show if it's downloading a wheel or compiling) +pip install llama-cpp-python[cuda] --verbose +``` + +**Signs a pre-built wheel is available:** +- Installation completes in 1-2 minutes +- You see "Downloading" messages for `.whl` files +- No compilation messages (no "Building wheel", "CMake", etc.) + +**Signs you'll need to compile from source:** +- Installation takes 20-30 minutes +- You see "Building wheel" or "CMake" messages +- Requires Visual Studio Build Tools (Windows) or build-essential (Linux) + +### Troubleshooting CUDA Installation + +**Problem: "No module named 'llama_cpp'" after installation** +- Solution: Make sure you installed the CUDA version, not the CPU version +- Reinstall: `pip uninstall llama-cpp-python && pip install llama-cpp-python[cuda]` + +**Problem: Installation takes a very long time (20-30 minutes)** +- This means it's compiling from source (no pre-built wheel available) +- Ensure you have Visual Studio Build Tools (Windows) or build-essential (Linux) installed + +**Problem: "CUDA not found" or "CUDA version mismatch"** +- Check your CUDA version: `python -c "import torch; print(torch.version.cuda)"` +- Ensure CUDA is properly installed and accessible +- See [INSTALL_LLAMACPP.md](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/blob/main/INSTALL_LLAMACPP.md) for detailed troubleshooting + +**Problem: "CMake not found"** +- Install CMake: `pip install cmake` or download from https://cmake.org/ +- On Windows, you may also need Visual Studio Build Tools + +### Portable ComfyUI Installation + +For portable ComfyUI installations (Windows), use the embedded Python: + +```bash +cd +.\python_embeded\python.exe -s -m pip install llama-cpp-python[cuda] +``` + +Or if compiling from source: +```batch +.\python_embeded\python.exe -s -m pip install cmake +set CMAKE_ARGS=-DLLAMA_CUBLAS=on +set FORCE_CMAKE=1 +.\python_embeded\python.exe -s -m pip install llama-cpp-python --no-cache-dir +``` + +--- + +[← Back to Home](Home) | [Next: Node Reference →](Node-Reference) + diff --git a/wiki/Instruction-Prompt-Usage.md b/wiki/Instruction-Prompt-Usage.md new file mode 100644 index 0000000..bd4ecc6 --- /dev/null +++ b/wiki/Instruction-Prompt-Usage.md @@ -0,0 +1,76 @@ +# Instruction Prompt Usage Across Steps + +This page explains how the instruction prompt is utilized across the different stages of the workflow generation pipeline. + +## Step 1 (WorkflowGenerator) + +The instruction prompt is used to generate the initial workflow diagram. This is the primary input that determines the overall workflow structure. + +## Step 2 (NodeValidator) + +The instruction prompt serves as **context** for the LLM when correcting invalid node names. It helps the LLM understand the workflow's intent and choose the most appropriate node from candidate matches. + +### How the Instruction is Used in Step 2 + +The instruction does **NOT** override the refinement prompt. Instead, it is **inserted as context** within the prompt template. + +When LLM refinement is enabled (`use_llm_refinement=True`), the instruction you provide is inserted into a structured prompt that asks the LLM to select the best node from candidate matches. The full prompt structure is: + +``` +"I would like you to act as an expert in ComfyUI platform. +I will provide a example, including a description about ComfyUI workflow +and a logical diagram in json format represents the comfyui workflow. +The logical diagram is a edges list [edge_1, edge_2, edge_3, ... , edge_n], +each edge is consist of [output_node_name,output_name,input_node_name,input_name], +represents a line between output node and input node. +Example: description: {YOUR_INSTRUCTION_HERE}. logical diagram: {diagram}. +Now, This logical diagram has one error node name. error name: {error_name}. +I will give you some candidate nodes. Please combine the above information +to select the most suitable candidate node. Candidate nodes: {candidate_nodes}. +You just need to return you choose node name. +Please return result in pure JSON format, including: +'''json{"candidate_node_name":"..."}'''" +``` + +**Key Points:** +- Your instruction replaces `{YOUR_INSTRUCTION_HERE}` in the prompt +- The instruction provides **context** about the workflow's intent +- The LLM uses this context to make better decisions when selecting from candidate nodes +- The prompt structure itself is fixed and cannot be overridden + +### How Different Instructions Affect Refinement + +When using **LLM refinement** (`use_llm_refinement=True`), the instruction prompt in Step 2 directly influences node selection: + +| Scenario | Effect on Refinement | +|----------|---------------------| +| **Same instruction as Step 1** | ✅ **Recommended** - Provides consistent context, helps LLM maintain original intent | +| **More specific/detailed instruction** | ✅ **Beneficial** - Can guide LLM toward more precise node choices (e.g., "use SDXL nodes" vs "use SD 1.5 nodes") | +| **Different focus/emphasis** | ⚠️ **Can be beneficial or problematic** - May bias workflow toward different goals (e.g., "optimize for speed" vs "optimize for quality") | +| **Contradictory instruction** | ❌ **Not recommended** - May cause confusion and incorrect node selections | +| **Empty instruction** | ⚠️ **Works but less effective** - LLM still works but has less context to make decisions | + +### Example Use Cases for Different Instructions + +1. **Refining for Specific Model Type:** + - Step 1: `"Create a text-to-image workflow"` + - Step 2: `"Create a text-to-image workflow using SDXL models"` → Guides LLM to select SDXL-specific nodes + +2. **Refining for Performance:** + - Step 1: `"Create an image upscaling workflow"` + - Step 2: `"Create a fast image upscaling workflow using lightweight models"` → May prefer faster upscaling nodes + +3. **Refining for Quality:** + - Step 1: `"Create an image generation workflow"` + - Step 2: `"Create a high-quality image generation workflow with advanced sampling"` → May prefer higher-quality sampling nodes + +### Best Practices + +- **Default**: Connect the `instruction` output from Step 1 to Step 2 for consistent context +- **Advanced**: Use a modified instruction in Step 2 only when you want to guide refinement in a specific direction +- **When LLM refinement is disabled** (`use_llm_refinement=False`): The instruction has no effect (only semantic search is used) + +--- + +[← Back to Home](Home) + diff --git a/wiki/Node-Reference.md b/wiki/Node-Reference.md new file mode 100644 index 0000000..f733fb5 --- /dev/null +++ b/wiki/Node-Reference.md @@ -0,0 +1,131 @@ +# Node Reference + +Detailed documentation for each node in ComfyUI-WorkflowGenerator. + +## Workflow Generator Pipeline + +**Category**: `WorkflowGenerator` + +**Purpose**: One-click solution for complete workflow generation. Processes your instruction through all three stages sequentially: generates the workflow diagram, validates and corrects node names, then builds the final ComfyUI workflow JSON. + +**Inputs**: +- `instruction` (string, required): Natural language description of the workflow +- `workflow_generator_model_path` (string): Path to WorkflowGenerator model (GGUF or HuggingFace, default: `workflow-generator-q8_0.gguf`) +- `workflow_generator_tokenizer_path` (string, optional): Path to tokenizer (auto-detected if not provided) +- `model_format` (string): "gguf" or "huggingface" (default: "gguf") +- `auto_gpu_layers` (bool): Auto-calculate GPU layers based on VRAM (default: True) +- `n_gpu_layers` (int/string): GPU layers for GGUF ("auto", "all", "none", or 0-50) +- `device_preference` (string): "auto", "cuda", or "cpu" +- `use_llm_refinement` (bool): Enable LLM refinement in NodeValidator (default: False) +- `node_validator_model_path` (string, optional): Path to NodeValidator model (required if use_llm_refinement=True) +- `catalog_directory` (string): Directory containing node_list.json and node_info.json +- `filename_prefix` (string): Prefix for saved workflow files +- `save_workflow_json` (bool): Whether to save workflow to file + +**Outputs**: +- `workflow_json` (string): Complete executable ComfyUI workflow JSON +- `instruction` (string): Pass-through of input instruction + +**Usage**: All-in-one solution. No other nodes required. + +**Best for**: Quick workflow creation without intermediate steps + +## WorkflowGenerator Node + +**Category**: `WorkflowGenerator` + +**Purpose**: Generates workflow diagrams from natural language instructions using a fine-tuned LLM. + +**Usage**: **Step 1** in the pipeline. Connects to NodeValidator (optional) or directly to WorkflowBuilder. + +**Inputs**: +- `instruction` (string, required): Natural language description of the workflow +- `model_path` (string): Path to WorkflowGenerator model (default: `workflow-generator-q8_0.gguf`) +- `tokenizer_path` (string, optional): Path to tokenizer (auto-detected if not provided) +- `model_format` (string): "gguf" or "huggingface" +- `auto_gpu_layers` (bool): Auto-calculate GPU layers (default: True) +- `n_gpu_layers` (int/string): GPU layers for GGUF +- `device_preference` (string): "auto", "cuda", or "cpu" +- `context_size` (int): Context window size (default: 4096) +- `top_p` (float): Top-p sampling parameter (default: 0.7) +- `max_new_tokens` (int): Maximum tokens to generate (default: 8192) +- `temperature` (float): Sampling temperature (default: 0.95) +- `allow_primitive_nodes` (bool): Whether to include primitive nodes in the diagram (default: False) + +**Outputs**: +- `workflow_edges` (string): Workflow diagram in JSON format (list of edges) +- `instruction` (string): Pass-through of input instruction + +### Understanding `allow_primitive_nodes` + +The `allow_primitive_nodes` parameter controls whether primitive nodes (like integers, strings, floats used as inputs) are included in the generated workflow diagram. + +- **False (Default & Recommended)**: Removes primitive nodes from the diagram. This results in a cleaner workflow where values are set directly in the node widgets rather than via external primitive nodes. This is usually preferred for readability. +- **True**: Includes primitive nodes. This can be useful if you prefer to have all inputs explicitly visible as separate nodes, but it can clutter the workflow diagram. + +## NodeValidator Node + +**Category**: `WorkflowGenerator` + +**Purpose**: Validates and corrects invalid node names in workflow diagrams using semantic search and optional LLM refinement. + +**Usage**: **Step 2** (Optional). Connects between WorkflowGenerator and WorkflowBuilder. Can be skipped if validation is not needed. + +**Inputs**: +- `workflow_edges` (string, required): Workflow diagram from WorkflowGenerator +- `instruction` (string, optional): Context instruction for LLM refinement +- `use_llm_refinement` (bool): Enable LLM refinement (default: False) +- `node_validator_model_path` (string, optional): Path to NodeValidator model (required if use_llm_refinement=True) +- `node_validator_tokenizer_path` (string, optional): Path to tokenizer +- `embedding_model_path` (string): Path to embedding model for semantic search +- `catalog_directory` (string): Directory containing node_list.json and node_info.json +- `top_k` (int): Number of candidate nodes to consider (default: 5) +- `similarity_threshold` (float): Minimum similarity score (default: 0.5) +- `seed` (int, lazy): Random seed (only if LLM enabled) + +**Outputs**: +- `workflow_edges (refined)` (string): Corrected workflow diagram with valid node names + +**Two modes**: +- **Semantic search only** (faster, deterministic): `use_llm_refinement=False` +- **LLM refinement** (slower, more accurate): `use_llm_refinement=True` + +See [Instruction Prompt Usage](Instruction-Prompt-Usage) for details on how the instruction is used in this node. + +## WorkflowBuilder Node + +**Category**: `WorkflowGenerator` + +**Purpose**: Converts workflow diagrams into executable ComfyUI workflow JSON with proper node configurations and connections. + +**Usage**: **Step 3** (Final). Connects to NodeValidator or WorkflowGenerator. + +**Inputs**: +- `workflow_edges` (string, required): Refined workflow diagram (connection-only input) +- `catalog_directory` (string): Directory containing node_list.json and node_info.json +- `filename_prefix` (string): Prefix for saved workflow files +- `save_workflow_json` (bool): Whether to save workflow to file + +**Outputs**: +- `workflow_json` (string): Executable ComfyUI workflow JSON + +## UpdateNodeCatalog Node + +**Category**: `WorkflowGenerator` + +**Purpose**: Scans and catalogs all available ComfyUI nodes (native and custom) for use in validation and workflow building. + +**Inputs**: +- `catalog_directory` (string): Directory to save node_list.json and node_info.json files (default: "catalog") +- `force_update` (bool): Force update even if catalog files already exist + +**Outputs**: +- `status` (string): Status message ("success" or error) +- `node_count` (int): Number of nodes cataloged + +**When to run**: After installing new custom nodes or updating ComfyUI + +--- + +[← Back to Home](Home) | [← Installation](Installation) | [Next: Configuration →](Configuration) + diff --git a/wiki/README.md b/wiki/README.md new file mode 100644 index 0000000..ae0ae0e --- /dev/null +++ b/wiki/README.md @@ -0,0 +1,78 @@ +# GitHub Wiki Setup Guide + +This directory contains the wiki pages for the ComfyUI-WorkflowGenerator GitHub Wiki. + +## How to Upload to GitHub Wiki + +GitHub Wikis are stored in a separate repository. Here's how to set it up: + +### Method 1: Using GitHub Web Interface (Easiest) + +1. Go to your repository: `https://github.com/danielpflorian/ComfyUI-WorkflowGenerator` +2. Click on the **"Wiki"** tab +3. Click **"Create the first page"** or **"New Page"** +4. Copy the content from each `.md` file in this directory +5. Paste into the GitHub Wiki editor +6. Use the filename (without `.md`) as the page title +7. Save the page + +### Method 2: Using Git (Recommended for Bulk Upload) + +1. **Clone the wiki repository:** + ```bash + git clone https://github.com/danielpflorian/ComfyUI-WorkflowGenerator.wiki.git + ``` + +2. **Copy the wiki files:** + ```bash + cp wiki/*.md ComfyUI-WorkflowGenerator.wiki/ + ``` + +3. **Commit and push:** + ```bash + cd ComfyUI-WorkflowGenerator.wiki + git add *.md + git commit -m "Add wiki pages" + git push origin master + ``` + +## Page Structure + +The wiki is organized as follows: + +- **Home.md** - Main landing page with navigation +- **Installation.md** - Detailed installation instructions +- **Node-Reference.md** - Complete node documentation +- **Configuration.md** - Configuration options and performance tuning +- **Advanced-Usage.md** - Advanced patterns and techniques +- **Troubleshooting.md** - Common issues and solutions +- **Instruction-Prompt-Usage.md** - How instructions work across pipeline steps + +## Linking Between Pages + +Wiki pages use simple markdown links without file extensions: + +```markdown +[Link Text](Page-Name) +``` + +For example: +- `[Home](Home)` links to the Home page +- `[Installation](Installation)` links to the Installation page +- `[Node Reference](Node-Reference)` links to the Node-Reference page + +## Updating the README + +After uploading the wiki, update the main repository README to link to the wiki: + +```markdown +For detailed documentation, see the [Wiki](https://github.com/danielpflorian/ComfyUI-WorkflowGenerator/wiki). +``` + +## Notes + +- GitHub Wikis automatically create a sidebar with all pages +- The first page you create becomes the home page (or you can name it `Home.md`) +- Page names are case-sensitive and spaces should be replaced with hyphens +- You can add a custom sidebar by creating a `_Sidebar.md` file + diff --git a/wiki/Troubleshooting.md b/wiki/Troubleshooting.md new file mode 100644 index 0000000..0b70ef9 --- /dev/null +++ b/wiki/Troubleshooting.md @@ -0,0 +1,110 @@ +# Troubleshooting + +Common issues, solutions, and frequently asked questions. + +## Common Issues + +### Models Not Found + +**Symptoms**: Error message about model file not found + +**Solutions**: +- Verify models are in `ComfyUI/models/LLM/` directory +- Check model paths are correct (use absolute paths if needed) +- Ensure tokenizer directories exist and contain required files +- Restart ComfyUI after moving files + +### Out of Memory (OOM) Errors + +**Symptoms**: CUDA out of memory errors, crashes during model loading + +**Solutions**: +- Reduce `n_gpu_layers` (try "auto" or lower number) +- Use smaller GGUF models (q8_0 or q4_0 quantization) +- Set `device_preference` to "cpu" for some processing stages (e.g., NodeValidator semantic search) +- Close other applications using GPU memory +- Use CPU-only mode if GPU memory is insufficient + +### Slow Performance + +**Symptoms**: Workflow generation takes a very long time + +**Solutions**: +- Use GGUF models (lower VRAM usage) +- Set `use_llm_refinement=False` (semantic search only) +- Increase `n_gpu_layers` if VRAM available +- Use smaller models if accuracy allows +- Ensure CUDA/Metal acceleration is properly installed + +### Invalid Node Names + +**Symptoms**: Errors about node types not found, incorrect node names in workflow + +**Solutions**: +- Run **Update Node Catalog** node to refresh database +- Enable `use_llm_refinement=True` for better accuracy +- Check that custom nodes are properly installed +- Verify catalog files are up to date + +### Import Errors + +**Symptoms**: Python import errors, missing modules + +**Solutions**: +- Verify dependencies: `pip install -r requirements.txt` +- Ensure `llama-cpp-python` is installed (see [Installation](Installation)) +- Check Python version (3.10+ required) +- Restart ComfyUI after installing dependencies +- Verify you're using the correct Python environment + +### Tokenizer Not Found + +**Symptoms**: Errors about missing tokenizer files + +**Solutions**: +- Ensure tokenizer directory exists with the correct name +- Verify tokenizer files (`tokenizer.json`, `tokenizer_config.json`, `vocab.json`) are present +- Check directory structure matches expected format +- Tokenizer directory name must match model name (without `.gguf` and quantization suffix) + +## Limitations & Prompting Tips + +**"Why isn't it generating the perfect workflow?"** + +It's important to understand that the system is not magic. It relies on a fine-tuned language model that has seen a specific set of training data (around 13,000 workflows). + +### Limitations +- **Training Data Bias**: The model is biased towards the workflows it was trained on (mostly SD 1.5, SDXL, SVD, etc.). It may struggle with very new nodes or niche custom node packs released after the dataset was created. +- **Hallucination**: Like all LLMs, it can hallucinate connections that don't exist or propose nodes that are similar to but not exactly what you have installed. +- **Complexity**: Extremely complex workflows with circular dependencies or very specific control flows may be simplified or incorrect. + +### Tips for Better Results +- **Be Specific**: Instead of "make an image", try "create a text to image workflow using SDXL with a refiner and simple styles". +- **Re-roll**: If the output looks wrong, try changing the seed or slightly rewording the prompt. +- **Step-by-Step**: Use the individual nodes. Generate the diagram first, inspect it, then validate. +- **Check Node Names**: If it's picking the wrong nodes, run the **Update Node Catalog** node to ensure it knows about your currently installed nodes. + +## FAQ + +**Q: Can I use HuggingFace models instead of GGUF?** +A: Yes, but they use more VRAM and may be slower. GGUF models are recommended for better performance. + +**Q: Do I need the NodeValidator model?** +A: Only if you want to use LLM refinement. Semantic search works without it but is less accurate. + +**Q: How often should I update the node catalog?** +A: After installing new custom nodes or updating ComfyUI. The catalog is cached for performance. + +**Q: Can I modify the generated workflow?** +A: Yes, the workflow JSON can be edited manually or loaded into ComfyUI for further modification. + +**Q: Why is my workflow generation slow?** +A: Check that you're using GGUF models, have GPU acceleration enabled, and LLM refinement is disabled if speed is priority. + +**Q: The quick install for llama-cpp-python didn't work. What should I do?** +A: Local compilation is often necessary, especially on Windows. See [Installation](Installation) for detailed compilation instructions. + +--- + +[← Back to Home](Home) | [← Configuration](Configuration) + diff --git a/workflows/WG_3-Chained-Nodes.json b/workflows/WG_3-Chained-Nodes.json new file mode 100644 index 0000000..ed3b6a6 --- /dev/null +++ b/workflows/WG_3-Chained-Nodes.json @@ -0,0 +1,319 @@ +{ + "id": "4a3dfff0-29c9-4460-aa79-fd3215d8fa87", + "revision": 0, + "last_node_id": 16, + "last_link_id": 10, + "nodes": [ + { + "id": 12, + "type": "NodeValidator", + "pos": [ + -853.6131510416667, + -580 + ], + "size": [ + 400, + 608 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "workflow_edges", + "type": "STRING", + "link": 7 + }, + { + "name": "instruction", + "type": "STRING", + "widget": { + "name": "instruction" + }, + "link": 9 + } + ], + "outputs": [ + { + "name": "workflow_edges (refined)", + "type": "STRING", + "links": [ + 8 + ] + }, + { + "name": "llm_prompts (debug)", + "type": "STRING", + "links": null + }, + { + "name": "candidate_nodes (debug)", + "type": "STRING", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "Node name for S&R": "NodeValidator", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.5.1" + } + }, + "widgets_values": [ + "", + true, + "Qwen2.5-7B-Instruct-q8_0.gguf", + "paraphrase-multilingual-MiniLM-L12-v2", + "catalog", + "auto", + "auto", + "auto", + true, + -1, + 4096, + 4096, + true, + false, + 512, + "Auto", + 0.95, + 5, + 0.7, + 3771414511, + "randomize" + ] + }, + { + "id": 13, + "type": "WorkflowBuilder", + "pos": [ + -417.2263020833334, + -580 + ], + "size": [ + 270.839453125, + 126 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "workflow_edges", + "type": "STRING", + "link": 8 + } + ], + "outputs": [ + { + "name": "workflow_json", + "type": "STRING", + "links": [] + }, + { + "name": "file_path", + "type": "STRING", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "Node name for S&R": "WorkflowBuilder", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.5.1" + } + }, + "widgets_values": [ + "catalog", + true, + "generated_workflow" + ] + }, + { + "id": 16, + "type": "Note", + "pos": [ + -110.00000000000011, + -580 + ], + "size": [ + 290, + 120 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": { + "ue_properties": { + "widget_ue_connectable": {}, + "version": "7.5.1", + "input_ue_unconnectable": {} + } + }, + "widgets_values": [ + "You don't have to connect anything to the Build and Save Workflow output, but can use a \"show text/string\" type of node here from another node package." + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 14, + "type": "Note", + "pos": [ + -850, + -810 + ], + "size": [ + 400, + 170 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": { + "ue_properties": { + "widget_ue_connectable": {}, + "version": "7.5.1", + "input_ue_unconnectable": {} + } + }, + "widgets_values": [ + "NodeValidator is optional but recommended. You do not have to use LLM Refinement, but it can help if the embedding model is not sufficient to fix edge inconsitencies in the output from WorkflowGenerator.\n\nYou can also either input the same instructions from WorkflowGenerator, add different instructions, or leave instructions blank, for the LLM model. These are only used if use_llm_refinement is enabled. I have not tested the affect of using different instructions thoroughly, so you will have to experiment on your own.\n\n" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 11, + "type": "WorkflowGenerator", + "pos": [ + -1290, + -580 + ], + "size": [ + 400, + 516 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "workflow_edges", + "type": "STRING", + "links": [ + 7 + ] + }, + { + "name": "instruction", + "type": "STRING", + "links": [ + 9 + ] + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "Node name for S&R": "WorkflowGenerator", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.5.1" + } + }, + "widgets_values": [ + "Create a flux1 image to image workflow with ultimate upscale support", + "workflow-generator-q8_0.gguf", + "auto", + "auto", + "auto", + true, + -1, + 8192, + 4096, + true, + false, + 512, + "Auto", + 0.95, + 0.7, + false, + 2818631848, + "randomize" + ] + } + ], + "links": [ + [ + 7, + 11, + 0, + 12, + 0, + "STRING" + ], + [ + 8, + 12, + 0, + 13, + 0, + "STRING" + ], + [ + 9, + 11, + 1, + 12, + 1, + "STRING" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ue_links": [], + "links_added_by_ue": [], + "ds": { + "scale": 0.9646149645000008, + "offset": [ + 1631.660607951843, + 853.986070950943 + ] + }, + "frontendVersion": "1.28.8", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/workflows/WG_Pipeline.json b/workflows/WG_Pipeline.json new file mode 100644 index 0000000..a084988 --- /dev/null +++ b/workflows/WG_Pipeline.json @@ -0,0 +1,143 @@ +{ + "id": "4a3dfff0-29c9-4460-aa79-fd3215d8fa87", + "revision": 0, + "last_node_id": 10, + "last_link_id": 4, + "nodes": [ + { + "id": 10, + "type": "WG_Pipeline", + "pos": [ + -1130, + -1020 + ], + "size": [ + 400, + 868 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "workflow_edges", + "type": "STRING", + "links": null + }, + { + "name": "workflow_edges (refined)", + "type": "STRING", + "links": null + }, + { + "name": "workflow_json", + "type": "STRING", + "links": null + }, + { + "name": "file_path", + "type": "STRING", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "Node name for S&R": "WG_Pipeline", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.5.1" + } + }, + "widgets_values": [ + "Create an SDXL text-to-image workflow with lora and depth controlnet support", + "workflow-generator-q8_0.gguf", + "auto", + "auto", + "auto", + true, + -1, + 8192, + 4096, + 0.95, + 0.7, + false, + true, + false, + "Qwen2.5-7B-Instruct-q8_0.gguf", + "paraphrase-multilingual-MiniLM-L12-v2", + "catalog", + "auto", + "auto", + "auto", + false, + -1, + 4096, + 4096, + 0.95, + 5, + 0.7, + true, + "generated_workflow", + 0, + "randomize" + ] + }, + { + "id": 4, + "type": "Note", + "pos": [ + -1600, + -1020 + ], + "size": [ + 430, + 290 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": { + "ue_properties": { + "widget_ue_connectable": {}, + "version": "7.5.1", + "input_ue_unconnectable": {} + } + }, + "widgets_values": [ + "### Workflow Generator Guide\n\n**Instruction Tips:**\n- Be specific: \"Create an SDXL text-to-image workflow with LoRA support\"\n- Mention key nodes: \"Use ControlNet with Canny preprocessor\"\n\n**Default Settings:**\n- `auto_gpu_layers` is ON by default to prevent Out-Of-Memory errors.\n\n**How to Use:**\n1. **Instruction:** Enter your workflow description.\n2. **Refine:** Keep OFF initially for speed. Turn ON (`use_llm_refinement`) only if node names are incorrect.\n3. **Output:** The workflow JSON is saved to `ComfyUI/output/`.\n\n**Models Required:**\n- Main: `workflow-generator-q8_0.gguf`\n- Embedding: `paraphrase-multilingual-MiniLM-L12-v2`\n- Refine: `Qwen2.5-7B-Instruct` (Optional)\n\nYou do not have to connect outputs to a show text/string node but can if want. " + ], + "color": "#432", + "bgcolor": "#653" + } + ], + "links": [], + "groups": [], + "config": {}, + "extra": { + "ue_links": [], + "links_added_by_ue": [], + "ds": { + "scale": 0.8769226950000011, + "offset": [ + 2178.610894239918, + 1170.963627634079 + ] + }, + "frontendVersion": "1.28.8", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/workflows/WG_UpdateNodeCatalog.json b/workflows/WG_UpdateNodeCatalog.json new file mode 100644 index 0000000..cf758a8 --- /dev/null +++ b/workflows/WG_UpdateNodeCatalog.json @@ -0,0 +1,113 @@ +{ + "id": "b8a29305-924f-4652-a6d5-6985ab095086", + "revision": 0, + "last_node_id": 33, + "last_link_id": 51, + "nodes": [ + { + "id": 31, + "type": "WG_UpdateNodeCatalog", + "pos": [ + -1370, + -1160 + ], + "size": [ + 270, + 102 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "status", + "type": "STRING", + "links": null + }, + { + "name": "node_count", + "type": "INT", + "links": null + } + ], + "properties": { + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65, + "Node name for S&R": "WG_UpdateNodeCatalog", + "ue_properties": { + "widget_ue_connectable": {}, + "input_ue_unconnectable": {}, + "version": "7.5.1" + } + }, + "widgets_values": [ + "catalog", + true + ] + }, + { + "id": 33, + "type": "Note", + "pos": [ + -1420, + -1400 + ], + "size": [ + 380, + 170 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": { + "ue_properties": { + "widget_ue_connectable": {}, + "version": "7.5.1", + "input_ue_unconnectable": {} + } + }, + "widgets_values": [ + "Run this first to build a database of your current nodes. You don't have to connect the outputs to anything unless you want to see status messages.\n\nBy default the catalog is saved to:\n\"comfyUI\\custom_nodes\\ComfyUI-WorkflowGenerator\\catalog\"\n\nThis folder container 2 files that are built automatically:\nnode_info.json\nnode_list.json" + ], + "color": "#432", + "bgcolor": "#653" + } + ], + "links": [], + "groups": [], + "config": {}, + "extra": { + "extra": { + "ds": { + "scale": 0.952486143572623, + "offset": [ + 163.72358533046707, + 22.018287636666063 + ] + } + }, + "ue_links": [], + "ds": { + "scale": 1.4122927695244514, + "offset": [ + 1942.8614244792802, + 1565.3844711413715 + ] + }, + "links_added_by_ue": [], + "frontendVersion": "1.28.8", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file