148 lines
4.9 KiB
Python
148 lines
4.9 KiB
Python
"""
|
||
LLS / ModelLoader
|
||
=================
|
||
功能域：模型加载与管理（对应功能分类总览第 1 节）
|
||
|
||
CATEGORY = "LLS/Model Loader"
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
# ---------- 防御性导入 ----------
|
||
|
||
try:
|
||
import folder_paths
|
||
except Exception as exc:
|
||
folder_paths = None
|
||
_FOLDER_PATHS_ERR = exc
|
||
else:
|
||
_FOLDER_PATHS_ERR = None
|
||
|
||
try:
|
||
import comfy.sd as comfy_sd
|
||
except Exception as exc:
|
||
comfy_sd = None
|
||
_COMFY_SD_ERR = exc
|
||
else:
|
||
_COMFY_SD_ERR = None
|
||
|
||
# 注：comfy.utils 在本模块未使用，无需导入
|
||
|
||
|
||
# ---------- 工具函数 ----------
|
||
|
||
def _get_checkpoint_names() -> list[str]:
|
||
if folder_paths is None:
|
||
return ["(ComfyUI not available)"]
|
||
try:
|
||
names = folder_paths.get_filename_list("checkpoints")
|
||
return names if names else ["(no checkpoints found)"]
|
||
except Exception:
|
||
return ["(no checkpoints found)"]
|
||
|
||
|
||
def _detect_family(ckpt_name: str, model_family: str) -> str:
|
||
"""根据 model_family 参数或文件名推断模型家族。"""
|
||
if model_family in ("SD1.5", "SDXL"):
|
||
return model_family
|
||
# Auto：通过文件名关键词简单判断
|
||
name_lower = ckpt_name.lower()
|
||
if any(kw in name_lower for kw in ("sdxl", "_xl", "-xl", "xl_")):
|
||
return "SDXL"
|
||
return "SD1.5"
|
||
|
||
|
||
# ---------- 节点类 ----------
|
||
|
||
class LLSSimpleCheckpointLoader:
|
||
"""
|
||
简化版 Checkpoint 加载器，支持 SD1.5 / SDXL 基础流程。
|
||
自动识别模型家族，输出 MODEL / CLIP / VAE 以及模型信息字符串。
|
||
内部复用 ComfyUI 原生 load_checkpoint_guess_config，已自动适配两种架构。
|
||
"""
|
||
|
||
CATEGORY = "LLS/Model Loader"
|
||
FUNCTION = "load_checkpoint"
|
||
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "STRING")
|
||
RETURN_NAMES = ("model", "clip", "vae", "model_info")
|
||
DESCRIPTION = "Load a SD1.5 or SDXL checkpoint. Outputs MODEL, CLIP, VAE and a model_info string."
|
||
|
||
_FAMILY_CHOICES = ["Auto", "SD1.5", "SDXL"]
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"ckpt_name": (_get_checkpoint_names(),),
|
||
"model_family": (cls._FAMILY_CHOICES, {"default": "Auto"}),
|
||
}
|
||
}
|
||
|
||
def load_checkpoint(self, ckpt_name: str, model_family: str):
|
||
if folder_paths is None:
|
||
raise RuntimeError(
|
||
"[LLS] folder_paths is not available. "
|
||
"Make sure this node runs inside a ComfyUI environment."
|
||
) from _FOLDER_PATHS_ERR
|
||
if comfy_sd is None:
|
||
raise RuntimeError(
|
||
"[LLS] comfy.sd is not available. "
|
||
"Make sure this node runs inside a ComfyUI environment."
|
||
) from _COMFY_SD_ERR
|
||
|
||
# 解析模型路径（使用 ComfyUI 原生方法，自动抛出文件不存在异常）
|
||
try:
|
||
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
|
||
except (AttributeError, TypeError):
|
||
# get_full_path_or_raise 在旧版 ComfyUI 中可能不存在，回退到 get_full_path
|
||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||
if not ckpt_path:
|
||
raise RuntimeError(
|
||
f"[LLS] Checkpoint '{ckpt_name}' not found in ComfyUI checkpoints directories. "
|
||
f"Please add the file to models/checkpoints/ and restart ComfyUI."
|
||
)
|
||
except FileNotFoundError as exc:
|
||
raise RuntimeError(
|
||
f"[LLS] Checkpoint '{ckpt_name}' not found: {exc}"
|
||
) from exc
|
||
|
||
# 加载 checkpoint（复用 ComfyUI 原生加载逻辑，自动适配 SD1.5 / SDXL）
|
||
try:
|
||
out = comfy_sd.load_checkpoint_guess_config(
|
||
ckpt_path,
|
||
output_vae=True,
|
||
output_clip=True,
|
||
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
||
)
|
||
except Exception as exc:
|
||
raise RuntimeError(
|
||
f"[LLS] Failed to load checkpoint '{ckpt_name}': {exc}"
|
||
) from exc
|
||
|
||
model, clip, vae = out[0], out[1], out[2]
|
||
|
||
if model is None:
|
||
raise RuntimeError(f"[LLS] Checkpoint '{ckpt_name}' did not produce a valid MODEL.")
|
||
if clip is None:
|
||
raise RuntimeError(f"[LLS] Checkpoint '{ckpt_name}' did not produce a valid CLIP.")
|
||
if vae is None:
|
||
raise RuntimeError(f"[LLS] Checkpoint '{ckpt_name}' did not produce a valid VAE.")
|
||
|
||
detected_family = _detect_family(ckpt_name, model_family)
|
||
|
||
model_info = (
|
||
f"ckpt={ckpt_name} | family={detected_family}"
|
||
)
|
||
|
||
return (model, clip, vae, model_info)
|
||
|
||
|
||
# ---------- 注册表 ----------
|
||
|
||
NODE_CLASS_MAPPINGS: dict[str, type] = {
|
||
"LLSSimpleCheckpointLoader": LLSSimpleCheckpointLoader,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = {
|
||
"LLSSimpleCheckpointLoader": "LLS Simple Checkpoint Loader",
|
||
}
|