last changes
This commit is contained in:
kmlbdh
@@ -2,16 +2,18 @@
|
||||
|
||||
A custom node pack for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) that allows you to run Large Language Models (LLMs) locally and use them for prompt generation and other text tasks directly within your ComfyUI workflows.
|
||||
|
||||
This pack provides nodes to connect to and utilize local LLMs (like Llama, Phi, Gemma, etc., in Hugging Face format) without needing external API calls. It's designed to integrate seamlessly with prompt generation workflows, such as those involving image description nodes like Florence-2.
|
||||
This pack provides nodes to connect to and utilize local LLMs in Hugging Face (PyTorch) or GGUF format, eliminating the need for external API calls. It's designed to integrate seamlessly with prompt generation workflows, such as those involving image description nodes like Florence-2.
|
||||
|
||||
## Features
|
||||
|
||||
* **Local LLM Execution:** Run powerful LLMs directly on your machine (CPU or GPU).
|
||||
* **Set Local LLM Service Connector Node:** Select and configure your local LLM model (models must be placed in `ComfyUI/models/LLM/`).
|
||||
* **Local Kontext Prompt Generator Node:** Generate detailed image prompts by combining descriptions and edit instructions, leveraging your local LLM.
|
||||
* **Local LLM Execution (Hugging Face & GGUF):** Run powerful LLMs directly on your machine (CPU or GPU) using either standard Hugging Face models or efficient GGUF models.
|
||||
* **Set Local LLM Service Connector Node (Hugging Face):** Select and configure your local Hugging Face format LLM model (models must be placed in `ComfyUI/models/LLM/`).
|
||||
* **Set Local GGUF LLM Service Connector Node:** Select and configure your local GGUF format LLM model file (`.gguf` files must be placed in `ComfyUI/models/LLM/`).
|
||||
* **Local Kontext Prompt Generator Node:** Generate detailed image prompts by combining descriptions and edit instructions, leveraging your connected local LLM.
|
||||
* **User Preset Management:** Add and remove custom prompt generation presets using dedicated nodes.
|
||||
* **Compatibility:** Designed to work with standard MieNodes prompt generators (e.g., `KontextPromptGenerator`) if needed, using the `LLMServiceConnector` type identifier.
|
||||
* **VRAM Optimization Ready:** Includes commented code examples for integrating quantization (4-bit/8-bit) using `bitsandbytes` to reduce memory footprint for running alongside large image models like Flux.
|
||||
* **Compatibility:** Designed to work with the `LLMServiceConnector` type identifier, ensuring compatibility with standard MieNodes prompt generators (e.g., `KontextPromptGenerator`) if needed.
|
||||
* **VRAM Optimization (Hugging Face):** Includes commented code examples for integrating Hugging Face model quantization (4-bit/8-bit) using `bitsandbytes` to reduce memory footprint.
|
||||
* **Efficient GGUF Models:** GGUF models are inherently quantized, offering lower memory usage and often good performance, especially on CPU.
|
||||
|
||||
## Installation
|
||||
|
||||
@@ -26,43 +28,53 @@ This pack provides nodes to connect to and utilize local LLMs (like Llama, Phi,
|
||||
```bash
|
||||
cd ComfyUI_LocalLLMNodes
|
||||
pip install -r requirements.txt
|
||||
# Or install directly:
|
||||
# pip install transformers torch
|
||||
# Optional for quantization: pip install bitsandbytes
|
||||
```
|
||||
*Note: Ensure you are installing these packages in the same Python environment that you use to run ComfyUI.*
|
||||
*Note on `llama-cpp-python`: Installing with GPU support (CUDA) requires specific environment variables during installation (e.g., `CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python --force-reinstall --no-cache-dir`). See the `llama-cpp-python` documentation for details. The CPU-only installation (`pip install llama-cpp-python`) is simpler and sufficient for CPU inference.
|
||||
|
||||
## Usage
|
||||
|
||||
1. **Download a Local LLM:**
|
||||
* Obtain a Hugging Face format LLM (e.g., `TinyLlama/TinyLlama-1.1B-Chat-v1.0`, `microsoft/Phi-3-mini-4k-instruct`, `google/gemma-2b-it`).
|
||||
* Download the model files into a subdirectory within your `ComfyUI/models/LLM/` folder.
|
||||
* Example: `ComfyUI/models/LLM/Phi-3-mini-4k-instruct/` should contain `config.json`, `pytorch_model.bin` (or `.safetensors`), `tokenizer_config.json`, etc.
|
||||
* **Hugging Face Format:**
|
||||
* Obtain a Hugging Face format LLM (e.g., `TinyLlama/TinyLlama-1.1B-Chat-v1.0`, `microsoft/Phi-3-mini-4k-instruct`, `NousResearch/Hermes-2-Pro-Llama-3-8B`).
|
||||
* Download the model files into a subdirectory within your `ComfyUI/models/LLM/` folder.
|
||||
* Example: `ComfyUI/models/LLM/Phi-3-mini-4k-instruct/` should contain `config.json`, `pytorch_model.bin` (or `.safetensors`), `tokenizer_config.json`, etc.
|
||||
* **GGUF Format:**
|
||||
* Obtain a GGUF format LLM file (e.g., `mistral-7b-instruct-v0.3.Q8_0.gguf`).
|
||||
* Place the `.gguf` file directly in your `ComfyUI/models/LLM/` folder, or within a subfolder (e.g., `ComfyUI/models/LLM/Mistral-7B-Instruct/` containing `mistral-7b-instruct-v0.3.Q8_0.gguf`).
|
||||
2. **Restart ComfyUI** to load the new nodes.
|
||||
3. **Find the Nodes:** Look for the new nodes in the ComfyUI node library under the categories:
|
||||
3. **Find the Nodes:** Look for the new nodes in the ComfyUI node library under the category:
|
||||
* `Local LLM Nodes/LLM Connectors`
|
||||
* `Local LLM Nodes/Prompt Generators`
|
||||
4. **Use the Nodes:**
|
||||
* Add the **"Set Local LLM Service Connector 🐑"** node to your graph.
|
||||
* Select your downloaded local LLM model from the dropdown menu.
|
||||
* Add the **"Local Kontext Prompt Generator 🐑"** node.
|
||||
* Connect the output of the "Set Local LLM Service Connector 🐑" node to the `llm_service_connector` input of the "Local Kontext Prompt Generator 🐑" node.
|
||||
* Provide inputs like `image1_description` (e.g., from Florence-2), `edit_instruction`, and select a `preset`.
|
||||
* Connect the `kontext_prompt` output to your desired node (e.g., an image generator).
|
||||
* **For Hugging Face Models:**
|
||||
* Add the **"Set Local LLM Service Connector 🐑 (HuggingFace)"** node to your graph.
|
||||
* Select your downloaded Hugging Face model directory from the dropdown menu.
|
||||
* **For GGUF Models:**
|
||||
* Add the **"Set Local GGUF LLM Service Connector 🐑"** node to your graph.
|
||||
* Select your downloaded GGUF model file (or its containing directory) from the dropdown menu.
|
||||
* **Common Steps:**
|
||||
* Add the **"Local Kontext Prompt Generator 🐑"** node.
|
||||
* Connect the output of the chosen "Set Local ... LLM Service Connector 🐑" node to the `llm_service_connector` input of the "Local Kontext Prompt Generator 🐑" node.
|
||||
* Provide inputs like `image1_description` (e.g., from Florence-2), `edit_instruction`, and select a `preset`.
|
||||
* Connect the `kontext_prompt` output to your desired node (e.g., an image generator).
|
||||
|
||||
## Memory Optimization (VRAM)
|
||||
## Memory Optimization
|
||||
|
||||
Running large LLMs alongside large image models (like SDXL or Flux) can strain GPU memory (VRAM).
|
||||
Running large LLMs alongside large image models (like SDXL or Flux) can strain system resources (RAM/VRAM).
|
||||
|
||||
* **Quantization:** The `local_llm_connector.py` file includes commented code examples showing how to implement 4-bit or 8-bit quantization using the `bitsandbytes` library. This can significantly reduce the LLM's VRAM usage.
|
||||
* To use quantization:
|
||||
1. Ensure `bitsandbytes` is installed (`pip install bitsandbytes`).
|
||||
2. Uncomment and adjust the quantization configuration section in the `_load_model` method within `local_llm_connector.py`.
|
||||
3. Restart ComfyUI.
|
||||
* **Hugging Face Models:**
|
||||
* **Quantization:** The `local_llm_connector.py` file includes commented code examples showing how to implement 4-bit or 8-bit quantization using the `bitsandbytes` library. This can significantly reduce the LLM's VRAM usage.
|
||||
* To use quantization:
|
||||
1. Ensure `bitsandbytes` is installed (`pip install bitsandbytes`).
|
||||
2. Uncomment and adjust the quantization configuration section in the `_load_model` method within `local_llm_connector.py`.
|
||||
3. Restart ComfyUI.
|
||||
* **GGUF Models:**
|
||||
* GGUF models are pre-quantized (e.g., Q4, Q5, Q8). Choosing a more quantized version (like Q8_0 vs. f16) inherently uses less memory. For GPU acceleration with GGUF models, configure the `n_gpu_layers` parameter during loading (if supported by your `llama-cpp-python` build).
|
||||
|
||||
## Nodes Included
|
||||
|
||||
* `SetLocalLLMServiceConnector`: Selects and prepares a connection to a local LLM model.
|
||||
* `SetLocalLLMServiceConnector`: Selects and prepares a connection to a local Hugging Face format LLM model.
|
||||
* `SetLocalGGUFLLMServiceConnector`: Selects and prepares a connection to a local GGUF format LLM model file.
|
||||
* `LocalKontextPromptGenerator`: Generates prompts using a connected local LLM based on descriptions and instructions.
|
||||
* `AddUserLocalKontextPreset`: Adds a custom preset for prompt generation.
|
||||
* `RemoveUserLocalKontextPreset`: Removes a custom preset.
|
||||
@@ -70,12 +82,13 @@ Running large LLMs alongside large image models (like SDXL or Flux) can strain G
|
||||
## Requirements
|
||||
|
||||
* [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
|
||||
* Python Libraries:
|
||||
* Python Libraries (see `requirements.txt` for versions):
|
||||
* `transformers`
|
||||
* `torch`
|
||||
* `bitsandbytes` (Optional, for quantization)
|
||||
* (See `requirements.txt`)
|
||||
* `bitsandbytes` (Optional, for Hugging Face model quantization)
|
||||
* `llama-cpp-python` (Optional, for GGUF model support)
|
||||
* Other dependencies as listed in `requirements.txt`
|
||||
|
||||
## Acknowledgements
|
||||
|
||||
This node pack builds upon concepts and structures found in the excellent [ComfyUI-MieNodes](https://github.com/MieMieeeee/ComfyUI-MieNodes) extension, particularly the `KontextPromptGenerator` and LLM service connector patterns.
|
||||
This node pack builds upon concepts and structures found in the excellent [ComfyUI-MieNodes](https://github.com/MieMieeeee/ComfyUI-MieNodes) extension, particularly the `KontextPromptGenerator` and LLM service connector patterns.
|
||||
+57
-26
@@ -1,8 +1,11 @@
|
||||
# ComfyUI/custom_nodes/ComfyUI_LocalLLMNodes/__init__.py
|
||||
|
||||
# --- Conditional Import Logic ---
|
||||
LOCAL_LLM_NODES_AVAILABLE = False
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
LOCAL_GGUF_NODES_AVAILABLE = False # Define this early
|
||||
|
||||
NODE_CLASS_MAPPINGS = {} # Initialize early
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {} # Initialize early
|
||||
|
||||
# --- Check for core dependencies ---
|
||||
# We need transformers and torch for the LLM nodes to work.
|
||||
@@ -16,6 +19,15 @@ except ImportError as e:
|
||||
print(f"[ComfyUI_LocalLLMNodes] Core dependencies (transformers, torch) not found or not importable: {e}")
|
||||
CORE_DEPS_AVAILABLE = False
|
||||
|
||||
# --- Check for GGUF dependency ---
|
||||
try:
|
||||
import llama_cpp
|
||||
GGUF_DEPS_AVAILABLE = True
|
||||
except ImportError as e:
|
||||
print(f"[ComfyUI_LocalLLMNodes] GGUF dependency (llama-cpp-python) not found: {e}")
|
||||
GGUF_DEPS_AVAILABLE = False
|
||||
|
||||
# --- Attempt to import Hugging Face based nodes ---
|
||||
if CORE_DEPS_AVAILABLE:
|
||||
try:
|
||||
# --- Import Node Classes ---
|
||||
@@ -25,33 +37,52 @@ if CORE_DEPS_AVAILABLE:
|
||||
AddUserLocalKontextPreset,
|
||||
RemoveUserLocalKontextPreset
|
||||
)
|
||||
|
||||
# --- Define Mappings ---
|
||||
# Using the class names directly as keys is standard.
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SetLocalLLMServiceConnector": SetLocalLLMServiceConnector,
|
||||
"LocalKontextPromptGenerator": LocalKontextPromptGenerator,
|
||||
"AddUserLocalKontextPreset": AddUserLocalKontextPreset,
|
||||
"RemoveUserLocalKontextPreset": RemoveUserLocalKontextPreset,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SetLocalLLMServiceConnector": "Set Local LLM Service Connector 🐑",
|
||||
"LocalKontextPromptGenerator": "Local Kontext Prompt Generator 🐑",
|
||||
"AddUserLocalKontextPreset": "Add User Local Kontext Preset 🐑",
|
||||
"RemoveUserLocalKontextPreset": "Remove User Local Kontext Preset 🐑",
|
||||
}
|
||||
|
||||
LOCAL_LLM_NODES_AVAILABLE = True
|
||||
print("[ComfyUI_LocalLLMNodes] All nodes are available.")
|
||||
|
||||
print("[ComfyUI_LocalLLMNodes] Hugging Face based LLM nodes are available.")
|
||||
except Exception as e:
|
||||
print(f"[ComfyUI_LocalLLMNodes] Error importing node classes: {e}")
|
||||
# If import fails, NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS remain empty,
|
||||
# and ComfyUI won't register any nodes from this package.
|
||||
print(f"[ComfyUI_LocalLLMNodes] Error importing transformers-based node classes: {e}")
|
||||
else:
|
||||
print("[ComfyUI_LocalLLMNodes] Nodes will NOT be available due to missing core dependencies.")
|
||||
print("[ComfyUI_LocalLLMNodes] Hugging Face based LLM nodes will NOT be available due to missing core dependencies (transformers, torch).")
|
||||
|
||||
# --- Attempt to import GGUF based nodes ---
|
||||
if GGUF_DEPS_AVAILABLE:
|
||||
try:
|
||||
from .local_gguf_llm_connector import SetLocalGGUFLLMServiceConnector
|
||||
LOCAL_GGUF_NODES_AVAILABLE = True
|
||||
print("[ComfyUI_LocalLLMNodes] GGUF LLM Connector node is available.")
|
||||
except Exception as e:
|
||||
print(f"[ComfyUI_LocalLLMNodes] Error importing GGUF node classes: {e}")
|
||||
else:
|
||||
print("[ComfyUI_LocalLLMNodes] GGUF LLM Connector node will NOT be available due to missing dependency (llama-cpp-python).")
|
||||
|
||||
# --- Define Mappings ---
|
||||
# Populate the dictionaries based on which nodes were successfully imported.
|
||||
|
||||
# Add Hugging Face based nodes if available
|
||||
if LOCAL_LLM_NODES_AVAILABLE:
|
||||
NODE_CLASS_MAPPINGS.update({
|
||||
"SetLocalLLMServiceConnector": SetLocalLLMServiceConnector,
|
||||
"LocalKontextPromptGenerator": LocalKontextPromptGenerator,
|
||||
"AddUserLocalKontextPreset": AddUserLocalKontextPreset,
|
||||
"RemoveUserLocalKontextPreset": RemoveUserLocalKontextPreset,
|
||||
})
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update({
|
||||
"SetLocalLLMServiceConnector": "Set Local LLM Service Connector 🐑 (HuggingFace)",
|
||||
"LocalKontextPromptGenerator": "Local Kontext Prompt Generator 🐑",
|
||||
"AddUserLocalKontextPreset": "Add User Local Kontext Preset 🐑",
|
||||
"RemoveUserLocalKontextPreset": "Remove User Local Kontext Preset 🐑",
|
||||
})
|
||||
|
||||
# Add GGUF based nodes if available
|
||||
if LOCAL_GGUF_NODES_AVAILABLE:
|
||||
NODE_CLASS_MAPPINGS.update({
|
||||
"SetLocalGGUFLLMServiceConnector": SetLocalGGUFLLMServiceConnector,
|
||||
})
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update({
|
||||
"SetLocalGGUFLLMServiceConnector": "Set Local GGUF LLM Service Connector 🐑",
|
||||
})
|
||||
|
||||
# --- Define what ComfyUI sees ---
|
||||
# ComfyUI looks for these specific dictionaries in __init__.py
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
# They must be defined at the top level of the module.
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,210 @@
|
||||
# ComfyUI/custom_nodes/ComfyUI_LocalLLMNodes/local_gguf_llm_connector.py
|
||||
import os
|
||||
import folder_paths
|
||||
# --- Library for GGUF local LLM inference ---
|
||||
try:
|
||||
import llama_cpp
|
||||
# import llama_cpp.llama_tokenizer # Optional, remove if not used
|
||||
LLAMA_CPP_AVAILABLE = True
|
||||
except ImportError:
|
||||
# --- Simple logging utility for this package (reuse existing one) ---
|
||||
# Import here to avoid potential issues if log isn't defined yet if imported at top in case of circular import risks,
|
||||
# though usually safe. Ensure local_llm_connector.py defines 'log' early.
|
||||
from .local_llm_connector import log
|
||||
log("[LocalGGUFLLMConnector] Warning: llama-cpp-python library not found.")
|
||||
LLAMA_CPP_AVAILABLE = False
|
||||
llama_cpp = None
|
||||
|
||||
# Reuse the logging function from local_llm_connector
|
||||
from .local_llm_connector import log # Ensure this is the final import for log
|
||||
|
||||
LOCAL_LLM_CATEGORY = "Local LLM Nodes/LLM Connectors" # Reuse or define new sub-category like "Local LLM Nodes/GGUF Connectors"
|
||||
|
||||
def get_local_gguf_model_names():
|
||||
"""Discovers .gguf files or directories containing .gguf files within models/LLM."""
|
||||
llm_models_dir = os.path.join(folder_paths.models_dir, "LLM")
|
||||
model_names = []
|
||||
|
||||
if os.path.exists(llm_models_dir):
|
||||
try:
|
||||
for item in os.listdir(llm_models_dir):
|
||||
item_path = os.path.join(llm_models_dir, item)
|
||||
if os.path.isfile(item_path) and item.endswith('.gguf'):
|
||||
# Add the filename without .gguf extension
|
||||
model_names.append(os.path.splitext(item)[0])
|
||||
elif os.path.isdir(item_path):
|
||||
# Check if directory contains a .gguf file
|
||||
for subitem in os.listdir(item_path):
|
||||
if subitem.endswith('.gguf'):
|
||||
model_names.append(item)
|
||||
break # Found one, add the directory name and move to next item
|
||||
except Exception as e:
|
||||
log(f"[LocalGGUFLLMConnector] Error scanning models/LLM directory: {e}")
|
||||
else:
|
||||
log(f"[LocalGGUFLLMConnector] models/LLM directory not found: {llm_models_dir}")
|
||||
|
||||
if not model_names:
|
||||
model_names = ["No_Local_GGUF_Models_Found"]
|
||||
return model_names
|
||||
|
||||
class SetLocalGGUFLLMServiceConnector:
|
||||
"""
|
||||
A node to select and prepare a connection to a local GGUF LLM model.
|
||||
Models should be placed in ComfyUI/models/LLM/your_model.gguf or ComfyUI/models/LLM/your_model_folder/your_model.gguf.
|
||||
Requires 'llama-cpp-python': pip install llama-cpp-python (consider CUDA flags for GPU support)
|
||||
"""
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
model_names = get_local_gguf_model_names()
|
||||
return {
|
||||
"required": {
|
||||
"local_gguf_model_name": (model_names, {"default": model_names[0] if model_names else "No_Local_GGUF_Models_Found"}),
|
||||
# Optional parameters for llama-cpp-python loading
|
||||
# "n_ctx": ("INT", {"default": 4096, "min": 1, "max": 100000}),
|
||||
# "n_gpu_layers": ("INT", {"default": 0, "min": -1, "max": 100}), # -1 = all
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LLMServiceConnector",) # Use the same type for compatibility
|
||||
FUNCTION = "get_connector"
|
||||
CATEGORY = LOCAL_LLM_CATEGORY
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def get_connector(self, local_gguf_model_name): # , n_ctx=4096, n_gpu_layers=0):
|
||||
"""Returns a connector object for the selected GGUF model."""
|
||||
if not LLAMA_CPP_AVAILABLE:
|
||||
raise Exception("The 'llama-cpp-python' library is required for the Local GGUF LLM node but is not installed.")
|
||||
|
||||
if local_gguf_model_name == "No_Local_GGUF_Models_Found":
|
||||
raise Exception("No local GGUF models found in models/LLM directory. Please place your .gguf files there.")
|
||||
|
||||
# Determine the base path selected by the user
|
||||
base_models_dir = os.path.join(folder_paths.models_dir, "LLM")
|
||||
model_path = os.path.join(base_models_dir, local_gguf_model_name)
|
||||
|
||||
gguf_file_path = None
|
||||
# Case 1: User selected a name derived from a file directly in models/LLM (e.g., 'model' for 'models/LLM/model.gguf')
|
||||
# model_path would be 'models/LLM/model' (not a .gguf file itself)
|
||||
if os.path.isfile(model_path) and model_path.endswith('.gguf'):
|
||||
gguf_file_path = model_path
|
||||
# Case 2: User selected a name derived from a directory (e.g., 'ModelDir' for 'models/LLM/ModelDir/')
|
||||
# model_path would be 'models/LLM/ModelDir' (a directory)
|
||||
elif os.path.isdir(model_path):
|
||||
# Search for the .gguf file inside the directory
|
||||
for file in os.listdir(model_path):
|
||||
if file.endswith('.gguf'):
|
||||
gguf_file_path = os.path.join(model_path, file)
|
||||
break # Use the first .gguf file found
|
||||
|
||||
# Final check: Ensure we found a valid .gguf file path
|
||||
if not gguf_file_path or not os.path.exists(gguf_file_path):
|
||||
raise FileNotFoundError(f"[LocalGGUFLLMConnector] GGUF model file not found for selection: '{local_gguf_model_name}'. Expected file: '{gguf_file_path}'")
|
||||
|
||||
# Create and return the GGUF connector instance, passing the resolved .gguf file path
|
||||
connector = LocalGGUFLLMServiceConnector(gguf_file_path) # Pass kwargs like n_ctx, n_gpu_layers if added
|
||||
return (connector,)
|
||||
|
||||
|
||||
class LocalGGUFLLMServiceConnector:
|
||||
"""
|
||||
Represents the connection to a specific local GGUF LLM using llama-cpp-python.
|
||||
The `gguf_file_path` should point directly to the .gguf file.
|
||||
"""
|
||||
def __init__(self, gguf_file_path): # , n_ctx=4096, n_gpu_layers=0):
|
||||
self.gguf_file_path = gguf_file_path
|
||||
self.model = None
|
||||
self.is_loaded = False
|
||||
# self.n_ctx = n_ctx
|
||||
# self.n_gpu_layers = n_gpu_layers
|
||||
|
||||
def _load_model(self):
|
||||
"""Loads the GGUF model using llama-cpp-python."""
|
||||
if self.is_loaded:
|
||||
return # Already loaded
|
||||
|
||||
try:
|
||||
log(f"[LocalGGUFLLMConnector] Loading GGUF model from: {self.gguf_file_path}")
|
||||
# --- Model Loading Configuration for llama-cpp-python ---
|
||||
# Adjust these parameters based on your needs and system capabilities.
|
||||
model_kwargs = {
|
||||
"n_ctx": 4096, # Context window size (adjust if needed, larger uses more memory)
|
||||
"n_threads": 8, # Number of CPU threads to use
|
||||
# "n_threads_batch": 8, # For batch processing (if applicable)
|
||||
# --- GPU Acceleration (if llama-cpp-python was built with CUDA support) ---
|
||||
# "n_gpu_layers": self.n_gpu_layers, # Number of layers to offload to GPU (e.g., 33 for 7B models)
|
||||
# --- Verbosity ---
|
||||
# "verbose": False, # Set to False to reduce llama.cpp logging
|
||||
}
|
||||
|
||||
# --- Load Model ---
|
||||
# Crucially, pass the full path to the .gguf file
|
||||
self.model = llama_cpp.Llama(model_path=self.gguf_file_path, **model_kwargs)
|
||||
|
||||
self.is_loaded = True
|
||||
log(f"[LocalGGUFLLMConnector] GGUF model loaded successfully from: {self.gguf_file_path}")
|
||||
except Exception as e:
|
||||
error_msg = f"[LocalGGUFLLMConnector] Failed to load GGUF model from '{self.gguf_file_path}': {e}"
|
||||
log(error_msg)
|
||||
# Re-raise the exception to halt the node execution
|
||||
raise Exception(error_msg) from e # Chain the exception
|
||||
|
||||
def invoke(self, messages, **generation_kwargs):
|
||||
"""
|
||||
Generates text using the local GGUF LLM based on the messages list.
|
||||
:param messages: List of message dictionaries (like OpenAI format).
|
||||
Example: [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello!"}]
|
||||
:param generation_kwargs: Additional arguments for text generation (e.g., max_tokens, temperature).
|
||||
These will be filtered for compatibility with llama-cpp-python.
|
||||
:return: The generated text string.
|
||||
"""
|
||||
try:
|
||||
if not self.is_loaded:
|
||||
self._load_model() # Load model on first invocation
|
||||
|
||||
if not self.model:
|
||||
raise Exception("[LocalGGUFLLMConnector] Local GGUF LLM model failed to load or is not initialized.")
|
||||
|
||||
# --- Prepare generation parameters for llama-cpp-python ---
|
||||
# Filter the provided kwargs to only include ones accepted by create_chat_completion
|
||||
# Commonly accepted kwargs for llama-cpp-python (check its documentation for the latest)
|
||||
accepted_kwargs = [
|
||||
'temperature', 'top_p', 'top_k', 'max_tokens', 'presence_penalty',
|
||||
'frequency_penalty', 'repeat_penalty', 'seed', 'stop', 'stream',
|
||||
'mirostat_mode', 'mirostat_tau', 'mirostat_eta'
|
||||
# Add others as needed/allowed by llama-cpp-python's API
|
||||
]
|
||||
filtered_kwargs = {k: v for k, v in generation_kwargs.items() if k in accepted_kwargs}
|
||||
|
||||
# Handle 'max_new_tokens' if passed (common in Transformers, convert to 'max_tokens' for llama-cpp)
|
||||
# Note: The logic here prioritizes 'max_tokens' if both are somehow passed.
|
||||
if 'max_new_tokens' in generation_kwargs and 'max_tokens' not in filtered_kwargs:
|
||||
filtered_kwargs['max_tokens'] = generation_kwargs['max_new_tokens']
|
||||
# Example of handling 'seed' if passed directly (llama-cpp might use it differently internally)
|
||||
# The `seed` is often handled by setting the global random state before generation
|
||||
# if 'seed' in generation_kwargs and generation_kwargs['seed'] is not None:
|
||||
# # llama_cpp.Llama.sample_seed might be relevant, or rely on global state
|
||||
# pass # llama-cpp-python often handles seed within the generation call if passed
|
||||
|
||||
# --- Call the LLM ---
|
||||
# Use the chat completion interface which is generally preferred and handles templates
|
||||
response = self.model.create_chat_completion(
|
||||
messages=messages,
|
||||
**filtered_kwargs # Pass the filtered and potentially adjusted kwargs
|
||||
)
|
||||
# --- Extract the generated text ---
|
||||
# llama-cpp-python's create_chat_completion returns a dict
|
||||
# response = {'id': '...', 'object': 'chat.completion', 'created': ..., 'model': '...', 'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': '...'}, 'finish_reason': 'stop'}], 'usage': {...}}
|
||||
generated_text = response['choices'][0]['message']['content']
|
||||
if generated_text is None:
|
||||
generated_text = "" # Handle potential None if model produces no content
|
||||
|
||||
# Ensure the output is a clean string
|
||||
final_text = generated_text.strip()
|
||||
return final_text # <-- Return ONLY the generated string
|
||||
|
||||
except Exception as e:
|
||||
# Catch any error that occurred within the try block and log it
|
||||
error_msg = f"[LocalGGUFLLMConnector] Error in invoke method: {str(e)}"
|
||||
log(error_msg)
|
||||
# Re-raise the exception so the calling node (e.g., LocalKontextPromptGenerator) knows it failed
|
||||
raise e # Or raise Exception(error_msg) from e
|
||||
+71
-28
@@ -1,6 +1,7 @@
|
||||
# ComfyUI/custom_nodes/ComfyUI_LocalLLMNodes/local_llm_connector.py
|
||||
import os
|
||||
import folder_paths
|
||||
import traceback # Add this import for better error reporting
|
||||
|
||||
# --- Simple logging utility for this package ---
|
||||
def log(message):
|
||||
@@ -90,6 +91,7 @@ class SetLocalLLMServiceConnector:
|
||||
connector = LocalLLMServiceConnector(model_path)
|
||||
return (connector,)
|
||||
|
||||
# --- Inside the LocalLLMServiceConnector class ---
|
||||
|
||||
class LocalLLMServiceConnector:
|
||||
"""
|
||||
@@ -103,6 +105,8 @@ class LocalLLMServiceConnector:
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.is_loaded = False
|
||||
# Optional: Store loading parameters if passed from the node
|
||||
# self.device = device
|
||||
|
||||
def _load_model(self):
|
||||
"""Loads the model and tokenizer if not already loaded."""
|
||||
@@ -110,21 +114,22 @@ class LocalLLMServiceConnector:
|
||||
return # Already loaded
|
||||
|
||||
try:
|
||||
log(f"Loading local LLM model from: {self.model_path}")
|
||||
log(f"[LocalLLMConnector] Loading local LLM model from: {self.model_path}")
|
||||
# --- Model Loading Configuration ---
|
||||
tokenizer_kwargs = {
|
||||
"trust_remote_code": True # Needed for some non-standard models
|
||||
}
|
||||
|
||||
|
||||
# --- Example Quantization Config (Uncomment and use if needed) ---
|
||||
# Requires 'bitsandbytes': pip install bitsandbytes
|
||||
# from transformers import BitsAndBytesConfig # Ensure this import is at the top or here
|
||||
# quantization_config = BitsAndBytesConfig(
|
||||
# load_in_4bit=True,
|
||||
# bnb_4bit_compute_dtype=torch.float16,
|
||||
# bnb_4bit_use_double_quant=True,
|
||||
# bnb_4bit_quant_type="nf4"
|
||||
# )
|
||||
|
||||
|
||||
model_kwargs = {
|
||||
"trust_remote_code": True,
|
||||
# --- Add Quantization Config if using ---
|
||||
@@ -140,7 +145,7 @@ class LocalLLMServiceConnector:
|
||||
# Handle models without a pad token (common with Llama-based models)
|
||||
if self.tokenizer.pad_token_id is None:
|
||||
self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
|
||||
log(f"Set pad_token_id to eos_token_id ({self.tokenizer.eos_token_id})")
|
||||
log(f"[LocalLLMConnector] Set pad_token_id to eos_token_id ({self.tokenizer.eos_token_id})")
|
||||
|
||||
# --- Load Model ---
|
||||
# Note: device_map="auto" is often crucial here, especially with quantization
|
||||
@@ -153,12 +158,18 @@ class LocalLLMServiceConnector:
|
||||
# self.model.to("cpu")
|
||||
|
||||
self.is_loaded = True
|
||||
log(f"Local LLM model loaded successfully from: {self.model_path}")
|
||||
log(f"[LocalLLMConnector] Local LLM model loaded successfully from: {self.model_path}")
|
||||
except Exception as e:
|
||||
error_msg = f"Failed to load local LLM model from {self.model_path}: {e}"
|
||||
# --- Improved Error Reporting for Loading ---
|
||||
error_msg = f"[LocalLLMConnector] Failed to load local LLM model from {self.model_path}"
|
||||
log(error_msg)
|
||||
# Re-raise to stop execution if loading fails
|
||||
raise Exception(error_msg) from e # Chain the exception
|
||||
# Print the full traceback to the console for detailed debugging
|
||||
log(f"[LocalLLMConnector] Traceback:\n{traceback.format_exc()}")
|
||||
# --- End of Improved Error Reporting ---
|
||||
|
||||
# Re-raise the original exception to halt the node execution
|
||||
# Using 'from e' preserves the original exception chain (though traceback.format_exc shows it)
|
||||
raise e # Or: raise Exception(error_msg) from e
|
||||
|
||||
def invoke(self, messages, **generation_kwargs):
|
||||
"""
|
||||
@@ -169,12 +180,16 @@ class LocalLLMServiceConnector:
|
||||
:param generation_kwargs: Additional arguments for text generation (e.g., max_new_tokens, temperature, seed).
|
||||
:return: The generated text string.
|
||||
"""
|
||||
# --- Log input for debugging (optional) ---
|
||||
# log(f"[LocalLLMConnector] invoke called with messages: {messages}")
|
||||
# log(f"[LocalLLMConnector] invoke called with generation_kwargs: {generation_kwargs}")
|
||||
|
||||
try:
|
||||
if not self.is_loaded:
|
||||
self._load_model() # Load model on first invocation
|
||||
|
||||
if not self.model or not self.tokenizer:
|
||||
raise Exception("Local LLM model or tokenizer failed to load.")
|
||||
raise Exception("[LocalLLMConnector] Local LLM model or tokenizer failed to load.")
|
||||
|
||||
# --- Format messages for the local model ---
|
||||
# Try using the tokenizer's chat template if available (more robust)
|
||||
@@ -187,7 +202,7 @@ class LocalLLMServiceConnector:
|
||||
raise ValueError("apply_chat_template did not return a string")
|
||||
except Exception as e:
|
||||
# Fallback to simple formatting if chat template fails or isn't available
|
||||
log(f"Falling back to simple prompt formatting: {e}")
|
||||
log(f"[LocalLLMConnector] Falling back to simple prompt formatting: {e}")
|
||||
prompt_parts = []
|
||||
for message in messages:
|
||||
role = message.get('role', 'user')
|
||||
@@ -201,48 +216,75 @@ class LocalLLMServiceConnector:
|
||||
prompt = "Assistant:" # Fallback if messages were empty
|
||||
|
||||
if not prompt.strip():
|
||||
log("Warning: Generated prompt is empty or whitespace.")
|
||||
log("[LocalLLMConnector] Warning: Generated prompt is empty or whitespace.")
|
||||
return "" # Return empty string if prompt is empty
|
||||
|
||||
# --- Tokenize the prompt ---
|
||||
try:
|
||||
inputs = self.tokenizer(prompt, return_tensors="pt")
|
||||
# Consider moving inputs to model's device if not using device_map="auto"
|
||||
# if hasattr(self.model, 'device') and self.model.device.type != 'meta':
|
||||
# inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
|
||||
# --- FIX 1: Move inputs to the model's device ---
|
||||
# Check if the model has a device attribute and move inputs accordingly.
|
||||
# This is essential when the model is on GPU (cuda) but inputs are on CPU.
|
||||
if hasattr(self.model, 'device') and self.model.device is not None:
|
||||
# Move each tensor in the inputs dictionary to the model's device
|
||||
inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
|
||||
# Note: If device_map="auto" places different parts on different devices,
|
||||
# this simple approach might need refinement, but it works for common cases
|
||||
# where the main model components are on a single device (like cuda:0).
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Error tokenizing prompt: {e}"
|
||||
error_msg = f"[LocalLLMConnector] Error tokenizing prompt or moving to device: {e}"
|
||||
log(error_msg)
|
||||
raise Exception(error_msg) from e
|
||||
|
||||
# --- Set default generation parameters ---
|
||||
# These defaults should be reasonable for prompt generation tasks
|
||||
default_kwargs = {
|
||||
"max_new_tokens": 250, # Slightly higher default for complex prompts
|
||||
"temperature": 0.7, # Default creativity
|
||||
"max_new_tokens": 350, # Slightly higher default for complex prompts
|
||||
"temperature": 0.6, # Default creativity
|
||||
"do_sample": True, # Enable sampling for variety
|
||||
"top_p": 0.9, # Nucleus sampling
|
||||
"repetition_penalty": 1.1, # Slight penalty (1.05-1.2) to discourage repeating words/phrases
|
||||
"length_penalty": 0.8, # Slightly penalize shorter sequences (if supported, encourages longer output)
|
||||
# "min_new_tokens": 50, # Optional: Force a minimum output length (if supported)
|
||||
"pad_token_id": self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id,
|
||||
# "eos_token_id": self.tokenizer.eos_token_id, # Optional: explicitly set EOS
|
||||
}
|
||||
|
||||
# --- FIX 2: Filter kwargs and handle 'seed' ---
|
||||
# Update defaults with any provided kwargs (e.g., from consuming nodes)
|
||||
# Filter out kwargs that might cause issues if not explicitly supported
|
||||
filtered_kwargs = {k: v for k, v in generation_kwargs.items() if k in ['max_new_tokens', 'temperature', 'top_p', 'top_k', 'do_sample', 'num_beams', 'early_stopping', 'pad_token_id', 'eos_token_id', 'seed']}
|
||||
# Handle 'seed' if passed - PyTorch manual seed (affects stochastic operations)
|
||||
if 'seed' in generation_kwargs and generation_kwargs['seed'] is not None:
|
||||
# Filter out kwargs that might cause issues or are handled separately.
|
||||
# Commonly accepted kwargs for transformers.generate (add/remove based on your models if needed)
|
||||
accepted_kwargs = [
|
||||
'max_new_tokens', 'min_new_tokens', 'max_length', 'min_length',
|
||||
'do_sample', 'temperature', 'top_k', 'top_p', 'typical_p',
|
||||
'repetition_penalty', 'length_penalty', 'no_repeat_ngram_size',
|
||||
'encoder_no_repeat_ngram_size', 'bad_words_ids', 'force_words_ids',
|
||||
'num_beams', 'num_beam_groups', 'penalty_alpha', 'use_cache',
|
||||
'output_attentions', 'output_hidden_states', 'return_dict_in_generate',
|
||||
'pad_token_id', 'bos_token_id', 'eos_token_id', 'exponential_decay_length_penalty'
|
||||
# Note: 'seed' is intentionally omitted as it's not a standard generate() kwarg
|
||||
]
|
||||
filtered_kwargs = {k: v for k, v in generation_kwargs.items() if k in accepted_kwargs}
|
||||
|
||||
# Handle 'seed' if passed - Set PyTorch manual seed (affects stochastic operations)
|
||||
# This must be done BEFORE calling model.generate()
|
||||
if generation_kwargs.get('seed') is not None: # Use .get() for safer access
|
||||
try:
|
||||
torch.manual_seed(generation_kwargs['seed'])
|
||||
# Optionally log if needed: log(f"[LocalLLMConnector] Set torch manual seed to {generation_kwargs['seed']}")
|
||||
except Exception as e:
|
||||
log(f"Warning: Could not set seed: {e}")
|
||||
log(f"[LocalLLMConnector] Warning: Could not set seed: {e}")
|
||||
|
||||
default_kwargs.update(filtered_kwargs)
|
||||
default_kwargs.update(filtered_kwargs) # Now only contains valid generate() kwargs
|
||||
|
||||
# --- Generate text ---
|
||||
try:
|
||||
self.model.eval()
|
||||
with torch.no_grad():
|
||||
outputs = self.model.generate(**inputs, **default_kwargs)
|
||||
outputs = self.model.generate(**inputs, **default_kwargs) # Should now work without device/seed errors
|
||||
except Exception as e:
|
||||
error_msg = f"Error during model.generate: {e}"
|
||||
error_msg = f"[LocalLLMConnector] Error during model.generate: {e}"
|
||||
log(error_msg)
|
||||
raise Exception(error_msg) from e
|
||||
|
||||
@@ -253,16 +295,17 @@ class LocalLLMServiceConnector:
|
||||
generated_text = self.tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
|
||||
# Ensure the output is a clean string
|
||||
final_text = generated_text.strip()
|
||||
# log(f"[LocalLLMConnector] Generated text (first 100 chars): {final_text[:100]}...")
|
||||
return final_text # <-- Return ONLY the generated string
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Error decoding generated tokens: {e}"
|
||||
error_msg = f"[LocalLLMConnector] Error decoding generated tokens: {e}"
|
||||
log(error_msg)
|
||||
raise Exception(error_msg) from e
|
||||
|
||||
except Exception as e:
|
||||
# Catch any error that occurred within the try block and log it
|
||||
error_msg = f"Error in invoke method: {str(e)}"
|
||||
error_msg = f"[LocalLLMConnector] Error in invoke method: {str(e)}"
|
||||
log(error_msg)
|
||||
# Re-raise the exception so the calling node knows it failed
|
||||
# Re-raise the exception so the calling node (e.g., KontextPromptGenerator) knows it failed
|
||||
raise e # Or raise Exception(error_msg) from e
|
||||
+147
-46
@@ -15,21 +15,87 @@ script_directory = os.path.dirname(os.path.abspath(__file__))
|
||||
USER_PRESETS_FILE = os.path.join(script_directory, "user_kontext_presets.json")
|
||||
|
||||
# --- Replicate Built-in Presets (from MieNodes' prompt_generator.py) ---
|
||||
# Ensure these match the original structure exactly.
|
||||
KONTEXT_PRESETS = {
|
||||
"Kontext Standard": {
|
||||
"system": "You are an expert AI assistant specializing in generating detailed, creative prompts for image generation models. Your task is to take user-provided image descriptions and edit instructions, then synthesize them into a single, highly descriptive prompt optimized for image generation. Focus on integrating key visual elements like character features, clothing details, scene setting, lighting, and artistic style. Ensure the final prompt is concise, avoids redundancy, and clearly conveys the user's intent for the desired image output."
|
||||
"Kid Clothes -> Flux Prompt": {
|
||||
"system": (
|
||||
"You are an expert prompt generator for Flux Kontext. Your task is to place clothing described by the user onto a baby or child with Arab facial features, while preserving all material, color, and style details from the original clothing description. "
|
||||
"Input: "
|
||||
"1. [IMAGE DESC]: A clean and specific description of a piece of baby or children’s clothing. "
|
||||
"2. [USER INTENT]: Typically vague, like 'put on baby' or 'make it worn'. "
|
||||
"Rules: "
|
||||
"Output only a single Flux prompt (no extra labels, quotes, or explanation). "
|
||||
"Always apply the clothes on a realistic Arab baby or toddler (chubby cheeks, warm skin tone, soft features, natural lighting). "
|
||||
"Retain all textures, cuts, stitching, patterns, and colors of the clothing exactly. "
|
||||
"Set the baby in a realistic environment: white blanket, nursery, soft studio setting. "
|
||||
"Describe lighting: soft diffused lighting, warm tones, shallow depth of field. "
|
||||
"Use fashion product photography terms and camera specs. "
|
||||
"Example: "
|
||||
"[IMAGE DESC]: A pale blue cotton baby romper with white cloud patterns, soft buttons and footies. "
|
||||
"[USER INTENT]: put it on a baby "
|
||||
"Output: "
|
||||
"A photorealistic Arab baby with round cheeks and light tan skin, lying on a soft white fleece blanket, wearing a pale blue cotton romper with stitched white cloud patterns and integrated footies, captured in soft natural lighting with shallow DOF, cozy nursery ambiance, 8K resolution, cinematic baby fashion catalog style."
|
||||
)
|
||||
},
|
||||
"Kontext Detailed": {
|
||||
"system": "You are a master prompt engineer for AI image generators. Your role is to meticulously craft prompts by deeply analyzing user inputs. Given descriptions of two images (e.g., a person and clothing) and specific edit instructions, you must seamlessly merge these elements. Prioritize descriptive keywords for physical attributes, textures, colors, background environments, and artistic influences. The resulting prompt should be rich in detail, logically structured, and highly effective at guiding the image generator to produce the envisioned scene with high fidelity."
|
||||
"Women’s Clothes -> Flux Prompt": {
|
||||
"system": (
|
||||
"You are an expert visual prompt composer for Flux Kontext. Your job is to place the clothing described by the user onto a realistic Muslim woman model in her 20s with Arab features, while preserving all original garment details. "
|
||||
"Input: "
|
||||
"1. [IMAGE DESC]: A detailed visual description of a clothing item (e.g., abaya, blouse, scarf). "
|
||||
"2. [USER INTENT]: Often short, like 'make her wear it' or 'put on a model'. "
|
||||
"Rules: "
|
||||
"Output a single photorealistic prompt. "
|
||||
"Always depict a modest Arab Muslim woman, 20s, wearing the clothes naturally and confidently. "
|
||||
"Keep all design details: fabric texture, color, embroidery, shape. "
|
||||
"Style should reflect high-end modest fashion or Islamic wear catalog. "
|
||||
"Scene: outdoor urban, soft studio light, clean modern backdrop, or lifestyle setting. "
|
||||
"Describe lighting, pose, camera angle, clothing drape and body fit. "
|
||||
"Example: "
|
||||
"[IMAGE DESC]: A beige linen abaya with gold-stitched cuffs and a waist tie. "
|
||||
"[USER INTENT]: put it on muslim woman "
|
||||
"Output: "
|
||||
"A young Arab Muslim woman in her 20s with soft olive skin and delicate features, wearing a beige linen abaya with subtle gold-stitched cuffs and a loosely tied waist sash, standing in a minimalist white studio under soft diffused lighting, side-facing pose with natural folds in the fabric, fashion editorial style with cinematic shallow DOF."
|
||||
)
|
||||
},
|
||||
"Kontext Minimalist": {
|
||||
"system": "You are an AI assistant focused on creating concise, clear prompts for image generation. Your task is to distill user-provided descriptions and edit instructions into a short, essential prompt. Identify the core subject, key action or interaction, and the most important visual style or setting. Eliminate unnecessary details and focus on the primary elements that define the scene. The output prompt should be direct and easy for the image generator to interpret accurately."
|
||||
"Beauty Product Use -> Flux Prompt": {
|
||||
"system": (
|
||||
"You are a creative image prompt generator for Flux Kontext focused on health, skincare, and beauty products. Your job is to place the described product in an artistic or commercial setting — either used by a model (Arab man or woman), or placed within a creative, beauty-inspired scene — while preserving all visual and material details of the product. "
|
||||
"Input: "
|
||||
"1. [IMAGE DESC]: A detailed image or design of a product (e.g., oil bottle, cream jar, serum tube). "
|
||||
"2. [USER INTENT]: Usually vague like 'make woman use it' or 'put it in beautiful place'. "
|
||||
"Rules: "
|
||||
"Output a single detailed visual prompt. "
|
||||
"Product must appear clearly: keep all branding, color, bottle design intact. "
|
||||
"If used: describe the Arab model using it naturally (e.g., applying to face, holding dropper). "
|
||||
"If placed: build a beautiful, creative background — spa setting, nature elements, mirror, floral. "
|
||||
"Use keywords: soft light, high-end spa, cinematic DOF, luxury textures, minimalist props. "
|
||||
"Example: "
|
||||
"[IMAGE DESC]: A small amber glass dropper bottle with golden label and black cap. "
|
||||
"[USER INTENT]: woman use it "
|
||||
"Output: "
|
||||
"An elegant Arab woman with natural tan skin and voluminous dark hair, holding an amber glass dropper bottle with a gold label and black cap, delicately applying serum to her cheek in front of a glowing vanity mirror, soft candle-lit ambiance, subtle shadows, cinematic close-up, bokeh background, luxury skincare ad style."
|
||||
)
|
||||
},
|
||||
"Beauty Product Display -> Flux Prompt": {
|
||||
"system": (
|
||||
"You are an expert in luxury beauty product visuals for Flux Kontext. Place the described product in a stunning, creative, brand-inspired environment (without any models), using beauty aesthetics and visual storytelling. "
|
||||
"Input: "
|
||||
"1. [IMAGE DESC]: Description of a product container or design (e.g., shampoo bottle, bar soap, lip balm). "
|
||||
"2. [USER INTENT]: Usually vague like 'put it in nature' or 'make it pretty'. "
|
||||
"Rules: "
|
||||
"Output only one final prompt. "
|
||||
"Focus on creative display — natural materials, elegant composition, light play. "
|
||||
"Use props like glass trays, greenery, water droplets, stone, sand, soft cloth. "
|
||||
"Mention reflections, shadows, lighting temperature. "
|
||||
"Product must stand out, sharply rendered. "
|
||||
"Example: "
|
||||
"[IMAGE DESC]: A matte pink soap bar with embossed logo. "
|
||||
"[USER INTENT]: put in creative background "
|
||||
"Output: "
|
||||
"A matte pink soap bar with embossed luxury logo resting on a smooth white marble tray, surrounded by pale rose petals and water droplets, gentle morning light filtering through a frosted glass window, soft shadows, minimalist spa aesthetic, 8K beauty product photography."
|
||||
)
|
||||
},
|
||||
"Kontext Artistic Style Focus": {
|
||||
"system": "You are an AI prompt specialist with an emphasis on artistic style and rendering techniques. Users will provide descriptions of elements and specific edits. Your goal is to construct a prompt that heavily emphasizes the desired artistic style (e.g., 'oil painting', 'cyberpunk', 'watercolor', 'photorealistic'). Integrate the provided subject and scene details, but frame them within the context of the specified artistic approach. Highlight relevant techniques, color palettes, brushwork, or visual effects associated with that style to guide the image generator effectively."
|
||||
}
|
||||
}
|
||||
# --- End of KONTEXT_PRESETS ---
|
||||
|
||||
def load_user_presets():
|
||||
"""Load user-defined presets from the JSON file."""
|
||||
@@ -87,44 +153,79 @@ class LocalKontextPromptGenerator(object):
|
||||
|
||||
# --- Replicate the core generate_kontext_prompt method exactly ---
|
||||
def generate_kontext_prompt(self, llm_service_connector, image1_description, image2_description, edit_instruction, preset, seed=None):
|
||||
"""
|
||||
Replicates the exact core logic from KontextPromptGenerator.generate_kontext_prompt.
|
||||
"""
|
||||
# --- Exact replication of the original logic ---
|
||||
all_presets = get_all_kontext_presets()
|
||||
preset_data = all_presets.get(preset)
|
||||
# --- Fixed Syntax Error ---
|
||||
if not preset_data: # <-- Corrected condition
|
||||
raise ValueError(f"Unknown preset: {preset}")
|
||||
"""
|
||||
Replicates the exact core logic from KontextPromptGenerator.generate_kontext_prompt.
|
||||
"""
|
||||
# --- Exact replication of the original logic ---
|
||||
all_presets = get_all_kontext_presets()
|
||||
preset_data = all_presets.get(preset)
|
||||
# --- Fixed Syntax Error ---
|
||||
if not preset_data: # <-- Corrected condition check
|
||||
raise ValueError(f"Unknown preset: {preset}")
|
||||
|
||||
# --- Refined user content construction for Flux prompt generation ---
|
||||
# --- Refined user content construction for strict interpretation ---
|
||||
def safe_str_convert(value):
|
||||
"""Converts input to string, handling lists and None."""
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, list):
|
||||
return " ".join(str(item) for item in value if item is not None)
|
||||
return str(value)
|
||||
|
||||
# 用户输入拼到user消息中，给LLM最大上下文
|
||||
user_content = ""
|
||||
if image1_description.strip():
|
||||
user_content += f"Image 1 (person) description: {image1_description.strip()}"
|
||||
if image2_description.strip():
|
||||
user_content += f"Image 2 (clothing) description: {image2_description.strip()}"
|
||||
if edit_instruction.strip():
|
||||
user_content += f"Edit instruction: {edit_instruction.strip()}"
|
||||
if not user_content.strip():
|
||||
user_content = "No additional image description or edit instruction provided."
|
||||
# Get and convert inputs
|
||||
raw_img1_desc = image1_description
|
||||
# raw_img2_desc = image2_description # Explicitly ignored for now
|
||||
raw_edit_inst = edit_instruction
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": preset_data["system"]},
|
||||
{"role": "user", "content": user_content},
|
||||
]
|
||||
img1_desc_str = safe_str_convert(raw_img1_desc).strip()
|
||||
# img2_desc_str = safe_str_convert(raw_img2_desc).strip() # Ignored
|
||||
edit_inst_str = safe_str_convert(raw_edit_inst).strip()
|
||||
|
||||
# --- Key Change: Sharper distinction and labeling ---
|
||||
# Construct user_content with very clear, labeled sections
|
||||
user_content_parts = []
|
||||
|
||||
if img1_desc_str:
|
||||
# Present the base description as the subject/context
|
||||
user_content_parts.append(f"IMAGE 1 DESCRIPTION: {img1_desc_str}")
|
||||
# user_content_parts.append(f"[SUBJECT]: {img1_desc_str}")
|
||||
|
||||
if edit_inst_str:
|
||||
# Directly frame edits as Flux adjustments
|
||||
user_content_parts.append(f"FLUX EDIT INSTRUCTION: {edit_inst_str}")
|
||||
else:
|
||||
# Default to a Flux-style improvement request
|
||||
user_content_parts.append("FLUX EDIT INSTRUCTION: Enhance details and realism.")
|
||||
|
||||
# Combine parts
|
||||
user_content = " ".join(user_content_parts) # Simple space separator
|
||||
|
||||
# Fallback if somehow inputs were empty
|
||||
if not user_content.strip():
|
||||
user_content = "IMAGE 1 DESCRIPTION: A product on a white background. FLUX EDIT INSTRUCTION: Make it photorealistic."
|
||||
|
||||
# --- End of Refined user content construction ---
|
||||
|
||||
# Rest of the method remains the same...
|
||||
messages = [
|
||||
{"role": "system", "content": preset_data["system"]},
|
||||
{"role": "user", "content": user_content},
|
||||
]
|
||||
|
||||
# --- Key Change: Use the local connector's invoke method ---
|
||||
# Pass the messages and seed.
|
||||
try:
|
||||
kontext_prompt = llm_service_connector.invoke(messages, seed=seed)
|
||||
# Ensure the output is stripped, like the original
|
||||
return (kontext_prompt.strip(),)
|
||||
except Exception as e:
|
||||
# Handle potential errors during local LLM invocation
|
||||
error_msg = f"Error generating kontext prompt with local LLM: {str(e)}"
|
||||
log(error_msg) # Log to console (using the 'log' function defined in this file)
|
||||
# Return the error message as the prompt string so the workflow doesn't crash silently
|
||||
return (error_msg,)
|
||||
|
||||
# --- Key Change: Use the local connector's invoke method ---
|
||||
# Pass the messages and seed.
|
||||
try:
|
||||
kontext_prompt = llm_service_connector.invoke(messages, seed=seed)
|
||||
# Ensure the output is stripped, like the original
|
||||
return (kontext_prompt.strip(),)
|
||||
except Exception as e:
|
||||
# Handle potential errors during local LLM invocation
|
||||
error_msg = f"Error generating kontext prompt with local LLM: {str(e)}"
|
||||
log(error_msg) # Log to console
|
||||
# Return the error message as the prompt string so the workflow doesn't crash silently
|
||||
return (error_msg,)
|
||||
|
||||
# --- Replicate the is_changed method for proper caching ---
|
||||
def is_changed(self, llm_service_connector, image1_description, image2_description, edit_instruction, preset, seed):
|
||||
@@ -252,4 +353,4 @@ class RemoveUserLocalKontextPreset: # <-- Changed class name prefix
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
# Always re-run when called to check for file changes
|
||||
return float("nan")
|
||||
return float("nan")
|
||||
+24
-5
@@ -1,5 +1,24 @@
|
||||
# ComfyUI/custom_nodes/ComfyUI_LocalLLMNodes/requirements.txt
|
||||
transformers
|
||||
torch
|
||||
# bitsandbytes # Uncomment if you implement quantization
|
||||
# accelerate # Often needed with transformers, especially for device_map
|
||||
# Core dependencies for Hugging Face transformers based nodes (e.g., Hermes-2-Pro-Llama-3-8B)
|
||||
transformers>=4.38.0 # Or a specific version you know works
|
||||
torch>=2.0.0 # Or a specific version you know works
|
||||
# Optional, for quantization with transformers (e.g., 4-bit)
|
||||
# bitsandbytes # Uncomment if you use BitsAndBytesConfig in local_llm_connector.py
|
||||
# accelerate # Often needed with transformers, especially for device_map
|
||||
|
||||
# Core dependency for GGUF based nodes (e.g., Mistral-7B-Instruct-Q8.gguf)
|
||||
llama-cpp-python>=0.1.0 # Or a specific version
|
||||
|
||||
# Optional: For downloading models from Hugging Face Hub
|
||||
huggingface_hub
|
||||
|
||||
# Optional: If your package uses any specific utility libraries
|
||||
# (Although deepdiff, pillow, etc. might be pulled in by ComfyUI or other nodes)
|
||||
# deepdiff
|
||||
# pillow
|
||||
|
||||
# Notes for Users:
|
||||
# For llama-cpp-python with GPU support (CUDA), installation is more complex.
|
||||
# Example for CUDA 11.8 (check llama-cpp-python docs for your version):
|
||||
# CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python
|
||||
# Example for CUDA 12.1:
|
||||
# CMAKE_ARGS="-DLLAMA_CUDA=on" pip install llama-cpp-python
|
||||
Reference in New Issue
Block a user