478 lines
22 KiB
Python
478 lines
22 KiB
Python
import os
|
|
import gc
|
|
import inspect
|
|
import torch
|
|
import sys
|
|
import json
|
|
from llama_cpp import Llama, llama_backend_free
|
|
from llama_cpp.llama_chat_format import LlamaChatCompletionHandlerRegistry
|
|
try:
|
|
from llama_cpp.llama_chat_format import MTMDChatHandler
|
|
except ImportError:
|
|
# Support for older llama-cpp-python versions
|
|
from llama_cpp.llama_chat_format import Llava15ChatHandler
|
|
MTMDChatHandler = Llava15ChatHandler
|
|
import folder_paths
|
|
from comfy_api.latest import io
|
|
from typing import Dict, Any, List
|
|
from .utils import image_to_data_uri
|
|
|
|
# Config file path
|
|
CONFIG_FILE = os.path.join(os.path.dirname(__file__), 'config.json')
|
|
|
|
def load_config() -> Dict[str, Any]:
|
|
"""Load config from config.json, return empty dict if not found or invalid."""
|
|
try:
|
|
with open(CONFIG_FILE, 'r') as f:
|
|
return json.load(f)
|
|
except (FileNotFoundError, json.JSONDecodeError):
|
|
return {}
|
|
|
|
def get_user_model_folders() -> List[str]:
|
|
"""Get user-specified model folders from config."""
|
|
config = load_config()
|
|
return config.get('model_folders', [])
|
|
|
|
def get_merged_model_folders() -> List[str]:
|
|
"""Merge ComfyUI text_encoders folders with user folders."""
|
|
try:
|
|
comfy_folders = folder_paths.get_folder_paths("text_encoders")
|
|
except:
|
|
comfy_folders = []
|
|
user_folders = get_user_model_folders()
|
|
all_folders = comfy_folders + user_folders
|
|
# Filter out non-existent paths
|
|
return [f for f in all_folders if os.path.exists(f)]
|
|
|
|
def scan_gguf_models_in_folders() -> List[str]:
|
|
"""Scan merged folders for GGUF model files."""
|
|
folders = get_merged_model_folders()
|
|
model_list = []
|
|
for folder in folders:
|
|
try:
|
|
files = os.listdir(folder)
|
|
model_list.extend([f for f in files if f.lower().endswith('.gguf')])
|
|
except:
|
|
pass # Skip inaccessible folders
|
|
return model_list
|
|
|
|
def find_model_path(model_name: str) -> str:
|
|
"""Find full path to model in merged folders."""
|
|
folders = get_merged_model_folders()
|
|
for folder in folders:
|
|
path = os.path.join(folder, model_name)
|
|
if os.path.exists(path):
|
|
return path
|
|
return None
|
|
|
|
# Global LLM instance for persistence
|
|
_global_llm = None
|
|
|
|
def _cleanup_global_llm(mode: str):
|
|
"""Helper function to cleanup the global LLM based on mode."""
|
|
global _global_llm
|
|
if mode == "persistent":
|
|
return # No cleanup
|
|
|
|
# Common cleanup for all non-persistent modes
|
|
if _global_llm is not None:
|
|
# Clean up chat_handler if it exists (for vision models)
|
|
if hasattr(_global_llm, "chat_handler") and _global_llm.chat_handler is not None:
|
|
try:
|
|
chat_handler = _global_llm.chat_handler
|
|
|
|
# Primary Fix: Close ExitStack if present (used in newer llama-cpp-python versions)
|
|
# This triggers the callback to mtmd_free -> mtmd_cpp.mtmd_free
|
|
if hasattr(chat_handler, "_exit_stack") and chat_handler._exit_stack is not None:
|
|
try:
|
|
chat_handler._exit_stack.close()
|
|
except Exception:
|
|
pass
|
|
|
|
# Secondary Fix: Manual cleanup of attributes as fallback
|
|
# Attempt to find and close the CLIP/mmproj model embedded in the handler
|
|
clip_attrs = ["clip_model", "_clip_model", "mmproj", "_mmproj", "clf"]
|
|
|
|
for attr_name in clip_attrs:
|
|
if hasattr(chat_handler, attr_name):
|
|
attr = getattr(chat_handler, attr_name)
|
|
if attr is not None:
|
|
# Try to close/free the model if a method exists
|
|
if hasattr(attr, "close"):
|
|
try:
|
|
attr.close()
|
|
except Exception:
|
|
pass
|
|
elif hasattr(attr, "__del__"):
|
|
try:
|
|
attr.__del__()
|
|
except Exception:
|
|
pass
|
|
|
|
# Explicitly remove the attribute to break reference cycles
|
|
setattr(chat_handler, attr_name, None)
|
|
|
|
_global_llm.chat_handler.close()
|
|
except Exception:
|
|
pass # Ignore cleanup errors for chat_handler
|
|
|
|
# Remove the handler callback/reference from the LLM if possible/accessible
|
|
if hasattr(_global_llm, "chat_handler"):
|
|
del _global_llm.chat_handler
|
|
# _global_llm might not have chat_handler attribute in strict sense if it was just in kwargs,
|
|
# but we are cleaning up the attribute we accessed.
|
|
|
|
gc.collect()
|
|
gc.collect()
|
|
try:
|
|
_global_llm.close()
|
|
except AttributeError:
|
|
pass # close() method may not be available in all versions
|
|
del _global_llm
|
|
_global_llm = None
|
|
gc.collect()
|
|
|
|
# Backend free cleanup for backend_free and full_cleanup modes
|
|
if mode in ["backend_free", "full_cleanup"]:
|
|
try:
|
|
llama_backend_free()
|
|
gc.collect()
|
|
except (NameError, AttributeError):
|
|
pass # llama_backend_free may not be available
|
|
|
|
# Full torch cleanup for full_cleanup mode
|
|
if mode == "full_cleanup":
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.synchronize()
|
|
gc.collect()
|
|
|
|
|
|
# Mapping of chat formats to vision chat handlers (dynamically discovered)
|
|
import llama_cpp.llama_chat_format as lcf
|
|
|
|
VISION_HANDLERS = {}
|
|
for name, obj in inspect.getmembers(lcf):
|
|
if inspect.isclass(obj) and issubclass(obj, MTMDChatHandler):
|
|
vision_name = f"vision-{name.lower().replace('chathandler', '')}"
|
|
VISION_HANDLERS[vision_name] = obj
|
|
|
|
|
|
class LlamaCPPModelLoader(io.ComfyNode):
|
|
"""Loads a GGUF model and optional mmproj for vision."""
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
model_list = scan_gguf_models_in_folders()
|
|
|
|
# Filter models based on criteria (case-insensitive)
|
|
model_name_list = [f for f in model_list if 'mmproj' not in f.lower() and 'draft' not in f.lower()]
|
|
mmproj_list = [f for f in model_list if 'mmproj' in f.lower()]
|
|
|
|
# Add "None" option to make mmproj truly optional
|
|
if mmproj_list:
|
|
mmproj_list.insert(0, "None")
|
|
else:
|
|
mmproj_list = ["None"]
|
|
|
|
# Available chat formats from llama-cpp-python
|
|
base_chat_formats = sorted(list(LlamaChatCompletionHandlerRegistry()._chat_handlers.keys()))
|
|
vision_formats = list(VISION_HANDLERS.keys())
|
|
chat_formats = sorted(base_chat_formats + vision_formats)
|
|
|
|
return io.Schema(
|
|
node_id="LlamaCPPModelLoader",
|
|
display_name="Llama CPP Model Loader",
|
|
category="LlamaCPP",
|
|
inputs=[
|
|
io.Combo.Input("model_name", options=model_name_list if model_name_list else ["No GGUF models found"], tooltip="Select GGUF model file"),
|
|
io.Combo.Input("chat_format", options=chat_formats, default="llama-2", tooltip="Chat format template", optional=True),
|
|
io.Combo.Input("mmproj_model_name", options=mmproj_list, default="None", tooltip="Multi-modal projector model for vision (select 'None' to disable)", optional=True),
|
|
],
|
|
outputs=[
|
|
io.Custom("LLAMA_MODEL").Output(display_name="MODEL")
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_name: str, chat_format: str = "llama-2", mmproj_model_name: str = "None") -> io.NodeOutput:
|
|
try:
|
|
model_path = find_model_path(model_name)
|
|
|
|
if model_path is None:
|
|
raise FileNotFoundError(f"Model file not found: {model_name}")
|
|
|
|
if not model_name.lower().endswith('.gguf'):
|
|
raise ValueError(f"Selected file is not a GGUF model: {model_name}")
|
|
|
|
model_info = {
|
|
"model_path": model_path,
|
|
"chat_format": chat_format,
|
|
}
|
|
|
|
# Handle mmproj if provided and not "None"
|
|
if mmproj_model_name and mmproj_model_name != "None":
|
|
mmproj_model_path = find_model_path(mmproj_model_name)
|
|
if mmproj_model_path is None:
|
|
raise FileNotFoundError(f"Multi-modal projector model not found: {mmproj_model_name}")
|
|
model_info["mmproj_model_path"] = mmproj_model_path
|
|
|
|
# Return model info dict
|
|
return io.NodeOutput(model_info)
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load model: {str(e)}")
|
|
|
|
|
|
class LlamaCPPOptions(io.ComfyNode):
|
|
"""Configures Llama-cpp-python initialization parameters."""
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="LlamaCPPOptions",
|
|
display_name="Llama CPP Options",
|
|
category="LlamaCPP",
|
|
inputs=[
|
|
io.Int.Input("n_gpu_layers", default=-1, min=-1, max=100, tooltip="Number of GPU layers to use", optional=True),
|
|
io.Int.Input("n_ctx", default=2048, min=-1, max=262144, tooltip="Context window size (0 for max, -1 for default)", optional=True),
|
|
io.Int.Input("n_threads", default=-1, min=-1, max=256, tooltip="Number of threads (-1 for auto)", optional=True),
|
|
io.Int.Input("n_threads_batch", default=-1, min=-1, max=256, tooltip="Number of threads per batch (-1 for auto)", optional=True),
|
|
io.Int.Input("n_batch", default=2048, min=-1, max=16384, tooltip="Batch size (-1 for default)", optional=True),
|
|
io.Int.Input("n_ubatch", default=512, min=-1, max=16384, tooltip="Micro batch size (-1 for default)", optional=True),
|
|
io.Int.Input("main_gpu", default=0, min=-1, max=100, tooltip="GPU ID for main device (-1 for default)", optional=True),
|
|
io.Boolean.Input("offload_kqv", default=True, tooltip="Enable offloading of K/Q/V tensors to GPU", optional=True),
|
|
io.Boolean.Input("numa", default=False, tooltip="Enable NUMA affinity", optional=True),
|
|
io.Boolean.Input("use_mmap", default=True, tooltip="Enable memory-mapped files", optional=True),
|
|
io.Boolean.Input("use_mlock", default=False, tooltip="Enable lock for memory-mapped files", optional=True),
|
|
io.Boolean.Input("use_direct_io", default=False, tooltip="Enable direct I/O for library (Linux only)", optional=True),
|
|
io.Boolean.Input("verbose", default=False, tooltip="Enable verbose logging", optional=True),
|
|
io.Boolean.Input("vision_use_gpu", default=True, tooltip="Vision: Enable GPU for vision handler", optional=True),
|
|
io.Int.Input("vision_image_min_tokens", default=-1, min=-1, max=16384, tooltip="Vision: Minimum image tokens (-1 for default)", optional=True),
|
|
io.Int.Input("vision_image_max_tokens", default=-1, min=-1, max=16384, tooltip="Vision: Maximum image tokens (-1 for default)", optional=True),
|
|
io.Boolean.Input("vision_enable_thinking", default=False, tooltip="Vision: Enable thinking", optional=True),
|
|
io.Boolean.Input("vision_force_reasoning", default=False, tooltip="Vision: Force reasoning", optional=True),
|
|
io.Boolean.Input("vision_add_vision_id", default=True, tooltip="Vision: Add vision ID", optional=True),
|
|
],
|
|
outputs=[
|
|
io.Custom("LLAMA_OPTIONS").Output(display_name="OPTIONS")
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
try:
|
|
# Filter out None/empty values
|
|
options = {k: v for k, v in kwargs.items() if v is not None and v != ""}
|
|
|
|
return io.NodeOutput(options)
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to process options: {str(e)}")
|
|
|
|
|
|
class LlamaCPPEngine(io.ComfyNode):
|
|
"""Executes chat completion using the loaded model."""
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="LlamaCPPEngine",
|
|
display_name="Llama CPP Engine",
|
|
category="LlamaCPP",
|
|
inputs=[
|
|
io.Custom("LLAMA_MODEL").Input("model", tooltip="Loaded Llama model from Model Loader"),
|
|
io.String.Input("prompt", multiline=True, tooltip="User prompt for chat completion"),
|
|
io.Image.Input("images", tooltip="Input image(s) for vision models (supports batches)", optional=True),
|
|
io.Custom("LLAMA_OPTIONS").Input("options", tooltip="Model options from Options node", optional=True),
|
|
io.String.Input("system_prompt", multiline=True, default="", tooltip="System prompt", optional=True),
|
|
io.Combo.Input("memory_cleanup", options=["close", "backend_free", "full_cleanup", "persistent"], default="close", tooltip="Memory cleanup method after generation", optional=True),
|
|
io.DynamicCombo.Input("response_format", options=[
|
|
io.DynamicCombo.Option("text", inputs=[]),
|
|
io.DynamicCombo.Option("json_object", inputs=[
|
|
io.String.Input("json_schema", multiline=True, default="{}", tooltip="JSON Schema to enforce if json_object is selected", optional=True),
|
|
])
|
|
], tooltip="Output format (json_object forces valid JSON)", optional=True),
|
|
io.Int.Input("max_tokens", default=512, min=1, max=262144, tooltip="Maximum tokens to generate", optional=True),
|
|
io.Float.Input("temperature", default=0.2, min=0.0, max=10.0, step=0.01, tooltip="Sampling temperature", optional=True),
|
|
io.Float.Input("top_p", default=0.95, min=0.0, max=1.0, step=0.01, tooltip="Top-p sampling", optional=True),
|
|
io.Int.Input("top_k", default=40, min=0, max=400, tooltip="Top-k sampling", optional=True),
|
|
io.Float.Input("min_p", default=0.05, min=0.0, max=1.0, step=0.01, tooltip="Min-p sampling", optional=True),
|
|
io.Float.Input("repeat_penalty", default=1.1, min=1.0, max=5.0, step=0.01, tooltip="Repeat penalty", optional=True),
|
|
io.Float.Input("present_penalty", default=0.0, min=0.0, max=5.0, step=0.01, tooltip="Present penalty", optional=True),
|
|
io.Float.Input("frequency_penalty", default=0.0, min=0.0, max=5.0, step=0.01, tooltip="Frequency penalty", optional=True),
|
|
io.Int.Input("seed", default=1, min=-sys.maxsize, max=sys.maxsize, control_after_generate=True, tooltip="Random seed", optional=True),
|
|
],
|
|
outputs=[
|
|
io.String.Output(display_name="RESPONSE")
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model: Dict[str, Any], prompt: str, images: torch.Tensor = None, options: Dict[str, Any] = None, system_prompt: str = "", memory_cleanup: str = "close", response_format: Dict[str, Any] = {"type": "text"}, max_tokens: int = 512, temperature: float = 0.2, top_p: float = 0.95, top_k: int = 40, min_p: float = 0.05, repeat_penalty: float = 1.1, present_penalty: float = 0.0, frequency_penalty: float = 0.0, seed: int = -1) -> io.NodeOutput:
|
|
global _global_llm
|
|
try:
|
|
# Validate inputs
|
|
if not isinstance(model, dict) or "model_path" not in model:
|
|
raise ValueError("Invalid model data received from Model Loader")
|
|
|
|
options = options or {}
|
|
if not isinstance(options, dict):
|
|
raise ValueError("Invalid options data received from Options node")
|
|
|
|
if not prompt.strip():
|
|
raise ValueError("Prompt cannot be empty")
|
|
|
|
# Extract model info
|
|
model_path = model["model_path"]
|
|
chat_format = model["chat_format"]
|
|
|
|
# Determine if vision handler is enabled
|
|
vision_enabled = "mmproj_model_path" in model and chat_format.startswith("vision-")
|
|
|
|
# Create messages for chat completion
|
|
messages = []
|
|
if system_prompt.strip():
|
|
messages.append({"role": "system", "content": system_prompt.strip()})
|
|
|
|
# Create user message content
|
|
if vision_enabled and images is not None:
|
|
user_content = [{"type": "text", "text": prompt.strip()}]
|
|
# images is (B, H, W, C)
|
|
B = images.shape[0]
|
|
for i in range(B):
|
|
img_tensor = images[i] # (H, W, C)
|
|
data_uri = image_to_data_uri(img_tensor)
|
|
user_content.append({
|
|
"type": "image_url",
|
|
"image_url": {"url": data_uri}
|
|
})
|
|
messages.append({"role": "user", "content": user_content})
|
|
else:
|
|
messages.append({"role": "user", "content": prompt.strip()})
|
|
|
|
# Prepare Llama initialization parameters
|
|
llama_kwargs = {
|
|
"model_path": model_path,
|
|
"chat_format": chat_format,
|
|
}
|
|
|
|
# Add options
|
|
for k, v in options.items():
|
|
if not k.startswith("vision_"):
|
|
# Whitelist for -1 values: only n_gpu_layers is allowed to pass -1
|
|
if v == -1 and k != "n_gpu_layers":
|
|
continue
|
|
llama_kwargs[k] = v
|
|
|
|
# Handle vision models: use chat_handler based on chat_format
|
|
if vision_enabled:
|
|
handler_class = VISION_HANDLERS.get(chat_format, MTMDChatHandler)
|
|
|
|
# Dynamically get parameters from the base class and actual class (MTMDChatHandler)
|
|
base_sig = inspect.signature(MTMDChatHandler)
|
|
handler_params = set(base_sig.parameters.keys())
|
|
handler_sig = inspect.signature(handler_class)
|
|
handler_params.update(handler_sig.parameters.keys())
|
|
|
|
handler_kwargs = {
|
|
"clip_model_path": model["mmproj_model_path"],
|
|
"verbose": options.get("verbose", False),
|
|
}
|
|
|
|
# Process vision_ prefixed options
|
|
for k, v in options.items():
|
|
if k.startswith("vision_"):
|
|
param_name = k.replace("vision_", "")
|
|
if param_name in handler_params:
|
|
# Skip -1 values for vision handlers
|
|
if v == -1:
|
|
continue
|
|
handler_kwargs[param_name] = v
|
|
|
|
chat_handler = handler_class(**handler_kwargs)
|
|
llama_kwargs["chat_handler"] = chat_handler
|
|
# Remove chat_format when using vision handler
|
|
llama_kwargs.pop("chat_format", None)
|
|
|
|
# Pre-generation LLM management
|
|
if memory_cleanup == "persistent":
|
|
if _global_llm is None:
|
|
_global_llm = Llama(**llama_kwargs)
|
|
else: # close or backend_free - always cleanup existing and create new
|
|
_cleanup_global_llm(memory_cleanup)
|
|
_global_llm = Llama(**llama_kwargs)
|
|
|
|
# Prepare response_format parameter
|
|
response_format_param = None
|
|
if isinstance(response_format, dict):
|
|
fmt_type = response_format.get("response_format", "text")
|
|
if fmt_type == "json_object":
|
|
response_format_param = {"type": "json_object"}
|
|
json_schema = response_format.get("json_schema", "{}").strip()
|
|
if json_schema:
|
|
try:
|
|
response_format_param["schema"] = json.loads(json_schema)
|
|
except json.JSONDecodeError:
|
|
response_format_param["schema"] = json_schema
|
|
else:
|
|
response_format_param = {"type": fmt_type}
|
|
else:
|
|
# Fallback for simple string if it happens
|
|
response_format_param = {"type": str(response_format)}
|
|
|
|
# Generate response using global LLM
|
|
response = _global_llm.create_chat_completion(
|
|
messages=messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
repeat_penalty=repeat_penalty,
|
|
present_penalty=present_penalty,
|
|
frequency_penalty=frequency_penalty,
|
|
min_p=min_p,
|
|
seed=seed,
|
|
response_format=response_format_param,
|
|
)
|
|
|
|
# Extract the response text
|
|
if not response or "choices" not in response or not response["choices"]:
|
|
raise RuntimeError("No response generated by the model")
|
|
|
|
response_text = response["choices"][0]["message"]["content"]
|
|
|
|
# Post-generation memory cleanup
|
|
_cleanup_global_llm(memory_cleanup)
|
|
|
|
except Exception as e:
|
|
response_text = f"Error generating response: {str(e)}"
|
|
print(f"LlamaCPP Engine Error: {str(e)}")
|
|
|
|
return io.NodeOutput(response_text)
|
|
|
|
|
|
class LlamaCPPMemoryCleanup(io.ComfyNode):
|
|
"""Manually triggers memory cleanup."""
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
template = io.MatchType.Template("passthrough")
|
|
return io.Schema(
|
|
node_id="LlamaCPPMemoryCleanup",
|
|
display_name="Llama CPP Memory Cleanup",
|
|
category="LlamaCPP",
|
|
inputs=[
|
|
io.Combo.Input("memory_cleanup", options=["close", "backend_free", "full_cleanup", "persistent"], default="close", tooltip="Memory cleanup method"),
|
|
io.MatchType.Input("passthrough", template=template, tooltip="Any input to pass through", optional=True),
|
|
],
|
|
outputs=[
|
|
io.MatchType.Output(template=template, display_name="PASSTHROUGH")
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, memory_cleanup: str, passthrough=None) -> io.NodeOutput:
|
|
try:
|
|
_cleanup_global_llm(memory_cleanup)
|
|
except Exception as e:
|
|
print(f"LlamaCPP Memory Cleanup Error: {str(e)}")
|
|
|
|
return io.NodeOutput(passthrough)
|