diff --git a/conditioning/__pycache__/nodes.cpython-312.pyc b/conditioning/__pycache__/nodes.cpython-312.pyc index 2de01a7..54d6a43 100644 Binary files a/conditioning/__pycache__/nodes.cpython-312.pyc and b/conditioning/__pycache__/nodes.cpython-312.pyc differ diff --git a/conditioning/nodes.py b/conditioning/nodes.py index 60c2717..4dff4ee 100644 --- a/conditioning/nodes.py +++ b/conditioning/nodes.py @@ -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, } diff --git a/image/__pycache__/nodes.cpython-312.pyc b/image/__pycache__/nodes.cpython-312.pyc index 4fa9c6f..5aa8a51 100644 Binary files a/image/__pycache__/nodes.cpython-312.pyc and b/image/__pycache__/nodes.cpython-312.pyc differ diff --git a/image/nodes.py b/image/nodes.py index 4ffc5f0..157c4e3 100644 --- a/image/nodes.py +++ b/image/nodes.py @@ -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) diff --git a/latent/__pycache__/nodes.cpython-312.pyc b/latent/__pycache__/nodes.cpython-312.pyc index cf5ea91..254c0b1 100644 Binary files a/latent/__pycache__/nodes.cpython-312.pyc and b/latent/__pycache__/nodes.cpython-312.pyc differ diff --git a/latent/nodes.py b/latent/nodes.py index c8ac511..bae5bf0 100644 --- a/latent/nodes.py +++ b/latent/nodes.py @@ -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, } diff --git a/lls_universal/__pycache__/backend_base.cpython-312.pyc b/lls_universal/__pycache__/backend_base.cpython-312.pyc index 387c47a..9621284 100644 Binary files a/lls_universal/__pycache__/backend_base.cpython-312.pyc and b/lls_universal/__pycache__/backend_base.cpython-312.pyc differ diff --git a/lls_universal/__pycache__/backend_flux.cpython-312.pyc b/lls_universal/__pycache__/backend_flux.cpython-312.pyc index 6118449..60e25ba 100644 Binary files a/lls_universal/__pycache__/backend_flux.cpython-312.pyc and b/lls_universal/__pycache__/backend_flux.cpython-312.pyc differ diff --git a/lls_universal/__pycache__/backend_sdxl.cpython-312.pyc b/lls_universal/__pycache__/backend_sdxl.cpython-312.pyc index dcfd890..7cf1acc 100644 Binary files a/lls_universal/__pycache__/backend_sdxl.cpython-312.pyc and b/lls_universal/__pycache__/backend_sdxl.cpython-312.pyc differ diff --git a/model_loader/__pycache__/nodes.cpython-312.pyc b/model_loader/__pycache__/nodes.cpython-312.pyc index 6a50fe8..90ab0d1 100644 Binary files a/model_loader/__pycache__/nodes.cpython-312.pyc and b/model_loader/__pycache__/nodes.cpython-312.pyc differ diff --git a/model_loader/nodes.py b/model_loader/nodes.py index cb4b81f..23eb8c3 100644 --- a/model_loader/nodes.py +++ b/model_loader/nodes.py @@ -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, diff --git a/sampling/__pycache__/nodes.cpython-312.pyc b/sampling/__pycache__/nodes.cpython-312.pyc index 47101ea..d2cfe58 100644 Binary files a/sampling/__pycache__/nodes.cpython-312.pyc and b/sampling/__pycache__/nodes.cpython-312.pyc differ diff --git a/sampling/nodes.py b/sampling/nodes.py index 2924c59..c55264b 100644 --- a/sampling/nodes.py +++ b/sampling/nodes.py @@ -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) diff --git a/tests/__pycache__/test_conditioning_compat.cpython-312.pyc b/tests/__pycache__/test_conditioning_compat.cpython-312.pyc new file mode 100644 index 0000000..81e6e50 Binary files /dev/null and b/tests/__pycache__/test_conditioning_compat.cpython-312.pyc differ diff --git a/tests/__pycache__/test_loader_prompt_refactor.cpython-312.pyc b/tests/__pycache__/test_loader_prompt_refactor.cpython-312.pyc new file mode 100644 index 0000000..31f6e85 Binary files /dev/null and b/tests/__pycache__/test_loader_prompt_refactor.cpython-312.pyc differ diff --git a/tests/__pycache__/test_upscale_switcher.cpython-312.pyc b/tests/__pycache__/test_upscale_switcher.cpython-312.pyc new file mode 100644 index 0000000..75829f1 Binary files /dev/null and b/tests/__pycache__/test_upscale_switcher.cpython-312.pyc differ diff --git a/tests/test_loader_prompt_refactor.py b/tests/test_loader_prompt_refactor.py new file mode 100644 index 0000000..7c40000 --- /dev/null +++ b/tests/test_loader_prompt_refactor.py @@ -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() diff --git a/upscale/__pycache__/nodes.cpython-312.pyc b/upscale/__pycache__/nodes.cpython-312.pyc index 6df1478..e210813 100644 Binary files a/upscale/__pycache__/nodes.cpython-312.pyc and b/upscale/__pycache__/nodes.cpython-312.pyc differ diff --git a/utils/__pycache__/model_info.cpython-312.pyc b/utils/__pycache__/model_info.cpython-312.pyc new file mode 100644 index 0000000..00bd536 Binary files /dev/null and b/utils/__pycache__/model_info.cpython-312.pyc differ diff --git a/utils/model_info.py b/utils/model_info.py new file mode 100644 index 0000000..b1fae1f --- /dev/null +++ b/utils/model_info.py @@ -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)