542 lines
19 KiB
Python
542 lines
19 KiB
Python
"""
|
|
LLS / Image
|
|
===========
|
|
功能域：图像处理与后处理（对应功能分类总览第 5 节）
|
|
|
|
CATEGORY = "LLS/Image"
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
import importlib
|
|
|
|
from ..mask.mask_utils import mask_to_image
|
|
from ..utils.model_info import (
|
|
MODEL_FAMILY_CHOICES,
|
|
canonicalize_family,
|
|
get_family_defaults,
|
|
get_latent_spec,
|
|
info_to_json,
|
|
parse_jsonish_info,
|
|
resolve_model_family,
|
|
resolve_vae_name,
|
|
)
|
|
|
|
|
|
try:
|
|
comfy_core_nodes = importlib.import_module("nodes")
|
|
except Exception:
|
|
comfy_core_nodes = None
|
|
|
|
try:
|
|
import comfy.utils as comfy_utils
|
|
except Exception as exc:
|
|
comfy_utils = None
|
|
_COMFY_UTILS_ERR = exc
|
|
else:
|
|
_COMFY_UTILS_ERR = None
|
|
|
|
|
|
_RESIZE_MODES = ["keep_aspect", "crop_center", "stretch", "none"]
|
|
_SIZE_SOURCES = ["input_image", "custom", "model_recommended"]
|
|
_SAVE_OUTPUT_MODES = ["save", "preview_only"]
|
|
|
|
|
|
def _round_to_multiple(value: int, multiple: int) -> int:
|
|
return max(multiple, ((int(value) + multiple - 1) // multiple) * multiple)
|
|
|
|
|
|
def _first_non_empty(*values):
|
|
for value in values:
|
|
if value not in (None, ""):
|
|
return value
|
|
return ""
|
|
|
|
|
|
def _get_image_dimensions(image) -> tuple[int, int, int]:
|
|
shape = tuple(getattr(image, "shape", ()))
|
|
if len(shape) < 4:
|
|
raise RuntimeError(
|
|
"[LLS] IMAGE input must be a tensor shaped like [batch, height, width, channels]."
|
|
)
|
|
return int(shape[0]), int(shape[2]), int(shape[1])
|
|
|
|
|
|
def _resolve_requested_size(
|
|
size_source: str,
|
|
width: int,
|
|
height: int,
|
|
source_width: int,
|
|
source_height: int,
|
|
model_info: dict,
|
|
) -> tuple[int, int]:
|
|
if size_source == "input_image":
|
|
return source_width, source_height
|
|
if size_source == "custom":
|
|
return int(width), int(height)
|
|
if size_source == "model_recommended":
|
|
return int(model_info["default_width"]), int(model_info["default_height"])
|
|
raise RuntimeError(f"[LLS] Unsupported size_source '{size_source}'.")
|
|
|
|
|
|
def _fit_inside_box(
|
|
source_width: int,
|
|
source_height: int,
|
|
target_width: int,
|
|
target_height: int,
|
|
multiple: int,
|
|
) -> tuple[int, int]:
|
|
if source_width <= 0 or source_height <= 0:
|
|
raise RuntimeError("[LLS] IMAGE input has invalid dimensions.")
|
|
|
|
scale = min(target_width / source_width, target_height / source_height)
|
|
fitted_width = max(multiple, int(round(source_width * scale)))
|
|
fitted_height = max(multiple, int(round(source_height * scale)))
|
|
fitted_width = min(target_width, _round_to_multiple(fitted_width, multiple))
|
|
fitted_height = min(target_height, _round_to_multiple(fitted_height, multiple))
|
|
return fitted_width, fitted_height
|
|
|
|
|
|
def _resize_image_to(image, target_width: int, target_height: int, resize_mode: str):
|
|
_batch_size, source_width, source_height = _get_image_dimensions(image)
|
|
if source_width == target_width and source_height == target_height:
|
|
return image
|
|
|
|
if comfy_utils is None:
|
|
raise RuntimeError(
|
|
"[LLS] comfy.utils is not available, so IMAGE resizing cannot run outside a ComfyUI environment."
|
|
) from _COMFY_UTILS_ERR
|
|
|
|
crop_mode = "center" if resize_mode == "crop_center" else "disabled"
|
|
try:
|
|
channel_first = image.movedim(-1, 1)
|
|
resized = comfy_utils.common_upscale(
|
|
channel_first,
|
|
target_width,
|
|
target_height,
|
|
"bilinear",
|
|
crop_mode,
|
|
)
|
|
return resized.movedim(1, -1)
|
|
except Exception as exc:
|
|
raise RuntimeError(f"[LLS] Failed to resize IMAGE for VAE encode: {exc}") from exc
|
|
|
|
|
|
def _infer_family_from_latent_or_vae(samples, vae=None) -> str:
|
|
tagged_family = getattr(vae, "_lls_family", None)
|
|
if tagged_family and tagged_family not in {"Auto", "auto", ""}:
|
|
return canonicalize_family(tagged_family)
|
|
|
|
if isinstance(samples, dict):
|
|
try:
|
|
ratio = int(samples.get("downscale_ratio_spacial", 0))
|
|
if ratio >= 16:
|
|
return "FLUX_DEV"
|
|
except Exception:
|
|
pass
|
|
latent_tensor = samples.get("samples")
|
|
latent_shape = tuple(getattr(latent_tensor, "shape", ()))
|
|
if len(latent_shape) >= 2:
|
|
try:
|
|
if int(latent_shape[1]) >= 16:
|
|
return "FLUX_DEV"
|
|
except Exception:
|
|
pass
|
|
|
|
return "SD1.5"
|
|
|
|
|
|
class LLSSimpleVAEEncode:
|
|
"""
|
|
把 IMAGE 编码为 LATENT，供 img2img 直接接入现有 KSampler。
|
|
编码本身复用 ComfyUI 原生 vae.encode，只在外层补尺寸处理与元信息。
|
|
"""
|
|
|
|
CATEGORY = "LLS/Image"
|
|
FUNCTION = "encode"
|
|
RETURN_TYPES = ("LATENT", "INT", "INT", "STRING")
|
|
RETURN_NAMES = ("latent", "width", "height", "latent_info")
|
|
DESCRIPTION = "Encode an IMAGE into a LATENT tensor for img2img workflows."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"vae": ("VAE",),
|
|
"resize_mode": (_RESIZE_MODES, {"default": "keep_aspect"}),
|
|
"size_source": (_SIZE_SOURCES, {"default": "input_image"}),
|
|
"width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8}),
|
|
"height": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8}),
|
|
"model_family": (MODEL_FAMILY_CHOICES, {"default": "Auto"}),
|
|
},
|
|
"optional": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
},
|
|
}
|
|
|
|
def encode(
|
|
self,
|
|
image,
|
|
vae,
|
|
resize_mode: str,
|
|
size_source: str,
|
|
width: int,
|
|
height: int,
|
|
model_family: str = "Auto",
|
|
model=None,
|
|
clip=None,
|
|
):
|
|
if image is None:
|
|
raise RuntimeError("[LLS] IMAGE input is required for LLS Simple VAE Encode.")
|
|
if vae is None:
|
|
raise RuntimeError(
|
|
"[LLS] Missing VAE. Connect the Loader VAE output or choose an external VAE in the loader."
|
|
)
|
|
|
|
family = resolve_model_family(model_family, model=model, clip=clip)
|
|
defaults = get_family_defaults(family)
|
|
latent_spec = get_latent_spec(defaults)
|
|
alignment = max(8, int(latent_spec["downscale_ratio"]))
|
|
batch_size, source_width, source_height = _get_image_dimensions(image)
|
|
|
|
requested_width, requested_height = _resolve_requested_size(
|
|
size_source=size_source,
|
|
width=width,
|
|
height=height,
|
|
source_width=source_width,
|
|
source_height=source_height,
|
|
model_info={
|
|
"default_width": defaults["default_width"],
|
|
"default_height": defaults["default_height"],
|
|
},
|
|
)
|
|
requested_width = _round_to_multiple(requested_width, alignment)
|
|
requested_height = _round_to_multiple(requested_height, alignment)
|
|
|
|
if resize_mode == "none":
|
|
final_width = _round_to_multiple(source_width, alignment)
|
|
final_height = _round_to_multiple(source_height, alignment)
|
|
elif resize_mode == "keep_aspect":
|
|
final_width, final_height = _fit_inside_box(
|
|
source_width=source_width,
|
|
source_height=source_height,
|
|
target_width=requested_width,
|
|
target_height=requested_height,
|
|
multiple=alignment,
|
|
)
|
|
elif resize_mode in {"crop_center", "stretch"}:
|
|
final_width = requested_width
|
|
final_height = requested_height
|
|
else:
|
|
raise RuntimeError(f"[LLS] Unsupported resize_mode '{resize_mode}'.")
|
|
|
|
processed_image = _resize_image_to(
|
|
image=image,
|
|
target_width=final_width,
|
|
target_height=final_height,
|
|
resize_mode=resize_mode,
|
|
)
|
|
|
|
try:
|
|
latent = vae.encode(processed_image)
|
|
except Exception as exc:
|
|
raise RuntimeError(f"[LLS] VAE encode failed: {exc}") from exc
|
|
|
|
latent_channels = int(getattr(latent, "shape", [0, latent_spec["latent_channels"]])[1])
|
|
latent_payload = {
|
|
"samples": latent,
|
|
"downscale_ratio_spacial": int(latent_spec["downscale_ratio"]),
|
|
"source": "image_encode",
|
|
}
|
|
latent_info = info_to_json(
|
|
{
|
|
"model_family": family,
|
|
"task_mode": "img2img",
|
|
"latent_source": "image_encode",
|
|
"input_image_width": source_width,
|
|
"input_image_height": source_height,
|
|
"width": final_width,
|
|
"height": final_height,
|
|
"batch_size": batch_size,
|
|
"resize_mode": resize_mode,
|
|
"size_source": size_source,
|
|
"vae_name": resolve_vae_name(vae, fallback="unknown"),
|
|
"latent_channels": latent_channels,
|
|
"downscale_ratio": int(latent_spec["downscale_ratio"]),
|
|
}
|
|
)
|
|
return (latent_payload, final_width, final_height, latent_info)
|
|
|
|
|
|
class LLSSimpleVAEDecode:
|
|
"""
|
|
简化版 VAE Decode 节点。
|
|
内部复用 ComfyUI 原生 VAE.decode。
|
|
"""
|
|
|
|
CATEGORY = "LLS/Image"
|
|
FUNCTION = "decode"
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "decode_info")
|
|
DESCRIPTION = "Decode a LATENT tensor into a pixel IMAGE using the VAE."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"samples": ("LATENT",),
|
|
"vae": ("VAE",),
|
|
},
|
|
}
|
|
|
|
def decode(self, samples, vae):
|
|
family = _infer_family_from_latent_or_vae(samples, vae)
|
|
|
|
if vae is None:
|
|
if family.startswith("FLUX"):
|
|
raise RuntimeError(
|
|
"[LLS] Missing FLUX AE/VAE. Place ae.safetensors in ComfyUI/models/vae/ "
|
|
"or set Loader.vae_source to 'external'."
|
|
)
|
|
raise RuntimeError(
|
|
"[LLS] Missing VAE. Connect the Loader VAE output or choose an external VAE in the loader."
|
|
)
|
|
if samples is None:
|
|
raise RuntimeError(
|
|
"[LLS] LATENT samples is None. Connect the output of KSampler to this node."
|
|
)
|
|
|
|
latent_tensor = samples.get("samples")
|
|
if latent_tensor is None:
|
|
raise RuntimeError(
|
|
"[LLS] LATENT dict does not contain 'samples'. Make sure the upstream node outputs a valid LATENT."
|
|
)
|
|
if getattr(latent_tensor, "is_nested", False):
|
|
latent_tensor = latent_tensor.unbind()[0]
|
|
|
|
try:
|
|
image = vae.decode(latent_tensor)
|
|
except Exception as exc:
|
|
raise RuntimeError(f"[LLS] VAE decode failed: {exc}") from exc
|
|
|
|
if len(getattr(image, "shape", ())) == 5:
|
|
image = image.reshape(-1, image.shape[-3], image.shape[-2], image.shape[-1])
|
|
|
|
defaults = get_family_defaults(family)
|
|
downscale_ratio = int(samples.get("downscale_ratio_spacial", defaults["downscale_ratio"]))
|
|
height = latent_tensor.shape[2] * downscale_ratio
|
|
width = latent_tensor.shape[3] * downscale_ratio
|
|
decode_info = info_to_json(
|
|
{
|
|
"model_family": family,
|
|
"width": width,
|
|
"height": height,
|
|
"batch_size": latent_tensor.shape[0],
|
|
"vae_name": resolve_vae_name(vae, fallback="unknown"),
|
|
"decode_stage": "vae_decode",
|
|
}
|
|
)
|
|
return (image, decode_info)
|
|
|
|
|
|
class LLSSaveImage:
|
|
"""
|
|
复用 ComfyUI 原生 SaveImage，并把 LLS 生成链路信息合并进 PNG metadata。
|
|
"""
|
|
|
|
CATEGORY = "LLS/Image"
|
|
FUNCTION = "save"
|
|
OUTPUT_NODE = True
|
|
RETURN_TYPES = ()
|
|
DESCRIPTION = "Save images with LLS generation metadata."
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"filename_prefix": ("STRING", {"default": "LLS"}),
|
|
"output_mode": (_SAVE_OUTPUT_MODES, {"default": "save"}),
|
|
"save_metadata": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
"prompt_info": ("STRING", {"forceInput": True}),
|
|
"latent_info": ("STRING", {"forceInput": True}),
|
|
"sample_info": ("STRING", {"forceInput": True}),
|
|
"decode_info": ("STRING", {"forceInput": True}),
|
|
"upscale_info": ("STRING", {"forceInput": True}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
def _merge_ui_results(self, primary: dict, secondary: dict | None = None) -> dict:
|
|
if not primary:
|
|
return secondary or {}
|
|
if not secondary:
|
|
return primary
|
|
|
|
merged = dict(primary or {})
|
|
merged_ui = dict(merged.get("ui", {}))
|
|
merged_images = list(merged_ui.get("images", []))
|
|
merged_images.extend(list((secondary.get("ui", {}) or {}).get("images", [])))
|
|
merged_ui["images"] = merged_images
|
|
merged["ui"] = merged_ui
|
|
return merged
|
|
|
|
def _build_metadata(
|
|
self,
|
|
image,
|
|
prompt_info: str | None = None,
|
|
latent_info: str | None = None,
|
|
sample_info: str | None = None,
|
|
decode_info: str | None = None,
|
|
upscale_info: str | None = None,
|
|
) -> dict:
|
|
prompt = parse_jsonish_info(prompt_info)
|
|
latent = parse_jsonish_info(latent_info)
|
|
sample = parse_jsonish_info(sample_info)
|
|
decode = parse_jsonish_info(decode_info)
|
|
upscale = parse_jsonish_info(upscale_info)
|
|
|
|
family = _first_non_empty(
|
|
prompt.get("model_family"),
|
|
latent.get("model_family"),
|
|
decode.get("model_family"),
|
|
sample.get("model_family"),
|
|
sample.get("family"),
|
|
)
|
|
checkpoint_name = _first_non_empty(
|
|
prompt.get("checkpoint_name"),
|
|
latent.get("checkpoint_name"),
|
|
decode.get("checkpoint_name"),
|
|
sample.get("checkpoint_name"),
|
|
sample.get("model_name"),
|
|
)
|
|
text_encoder_name = _first_non_empty(
|
|
prompt.get("text_encoder_name"),
|
|
prompt.get("text_encoder_name_1"),
|
|
prompt.get("text_encoder_name_2"),
|
|
latent.get("text_encoder_name"),
|
|
decode.get("text_encoder_name"),
|
|
)
|
|
vae_name = _first_non_empty(
|
|
decode.get("vae_name"),
|
|
latent.get("vae_name"),
|
|
prompt.get("vae_name"),
|
|
)
|
|
|
|
metadata = {
|
|
"positive_prompt": prompt.get("positive_prompt", ""),
|
|
"negative_prompt": prompt.get("negative_prompt", ""),
|
|
"seed": sample.get("seed"),
|
|
"model_family": family,
|
|
"checkpoint_name": checkpoint_name,
|
|
"vae_name": vae_name,
|
|
"text_encoder_name": text_encoder_name,
|
|
"steps": sample.get("steps"),
|
|
"cfg": sample.get("cfg"),
|
|
"guidance": sample.get("guidance"),
|
|
"sampler_name": sample.get("sampler_name"),
|
|
"scheduler": sample.get("scheduler"),
|
|
"denoise": sample.get("denoise"),
|
|
"width": decode.get("width") or latent.get("width") or image.shape[2],
|
|
"height": decode.get("height") or latent.get("height") or image.shape[1],
|
|
"batch_size": decode.get("batch_size") or latent.get("batch_size") or image.shape[0],
|
|
"upscale_mode": upscale.get("mode") or upscale.get("upscale_mode"),
|
|
"scale": upscale.get("scale") or upscale.get("upscale_scale"),
|
|
"timestamp": datetime.now().isoformat(timespec="seconds"),
|
|
}
|
|
return metadata
|
|
|
|
def save(
|
|
self,
|
|
image=None,
|
|
filename_prefix: str = "LLS",
|
|
output_mode: str = "save",
|
|
save_metadata: bool = True,
|
|
mask=None,
|
|
prompt_info: str | None = None,
|
|
latent_info: str | None = None,
|
|
sample_info: str | None = None,
|
|
decode_info: str | None = None,
|
|
upscale_info: str | None = None,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
if image is None and mask is None:
|
|
raise RuntimeError("[LLS] At least one of IMAGE or MASK must be connected for LLS Save Image.")
|
|
|
|
if output_mode == "preview_only":
|
|
if comfy_core_nodes is None or not hasattr(comfy_core_nodes, "PreviewImage"):
|
|
raise RuntimeError("[LLS] ComfyUI core PreviewImage node is not available.")
|
|
previewer = comfy_core_nodes.PreviewImage()
|
|
image_result = None
|
|
if image is not None:
|
|
image_result = previewer.save_images(
|
|
image,
|
|
prompt=prompt,
|
|
extra_pnginfo=None,
|
|
)
|
|
|
|
mask_result = None
|
|
if mask is not None:
|
|
mask_result = previewer.save_images(
|
|
mask_to_image(mask),
|
|
prompt=prompt,
|
|
extra_pnginfo=None,
|
|
)
|
|
return self._merge_ui_results(image_result, mask_result)
|
|
|
|
if comfy_core_nodes is None or not hasattr(comfy_core_nodes, "SaveImage"):
|
|
raise RuntimeError("[LLS] ComfyUI core SaveImage node is not available.")
|
|
|
|
saver = comfy_core_nodes.SaveImage()
|
|
image_result = None
|
|
if image is not None:
|
|
merged_extra_pnginfo = dict(extra_pnginfo or {})
|
|
if save_metadata:
|
|
merged_extra_pnginfo["lls_metadata"] = self._build_metadata(
|
|
image=image,
|
|
prompt_info=prompt_info,
|
|
latent_info=latent_info,
|
|
sample_info=sample_info,
|
|
decode_info=decode_info,
|
|
upscale_info=upscale_info,
|
|
)
|
|
image_result = saver.save_images(
|
|
image,
|
|
filename_prefix=filename_prefix,
|
|
prompt=prompt,
|
|
extra_pnginfo=merged_extra_pnginfo,
|
|
)
|
|
|
|
mask_result = None
|
|
if mask is not None:
|
|
mask_filename_prefix = filename_prefix if image is None else f"{filename_prefix}_mask"
|
|
mask_result = saver.save_images(
|
|
mask_to_image(mask),
|
|
filename_prefix=mask_filename_prefix,
|
|
prompt=prompt,
|
|
extra_pnginfo={},
|
|
)
|
|
return self._merge_ui_results(image_result, mask_result)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS: dict[str, type] = {
|
|
"LLSSimpleVAEEncode": LLSSimpleVAEEncode,
|
|
"LLSSimpleVAEDecode": LLSSimpleVAEDecode,
|
|
"LLSSaveImage": LLSSaveImage,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = {
|
|
"LLSSimpleVAEEncode": "LLS Simple VAE Encode",
|
|
"LLSSimpleVAEDecode": "LLS Simple VAE Decode",
|
|
"LLSSaveImage": "LLS Save Image",
|
|
}
|