404 lines
17 KiB
Python
404 lines
17 KiB
Python
import multiprocessing
|
|
import os
|
|
import re
|
|
import json
|
|
import base64
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
from glob import glob
|
|
from rapidfuzz import fuzz, process
|
|
import llama_cpp
|
|
from llama_cpp import Llama
|
|
from comfy_api.latest import io
|
|
from pathlib import Path
|
|
import torch
|
|
import folder_paths
|
|
import comfy.model_management
|
|
|
|
from .utils import *
|
|
|
|
# Standard GGML enum integer values: 1 = F16, 8 = Q8_0, 2 = Q4_0
|
|
KV_CACHE_TYPES = {
|
|
"f16": getattr(getattr(llama_cpp, "llama_cpp", llama_cpp), "GGML_TYPE_F16", 1),
|
|
"q8_0": getattr(getattr(llama_cpp, "llama_cpp", llama_cpp), "GGML_TYPE_Q8_0", 8),
|
|
"q4_0": getattr(getattr(llama_cpp, "llama_cpp", llama_cpp), "GGML_TYPE_Q4_0", 2),
|
|
}
|
|
|
|
class LlamaModel:
|
|
def __init__(self):
|
|
self.llm = None
|
|
|
|
class LlamaPatcher:
|
|
"""Minimal patcher so comfy.model_management can load and unload the llama.cpp model like any other model."""
|
|
def __init__(self, model_file, mmproj_file, n_gpu_layers, llama_kwargs):
|
|
# Stable object tracked by LoadedModel; the Llama instance inside is created and freed on demand
|
|
self.model = LlamaModel()
|
|
self.parent = None
|
|
self.model_file = model_file
|
|
self.mmproj_file = mmproj_file
|
|
self.llama_kwargs = llama_kwargs
|
|
self.size = os.path.getsize(model_file) + (os.path.getsize(mmproj_file) if mmproj_file else 0)
|
|
self.offload_device = torch.device("cpu")
|
|
self.load_device = self.offload_device if n_gpu_layers == 0 else comfy.model_management.get_torch_device()
|
|
|
|
def load(self):
|
|
comfy.model_management.load_models_gpu([self])
|
|
return self.model.llm
|
|
|
|
def model_size(self):
|
|
return self.size
|
|
|
|
def loaded_size(self):
|
|
return self.size if self.model.llm is not None else 0
|
|
|
|
def current_loaded_device(self):
|
|
return self.load_device if self.model.llm is not None else self.offload_device
|
|
|
|
def model_dtype(self):
|
|
return None
|
|
|
|
def model_patches_to(self, device):
|
|
pass
|
|
|
|
def model_patches_models(self):
|
|
return []
|
|
|
|
def is_dynamic(self):
|
|
return False
|
|
|
|
def is_clone(self, other):
|
|
return other.model is self.model
|
|
|
|
def lowvram_patch_counter(self):
|
|
return 0
|
|
|
|
def partially_load(self, device, extra_memory, force_patch_weights=False):
|
|
if self.model.llm is not None:
|
|
return 0
|
|
chat_handler = get_chat_handler(self.model_file, self.mmproj_file) if self.mmproj_file else None
|
|
self.model.llm = Llama(model_path=self.model_file, chat_handler=chat_handler, **self.llama_kwargs)
|
|
return self.size
|
|
|
|
def partially_unload(self, device, memory_to_free):
|
|
return 0
|
|
|
|
def detach(self, unpatch_all=True):
|
|
if unpatch_all and self.model.llm is not None:
|
|
self.model.llm.close()
|
|
self.model.llm = None
|
|
return self.model
|
|
|
|
class DavchaLLMLoader(io.ComfyNode):
|
|
gguf_folder = os.path.join(folder_paths.models_dir, "LLM", "GGUF")
|
|
|
|
# Cache pour les métadonnées (pour accélérer define_schema et F5)
|
|
_metadata_cache = {}
|
|
_path_map = {}
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
path = os.path.join(cls.gguf_folder, "**", "*.gguf")
|
|
files = list(filter(lambda x: "mmproj" not in x, glob(path, recursive=True)))
|
|
|
|
# Nettoyage du cache si on a supprimé un fichier du dossier
|
|
cls._metadata_cache = {k: v for k, v in cls._metadata_cache.items() if k in files}
|
|
|
|
options = []
|
|
cls._path_map.clear()
|
|
|
|
for file in files:
|
|
# On vérifie la date de modification du fichier
|
|
mtime = os.path.getmtime(file)
|
|
|
|
# On ne lit le GGUF que si c'est un nouveau fichier ou s'il a été modifié
|
|
if file not in cls._metadata_cache or cls._metadata_cache[file]['mtime'] != mtime:
|
|
try:
|
|
llm_meta = Llama(file, vocab_only=True, verbose=False)
|
|
arch = llm_meta.metadata.get("general.architecture", "llama")
|
|
max_n_ctx = int(llm_meta.metadata.get(f"{arch}.context_length", 8192))
|
|
max_n_gpu_layers = int(llm_meta.metadata.get(f"{arch}.block_count", 99))
|
|
|
|
# Détection si c'est un modèle MoE (ex: DeepSeek, Mixtral, Qwen MoE, etc.)
|
|
expert_count = int(llm_meta.metadata.get(f"{arch}.expert_count", 0))
|
|
expert_used_count = int(llm_meta.metadata.get(f"{arch}.expert_used_count", 0))
|
|
is_moe = expert_count > 0 or expert_used_count > 0 or any(kw in arch.lower() for kw in ["moe", "mixtral", "dbrx"])
|
|
|
|
cls._metadata_cache[file] = {
|
|
'mtime': mtime,
|
|
'max_n_ctx': max_n_ctx,
|
|
'max_n_gpu_layers': max_n_gpu_layers,
|
|
'is_moe': is_moe,
|
|
'expert_count': expert_count,
|
|
'expert_used_count': expert_used_count,
|
|
}
|
|
except Exception as e:
|
|
print(f"[DavchaLLM] Erreur lecture métadonnées {file}: {e}")
|
|
cls._metadata_cache[file] = {'mtime': mtime, 'max_n_ctx': 8192, 'max_n_gpu_layers': 99, 'is_moe': False, 'expert_count': 0, 'expert_used_count': 0}
|
|
|
|
cached_data = cls._metadata_cache[file]
|
|
|
|
# Sub-inputs per selected model
|
|
inputs = [
|
|
io.Int.Input("n_ctx", min=512, max=cached_data['max_n_ctx'], default=4096),
|
|
io.Int.Input("n_gpu_layers", min=-1, max=cached_data['max_n_gpu_layers'], default=-1),
|
|
]
|
|
|
|
# Dynamically add n_cpu_moe ONLY for MoE models
|
|
if cached_data['is_moe']:
|
|
inputs.append(io.Int.Input("n_cpu_moe", min=0, max=cached_data['max_n_gpu_layers'], default=0))
|
|
inputs.append(io.Int.Input("expert_used_count", min=1, max=cached_data['expert_count'], default=cached_data['expert_used_count']))
|
|
|
|
name = Path(file).stem
|
|
cls._path_map[name] = file
|
|
option = io.DynamicCombo.Option(name, inputs)
|
|
options.append(option)
|
|
|
|
default_threads = max(1, multiprocessing.cpu_count() // 2)
|
|
|
|
return io.Schema(
|
|
node_id="DavchaLLMLoader",
|
|
category="davcha/llm",
|
|
inputs=[
|
|
io.DynamicCombo.Input("model", options=options),
|
|
io.Int.Input("n_batch", min=64, max=32768, default=2048),
|
|
io.Int.Input("n_ubatch", min=64, max=4096, default=512),
|
|
io.Int.Input("n_threads", min=1, max=128, default=default_threads),
|
|
io.Combo.Input("kv_cache_type", options=["f16", "q8_0", "q4_0"], default="q8_0"),
|
|
io.Boolean.Input("flash_attn", default=True),
|
|
],
|
|
outputs=[
|
|
io.Custom("DavchaLLMModel").Output(),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, n_batch, n_ubatch, n_threads, kv_cache_type, flash_attn):
|
|
n_ctx = model.get("n_ctx", 4096)
|
|
n_gpu_layers = model.get("n_gpu_layers", -1)
|
|
# Defaults to 0 safely for non-MoE models where "n_cpu_moe" is not present
|
|
n_cpu_moe = model.get("n_cpu_moe", 0)
|
|
expert_used_count = model.get("expert_used_count", None)
|
|
|
|
model_name = model.get("model", None)
|
|
if model_name is None:
|
|
raise FileExistsError("model")
|
|
|
|
model_file = cls._path_map[model_name]
|
|
mmproj_files = glob(os.path.join(os.path.dirname(model_file), "*mmproj*.gguf"))
|
|
mmproj = mmproj_files[0] if mmproj_files else None
|
|
|
|
kv_type = KV_CACHE_TYPES.get(kv_cache_type, KV_CACHE_TYPES["q8_0"])
|
|
|
|
try:
|
|
meta_llm = Llama(model_path=model_file, vocab_only=True, verbose=False)
|
|
arch = meta_llm.metadata.get("general.architecture", "llama")
|
|
except:
|
|
arch = "llama"
|
|
|
|
kv_overrides = {}
|
|
if expert_used_count is not None:
|
|
override_key = f"{arch}.expert_used_count"
|
|
kv_overrides[override_key] = int(expert_used_count)
|
|
print(f"[DavchaLLM] Overriding MoE Top-K: {override_key} = {expert_used_count}")
|
|
|
|
llama_kwargs = dict(
|
|
n_ctx=n_ctx,
|
|
n_gpu_layers=n_gpu_layers,
|
|
n_cpu_moe=n_cpu_moe,
|
|
kv_overrides=kv_overrides,
|
|
type_k=kv_type,
|
|
type_v=kv_type,
|
|
flash_attn=flash_attn,
|
|
swa_full=True,
|
|
n_batch=n_batch,
|
|
n_ubatch=n_ubatch,
|
|
n_threads=n_threads,
|
|
n_threads_batch=n_threads,
|
|
offload_kqv=True,
|
|
logits_all=False,
|
|
verbose=False
|
|
)
|
|
|
|
return io.NodeOutput(LlamaPatcher(model_file, mmproj, n_gpu_layers, llama_kwargs))
|
|
|
|
class DavchaLLM(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="DavchaLLM",
|
|
category="davcha/llm",
|
|
inputs=[
|
|
io.Custom("DavchaLLMModel").Input("llm"),
|
|
io.String.Input("system", multiline=True, dynamic_prompts=False),
|
|
io.String.Input("prompt", multiline=True, dynamic_prompts=False),
|
|
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
|
|
io.Int.Input("max_tokens", min=1, default=512),
|
|
io.Float.Input("temperature", min=0.0, max=2.0, default=0.9),
|
|
io.Float.Input("top_p", min=0.0, max=1.0, default=0.95),
|
|
io.Int.Input("top_k", min=0, max=100, default=40),
|
|
io.Float.Input("repeat_penalty", min=0.5, max=2.0, default=1.2),
|
|
io.Boolean.Input("enable_thinking", default=True),
|
|
io.String.Input("response_format", multiline=True, dynamic_prompts=False, default=""),
|
|
io.Image.Input("images", optional=True),
|
|
],
|
|
outputs=[
|
|
io.String.Output()
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, llm, system, prompt, seed, max_tokens, temperature, top_p, top_k, repeat_penalty, enable_thinking, response_format, images=None):
|
|
if images is not None:
|
|
content = []
|
|
if not isinstance(images, list):
|
|
images = [images]
|
|
for image in images:
|
|
array = (image[0]*255).clamp(0, 255).byte().cpu().numpy()
|
|
pil_img = Image.fromarray(array, mode="RGB")
|
|
buf = BytesIO()
|
|
pil_img.save(buf, format="PNG")
|
|
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
content.append({'type': 'image_url', 'image_url': {'url': f"data:image/png;base64,{b64}"}})
|
|
|
|
content.append({'type': 'text', 'text': prompt})
|
|
|
|
messages = [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": content}
|
|
]
|
|
else:
|
|
messages = [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": prompt}
|
|
]
|
|
|
|
response_format = json.loads(response_format) if response_format.strip() != "" else None
|
|
|
|
llm = llm.load()
|
|
|
|
with disable_thinking_context(llm, enable_thinking):
|
|
result = llm.create_chat_completion(
|
|
seed=seed,
|
|
messages=messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
repeat_penalty=repeat_penalty,
|
|
response_format=response_format,
|
|
)
|
|
|
|
result = strip_via_fuzzy_tags(result['choices'][0]['message']['content'], llm)
|
|
|
|
return io.NodeOutput(result)
|
|
|
|
class DavchaPromptEnricher(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="DavchaPromptEnricher",
|
|
category="davcha/llm",
|
|
inputs=[
|
|
io.Custom("DavchaLLMModel").Input("llm"),
|
|
io.String.Input("system", multiline=True, dynamic_prompts=False, default="You are a headless REST API server. You receive requests and output strictly valid JSON. You do not have a personality. You do not output markdown formatting. Output a single JSON object starting with {."),
|
|
io.String.Input("user_instruction", multiline=True, dynamic_prompts=False, default="Here is a prompt with some variables:\n{prompt}\n\n---\n\nWrite detailed visual descriptions for the following variables: {keys}. Ensure the generated values can be replaced in the prompt, ensuring grammatical correctness and semantic coherence. Ensure no details from the variables are missing. Minimum length per key > 0. Result in the following JSON format: {json_format}"),
|
|
io.String.Input("prompt", multiline=True, dynamic_prompts=False),
|
|
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
|
|
io.Int.Input("max_tokens", min=1, default=512),
|
|
io.Float.Input("temperature", min=0.0, max=2.0, default=0.9, step=0.01),
|
|
io.Float.Input("top_p", min=0.0, max=1.0, default=0.95, step=0.01),
|
|
io.Int.Input("top_k", min=0, max=100, default=40),
|
|
io.Float.Input("repeat_penalty", min=0.5, max=2.0, default=1.2, step=0.01),
|
|
io.Boolean.Input("enable_thinking", default=True),
|
|
io.Boolean.Input("force_json_output", default=True),
|
|
io.Image.Input("images", optional=True),
|
|
],
|
|
outputs=[
|
|
io.String.Output()
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, llm, system, user_instruction, prompt, seed, max_tokens, temperature, top_p, top_k, repeat_penalty, enable_thinking, force_json_output, images=None):
|
|
if not re.search(r'\{([^}]+)\}', prompt):
|
|
return io.NodeOutput(prompt)
|
|
|
|
m = re.findall(r'\{([^}]+)\}', prompt)
|
|
keys = {x: "" for x in m}
|
|
json_format = {
|
|
"type": "json_object",
|
|
"schema": {
|
|
"type": "object",
|
|
"properties": {key:{"type":"string"} for key in m},
|
|
"required": m
|
|
}
|
|
}
|
|
|
|
p = user_instruction
|
|
p = p.replace("{keys}", str(list(keys.keys())))
|
|
p = p.replace("{json_format}", str(keys))
|
|
p = p.replace("{prompt}", prompt)
|
|
|
|
messages = [{"role": "system", "content": system or ""}] if system else []
|
|
|
|
if images is not None:
|
|
content = []
|
|
if not isinstance(images, list):
|
|
images = [images]
|
|
for image in images:
|
|
array = (image[0]*255).clamp(0, 255).byte().cpu().numpy()
|
|
pil_img = Image.fromarray(array, mode="RGB")
|
|
buf = BytesIO()
|
|
pil_img.save(buf, format="PNG")
|
|
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
content.append({'type': 'image_url', 'image_url': {'url': f"data:image/png;base64,{b64}"}})
|
|
|
|
content.append({'type': 'text', 'text': p})
|
|
|
|
messages.append(
|
|
{"role": "user", "content": content}
|
|
)
|
|
else:
|
|
messages.append(
|
|
{"role": "user", "content": p}
|
|
)
|
|
|
|
llm = llm.load()
|
|
with disable_thinking_context(llm, enable_thinking):
|
|
result = llm.create_chat_completion(
|
|
seed=seed,
|
|
messages=messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
repeat_penalty=repeat_penalty,
|
|
response_format=json_format if force_json_output else None,
|
|
)
|
|
groups = result['choices'][0]['message']['content']
|
|
print(f'------------------\nLLM Response\n------------------\n\n{groups}\n\n------------------')
|
|
groups = strip_via_fuzzy_tags(groups, llm)
|
|
|
|
import json_repair
|
|
p = prompt
|
|
j = json_repair.loads(groups)
|
|
d = {}
|
|
if isinstance(j, list):
|
|
for item in j:
|
|
d.update(item)
|
|
else:
|
|
d = j
|
|
for k, v in d.items():
|
|
# find correct key
|
|
best_match = process.extractOne(
|
|
query=k,
|
|
choices=list(keys.keys()),
|
|
scorer=fuzz.partial_ratio
|
|
)
|
|
if best_match:
|
|
matched_key, *_ = best_match
|
|
p = p.replace(f"{{{matched_key}}}", "\n".join([str(x) for x in v]) if isinstance(v, list) else str(v))
|
|
else:
|
|
print(f"unmatched key: {k}")
|
|
|
|
return io.NodeOutput(p) |