208 lines
9.7 KiB
Python
208 lines
9.7 KiB
Python
"""PromptForge Loader: one node for checkpoint / diffusion model / GGUF + text encoders + VAEs.
|
|
|
|
It calls ComfyUI's own loaders (GGUF files go through gguf_native, the GGUF loading code of ComfyUI-GGUF), so new
|
|
model types and loader fixes in ComfyUI apply here automatically. Its own jobs are choosing the text encoder type
|
|
from the model family, picking baked-in components, and choosing weight settings for this machine.
|
|
|
|
File choices carry a source tag: "[ckpt] x", "[unet] x", "[gguf] x" for models and "[te] x", "[gguf] x" for
|
|
text encoders. web/promptforge_loader.js narrows the lists by family; the node accepts any tagged name.
|
|
"""
|
|
import os
|
|
import re
|
|
|
|
import torch
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import folder_paths
|
|
import nodes
|
|
|
|
from . import autotune, families, gguf_native
|
|
|
|
AUTO = "auto"
|
|
BAKED = "(baked-in)"
|
|
NONE = "(none)"
|
|
CLIP_SLOTS = 4
|
|
TAG = re.compile(r"^\[(ckpt|unet|gguf|te)\] (.+)$")
|
|
# (tag, folder key) per kind of file; the gguf folders are made by gguf_native
|
|
MODEL_SOURCES = (("ckpt", "checkpoints"), ("unet", "diffusion_models"), ("gguf", "unet_gguf"))
|
|
CLIP_SOURCES = (("te", "text_encoders"), ("gguf", "clip_gguf"))
|
|
WEIGHT_DTYPES = (AUTO, "default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2", "bf16", "fp16", "fp32")
|
|
# weight_dtype -> comfy model_options
|
|
DTYPE_OPTIONS = {
|
|
"bf16": {"dtype": torch.bfloat16},
|
|
"fp16": {"dtype": torch.float16},
|
|
"fp32": {"dtype": torch.float32},
|
|
"fp8_e4m3fn": {"dtype": torch.float8_e4m3fn},
|
|
"fp8_e4m3fn_fast": {"dtype": torch.float8_e4m3fn, "fp8_optimizations": True},
|
|
"fp8_e5m2": {"dtype": torch.float8_e5m2},
|
|
}
|
|
|
|
|
|
def _tagged(sources):
|
|
names = []
|
|
for tag, folder in sources:
|
|
if folder in folder_paths.folder_names_and_paths:
|
|
names.extend(f"[{tag}] {name}" for name in folder_paths.get_filename_list(folder))
|
|
return names
|
|
|
|
|
|
def _resolve(choice, sources):
|
|
"""(tag, folder, name, full path) for a tagged choice."""
|
|
match = TAG.match(choice)
|
|
folders = dict(sources)
|
|
if not match or match.group(1) not in folders:
|
|
raise ValueError(f"[PromptForge] '{choice}' is not a file choice of this loader.")
|
|
tag, name = match.groups()
|
|
folder = folders[tag]
|
|
return tag, folder, name, folder_paths.get_full_path_or_raise(folder, name)
|
|
|
|
|
|
def _vae_names():
|
|
return nodes.VAELoader.vae_list(nodes.VAELoader)
|
|
|
|
|
|
_model_families = {} # full path -> ((mtime, size), family key)
|
|
|
|
|
|
def _model_family_key(name, folder):
|
|
path = folder_paths.get_full_path(folder, name)
|
|
stat = os.stat(path)
|
|
signature = (stat.st_mtime_ns, stat.st_size)
|
|
cached = _model_families.get(path)
|
|
if cached is None or cached[0] != signature:
|
|
family = families.family_of_model(name, folder, path)
|
|
cached = _model_families[path] = (signature, family.key if family else "")
|
|
return cached[1]
|
|
|
|
|
|
def catalog():
|
|
"""Everything the UI needs: files with their family, and what each family expects."""
|
|
models = []
|
|
for tag, folder in MODEL_SOURCES:
|
|
if folder not in folder_paths.folder_names_and_paths:
|
|
continue
|
|
for name in folder_paths.get_filename_list(folder):
|
|
models.append({"value": f"[{tag}] {name}", "family": _model_family_key(name, folder)})
|
|
return {
|
|
"models": models,
|
|
"families": [{"key": f.key, "label": f.label, "clips": list(f.clips), "vaes": list(f.vaes), "kind": f.kind} for f in families.FAMILIES],
|
|
}
|
|
|
|
|
|
@PromptServer.instance.routes.get("/promptforge/catalog")
|
|
async def _catalog_route(request):
|
|
return web.json_response(catalog())
|
|
|
|
|
|
class PROMPTFORGE_Loader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
models = _tagged(MODEL_SOURCES) or ["(no models found)"]
|
|
clips = [BAKED, NONE] + _tagged(CLIP_SOURCES)
|
|
vaes = [BAKED, NONE] + _vae_names()
|
|
return {
|
|
"required": {
|
|
"model": (models, {"tooltip": "[ckpt] checkpoint, [unet] diffusion model, [gguf] GGUF model."}),
|
|
"family": ([AUTO] + [family.label for family in families.FAMILIES], {"default": AUTO, "tooltip": "auto detects it from the file. Sets the text encoder type and the model list filter."}),
|
|
"weight_dtype": (list(WEIGHT_DTYPES), {"default": AUTO, "tooltip": "auto picks for this machine using optimize_for."}),
|
|
"optimize_for": (["balanced", "quality", "speed"], {"default": "balanced", "tooltip": "Used by weight_dtype auto. quality keeps stored precision; balanced uses fp8 storage only when the model does not fit in VRAM; speed also uses fp8 compute when a quick benchmark shows it is faster on this GPU."}),
|
|
"clip_1": (clips, {"default": BAKED, "tooltip": "Text encoder. (baked-in) uses the one inside the checkpoint."}),
|
|
"clip_2": (clips, {"default": NONE}),
|
|
"clip_3": (clips, {"default": NONE}),
|
|
"clip_4": (clips, {"default": NONE}),
|
|
"clip_type": ([AUTO] + [member.name.lower() for member in comfy.sd.CLIPType], {"default": AUTO, "tooltip": "auto uses the family's type."}),
|
|
"clip_device": (["default", "cpu"], {"default": "default"}),
|
|
"vae": (vaes, {"default": BAKED, "tooltip": "Main VAE (image or video)."}),
|
|
"audio_vae": (vaes, {"default": NONE}),
|
|
"extra_vae": (vaes, {"default": NONE, "tooltip": "Optional second VAE, for example an upscaling one to switch to later."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "VAE", "VAE", "STRING", "PROMPTFORGE_PIPE")
|
|
RETURN_NAMES = ("MODEL", "CLIP", "VAE", "AUDIO_VAE", "EXTRA_VAE", "INFO", "PIPE")
|
|
FUNCTION = "load"
|
|
CATEGORY = "🔮 PromptForge/Load"
|
|
DESCRIPTION = "Loads a model, its text encoders and VAEs in one node, choosing the text encoder type and weight settings automatically."
|
|
|
|
def load(self, model, family, weight_dtype, optimize_for, clip_1, clip_2, clip_3, clip_4, clip_type, clip_device, vae, audio_vae, extra_vae):
|
|
tag, folder, name, path = _resolve(model, MODEL_SOURCES)
|
|
chosen = next((f for f in families.FAMILIES if f.label == family), None)
|
|
if chosen is None:
|
|
chosen = families.family_of_model(name, folder, path)
|
|
info = [f"model: {name} ({chosen.label if chosen else 'unknown family'})"]
|
|
|
|
model_options = {}
|
|
if tag == "gguf":
|
|
if weight_dtype not in (AUTO, "default"):
|
|
info.append("weight_dtype is ignored for GGUF models")
|
|
else:
|
|
dtype, reason = self._weight_dtype(weight_dtype, optimize_for, path)
|
|
model_options.update(DTYPE_OPTIONS.get(dtype, {}))
|
|
info.append(f"weights: {dtype} ({reason})")
|
|
|
|
clip_names = [choice for choice in (clip_1, clip_2, clip_3, clip_4) if choice not in (BAKED, NONE)]
|
|
want_baked_clip = clip_1 == BAKED and not clip_names
|
|
want_baked_vae = vae == BAKED
|
|
|
|
loaded_clip = loaded_vae = None
|
|
if tag == "ckpt":
|
|
unet, loaded_clip, loaded_vae, _ = comfy.sd.load_checkpoint_guess_config(
|
|
path, output_vae=want_baked_vae, output_clip=want_baked_clip,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"), model_options=model_options)
|
|
elif tag == "unet":
|
|
unet = comfy.sd.load_diffusion_model(path, model_options=model_options)
|
|
else:
|
|
unet = self._load_gguf(name)
|
|
|
|
if clip_names:
|
|
loaded_clip = self._load_clip(clip_names, clip_type, chosen, clip_device)
|
|
info.append("clip: " + " + ".join(TAG.match(choice).group(2) for choice in clip_names))
|
|
elif loaded_clip is not None:
|
|
info.append("clip: baked-in")
|
|
elif chosen is not None and chosen.clips and clip_1 != NONE:
|
|
info.append(f"clip: none loaded; {chosen.label} needs {' + '.join(chosen.clips)}")
|
|
|
|
vaes = []
|
|
for slot, choice in (("vae", vae), ("audio_vae", audio_vae), ("extra_vae", extra_vae)):
|
|
if slot == "vae" and choice == BAKED:
|
|
vaes.append(loaded_vae)
|
|
elif choice in (BAKED, NONE):
|
|
vaes.append(None)
|
|
else:
|
|
vaes.append(nodes.VAELoader().load_vae(choice)[0])
|
|
info.extend(f"{label}: {'loaded' if loaded is not None else 'none'}" for label, loaded in zip(("vae", "audio_vae", "extra_vae"), vaes))
|
|
pipe = {"family": chosen.key if chosen else "", "model": name, "folder": folder}
|
|
return (unet, loaded_clip, *vaes, "\n".join(info), pipe)
|
|
|
|
@staticmethod
|
|
def _weight_dtype(weight_dtype, optimize_for, path):
|
|
if weight_dtype != AUTO:
|
|
return weight_dtype, "chosen"
|
|
header, _ = families.read_header(path) if path.endswith(".safetensors") else ({}, {})
|
|
stored = {dtype for _, dtype in header.values()}
|
|
return autotune.choose_weight_dtype(stored, os.path.getsize(path), optimize_for)
|
|
|
|
def _load_gguf(self, name):
|
|
return gguf_native.load_unet(folder_paths.get_full_path_or_raise("unet_gguf", name))
|
|
|
|
def _load_clip(self, choices, clip_type, chosen, device):
|
|
resolved = [_resolve(choice, CLIP_SOURCES) for choice in choices]
|
|
paths = [entry[3] for entry in resolved]
|
|
if clip_type == AUTO:
|
|
type_ = families.clip_type(chosen)
|
|
else:
|
|
type_ = getattr(comfy.sd.CLIPType, clip_type.upper())
|
|
if any(entry[0] == "gguf" for entry in resolved):
|
|
return gguf_native.load_clip(paths, type_)
|
|
options = {}
|
|
if device == "cpu":
|
|
options["load_device"] = options["offload_device"] = torch.device("cpu")
|
|
return comfy.sd.load_clip(ckpt_paths=paths, embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=type_, model_options=options)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"PROMPTFORGE_Loader": PROMPTFORGE_Loader}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"PROMPTFORGE_Loader": "PromptForge Loader"}
|