304 lines
9.2 KiB
Python
304 lines
9.2 KiB
Python
"""
|
||
node/nodes.py — 根级共享工具与统一门面节点
|
||
==========================================
|
||
|
||
职责：
|
||
1. 提供可被各子包复用的通用工具函数。
|
||
2. 注册根级的 `LLS Universal Image Generator` 节点。
|
||
|
||
现有子包仍可继续从本文件导入工具函数；新增的统一节点只作为门面，
|
||
具体的 SD1.5 / SDXL / FLUX 差异全部下沉到 `lls_universal/` 后端适配器。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import contextlib
|
||
import inspect
|
||
import importlib
|
||
from typing import Any, Iterable
|
||
|
||
# ---------- 防御性导入 ----------
|
||
|
||
try:
|
||
import torch
|
||
except Exception as exc:
|
||
torch = None
|
||
_TORCH_IMPORT_ERROR = exc
|
||
else:
|
||
_TORCH_IMPORT_ERROR = None
|
||
|
||
try:
|
||
import folder_paths
|
||
except Exception as exc:
|
||
folder_paths = None
|
||
_FOLDER_PATHS_IMPORT_ERROR = exc
|
||
else:
|
||
_FOLDER_PATHS_IMPORT_ERROR = None
|
||
|
||
try:
|
||
import comfy.model_management as model_management
|
||
import comfy.utils as comfy_utils
|
||
except Exception as exc:
|
||
model_management = None
|
||
comfy_utils = None
|
||
_COMFY_IMPORT_ERROR = exc
|
||
else:
|
||
_COMFY_IMPORT_ERROR = None
|
||
|
||
try:
|
||
comfy_core_nodes = importlib.import_module("nodes")
|
||
except Exception:
|
||
comfy_core_nodes = None
|
||
|
||
try:
|
||
from .lls_universal.dispatcher import get_backend
|
||
from .lls_universal.request import (
|
||
MODEL_FAMILY_CHOICES,
|
||
TASK_MODE_CHOICES,
|
||
LLSUniversalGenerationRequest,
|
||
)
|
||
except Exception as exc:
|
||
get_backend = None
|
||
MODEL_FAMILY_CHOICES = ("SD1.5", "SDXL", "FLUX")
|
||
TASK_MODE_CHOICES = ("txt2img",)
|
||
LLSUniversalGenerationRequest = None
|
||
_UNIVERSAL_IMPORT_ERROR = exc
|
||
else:
|
||
_UNIVERSAL_IMPORT_ERROR = None
|
||
|
||
try:
|
||
from .sampling.nodes import _get_samplers, _get_schedulers
|
||
except Exception:
|
||
_get_samplers = None
|
||
_get_schedulers = None
|
||
|
||
|
||
_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",
|
||
]
|
||
_DEFAULT_SCHEDULERS = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
|
||
_MAX_RESOLUTION = int(getattr(comfy_core_nodes, "MAX_RESOLUTION", 8192))
|
||
|
||
|
||
# ---------- 共享工具函数 ----------
|
||
|
||
def _torch_inference_mode():
|
||
"""返回 torch.inference_mode() context manager； torch 不可用时返回空上下文。"""
|
||
if torch is None or not hasattr(torch, "inference_mode"):
|
||
return contextlib.nullcontext()
|
||
return torch.inference_mode()
|
||
|
||
|
||
def _require_torch() -> None:
|
||
"""断言 torch 可用，否则抛出 RuntimeError。"""
|
||
if torch is None:
|
||
raise RuntimeError(
|
||
"PyTorch is not available in the current environment. "
|
||
"Load this custom node inside ComfyUI's Python environment."
|
||
) from _TORCH_IMPORT_ERROR
|
||
|
||
|
||
def _require_comfy_modules() -> None:
|
||
"""断言 ComfyUI 核心模块均可用，否则抛出 RuntimeError。"""
|
||
if folder_paths is None:
|
||
raise RuntimeError(
|
||
"ComfyUI folder_paths could not be imported. "
|
||
"Place this plugin directory under ComfyUI/custom_nodes/ and restart ComfyUI."
|
||
) from _FOLDER_PATHS_IMPORT_ERROR
|
||
if comfy_utils is None or model_management is None:
|
||
raise RuntimeError(
|
||
"ComfyUI runtime modules could not be imported. "
|
||
"This node must run inside a ComfyUI environment."
|
||
) from _COMFY_IMPORT_ERROR
|
||
|
||
|
||
def _clamp_image(image: Any) -> Any:
|
||
"""将图像 tensor clamp 到 [0.0, 1.0]。"""
|
||
if torch is None:
|
||
return image
|
||
return torch.clamp(image, 0.0, 1.0)
|
||
|
||
|
||
def _sorted_unique(items: Iterable[str]) -> list[str]:
|
||
"""去重并排序字符串列表（保持确定性顺序）。"""
|
||
return sorted(dict.fromkeys(items))
|
||
|
||
|
||
def _is_oom_exception(exc: Exception) -> bool:
|
||
"""判断异常是否为显存不足（OOM）。"""
|
||
oom_types: list[type[BaseException]] = []
|
||
if model_management is not None:
|
||
oom_type = getattr(model_management, "OOM_EXCEPTION", None)
|
||
if isinstance(oom_type, type):
|
||
oom_types.append(oom_type)
|
||
if torch is not None:
|
||
cuda = getattr(torch, "cuda", None)
|
||
oom_type = getattr(cuda, "OutOfMemoryError", None)
|
||
if isinstance(oom_type, type):
|
||
oom_types.append(oom_type)
|
||
if any(isinstance(exc, t) for t in oom_types):
|
||
return True
|
||
return "out of memory" in str(exc).lower()
|
||
|
||
|
||
def _make_progress_bar(steps: int):
|
||
"""创建 ComfyUI ProgressBar；comfy_utils 不可用时返回 None。"""
|
||
if comfy_utils is None or not hasattr(comfy_utils, "ProgressBar"):
|
||
return None
|
||
return comfy_utils.ProgressBar(steps)
|
||
|
||
|
||
def _supports_argument(fn: Any, argument_name: str) -> bool:
|
||
"""检查函数签名中是否包含指定参数名。"""
|
||
try:
|
||
signature = inspect.signature(fn)
|
||
except (TypeError, ValueError):
|
||
return False
|
||
return argument_name in signature.parameters
|
||
|
||
|
||
def _get_checkpoint_names() -> list[str]:
|
||
"""读取 ComfyUI checkpoints 列表；脱离 ComfyUI 时返回占位提示。"""
|
||
if folder_paths is None:
|
||
return ["(ComfyUI not available)"]
|
||
try:
|
||
names = folder_paths.get_filename_list("checkpoints")
|
||
except Exception:
|
||
names = []
|
||
return names if names else ["(no checkpoints found)"]
|
||
|
||
|
||
def _get_sampler_choices() -> list[str]:
|
||
if callable(_get_samplers):
|
||
try:
|
||
names = list(_get_samplers())
|
||
if names:
|
||
return names
|
||
except Exception:
|
||
pass
|
||
return _DEFAULT_SAMPLERS
|
||
|
||
|
||
def _get_scheduler_choices() -> list[str]:
|
||
if callable(_get_schedulers):
|
||
try:
|
||
names = list(_get_schedulers())
|
||
if names:
|
||
return names
|
||
except Exception:
|
||
pass
|
||
return _DEFAULT_SCHEDULERS
|
||
|
||
|
||
# ---------- 根级统一节点 ----------
|
||
|
||
class LLSUniversalImageGenerator:
|
||
"""
|
||
统一 txt2img 入口。
|
||
|
||
节点本身不拼接任何“万能采样流程”；
|
||
它只负责收集请求并交给 dispatcher 选择对应 backend。
|
||
"""
|
||
|
||
CATEGORY = "LLS/Image"
|
||
FUNCTION = "generate"
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("image",)
|
||
DESCRIPTION = (
|
||
"Unified txt2img node that dispatches SD1.5, SDXL, and FLUX generation "
|
||
"to family-specific backend adapters."
|
||
)
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"model_family": (list(MODEL_FAMILY_CHOICES), {"default": "SD1.5"}),
|
||
"task_mode": (list(TASK_MODE_CHOICES), {"default": "txt2img"}),
|
||
"model_name": (_get_checkpoint_names(),),
|
||
"positive_prompt": ("STRING", {"default": "", "multiline": True}),
|
||
"negative_prompt": ("STRING", {"default": "", "multiline": True}),
|
||
"width": ("INT", {"default": 1024, "min": 64, "max": _MAX_RESOLUTION, "step": 8}),
|
||
"height": ("INT", {"default": 1024, "min": 64, "max": _MAX_RESOLUTION, "step": 8}),
|
||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
||
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.1}),
|
||
"seed": (
|
||
"INT",
|
||
{
|
||
"default": -1,
|
||
"min": -1,
|
||
"max": 0xFFFFFFFFFFFFFFFF,
|
||
"control_after_generate": True,
|
||
},
|
||
),
|
||
"sampler_name": (_get_sampler_choices(), {"default": "euler"}),
|
||
"scheduler": (_get_scheduler_choices(), {"default": "normal"}),
|
||
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||
}
|
||
}
|
||
|
||
def generate(
|
||
self,
|
||
model_family: str,
|
||
task_mode: str,
|
||
model_name: str,
|
||
positive_prompt: str,
|
||
negative_prompt: str,
|
||
width: int,
|
||
height: int,
|
||
steps: int,
|
||
cfg: float,
|
||
seed: int,
|
||
sampler_name: str,
|
||
scheduler: str,
|
||
denoise: float,
|
||
):
|
||
if get_backend is None or LLSUniversalGenerationRequest is None:
|
||
raise RuntimeError(
|
||
"[LLS] Universal generator modules could not be imported. "
|
||
"Check the lls_universal package for syntax or import errors."
|
||
) from _UNIVERSAL_IMPORT_ERROR
|
||
|
||
request = LLSUniversalGenerationRequest(
|
||
model_family=model_family,
|
||
task_mode=task_mode,
|
||
model_name=model_name,
|
||
positive_prompt=positive_prompt,
|
||
negative_prompt=negative_prompt,
|
||
width=width,
|
||
height=height,
|
||
steps=steps,
|
||
cfg=cfg,
|
||
seed=seed,
|
||
sampler_name=sampler_name,
|
||
scheduler=scheduler,
|
||
denoise=denoise,
|
||
).validate()
|
||
|
||
backend = get_backend(request.model_family)
|
||
image = backend.generate(request)
|
||
return (image,)
|
||
|
||
|
||
# ---------- 根级注册表 ----------
|
||
|
||
NODE_CLASS_MAPPINGS: dict[str, type] = {
|
||
"LLSUniversalImageGenerator": LLSUniversalImageGenerator,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = {
|
||
"LLSUniversalImageGenerator": "LLS Universal Image Generator",
|
||
}
|