nodes/Gin3601-LLS-node
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

重构核心模块

Browse Source
This commit is contained in:
Gin3601
2026-05-19 16:39:41 +08:00
parent fa36d0329f
commit b02dd7a1a3
20 changed files with 1176 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
conditioning/__pycache__/nodes.cpython-312.pyc
BIN
View File
Binary file not shown.
conditioning/nodes.py
+114 -60
View File
@@ -7,31 +7,38 @@ CATEGORY = "LLS/Conditioning"
"""
from __future__ import annotations
# ---------- 防御性导入 ----------
# （本节点直接调用 clip 对象方法，无需额外导入 comfy.sd）
from ..utils.model_info import is_flux_family, is_sdxl_family, parse_model_info
_CLIP_SKIP_CHOICES = [None] + list(range(-1, -25, -1))
def _encode_conditioning(clip, prompt: str):
"""
兼容新旧 ComfyUI 的文本编码入口。
def _normalize_clip_skip(value) -> int:
if value in (None, "", "None"):
return -1
try:
normalized = int(value)
except (TypeError, ValueError) as exc:
raise RuntimeError(f"[LLS] Invalid clip_skip value: {value!r}") from exc
if normalized < -24 or normalized > -1:
raise RuntimeError(f"[LLS] clip_skip out of supported range [-24, -1]: {normalized}")
return normalized
新版本优先使用 `encode_from_tokens_scheduled()`；
若运行环境较旧，仅存在 `encode_from_tokens()`，则手动包装成
ComfyUI 约定的 CONDITIONING 结构。
"""
tokens = clip.tokenize(prompt)
def _encode_tokens(clip, tokens, add_dict=None):
scheduled = getattr(clip, "encode_from_tokens_scheduled", None)
if callable(scheduled):
return scheduled(tokens)
try:
if add_dict:
return scheduled(tokens, add_dict=add_dict)
return scheduled(tokens)
except TypeError:
return scheduled(tokens)
plain_encode = getattr(clip, "encode_from_tokens", None)
if not callable(plain_encode):
raise AttributeError(
"CLIP object does not expose encode_from_tokens_scheduled() "
"or encode_from_tokens()."
"CLIP object does not expose encode_from_tokens_scheduled() or encode_from_tokens()."
)
try:
@@ -51,30 +58,65 @@ def _encode_conditioning(clip, prompt: str):
return [[cond, pooled_dict]]
def _normalize_clip_skip(value) -> int:
"""
兼容旧工作流中的 null / None，以及字符串形式的 clip_skip。
"""
if value in (None, "", "None"):
return -1
try:
normalized = int(value)
except (TypeError, ValueError) as exc:
raise RuntimeError(f"[LLS] Invalid clip_skip value: {value!r}") from exc
if normalized < -24 or normalized > -1:
raise RuntimeError(f"[LLS] clip_skip out of supported range [-24, -1]: {normalized}")
return normalized
def _encode_standard_prompt(clip, prompt: str):
tokens = clip.tokenize(prompt)
return _encode_tokens(clip, tokens)
# ---------- 节点类 ----------
def _build_sdxl_tokens(clip, prompt: str):
tokens = clip.tokenize(prompt)
local_tokens = clip.tokenize(prompt)
if "l" in local_tokens:
tokens["l"] = local_tokens["l"]
if "g" not in tokens and "g" in local_tokens:
tokens["g"] = local_tokens["g"]
if "l" in tokens and "g" in tokens and len(tokens["l"]) != len(tokens["g"]):
empty = clip.tokenize("")
while len(tokens["l"]) < len(tokens["g"]):
tokens["l"] += empty.get("l", [])
while len(tokens["g"]) < len(tokens["l"]):
tokens["g"] += empty.get("g", [])
return tokens
def _encode_sdxl_prompt(clip, prompt: str, width: int, height: int):
tokens = _build_sdxl_tokens(clip, prompt)
add_dict = {
"width": width,
"height": height,
"crop_w": 0,
"crop_h": 0,
"target_width": width,
"target_height": height,
}
return _encode_tokens(clip, tokens, add_dict=add_dict)
def _build_flux_tokens(clip, prompt: str):
tokens = clip.tokenize(prompt)
t5_tokens = clip.tokenize(prompt)
if "t5xxl" in t5_tokens:
tokens["t5xxl"] = t5_tokens["t5xxl"]
return tokens
def _encode_flux_prompt(clip, prompt: str, guidance):
tokens = _build_flux_tokens(clip, prompt)
add_dict = {}
if guidance is not None:
add_dict["guidance"] = guidance
return _encode_tokens(clip, tokens, add_dict=add_dict or None)
class LLSSimplePromptEncode:
"""
简化版提示词编码节点。
优先复用 ComfyUI 原生 clip.encode_from_tokens_scheduled()，
旧版环境下自动回退到 encode_from_tokens()，
以兼容 SD1.5 / SDXL 的常见 CLIP 编码路径。
- 默认仍兼容原有 SD 风格 CLIP 编码。
- 若连接 model_info，则按 family 自动切换到 SDXL / FLUX 编码路径。
"""
CATEGORY = "LLS/Conditioning"
@@ -82,67 +124,79 @@ class LLSSimplePromptEncode:
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "STRING")
RETURN_NAMES = ("positive", "negative", "prompt_info")
DESCRIPTION = (
"Encode positive and negative prompts into CONDITIONING tensors. "
"Prefers ComfyUI native encode_from_tokens_scheduled and falls back to "
"encode_from_tokens for older runtimes."
"Encode prompts into CONDITIONING tensors. Uses model_info when available to "
"dispatch between SD15, SDXL, and FLUX text-encoding paths."
)
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"clip": ("CLIP",),
"text_encoder": ("CLIP",),
"positive_prompt": ("STRING", {"default": "", "multiline": True}),
"negative_prompt": ("STRING", {"default": "", "multiline": True}),
"clip_skip": (_CLIP_SKIP_CHOICES, {"default": -1}),
}
},
"optional": {
"model_info": ("STRING", {"default": ""}),
},
}
def encode(self, clip, positive_prompt: str, negative_prompt: str, clip_skip):
if clip is None:
def encode(
self,
text_encoder,
positive_prompt: str,
negative_prompt: str,
clip_skip,
model_info: str | None = None,
):
if text_encoder is None:
raise RuntimeError(
"[LLS] CLIP is None. Make sure a checkpoint is loaded before connecting to this node."
"[LLS] text_encoder is None. Connect the loader output or load a compatible text encoder first."
)
info = parse_model_info(model_info)
family = info["family"]
clip_skip = _normalize_clip_skip(clip_skip)
# 应用 clip_skip（ComfyUI CLIP 对象支持 clip_layer）
if clip_skip != -1:
try:
clip = clip.clone()
clip.clip_layer(clip_skip)
text_encoder = text_encoder.clone()
text_encoder.clip_layer(clip_skip)
except Exception:
# clip_layer 不可用时静默忽略，不阻断流程
pass
# 编码 positive — 复用 ComfyUI 原生编码入口
try:
positive = _encode_conditioning(clip, positive_prompt)
if is_flux_family(family):
guidance = info.get("guidance") if info.get("guidance_embed") else None
positive = _encode_flux_prompt(text_encoder, positive_prompt, guidance)
negative = _encode_flux_prompt(text_encoder, "", guidance)
negative_note = "negative_ignored_for_flux"
elif is_sdxl_family(family):
width = int(info.get("base_width", 1024))
height = int(info.get("base_height", 1024))
positive = _encode_sdxl_prompt(text_encoder, positive_prompt, width, height)
negative = _encode_sdxl_prompt(text_encoder, negative_prompt, width, height)
negative_note = "negative_used"
else:
positive = _encode_standard_prompt(text_encoder, positive_prompt)
negative = _encode_standard_prompt(text_encoder, negative_prompt)
negative_note = "negative_used"
except Exception as exc:
raise RuntimeError(
f"[LLS] Failed to encode positive prompt: {exc}"
) from exc
# 编码 negative — 复用 ComfyUI 原生编码入口
try:
negative = _encode_conditioning(clip, negative_prompt)
except Exception as exc:
raise RuntimeError(
f"[LLS] Failed to encode negative prompt: {exc}"
) from exc
raise RuntimeError(f"[LLS] Failed to encode prompts for family {family}: {exc}") from exc
pos_preview = positive_prompt[:60].replace("\n", " ")
neg_preview = negative_prompt[:40].replace("\n", " ")
if is_flux_family(family):
neg_preview = ""
prompt_info = (
f"clip_skip={clip_skip} "
f"| pos=\"{pos_preview}...\" | neg=\"{neg_preview}...\""
f"family={family} | clip_skip={clip_skip} | "
f"pos=\"{pos_preview}...\" | neg=\"{neg_preview}...\" | {negative_note}"
)
return (positive, negative, prompt_info)
# ---------- 注册表 ----------
NODE_CLASS_MAPPINGS: dict[str, type] = {
"LLSSimplePromptEncode": LLSSimplePromptEncode,
}
image/__pycache__/nodes.cpython-312.pyc
BIN
View File
Binary file not shown.
image/nodes.py
+16 -4
View File
@@ -7,6 +7,8 @@ CATEGORY = "LLS/Image"
"""
from __future__ import annotations
from ..utils.model_info import get_latent_spec, parse_model_info
# ---------- 节点类 ----------
@@ -32,10 +34,13 @@ class LLSSimpleVAEDecode:
"required": {
"samples": ("LATENT",),
"vae": ("VAE",),
}
},
"optional": {
"model_info": ("STRING", {"default": ""}),
},
}
def decode(self, samples, vae):
def decode(self, samples, vae, model_info: str | None = None):
if vae is None:
raise RuntimeError(
"[LLS] VAE is None. Make sure a checkpoint is loaded and VAE is connected."
@@ -65,9 +70,16 @@ class LLSSimpleVAEDecode:
if len(getattr(image, "shape", ())) == 5:
image = image.reshape(-1, image.shape[-3], image.shape[-2], image.shape[-1])
h, w = latent_tensor.shape[2] * 8, latent_tensor.shape[3] * 8
info = parse_model_info(model_info)
latent_spec = get_latent_spec(info)
downscale_ratio = int(samples.get("downscale_ratio_spacial", latent_spec["downscale_ratio"]))
h, w = latent_tensor.shape[2] * downscale_ratio, latent_tensor.shape[3] * downscale_ratio
batch = latent_tensor.shape[0]
decode_info = f"decoded={w}x{h} | batch={batch} | latent_shape={list(latent_tensor.shape)}"
decode_info = (
f"family={info['family']} | decoded={w}x{h} | batch={batch} "
f"| downscale={downscale_ratio} | latent_shape={list(latent_tensor.shape)}"
)
return (image, decode_info)
latent/__pycache__/nodes.cpython-312.pyc
BIN
View File
Binary file not shown.
latent/nodes.py
+36 -33
View File
@@ -7,8 +7,6 @@ CATEGORY = "LLS/Latent"
"""
from __future__ import annotations
# ---------- 防御性导入 ----------
try:
import torch
except Exception as exc:
@@ -22,16 +20,14 @@ try:
except Exception:
model_management = None
from ..utils.model_info import get_latent_spec, parse_model_info
# ---------- 工具函数 ----------
_SIZE_PRESETS = [
"Custom",
# SD1.5 推荐尺寸
"512x512",
"512x768",
"768x512",
# SDXL 推荐尺寸
"1024x1024",
"832x1216",
"1216x832",
@@ -42,17 +38,16 @@ _SIZE_PRESETS = [
]
def _round_to_8(value: int) -> int:
"""将整数修正为最近的 8 的倍数。"""
return max(8, round(value / 8) * 8)
def _round_to_multiple(value: int, multiple: int) -> int:
return max(multiple, ((value + multiple - 1) // multiple) * multiple)
# ---------- 节点类 ----------
class LLSSimpleEmptyLatent:
"""
生成空白 Latent，作为 txt2img 推理的起点。
支持常用分辨率预设，自动修正宿高为 8 的倍数。
默认保持 SD 风格 4 通道 / 8x 下采样；
若提供 model_info，则可自动切换到 FLUX 的 128 通道 / 16x 下采样。
"""
CATEGORY = "LLS/Latent"
@@ -61,9 +56,7 @@ class LLSSimpleEmptyLatent:
RETURN_NAMES = ("latent", "width", "height", "latent_info")
DESCRIPTION = (
"Create an empty latent image as the starting point for txt2img. "
"Select a preset or enter custom width/height. "
"Width and height are automatically rounded to multiples of 8. "
"SD1.5 recommended: 512x512; SDXL recommended: 1024x1024 or 896x1152."
"When model_info is connected, latent channels and downscale ratio adapt to the model family."
)
@classmethod
@@ -74,30 +67,39 @@ class LLSSimpleEmptyLatent:
"width": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}),
"height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
}
},
"optional": {
"model_info": ("STRING", {"default": ""}),
},
}
def create_empty_latent(self, size_preset: str, width: int, height: int, batch_size: int):
def create_empty_latent(
self,
size_preset: str,
width: int,
height: int,
batch_size: int,
model_info: str | None = None,
):
if torch is None:
raise RuntimeError(
"[LLS] PyTorch is not available. "
"Make sure this node runs inside a ComfyUI environment."
"[LLS] PyTorch is not available. Make sure this node runs inside a ComfyUI environment."
) from _TORCH_ERR
info = parse_model_info(model_info)
latent_spec = get_latent_spec(info)
latent_channels = latent_spec["latent_channels"]
downscale_ratio = latent_spec["downscale_ratio"]
notes = []
# 如果选择了预设，覆盖 width / height
if size_preset != "Custom":
try:
w_str, h_str = size_preset.split("x")
width = int(w_str)
height = int(h_str)
width, height = [int(part) for part in size_preset.split("x", 1)]
except Exception:
notes.append(f"preset parse failed, using custom {width}x{height}")
# 修正为 8 的倍数
corrected_w = _round_to_8(width)
corrected_h = _round_to_8(height)
corrected_w = _round_to_multiple(width, downscale_ratio)
corrected_h = _round_to_multiple(height, downscale_ratio)
if corrected_w != width or corrected_h != height:
notes.append(f"size corrected {width}x{height} -> {corrected_w}x{corrected_h}")
width, height = corrected_w, corrected_h
@@ -114,25 +116,26 @@ class LLSSimpleEmptyLatent:
except Exception:
dtype = torch.float32
# 生成空白 Latent（与 ComfyUI EmptyLatentImage 格式一致）
# Latent 尺寸是图像尺寸的 1/8
latent = torch.zeros(
[batch_size, 4, height // 8, width // 8],
[batch_size, latent_channels, height // downscale_ratio, width // downscale_ratio],
device=device,
dtype=dtype,
)
note_str = " | ".join(notes) if notes else "ok"
latent_info = (
f"size={width}x{height} | batch={batch_size} "
f"| preset={size_preset} | {note_str}"
f"family={info['family']} | size={width}x{height} | batch={batch_size} | "
f"channels={latent_channels} | downscale={downscale_ratio} | preset={size_preset} | {note_str}"
)
return ({"samples": latent, "downscale_ratio_spacial": 8}, width, height, latent_info)
return (
{"samples": latent, "downscale_ratio_spacial": downscale_ratio},
width,
height,
latent_info,
)
# ---------- 注册表 ----------
NODE_CLASS_MAPPINGS: dict[str, type] = {
"LLSSimpleEmptyLatent": LLSSimpleEmptyLatent,
}
lls_universal/__pycache__/backend_base.cpython-312.pyc
BIN
View File
Binary file not shown.
lls_universal/__pycache__/backend_flux.cpython-312.pyc
BIN
View File
Binary file not shown.
lls_universal/__pycache__/backend_sdxl.cpython-312.pyc
BIN
View File
Binary file not shown.
model_loader/__pycache__/nodes.cpython-312.pyc
BIN
View File
Binary file not shown.
model_loader/nodes.py
+343 -71
View File
@@ -7,7 +7,7 @@ CATEGORY = "LLS/Model Loader"
"""
from __future__ import annotations
# ---------- 防御性导入 ----------
import importlib
try:
import folder_paths
@@ -25,118 +25,390 @@ except Exception as exc:
else:
_COMFY_SD_ERR = None
# 注：comfy.utils 在本模块未使用，无需导入
try:
comfy_core_nodes = importlib.import_module("nodes")
except Exception as exc:
comfy_core_nodes = None
_CORE_NODES_ERR = exc
else:
_CORE_NODES_ERR = None
from ..utils.model_info import (
MODEL_FAMILY_CHOICES,
build_model_info,
canonicalize_family,
family_requires_guidance,
get_latent_spec,
infer_family_from_name,
is_flux_family,
is_sdxl_family,
)
# ---------- 工具函数 ----------
AUTO_PLACEHOLDER = "(auto)"
NO_MODELS_PLACEHOLDER = "(no models found)"
NO_VAE_PLACEHOLDER = "(no vae found)"
NO_TEXT_ENCODER_PLACEHOLDER = "(no text encoders found)"
LOAD_MODE_CHOICES = ["simple", "advanced"]
VAE_SOURCE_CHOICES = ["auto", "embedded", "external", "none"]
TEXT_ENCODER_SOURCE_CHOICES = ["auto", "embedded", "external", "manual"]
def _get_checkpoint_names() -> list[str]:
def _get_filename_list(category: str) -> list[str]:
if folder_paths is None:
return ["(ComfyUI not available)"]
return []
try:
names = folder_paths.get_filename_list("checkpoints")
return names if names else ["(no checkpoints found)"]
return list(folder_paths.get_filename_list(category))
except Exception:
return ["(no checkpoints found)"]
return []
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"
def _with_placeholder(names: list[str], placeholder: str) -> list[str]:
cleaned = [name for name in names if name]
if not cleaned:
return [AUTO_PLACEHOLDER, placeholder]
return [AUTO_PLACEHOLDER] + cleaned
# ---------- 节点类 ----------
def _get_model_name_choices() -> list[str]:
checkpoints = _get_filename_list("checkpoints")
diffusion_models = [f"diffusion_models/{name}" for name in _get_filename_list("diffusion_models")]
names = list(dict.fromkeys(checkpoints + diffusion_models))
return names if names else [NO_MODELS_PLACEHOLDER]
def _get_vae_choices() -> list[str]:
if comfy_core_nodes is not None:
vae_loader = getattr(comfy_core_nodes, "VAELoader", None)
if vae_loader is not None and hasattr(vae_loader, "vae_list"):
try:
names = list(vae_loader.vae_list(vae_loader))
return _with_placeholder(names, NO_VAE_PLACEHOLDER)
except Exception:
pass
return _with_placeholder(_get_filename_list("vae"), NO_VAE_PLACEHOLDER)
def _get_text_encoder_choices() -> list[str]:
return _with_placeholder(_get_filename_list("text_encoders"), NO_TEXT_ENCODER_PLACEHOLDER)
def _normalize_choice(value: str | None, placeholder: str) -> str | None:
if value in (None, "", AUTO_PLACEHOLDER, placeholder):
return None
return value
def _get_full_path(category: str, name: str) -> str | None:
if folder_paths is None:
return None
try:
return folder_paths.get_full_path_or_raise(category, name)
except AttributeError:
return folder_paths.get_full_path(category, name)
except FileNotFoundError:
return None
def _resolve_model_path(model_name: str, family: str) -> tuple[str, str]:
if model_name == NO_MODELS_PLACEHOLDER:
raise RuntimeError(
"[LLS] No models were found in ComfyUI models/checkpoints or models/diffusion_models."
)
if model_name.startswith("diffusion_models/"):
if not is_flux_family(family):
raise RuntimeError(
f"[LLS] {family} expects a checkpoint file. Select a checkpoint from models/checkpoints/."
)
short_name = model_name.split("/", 1)[1]
model_path = _get_full_path("diffusion_models", short_name)
if model_path:
return ("diffusion_models", model_path)
checkpoint_path = _get_full_path("checkpoints", model_name)
if checkpoint_path:
return ("checkpoints", checkpoint_path)
if is_flux_family(family):
diffusion_path = _get_full_path("diffusion_models", model_name)
if diffusion_path:
return ("diffusion_models", diffusion_path)
raise RuntimeError(
f"[LLS] Model '{model_name}' was not found in the expected ComfyUI model directories."
)
def _load_checkpoint_bundle(model_path: str):
try:
return comfy_sd.load_checkpoint_guess_config(
model_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 '{model_path}': {exc}") from exc
def _load_model(model_source: str, model_path: str):
if model_source == "checkpoints":
loaded = _load_checkpoint_bundle(model_path)
return loaded[0], loaded[1], loaded[2]
try:
model = comfy_sd.load_diffusion_model(model_path)
except Exception as exc:
raise RuntimeError(f"[LLS] Failed to load diffusion model '{model_path}': {exc}") from exc
return model, None, None
def _load_external_vae(vae_name: str):
if comfy_core_nodes is None:
raise RuntimeError(
"[LLS] ComfyUI core VAELoader is unavailable, so external VAE loading cannot be used."
) from _CORE_NODES_ERR
vae_loader_cls = getattr(comfy_core_nodes, "VAELoader", None)
if vae_loader_cls is None:
raise RuntimeError("[LLS] ComfyUI core VAELoader class was not found.")
return vae_loader_cls().load_vae(vae_name)[0]
def _load_external_text_encoder(family: str, name1: str | None, name2: str | None):
if comfy_core_nodes is None:
raise RuntimeError(
"[LLS] ComfyUI core CLIP loaders are unavailable, so external text encoder loading cannot be used."
) from _CORE_NODES_ERR
if family == "SD15":
if not name1:
raise RuntimeError(
"[LLS] SD15 external text encoder loading requires external_text_encoder_1."
)
clip_loader_cls = getattr(comfy_core_nodes, "CLIPLoader", None)
if clip_loader_cls is None:
raise RuntimeError("[LLS] ComfyUI core CLIPLoader class was not found.")
return clip_loader_cls().load_clip(name1, type="stable_diffusion")[0]
if is_sdxl_family(family):
if not name1 or not name2:
raise RuntimeError(
"[LLS] SDXL external text encoder loading requires clip_l and clip_g via "
"external_text_encoder_1/external_text_encoder_2."
)
dual_loader_cls = getattr(comfy_core_nodes, "DualCLIPLoader", None)
if dual_loader_cls is None:
raise RuntimeError("[LLS] ComfyUI core DualCLIPLoader class was not found.")
return dual_loader_cls().load_clip(name1, name2, type="sdxl")[0]
if not name1 or not name2:
raise RuntimeError(
"[LLS] FLUX text encoder loading requires clip_l and t5xxl via "
"external_text_encoder_1/external_text_encoder_2."
)
dual_loader_cls = getattr(comfy_core_nodes, "DualCLIPLoader", None)
if dual_loader_cls is None:
raise RuntimeError("[LLS] ComfyUI core DualCLIPLoader class was not found.")
return dual_loader_cls().load_clip(name1, name2, type="flux")[0]
def _resolve_text_encoder(
family: str,
source: str,
embedded_text_encoder,
external_name_1: str | None,
external_name_2: str | None,
) -> tuple[object | None, str]:
if source == "manual":
return None, "manual"
if source == "embedded":
if embedded_text_encoder is None:
raise RuntimeError(
f"[LLS] {family} was set to use embedded text encoders, but the selected model did not provide them."
)
return embedded_text_encoder, "embedded"
if source == "external":
return _load_external_text_encoder(family, external_name_1, external_name_2), "external"
if embedded_text_encoder is not None:
return embedded_text_encoder, "embedded"
if external_name_1:
return _load_external_text_encoder(family, external_name_1, external_name_2), "external"
if is_flux_family(family):
raise RuntimeError(
"[LLS] FLUX models need text encoders. Provide clip_l and t5xxl via "
"external_text_encoder_1/external_text_encoder_2, or use a checkpoint that embeds them."
)
raise RuntimeError(
f"[LLS] {family} did not expose an embedded text encoder. "
"Switch text_encoder_source to 'external' and choose the required encoder files."
)
def _resolve_vae(
family: str,
source: str,
embedded_vae,
external_vae_name: str | None,
) -> tuple[object | None, str]:
if source == "none":
if is_flux_family(family):
raise RuntimeError(
"[LLS] FLUX generation requires a VAE/AE. Provide external_vae_name or use a model with an embedded VAE."
)
return None, "none"
if source == "embedded":
if embedded_vae is None:
raise RuntimeError(
f"[LLS] {family} was set to use an embedded VAE, but the selected model did not provide one."
)
return embedded_vae, "embedded"
if source == "external":
if not external_vae_name:
raise RuntimeError("[LLS] External VAE loading requires external_vae_name.")
return _load_external_vae(external_vae_name), "external"
if embedded_vae is not None:
return embedded_vae, "embedded"
if external_vae_name:
return _load_external_vae(external_vae_name), "external"
if is_flux_family(family):
raise RuntimeError(
"[LLS] FLUX models usually need an external AE/VAE. Provide external_vae_name or use a model with an embedded VAE."
)
return None, "none"
def _text_encoder_type_for_family(family: str) -> str:
if family == "SD15":
return "clip"
if is_sdxl_family(family):
return "sdxl_dual_clip"
return "flux_clip_l_t5xxl"
def _required_text_encoders_for_family(family: str) -> list[str]:
if family == "SD15":
return ["clip"]
if is_sdxl_family(family):
return ["clip_l", "clip_g"]
return ["clip_l", "t5xxl"]
def _required_vae_for_family(family: str) -> str:
return "required" if is_flux_family(family) else "optional"
class LLSSimpleCheckpointLoader:
"""
简化版 Checkpoint 加载器，支持 SD1.5 / SDXL 基础流程。
自动识别模型家族，输出 MODEL / CLIP / VAE 以及模型信息字符串。
内部复用 ComfyUI 原生 load_checkpoint_guess_config，已自动适配两种架构。
统一外观的模型加载节点。
对外仍然输出 MODEL / CLIP / VAE / STRING 以兼容旧工作流；
但第二个端口实际语义已经升级为 text_encoder，资源分派策略则由
model_family + source 选项控制。
"""
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"]
RETURN_NAMES = ("model", "text_encoder", "vae", "model_info")
DESCRIPTION = (
"Load SD15, SDXL, SDXL Turbo, and FLUX families with family-specific resource "
"dispatch. The CLIP-typed output now represents the resolved text encoder bundle."
)
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"ckpt_name": (_get_checkpoint_names(),),
"model_family": (cls._FAMILY_CHOICES, {"default": "Auto"}),
"ckpt_name": (_get_model_name_choices(),),
"model_family": (MODEL_FAMILY_CHOICES, {"default": "Auto"}),
"load_mode": (LOAD_MODE_CHOICES, {"default": "simple"}),
"vae_source": (VAE_SOURCE_CHOICES, {"default": "auto"}),
"text_encoder_source": (TEXT_ENCODER_SOURCE_CHOICES, {"default": "auto"}),
"external_vae_name": (_get_vae_choices(), {"default": AUTO_PLACEHOLDER}),
"external_text_encoder_1": (_get_text_encoder_choices(), {"default": AUTO_PLACEHOLDER}),
"external_text_encoder_2": (_get_text_encoder_choices(), {"default": AUTO_PLACEHOLDER}),
}
}
def load_checkpoint(self, ckpt_name: str, model_family: str):
def load_checkpoint(
self,
ckpt_name: str,
model_family: str,
load_mode: str,
vae_source: str,
text_encoder_source: str,
external_vae_name: str,
external_text_encoder_1: str,
external_text_encoder_2: str,
):
if folder_paths is None:
raise RuntimeError(
"[LLS] folder_paths is not available. "
"Make sure this node runs inside a ComfyUI environment."
"[LLS] folder_paths is not available. Make sure this node runs inside ComfyUI."
) 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."
"[LLS] comfy.sd is not available. Make sure this node runs inside ComfyUI."
) 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
family = canonicalize_family(model_family)
if model_family == "Auto":
family = infer_family_from_name(ckpt_name, family)
# 加载 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]
model_source, model_path = _resolve_model_path(ckpt_name, family)
model, embedded_text_encoder, embedded_vae = _load_model(model_source, model_path)
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.")
raise RuntimeError(f"[LLS] The selected model '{ckpt_name}' did not produce a valid MODEL.")
detected_family = _detect_family(ckpt_name, model_family)
external_vae_name = _normalize_choice(external_vae_name, NO_VAE_PLACEHOLDER)
external_text_encoder_1 = _normalize_choice(external_text_encoder_1, NO_TEXT_ENCODER_PLACEHOLDER)
external_text_encoder_2 = _normalize_choice(external_text_encoder_2, NO_TEXT_ENCODER_PLACEHOLDER)
model_info = (
f"ckpt={ckpt_name} | family={detected_family}"
text_encoder, resolved_text_encoder_source = _resolve_text_encoder(
family=family,
source=text_encoder_source,
embedded_text_encoder=embedded_text_encoder,
external_name_1=external_text_encoder_1,
external_name_2=external_text_encoder_2,
)
vae, resolved_vae_source = _resolve_vae(
family=family,
source=vae_source,
embedded_vae=embedded_vae,
external_vae_name=external_vae_name,
)
return (model, clip, vae, model_info)
latent_spec = get_latent_spec({"family": family})
model_info = build_model_info(
family=family,
model_name=ckpt_name,
model_source=model_source,
load_mode=load_mode,
text_encoder_source=resolved_text_encoder_source,
vae_source=resolved_vae_source,
text_encoder_type=_text_encoder_type_for_family(family),
has_embedded_vae=embedded_vae is not None,
has_embedded_text_encoder=embedded_text_encoder is not None,
required_text_encoders=_required_text_encoders_for_family(family),
required_vae=_required_vae_for_family(family),
is_turbo=(family == "SDXL_TURBO"),
guidance_embed=family_requires_guidance(family),
latent_channels=latent_spec["latent_channels"],
downscale_ratio=latent_spec["downscale_ratio"],
)
return (model, text_encoder, vae, model_info)
# ---------- 注册表 ----------
NODE_CLASS_MAPPINGS: dict[str, type] = {
"LLSSimpleCheckpointLoader": LLSSimpleCheckpointLoader,
sampling/__pycache__/nodes.cpython-312.pyc
BIN
View File
Binary file not shown.
sampling/nodes.py
+15 -13
View File
@@ -9,6 +9,8 @@ from __future__ import annotations
import random
from ..utils.model_info import get_sampling_preset, parse_model_info
# ---------- 防御性导入 ----------
try:
@@ -53,12 +55,6 @@ except Exception:
_QUALITY_PRESETS = ["Manual", "Fast", "Balanced", "High Quality"]
_PRESET_PARAMS = {
"Fast": {"steps": 12, "cfg": 6.5},
"Balanced": {"steps": 20, "cfg": 7.0},
"High Quality": {"steps": 30, "cfg": 7.5},
}
_DEFAULT_SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral",
"lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral",
"dpmpp_sde", "dpmpp_2m", "dpmpp_2m_sde", "ddim", "uni_pc"]
@@ -208,7 +204,10 @@ class LLSSimpleKSampler:
"sampler_name": (_get_samplers(), {"default": "euler_ancestral"}),
"scheduler": (_get_schedulers(), {"default": "karras"}),
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
}
},
"optional": {
"model_info": ("STRING", {"default": ""}),
},
}
def sample(
@@ -224,6 +223,7 @@ class LLSSimpleKSampler:
sampler_name: str,
scheduler: str,
denoise: float,
model_info: str | None = None,
):
if comfy_sample is None:
raise RuntimeError(
@@ -236,11 +236,13 @@ class LLSSimpleKSampler:
"Make sure this node runs inside a ComfyUI environment."
) from _COMFY_SAMPLERS_ERR
# 应用 quality_preset（Manual 不覆盖用户输入）
if quality_preset in _PRESET_PARAMS:
preset = _PRESET_PARAMS[quality_preset]
steps = preset["steps"]
cfg = preset["cfg"]
info = parse_model_info(model_info)
family = info["family"]
preset = get_sampling_preset(info, quality_preset)
if preset is not None:
steps = int(preset["steps"])
cfg = float(preset["cfg"])
# seed = -1 时随机生成
actual_seed = seed
@@ -269,7 +271,7 @@ class LLSSimpleKSampler:
sample_info = (
f"seed={actual_seed} | steps={steps} | cfg={cfg} "
f"| sampler={sampler_name} | scheduler={scheduler} "
f"| denoise={denoise} | preset={quality_preset}"
f"| denoise={denoise} | preset={quality_preset} | family={family}"
)
return (result_latent, sample_info)
tests/__pycache__/test_conditioning_compat.cpython-312.pyc
BIN
View File
Binary file not shown.
tests/__pycache__/test_loader_prompt_refactor.cpython-312.pyc
BIN
View File
Binary file not shown.
tests/__pycache__/test_upscale_switcher.cpython-312.pyc
BIN
View File
Binary file not shown.
tests/test_loader_prompt_refactor.py
+418
View File
@@ -0,0 +1,418 @@
import importlib.util
import json
import pathlib
import sys
import types
import unittest
from unittest import mock
ROOT = pathlib.Path(__file__).resolve().parents[1]
def load_plugin_package():
spec = importlib.util.spec_from_file_location(
"lls_node_test_refactor",
ROOT / "__init__.py",
submodule_search_locations=[str(ROOT)],
)
module = importlib.util.module_from_spec(spec)
sys.modules["lls_node_test_refactor"] = module
spec.loader.exec_module(module)
return module
class FolderPathsStub:
def __init__(self):
self._files = {
"checkpoints": [
"sd15.safetensors",
"sdxl_turbo.safetensors",
"flux-dev-checkpoint.safetensors",
],
"diffusion_models": [
"flux-schnell.safetensors",
],
"vae": [
"ae.safetensors",
],
"text_encoders": [
"clip_l.safetensors",
"clip_g.safetensors",
"t5xxl_fp16.safetensors",
],
}
def get_filename_list(self, category):
return list(self._files.get(category, []))
def get_full_path_or_raise(self, category, name):
if name in self._files.get(category, []):
return f"/fake/{category}/{name}"
raise FileNotFoundError(name)
def get_full_path(self, category, name):
if name in self._files.get(category, []):
return f"/fake/{category}/{name}"
return None
def get_folder_paths(self, category):
return [f"/fake/{category}"]
class ComfySDStub:
class CLIPType:
STABLE_DIFFUSION = "stable_diffusion"
FLUX = "flux"
def load_checkpoint_guess_config(
self,
ckpt_path,
output_vae=True,
output_clip=True,
embedding_directory=None,
):
if ckpt_path.endswith("sd15.safetensors"):
return ("MODEL::SD15", "CLIP::SD15", "VAE::SD15")
if ckpt_path.endswith("sdxl_turbo.safetensors"):
return ("MODEL::SDXL_TURBO", None, None)
if ckpt_path.endswith("flux-dev-checkpoint.safetensors"):
return ("MODEL::FLUX_DEV", None, None)
raise AssertionError(f"unexpected checkpoint path: {ckpt_path}")
def load_diffusion_model(self, unet_path, model_options=None, disable_dynamic=False):
return f"MODEL::{pathlib.Path(unet_path).name}"
def load_clip(
self,
ckpt_paths,
embedding_directory=None,
clip_type=None,
model_options=None,
disable_dynamic=False,
):
joined = "+".join(pathlib.Path(path).name for path in ckpt_paths)
return f"CLIP::{clip_type}::{joined}"
class CoreNodesStub:
class VAELoader:
@staticmethod
def vae_list(_self):
return ["ae.safetensors"]
def load_vae(self, vae_name):
return (f"VAE::{vae_name}",)
class CLIPLoader:
def load_clip(self, clip_name, type="stable_diffusion", device="default"):
return (f"CLIP::{type}::{clip_name}",)
class DualCLIPLoader:
def load_clip(self, clip_name1, clip_name2, type, device="default"):
return (f"CLIP::{type}::{clip_name1}+{clip_name2}",)
class RecordingClip:
def __init__(self, calls=None):
self.calls = [] if calls is None else calls
self.layer = None
def clone(self):
cloned = RecordingClip(self.calls)
cloned.layer = self.layer
return cloned
def clip_layer(self, layer):
self.layer = layer
def tokenize(self, text):
return {
"text": text,
"g": [f"g::{text}"],
"l": [f"l::{text}"],
"t5xxl": [f"t5::{text}"],
}
def encode_from_tokens_scheduled(self, tokens, add_dict=None):
self.calls.append(
{
"tokens": dict(tokens),
"add_dict": dict(add_dict or {}),
"layer": self.layer,
}
)
return [[f"cond::{tokens.get('text', '')}", {"pooled_output": dict(add_dict or {})}]]
class FakeVAE:
def decode(self, latent_tensor):
batch, _channels, height, width = latent_tensor.shape
return FakeTensor((batch, height, width, 3))
class FakeTensor:
def __init__(self, shape):
self.shape = tuple(shape)
self.dtype = "float32"
self.layout = "strided"
self.device = "cpu"
self.is_nested = False
def size(self):
return self.shape
def reshape(self, *shape):
return FakeTensor(shape)
class FakeTorch:
float32 = "float32"
@staticmethod
def zeros(shape, device=None, dtype=None, layout=None):
return FakeTensor(shape)
class TestLoaderPromptRefactor(unittest.TestCase):
def test_loader_schema_exposes_new_family_and_source_controls(self):
plugin = load_plugin_package()
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleCheckpointLoader"]
required = node_cls.INPUT_TYPES()["required"]
self.assertEqual(
tuple(required.keys()),
(
"ckpt_name",
"model_family",
"load_mode",
"vae_source",
"text_encoder_source",
"external_vae_name",
"external_text_encoder_1",
"external_text_encoder_2",
),
)
family_choices = required["model_family"][0]
self.assertIn("SD15", family_choices)
self.assertIn("SDXL_TURBO", family_choices)
self.assertIn("FLUX_SCHNELL", family_choices)
self.assertIn("FLUX_DEV", family_choices)
def test_loader_returns_json_model_info_for_sd15_embedded_resources(self):
load_plugin_package()
from lls_node_test_refactor.model_loader import nodes as loader_nodes
with mock.patch.object(loader_nodes, "folder_paths", FolderPathsStub()), mock.patch.object(
loader_nodes,
"comfy_sd",
ComfySDStub(),
), mock.patch.object(loader_nodes, "comfy_core_nodes", CoreNodesStub()):
node = loader_nodes.LLSSimpleCheckpointLoader()
model, text_encoder, vae, model_info = node.load_checkpoint(
"sd15.safetensors",
"SD15",
"simple",
"auto",
"auto",
"(auto)",
"(auto)",
"(auto)",
)
info = json.loads(model_info)
self.assertEqual(model, "MODEL::SD15")
self.assertEqual(text_encoder, "CLIP::SD15")
self.assertEqual(vae, "VAE::SD15")
self.assertEqual(info["family"], "SD15")
self.assertEqual(info["text_encoder_type"], "clip")
self.assertTrue(info["has_embedded_vae"])
self.assertEqual(info["required_text_encoders"], ["clip"])
self.assertEqual(info["required_vae"], "optional")
self.assertFalse(info["is_turbo"])
def test_loader_builds_flux_resources_from_external_models(self):
load_plugin_package()
from lls_node_test_refactor.model_loader import nodes as loader_nodes
with mock.patch.object(loader_nodes, "folder_paths", FolderPathsStub()), mock.patch.object(
loader_nodes,
"comfy_sd",
ComfySDStub(),
), mock.patch.object(loader_nodes, "comfy_core_nodes", CoreNodesStub()):
node = loader_nodes.LLSSimpleCheckpointLoader()
model, text_encoder, vae, model_info = node.load_checkpoint(
"diffusion_models/flux-schnell.safetensors",
"FLUX_SCHNELL",
"advanced",
"external",
"external",
"ae.safetensors",
"clip_l.safetensors",
"t5xxl_fp16.safetensors",
)
info = json.loads(model_info)
self.assertEqual(model, "MODEL::flux-schnell.safetensors")
self.assertEqual(text_encoder, "CLIP::flux::clip_l.safetensors+t5xxl_fp16.safetensors")
self.assertEqual(vae, "VAE::ae.safetensors")
self.assertEqual(info["family"], "FLUX_SCHNELL")
self.assertEqual(info["text_encoder_type"], "flux_clip_l_t5xxl")
self.assertEqual(info["required_text_encoders"], ["clip_l", "t5xxl"])
self.assertEqual(info["required_vae"], "required")
def test_loader_raises_clear_error_when_flux_resources_are_missing(self):
load_plugin_package()
from lls_node_test_refactor.model_loader import nodes as loader_nodes
with mock.patch.object(loader_nodes, "folder_paths", FolderPathsStub()), mock.patch.object(
loader_nodes,
"comfy_sd",
ComfySDStub(),
), mock.patch.object(loader_nodes, "comfy_core_nodes", CoreNodesStub()):
node = loader_nodes.LLSSimpleCheckpointLoader()
with self.assertRaisesRegex(RuntimeError, "clip_l.*t5xxl"):
node.load_checkpoint(
"diffusion_models/flux-schnell.safetensors",
"FLUX_DEV",
"simple",
"external",
"external",
"ae.safetensors",
"(auto)",
"(auto)",
)
def test_prompt_encode_dispatches_sdxl_using_model_info_dimensions(self):
plugin = load_plugin_package()
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimplePromptEncode"]
node = node_cls()
clip = RecordingClip()
model_info = json.dumps({"family": "SDXL", "base_width": 1216, "base_height": 832})
positive, negative, prompt_info = node.encode(
clip,
"a castle",
"low quality",
-1,
model_info=model_info,
)
self.assertEqual(len(clip.calls), 2)
self.assertIn("g", clip.calls[0]["tokens"])
self.assertIn("l", clip.calls[0]["tokens"])
self.assertEqual(
clip.calls[0]["add_dict"],
{
"width": 1216,
"height": 832,
"crop_w": 0,
"crop_h": 0,
"target_width": 1216,
"target_height": 832,
},
)
self.assertEqual(positive[0][0], "cond::a castle")
self.assertEqual(negative[0][0], "cond::low quality")
self.assertIn("family=SDXL", prompt_info)
def test_prompt_encode_dispatches_flux_and_neutralizes_negative_prompt(self):
plugin = load_plugin_package()
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimplePromptEncode"]
node = node_cls()
clip = RecordingClip()
model_info = json.dumps({"family": "FLUX_DEV", "guidance": 3.5, "guidance_embed": True})
positive, negative, prompt_info = node.encode(
clip,
"a robot in a forest",
"ugly, blurry",
-1,
model_info=model_info,
)
self.assertEqual(len(clip.calls), 2)
self.assertIn("t5xxl", clip.calls[0]["tokens"])
self.assertEqual(clip.calls[0]["add_dict"]["guidance"], 3.5)
self.assertEqual(clip.calls[1]["tokens"]["text"], "")
self.assertEqual(positive[0][0], "cond::a robot in a forest")
self.assertEqual(negative[0][0], "cond::")
self.assertIn("negative_ignored_for_flux", prompt_info)
def test_empty_latent_uses_flux_shape_when_model_info_requests_it(self):
load_plugin_package()
from lls_node_test_refactor.latent import nodes as latent_nodes
node = latent_nodes.LLSSimpleEmptyLatent()
model_info = json.dumps({"family": "FLUX_SCHNELL", "latent_channels": 128, "downscale_ratio": 16})
with mock.patch.object(latent_nodes, "torch", FakeTorch()):
latent, width, height, latent_info = node.create_empty_latent(
"Custom",
1000,
1000,
2,
model_info=model_info,
)
self.assertEqual((width, height), (1008, 1008))
self.assertEqual(tuple(latent["samples"].shape), (2, 128, 63, 63))
self.assertEqual(latent["downscale_ratio_spacial"], 16)
self.assertIn("family=FLUX_SCHNELL", latent_info)
def test_ksampler_uses_turbo_preset_defaults_when_model_info_requests_it(self):
load_plugin_package()
from lls_node_test_refactor.sampling import nodes as sampling_nodes
recorded = {}
def fake_common_ksampler(**kwargs):
recorded.update(kwargs)
return kwargs["latent"]
with mock.patch.object(sampling_nodes, "comfy_sample", object()), mock.patch.object(
sampling_nodes,
"comfy_samplers",
object(),
), mock.patch.object(
sampling_nodes,
"_common_ksampler",
side_effect=fake_common_ksampler,
):
node = sampling_nodes.LLSSimpleKSampler()
latent, sample_info = node.sample(
model="model",
positive="positive",
negative="negative",
latent_image={"samples": FakeTensor((1, 4, 64, 64))},
quality_preset="Balanced",
seed=123,
steps=30,
cfg=7.0,
sampler_name="euler",
scheduler="normal",
denoise=1.0,
model_info=json.dumps({"family": "SDXL_TURBO", "is_turbo": True}),
)
self.assertEqual(recorded["steps"], 6)
self.assertEqual(recorded["cfg"], 1.5)
self.assertEqual(latent["samples"].shape, (1, 4, 64, 64))
self.assertIn("family=SDXL_TURBO", sample_info)
def test_vae_decode_uses_model_info_downscale_ratio_for_flux(self):
plugin = load_plugin_package()
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleVAEDecode"]
node = node_cls()
image, decode_info = node.decode(
{"samples": FakeTensor((1, 128, 64, 64))},
FakeVAE(),
model_info=json.dumps({"family": "FLUX_DEV", "downscale_ratio": 16}),
)
self.assertEqual(tuple(image.shape), (1, 64, 64, 3))
self.assertIn("decoded=1024x1024", decode_info)
if __name__ == "__main__":
unittest.main()
upscale/__pycache__/nodes.cpython-312.pyc
BIN
View File
Binary file not shown.
utils/__pycache__/model_info.cpython-312.pyc
BIN
View File
Binary file not shown.
utils/model_info.py
+234
View File
@@ -0,0 +1,234 @@
from __future__ import annotations
import json
import re
from typing import Any
AUTO_FAMILY = "Auto"
LEGACY_FAMILY_ALIASES: dict[str, str] = {
"SD1.5": "SD15",
"SD15": "SD15",
"SDXL": "SDXL",
"SDXL_TURBO": "SDXL_TURBO",
"FLUX": "FLUX_DEV",
"FLUX_SCHNELL": "FLUX_SCHNELL",
"FLUX_DEV": "FLUX_DEV",
}
MODEL_FAMILY_CHOICES: list[str] = [
AUTO_FAMILY,
"SD15",
"SD1.5",
"SDXL",
"SDXL_TURBO",
"FLUX_SCHNELL",
"FLUX_DEV",
]
TEXT_ENCODER_TYPES: dict[str, str] = {
"SD15": "clip",
"SDXL": "sdxl_dual_clip",
"SDXL_TURBO": "sdxl_dual_clip",
"FLUX_SCHNELL": "flux_clip_l_t5xxl",
"FLUX_DEV": "flux_clip_l_t5xxl",
}
REQUIRED_TEXT_ENCODERS: dict[str, list[str]] = {
"SD15": ["clip"],
"SDXL": ["clip_l", "clip_g"],
"SDXL_TURBO": ["clip_l", "clip_g"],
"FLUX_SCHNELL": ["clip_l", "t5xxl"],
"FLUX_DEV": ["clip_l", "t5xxl"],
}
REQUIRED_VAE: dict[str, str] = {
"SD15": "optional",
"SDXL": "optional",
"SDXL_TURBO": "optional",
"FLUX_SCHNELL": "required",
"FLUX_DEV": "required",
}
LATENT_SPECS: dict[str, dict[str, int]] = {
"SD15": {"latent_channels": 4, "downscale_ratio": 8},
"SDXL": {"latent_channels": 4, "downscale_ratio": 8},
"SDXL_TURBO": {"latent_channels": 4, "downscale_ratio": 8},
"FLUX_SCHNELL": {"latent_channels": 128, "downscale_ratio": 16},
"FLUX_DEV": {"latent_channels": 128, "downscale_ratio": 16},
}
BASE_RESOLUTIONS: dict[str, tuple[int, int]] = {
"SD15": (512, 512),
"SDXL": (1024, 1024),
"SDXL_TURBO": (1024, 1024),
"FLUX_SCHNELL": (1024, 1024),
"FLUX_DEV": (1024, 1024),
}
GUIDANCE_EMBED: dict[str, bool] = {
"SD15": False,
"SDXL": False,
"SDXL_TURBO": False,
"FLUX_SCHNELL": False,
"FLUX_DEV": True,
}
DEFAULT_GUIDANCE: dict[str, float | None] = {
"SD15": None,
"SDXL": None,
"SDXL_TURBO": None,
"FLUX_SCHNELL": None,
"FLUX_DEV": 3.5,
}
SAMPLING_PRESETS: dict[str, dict[str, dict[str, float | int]]] = {
"SD15": {
"Fast": {"steps": 12, "cfg": 6.5},
"Balanced": {"steps": 20, "cfg": 7.0},
"High Quality": {"steps": 30, "cfg": 7.5},
},
"SDXL": {
"Fast": {"steps": 16, "cfg": 5.5},
"Balanced": {"steps": 24, "cfg": 6.0},
"High Quality": {"steps": 36, "cfg": 6.5},
},
"SDXL_TURBO": {
"Fast": {"steps": 4, "cfg": 1.0},
"Balanced": {"steps": 6, "cfg": 1.5},
"High Quality": {"steps": 8, "cfg": 2.0},
},
"FLUX_SCHNELL": {
"Fast": {"steps": 4, "cfg": 1.0},
"Balanced": {"steps": 6, "cfg": 1.0},
"High Quality": {"steps": 8, "cfg": 1.0},
},
"FLUX_DEV": {
"Fast": {"steps": 18, "cfg": 1.0},
"Balanced": {"steps": 28, "cfg": 1.0},
"High Quality": {"steps": 36, "cfg": 1.0},
},
}
_LEGACY_MODEL_INFO_PATTERN = re.compile(r"(\w+)=([^|]+)")
def canonicalize_family(family: str | None) -> str:
if family is None:
return "SD15"
normalized = str(family).strip()
if normalized == AUTO_FAMILY:
return "SD15"
return LEGACY_FAMILY_ALIASES.get(normalized, normalized or "SD15")
def infer_family_from_name(model_name: str | None, fallback: str = "SD15") -> str:
if not model_name:
return canonicalize_family(fallback)
name = str(model_name).lower()
if "flux" in name:
if "schnell" in name:
return "FLUX_SCHNELL"
return "FLUX_DEV"
if "turbo" in name and "xl" in name:
return "SDXL_TURBO"
if any(token in name for token in ("sdxl", "_xl", "-xl", "xl_")):
return "SDXL"
return canonicalize_family(fallback)
def is_flux_family(family: str | None) -> bool:
return canonicalize_family(family) in {"FLUX_SCHNELL", "FLUX_DEV"}
def is_sdxl_family(family: str | None) -> bool:
return canonicalize_family(family) in {"SDXL", "SDXL_TURBO"}
def family_requires_guidance(family: str | None) -> bool:
return bool(GUIDANCE_EMBED.get(canonicalize_family(family), False))
def get_latent_spec(model_info: dict[str, Any] | str | None) -> dict[str, int]:
info = parse_model_info(model_info)
family = info["family"]
return {
"latent_channels": int(info.get("latent_channels", LATENT_SPECS[family]["latent_channels"])),
"downscale_ratio": int(info.get("downscale_ratio", LATENT_SPECS[family]["downscale_ratio"])),
}
def get_sampling_preset(model_info: dict[str, Any] | str | None, quality_preset: str) -> dict[str, float | int] | None:
if quality_preset == "Manual":
return None
info = parse_model_info(model_info)
family = info["family"]
family_presets = SAMPLING_PRESETS.get(family, SAMPLING_PRESETS["SD15"])
return family_presets.get(quality_preset)
def parse_model_info(model_info: dict[str, Any] | str | None) -> dict[str, Any]:
raw: dict[str, Any] = {}
if isinstance(model_info, dict):
raw = dict(model_info)
elif isinstance(model_info, str):
stripped = model_info.strip()
if stripped:
try:
parsed = json.loads(stripped)
except json.JSONDecodeError:
raw = {
key: value.strip()
for key, value in _LEGACY_MODEL_INFO_PATTERN.findall(stripped)
}
else:
if isinstance(parsed, dict):
raw = dict(parsed)
family = canonicalize_family(raw.get("family") or infer_family_from_name(raw.get("model_name") or raw.get("ckpt_name") or raw.get("ckpt"), "SD15"))
width, height = BASE_RESOLUTIONS.get(family, BASE_RESOLUTIONS["SD15"])
latent_spec = LATENT_SPECS.get(family, LATENT_SPECS["SD15"])
required_text_encoders = raw.get("required_text_encoders", REQUIRED_TEXT_ENCODERS.get(family, ["clip"]))
if isinstance(required_text_encoders, str):
required_text_encoders = [required_text_encoders]
info: dict[str, Any] = {
"family": family,
"text_encoder_type": raw.get("text_encoder_type", TEXT_ENCODER_TYPES.get(family, "clip")),
"has_embedded_vae": bool(raw.get("has_embedded_vae", False)),
"required_text_encoders": list(required_text_encoders),
"required_vae": raw.get("required_vae", REQUIRED_VAE.get(family, "optional")),
"is_turbo": bool(raw.get("is_turbo", family == "SDXL_TURBO")),
"guidance_embed": bool(raw.get("guidance_embed", GUIDANCE_EMBED.get(family, False))),
"guidance": raw.get("guidance", DEFAULT_GUIDANCE.get(family)),
"latent_channels": int(raw.get("latent_channels", latent_spec["latent_channels"])),
"downscale_ratio": int(raw.get("downscale_ratio", latent_spec["downscale_ratio"])),
"base_width": int(raw.get("base_width", width)),
"base_height": int(raw.get("base_height", height)),
}
for key in (
"model_name",
"model_source",
"text_encoder_source",
"vae_source",
"load_mode",
"has_embedded_text_encoder",
):
if key in raw:
info[key] = raw[key]
return info
def build_model_info(**kwargs: Any) -> str:
info = parse_model_info(kwargs)
for key, value in kwargs.items():
if value is not None:
info[key] = value
info["family"] = canonicalize_family(info.get("family"))
info.setdefault("text_encoder_type", TEXT_ENCODER_TYPES.get(info["family"], "clip"))
return json.dumps(info, ensure_ascii=True, sort_keys=True)