427 lines
18 KiB
Python
427 lines
18 KiB
Python
import os
|
|
import gc
|
|
import inspect
|
|
import torch
|
|
import base64
|
|
import io
|
|
import json
|
|
from PIL import Image
|
|
from torchvision.transforms import ToPILImage
|
|
from llama_cpp import Llama, llama_backend_free
|
|
from llama_cpp.llama_chat_format import (
|
|
Llava15ChatHandler,
|
|
LlamaChatCompletionHandlerRegistry,
|
|
)
|
|
import folder_paths
|
|
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
|
from typing import Dict, Any, List, Type
|
|
|
|
# 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 _convert_image_to_data_uri(image_tensor: torch.Tensor) -> str:
|
|
"""Convert a ComfyUI image tensor to a base64 data URI for vision models."""
|
|
try:
|
|
to_pil = ToPILImage()
|
|
# ComfyUI images are (B, H, W, C), select first batch and permute to (C, H, W)
|
|
pil_img = to_pil(image_tensor[0].permute(2, 0, 1))
|
|
# Convert to base64
|
|
buffered = io.BytesIO()
|
|
pil_img.save(buffered, format="PNG")
|
|
img_base64 = base64.b64encode(buffered.getvalue()).decode()
|
|
data_uri = f"data:image/png;base64,{img_base64}"
|
|
# Clean up memory
|
|
buffered.close()
|
|
del pil_img
|
|
return data_uri
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to convert image tensor to data URI: {str(e)}")
|
|
|
|
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, lcf.Llava15ChatHandler):
|
|
vision_name = f"vision-{name.lower().replace('chathandler', '')}"
|
|
VISION_HANDLERS[vision_name] = obj
|
|
|
|
|
|
class LlamaCPPModelLoader(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> InputTypeDict:
|
|
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 {
|
|
"required": {
|
|
"model_name": (model_name_list if model_name_list else ["No GGUF models found"], {"tooltip": "Select GGUF model file"}),
|
|
},
|
|
"optional": {
|
|
"chat_format": (chat_formats, {"default": "llama-2", "tooltip": "Chat format template"}),
|
|
"mmproj_model_name": (mmproj_list, {"default": "None", "tooltip": "Multi-modal projector model for vision (select 'None' to disable)"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LLAMA_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "LlamaCPP"
|
|
|
|
def load_model(self, model_name: str, chat_format: str = "llama-2", mmproj_model_name: str = "None") -> tuple:
|
|
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 (model_info,)
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to load model: {str(e)}")
|
|
|
|
|
|
class LlamaCPPOptions(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> InputTypeDict:
|
|
return {
|
|
"optional": {
|
|
"n_gpu_layers": ("INT", {"default": -1, "min": -1, "max": 100, "tooltip": "Number of GPU layers to use"}),
|
|
"n_ctx": ("INT", {"default": 2048, "min": 0, "max": 262144, "tooltip": "Context window size (0 for max)"}),
|
|
"n_threads": ("INT", {"default": -1, "min": -1, "max": 256, "tooltip": "Number of threads (-1 for auto)"}),
|
|
"n_threads_batch": ("INT", {"default": -1, "min": -1, "max": 256, "tooltip": "Number of threads per batch (-1 for auto)"}),
|
|
"n_batch": ("INT", {"default": 512, "min": 1, "max": 16384, "tooltip": "Batch size"}),
|
|
"n_ubatch": ("INT", {"default": 512, "min": 1, "max": 16384, "tooltip": "Micro batch size"}),
|
|
"main_gpu": ("INT", {"default": 0, "min": 0, "max": 100, "tooltip": "GPU ID for main device"}),
|
|
"offload_kqv": (IO.BOOLEAN, {"default": True, "tooltip": "Enable offloading of K/Q/V tensors to GPU", "label_on": "Enabled", "label_off": "Disabled"}),
|
|
"numa": (IO.BOOLEAN, {"default": False, "tooltip": "Enable NUMA affinity", "label_on": "Enabled", "label_off": "Disabled"}),
|
|
"use_mmap": (IO.BOOLEAN, {"default": True, "tooltip": "Enable memory-mapped files", "label_on": "Enabled", "label_off": "Disabled"}),
|
|
"use_mlock": (IO.BOOLEAN, {"default": False, "tooltip": "Enable lock for memory-mapped files", "label_on": "Enabled", "label_off": "Disabled"}),
|
|
"verbose": (IO.BOOLEAN, {"default": False, "tooltip": "Enable verbose logging", "label_on": "Enabled", "label_off": "Disabled"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LLAMA_OPTIONS",)
|
|
RETURN_NAMES = ("options",)
|
|
FUNCTION = "get_options"
|
|
CATEGORY = "LlamaCPP"
|
|
|
|
def get_options(self, **kwargs) -> tuple:
|
|
try:
|
|
# Filter out None/empty values
|
|
options = {k: v for k, v in kwargs.items() if v is not None and v != ""}
|
|
|
|
return (options,)
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to process options: {str(e)}")
|
|
|
|
|
|
class LlamaCPPEngine(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> InputTypeDict:
|
|
return {
|
|
"required": {
|
|
"model": ("LLAMA_MODEL", {"tooltip": "Loaded Llama model from Model Loader"}),
|
|
"prompt": ("STRING", {"multiline": True, "tooltip": "User prompt for chat completion"}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE", {"tooltip": "Input image for vision models"}),
|
|
"options": ("LLAMA_OPTIONS", {"tooltip": "Model options from Options node"}),
|
|
"system_prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "System prompt"}),
|
|
"memory_cleanup": (["close", "backend_free", "full_cleanup", "persistent"], {"default": "close", "tooltip": "Memory cleanup method after generation"}),
|
|
"response_format": (["text", "json_object"], {"default": "text", "tooltip": "Output format (json_object forces valid JSON)"}),
|
|
"max_tokens": ("INT", {"default": 512, "min": 1, "max": 262144, "tooltip": "Maximum tokens to generate"}),
|
|
"temperature": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Sampling temperature"}),
|
|
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Top-p sampling"}),
|
|
"top_k": ("INT", {"default": 100, "min": 0, "max": 400, "tooltip": "Top-k sampling"}),
|
|
"repeat_penalty": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.01, "tooltip": "Repeat penalty"}),
|
|
"seed": ("INT", {"default": -1, "min": -1, "tooltip": "Random seed (-1 for random)"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("response",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "LlamaCPP"
|
|
|
|
def generate(self, model: Dict[str, Any], prompt: str, image: torch.Tensor = None, options: Dict[str, Any] = None, system_prompt: str = "", memory_cleanup: str = "close", response_format: str = "text", max_tokens: int = 128, temperature: float = 0.8, top_p: float = 0.95, top_k: int = 40, repeat_penalty: float = 1.1, seed: int = -1) -> tuple:
|
|
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
|
|
user_content = prompt.strip()
|
|
|
|
# If vision is enabled and image is provided, convert image and create structured content
|
|
if vision_enabled and image is not None:
|
|
image_data_uri = _convert_image_to_data_uri(image)
|
|
user_content = [
|
|
{"type": "image_url", "image_url": image_data_uri},
|
|
{"type": "text", "text": prompt.strip()}
|
|
]
|
|
|
|
messages.append({"role": "user", "content": user_content})
|
|
|
|
# Prepare Llama initialization parameters
|
|
llama_kwargs = {
|
|
"model_path": model_path,
|
|
"chat_format": chat_format,
|
|
}
|
|
|
|
# Add options
|
|
for k, v in options.items():
|
|
llama_kwargs[k] = v
|
|
|
|
# Handle vision models: use chat_handler based on chat_format
|
|
if vision_enabled:
|
|
handler_class = VISION_HANDLERS.get(chat_format, Llava15ChatHandler)
|
|
chat_handler = handler_class(clip_model_path=model["mmproj_model_path"])
|
|
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 = {"type": response_format} if response_format is not None else None
|
|
|
|
# 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,
|
|
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 (response_text,)
|
|
|
|
|
|
class LlamaCPPMemoryCleanup(ComfyNodeABC):
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> InputTypeDict:
|
|
return {
|
|
"required": {
|
|
"memory_cleanup": (["close", "backend_free", "full_cleanup", "persistent"], {"default": "close", "tooltip": "Memory cleanup method"}),
|
|
},
|
|
"optional": {
|
|
"passthrough": (IO.ANY, {"tooltip": "Any input to pass through"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (IO.ANY,)
|
|
RETURN_NAMES = ("passthrough",)
|
|
FUNCTION = "cleanup"
|
|
CATEGORY = "LlamaCPP"
|
|
|
|
def cleanup(self, memory_cleanup: str, passthrough=None) -> tuple:
|
|
try:
|
|
_cleanup_global_llm(memory_cleanup)
|
|
except Exception as e:
|
|
print(f"LlamaCPP Memory Cleanup Error: {str(e)}")
|
|
|
|
return (passthrough,)
|