350 lines
13 KiB
Python
350 lines
13 KiB
Python
import os
|
||
import gc
|
||
import threading
|
||
from typing import Optional
|
||
|
||
import torch
|
||
from llama_cpp import Llama
|
||
from comfy import model_management
|
||
|
||
import folder_paths
|
||
|
||
# -------------------------------------------------------------------
|
||
# 1. Helpers for GGUF file listing & resolution (unchanged)
|
||
# -------------------------------------------------------------------
|
||
_CACHE = {}
|
||
_CACHE_LOCK = threading.Lock()
|
||
|
||
def _list_gguf_files():
|
||
candidates = []
|
||
base_dirs = []
|
||
try:
|
||
base_dirs.extend(folder_paths.get_folder_paths("text_encoders"))
|
||
except Exception:
|
||
pass
|
||
|
||
for base in base_dirs:
|
||
if not base or not os.path.isdir(base):
|
||
continue
|
||
for root, _, files in os.walk(base):
|
||
for name in files:
|
||
if name.lower().endswith(".gguf"):
|
||
full = os.path.join(root, name)
|
||
try:
|
||
rel = os.path.relpath(full, base)
|
||
if rel not in candidates:
|
||
candidates.append(rel)
|
||
except Exception:
|
||
if full not in candidates:
|
||
candidates.append(full)
|
||
|
||
if "<manual_path>" not in candidates:
|
||
candidates.append("<manual_path>")
|
||
return candidates
|
||
|
||
def _resolve_model_path(model_name: str, manual_model_path: str = "") -> str:
|
||
if os.path.isabs(model_name) and os.path.isfile(model_name):
|
||
return model_name
|
||
if model_name == "<manual_path>":
|
||
p = (manual_model_path or "").strip()
|
||
if not p:
|
||
raise ValueError("model is <manual_path> but manual_model_path is empty")
|
||
if not os.path.isfile(p):
|
||
raise FileNotFoundError(f"GGUF model not found: {p}")
|
||
return p
|
||
full = folder_paths.get_full_path("text_encoders", model_name)
|
||
if full and os.path.isfile(full):
|
||
return full
|
||
try:
|
||
for base in folder_paths.get_folder_paths("text_encoders"):
|
||
probe = os.path.join(base, model_name)
|
||
if os.path.isfile(probe):
|
||
return probe
|
||
except Exception:
|
||
pass
|
||
raise FileNotFoundError(f"Could not resolve GGUF model path for: {model_name}")
|
||
|
||
def _build_messages(system_prompt: str, user_prompt: str):
|
||
messages = []
|
||
if system_prompt.strip():
|
||
messages.append({"role": "system", "content": system_prompt.strip()})
|
||
messages.append({"role": "user", "content": user_prompt.strip()})
|
||
return messages
|
||
|
||
def _assemble_prompt(messages: list, model_family: str) -> str:
|
||
prompt = ""
|
||
if model_family == "granite":
|
||
for msg in messages:
|
||
prompt += f"<|start_of_role|>{msg['role']}<|end_of_role|>{msg['content']}<|end_of_text|>\n"
|
||
prompt += "<|start_of_role|>assistant<|end_of_role|><think>\n"
|
||
else:
|
||
for msg in messages:
|
||
prompt += f"<|im_start|>{msg['role']}\n{msg['content']}<|im_end|>\n"
|
||
prompt += "<|im_start|>assistant\n<think>"
|
||
return prompt
|
||
|
||
|
||
# -------------------------------------------------------------------
|
||
# 2. Lookup table for Mordant models (unchanged)
|
||
# -------------------------------------------------------------------
|
||
MORDANT_INFO = {
|
||
"mordant-1.2b": {
|
||
"layers": 16,
|
||
"bf16_total_gb": 2.35,
|
||
"per_layer": {
|
||
"bf16": 2.35 / 16,
|
||
"q8_0": 1.27 / 16,
|
||
"q6_k_m": 1.01682 / 16,
|
||
"q5_k_m": None,
|
||
"q4_k_m": None,
|
||
}
|
||
},
|
||
"mordant-3b": {
|
||
"layers": 40,
|
||
"bf16_total_gb": 6.99,
|
||
"per_layer": {
|
||
"bf16": 6.99 / 40,
|
||
"q8_0": 3.86 / 40,
|
||
"q6_k_m": 3.05 / 40,
|
||
"q5_k_m": 2.70 / 40,
|
||
"q4_k_m": None,
|
||
}
|
||
},
|
||
"mordant-7b": {
|
||
"layers": 32,
|
||
"bf16_total_gb": 16.29,
|
||
"per_layer": {
|
||
"bf16": 16.29 / 32,
|
||
"q8_0": 9.58 / 32,
|
||
"q6_k_m": 7.85 / 32,
|
||
"q5_k_m": 7.08 / 32,
|
||
"q4_k_m": None,
|
||
}
|
||
},
|
||
"mordant-12b": {
|
||
"layers": 40,
|
||
"bf16_total_gb": 24.81,
|
||
"per_layer": {
|
||
"bf16": 24.81 / 40,
|
||
"q8_0": 13.55 / 40,
|
||
"q6_k_m": 10.64 / 40,
|
||
"q5_k_m": 9.33 / 40,
|
||
"q4_k_m": 8.11 / 40,
|
||
}
|
||
},
|
||
}
|
||
|
||
def _detect_model_info(model_path: str):
|
||
try:
|
||
llm_meta = Llama(model_path=model_path, vocab_only=True, n_gpu_layers=0, verbose=False)
|
||
meta = llm_meta.metadata
|
||
arch = meta.get("general.architecture", "unknown")
|
||
total_layers = int(meta.get(f"{arch}.block_count", 0))
|
||
quant_version = meta.get("general.quantization_version", "unknown").lower()
|
||
if quant_version == "unknown":
|
||
quant_version = meta.get("tokenizer.ggml.quantization_version", "unknown").lower()
|
||
del llm_meta
|
||
except Exception:
|
||
fname = os.path.basename(model_path).lower()
|
||
for name in MORDANT_INFO:
|
||
if name in fname:
|
||
info = MORDANT_INFO[name]
|
||
total_layers = info["layers"]
|
||
quant_version = "bf16"
|
||
for q in ["q8_0", "q6_k_m", "q5_k_m", "q4_k_m", "q3_k_m", "q2_k_m"]:
|
||
if q in fname:
|
||
quant_version = q
|
||
break
|
||
break
|
||
else:
|
||
total_layers = 32
|
||
quant_version = "bf16"
|
||
|
||
model_name = None
|
||
fname_lower = os.path.basename(model_path).lower()
|
||
for name in MORDANT_INFO:
|
||
if name in fname_lower:
|
||
model_name = name
|
||
break
|
||
|
||
if model_name is None:
|
||
file_size_gb = os.path.getsize(model_path) / (1024**3)
|
||
per_layer_vram = file_size_gb / total_layers if total_layers > 0 else 0.5
|
||
return (model_name or "unknown", total_layers, per_layer_vram, quant_version)
|
||
|
||
info = MORDANT_INFO[model_name]
|
||
layers = info["layers"]
|
||
per_layer_dict = info.get("per_layer", {})
|
||
per_layer_vram = None
|
||
quant_clean = quant_version.replace("_0", "").replace("_1", "").lower()
|
||
if quant_version in per_layer_dict and per_layer_dict[quant_version] is not None:
|
||
per_layer_vram = per_layer_dict[quant_version]
|
||
elif quant_clean in per_layer_dict and per_layer_dict[quant_clean] is not None:
|
||
per_layer_vram = per_layer_dict[quant_clean]
|
||
else:
|
||
bf16_per_layer = per_layer_dict.get("bf16")
|
||
if bf16_per_layer is None:
|
||
bf16_total = info["bf16_total_gb"]
|
||
bf16_per_layer = bf16_total / layers
|
||
actual_file_size_gb = os.path.getsize(model_path) / (1024**3)
|
||
per_layer_vram = actual_file_size_gb / layers
|
||
|
||
return model_name, layers, per_layer_vram, quant_version
|
||
|
||
|
||
def _auto_n_gpu_layers(model_path: str, verbose: bool = False) -> int:
|
||
if not torch.cuda.is_available():
|
||
return 0
|
||
|
||
total_vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
|
||
usable_vram_gb = max(0, total_vram_gb - 2.0)
|
||
|
||
model_name, max_layers, per_layer_vram, quant = _detect_model_info(model_path)
|
||
if per_layer_vram <= 0:
|
||
return 0
|
||
|
||
n_layers = int(usable_vram_gb / per_layer_vram)
|
||
n_layers = max(0, min(max_layers, n_layers))
|
||
|
||
if verbose:
|
||
print(f"[Mordant Enhancer] GPU VRAM: {total_vram_gb:.2f} GB, "
|
||
f"Reserved: 2.00 GB, Usable: {usable_vram_gb:.2f} GB")
|
||
print(f"[Mordant Enhancer] Model: {model_name}, Layers: {max_layers}, "
|
||
f"Quant: {quant}, Per‑layer VRAM: {per_layer_vram:.4f} GB")
|
||
print(f"[Mordant Enhancer] Auto‑selected GPU layers: {n_layers}")
|
||
|
||
return n_layers
|
||
|
||
|
||
# -------------------------------------------------------------------
|
||
# 3. Main node class – Mordant Prompt Enhancer
|
||
# -------------------------------------------------------------------
|
||
class MordantPromptEnhancer:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"model": (_list_gguf_files(),),
|
||
"text": ("STRING", {"multiline": True, "default": "Rewrite this image prompt conservatively."}),
|
||
"mode": (["enhancer", "general"], {"default": "enhancer"}),
|
||
"enable_enhancement": ("BOOLEAN", {"default": True}),
|
||
"sampling_mode": ("BOOLEAN", {"default": True}),
|
||
"max_new_tokens": ("INT", {"default": 8192, "min": 1, "max": 8192}),
|
||
"seed": ("INT", {"default": 0, "min": -1, "max": 2147483647}),
|
||
"temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 5.0, "step": 0.05}),
|
||
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
"top_k": ("INT", {"default": 50, "min": 0, "max": 1000}),
|
||
"repeat_penalty": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.01}),
|
||
"n_batch": ("INT", {"default": 2048, "min": 32, "max": 4096, "step": 32}),
|
||
"keep_loaded": ("BOOLEAN", {"default": False}),
|
||
"offload_kqv": ("BOOLEAN", {"default": True}),
|
||
"flash_attn": ("BOOLEAN", {"default": True}),
|
||
"verbose": ("BOOLEAN", {"default": False}),
|
||
},
|
||
"optional": {
|
||
"manual_model_path": ("STRING", {"default": "", "multiline": False}),
|
||
"custom_system_prompt": ("STRING", {"default": "", "multiline": True}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING", "STRING")
|
||
RETURN_NAMES = ("composition", "thinking")
|
||
FUNCTION = "generate"
|
||
CATEGORY = "text/llm"
|
||
|
||
def generate(self, model, text, mode, enable_enhancement,
|
||
sampling_mode, max_new_tokens, seed, temperature, top_p, top_k,
|
||
repeat_penalty, n_batch, keep_loaded, offload_kqv, flash_attn,
|
||
verbose, manual_model_path="", custom_system_prompt=""):
|
||
|
||
if not enable_enhancement:
|
||
return (text, "")
|
||
|
||
model_path = _resolve_model_path(model, manual_model_path)
|
||
system_prompt = (custom_system_prompt or "").strip()
|
||
|
||
n_ctx = 8192
|
||
n_threads = os.cpu_count() or 4
|
||
n_gpu_layers = _auto_n_gpu_layers(model_path, verbose=verbose)
|
||
|
||
if not sampling_mode:
|
||
eff_temp, eff_top_p, eff_top_k, eff_repeat_penalty = 0.2, 0.9, 20, max(repeat_penalty, 1.05)
|
||
else:
|
||
eff_temp, eff_top_p, eff_top_k, eff_repeat_penalty = temperature, top_p, top_k, repeat_penalty
|
||
|
||
try:
|
||
from llama_cpp import Llama
|
||
except ImportError:
|
||
raise RuntimeError("llama-cpp-python not installed.")
|
||
|
||
def _create_model():
|
||
kwargs = dict(
|
||
model_path=model_path,
|
||
n_ctx=n_ctx,
|
||
n_gpu_layers=n_gpu_layers,
|
||
n_batch=n_batch,
|
||
verbose=verbose,
|
||
cuda_graphs=False,
|
||
use_mmap=False,
|
||
use_mlock=False,
|
||
)
|
||
if n_threads is not None:
|
||
kwargs["n_threads"] = n_threads
|
||
try:
|
||
kwargs["offload_kqv"] = offload_kqv
|
||
kwargs["flash_attn"] = flash_attn
|
||
return Llama(**kwargs)
|
||
except TypeError:
|
||
kwargs.pop("offload_kqv", None)
|
||
kwargs.pop("flash_attn", None)
|
||
return Llama(**kwargs)
|
||
|
||
if keep_loaded:
|
||
key = (model_path, n_ctx, n_gpu_layers, n_threads, n_batch, offload_kqv, flash_attn)
|
||
with _CACHE_LOCK:
|
||
llm = _CACHE.get(key)
|
||
if llm is None:
|
||
llm = _create_model()
|
||
_CACHE[key] = llm
|
||
else:
|
||
llm = _create_model()
|
||
|
||
actual_family = "chatml"
|
||
if hasattr(llm, 'metadata'):
|
||
chat_template = llm.metadata.get('tokenizer.chat_template', '').lower()
|
||
if 'granite' in chat_template or 'start_of_role' in chat_template:
|
||
actual_family = "granite"
|
||
|
||
messages = _build_messages(system_prompt, text)
|
||
prompt = _assemble_prompt(messages, actual_family)
|
||
|
||
out = llm.create_completion(
|
||
prompt=prompt,
|
||
max_tokens=max_new_tokens,
|
||
temperature=eff_temp,
|
||
top_p=eff_top_p,
|
||
top_k=eff_top_k,
|
||
repeat_penalty=eff_repeat_penalty,
|
||
seed=None if seed < 0 else seed,
|
||
)
|
||
|
||
result_text = out["choices"][0]["text"]
|
||
if "</think>" in (result_text or ""):
|
||
parts = (result_text or "").rsplit("</think>", 1)
|
||
analysis_text, final_text = parts[0].strip(), parts[1].strip()
|
||
else:
|
||
analysis_text, final_text = "", (result_text or "").strip()
|
||
|
||
return (final_text, analysis_text)
|
||
|
||
|
||
# -------------------------------------------------------------------
|
||
# 4. Node registration
|
||
# -------------------------------------------------------------------
|
||
NODE_CLASS_MAPPINGS = {
|
||
"MordantPromptEnhancer": MordantPromptEnhancer,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"MordantPromptEnhancer": "Mordant Prompt Enhancer",
|
||
} |