2942 lines
140 KiB
Python
2942 lines
140 KiB
Python
"""ComfyUI PromptForge - prompt enhancement nodes.
|
||
|
||
- PromptForge Model: picks a connected CLIP, a safetensors model or a GGUF
|
||
and plans quant / context / GPU layers for the current GPU
|
||
- PromptForge Enhancer: writes the prompt with a preset system prompt,
|
||
optional images / video frames, and output cleanup
|
||
- PromptForge Config Finder: reports the planned settings for local models
|
||
"""
|
||
|
||
import base64
|
||
import ctypes
|
||
import glob
|
||
import io
|
||
import json
|
||
import gc
|
||
import importlib.util
|
||
import inspect
|
||
import os
|
||
import re
|
||
import struct
|
||
import time
|
||
from collections import Counter
|
||
from pathlib import Path
|
||
|
||
import psutil
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
|
||
def _preload_nvrtc_builtins():
|
||
"""PyTorch compiles a few CUDA kernels (for example an int64 prod) at run time with NVRTC, which opens its builtins library by name; a pip install puts it
|
||
where the loader does not look, and the kernel fails with "failed to open libnvrtc-builtins". Loading it by path makes the name resolve."""
|
||
nvidia = importlib.util.find_spec("nvidia")
|
||
for folder in (nvidia.submodule_search_locations if nvidia else None) or []:
|
||
for path in glob.glob(os.path.join(folder, "*", "lib", "libnvrtc-builtins.so.*")):
|
||
if ".alt." not in path:
|
||
ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
|
||
|
||
|
||
if os.name == "posix":
|
||
_preload_nvrtc_builtins()
|
||
|
||
from .frame_interpolation import NODE_CLASS_MAPPINGS as VFI_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as VFI_DISPLAY_NAMES
|
||
from .video_save import NODE_CLASS_MAPPINGS as SAVE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as SAVE_DISPLAY_NAMES
|
||
from .lora_loader import NODE_CLASS_MAPPINGS as LORA_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as LORA_DISPLAY_NAMES
|
||
from .unified_loader import NODE_CLASS_MAPPINGS as LOADER_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as LOADER_DISPLAY_NAMES
|
||
from .empty_latent import NODE_CLASS_MAPPINGS as LATENT_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as LATENT_DISPLAY_NAMES
|
||
from .sampler import NODE_CLASS_MAPPINGS as SAMPLER_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as SAMPLER_DISPLAY_NAMES
|
||
from .quantize_model import NODE_CLASS_MAPPINGS as QUANT_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as QUANT_DISPLAY_NAMES
|
||
from .minimax_decode import NODE_CLASS_MAPPINGS as DECODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as DECODE_DISPLAY_NAMES
|
||
from . import image_library, image_actions
|
||
from .image_actions_node import PROMPTFORGE_ImageAction
|
||
from .moss_nodes import NODE_CLASS_MAPPINGS as MOSS_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as MOSS_DISPLAY_NAMES
|
||
from .h3_generate import NODE_CLASS_MAPPINGS as ENDLESS_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as ENDLESS_DISPLAY_NAMES
|
||
from .h3_settings import NODE_CLASS_MAPPINGS as H3_SETTINGS_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as H3_SETTINGS_DISPLAY_NAMES
|
||
from .preview_override import NODE_CLASS_MAPPINGS as PREVIEW_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as PREVIEW_DISPLAY_NAMES
|
||
from .load_image import NODE_CLASS_MAPPINGS as IMAGE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as IMAGE_DISPLAY_NAMES
|
||
from . import h3_ideas # noqa: F401 (registers the saved-ideas routes)
|
||
from .load_video import NODE_CLASS_MAPPINGS as VIDEO_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as VIDEO_DISPLAY_NAMES
|
||
from .load_audio import NODE_CLASS_MAPPINGS as AUDIO_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as AUDIO_DISPLAY_NAMES
|
||
import transformers
|
||
from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer, BitsAndBytesConfig
|
||
from transformers import LogitsProcessor, LogitsProcessorList, StoppingCriteria, StoppingCriteriaList
|
||
from transformers.utils.chat_template_utils import render_jinja_template
|
||
|
||
try:
|
||
from transformers import AutoModelForImageTextToText as AutoVision
|
||
except Exception: # fallback for older transformers naming
|
||
from transformers import AutoModelForVision2Seq as AutoVision
|
||
|
||
try:
|
||
import folder_paths
|
||
except Exception: # ComfyUI injects this module; keep a safe fallback
|
||
folder_paths = None
|
||
|
||
PACK_DIR = Path(__file__).resolve().parent
|
||
PRESETS_PATH = PACK_DIR / "presets" / "promptforge_presets.json"
|
||
MODELS_PATH = PACK_DIR / "catalog" / "models.json"
|
||
|
||
with open(PRESETS_PATH, "r", encoding="utf-8") as _f:
|
||
PROMPTFORGE_PRESETS = json.load(_f).get("presets") or {}
|
||
|
||
with open(MODELS_PATH, "r", encoding="utf-8") as _f:
|
||
PROMPTFORGE_MODELS = json.load(_f) or {}
|
||
|
||
OUTPUT_GUARD = (
|
||
"Output rules that override all style preferences: return only the final usable prompt text. "
|
||
"Do not include analysis, planning, bullet points, headings, markdown, JSON, explanations, alternatives, "
|
||
"self-corrections, notes, apologies, or messages to the user. Do not write phrases such as Final Plan, Final Check, "
|
||
"Wait, Okay, First, Next, Then, I will, or I need. "
|
||
"Never write in first person and never comment on your own work: do not say you followed the instructions, "
|
||
"did a good job, respected the rules, avoided something, or met the request. Never add a closing note, "
|
||
"self-assessment, or summary about what you wrote. Start directly with the required prompt text. "
|
||
"Never repeat, quote, paraphrase, or answer any instruction, rule, constraint, policy, or note from this system prompt; "
|
||
"the output must contain only the prompt text for the target model."
|
||
)
|
||
|
||
_META_PATTERNS = (
|
||
r"^#{1,6}\s",
|
||
r"^(note|remember|important|rule|constraint|instruction|instructions|output|response)\b\s*[:\-]",
|
||
r"^(here is|here's|this is)\b[^.!?]{0,40}\bprompt\b",
|
||
r"^(okay|ok|sure|certainly|alright|absolutely|great|perfect)\s*[,!:.]\s*",
|
||
r"\b(explicitly stated|as requested|as instructed|anything other than)\b",
|
||
r"\b(the user's (request|text|wording|prompt)|system prompt|this prompt|the generated prompt|output contract|content policy|negative constraints?)\b",
|
||
r"\b(output only|only output|in the output)\b",
|
||
r"\b(the prompt|the response|the answer)\s+(should|must|is|will|needs to)\b",
|
||
r"\b(do not|don't|never|avoid|must not|should not)\b[^.!?]{0,80}\b(add|include|mention|name|output|write|describe|use|state|say|deviate|copy|repeat|paraphrase)\b",
|
||
r"\b(I will|I'll|we need to|we should|the goal is|my task|the task is)\b",
|
||
r"\b(here is|here's|below is)\s+(the|your|an?)\s+(enhanced|final|generated|resulting)?\s*(prompt|output|response)\b",
|
||
r"\b(i'?ve|i have|we'?ve|we have)\b[^.!?]{0,90}\b(done|followed|listened|obeyed|applied|respected|honou?red|ensured|made sure|crafted|created|generated|written|wrote|included|avoided|kept|used|focused|paid attention)\b",
|
||
r"\b(i|we)\b[^.!?]{0,90}\b(followed|listened|obeyed|applied|respected|honou?red|ensured|made sure|crafted|generated|avoided|paid attention)\b[^.!?]{0,60}\b(instruction|instructions|rules?|constraints?|guidelines?|request|prompt)\b",
|
||
r"\b(i|we)\s+(made sure|ensured)\b",
|
||
r"\b(done a good job|did a good job|good job|well done|task complete|successfully completed)\b",
|
||
r"\b(the|this|your)\s+(enhanced|final|generated|resulting)?\s*(prompt|output|response)\s+(now\s+)?(is|includes|contains|follows|respects|meets|has been|was|will)\b",
|
||
r"\b(the user|user'?s)\s+(asked|requested|wants|wanted|instructions?)\b",
|
||
r"\b(i hope|hope (this|that) helps|let me know|feel free|if you (want|need|would like))\b",
|
||
r"\b(note|disclaimer|warning)\s*[:\-]",
|
||
r"\b(i|we)\s+(will|would|can|am going to|are going to)\s+(now\s+)?(write|create|generate|provide|describe)\b",
|
||
)
|
||
_META_RE = re.compile("|".join(f"(?:{p})" for p in _META_PATTERNS), re.I)
|
||
|
||
_CHAT_MARKER_RE = re.compile(
|
||
r"<\|(?:im_start|im_end|endoftext|begin_of_text|start_header_id|end_header_id|eot_id|end|system|user|assistant|model|think|/think)\|>\s*(?:system|user|assistant|model)?\s*>?",
|
||
re.I,
|
||
)
|
||
_CHAT_TAG_RE = re.compile(
|
||
r"</?(?:start_of_turn|end_of_turn|bos|eos|pad|unk|s)>|\[/?INST\]|<<?/?SYS>>|<\|?/?channel\|?>",
|
||
re.I,
|
||
)
|
||
_ROLE_LINE_RE = re.compile(r"^(system|user|assistant|model)\s*:\s*", re.I)
|
||
|
||
|
||
def _extract_last_turn(text):
|
||
"""Keep only the final assistant/model turn when template markers are present."""
|
||
patterns = (
|
||
(r"<\|im_start\|>\s*(system|user|assistant|model)\s*>?", r"<\|im_end\|>", ("assistant", "model")),
|
||
(r"<start_of_turn>\s*(system|user|model|assistant)\s*>?", r"<end_of_turn>", ("model", "assistant")),
|
||
(r"<\|start_header_id\|>\s*(system|user|assistant)\s*<\|end_header_id\|>[ \t]*", r"<\|eot_id\|>", ("assistant",)),
|
||
)
|
||
for start_pattern, end_pattern, roles in patterns:
|
||
matches = list(re.finditer(start_pattern, text, re.I))
|
||
if not matches:
|
||
continue
|
||
for match in reversed(matches):
|
||
if match.group(1).lower() in roles:
|
||
tail = text[match.end():]
|
||
end = re.search(end_pattern, tail)
|
||
return tail[:end.start()] if end else tail
|
||
return ""
|
||
return text
|
||
|
||
|
||
def _strip_chat_tokens(text):
|
||
"""Remove ChatML/Gemma/Llama role markers and keep the assistant answer."""
|
||
text = _extract_last_turn(text or "")
|
||
text = _CHAT_MARKER_RE.sub("", text)
|
||
text = _CHAT_TAG_RE.sub("", text)
|
||
lines = text.splitlines()
|
||
role_positions = []
|
||
assistant_positions = []
|
||
for index, line in enumerate(lines):
|
||
role = line.strip().lower().rstrip(":")
|
||
if role in ("system", "user", "assistant", "model"):
|
||
role_positions.append(index)
|
||
if role in ("assistant", "model"):
|
||
assistant_positions.append(index)
|
||
if role_positions:
|
||
if not assistant_positions:
|
||
return ""
|
||
lines = lines[assistant_positions[-1] + 1:]
|
||
kept_lines = []
|
||
for line in lines:
|
||
stripped = line.strip()
|
||
if not stripped:
|
||
kept_lines.append("")
|
||
continue
|
||
if stripped.lower().rstrip(":") in ("system", "user", "assistant", "model"):
|
||
continue
|
||
kept_lines.append(_ROLE_LINE_RE.sub("", stripped))
|
||
cleaned = "\n".join(kept_lines)
|
||
return re.sub(r"\n{3,}", "\n\n", cleaned).strip()
|
||
|
||
_NEGATIVE_PROMPT_TERMS = (
|
||
"blur", "blurry", "watermark", "logo", "extra fingers", "extra limbs",
|
||
"bad anatomy", "low quality", "worst quality", "deformed", "sweat",
|
||
"wetness", "moisture", "droplets",
|
||
)
|
||
_NEGATIVE_PROMPT_RE = re.compile(
|
||
r"\b(no|without|avoid|free of|don't|do not|must not|should not)\b[^.!?]{0,50}\b("
|
||
+ "|".join(re.escape(term) for term in _NEGATIVE_PROMPT_TERMS)
|
||
+ r")\b",
|
||
re.I,
|
||
)
|
||
|
||
|
||
def _clean_prompt_output(text):
|
||
"""Remove leaked instructions, reasoning and self-commentary while keeping the prompt."""
|
||
text = _strip_chat_tokens(text)
|
||
kept_lines = []
|
||
for line in text.splitlines():
|
||
stripped = line.strip()
|
||
if not stripped:
|
||
kept_lines.append("")
|
||
continue
|
||
if re.match(r"^negative prompt\s*[:\-]", stripped, re.I):
|
||
continue
|
||
stripped = re.sub(r"^(?:enhanced prompt|positive prompt|prompt)\s*[:\-]\s*", "", stripped, flags=re.I)
|
||
stripped = re.sub(r"^(?:okay|ok|sure|certainly|alright|absolutely|great|perfect)\s*[,!:.]\s*", "", stripped, flags=re.I)
|
||
stripped = re.sub(r"^(?:here is|here's|below is)\s+(?:the|your|an?)\s+(?:enhanced|final|generated|resulting)?\s*(?:prompt|output|response)\s*[:\-]?\s*", "", stripped, flags=re.I)
|
||
if not stripped:
|
||
continue
|
||
sentences = re.split(r"(?<=[.!?])\s+", stripped)
|
||
kept = [sentence.strip() for sentence in sentences if sentence.strip() and not _META_RE.search(sentence)]
|
||
if kept:
|
||
kept_lines.append(" ".join(kept))
|
||
cleaned = "\n".join(kept_lines)
|
||
cleaned = "\n".join(line for line in cleaned.splitlines() if not line.strip() or re.search(r"[A-Za-z0-9]", line))
|
||
cleaned = cleaned.replace("**", "").replace("__", "")
|
||
cleaned = re.sub(r"^\s*[-*•]\s+", "", cleaned, flags=re.M)
|
||
cleaned = re.sub(r"\n{3,}", "\n\n", cleaned).strip()
|
||
if not re.search(r"[A-Za-z0-9]", cleaned):
|
||
return ""
|
||
return cleaned.strip('"').strip("`").strip()
|
||
|
||
THINK_TAGS = (
|
||
("\u27e8think\u27e9", "\u27e8/think\u27e9"),
|
||
("<think>", "</think>"),
|
||
("<thinking>", "</thinking>"),
|
||
)
|
||
|
||
REFUSAL_MARKERS = (
|
||
"i cannot fulfill", "i can't fulfill", "i cannot comply", "i can't comply",
|
||
"i cannot assist", "i can't assist", "i cannot help", "i can't help",
|
||
"i cannot create", "i can't create", "i cannot generate", "i can't generate",
|
||
"cannot provide", "can't provide", "unable to comply", "unable to help",
|
||
"i must decline", "i have to decline", "i won't", "i will not",
|
||
"as an ai", "against my guidelines", "violates my", "not able to help",
|
||
"not appropriate", "inappropriate content", "sorry, but i can",
|
||
"i'm sorry, but", "i am sorry, but", "i apologize, but",
|
||
)
|
||
|
||
|
||
def _hard_loop(text):
|
||
"""Only the unmistakable loops, used to stop a running sampler without cutting good text."""
|
||
normalized = re.sub(r"\s+", " ", text.lower()).strip()[-6000:]
|
||
words = re.findall(r"\w+", normalized)
|
||
if len(words) < 80:
|
||
return False
|
||
grams = [" ".join(words[i:i + 12]) for i in range(len(words) - 11)]
|
||
counts = Counter(grams)
|
||
return any(count >= 4 for count in counts.values())
|
||
|
||
|
||
def _is_repetitive_text(text):
|
||
"""True for degenerate loops; structured prompts that reuse a few phrases are fine."""
|
||
if not text or len(text) < 60:
|
||
return False
|
||
normalized = re.sub(r"\s+", " ", text.lower()).strip()[:8000]
|
||
sentences = [s.strip() for s in re.split(r"(?<=[.!?])\s+", normalized) if len(s.strip()) >= 30]
|
||
if sentences:
|
||
counts = Counter(sentences)
|
||
# a long prompt may reuse a sentence or two; a loop is mostly copies
|
||
if max(counts.values()) >= 5 or sum(len(s) * (c - 1) for s, c in counts.items()) > 0.3 * len(normalized):
|
||
return True
|
||
words = re.findall(r"\w+", normalized)
|
||
for size, min_count, share in ((8, 3, 0.5), (5, 4, 0.75)):
|
||
if len(words) < size * 4:
|
||
continue
|
||
grams = [" ".join(words[i:i + size]) for i in range(len(words) - size + 1)]
|
||
counts = Counter(grams)
|
||
covered = set()
|
||
for i, gram in enumerate(grams):
|
||
if counts[gram] >= min_count:
|
||
covered.update(range(i, i + size))
|
||
if len(covered) > share * len(words):
|
||
return True
|
||
lines = [line.strip().lower() for line in text.splitlines() if len(line.strip()) >= 30]
|
||
if lines:
|
||
counts = Counter(lines)
|
||
if max(counts.values()) >= 5 or sum(len(l) * (c - 1) for l, c in counts.items()) > 0.3 * len(text):
|
||
return True
|
||
return False
|
||
|
||
|
||
def _dedupe_repeats(text):
|
||
"""Drop repeated sentences and lines, keeping the first occurrence."""
|
||
seen = set()
|
||
lines = []
|
||
for line in text.splitlines():
|
||
kept = []
|
||
for sentence in re.split(r"(?<=[.!?])\s+", line.strip()):
|
||
key = re.sub(r"\s+", " ", sentence.lower()).strip()
|
||
if len(key) >= 30:
|
||
if key in seen:
|
||
continue
|
||
seen.add(key)
|
||
kept.append(sentence)
|
||
lines.append(" ".join(kept))
|
||
return re.sub(r"\n{3,}", "\n\n", "\n".join(lines)).strip()
|
||
|
||
|
||
_VIDEO_FIELD_RE = re.compile(
|
||
r"^(?:integrated_multimodal_description|overall_soundscape|non_diegetic_music|soundscape|audio(?:_design)?|sound_effects?|music)\s*:",
|
||
re.I,
|
||
)
|
||
_VIDEO_SENTENCE_RE = re.compile(
|
||
r"\b(?:non-diegetic music|soundscape|sound effects?|audio cue|voice-?over|background music|diegetic sound)\b",
|
||
re.I,
|
||
)
|
||
_SHOT_LINE_RE = re.compile(r"^\s*(?:\[\s*)?shot\s+\d+", re.I)
|
||
|
||
|
||
def _strip_video_meta(text):
|
||
"""Remove video/audio fields from image-only presets."""
|
||
lines = []
|
||
for line in text.splitlines():
|
||
stripped = line.strip()
|
||
if not stripped:
|
||
lines.append("")
|
||
continue
|
||
if _VIDEO_FIELD_RE.match(stripped) or _SHOT_LINE_RE.match(stripped):
|
||
continue
|
||
sentences = re.split(r"(?<=[.!?])\s+", stripped)
|
||
kept = [sentence for sentence in sentences if not _VIDEO_SENTENCE_RE.search(sentence)]
|
||
if kept:
|
||
lines.append(" ".join(kept))
|
||
cleaned = "\n".join(lines)
|
||
return re.sub(r"\n{3,}", "\n\n", cleaned).strip()
|
||
|
||
|
||
def _next_seed(seed):
|
||
return (int(seed) + 1) % (2 ** 32 - 1) or 1
|
||
|
||
|
||
_TOP_P_WARNED = False
|
||
|
||
|
||
def _clip_top_p_supported():
|
||
try:
|
||
major, minor = (int(part) for part in str(torch.__version__).split("+")[0].split(".")[:2])
|
||
return (major, minor) < (2, 13)
|
||
except Exception:
|
||
return True
|
||
|
||
|
||
_REASONING_OPENER_RE = re.compile(
|
||
r"^\W*(?:okay|ok|alright|let me|let's|first,|the user (?:wants|asked|is asking|provided)|i need to|we need to|i should|"
|
||
r"breaking (?:this|it) down|to (?:start|begin) with|so,? the)\b",
|
||
re.I,
|
||
)
|
||
|
||
|
||
def _bad_output_reason(text, min_chars=20):
|
||
"""Return a short reason when output is empty, short, degenerate, leaked reasoning or a refusal."""
|
||
answer, _ = _split_thinking(text or "")
|
||
if _REASONING_OPENER_RE.match(answer):
|
||
return "reasoning leak"
|
||
clean = _clean_prompt_output(answer)
|
||
if len(clean) < min_chars:
|
||
return "too short"
|
||
if _is_repetitive_text(clean):
|
||
return "repetitive"
|
||
if _NEGATIVE_PROMPT_RE.search(clean):
|
||
return "negative-prompt wording"
|
||
if re.search(r"(.)\1{15,}", clean):
|
||
return "runaway repeat"
|
||
if len(clean) > 200 and len(set(clean)) < 8:
|
||
return "low variety"
|
||
words = re.findall(r"\w+", clean.lower())
|
||
if len(words) > 30:
|
||
if Counter(words).most_common(1)[0][1] / len(words) > 0.45:
|
||
return "repeated words"
|
||
lines = [line.strip() for line in clean.splitlines() if line.strip()]
|
||
if len(lines) >= 4:
|
||
if Counter(lines).most_common(1)[0][1] >= 4:
|
||
return "repeated lines"
|
||
lowered = clean.lower()
|
||
if len(clean) < 600 and any(marker in lowered for marker in REFUSAL_MARKERS):
|
||
return "refusal"
|
||
if any(lowered.startswith(marker) for marker in REFUSAL_MARKERS):
|
||
return "refusal"
|
||
return ""
|
||
|
||
|
||
_MINIMAX_SECTIONS = {
|
||
"ref2va": ("subject_definitions", "summary", "retention_analysis", "detailed_description", "overall_soundscape", "non_diegetic_music"),
|
||
"fl2va": ("integrated_multimodal_description", "overall_soundscape", "non_diegetic_music"),
|
||
"i2va": ("integrated_multimodal_description", "overall_soundscape", "non_diegetic_music"),
|
||
"l2va": ("integrated_multimodal_description", "overall_soundscape", "non_diegetic_music"),
|
||
"t2va": ("integrated_multimodal_description", "overall_soundscape", "non_diegetic_music"),
|
||
}
|
||
|
||
|
||
_MINIMAX_FALLBACK = {
|
||
"subject_definitions": "<Subject 1> is the referenced person.",
|
||
"summary": "The clip shows the requested action.",
|
||
"retention_analysis": "<Subject 1>: fully_preserved - appearance retained.",
|
||
"detailed_description": "The requested action runs from the first frame to the last.",
|
||
"integrated_multimodal_description": "[Shot 1] The requested action runs from the first frame to the last.",
|
||
"overall_soundscape": "Breathing, skin contact and movement sounds throughout.",
|
||
"non_diegetic_music": "N/A",
|
||
}
|
||
|
||
|
||
def _ensure_minimax_sections(text, mode):
|
||
"""Append any section the model left out, so the prompt keeps the whole contract."""
|
||
for name in _MINIMAX_SECTIONS.get(mode, ()):
|
||
if not re.search(rf"(?im)^{name}:", text):
|
||
text = text.rstrip() + f"\n\n{name}:\n{_MINIMAX_FALLBACK[name]}"
|
||
return text
|
||
|
||
|
||
def _normalize_minimax_sections(text, mode):
|
||
"""Rewrite the section labels to the exact lowercase contract, one per line."""
|
||
for name in _MINIMAX_SECTIONS.get(mode, ()):
|
||
text = re.sub(rf"(?im)^{name}:\s*", name + ":\n", text, count=1)
|
||
return text
|
||
|
||
|
||
def _strip_inner_timestamps(text, mode="ref2va"):
|
||
"""Fold timestamps inside a single-shot action prompt into plain prose and drop repeated [Shot 1] labels."""
|
||
marker = "detailed_description:" if mode == "ref2va" else "integrated_multimodal_description:"
|
||
cut = text.find(marker)
|
||
head, body = (text[:cut], text[cut:]) if cut >= 0 else ("", text)
|
||
body = body.replace(marker, marker + "\n", 1)
|
||
start = body.find("\n") + 1
|
||
segment = body[start:]
|
||
first = segment.find("[Shot 1]")
|
||
if first >= 0:
|
||
segment = segment[:first + 8] + re.sub(r"\[Shot 1\]", "", segment[first + 8:])
|
||
else:
|
||
segment = "[Shot 1] " + segment.lstrip()
|
||
body = body[:start] + segment
|
||
body = re.sub(r"\b(?:From|Between)\s+\d{2}:\d{2}\.\d{3}\s+(?:to|and)\s+\d{2}:\d{2}\.\d{3}\s*,?\s*", "Then ", body)
|
||
body = re.sub(r"\b(?:[Aa]t|[Bb]y|[Uu]ntil|[Bb]efore|[Aa]fter)\s+\d{2}:\d{2}\.\d{3}\s*,?\s*", "", body)
|
||
body = re.sub(r"\b(?:[Aa]t|[Aa]round|[Bb]y|[Aa]fter|[Bb]efore|[Uu]ntil|[Ff]rom)\s+\d{1,2}(?:\.\d+)?\s*(?:seconds?|s)\b\s*,?\s*", "", body)
|
||
body = re.sub(r"\b(?:[Aa]pproximately|[Rr]oughly|[Aa]bout)\s+\d{1,2}(?:\.\d+)?\s*seconds?\b\s*,?\s*", "", body)
|
||
body = re.sub(r"\b(?:[Aa]t|[Aa]round|[Bb]y|[Uu]ntil|[Bb]efore|[Aa]fter)\s*,", "", body)
|
||
body = re.sub(r"\(\s*[–-]?\s*\)", "", body)
|
||
body = re.sub(r"\d{2}:\d{2}\.\d{3}", "", body)
|
||
body = re.sub(r"((?:^|[.!?]\s+|\n)|\[Shot 1\]\s+)([a-z])", lambda match: match.group(1) + match.group(2).upper(), body)
|
||
return head + re.sub(r"[ \t]{2,}", " ", body)
|
||
|
||
|
||
def _inject_reference_labels(text, user_text):
|
||
"""Attach the source picture/video to every subject line that forgot to cite it."""
|
||
sources = {}
|
||
for match in re.finditer(r"(?im)^Subject (\d+)\b[^:\n]*:\s*.*?(<Picture\s+\d+>|<Video\s+\d+>)", user_text):
|
||
sources[int(match.group(1))] = match.group(2)
|
||
if not sources:
|
||
return text
|
||
start = re.search(r"(?im)^subject_definitions:", text)
|
||
if not start:
|
||
return text
|
||
end = re.search(r"(?im)^summary:", text[start.end():])
|
||
body = text[start.end():start.end() + end.start()] if end else text[start.end():]
|
||
|
||
def fix(match):
|
||
number = int(match.group(1))
|
||
line = match.group(0)
|
||
if number in sources and not re.search(r"<(?:Picture|Video)\s+\d+>", line):
|
||
line = line.rstrip(" .") + f" from {sources[number]}."
|
||
return line
|
||
|
||
body = re.sub(r"(?im)^<Subject (\d+)>[^\n]*", fix, body)
|
||
return text[:start.end()] + body + (text[start.end() + end.start():] if end else "")
|
||
|
||
|
||
def _minimax_format_problem(text, mode, one_shot, hint="", picture_count=0):
|
||
"""Return what breaks the H3 output contract, or '' when the prompt is usable."""
|
||
sections = _MINIMAX_SECTIONS.get(mode)
|
||
if not sections:
|
||
return ""
|
||
positions = []
|
||
for name in sections:
|
||
match = re.search(rf"^{name}:", text, re.M | re.I)
|
||
if not match:
|
||
return f"missing section {name}"
|
||
positions.append(match.start())
|
||
if positions != sorted(positions):
|
||
return "sections out of order"
|
||
if mode == "ref2va" and picture_count:
|
||
block = text[positions[0]:positions[1]] if len(positions) > 1 else text[positions[0]:]
|
||
first = re.search(r"(?im)^<Subject 1>.*$", block)
|
||
if first and not re.search(r"<(?:Picture|Video)\s+\d+>", first.group(0)):
|
||
return "Subject 1 does not cite its reference picture"
|
||
if not re.search(r"<Picture\s+\d+>", block):
|
||
return "subject_definitions cites no reference picture"
|
||
body = text[positions[sections.index("detailed_description")]:] if "detailed_description" in sections else text[positions[0]:]
|
||
body = body.split("overall_soundscape:")[0]
|
||
shots = [int(number) for number in re.findall(r"\[Shot (\d+)\]", body)]
|
||
if one_shot:
|
||
if shots != [1]:
|
||
return "the action must be a single [Shot 1]"
|
||
if re.search(r"\d{2}:\d{2}\.\d{3}", text):
|
||
return "timestamps inside the single shot"
|
||
elif shots != list(range(1, len(shots) + 1)):
|
||
return "shot numbers must start at 1 and rise by one without repeats"
|
||
elif len(shots) < 2 and re.search(r"\d{2}:\d{2}\.\d{3}", body):
|
||
return "timestamps without a new shot label"
|
||
if one_shot and "striptease" not in hint.lower():
|
||
opening = " ".join(re.findall(r"\w+", text.split("[Shot 1]", 1)[-1])[:150]).lower()
|
||
if not any(word in opening for word in ("penis", "vagina", "vulva", "clitoris", "shaft", "glans", "genitals", "tongue", "mouth")):
|
||
return "the act starts too late (contact must be in the opening sentences)"
|
||
words = len(re.findall(r"\w+", body))
|
||
if words < 450:
|
||
return f"body too short ({words} words, 500-800 required)"
|
||
if words > 950:
|
||
return f"body too long ({words} words, 500-800 required)"
|
||
return ""
|
||
|
||
|
||
RUNAWAY_LOOP = 6 # the same word this many times in a row
|
||
RUNAWAY_SENTENCE = 110 # words in one sentence; real prompt sentences are far shorter
|
||
RUNAWAY_KEEP = 35 # words of a runaway sentence that are still kept
|
||
|
||
|
||
def _land_text(text):
|
||
"""End on a complete sentence, so a token cap never leaves a half word or clause."""
|
||
text = text.rstrip()
|
||
if not text or text[-1] in '.!?…"\')]':
|
||
return text
|
||
for index in range(len(text) - 1, 0, -1):
|
||
if text[index] in '.!?':
|
||
return text[:index + 1].rstrip()
|
||
return text + "."
|
||
|
||
|
||
def _cut_runaway_tail(text):
|
||
"""Cut the degenerate tail of a prompt the model lost track of: a looping word, or one endless run-on sentence. Returns (text, whether something was cut)."""
|
||
cut = None
|
||
words = list(re.finditer(r"\S+", text))
|
||
run = 1
|
||
for index in range(1, len(words)):
|
||
run = run + 1 if words[index].group().lower().strip(".,;:!?") == words[index - 1].group().lower().strip(".,;:!?") else 1
|
||
if run >= RUNAWAY_LOOP:
|
||
cut = words[index - run + 2].start() # one copy of the word stays
|
||
break
|
||
for sentence in re.finditer(r"[^.!?\n]+", text):
|
||
inside = list(re.finditer(r"\S+", sentence.group()))
|
||
if len(inside) > RUNAWAY_SENTENCE:
|
||
start = sentence.start() + inside[RUNAWAY_KEEP].start()
|
||
cut = start if cut is None else min(cut, start)
|
||
break
|
||
if cut is None:
|
||
return text, False
|
||
return text[:cut].rstrip(" ,;:-") + ".", True
|
||
|
||
|
||
def _generate_checked(attempt, seed, min_chars, auto_retry, tries=3):
|
||
"""Run attempt(seed) until the output passes the checks, up to `tries` seeds; keep the best on failure."""
|
||
texts = []
|
||
for index in range(tries if auto_retry else 1):
|
||
text = attempt(seed)
|
||
reason = _bad_output_reason(_split_thinking(text)[0], min_chars)
|
||
if not reason:
|
||
return text
|
||
texts.append(text)
|
||
seed = _next_seed(seed)
|
||
if index + 1 < tries and auto_retry:
|
||
print(f"[PromptForge] output rejected ({reason}); retrying with a new seed ({index + 2}/{tries}).")
|
||
if reason == "repetitive":
|
||
return _dedupe_repeats(max(texts, key=len))
|
||
return texts[-1]
|
||
|
||
|
||
# MiniMax H3 modes: how many images the mode can use, whether a reference video is allowed, and each image's role.
|
||
# Labels follow the H3 text encoder: <Picture i> / <Video k> are 1-based per type in connection order, gaps closed up.
|
||
_MINIMAX_MODES = {
|
||
"fl2va": (2, False, ("the exact first frame", "the exact last frame")),
|
||
"i2va": (1, False, ("the exact first frame",)),
|
||
"l2va": (1, False, ("the exact last frame",)),
|
||
"ref2va": (None, True, ()),
|
||
"t2va": (0, False, ()),
|
||
}
|
||
_MINIMAX_MODE_RE = re.compile(r"minimax h3 (" + "|".join(_MINIMAX_MODES) + r")\b", re.I)
|
||
|
||
|
||
# languages a video family speaks reliably (third-party model cards; other languages work partially)
|
||
DIALOGUE_LANGUAGES = ("Arabic", "Chinese", "English", "French", "German", "Italian", "Japanese", "Korean", "Portuguese", "Russian", "Spanish")
|
||
DIALOGUE_AUTO = "auto (from the instructions)"
|
||
DIALOGUE_OTHER = "other (type below)"
|
||
# these are counted in characters; every other language in words
|
||
CHARACTER_LANGUAGES = {"Chinese", "Japanese", "Korean"}
|
||
|
||
|
||
def _dialogue_language(language, other):
|
||
"""The language name to write in, or '' when it is left to the instructions."""
|
||
if language == DIALOGUE_OTHER:
|
||
return (other or "").strip()
|
||
return "" if language == DIALOGUE_AUTO else language
|
||
|
||
|
||
def _dialogue_block(instructions, language, other, seconds):
|
||
"""System prompt section asking for dialogue in the video prompt, sized to what fits in `seconds`; '' when off."""
|
||
instructions = (instructions or "").strip()
|
||
chosen = _dialogue_language(language, other)
|
||
if not instructions and not chosen:
|
||
return ""
|
||
unit, rate = ("characters", 4.0) if chosen in CHARACTER_LANGUAGES else ("words", 2.3)
|
||
budget = max(3, round(max(seconds - 1.0, 1.0) * rate))
|
||
name = chosen or "LANGUAGE"
|
||
where = f"in {chosen}" if chosen else "in the language the dialogue instructions imply (English when they do not say)"
|
||
topic = f"Dialogue instructions from the user: {instructions}" if instructions else "The user gave no dialogue instructions, so dialogue is optional: write spoken lines only when the scene calls for them. A video without speech is fine."
|
||
return f"""# Dialogue ({"required in this video" if instructions else "optional, only the language is set"})
|
||
{topic}
|
||
Write every spoken line {where}. Only the spoken words are in that language; the rest of the prompt stays in English. The words are what a native speaker would really say, never a translation.
|
||
|
||
Exact shape of every spoken line, copy it:
|
||
the young woman with a quiet, breathy voice (S1) says: <d>[{name}] the spoken words in {name}</d>
|
||
|
||
Rules, all mandatory:
|
||
1. The words are wrapped in <d> and </d>. The closing </d> is never left out.
|
||
2. Directly after <d> comes the language tag in square brackets, [{name}], then a space, then the words. The tag always holds the English name of the language, spelled the same way every time, even when the words are not English.
|
||
3. Between the tags there is only speech in {name}: no quotation marks, no stage directions, no translation, no English words, no emoji.
|
||
4. Name the speaker before every line with the ID: first appearance with a short voice description ("the tired baker with a low, gravelly voice (S1) says:"), then "(S1) says:" or "(S2) answers:". People who do not speak get no ID.
|
||
5. Write each line inside the shot where it is spoken, at its timestamp.
|
||
6. Timing: the video lasts {seconds:g} seconds and people speak about {rate:g} {unit} per second. Leave a second before the first line and a short pause between speakers, so all speakers together say at most {budget} {unit}. Every line ends before the video does. Short natural exchanges beat speeches; say less when the topic needs more than fits.
|
||
|
||
Wrong, never write: She says "<words>" | <d><words></d> (no language tag) | <d>[{name}] <words> (no closing tag) | <d>[{name}] <words in English></d> (words in the wrong language) | [{name}] <d><words></d> (tag outside)
|
||
Before you answer, check every <d>: it has [{name}] right after it, the words are {name}, and </d> closes it."""
|
||
|
||
|
||
DIALOGUE_TAG = re.compile(r"<d>\s*\[([^\]]*)\]\s*(.+?)\s*</d>", re.S)
|
||
|
||
|
||
def _dialogue_problem(text, language):
|
||
"""What is wrong with the dialogue tags in `text`, or '' when they are right. `language` '' skips the language check."""
|
||
tags = DIALOGUE_TAG.findall(text)
|
||
if not tags:
|
||
return "no <d>[Language] words</d> dialogue line was written" if "<d>" not in text and "</d>" not in text else "the <d> tags are incomplete: every line needs <d>[Language] words</d>"
|
||
if text.count("<d>") != text.count("</d>") or text.count("<d>") != len(tags):
|
||
return "some <d> tags are missing their [Language] tag or their closing </d>"
|
||
if language:
|
||
wrong = next((tag for tag, _ in tags if tag.strip().lower() != language.lower()), None)
|
||
if wrong is not None:
|
||
return f"a language tag says [{wrong}] but it must be [{language}]"
|
||
return ""
|
||
|
||
|
||
FAMILY_PRESETS = {
|
||
"krea2": "Krea 2 T2I",
|
||
"moss_image": "MOSS Image T2I",
|
||
"qwen_image": "Qwen Image T2I",
|
||
"anima": "Anima Tags",
|
||
"ace": "ACE-Step 1.5",
|
||
}
|
||
|
||
|
||
def _family_preset(family, has_images, has_video):
|
||
"""Preset for a model family when the preset is on auto; MiniMax H3 follows the connected media."""
|
||
if family == "minimax_h3":
|
||
return "MiniMax H3 Ref2VA" if has_images or has_video else "MiniMax H3 T2VA"
|
||
return FAMILY_PRESETS.get(family, DEFAULT_PRESET)
|
||
|
||
|
||
def _minimax_mode(preset_name):
|
||
match = _MINIMAX_MODE_RE.search(preset_name or "")
|
||
return match.group(1).lower() if match else None
|
||
|
||
|
||
def _minimax_alignment_line(mode_key, picture_count, seconds):
|
||
"""The frame-alignment line the frame modes open with, limited to the pictures that are attached ('' = omit it)."""
|
||
if not picture_count:
|
||
return ""
|
||
head = "How the reference pictures align with the target video — "
|
||
if mode_key == "i2va":
|
||
return "For the target video, at 0.00 seconds into the target video, <Picture 1> (from [Shot 1]) is fully referenced."
|
||
if mode_key == "l2va":
|
||
return head + f"<Picture 1> (from [Shot N]) aligns with the {seconds}-second mark of the target video."
|
||
if mode_key == "fl2va":
|
||
first = "Picture 1 (from Shot 1) aligns with the 0.00-second mark of the target video"
|
||
if picture_count == 1:
|
||
return head + first + "."
|
||
return head + first + f"; Picture 2 (from Shot N) aligns with the {seconds}-second mark of the target video."
|
||
return ""
|
||
|
||
|
||
_MINIMAX_REF_RE = re.compile(r"\b(Picture|Video|Audio)\s+(\d+)\b")
|
||
|
||
|
||
def _drop_absent_references(text, picture_count, has_video, user_text):
|
||
"""Remove sentences citing a picture, video or audio label that is not attached (audio the user's own text names is allowed)."""
|
||
def absent(match):
|
||
kind, number = match.group(1), int(match.group(2))
|
||
if kind == "Picture":
|
||
return number > picture_count
|
||
if kind == "Video":
|
||
return number > (1 if has_video else 0)
|
||
return not re.search(rf"\bAudio\s+{number}\b", user_text)
|
||
|
||
removed = 0
|
||
lines = []
|
||
for line in text.splitlines():
|
||
kept = []
|
||
for sentence in re.split(r"(?<=[.!?])\s+", line):
|
||
if any(absent(m) for m in _MINIMAX_REF_RE.finditer(sentence)):
|
||
removed += 1
|
||
else:
|
||
kept.append(sentence)
|
||
lines.append(" ".join(kept))
|
||
if removed:
|
||
print(f"[PromptForge] removed {removed} sentence(s) that cited media which is not connected.")
|
||
return re.sub(r"\n{3,}", "\n\n", "\n".join(lines)).strip()
|
||
|
||
|
||
def _minimax_inputs_block(mode_key, picture_count, video_frames, dropped, seconds):
|
||
"""Authoritative list of what exists for this MiniMax mode; everything else must never be mentioned."""
|
||
max_images, allows_video, roles = _MINIMAX_MODES[mode_key]
|
||
lines = ["# MiniMax inputs (authoritative)", f"Mode: {mode_key.upper()}."]
|
||
present = []
|
||
for index in range(picture_count):
|
||
role = f" ({roles[index]})" if index < len(roles) else ""
|
||
present.append(f"<Picture {index + 1}>{role}")
|
||
if video_frames:
|
||
present.append(f"<Video 1> (one reference video, {video_frames} frames sampled in time order)")
|
||
lines.append("Attached: " + (", ".join(present) if present else "nothing") + ".")
|
||
absent = []
|
||
if max_images is None:
|
||
absent.append(f"<Picture {picture_count + 1}> and any higher picture number")
|
||
else:
|
||
absent.extend(f"<Picture {index + 1}> ({roles[index]})" for index in range(picture_count, max_images))
|
||
if max_images == 0 or picture_count >= max_images:
|
||
absent.append(f"<Picture {max(picture_count, max_images) + 1}> and any higher picture number")
|
||
if not video_frames:
|
||
absent.append("<Video 1> and any other video")
|
||
absent.append("any <Audio N> that the user's text does not itself describe (this node attaches no audio)")
|
||
lines.append("Not attached, so never mention, cite or invent: " + "; ".join(absent) + ".")
|
||
lines.append(
|
||
"Where the instructions above refer to something that is not attached, ignore that part and work only from the user's text. "
|
||
"Use each attached label exactly as written and never renumber it."
|
||
)
|
||
if mode_key in ("fl2va", "i2va", "l2va"):
|
||
alignment = _minimax_alignment_line(mode_key, picture_count, seconds)
|
||
if alignment:
|
||
lines.append(f"Alignment line: write exactly this, replacing N with the number of the last shot, and name no other picture in it: {alignment}")
|
||
else:
|
||
lines.append("Alignment line: nothing is attached, so omit the alignment line entirely and start with the first field.")
|
||
if dropped:
|
||
lines.append("Connected but unused by this mode (do not mention): " + ", ".join(dropped) + ".")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def _tokens_text(clip, tokens):
|
||
batch = next(iter(tokens.values())) if isinstance(tokens, dict) else tokens
|
||
return clip.decode([t[0] for row in batch for t in row if isinstance(t[0], int)], skip_special_tokens=False)
|
||
|
||
|
||
def _tokens_have_images(tokens):
|
||
if isinstance(tokens, dict):
|
||
if tokens.get("type") == "image":
|
||
return True
|
||
return any(_tokens_have_images(value) for value in tokens.values())
|
||
if isinstance(tokens, (list, tuple)):
|
||
return any(_tokens_have_images(value) for value in tokens)
|
||
return False
|
||
|
||
|
||
def _split_thinking(text):
|
||
"""Remove reasoning blocks from model output and return (answer, thinking).
|
||
|
||
The opening think tag can come from the prompt template rather than the
|
||
generated text, so a stray closing tag also splits reasoning from answer.
|
||
"""
|
||
answer = text or ""
|
||
thoughts = []
|
||
channel_thought = re.findall(
|
||
r"(?is)<\|channel>\s*thought\b(.*?)(?=<\|channel>\s*final\b|<channel\|>|$)", answer
|
||
)
|
||
thoughts.extend(part.strip() for part in channel_thought if part.strip())
|
||
answer = re.sub(
|
||
r"(?is)<\|channel>\s*thought\b.*?(?=<\|channel>\s*final\b|<channel\|>|$)", "", answer
|
||
)
|
||
channel_final = re.search(r"(?is)<\|channel>\s*final\b(.*?)(?:<channel\|>|$)", answer)
|
||
if channel_final:
|
||
answer = channel_final.group(1)
|
||
answer = answer.replace("<|channel>", "").replace("<channel|>", "")
|
||
for open_tag, close_tag in THINK_TAGS:
|
||
pattern = r"(?is)" + re.escape(open_tag) + r".*?" + re.escape(close_tag)
|
||
for match in re.finditer(pattern, answer):
|
||
inner = match.group(0).replace(open_tag, "").replace(close_tag, "").strip()
|
||
if inner:
|
||
thoughts.append(inner)
|
||
answer = re.sub(pattern, "", answer)
|
||
|
||
lowered = answer.lower()
|
||
split_idx = -1
|
||
split_len = 0
|
||
for _, close_tag in THINK_TAGS:
|
||
idx = lowered.rfind(close_tag.lower())
|
||
if idx > split_idx:
|
||
split_idx, split_len = idx, len(close_tag)
|
||
if split_idx != -1:
|
||
head = answer[:split_idx].strip()
|
||
if head:
|
||
thoughts.append(head)
|
||
answer = answer[split_idx + split_len:]
|
||
else:
|
||
open_idx = -1
|
||
open_len = 0
|
||
for open_tag, _ in THINK_TAGS:
|
||
idx = lowered.rfind(open_tag.lower())
|
||
if idx > open_idx:
|
||
open_idx, open_len = idx, len(open_tag)
|
||
if open_idx != -1:
|
||
tail = answer[open_idx + open_len:].strip()
|
||
if tail:
|
||
thoughts.append(tail)
|
||
answer = answer[:open_idx]
|
||
return answer.strip(), "\n\n".join(thoughts).strip()
|
||
|
||
|
||
DEFAULT_PRESET = "Enhance"
|
||
|
||
HERETIC_MODELS = PROMPTFORGE_MODELS.get("heretic_hf_models") or []
|
||
GGUF_MODELS = PROMPTFORGE_MODELS.get("gguf_models") or []
|
||
LOCAL_SCAN_DIRS = PROMPTFORGE_MODELS.get("local_scan_dirs") or []
|
||
|
||
def _safetensors_is_decoder(path):
|
||
"""True when a single safetensors file holds a decoder LLM (has token embeddings)."""
|
||
try:
|
||
with open(path, "rb") as f:
|
||
header = json.loads(f.read(struct.unpack("<Q", f.read(8))[0]))
|
||
except (OSError, ValueError, struct.error):
|
||
return False
|
||
return any(key.endswith("embed_tokens.weight") for key in header)
|
||
|
||
|
||
def _scan_roots():
|
||
"""Directories to scan: every registered text_encoders folder plus the catalog subdirs of models_dir."""
|
||
roots = []
|
||
if folder_paths is not None:
|
||
for folder in folder_paths.get_folder_paths("text_encoders"):
|
||
roots.append(Path(folder))
|
||
models_dir = Path(folder_paths.models_dir)
|
||
roots.extend(models_dir / sub for sub in LOCAL_SCAN_DIRS)
|
||
seen = set()
|
||
unique = []
|
||
for root in roots:
|
||
if root.is_dir() and str(root) not in seen:
|
||
seen.add(str(root))
|
||
unique.append(root)
|
||
return unique
|
||
|
||
|
||
def _scan_local_text_models():
|
||
"""Scan ComfyUI text-model directories for local LLM / text encoder models.
|
||
|
||
Multi-part checkpoints collapse to their containing HF directory, mmproj
|
||
vision sidecars and non-generative encoders (CLIP, T5, projections) are
|
||
skipped. The picker option is the path relative to models_dir so a
|
||
workflow saved on one machine resolves on another with the same layout.
|
||
"""
|
||
found = []
|
||
seen_paths = set()
|
||
models_dir = Path(folder_paths.models_dir) if folder_paths is not None else None
|
||
for root in _scan_roots():
|
||
for p in sorted(root.rglob("*")):
|
||
if not p.is_file() or p.suffix.lower() not in (".safetensors", ".gguf", ".bin"):
|
||
continue
|
||
if p.suffix.lower() == ".gguf" and "mmproj" in p.name.lower():
|
||
continue
|
||
if p.suffix.lower() in (".safetensors", ".bin") and (p.parent / "config.json").is_file():
|
||
model_path = p.parent
|
||
elif p.suffix.lower() == ".safetensors" and not _safetensors_is_decoder(p):
|
||
continue
|
||
else:
|
||
model_path = p
|
||
key = str(model_path)
|
||
if key in seen_paths:
|
||
continue
|
||
seen_paths.add(key)
|
||
option = key
|
||
if models_dir is not None and model_path.is_relative_to(models_dir):
|
||
option = str(model_path.relative_to(models_dir))
|
||
found.append({"path": key, "option": option, "name": model_path.name, "kind": "gguf" if p.suffix.lower() == ".gguf" else "safetensors"})
|
||
return found
|
||
|
||
|
||
LOCAL_TEXT_MODELS = _scan_local_text_models()
|
||
_SCANNED_GGUF_NAMES = {Path(m.get("path") or "").name for m in LOCAL_TEXT_MODELS if str(m.get("path") or "").lower().endswith(".gguf")}
|
||
|
||
def _catalog_entry(item, kind):
|
||
return {
|
||
"option": item.get("label") or item.get("name") or item.get("key"),
|
||
"kind": kind,
|
||
"key": item.get("key"),
|
||
"params_b": item.get("params_b"),
|
||
"avail_match": item.get("avail_match") or [],
|
||
"repo_id": item.get("repo_id"),
|
||
"base_dir": item.get("base_dir"),
|
||
"quant_files": item.get("quant_files") or {},
|
||
"mmproj_file": item.get("mmproj_file"),
|
||
"local_path": item.get("local_path") or "",
|
||
}
|
||
|
||
|
||
HERETIC_MODEL_ENTRIES = (
|
||
[_catalog_entry(item, "hf") for item in HERETIC_MODELS]
|
||
+ [_catalog_entry(item, "gguf") for item in GGUF_MODELS]
|
||
)
|
||
|
||
|
||
def _cache_has(repo_id, filename):
|
||
if not repo_id or not filename:
|
||
return False
|
||
cache = os.environ.get("HF_HUB_CACHE") or os.environ.get("HUGGINGFACE_HUB_CACHE")
|
||
hub = Path(cache) if cache else Path(os.environ.get("HF_HOME") or Path.home() / ".cache" / "huggingface") / "hub"
|
||
snapshots = hub / f"models--{repo_id.replace('/', '--')}" / "snapshots"
|
||
if not snapshots.is_dir():
|
||
return False
|
||
return any((snapshot / filename).is_file() for snapshot in snapshots.iterdir())
|
||
|
||
|
||
def _gguf_local_quants(entry):
|
||
quant_files = entry.get("quant_files") or {}
|
||
if not quant_files:
|
||
return []
|
||
base_dir = Path(entry.get("base_dir") or "LLM/GGUF")
|
||
if folder_paths is not None:
|
||
models_dir = getattr(folder_paths, "models_dir", "") or ""
|
||
if models_dir:
|
||
base_dir = Path(models_dir) / base_dir
|
||
scanned = _SCANNED_GGUF_NAMES
|
||
repo_id = entry.get("repo_id")
|
||
return [
|
||
quant for quant, filename in quant_files.items()
|
||
if (base_dir / filename).is_file() or filename in scanned or _cache_has(repo_id, filename)
|
||
]
|
||
|
||
|
||
def _entry_local_path(entry):
|
||
direct = entry.get("local_path") or ""
|
||
if direct and Path(direct).exists():
|
||
return direct
|
||
matches = entry.get("avail_match") or []
|
||
if isinstance(matches, str):
|
||
matches = [matches]
|
||
matches = [m.lower() for m in matches if m]
|
||
if not matches:
|
||
return ""
|
||
for m in LOCAL_TEXT_MODELS:
|
||
name = (m.get("name") or "").lower()
|
||
if not all(token in name for token in matches):
|
||
continue
|
||
path = m.get("path") or ""
|
||
if entry.get("kind") == "hf" and path.lower().endswith(".gguf"):
|
||
continue
|
||
if entry.get("kind") == "gguf" and not path.lower().endswith(".gguf"):
|
||
continue
|
||
return path
|
||
return ""
|
||
|
||
|
||
def _vision_tag(entry):
|
||
"""Picker tag: does this GGUF have a matching vision projector, locally or in the catalog?"""
|
||
local = _entry_local_path(entry)
|
||
if local and any(Path(local).parent.glob("*mmproj*.gguf")):
|
||
return " [VISION]"
|
||
mmproj = entry.get("mmproj_file")
|
||
if mmproj and _cache_has(entry.get("repo_id"), mmproj):
|
||
return " [VISION]"
|
||
return " [VISION dl]" if mmproj else " [TEXT]"
|
||
|
||
|
||
def _entry_label(entry):
|
||
"""(picker label, is_local): [LOCAL ...] for models on disk or in the HF cache, [DOWNLOAD] for the rest."""
|
||
label = entry["option"]
|
||
if entry.get("kind") == "gguf":
|
||
quants = _gguf_local_quants(entry)
|
||
tag = _vision_tag(entry)
|
||
if quants:
|
||
ordered = [q for q in ("q4", "q5", "q6", "q8") if q in quants]
|
||
return f"[LOCAL {'/'.join(q.upper() for q in ordered)}]{tag} {label}", True
|
||
return f"[DOWNLOAD]{tag} {label}", False
|
||
if _entry_local_path(entry) or _cache_has(entry.get("repo_id"), "config.json"):
|
||
return f"[LOCAL] {label}", True
|
||
return f"[DOWNLOAD] {label}", False
|
||
|
||
|
||
def _source_options(kind):
|
||
"""Picker labels for one source, local models first: catalog entries on disk, scanned files, then downloads."""
|
||
labelled = [_entry_label(e) for e in HERETIC_MODEL_ENTRIES if e["kind"] == kind]
|
||
local = [label for label, is_local in labelled if is_local]
|
||
for entry in LOCAL_TEXT_MODELS:
|
||
if entry["kind"] != ("gguf" if kind == "gguf" else "safetensors"):
|
||
continue
|
||
tag = ""
|
||
if kind == "gguf":
|
||
tag = " [VISION]" if any(Path(entry["path"]).parent.glob("*mmproj*.gguf")) else " [TEXT]"
|
||
local.append(f"[LOCAL]{tag} {entry['option']}")
|
||
remote = [label for label, is_local in labelled if not is_local]
|
||
return local, remote
|
||
|
||
|
||
LOCAL_SAFETENSORS_OPTIONS, REMOTE_SAFETENSORS_OPTIONS = _source_options("hf")
|
||
LOCAL_GGUF_OPTIONS, REMOTE_GGUF_OPTIONS = _source_options("gguf")
|
||
SAFETENSORS_MODEL_OPTIONS = (LOCAL_SAFETENSORS_OPTIONS + REMOTE_SAFETENSORS_OPTIONS) or ["none"]
|
||
GGUF_MODEL_OPTIONS = (LOCAL_GGUF_OPTIONS + REMOTE_GGUF_OPTIONS) or ["none"]
|
||
|
||
PRESET_OPTIONS = ["auto"] + sorted(
|
||
k for k, v in PROMPTFORGE_PRESETS.items()
|
||
if isinstance(v, dict) and (v.get("text") or v.get("system_prompt"))
|
||
)
|
||
|
||
MODE_OPTIONS = ["text_only", "inputs"]
|
||
QUANT_OPTIONS = ["auto", "fp16", "q8", "q6", "q5", "q4"]
|
||
GGUF_QUANT_ORDER = ("q8", "q6", "q5", "q4")
|
||
|
||
|
||
def _gguf_quant_order(quantization, quant_files):
|
||
requested = {
|
||
"q8": "q8", "8bit": "q8", "q6": "q6", "6bit": "q6",
|
||
"q5": "q5", "5bit": "q5", "q4": "q4", "4bit": "q4",
|
||
"fp16": "q8", "auto": "q8",
|
||
}.get(quantization or "auto", "q8")
|
||
order = [requested] + [key for key in GGUF_QUANT_ORDER if key != requested]
|
||
return [key for key in order if key in quant_files]
|
||
|
||
|
||
_STATUS_PREFIX_RE = re.compile(r"^\[(?:AVAILABLE|LOCAL|DOWNLOAD)[^\]]*\]\s*", re.I)
|
||
|
||
|
||
def _bare_option(option):
|
||
text = _STATUS_PREFIX_RE.sub("", str(option or ""))
|
||
return re.sub(r"^\[(?:VISION(?: dl)?|TEXT)\]\s*", "", text, flags=re.I)
|
||
|
||
|
||
def _resolve_model_entry(option):
|
||
"""Map a friendly model option to a runtime descriptor; the '[LOCAL ...] ' prefix is stripped first."""
|
||
if not option or option == "auto":
|
||
ent = HERETIC_MODELS[0] if HERETIC_MODELS else {}
|
||
first = HERETIC_MODEL_ENTRIES[0] if HERETIC_MODEL_ENTRIES else {}
|
||
return {"kind": "hf", "key": (ent or {}).get("key") or "heretic_qwen35_9b", "params_b": (ent or {}).get("params_b") or 9, "repo_id": (ent or {}).get("repo_id"), "local_path": _entry_local_path(first) if first else ""}
|
||
label = _bare_option(option)
|
||
for e in HERETIC_MODEL_ENTRIES:
|
||
if e["option"] == label:
|
||
if e["kind"] == "hf":
|
||
ent = next((m for m in HERETIC_MODELS if m.get("key") == e["key"]), {})
|
||
return {"kind": "hf", "key": e["key"], "params_b": e["params_b"], "repo_id": ent.get("repo_id"), "local_path": _entry_local_path(e)}
|
||
ent = next((g for g in GGUF_MODELS if g.get("key") == e["key"]), {})
|
||
return {"kind": "gguf", "key": e["key"], "params_b": e["params_b"], "entry": ent}
|
||
return None
|
||
|
||
|
||
def _local_model_option(option):
|
||
"""Return the absolute path when the picker option is a scanned local model (relative label or old absolute path)."""
|
||
opt = _bare_option(option)
|
||
if not opt or opt == "auto":
|
||
return ""
|
||
for e in LOCAL_TEXT_MODELS:
|
||
if opt in (e["option"], e["path"]):
|
||
return e["path"]
|
||
return ""
|
||
|
||
|
||
def _hf_dtype_kwargs(dtype):
|
||
"""transformers 4.56+ renamed `torch_dtype` to `dtype` on from_pretrained."""
|
||
try:
|
||
major, minor = (int(part) for part in transformers.__version__.split(".")[:2])
|
||
if (major, minor) >= (4, 56):
|
||
return {"dtype": dtype}
|
||
except Exception:
|
||
pass
|
||
return {"torch_dtype": dtype}
|
||
|
||
|
||
def _hf_compute_dtype():
|
||
"""bf16 where the GPU supports it; Gemma/Qwen logits overflow in fp16."""
|
||
if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
|
||
return torch.bfloat16
|
||
return torch.float16
|
||
|
||
|
||
def _hf_max_memory(reserve=3 * 1024 ** 3):
|
||
"""GPU budget for device_map='auto', leaving headroom for activations and KV cache."""
|
||
try:
|
||
free_bytes = torch.cuda.mem_get_info()[0]
|
||
except Exception:
|
||
return None
|
||
budget = max(int(free_bytes) - reserve, 768 * 1024 ** 2)
|
||
return {torch.cuda.current_device(): budget, "cpu": psutil.virtual_memory().available}
|
||
|
||
|
||
def _vision_size(width, height, max_pixels=1024 * 1024):
|
||
"""Target (w, h) for references above ~1 megapixel, else None. Vision encoders cost about 1k tokens per megapixel, and every token also slows each generated token."""
|
||
if width * height <= max_pixels:
|
||
return None
|
||
scale = (max_pixels / (width * height)) ** 0.5
|
||
return max(32, round(width * scale)), max(32, round(height * scale))
|
||
|
||
|
||
def _free_vram(keep=None, needed=10 ** 30):
|
||
"""Unload ComfyUI's models (least recently used first) until `needed` bytes are free."""
|
||
try:
|
||
import comfy.model_management as model_management
|
||
keep_loaded = []
|
||
if keep is not None:
|
||
keep_loaded = [m for m in model_management.current_loaded_models if m.model is keep]
|
||
model_management.free_memory(needed, model_management.get_torch_device(), keep_loaded=keep_loaded)
|
||
if needed >= 10 ** 30:
|
||
model_management.free_memory(needed, torch.device("cpu"))
|
||
model_management.soft_empty_cache()
|
||
except Exception:
|
||
pass
|
||
gc.collect()
|
||
try:
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
_GGUF_VALUE_SIZE = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8}
|
||
|
||
|
||
def _gguf_skip(f, vtype):
|
||
if vtype == 8:
|
||
f.seek(struct.unpack("<Q", f.read(8))[0], 1)
|
||
elif vtype == 9:
|
||
elem = struct.unpack("<I", f.read(4))[0]
|
||
count = struct.unpack("<Q", f.read(8))[0]
|
||
if elem == 8:
|
||
for _ in range(count):
|
||
f.seek(struct.unpack("<Q", f.read(8))[0], 1)
|
||
else:
|
||
f.seek(_GGUF_VALUE_SIZE[elem] * count, 1)
|
||
else:
|
||
f.seek(_GGUF_VALUE_SIZE[vtype], 1)
|
||
|
||
|
||
def _gguf_read(f, vtype):
|
||
if vtype == 8:
|
||
return f.read(struct.unpack("<Q", f.read(8))[0]).decode("utf-8", "replace")
|
||
if vtype == 9:
|
||
elem = struct.unpack("<I", f.read(4))[0]
|
||
return [_gguf_read(f, elem) for _ in range(struct.unpack("<Q", f.read(8))[0])]
|
||
fmt = {0: "<B", 1: "<b", 2: "<H", 3: "<h", 4: "<I", 5: "<i", 6: "<f", 7: "<?", 10: "<Q", 11: "<q", 12: "<d"}[vtype]
|
||
return struct.unpack(fmt, f.read(struct.calcsize(fmt)))[0]
|
||
|
||
|
||
def _gguf_metadata(path):
|
||
"""Read GGUF metadata without mapping the tensor data."""
|
||
meta = {}
|
||
needed = None
|
||
with open(path, "rb") as f:
|
||
if f.read(4) != b"GGUF":
|
||
return meta
|
||
f.read(4)
|
||
f.read(8)
|
||
count = struct.unpack("<Q", f.read(8))[0]
|
||
for _ in range(count):
|
||
key = f.read(struct.unpack("<Q", f.read(8))[0]).decode("utf-8", "replace")
|
||
vtype = struct.unpack("<I", f.read(4))[0]
|
||
if needed is not None and key not in needed:
|
||
_gguf_skip(f, vtype)
|
||
continue
|
||
value = _gguf_read(f, vtype)
|
||
meta[key] = value
|
||
if key == "general.architecture":
|
||
needed = {f"{value}.{k}" for k in (
|
||
"block_count", "attention.head_count", "attention.head_count_kv",
|
||
"embedding_length", "attention.key_length", "attention.value_length",
|
||
"attention.sliding_window",
|
||
)}
|
||
needed.add("tokenizer.chat_template")
|
||
if needed is not None and needed <= meta.keys():
|
||
break
|
||
return meta
|
||
|
||
|
||
def _gguf_kv_profile(path):
|
||
"""Return (layers, kv bytes per token, swa window or None) from GGUF metadata."""
|
||
try:
|
||
meta = _gguf_metadata(path)
|
||
prefix = f"{meta['general.architecture']}."
|
||
swa = meta.get(f"{prefix}attention.sliding_window")
|
||
n_layers = int(meta[f"{prefix}block_count"])
|
||
n_heads = int(meta[f"{prefix}attention.head_count"])
|
||
kv_heads = meta[f"{prefix}attention.head_count_kv"]
|
||
embd = int(meta[f"{prefix}embedding_length"])
|
||
key_len = meta.get(f"{prefix}attention.key_length", embd // n_heads)
|
||
value_len = meta.get(f"{prefix}attention.value_length", embd // n_heads)
|
||
def per_layer(value):
|
||
if isinstance(value, list):
|
||
return [int(v) for v in value] if len(value) == n_layers else None
|
||
return [int(value)] * n_layers
|
||
kv_heads_l = per_layer(kv_heads)
|
||
key_l = per_layer(key_len)
|
||
value_l = per_layer(value_len)
|
||
if kv_heads_l is None or key_l is None or value_l is None:
|
||
return None
|
||
kv_bytes = sum(kv_heads_l[i] * (key_l[i] + value_l[i]) * 2 for i in range(n_layers))
|
||
return n_layers, kv_bytes, int(swa) if swa is not None else None
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _gguf_memory_estimate(path, n_ctx, layers, mmproj_path=None):
|
||
"""Approximate VRAM needed for the GPU-resident part of a GGUF model."""
|
||
weights = path.stat().st_size
|
||
extra = mmproj_path.stat().st_size + 1024 ** 3 if mmproj_path is not None else 0 # + image encoder compute buffers
|
||
profile = _gguf_kv_profile(path)
|
||
if profile is None:
|
||
return weights + extra + 768 * 1024 ** 2
|
||
n_layers, kv_per_token, swa = profile
|
||
if swa is not None:
|
||
kv_bytes = kv_per_token * (min(n_ctx, swa) * 0.95 + n_ctx * 0.05)
|
||
else:
|
||
kv_bytes = kv_per_token * n_ctx
|
||
if layers == 0:
|
||
return extra + 768 * 1024 ** 2
|
||
fraction = 1.0 if layers < 0 else min(1.0, max(0.0, layers / n_layers))
|
||
return (weights + kv_bytes) * fraction + extra + 768 * 1024 ** 2
|
||
|
||
|
||
def _gguf_layers_that_fit(path, n_ctx, mmproj_path=None):
|
||
"""GPU layer count for the VRAM that is free right now (other processes included)."""
|
||
try:
|
||
free_bytes = torch.cuda.mem_get_info()[0]
|
||
except Exception:
|
||
return 0
|
||
profile = _gguf_kv_profile(path)
|
||
if profile is None:
|
||
return 0
|
||
n_layers, kv_per_token, swa = profile
|
||
kv = kv_per_token * (min(n_ctx, swa) * 0.95 + n_ctx * 0.05) if swa is not None else kv_per_token * n_ctx
|
||
extra = mmproj_path.stat().st_size if mmproj_path is not None else 0
|
||
cost = (path.stat().st_size + kv) / n_layers
|
||
return max(0, min(n_layers, int((free_bytes - extra - 768 * 1024 ** 2) / cost)))
|
||
|
||
|
||
def _gpu_budget_bytes():
|
||
try:
|
||
total = torch.cuda.get_device_properties(0).total_memory
|
||
except Exception:
|
||
return 0
|
||
return max(int(total) - 2 * 1024 ** 3, 1024 ** 3)
|
||
|
||
|
||
def _auto_batch(budget):
|
||
if budget >= 12 * 1024 ** 3:
|
||
return 2048, 512
|
||
if budget >= 8 * 1024 ** 3:
|
||
return 1024, 256
|
||
return 512, 128
|
||
|
||
|
||
def _auto_gguf_fit(path, mmproj_path, budget):
|
||
profile = _gguf_kv_profile(path)
|
||
if profile is None:
|
||
return None
|
||
n_layers, kv_per_token, swa = profile
|
||
weights = path.stat().st_size
|
||
extra = mmproj_path.stat().st_size if mmproj_path is not None else 0
|
||
scales = {"f16": 1.0, "q8_0": 0.5, "q4_0": 0.28}
|
||
for ctx in (16384, 12288, 8192, 6144, 4096):
|
||
for kind in ("f16", "q8_0", "q4_0"):
|
||
kv = kv_per_token * (min(ctx, swa) * 0.95 + ctx * 0.05) if swa is not None else kv_per_token * ctx
|
||
if weights + kv * scales[kind] + extra <= budget:
|
||
return {"context_size": ctx, "gpu_layers": -1, "kv_cache_type": kind, "offload_kqv": True}
|
||
ctx, kind = 4096, "q8_0"
|
||
kv = kv_per_token * (min(ctx, swa) * 0.95 + ctx * 0.05) if swa is not None else kv_per_token * ctx
|
||
cost = (weights + kv * scales[kind] + extra) / n_layers
|
||
layers = max(0, int((budget - extra - 512 * 1024 ** 2) / cost)) if cost > 0 else 0
|
||
return {"context_size": ctx, "gpu_layers": layers, "kv_cache_type": kind, "offload_kqv": True}
|
||
|
||
|
||
def _auto_gguf_settings(entry, quantization, mmproj_path=None):
|
||
budget = _gpu_budget_bytes()
|
||
if budget <= 0:
|
||
return {}
|
||
candidates = []
|
||
if quantization in ("auto", "fp16", ""):
|
||
for q in ("q8", "q6", "q5", "q4"):
|
||
path = _find_gguf(entry, q, download=False, exact=True)
|
||
if path is not None:
|
||
candidates.append((q, path))
|
||
else:
|
||
path = _find_gguf(entry, quantization, download=False, exact=True)
|
||
if path is not None:
|
||
candidates.append((quantization, path))
|
||
else:
|
||
for q in ("q8", "q6", "q5", "q4"):
|
||
local = _find_gguf(entry, q, download=False, exact=True)
|
||
if local is not None:
|
||
candidates.append((q, local))
|
||
if candidates:
|
||
print(f"[PromptForge] {quantization} is not local; using local {candidates[0][0]}.")
|
||
chosen = None
|
||
for q, path in candidates:
|
||
fit = _auto_gguf_fit(path, mmproj_path, budget)
|
||
if fit is not None:
|
||
chosen = (q, fit)
|
||
break
|
||
if chosen is None:
|
||
return {"quantization": "q4" if quantization in ("auto", "fp16", "") else quantization}
|
||
q, fit = chosen
|
||
batch, ubatch = _auto_batch(budget)
|
||
settings = {
|
||
"quantization": q, "batch_size": batch, "ubatch_size": ubatch,
|
||
"swa_full": False, "flash_attn": "auto",
|
||
}
|
||
settings.update(fit)
|
||
return settings
|
||
|
||
|
||
def _auto_gguf_path_settings(path, quantization, mmproj=None):
|
||
budget = _gpu_budget_bytes()
|
||
batch, ubatch = _auto_batch(budget)
|
||
file_quant = re.search(r"[-_.](q[2-8])_", path.name.lower())
|
||
settings = {
|
||
"quantization": file_quant.group(1) if file_quant else quantization,
|
||
"batch_size": batch, "ubatch_size": ubatch, "swa_full": False, "flash_attn": "auto",
|
||
}
|
||
settings.update(_auto_gguf_fit(path, mmproj, budget) or {})
|
||
return settings
|
||
|
||
|
||
def _dir_params_b(path):
|
||
"""Parameter count in billions from the safetensors bytes of an HF model directory (bf16 weights)."""
|
||
path = Path(str(path or ""))
|
||
if not path.is_dir():
|
||
return None
|
||
total = sum(f.stat().st_size for f in path.glob("*.safetensors") if "mtp" not in f.name.lower())
|
||
return round(total / 2 / 1e9, 1) if total else None
|
||
|
||
|
||
def _safetensors_params_b(resolved, src=""):
|
||
return (resolved or {}).get("params_b") or _dir_params_b((resolved or {}).get("local_path") or src)
|
||
|
||
|
||
def _auto_safetensors_settings(resolved, src=""):
|
||
budget = _gpu_budget_bytes()
|
||
params_b = _safetensors_params_b(resolved, src)
|
||
if not params_b or budget <= 0:
|
||
return {}
|
||
if params_b * 2.0 * 1024 ** 3 <= budget * 0.85:
|
||
return {"quantization": "auto"}
|
||
if params_b * 1.35 * 1024 ** 3 + 2 * 1024 ** 3 <= budget * 0.95:
|
||
return {"quantization": "q8"}
|
||
return {"quantization": "q4"}
|
||
|
||
|
||
def _gguf_vision_handler(name):
|
||
low = (name or "").lower()
|
||
if "gemma-4" in low or "gemma4" in low:
|
||
return "Gemma4ChatHandler"
|
||
if "gemma-3" in low or "gemma3" in low:
|
||
return "Gemma3ChatHandler"
|
||
if "qwen2.5" in low or "qwen25" in low:
|
||
return "Qwen25VLChatHandler"
|
||
if "qwen3.5" in low or "qwen3.6" in low or "qwen35" in low:
|
||
return "Qwen35ChatHandler"
|
||
return "Qwen3VLChatHandler"
|
||
|
||
|
||
def _gguf_chat_format(name):
|
||
low = (name or "").lower()
|
||
if "gemma" in low:
|
||
return "gemma"
|
||
if "llama" in low:
|
||
return "llama-3"
|
||
if "ministral" in low or "mistral" in low:
|
||
return "mistral-instruct"
|
||
if "qwen" in low or "heretic" in low or "phi" in low:
|
||
return "chatml"
|
||
return None
|
||
|
||
|
||
def _gguf_avail_tokens(entry):
|
||
key = (entry or {}).get("key")
|
||
for item in HERETIC_MODEL_ENTRIES:
|
||
if item.get("key") == key:
|
||
return [t.lower() for t in (item.get("avail_match") or [])]
|
||
return []
|
||
|
||
|
||
def _scanned_ggufs():
|
||
return [m["path"] for m in LOCAL_TEXT_MODELS if (m.get("path") or "").lower().endswith(".gguf")]
|
||
|
||
|
||
def _find_gguf(entry, quantization, download=True, exact=False):
|
||
"""Find or download a GGUF for a catalog entry, honouring the requested quant."""
|
||
qf = (entry or {}).get("quant_files") or {}
|
||
if not qf:
|
||
return None
|
||
order = _gguf_quant_order(quantization, qf) or list(qf)
|
||
if exact:
|
||
order = order[:1]
|
||
|
||
base_dir = Path(entry.get("base_dir") or "LLM/GGUF")
|
||
if folder_paths is not None:
|
||
models_dir = getattr(folder_paths, "models_dir", "") or ""
|
||
if models_dir:
|
||
base_dir = Path(models_dir) / base_dir
|
||
scanned = _scanned_ggufs()
|
||
repo_id = (entry or {}).get("repo_id")
|
||
|
||
def local(name):
|
||
candidate = base_dir / name
|
||
if candidate.is_file():
|
||
return candidate
|
||
for path in scanned:
|
||
if Path(path).name == name:
|
||
return Path(path)
|
||
return None
|
||
|
||
def download_q(name):
|
||
if not repo_id:
|
||
return None
|
||
from huggingface_hub import hf_hub_download
|
||
try:
|
||
return Path(hf_hub_download(repo_id=repo_id, filename=name, local_files_only=True))
|
||
except Exception:
|
||
pass
|
||
if not download:
|
||
return None
|
||
try:
|
||
print("[PromptForge] downloading GGUF model (first use).")
|
||
return Path(hf_hub_download(repo_id=repo_id, filename=name))
|
||
except Exception:
|
||
return None
|
||
|
||
found = local(qf[order[0]])
|
||
if found is not None:
|
||
return found
|
||
found = download_q(qf[order[0]])
|
||
if found is not None:
|
||
return found
|
||
|
||
for key in order[1:]:
|
||
found = local(qf[key])
|
||
if found is not None:
|
||
return found
|
||
found = download_q(qf[key])
|
||
if found is not None:
|
||
return found
|
||
|
||
if exact:
|
||
return None
|
||
|
||
tokens = _gguf_avail_tokens(entry)
|
||
if tokens:
|
||
matches = [p for p in scanned if all(t in Path(p).name.lower() for t in tokens)]
|
||
if matches:
|
||
return Path(matches[0])
|
||
return None
|
||
|
||
|
||
def _find_mmproj(gguf_path, entry):
|
||
names = []
|
||
if (entry or {}).get("mmproj_file"):
|
||
names.append(entry["mmproj_file"])
|
||
dirs = [gguf_path.parent]
|
||
for path in _scanned_ggufs():
|
||
parent = Path(path).parent
|
||
if parent not in dirs:
|
||
dirs.append(parent)
|
||
for directory in dirs:
|
||
for name in names:
|
||
candidate = directory / name
|
||
if candidate.is_file():
|
||
return candidate
|
||
for candidate in sorted(directory.glob("*mmproj*.gguf")):
|
||
return candidate
|
||
return None
|
||
|
||
|
||
def _local_mmproj(gguf_path, entry):
|
||
name = (entry or {}).get("mmproj_file")
|
||
if name:
|
||
candidate = gguf_path.parent / name
|
||
if candidate.is_file():
|
||
return candidate
|
||
matches = sorted(gguf_path.parent.glob("*mmproj*.gguf"))
|
||
return matches[0] if matches else None
|
||
|
||
|
||
def _resolve_mmproj(gguf_path, entry, override=""):
|
||
"""Pick the mmproj for a vision GGUF: override, matching catalog file, HF download, then any local."""
|
||
if override:
|
||
path = Path(override)
|
||
if not path.is_file():
|
||
raise ValueError(f"[PromptForge] mmproj_path does not exist: {path}")
|
||
return path
|
||
name = (entry or {}).get("mmproj_file")
|
||
if name:
|
||
candidate = gguf_path.parent / name
|
||
if candidate.is_file():
|
||
return candidate
|
||
repo_id = (entry or {}).get("repo_id")
|
||
if repo_id and name:
|
||
try:
|
||
from huggingface_hub import hf_hub_download
|
||
return Path(hf_hub_download(repo_id=repo_id, filename=name))
|
||
except Exception as exc:
|
||
raise ValueError(f"[PromptForge] Could not download mmproj '{name}' from {repo_id}: {exc}") from exc
|
||
return _find_mmproj(gguf_path, entry)
|
||
|
||
|
||
def _gguf_target(option, quantization):
|
||
"""(catalog entry, gguf path) for a picker option: scanned local file (relative label or absolute path) or catalog model."""
|
||
text = str(option or "")
|
||
local = _local_model_option(text) or (text if text.lower().endswith(".gguf") and Path(text).is_file() else "")
|
||
if local:
|
||
return None, Path(local)
|
||
resolved = _resolve_model_entry(text)
|
||
entry = resolved.get("entry") if resolved else None
|
||
path = _find_gguf(entry, quantization, download=False) if entry else None
|
||
return entry, path
|
||
|
||
|
||
def _option_for_key(key):
|
||
for entry in HERETIC_MODEL_ENTRIES:
|
||
if entry.get("key") == key:
|
||
return entry.get("option")
|
||
return None
|
||
|
||
|
||
LOCAL_HEADER = "────── local models (on disk, ready) ──────"
|
||
REMOTE_HEADER = "────── not downloaded (fetched on first use) ──────"
|
||
MODEL_OPTIONS = (
|
||
["connected clip", LOCAL_HEADER]
|
||
+ [f"[gguf] {option}" for option in LOCAL_GGUF_OPTIONS]
|
||
+ [f"[safetensors] {option}" for option in LOCAL_SAFETENSORS_OPTIONS]
|
||
+ [REMOTE_HEADER]
|
||
+ [f"[gguf] {option}" for option in REMOTE_GGUF_OPTIONS]
|
||
+ [f"[safetensors] {option}" for option in REMOTE_SAFETENSORS_OPTIONS]
|
||
)
|
||
print(
|
||
f"[PromptForge] {len(LOCAL_GGUF_OPTIONS)} local GGUF and {len(LOCAL_SAFETENSORS_OPTIONS)} local safetensors models ready; "
|
||
f"{len(REMOTE_GGUF_OPTIONS) + len(REMOTE_SAFETENSORS_OPTIONS)} catalog models would download on first use."
|
||
)
|
||
|
||
|
||
def _check_not_header(option):
|
||
if option in (LOCAL_HEADER, REMOTE_HEADER):
|
||
raise ValueError("[PromptForge] that picker entry is a section header; choose a model below it.")
|
||
|
||
|
||
def _default_model_option():
|
||
"""A local GGUF is the fastest path, so prefer one when llama.cpp is installed; else the first safetensors entry."""
|
||
if importlib.util.find_spec("llama_cpp") is not None:
|
||
for option in MODEL_OPTIONS:
|
||
if option.startswith("[gguf] [LOCAL"):
|
||
return option
|
||
return next((option for option in MODEL_OPTIONS if option.startswith("[safetensors] ")), MODEL_OPTIONS[0])
|
||
|
||
|
||
_DEFAULT_MODEL = _default_model_option()
|
||
|
||
|
||
def _parse_model_option(option):
|
||
text = str(option or "")
|
||
if text == "connected clip":
|
||
return "connected clip", "auto", "auto"
|
||
if text.startswith("[safetensors] "):
|
||
return "safetensors", text[len("[safetensors] "):], "auto"
|
||
if text.startswith("[gguf] "):
|
||
return "GGUF", "auto", text[len("[gguf] "):]
|
||
return "safetensors", text, "auto"
|
||
|
||
|
||
_KV_TYPE_VALUES = {"f16": 1, "q8_0": 8, "q5_1": 7, "q5_0": 6, "q4_1": 3, "q4_0": 2}
|
||
|
||
|
||
def _kv_type_value(name):
|
||
return _KV_TYPE_VALUES.get(str(name).lower(), 1)
|
||
|
||
|
||
def _display_model(name, source):
|
||
"""Re-check local availability at run time so the report never shows a stale widget label."""
|
||
bare = _bare_option(name)
|
||
if _local_model_option(bare) or Path(bare).exists():
|
||
return f"[LOCAL] {bare}"
|
||
resolved = _resolve_model_entry(bare)
|
||
if source == "GGUF":
|
||
entry = resolved.get("entry") if resolved else None
|
||
quants = _gguf_local_quants(entry) if entry else []
|
||
if quants:
|
||
ordered = [q for q in ("q4", "q5", "q6", "q8") if q in quants]
|
||
return f"[LOCAL {'/'.join(q.upper() for q in ordered)}] {bare}"
|
||
return f"[DOWNLOAD] {bare}"
|
||
local = (resolved or {}).get("local_path") or ""
|
||
repo = (resolved or {}).get("repo_id") or ""
|
||
if local or (repo and _cache_has(repo, "config.json")):
|
||
return f"[LOCAL] {bare}"
|
||
return f"[DOWNLOAD] {bare}"
|
||
|
||
|
||
def _settings_report(settings):
|
||
source = settings.get("source", "safetensors")
|
||
if source == "GGUF":
|
||
return (
|
||
f"GGUF {_display_model(settings.get('gguf_model'), 'GGUF')} | quant {settings.get('quantization')} | "
|
||
f"ctx {settings.get('context_size')} | gpu_layers {settings.get('gpu_layers')} | "
|
||
f"kv {settings.get('kv_cache_type')} | batch {settings.get('batch_size')}/{settings.get('ubatch_size')}"
|
||
)
|
||
if source == "safetensors":
|
||
return f"safetensors {_display_model(settings.get('safetensors_model'), 'safetensors')} | quant {settings.get('quantization')}"
|
||
return "connected clip | no model loaded"
|
||
|
||
|
||
def _recommend_settings(data):
|
||
_check_not_header(data.get("model", ""))
|
||
source, safetensors_model, gguf_model = _parse_model_option(data.get("model", ""))
|
||
quantization = data.get("quantization", "auto")
|
||
settings = {
|
||
"source": source, "safetensors_model": safetensors_model, "gguf_model": gguf_model,
|
||
"quantization": quantization,
|
||
}
|
||
if source == "GGUF":
|
||
entry, path = _gguf_target(gguf_model, quantization)
|
||
mmproj = _local_mmproj(path, entry) if path is not None else None
|
||
if entry is None and path is not None:
|
||
settings.update(_auto_gguf_path_settings(path, quantization, mmproj))
|
||
else:
|
||
settings.update(_auto_gguf_settings(entry, quantization, mmproj))
|
||
elif source == "safetensors":
|
||
settings.update(_auto_safetensors_settings(_resolve_model_entry(safetensors_model), _local_model_option(safetensors_model)))
|
||
for key, value in (
|
||
("context_size", 16384), ("gpu_layers", -1), ("batch_size", 2048), ("ubatch_size", 512),
|
||
("kv_cache_type", "f16"), ("offload_kqv", True), ("swa_full", False), ("flash_attn", "auto"),
|
||
):
|
||
settings.setdefault(key, value)
|
||
return settings
|
||
|
||
|
||
def _test_settings(data):
|
||
_check_not_header(data.get("model", ""))
|
||
source, safetensors_model, gguf_model = _parse_model_option(data.get("model", ""))
|
||
quantization = data.get("quantization", "auto")
|
||
if source == "connected clip":
|
||
return {"ok": True, "message": "connected clip is verified when the Model node runs.", "settings": {}}
|
||
if source == "GGUF":
|
||
return _test_gguf(gguf_model, quantization)
|
||
return _test_safetensors(safetensors_model, quantization)
|
||
|
||
|
||
def _test_gguf(option, quantization):
|
||
entry, path = _gguf_target(option, quantization)
|
||
if path is None:
|
||
return {"ok": False, "message": "model is not downloaded yet; run it once, then test."}
|
||
mmproj = _local_mmproj(path, entry)
|
||
plan = _auto_gguf_path_settings(path, quantization, mmproj) if entry is None else _auto_gguf_settings(entry, quantization, mmproj)
|
||
ctx = int(plan.get("context_size", 16384))
|
||
kv = plan.get("kv_cache_type", "f16")
|
||
layers = int(plan.get("gpu_layers", -1))
|
||
batch = int(plan.get("batch_size", 2048))
|
||
ubatch = int(plan.get("ubatch_size", 512))
|
||
candidates = [(ctx, kv)]
|
||
for fallback_ctx in (12288, 8192, 6144, 4096):
|
||
if fallback_ctx < ctx:
|
||
candidates.append((fallback_ctx, kv))
|
||
for fallback_kv in ("q8_0", "q4_0"):
|
||
if fallback_kv != kv:
|
||
candidates.append((ctx, fallback_kv))
|
||
handler = None
|
||
if mmproj is not None:
|
||
try:
|
||
from llama_cpp import llama_chat_format
|
||
handler_name = _gguf_vision_handler(f"{path.name} {(entry or {}).get('name', '')}")
|
||
handler_cls = getattr(llama_chat_format, handler_name, None)
|
||
if handler_cls is not None:
|
||
handler = handler_cls(clip_model_path=str(mmproj), verbose=False)
|
||
except Exception:
|
||
handler = None
|
||
from llama_cpp import Llama
|
||
_free_vram()
|
||
last_error = ""
|
||
for test_ctx, test_kv in candidates[:4]:
|
||
kwargs = {
|
||
"model_path": str(path),
|
||
"n_ctx": test_ctx,
|
||
"n_batch": min(batch, test_ctx),
|
||
"n_ubatch": min(ubatch, min(batch, test_ctx)),
|
||
"n_gpu_layers": layers,
|
||
"offload_kqv": True,
|
||
"swa_full": False,
|
||
"type_k": _kv_type_value(test_kv),
|
||
"type_v": _kv_type_value(test_kv),
|
||
"use_mmap": True,
|
||
"verbose": False,
|
||
}
|
||
if handler is not None:
|
||
kwargs["chat_handler"] = handler
|
||
try:
|
||
llm = Llama(**kwargs)
|
||
llm.close()
|
||
return {
|
||
"ok": True,
|
||
"message": f"loaded {path.name} at {test_ctx} tokens, kv {test_kv}, gpu_layers {layers}.",
|
||
"settings": {
|
||
"quantization": plan.get("quantization", quantization),
|
||
"context_size": test_ctx, "kv_cache_type": test_kv, "gpu_layers": layers,
|
||
"batch_size": batch, "ubatch_size": ubatch, "offload_kqv": True,
|
||
"swa_full": False, "flash_attn": "auto",
|
||
},
|
||
}
|
||
except Exception as exc:
|
||
last_error = str(exc)
|
||
gc.collect()
|
||
return {"ok": False, "message": last_error or "context creation failed"}
|
||
|
||
|
||
def _load_hf_model(src, load_kwargs):
|
||
"""Load with the vision class when the config has one (text-only prompts work there too), else as a causal LM."""
|
||
try:
|
||
return AutoVision.from_pretrained(src, **load_kwargs)
|
||
except ValueError as exc:
|
||
if "Unrecognized configuration class" not in str(exc):
|
||
raise
|
||
return AutoModelForCausalLM.from_pretrained(src, **load_kwargs)
|
||
|
||
|
||
def _test_safetensors(option, quantization):
|
||
resolved = _resolve_model_entry(option)
|
||
src = (resolved or {}).get("local_path") or (resolved or {}).get("repo_id") or _local_model_option(option)
|
||
path = Path(src) if src else None
|
||
if path is None or not path.exists():
|
||
return {"ok": False, "message": "model is not downloaded yet; run it once, then test."}
|
||
if not path.is_dir():
|
||
return {"ok": False, "message": "single-file encoders are verified on first generation."}
|
||
plan = _auto_safetensors_settings(resolved, src)
|
||
quant = plan.get("quantization") or ("q4" if quantization in ("auto", "fp16", "") else quantization)
|
||
dtype = _hf_compute_dtype()
|
||
quant_cfg = None
|
||
if quant in ("q8", "q4"):
|
||
quant_cfg = BitsAndBytesConfig(
|
||
load_in_8bit=(quant == "q8"),
|
||
load_in_4bit=(quant == "q4"),
|
||
bnb_4bit_quant_type="nf4",
|
||
bnb_4bit_use_double_quant=True,
|
||
bnb_4bit_compute_dtype=dtype,
|
||
llm_int8_enable_fp32_cpu_offload=True,
|
||
)
|
||
kwargs = {"quantization_config": quant_cfg, "device_map": "auto", **_hf_dtype_kwargs(dtype)}
|
||
max_memory = _hf_max_memory()
|
||
if max_memory is not None:
|
||
kwargs["max_memory"] = max_memory
|
||
_free_vram()
|
||
try:
|
||
model = _load_hf_model(str(path), kwargs)
|
||
del model
|
||
gc.collect()
|
||
torch.cuda.empty_cache()
|
||
return {"ok": True, "message": f"loaded {path.name} in {quant}.", "settings": {"quantization": quant}}
|
||
except Exception as exc:
|
||
return {"ok": False, "message": str(exc)}
|
||
|
||
|
||
def _scan_configs(data):
|
||
scope = data.get("scope", "all local")
|
||
configs = {}
|
||
lines = []
|
||
if scope == "all local":
|
||
for entry in GGUF_MODELS:
|
||
path = _find_gguf(entry, "auto", download=False)
|
||
if path is None:
|
||
continue
|
||
key = entry.get("key") or path.stem
|
||
mmproj = _local_mmproj(path, entry)
|
||
cfg = _auto_gguf_settings(entry, "auto", mmproj)
|
||
cfg.update({"source": "GGUF", "gguf_model": _option_for_key(key) or key, "safetensors_model": "auto"})
|
||
configs[key] = cfg
|
||
lines.append(f"{key}: {_settings_report(cfg)}")
|
||
for option in LOCAL_SAFETENSORS_OPTIONS:
|
||
local = _local_model_option(option)
|
||
if local and not Path(local).is_dir():
|
||
continue
|
||
cfg = _auto_safetensors_settings(_resolve_model_entry(option), local)
|
||
cfg.update({"source": "safetensors", "safetensors_model": option, "gguf_model": "auto"})
|
||
key = _bare_option(option)
|
||
configs[key] = cfg
|
||
lines.append(f"{key}: {_settings_report(cfg)}")
|
||
lines.append(f"not downloaded: {len(REMOTE_GGUF_OPTIONS)} GGUF, {len(REMOTE_SAFETENSORS_OPTIONS)} safetensors catalog models (they fetch on first use).")
|
||
else:
|
||
cfg = _recommend_settings(data)
|
||
configs[cfg["source"]] = cfg
|
||
lines.append(_settings_report(cfg))
|
||
if not lines:
|
||
lines.append("No local models found. Download a model first or use 'selected'.")
|
||
return {"report": "\n".join(lines), "configs": configs}
|
||
|
||
|
||
class PROMPTFORGE_Model:
|
||
CATEGORY = "🔮 PromptForge/Prompts"
|
||
RETURN_TYPES = ("PROMPTFORGE_MODEL", "STRING")
|
||
RETURN_NAMES = ("MODEL", "REPORT")
|
||
FUNCTION = "build"
|
||
DESCRIPTION = "PromptForge model and runtime settings. One picker for clip / safetensors / GGUF; everything else is configured for the current GPU automatically."
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"model": (MODEL_OPTIONS, {"default": _DEFAULT_MODEL, "tooltip": "What to run. The list is split into local models (on disk or in the HF cache, ready to use; [LOCAL Q4/Q8] names the GGUF quants present) and [DOWNLOAD] catalog models that fetch on first use."}),
|
||
"quantization": (QUANT_OPTIONS, {"default": "auto", "tooltip": "Quantization. auto picks the best local option for this GPU."}),
|
||
"context_size": ("INT", {"default": 0, "min": 0, "max": 262144, "step": 512, "tooltip": "GGUF context length. 0 = auto (largest that fits, 16k first)."}),
|
||
"local_model_path": ("STRING", {"default": "", "multiline": True, "tooltip": "Optional local safetensors dir/file or GGUF file; overrides the picker."}),
|
||
"mmproj_path": ("STRING", {"default": "", "tooltip": "Optional mmproj (.gguf) for GGUF image/video mode."}),
|
||
},
|
||
"optional": {
|
||
"clip": ("CLIP", {"tooltip": "Connected ComfyUI text encoder. Required when model is 'connected clip'."}),
|
||
},
|
||
}
|
||
|
||
def build(self, model, quantization, context_size, local_model_path, mmproj_path, clip=None, **kwargs):
|
||
_check_not_header(model)
|
||
source, safetensors_model, gguf_model = _parse_model_option(model)
|
||
local = (local_model_path or "").strip()
|
||
if local:
|
||
if not Path(local).exists():
|
||
raise ValueError(f"[PromptForge] local_model_path does not exist on this machine: {local}")
|
||
source = "GGUF" if local.lower().endswith(".gguf") else "safetensors"
|
||
if source == "GGUF":
|
||
gguf_model = local
|
||
else:
|
||
safetensors_model = local
|
||
settings = {
|
||
"source": source,
|
||
"safetensors_model": safetensors_model,
|
||
"gguf_model": gguf_model,
|
||
"quantization": quantization,
|
||
"local_model_path": local_model_path,
|
||
"mmproj_path": mmproj_path,
|
||
"clip": clip,
|
||
}
|
||
if source == "GGUF":
|
||
entry, path = _gguf_target(gguf_model, quantization)
|
||
mmproj = _local_mmproj(path, entry) if path is not None else None
|
||
if entry is None and path is not None:
|
||
settings.update(_auto_gguf_path_settings(path, quantization, mmproj))
|
||
else:
|
||
settings.update(_auto_gguf_settings(entry, quantization, mmproj))
|
||
elif source == "safetensors":
|
||
settings.update(_auto_safetensors_settings(_resolve_model_entry(safetensors_model), _local_model_option(safetensors_model) or local))
|
||
for key, value in (
|
||
("context_size", 16384), ("gpu_layers", -1), ("batch_size", 2048), ("ubatch_size", 512),
|
||
("kv_cache_type", "f16"), ("offload_kqv", True), ("swa_full", False), ("flash_attn", "auto"),
|
||
):
|
||
settings.setdefault(key, value)
|
||
if int(context_size) > 0:
|
||
settings["context_size"] = int(context_size)
|
||
if source == "GGUF":
|
||
entry, _ = _gguf_target(gguf_model, settings.get("quantization", quantization))
|
||
chosen = settings.get("quantization", quantization)
|
||
if entry is not None and _find_gguf(entry, chosen, download=False, exact=True) is None:
|
||
print(f"[PromptForge] {gguf_model} ({chosen}) is not local; it will download on first use.")
|
||
elif source == "safetensors":
|
||
resolved = _resolve_model_entry(safetensors_model)
|
||
repo = (resolved or {}).get("repo_id") or ""
|
||
local = (resolved or {}).get("local_path") or ""
|
||
if repo and not local and not _cache_has(repo, "config.json"):
|
||
print(f"[PromptForge] {safetensors_model} is not local; it will download on first use.")
|
||
if source == "connected clip":
|
||
if clip is None:
|
||
raise ValueError("[PromptForge] model is 'connected clip' but no CLIP input is connected.")
|
||
try:
|
||
clip.tokenize("test")
|
||
except Exception as exc:
|
||
raise ValueError(f"[PromptForge] connected CLIP failed to tokenize: {exc}") from exc
|
||
report = _settings_report(settings)
|
||
print(f"[PromptForge] model: {report}")
|
||
return (settings, report)
|
||
|
||
|
||
class PROMPTFORGE_ConfigFinder:
|
||
CATEGORY = "🔮 PromptForge/Prompts"
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("REPORT",)
|
||
FUNCTION = "scan"
|
||
OUTPUT_NODE = True
|
||
DESCRIPTION = "Scan local models and compute the best PromptForge settings for this GPU. Runs on its own; no generation."
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"scope": (["all local", "selected"], {"default": "all local", "tooltip": "Scan every locally available model, or only the selected one."}),
|
||
"model": (MODEL_OPTIONS, {"default": _DEFAULT_MODEL, "tooltip": "Model to plan for when scope is 'selected'."}),
|
||
"quantization": (QUANT_OPTIONS, {"default": "auto", "tooltip": "Quantization to plan for; auto picks the best local quant."}),
|
||
},
|
||
}
|
||
|
||
def scan(self, scope, model, quantization, **kwargs):
|
||
result = _scan_configs({"scope": scope, "model": model, "quantization": quantization})
|
||
for line in result["report"].splitlines():
|
||
print(f"[PromptForge] {line}")
|
||
return {"ui": {"text": [result["report"]]}, "result": (result["report"],)}
|
||
|
||
|
||
try:
|
||
import asyncio
|
||
from aiohttp import web
|
||
from server import PromptServer
|
||
|
||
@PromptServer.instance.routes.post("/promptforge/test")
|
||
async def _promptforge_test(request):
|
||
data = await request.json()
|
||
return web.json_response(await asyncio.to_thread(_test_settings, data))
|
||
|
||
@PromptServer.instance.routes.post("/promptforge/scan")
|
||
async def _promptforge_scan(request):
|
||
data = await request.json()
|
||
return web.json_response(await asyncio.to_thread(_scan_configs, data))
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
class PROMPTFORGE_PromptEnhancer:
|
||
CATEGORY = "🔮 PromptForge/Prompts"
|
||
RETURN_TYPES = ("STRING", "STRING")
|
||
RETURN_NAMES = ("ENHANCED_PROMPT", "THINKING")
|
||
FUNCTION = "enhance"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"model": ("PROMPTFORGE_MODEL", {"tooltip": "MODEL output from a PromptForge Model node (source, model, quant and runtime settings)."}),
|
||
"prompt_text": ("STRING", {"default": "", "multiline": True, "tooltip": "Raw user request / media-derived text to enhance. Can be loose natural language, random idea, detailed description, keyword pile, anything."}),
|
||
"mode": (MODE_OPTIONS, {"default": "text_only", "tooltip": "text_only ignores all media; inputs uses every connected image (image / image2 / image3) and the sampled video frames."}),
|
||
"enhancement_preset": (PRESET_OPTIONS, {"default": "auto", "tooltip": "Preset system prompt: auto, or one of the PromptForge presets (Krea 2 T2I, MOSS Image T2I, Qwen Image T2I, generic styles...)."}),
|
||
"seed": ("INT", {"default": 1, "min": 1, "max": 2 ** 32 - 1, "tooltip": "Sampling seed."}),
|
||
"max_tokens": ("INT", {"default": 3072, "min": 32, "max": 65536, "tooltip": "Lower bound for the prompt length. Presets raise it automatically to a safe budget (H3 action 1600, idea 1200, image 700, tags 220) so the prompt is not cut off."}),
|
||
"temperature": ("FLOAT", {"default": 0.7, "min": 0.1, "max": 1.0, "step": 0.01, "tooltip": "Sampling temperature. Qwen3-VL instruct is tuned for 0.7; 0.6 for thinking mode."}),
|
||
"top_p": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Nucleus cutoff. Qwen3-VL instruct: 0.8 (thinking: 0.95)."}),
|
||
"top_k": ("INT", {"default": 20, "min": 0, "max": 100, "tooltip": "Qwen3-VL is tuned for 20."}),
|
||
"num_beams": ("INT", {"default": 1, "min": 1, "max": 8, "tooltip": "Beam search width. 1 = fast sampling; >1 trades speed for stability."}),
|
||
"repetition_penalty": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 2.0, "step": 0.01, "tooltip": "Qwen3-VL is tuned for 1.0; raise to 1.05-1.1 only if a model still loops."}),
|
||
"presence_penalty": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "Discourages tokens already seen. 1.5 is the Qwen serving value; 0.5-1.0 is safer for long structured prompts. 0 = off."}),
|
||
"system_prompt_override": ("STRING", {"default": "", "multiline": True, "tooltip": "Custom system prompt override. Leave empty to use the selected preset."}),
|
||
},
|
||
"optional": {
|
||
"video_seconds": ("INT", {"default": 5, "min": 1, "max": 60, "tooltip": "Target video duration in seconds."}),
|
||
"image": ("IMAGE", {"tooltip": "First reference image."}),
|
||
"image2": ("IMAGE", {"tooltip": "Second reference image / edit source."}),
|
||
"image3": ("IMAGE", {"tooltip": "Third reference image / edit source."}),
|
||
"video": ("IMAGE", {"tooltip": "Video frames input. Use frame_count to control how many frames are sampled."}),
|
||
"frame_count": ("INT", {"default": 16, "min": 1, "max": 64, "tooltip": "Frames sampled from video input."}),
|
||
"negative_prompt": ("STRING", {"default": "", "multiline": True, "tooltip": "What the generated prompt/scene must not contain."}),
|
||
"pipe": ("PROMPTFORGE_PIPE", {"tooltip": "From a PromptForge Loader. With enhancement_preset on auto, the preset for that model family is used."}),
|
||
"dialogue": ("STRING", {"default": "", "multiline": True, "tooltip": "Video presets only (MiniMax H3). What the dialogue should be about and how it should sound, in your own words. The enhancer writes the lines, speaker IDs and tags, sized to video_seconds. Empty (with language on auto) = no dialogue instructions."}),
|
||
"dialogue_language": ([DIALOGUE_AUTO, *DIALOGUE_LANGUAGES, DIALOGUE_OTHER], {"default": DIALOGUE_AUTO, "tooltip": "Language of the spoken lines. The list is what MiniMax H3 speaks reliably."}),
|
||
"dialogue_language_other": ("STRING", {"default": "", "tooltip": "Used when dialogue_language is 'other', for example Polish. Languages outside the list are only partially supported by the video model."}),
|
||
"auto_retry": ("BOOLEAN", {"default": True, "tooltip": "Regenerate with a new seed (up to 3 tries) when the output looks empty, repetitive, leaked reasoning or a refusal. Turn off to always keep the first output."}),
|
||
"continuation": ("STRING", {"default": "", "multiline": True, "tooltip": "Video presets only (MiniMax H3). A continuation brief added to the system prompt; the Endless Video node fills this in."}),
|
||
"context_video": ("IMAGE", {"tooltip": "Video presets only (MiniMax H3). Frames the generator will start from. The model looks at them to continue the action, but they are never cited as <Video 1>."}),
|
||
"rules": ("STRING", {"default": "", "multiline": True, "tooltip": "Video presets only (MiniMax H3). Extra binding rules appended to the system prompt, for example the one-shot action rules of the H3 Generate node."}),
|
||
"action": ("PROMPTFORGE_ACTION", {"tooltip": "From a PromptForge Action node: a pose, place, camera, light and outfit as a binding brief."}),
|
||
"use_action": ("BOOLEAN", {"default": True, "tooltip": "Off ignores the connected Action node and writes from your own text alone; use it for custom scenes while the node stays connected."}),
|
||
},
|
||
}
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Node entry point
|
||
# ------------------------------------------------------------------ #
|
||
def enhance(
|
||
self,
|
||
model,
|
||
prompt_text,
|
||
mode,
|
||
enhancement_preset,
|
||
seed,
|
||
max_tokens,
|
||
temperature,
|
||
top_p,
|
||
top_k,
|
||
num_beams,
|
||
repetition_penalty,
|
||
presence_penalty,
|
||
system_prompt_override,
|
||
video_seconds=5,
|
||
image=None,
|
||
image2=None,
|
||
image3=None,
|
||
video=None,
|
||
frame_count=16,
|
||
negative_prompt="",
|
||
auto_retry=True,
|
||
pipe=None,
|
||
dialogue="",
|
||
dialogue_language=DIALOGUE_AUTO,
|
||
dialogue_language_other="",
|
||
continuation="",
|
||
context_video=None,
|
||
rules="",
|
||
action=None,
|
||
use_action=True,
|
||
**kwargs,
|
||
):
|
||
if enhancement_preset == "auto" and pipe:
|
||
enhancement_preset = _family_preset(pipe.get("family"), any(value is not None for value in (image, image2, image3)), video is not None)
|
||
print("[PromptForge] request received.")
|
||
cfg = model if isinstance(model, dict) else {}
|
||
source = cfg.get("source", "safetensors")
|
||
safetensors_model = cfg.get("safetensors_model", "auto")
|
||
gguf_model = cfg.get("gguf_model", "auto")
|
||
quantization = cfg.get("quantization", "auto")
|
||
context_size = cfg.get("context_size", 16384)
|
||
gpu_layers = cfg.get("gpu_layers", -1)
|
||
batch_size = cfg.get("batch_size", 2048)
|
||
ubatch_size = cfg.get("ubatch_size", 512)
|
||
kv_cache_type = cfg.get("kv_cache_type", "f16")
|
||
offload_kqv = cfg.get("offload_kqv", True)
|
||
swa_full = cfg.get("swa_full", False)
|
||
flash_attn = cfg.get("flash_attn", "auto")
|
||
local_model_path = cfg.get("local_model_path", "")
|
||
mmproj_path = cfg.get("mmproj_path", "")
|
||
clip = cfg.get("clip")
|
||
override_text = (system_prompt_override or "").strip()
|
||
sys_text = override_text
|
||
if not sys_text:
|
||
key = enhancement_preset if enhancement_preset and enhancement_preset in PROMPTFORGE_PRESETS else DEFAULT_PRESET
|
||
if enhancement_preset and enhancement_preset not in PRESET_OPTIONS:
|
||
print(f"[PromptForge] preset '{enhancement_preset}' not found; using {DEFAULT_PRESET}.")
|
||
entry = PROMPTFORGE_PRESETS.get(key) or {}
|
||
sys_text = entry.get("text") or entry.get("system_prompt") or (PROMPTFORGE_PRESETS.get(DEFAULT_PRESET) or {}).get("text") or ""
|
||
|
||
user_text = (prompt_text or "").strip()
|
||
action_block = ""
|
||
if use_action and isinstance(action, dict) and action.get("action") in image_actions.ACTIONS:
|
||
if action.get("source") == "input image" and not any(value is not None for value in (image, image2, image3)):
|
||
print("[PromptForge] the Action node expects an input image but none is connected; writing the person from text instead.")
|
||
action = dict(action, source="text description")
|
||
action_block = image_actions.brief(
|
||
action["action"], (action.get("subject_1", ""), action.get("subject_2", "")),
|
||
action.get("camera", "auto"), action.get("place", "auto"), action.get("lighting", "auto"),
|
||
action.get("outfit", "as in the image"), action.get("style", "auto"),
|
||
action.get("notes", ""), action.get("overrides", ""), action.get("source", "input image"), action.get("instructions", ""))
|
||
duration_text = f"{float(video_seconds):.2f}"
|
||
duration_line = ""
|
||
if sys_text and "{video_seconds}" in sys_text:
|
||
sys_text = sys_text.replace("{video_seconds}", duration_text)
|
||
elif "minimax" in (enhancement_preset or "").lower():
|
||
duration_line = f"Target video duration: {duration_text} seconds."
|
||
guard = OUTPUT_GUARD
|
||
guard += " Do not output chain-of-thought or reasoning; start with the final prompt."
|
||
minimax_key = None if override_text else _minimax_mode(enhancement_preset)
|
||
minimax_block = ""
|
||
if minimax_key:
|
||
image, image2, image3, video, minimax_block = self._minimax_inputs(minimax_key, mode, image, image2, image3, video, frame_count, duration_text)
|
||
media_line = self._media_inventory(mode, image, image2, image3, video, frame_count)
|
||
labelled_video = video is not None
|
||
if minimax_key and mode != "text_only" and self._has_media(context_video):
|
||
video = context_video
|
||
dialogue_block = _dialogue_block(dialogue, dialogue_language, dialogue_language_other, float(video_seconds)) if minimax_key else ""
|
||
negative_text = (negative_prompt or "").strip()
|
||
negative_line = ""
|
||
if negative_text:
|
||
negative_line = (
|
||
"# Negative constraints\n"
|
||
"The generated prompt and the scene it describes must NOT contain any of the following: "
|
||
+ negative_text +
|
||
". Write the scene so they are absent, converting each exclusion into a positive description "
|
||
"of the desired state; never name them inside the prompt itself."
|
||
)
|
||
continuation_block = (continuation or "").strip() if minimax_key else ""
|
||
system_parts = [p for p in (sys_text, media_line, minimax_block, continuation_block, dialogue_block, rules, negative_line, guard) if p and p.strip()]
|
||
action_active = bool(action_block)
|
||
user_parts = [p for p in ("" if action_active else user_text, action_block, duration_line) if p and p.strip()]
|
||
system_prompt = "\n\n".join(system_parts)
|
||
user_prompt = "\n\n".join(user_parts)
|
||
full_prompt = "\n\n".join(p for p in (system_prompt, user_prompt) if p)
|
||
self.last_request = {
|
||
"preset": enhancement_preset, "mode": mode, "images": sum(value is not None for value in (image, image2, image3)), "video": video is not None, "user": user_prompt,
|
||
"sections": [(name, text) for name, text in (("media", media_line), ("MiniMax mode", minimax_block), ("continuation", continuation_block), ("dialogue", dialogue_block), ("rules", rules), ("negative prompt", negative_line),
|
||
("output rules", guard), ("template", sys_text)) if text and text.strip()],
|
||
}
|
||
|
||
eff_local = (local_model_path or "").strip()
|
||
if source == "connected clip":
|
||
if clip is None:
|
||
raise ValueError("[PromptForge] Source is 'connected clip' but no CLIP input is connected.")
|
||
engine = "text_encoder"
|
||
model_name = "auto"
|
||
clip_for_generate = clip
|
||
elif source == "GGUF":
|
||
engine = "gguf"
|
||
model_name = gguf_model
|
||
clip_for_generate = None
|
||
else:
|
||
engine = "text_encoder"
|
||
model_name = safetensors_model
|
||
clip_for_generate = None
|
||
|
||
if eff_local:
|
||
is_gguf = eff_local.lower().endswith(".gguf")
|
||
if source == "GGUF" and not is_gguf:
|
||
raise ValueError("[PromptForge] Source is GGUF but local_model_path is not a .gguf file.")
|
||
if source != "GGUF" and is_gguf:
|
||
raise ValueError("[PromptForge] Source is safetensors but local_model_path points to a .gguf file; select the GGUF source.")
|
||
elif source == "safetensors" and model_name == "none":
|
||
raise ValueError("[PromptForge] No safetensors model available; add one to the model dirs or set local_model_path.")
|
||
elif source == "GGUF" and model_name == "none":
|
||
raise ValueError("[PromptForge] No GGUF model available; add a .gguf to the model dirs or set local_model_path.")
|
||
|
||
preset_l = (enhancement_preset or "").lower()
|
||
short_form = any(word in preset_l for word in ("tags", "refine", "technical"))
|
||
if short_form:
|
||
minimum = 220
|
||
elif any(word in preset_l for word in ("minimax", "h3")):
|
||
minimum = 1600 if use_action else 1200
|
||
elif any(word in preset_l for word in ("moss", "krea", "qwen", "anima", "pony", "illustrious", "flux", "image")):
|
||
minimum = 700
|
||
else:
|
||
minimum = 900
|
||
if max_tokens < minimum:
|
||
print(f"[PromptForge] max_tokens {max_tokens} is below the {minimum} this preset needs; using {minimum}.")
|
||
max_tokens = minimum
|
||
min_chars = 40 if short_form else (400 if max_tokens >= 192 else 120)
|
||
def run(attempt_seed, attempt_system):
|
||
return self._generate(
|
||
full_prompt=full_prompt,
|
||
system_prompt=attempt_system,
|
||
user_prompt=user_prompt,
|
||
mode=mode,
|
||
engine=engine,
|
||
model_name=model_name,
|
||
quantization=quantization,
|
||
context_size=context_size,
|
||
gpu_layers=gpu_layers,
|
||
batch_size=batch_size,
|
||
ubatch_size=ubatch_size,
|
||
kv_cache_type=kv_cache_type,
|
||
offload_kqv=offload_kqv,
|
||
swa_full=swa_full,
|
||
flash_attn=flash_attn,
|
||
max_tokens=max_tokens,
|
||
temperature=temperature,
|
||
top_p=top_p,
|
||
top_k=top_k,
|
||
num_beams=num_beams,
|
||
repetition_penalty=repetition_penalty,
|
||
presence_penalty=presence_penalty,
|
||
image=image,
|
||
image2=image2,
|
||
image3=image3,
|
||
video=video,
|
||
frame_count=frame_count,
|
||
seed=attempt_seed,
|
||
clip=clip_for_generate,
|
||
local_model_path=local_model_path,
|
||
mmproj_path=mmproj_path,
|
||
disable_thinking=True,
|
||
auto_retry=auto_retry,
|
||
min_chars=min_chars,
|
||
)
|
||
answer, thinking, best_answer = "", "", ""
|
||
best_words = best_complete = 0
|
||
for attempt in range(3 if auto_retry else 1):
|
||
enhanced = run(seed, system_prompt)
|
||
answer, thinking = _split_thinking(enhanced)
|
||
if minimax_key:
|
||
answer = _normalize_minimax_sections(answer, minimax_key)
|
||
if minimax_key and rules.strip():
|
||
answer = _strip_inner_timestamps(answer, minimax_key)
|
||
if minimax_key == "ref2va":
|
||
answer = _inject_reference_labels(answer, user_text)
|
||
problem = _minimax_format_problem(answer, minimax_key, bool(rules.strip()), user_text,
|
||
sum(value is not None for value in (image, image2, image3)))
|
||
words = len(re.findall(r"\w+", answer))
|
||
complete = len(re.findall(r"(?im)^(?:" + "|".join(_MINIMAX_SECTIONS.get(minimax_key, ())) + r"):", answer))
|
||
if (complete, words) > (best_complete, best_words):
|
||
best_answer, best_words, best_complete = answer, words, complete
|
||
if not problem or not auto_retry or attempt == 2:
|
||
break
|
||
print(f"[PromptForge] output rejected ({problem}); retrying with a new seed.")
|
||
seed = _next_seed(seed)
|
||
system_prompt += (f"\n\n# Correction\nYour previous answer broke the required format: {problem}. "
|
||
"Write the whole prompt again and follow the output contract and the action rules exactly.")
|
||
if problem.startswith("body too short") and best_answer:
|
||
answer = best_answer
|
||
if minimax_key:
|
||
answer = _ensure_minimax_sections(answer, minimax_key)
|
||
answer = _normalize_minimax_sections(answer, minimax_key)
|
||
if minimax_key == "ref2va":
|
||
answer = _inject_reference_labels(answer, user_text)
|
||
if dialogue_block and auto_retry and (dialogue or "").strip():
|
||
language = _dialogue_language(dialogue_language, dialogue_language_other)
|
||
for attempt in range(1, 2):
|
||
problem = _dialogue_problem(answer, language)
|
||
if not problem:
|
||
break
|
||
print(f"[PromptForge] dialogue rejected ({problem}); retrying once.")
|
||
seed = _next_seed(seed)
|
||
enhanced = run(seed, system_prompt + f"\n\n# Correction\nYour previous answer broke the dialogue format: {problem}. Write the whole prompt again and follow the Dialogue section exactly.")
|
||
answer, thinking = _split_thinking(enhanced)
|
||
answer = _clean_prompt_output(answer)
|
||
answer, cut = _cut_runaway_tail(answer)
|
||
answer = _land_text(answer)
|
||
if cut:
|
||
print("[PromptForge] the prompt writer lost track near the end (a looping word or an endless sentence); the runaway tail was cut. A lower temperature, a higher repetition penalty or a smaller max_tokens helps.")
|
||
if minimax_key:
|
||
answer = _drop_absent_references(answer, sum(value is not None for value in (image, image2, image3)), labelled_video, user_text)
|
||
preset_lower = (enhancement_preset or "").lower()
|
||
if not any(marker in preset_lower for marker in ("minimax", "wan", "ltx", "ace-step")):
|
||
answer = _strip_video_meta(answer)
|
||
if not answer.strip():
|
||
detail = "only reasoning/thinking" if thinking.strip() else "no usable text"
|
||
raise ValueError(
|
||
f"[PromptForge] the model returned {detail} and no final prompt. "
|
||
"Use a non-thinking model (Qwen3-VL for vision), raise max_tokens, or change the preset."
|
||
)
|
||
return (answer, thinking)
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Model lifecycle: unload after the prompt has been generated
|
||
# ------------------------------------------------------------------ #
|
||
def _release(self, llm=None):
|
||
if llm is not None:
|
||
try:
|
||
llm.close()
|
||
except Exception:
|
||
pass
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Generation dispatch
|
||
# ------------------------------------------------------------------ #
|
||
def _generate(self, full_prompt, system_prompt, user_prompt, mode, engine, model_name, quantization, context_size, gpu_layers, batch_size, ubatch_size, kv_cache_type, offload_kqv, swa_full, flash_attn, max_tokens, temperature, top_p, top_k, num_beams, repetition_penalty, presence_penalty, image, image2, image3, video, frame_count, seed, clip, local_model_path, mmproj_path, disable_thinking, auto_retry, min_chars=20):
|
||
try:
|
||
if engine == "gguf":
|
||
return self._generate_gguf(
|
||
full_prompt=full_prompt, system_prompt=system_prompt, user_prompt=user_prompt,
|
||
mode=mode, model_name=model_name, quantization=quantization,
|
||
context_size=context_size, gpu_layers=gpu_layers, batch_size=batch_size, ubatch_size=ubatch_size,
|
||
kv_cache_type=kv_cache_type, offload_kqv=offload_kqv, swa_full=swa_full,
|
||
flash_attn=flash_attn,
|
||
max_tokens=max_tokens, temperature=temperature, top_p=top_p, top_k=top_k,
|
||
num_beams=num_beams, repetition_penalty=repetition_penalty, presence_penalty=presence_penalty, seed=seed,
|
||
image=image, image2=image2, image3=image3, video=video, frame_count=frame_count,
|
||
local_model_path=local_model_path, mmproj_path=mmproj_path,
|
||
disable_thinking=disable_thinking, auto_retry=auto_retry, min_chars=min_chars,
|
||
)
|
||
return self._generate_text_encoder(
|
||
full_prompt=full_prompt, system_prompt=system_prompt, user_prompt=user_prompt,
|
||
mode=mode, model_name=model_name, quantization=quantization,
|
||
max_tokens=max_tokens, temperature=temperature, top_p=top_p, top_k=top_k,
|
||
num_beams=num_beams, repetition_penalty=repetition_penalty, presence_penalty=presence_penalty,
|
||
image=image, image2=image2, image3=image3, video=video, frame_count=frame_count,
|
||
seed=seed, clip=clip, local_model_path=local_model_path,
|
||
disable_thinking=disable_thinking, auto_retry=auto_retry, min_chars=min_chars,
|
||
)
|
||
finally:
|
||
gc.collect()
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# GGUF / llama.cpp engine
|
||
# ------------------------------------------------------------------ #
|
||
def _generate_gguf(self, full_prompt, system_prompt, user_prompt, mode, model_name, quantization, context_size, gpu_layers, batch_size, ubatch_size, kv_cache_type, offload_kqv, swa_full, flash_attn, max_tokens, temperature, top_p, top_k, num_beams, repetition_penalty, presence_penalty, seed, image, image2, image3, video, frame_count, local_model_path, mmproj_path, disable_thinking, auto_retry, min_chars=20):
|
||
entry = None
|
||
gguf_path = None
|
||
eff_local = (local_model_path or "").strip() or _local_model_option(model_name)
|
||
if not eff_local and str(model_name).lower().endswith(".gguf") and Path(str(model_name)).is_file():
|
||
eff_local = str(model_name)
|
||
if eff_local and str(eff_local).lower().endswith(".gguf"):
|
||
gguf_path = Path(str(eff_local))
|
||
if gguf_path is None:
|
||
resolved = _resolve_model_entry(model_name) if (model_name and model_name != "auto") else _resolve_model_entry("auto")
|
||
if resolved is None:
|
||
raise ValueError(f"[PromptForge] GGUF model '{model_name}' is not in the catalog and was not found in the model folders on this machine.")
|
||
entry = resolved.get("entry") if resolved.get("kind") == "gguf" else None
|
||
if entry is None:
|
||
raise ValueError(f"[PromptForge] '{model_name}' is a safetensors catalog entry; pick a [gguf] option for the GGUF engine.")
|
||
if entry is not None:
|
||
exact = _find_gguf(entry, quantization, download=False, exact=True)
|
||
if exact is not None:
|
||
gguf_path = exact
|
||
else:
|
||
local_any = _find_gguf(entry, "auto", download=False)
|
||
if local_any is not None:
|
||
print(f"[PromptForge] {quantization} is not local; using {local_any.name}.")
|
||
gguf_path = local_any
|
||
else:
|
||
gguf_path = _find_gguf(entry, quantization)
|
||
if gguf_path is None:
|
||
raise ValueError("[PromptForge] No GGUF model selected / available. Put the .gguf in a scanned models dir (text_encoders / LLM / gguf) or set local_model_path.")
|
||
|
||
images = self._mode_images(mode, image, image2, image3, video, frame_count) if mode != "text_only" else []
|
||
mmproj = _resolve_mmproj(gguf_path, entry, mmproj_path) if images else None
|
||
if images and mmproj is None:
|
||
raise ValueError("[PromptForge] Vision GGUF needs an mmproj file next to the model (e.g. mmproj-F16.gguf). Put one in the same folder, set mmproj_path, or switch mode to text_only.")
|
||
|
||
from llama_cpp import Llama
|
||
|
||
n_ctx = max(512, int(context_size))
|
||
estimated = len(full_prompt) // 3 + len(images) * 1024 + int(max_tokens) + 512
|
||
if estimated > n_ctx:
|
||
print(f"[PromptForge] context_size={n_ctx} may truncate this request (about {estimated} tokens needed).")
|
||
|
||
kv_type = _kv_type_value(kv_cache_type)
|
||
flash_types = {"auto": -1, "on": 1, "off": 0}
|
||
flash_type = flash_types.get(flash_attn, flash_types["auto"])
|
||
n_batch = min(max(64, int(batch_size)), n_ctx)
|
||
n_ubatch = max(32, min(int(ubatch_size), n_batch))
|
||
|
||
llama_kwargs = {
|
||
"model_path": str(gguf_path),
|
||
"n_ctx": n_ctx,
|
||
"n_batch": n_batch,
|
||
"n_ubatch": n_ubatch,
|
||
"n_gpu_layers": int(gpu_layers),
|
||
"offload_kqv": bool(offload_kqv),
|
||
"swa_full": bool(swa_full),
|
||
"type_k": kv_type,
|
||
"type_v": kv_type,
|
||
"use_mmap": True,
|
||
"verbose": False,
|
||
}
|
||
supports_flash = "flash_attn_type" in inspect.signature(Llama.__init__).parameters
|
||
if supports_flash:
|
||
llama_kwargs["flash_attn_type"] = flash_type
|
||
if images:
|
||
from llama_cpp import llama_chat_format
|
||
model_name_lower = f"{gguf_path.name} {(entry or {}).get('name', '')}".lower()
|
||
handler_name = _gguf_vision_handler(model_name_lower)
|
||
handler_cls = getattr(llama_chat_format, handler_name, None)
|
||
if handler_cls is None:
|
||
raise ValueError(f"[PromptForge] this llama-cpp-python build has no {handler_name}; upgrade llama-cpp-python for this vision architecture.")
|
||
llama_kwargs["chat_handler"] = handler_cls(clip_model_path=str(mmproj), verbose=False)
|
||
else:
|
||
metadata = _gguf_metadata(gguf_path)
|
||
if "tokenizer.chat_template" not in metadata:
|
||
chat_format = _gguf_chat_format(f"{gguf_path.name} {(entry or {}).get('name', '')}")
|
||
if chat_format:
|
||
llama_kwargs["chat_format"] = chat_format
|
||
if llama_kwargs["n_gpu_layers"] != 0:
|
||
try:
|
||
free_bytes = torch.cuda.mem_get_info()[0]
|
||
except Exception:
|
||
free_bytes = 0
|
||
required = _gguf_memory_estimate(gguf_path, n_ctx, llama_kwargs["n_gpu_layers"], mmproj)
|
||
if free_bytes < required:
|
||
_free_vram(needed=int(required))
|
||
|
||
fallbacks = []
|
||
if kv_type != _kv_type_value("f16"):
|
||
fallbacks.append(("kv_cache_type", {"type_k": _kv_type_value("f16"), "type_v": _kv_type_value("f16")}))
|
||
if supports_flash and flash_type == flash_types["on"]:
|
||
fallbacks.append(("flash_attn", {"flash_attn_type": flash_types["auto"]}))
|
||
if bool(swa_full):
|
||
fallbacks.append(("swa_full", {"swa_full": False}))
|
||
|
||
llm = None
|
||
last_error = None
|
||
retried_layers = False
|
||
index = 0
|
||
while index <= len(fallbacks):
|
||
try:
|
||
llm = Llama(**llama_kwargs)
|
||
break
|
||
except (RuntimeError, ValueError) as exc:
|
||
message = str(exc)
|
||
if "Failed to load model" in message and llama_kwargs["n_gpu_layers"] != 0 and not retried_layers:
|
||
retried_layers = True
|
||
llama_kwargs["n_gpu_layers"] = _gguf_layers_that_fit(gguf_path, n_ctx, mmproj)
|
||
print(f"[PromptForge] not enough free VRAM for the whole GGUF; retrying with {llama_kwargs['n_gpu_layers']} GPU layers.")
|
||
gc.collect()
|
||
continue
|
||
if "Failed to create llama context" not in message:
|
||
raise
|
||
last_error = message
|
||
if index >= len(fallbacks):
|
||
break
|
||
label, overrides = fallbacks[index]
|
||
index += 1
|
||
print(f"[PromptForge] llama.cpp rejected {label}; retrying with the safer setting.")
|
||
llama_kwargs.update(overrides)
|
||
gc.collect()
|
||
if llm is None:
|
||
raise ValueError(
|
||
f"[PromptForge] llama.cpp could not create a {n_ctx}-token context. "
|
||
"Lower context_size, gpu_layers or batch_size, set kv_cache_type to q8_0/q4_0, or turn offload_kqv off. "
|
||
+ (last_error or "")
|
||
)
|
||
print(f"[PromptForge] GGUF ready: ctx {n_ctx}, gpu_layers {llama_kwargs['n_gpu_layers']}, kv {kv_cache_type}.")
|
||
|
||
model_label_lower = f"{gguf_path.name} {(entry or {}).get('name', '')}".lower()
|
||
gemma4 = "gemma-4" in model_label_lower or "gemma4" in model_label_lower
|
||
chat_template = llm.metadata.get("tokenizer.chat_template") if not images else None
|
||
user_text_value = user_prompt or full_prompt
|
||
if images and gemma4:
|
||
# the vision chat handler renders its own template: fold the system turn into the user turn and ask for no thought channel
|
||
if (system_prompt or "").strip():
|
||
user_text_value = f"{system_prompt.strip()}\n\n{user_text_value}".strip()
|
||
user_text_value += "\n\nIMPORTANT: Return the final prompt directly. Do not use a thought channel. Do not write reasoning."
|
||
if images:
|
||
content = [{"type": "text", "text": user_text_value}]
|
||
content.extend({"type": "image_url", "image_url": {"url": self._pil_data_url(p)}} for p in images)
|
||
else:
|
||
content = user_text_value
|
||
messages = []
|
||
if (system_prompt or "").strip() and not (images and gemma4):
|
||
messages.append({"role": "system", "content": system_prompt})
|
||
messages.append({"role": "user", "content": content})
|
||
completion_kwargs = {
|
||
"max_tokens": max_tokens,
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
"seed": seed,
|
||
}
|
||
if top_k:
|
||
completion_kwargs["top_k"] = top_k
|
||
# older llama-cpp-python builds lack penalty_last_n / reasoning_budget
|
||
supported = inspect.signature(llm.create_completion if chat_template else llm.create_chat_completion).parameters
|
||
for key, value in (("frequency_penalty", 0.1), ("presence_penalty", presence_penalty), ("penalty_last_n", 256)):
|
||
if key in supported:
|
||
completion_kwargs[key] = value
|
||
if chat_template:
|
||
# The model's own template with thinking disabled (empty think block / primed thought channel) is the
|
||
# reliable non-thinking mode; llama.cpp adds BOS itself, so the template gets an empty bos_token.
|
||
prompt_text = render_jinja_template(
|
||
[messages], chat_template=chat_template, add_generation_prompt=True, bos_token="", eos_token="",
|
||
enable_thinking=not disable_thinking,
|
||
)[0][0]
|
||
completion_kwargs["prompt"] = prompt_text
|
||
completion_kwargs["stop"] = ["<|im_end|>", "<turn|>", "<end_of_turn>", "<|eot_id|>", "<|end|>"]
|
||
else:
|
||
completion_kwargs["messages"] = messages
|
||
if disable_thinking and "reasoning_budget" in supported:
|
||
completion_kwargs["reasoning_budget"] = 0
|
||
|
||
def attempt(seed_value):
|
||
import comfy.utils as comfy_utils
|
||
kwargs = dict(completion_kwargs)
|
||
kwargs["seed"] = seed_value
|
||
kwargs["stream"] = True
|
||
started = time.time()
|
||
progress = comfy_utils.ProgressBar(max_tokens)
|
||
parts, tokens = [], 0
|
||
if chat_template:
|
||
stream = llm.create_completion(**kwargs)
|
||
else:
|
||
stream = llm.create_chat_completion(**kwargs)
|
||
for chunk in stream:
|
||
choice = (chunk.get("choices") or [{}])[0]
|
||
if chat_template:
|
||
piece = choice.get("text") or ""
|
||
else:
|
||
piece = (choice.get("delta") or {}).get("content") or ""
|
||
if piece:
|
||
parts.append(piece)
|
||
tokens += 1
|
||
if tokens % 8 == 0:
|
||
progress.update_absolute(min(tokens, max_tokens), max_tokens)
|
||
used = (chunk.get("usage") or {}).get("completion_tokens")
|
||
if used:
|
||
tokens = used
|
||
progress.update_absolute(max_tokens, max_tokens)
|
||
print(f"[PromptForge] GGUF generated {tokens} tokens in {time.time() - started:.1f}s.")
|
||
return "".join(parts).strip()
|
||
|
||
try:
|
||
enhanced = _generate_checked(attempt, seed, min_chars, auto_retry)
|
||
except ValueError as exc:
|
||
if images and "mmproj" in str(exc).lower():
|
||
expected = (entry or {}).get("mmproj_file") or "mmproj-*.gguf"
|
||
raise ValueError(f"[PromptForge] Vision GGUF could not load mmproj '{mmproj}'. It must match the text model {gguf_path.name}. Download {expected} next to the model or set mmproj_path.") from exc
|
||
raise
|
||
finally:
|
||
self._release(llm=llm)
|
||
return enhanced
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Text encoder engines (native ComfyUI CLIP route / bitsandbytes route)
|
||
# ------------------------------------------------------------------ #
|
||
def _generate_text_encoder(self, full_prompt, system_prompt, user_prompt, mode, model_name, quantization, max_tokens, temperature, top_p, top_k, num_beams, repetition_penalty, presence_penalty, image, image2, image3, video, frame_count, seed, clip, local_model_path, disable_thinking, auto_retry, min_chars=20):
|
||
if clip is not None:
|
||
enhanced = self._generate_with_clip(
|
||
system_prompt=system_prompt, user_prompt=user_prompt or full_prompt, mode=mode, max_tokens=max_tokens, temperature=temperature,
|
||
top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, presence_penalty=presence_penalty,
|
||
seed=seed, image=image, image2=image2, image3=image3, video=video, frame_count=frame_count,
|
||
clip=clip, disable_thinking=disable_thinking, auto_retry=auto_retry, min_chars=min_chars,
|
||
)
|
||
return enhanced
|
||
|
||
eff_local = (local_model_path or "").strip() or _local_model_option(model_name)
|
||
if not eff_local and Path(str(model_name)).exists():
|
||
eff_local = str(model_name)
|
||
src = eff_local or ""
|
||
resolved = _resolve_model_entry(model_name) if (model_name and model_name != "auto") else _resolve_model_entry("auto")
|
||
if not src and resolved is not None and resolved.get("kind") == "hf":
|
||
src = resolved.get("local_path") or resolved.get("repo_id") or ""
|
||
if not src:
|
||
raise ValueError(f"[PromptForge] safetensors model '{model_name}' is not in the catalog and was not found in the model folders on this machine.")
|
||
params_b = _safetensors_params_b(resolved, src)
|
||
|
||
quant = {"8bit": "q8", "4bit": "q4"}.get(quantization or "auto", quantization or "auto")
|
||
if quant in ("q5", "q6"):
|
||
print(f"[PromptForge] {quant} is a GGUF quant level; loading safetensors in q8 instead.")
|
||
quant = "q8"
|
||
if quant == "q8" and params_b and torch.cuda.is_available():
|
||
try:
|
||
vram = torch.cuda.get_device_properties(0).total_memory
|
||
except Exception:
|
||
vram = 0
|
||
if vram and float(params_b) * 1.15 * 1024 ** 3 > vram * 0.85:
|
||
print("[PromptForge] q8 does not fit this GPU; using q4 instead.")
|
||
quant = "q4"
|
||
src_path = Path(src)
|
||
if not src_path.exists():
|
||
from huggingface_hub import snapshot_download
|
||
src = snapshot_download(repo_id=src, local_dir=None)
|
||
src_path = Path(src)
|
||
|
||
if quant in ("q8", "q4") and not src_path.is_dir():
|
||
print(f"[PromptForge] {quant} quantization needs an HF model directory; loading this single file in fp16 instead.")
|
||
quant = "auto"
|
||
|
||
if src_path.is_dir():
|
||
if torch.cuda.is_available():
|
||
_free_vram()
|
||
if (src_path / "config.json").is_file() and "qwen3_5" in (src_path / "config.json").read_text(encoding="utf-8") and importlib.util.find_spec("fla") is None:
|
||
print("[PromptForge] this Qwen3.5 folder runs on slow fallback kernels (3x slower than the single-file or GGUF versions). pip install flash-linear-attention in ComfyUI's environment makes it much faster.")
|
||
if any(src_path.glob("*mtp*.safetensors")):
|
||
print("[PromptForge] this folder has a multi-token-prediction head: PromptForge Quantize Text Encoder makes a single file from it that writes about twice as fast.")
|
||
hf_dtype = torch.float16 if quant == "fp16" else _hf_compute_dtype()
|
||
if quant in ("auto", "q8") and params_b:
|
||
try:
|
||
free_bytes = torch.cuda.mem_get_info()[0]
|
||
except Exception:
|
||
free_bytes = 0
|
||
per_param = 2.0 if quant == "auto" else 1.35
|
||
if free_bytes and float(params_b) * per_param * 1024 ** 3 + 1.5 * 1024 ** 3 > free_bytes:
|
||
print(f"[PromptForge] {'bf16' if quant == 'auto' else quant} does not fit the free VRAM; using q4 instead.")
|
||
quant = "q4"
|
||
quant_cfg = None
|
||
if quant in ("q8", "q4"):
|
||
quant_cfg = BitsAndBytesConfig(
|
||
load_in_8bit=(quant == "q8"),
|
||
load_in_4bit=(quant == "q4"),
|
||
bnb_4bit_quant_type="nf4",
|
||
bnb_4bit_use_double_quant=True,
|
||
bnb_4bit_compute_dtype=hf_dtype,
|
||
llm_int8_enable_fp32_cpu_offload=True,
|
||
)
|
||
device_map = "auto"
|
||
if quant in ("q8", "q4") and torch.cuda.is_available():
|
||
device_map = {"": torch.cuda.current_device()}
|
||
load_kwargs = {
|
||
"quantization_config": quant_cfg,
|
||
"device_map": device_map,
|
||
"attn_implementation": "sdpa",
|
||
**_hf_dtype_kwargs(hf_dtype),
|
||
}
|
||
if torch.cuda.is_available() and device_map == "auto":
|
||
max_memory = _hf_max_memory()
|
||
if max_memory is not None:
|
||
load_kwargs["max_memory"] = max_memory
|
||
model = processor = tokenizer = None
|
||
|
||
def _load():
|
||
try:
|
||
return _load_hf_model(src, load_kwargs)
|
||
except (TypeError, ValueError) as exc:
|
||
message = str(exc).lower()
|
||
if "attn_implementation" in message or "sdpa" in message:
|
||
load_kwargs.pop("attn_implementation", None)
|
||
return _load_hf_model(src, load_kwargs)
|
||
raise
|
||
|
||
try:
|
||
model = _load()
|
||
device_map = getattr(model, "hf_device_map", None) or {}
|
||
if any(str(device) in ("cpu", "disk") for device in device_map.values()):
|
||
print("[PromptForge] model is partly on CPU/disk; generation will be slower (use q4 or a smaller model).")
|
||
processor = AutoProcessor.from_pretrained(src)
|
||
tokenizer = AutoTokenizer.from_pretrained(src)
|
||
conversation = []
|
||
if (system_prompt or "").strip():
|
||
conversation.append({"role": "system", "content": system_prompt})
|
||
refs = self._mode_images(mode, image, image2, image3, video, frame_count)
|
||
if refs and getattr(processor, "image_processor", None) is None:
|
||
raise ValueError("[PromptForge] The selected safetensors model has no image processor; image/video mode needs a vision-capable model.")
|
||
user_content = user_prompt or full_prompt
|
||
if refs:
|
||
user_content = [{"type": "image", "image": r} for r in refs if r is not None]
|
||
user_content.append({"type": "text", "text": user_prompt or full_prompt})
|
||
conversation.append({"role": "user", "content": user_content})
|
||
chat = processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True, enable_thinking=not disable_thinking)
|
||
processed = processor(text=chat, images=[r for r in refs if r is not None] or None, return_tensors="pt")
|
||
try:
|
||
model_device = model.get_input_embeddings().weight.device
|
||
except Exception:
|
||
model_device = next(
|
||
(param.device for param in model.parameters() if param.device.type != "meta"),
|
||
torch.device("cuda" if torch.cuda.is_available() else "cpu"),
|
||
)
|
||
if model_device.type == "meta":
|
||
model_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
allowed_inputs = set(inspect.signature(model.forward).parameters)
|
||
prepare = getattr(model, "prepare_inputs_for_generation", None)
|
||
if prepare is not None:
|
||
allowed_inputs |= set(inspect.signature(prepare).parameters)
|
||
allowed_inputs.discard("kwargs")
|
||
allowed_inputs.discard("model_kwargs")
|
||
model_inputs = {
|
||
k: (v.to(model_device, non_blocking=True) if torch.is_tensor(v) else v)
|
||
for k, v in processed.items()
|
||
if k in allowed_inputs
|
||
}
|
||
stop_tokens = [tokenizer.eos_token_id]
|
||
if getattr(tokenizer, "eot_id", None) is not None:
|
||
stop_tokens.append(tokenizer.eot_id)
|
||
kwargs = {
|
||
"max_new_tokens": max_tokens,
|
||
"repetition_penalty": repetition_penalty,
|
||
"num_beams": num_beams,
|
||
"eos_token_id": stop_tokens,
|
||
"pad_token_id": tokenizer.pad_token_id,
|
||
}
|
||
if num_beams == 1:
|
||
kwargs.update({"do_sample": True, "temperature": temperature, "top_p": top_p})
|
||
if params_b is not None and params_b >= 3.5:
|
||
kwargs["top_k"] = top_k or 20
|
||
if "presence_penalty" in inspect.signature(model.generate).parameters:
|
||
kwargs["presence_penalty"] = presence_penalty
|
||
|
||
input_len = model_inputs["input_ids"].shape[-1]
|
||
|
||
class _FiniteLogits(LogitsProcessor):
|
||
def __call__(self, input_ids, scores):
|
||
return torch.nan_to_num(scores, nan=-1e4, posinf=1e4, neginf=-1e4)
|
||
|
||
class _LoopStop(StoppingCriteria):
|
||
def __init__(self):
|
||
self.step = 0
|
||
|
||
def __call__(self, input_ids, scores, **kw):
|
||
self.step += 1
|
||
if self.step < 64 or self.step % 64:
|
||
return False
|
||
text = tokenizer.decode(input_ids[0, input_len:], skip_special_tokens=True)
|
||
return len(text) > 900 and _hard_loop(text)
|
||
|
||
kwargs["stopping_criteria"] = StoppingCriteriaList([_LoopStop()])
|
||
kwargs["logits_processor"] = LogitsProcessorList([_FiniteLogits()])
|
||
|
||
def attempt(seed_value):
|
||
torch.manual_seed(seed_value)
|
||
started = time.time()
|
||
outputs = model.generate(**model_inputs, **kwargs)
|
||
tokens = outputs.shape[-1] - input_len
|
||
if tokens:
|
||
print(f"[PromptForge] safetensors generated {tokens} tokens in {time.time() - started:.1f}s.")
|
||
return tokenizer.decode(outputs[0, input_len:], skip_special_tokens=True)
|
||
|
||
return _generate_checked(attempt, seed, min_chars, auto_retry)
|
||
except RuntimeError as exc:
|
||
if "out of memory" in str(exc).lower():
|
||
raise ValueError(
|
||
"[PromptForge] not enough free VRAM for this safetensors quantization. "
|
||
"Use q4, a smaller model, or the GGUF version."
|
||
) from exc
|
||
raise
|
||
finally:
|
||
del model, processor, tokenizer
|
||
|
||
import comfy.sd
|
||
if torch.cuda.is_available() and src_path.stat().st_size > torch.cuda.mem_get_info()[0] - 1.5 * 1024 ** 3:
|
||
print(f"[PromptForge] this writer file is {src_path.stat().st_size / 1024 ** 3:.1f} GB and only {torch.cuda.mem_get_info()[0] / 1024 ** 3:.1f} GB of VRAM is free: it streams from RAM and writes slowly. A smaller file (NVFP4, about 9 GB) or a GGUF is much faster.")
|
||
try:
|
||
clip_obj = comfy.sd.load_clip(ckpt_paths=[src])
|
||
except Exception as exc:
|
||
raise ValueError(f"[PromptForge] ComfyUI could not load this text encoder: {exc}") from exc
|
||
enhanced = self._generate_with_clip(
|
||
system_prompt=system_prompt, user_prompt=user_prompt or full_prompt, mode=mode, max_tokens=max_tokens, temperature=temperature,
|
||
top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, presence_penalty=presence_penalty,
|
||
seed=seed, image=image, image2=image2, image3=image3, video=video, frame_count=frame_count,
|
||
clip=clip_obj, disable_thinking=disable_thinking, auto_retry=auto_retry, min_chars=min_chars,
|
||
)
|
||
return enhanced
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Native ComfyUI CLIP generation (connected clip or locally loaded clip)
|
||
# ------------------------------------------------------------------ #
|
||
def _generate_with_clip(self, system_prompt, user_prompt, mode, max_tokens, temperature, top_p, top_k, repetition_penalty, presence_penalty, seed, image, image2, image3, video, frame_count, clip, disable_thinking, auto_retry, min_chars=20):
|
||
global _TOP_P_WARNED
|
||
images = self._clip_images(mode, image, image2, image3, video, frame_count)
|
||
# Gemma/LTX tokenizers take one [B, H, W, C] batch as `image`; Qwen tokenizers also accept an `images` list.
|
||
tokenize_kwargs = {"min_length": 1, "thinking": not disable_thinking, "system_prompt": system_prompt or ""}
|
||
if images:
|
||
if all(t.shape == images[0].shape for t in images):
|
||
tokenize_kwargs["image"] = torch.cat(images, dim=0)
|
||
else:
|
||
tokenize_kwargs["images"] = images
|
||
tokens = clip.tokenize(user_prompt, **tokenize_kwargs)
|
||
if system_prompt and system_prompt[:40] not in _tokens_text(clip, tokens):
|
||
# this encoder's tokenizer has no chat template (Anima / ACE Qwen encoders): wrap the turns in ChatML ourselves
|
||
chat = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
|
||
if disable_thinking:
|
||
chat += "<think>\n\n</think>\n\n"
|
||
tokens = clip.tokenize(chat, **tokenize_kwargs)
|
||
if images and not _tokens_have_images(tokens):
|
||
raise ValueError("[PromptForge] The selected text encoder has no vision support; the provided image/video would be ignored. Connect a vision-capable encoder or use text_only.")
|
||
|
||
effective_top_p = top_p
|
||
if top_p < 1.0 and not _clip_top_p_supported():
|
||
effective_top_p = 1.0
|
||
if not _TOP_P_WARNED:
|
||
print("[PromptForge] ComfyUI's CLIP sampler top_p is incompatible with this torch build; using top_p=1.0 (temperature, top_k, repetition and seed still apply).")
|
||
_TOP_P_WARNED = True
|
||
|
||
if torch.cuda.is_available() and torch.cuda.mem_get_info()[0] < 2 * 1024 ** 3:
|
||
_free_vram(getattr(clip, "patcher", None))
|
||
|
||
def run(top_p_value, seed_value):
|
||
return clip.generate(
|
||
tokens,
|
||
do_sample=True,
|
||
max_length=max_tokens,
|
||
temperature=temperature,
|
||
top_k=top_k,
|
||
top_p=top_p_value,
|
||
min_p=0.0,
|
||
repetition_penalty=repetition_penalty,
|
||
seed=seed_value,
|
||
presence_penalty=presence_penalty,
|
||
)
|
||
|
||
def attempt(seed_value):
|
||
try:
|
||
return clip.decode(run(effective_top_p, seed_value), skip_special_tokens=True)
|
||
except AttributeError as exc:
|
||
raise ValueError("[PromptForge] this text encoder cannot generate text (CLIP/T5-style encoder). Use a Qwen / Gemma LLM encoder, or a safetensors/GGUF model.") from exc
|
||
except RuntimeError as exc:
|
||
if effective_top_p < 1.0 and "single memory location" in str(exc):
|
||
print("[PromptForge] ComfyUI's CLIP sampler hit a torch overlap error on top_p; retrying with top_p=1.0 (temperature, top_k, repetition and seed still apply).")
|
||
return clip.decode(run(1.0, seed_value), skip_special_tokens=True)
|
||
raise
|
||
|
||
return _generate_checked(attempt, seed, min_chars, auto_retry)
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Input adapters
|
||
# ------------------------------------------------------------------ #
|
||
def _clip_images(self, mode, image, image2, image3, video, frame_count):
|
||
"""Build the [1, H, W, C] tensor list ComfyUI multimodal tokenizers expect."""
|
||
if mode == "text_only":
|
||
return []
|
||
images = []
|
||
for tensor in (image, image2, image3):
|
||
if torch.is_tensor(tensor) and tensor.ndim == 4:
|
||
images.append(tensor[:1])
|
||
elif torch.is_tensor(tensor) and tensor.ndim == 3:
|
||
images.append(tensor.unsqueeze(0))
|
||
if torch.is_tensor(video) and video.ndim == 4 and video.shape[0] > 0:
|
||
steps = min(max(1, frame_count), video.shape[0])
|
||
idx = torch.linspace(0, video.shape[0] - 1, steps=steps).round().long()
|
||
images.extend(video[i:i + 1] for i in idx)
|
||
fitted = []
|
||
for t in images:
|
||
t = t.detach().float().cpu().clamp(0.0, 1.0)
|
||
size = _vision_size(t.shape[2], t.shape[1])
|
||
if size is not None:
|
||
t = F.interpolate(t.permute(0, 3, 1, 2), size=(size[1], size[0]), mode="area").permute(0, 2, 3, 1)
|
||
fitted.append(t)
|
||
return fitted
|
||
|
||
def _media_inventory(self, mode, image, image2, image3, video, frame_count):
|
||
"""Describe only the media that actually reached the node."""
|
||
if mode == "text_only":
|
||
return (
|
||
"# Provided inputs (authoritative)\n"
|
||
"No media is used in text_only mode. Write a text-only prompt. "
|
||
"Do not mention images, references, video, or frames."
|
||
)
|
||
images = sum(1 for value in (image, image2, image3) if self._has_media(value))
|
||
frames = 0
|
||
if self._has_media(video):
|
||
if torch.is_tensor(video) and video.ndim == 4:
|
||
frames = min(max(1, int(frame_count)), int(video.shape[0]))
|
||
elif isinstance(video, (list, tuple)):
|
||
frames = min(max(1, int(frame_count)), len(video))
|
||
else:
|
||
frames = 1
|
||
media = []
|
||
if images == 1:
|
||
media.append("1 image")
|
||
elif images > 1:
|
||
media.append(f"{images} images")
|
||
if frames:
|
||
media.append(f"1 video ({frames} sampled frame{'s' if frames != 1 else ''})")
|
||
if not media:
|
||
return (
|
||
"# Provided inputs (authoritative)\n"
|
||
"No image or video is connected. Write a text-only prompt. "
|
||
"Do not mention images, references, video, or frames."
|
||
)
|
||
return (
|
||
"# Provided inputs (authoritative)\n"
|
||
f"The only media connected to this node: {' and '.join(media)}. "
|
||
"Nothing else is connected: no additional images, references, video, or frames. "
|
||
"Write the prompt for exactly these inputs. Never mention, imply, reference, or invent "
|
||
"any other image, reference, video, frame, or media."
|
||
)
|
||
|
||
def _minimax_inputs(self, mode_key, mode, image, image2, image3, video, frame_count, seconds):
|
||
"""Trim the connected media to what the MiniMax mode can use and describe it with the H3 labels."""
|
||
max_images, allows_video, _ = _MINIMAX_MODES[mode_key]
|
||
pictures = [value for value in (image, image2, image3) if mode != "text_only" and self._has_media(value)]
|
||
has_video = mode != "text_only" and self._has_media(video)
|
||
dropped = []
|
||
if max_images is not None and len(pictures) > max_images:
|
||
dropped.extend(f"image {index + 1}" for index in range(max_images, len(pictures)))
|
||
pictures = pictures[:max_images]
|
||
if has_video and not allows_video:
|
||
dropped.append("video")
|
||
has_video = False
|
||
if dropped:
|
||
print(f"[PromptForge] MiniMax {mode_key.upper()} cannot use: {', '.join(dropped)}; ignoring.")
|
||
pictures += [None] * (3 - len(pictures))
|
||
video = video if has_video else None
|
||
frames = 0
|
||
if has_video:
|
||
frames = min(max(1, int(frame_count)), int(video.shape[0])) if torch.is_tensor(video) and video.ndim == 4 else 1
|
||
block = _minimax_inputs_block(mode_key, sum(value is not None for value in pictures), frames, dropped, seconds)
|
||
return pictures[0], pictures[1], pictures[2], video, block
|
||
|
||
def _has_media(self, value):
|
||
if value is None:
|
||
return False
|
||
if torch.is_tensor(value):
|
||
return value.numel() > 0
|
||
if isinstance(value, (list, tuple)):
|
||
return len(value) > 0
|
||
return True
|
||
|
||
def _mode_images(self, mode, image, image2, image3, video, frame_count):
|
||
if mode == "text_only":
|
||
return []
|
||
images = [self._pil_from_input(image), self._pil_from_input(image2), self._pil_from_input(image3)]
|
||
images = [p for p in images if p is not None]
|
||
if video is not None:
|
||
images.extend(self._frames_from_input(video, frame_count))
|
||
return [self._fit_for_vision(p) for p in images]
|
||
|
||
def _fit_for_vision(self, im):
|
||
size = _vision_size(im.width, im.height)
|
||
if size is None:
|
||
return im
|
||
from PIL import Image
|
||
return im.resize(size, Image.LANCZOS)
|
||
|
||
def _pil_from_input(self, img):
|
||
from PIL import Image
|
||
if img is None:
|
||
return None
|
||
if isinstance(img, Image.Image):
|
||
return img
|
||
if torch.is_tensor(img):
|
||
img = img.detach().float().cpu()
|
||
if img.ndim == 4:
|
||
img = img[0]
|
||
arr = img.clamp(0.0, 1.0).mul(255.0).round().to(torch.uint8).numpy()
|
||
else:
|
||
import numpy as np
|
||
arr = np.asarray(img)
|
||
if arr.ndim == 4:
|
||
arr = arr[0]
|
||
if arr.dtype != np.uint8:
|
||
arr = (np.clip(arr, 0.0, 1.0) * 255.0).round().astype(np.uint8)
|
||
if arr.ndim == 2:
|
||
arr = arr[:, :, None]
|
||
im = Image.fromarray(arr)
|
||
if im.mode != "RGB":
|
||
im = im.convert("RGB")
|
||
return im
|
||
|
||
def _pil_data_url(self, image):
|
||
buffer = io.BytesIO()
|
||
image.save(buffer, format="JPEG", quality=92)
|
||
return "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode("ascii")
|
||
|
||
def _frames_from_input(self, video, frame_count):
|
||
from PIL import Image
|
||
frames = []
|
||
if video is None:
|
||
return frames
|
||
if isinstance(video, list) and all(isinstance(f, Image.Image) for f in video):
|
||
frames = list(video)
|
||
elif torch.is_tensor(video) and video.ndim == 4:
|
||
steps = min(max(1, frame_count), video.shape[0])
|
||
idx = torch.linspace(0, video.shape[0] - 1, steps=steps).round().long()
|
||
frames = [self._pil_from_input(video[i]) for i in idx]
|
||
return [f for f in frames if f is not None][: max(1, frame_count)]
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"PROMPTFORGE_Model": PROMPTFORGE_Model,
|
||
"PROMPTFORGE_PromptEnhancer": PROMPTFORGE_PromptEnhancer,
|
||
"PROMPTFORGE_ConfigFinder": PROMPTFORGE_ConfigFinder,
|
||
"PROMPTFORGE_ImageAction": PROMPTFORGE_ImageAction,
|
||
**VFI_CLASS_MAPPINGS,
|
||
**SAVE_CLASS_MAPPINGS,
|
||
**LORA_CLASS_MAPPINGS,
|
||
**LOADER_CLASS_MAPPINGS,
|
||
**LATENT_CLASS_MAPPINGS,
|
||
**SAMPLER_CLASS_MAPPINGS,
|
||
**QUANT_CLASS_MAPPINGS,
|
||
**DECODE_CLASS_MAPPINGS,
|
||
**IMAGE_CLASS_MAPPINGS,
|
||
**AUDIO_CLASS_MAPPINGS,
|
||
**VIDEO_CLASS_MAPPINGS,
|
||
**ENDLESS_CLASS_MAPPINGS,
|
||
**H3_SETTINGS_CLASS_MAPPINGS,
|
||
**PREVIEW_CLASS_MAPPINGS,
|
||
**MOSS_CLASS_MAPPINGS,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"PROMPTFORGE_Model": "PromptForge Prompt Model",
|
||
"PROMPTFORGE_PromptEnhancer": "PromptForge Prompt Enhancer",
|
||
"PROMPTFORGE_ConfigFinder": "PromptForge Config Finder",
|
||
"PROMPTFORGE_ImageAction": "PromptForge Action",
|
||
**VFI_DISPLAY_NAMES,
|
||
**SAVE_DISPLAY_NAMES,
|
||
**LORA_DISPLAY_NAMES,
|
||
**LOADER_DISPLAY_NAMES,
|
||
**LATENT_DISPLAY_NAMES,
|
||
**SAMPLER_DISPLAY_NAMES,
|
||
**QUANT_DISPLAY_NAMES,
|
||
**DECODE_DISPLAY_NAMES,
|
||
**IMAGE_DISPLAY_NAMES,
|
||
**AUDIO_DISPLAY_NAMES,
|
||
**VIDEO_DISPLAY_NAMES,
|
||
**ENDLESS_DISPLAY_NAMES,
|
||
**H3_SETTINGS_DISPLAY_NAMES,
|
||
**PREVIEW_DISPLAY_NAMES,
|
||
**MOSS_DISPLAY_NAMES,
|
||
}
|
||
|
||
WEB_DIRECTORY = "./web"
|