1984 lines
78 KiB
Python
1984 lines
78 KiB
Python
"""
|
|
ComfyUI-LLPS-v1.2
|
|
LiveLatentPreviewer & Saver v1.2
|
|
|
|
This v1 intentionally avoids global monkeypatching. It provides:
|
|
- LLPS Controller: workflow-level preview management declaration for v2 UI.
|
|
- Legacy LLPS Config and LLPS KSampler nodes for v1.2 callback-based preview saving.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import json
|
|
import time
|
|
import logging
|
|
import base64
|
|
import threading
|
|
import urllib.error
|
|
import urllib.request
|
|
import uuid
|
|
from io import BytesIO
|
|
from dataclasses import dataclass, asdict
|
|
from typing import Any, Dict, Optional, Tuple
|
|
|
|
import torch
|
|
from PIL import Image
|
|
|
|
import nodes as comfy_nodes
|
|
import comfy.sample
|
|
import comfy.samplers
|
|
import comfy.utils
|
|
import comfy.model_management
|
|
import comfy.nested_tensor
|
|
import comfy_extras.nodes_custom_sampler as comfy_custom_sampler_nodes
|
|
from comfy.cli_args import args
|
|
from comfy_execution.utils import get_executing_context
|
|
import folder_paths
|
|
import latent_preview
|
|
|
|
|
|
LLPS_CATEGORY = "LLPS/Live Latent Preview"
|
|
LLPS_LEGACY_CATEGORY = "LLPS/Legacy v1.2"
|
|
TAESD_PREVIEWER_CLASSES = {"TAESDPreviewerImpl", "TAEHVPreviewerImpl"}
|
|
TAESD_DECODER_URLS = {
|
|
"taesd_decoder": "https://raw.githubusercontent.com/madebyollin/taesd/main/taesd_decoder.pth",
|
|
"taesdxl_decoder": "https://raw.githubusercontent.com/madebyollin/taesd/main/taesdxl_decoder.pth",
|
|
"taesd3_decoder": "https://raw.githubusercontent.com/madebyollin/taesd/main/taesd3_decoder.pth",
|
|
"taef1_decoder": "https://raw.githubusercontent.com/madebyollin/taesd/main/taef1_decoder.pth",
|
|
}
|
|
LLPS_SAMPLER_TYPES = {"KSampler", "KSamplerAdvanced", "SamplerCustom", "SamplerCustomAdvanced"}
|
|
_LLPS_RUN_LOCK = threading.Lock()
|
|
_LLPS_PREVIEW_RUNS: Dict[Tuple[str, str], list[Dict[str, Any]]] = {}
|
|
|
|
|
|
def _sanitize_piece(text: str, fallback: str = "LLPS") -> str:
|
|
text = str(text or "").strip()
|
|
if not text:
|
|
return fallback
|
|
# Windows-safe and reasonably portable.
|
|
text = re.sub(r'[<>:"/\\|?*\x00-\x1f]', "_", text)
|
|
text = re.sub(r"\s+", " ", text).strip().strip(".")
|
|
return text[:120] or fallback
|
|
|
|
|
|
def _bool(v: Any) -> bool:
|
|
return bool(v) and str(v).lower() not in {"false", "0", "none", "off"}
|
|
|
|
|
|
@dataclass
|
|
class LLPSConfigData:
|
|
enabled: bool = True
|
|
live_preview_method: str = "latent2rgb" # server_default, none, auto, latent2rgb, taesd
|
|
show_preview: bool = True
|
|
save_preview: bool = False
|
|
save_base_path: str = ""
|
|
subfolder: str = ""
|
|
filename_prefix: str = "LLPS"
|
|
image_format: str = "jpg" # jpg, png, webp
|
|
save_every_n_steps: int = 1
|
|
save_metadata_json: bool = True
|
|
jpeg_quality: int = 90
|
|
preview_max_side: int = 512
|
|
|
|
@classmethod
|
|
def from_obj(cls, obj: Optional[Dict[str, Any]]) -> "LLPSConfigData":
|
|
if isinstance(obj, cls):
|
|
return obj
|
|
if not isinstance(obj, dict):
|
|
return cls()
|
|
data = cls()
|
|
for k in asdict(data).keys():
|
|
if k in obj:
|
|
setattr(data, k, obj[k])
|
|
data.enabled = _bool(data.enabled)
|
|
data.show_preview = _bool(data.show_preview)
|
|
data.save_preview = _bool(data.save_preview)
|
|
data.save_metadata_json = _bool(data.save_metadata_json)
|
|
data.save_every_n_steps = max(1, int(data.save_every_n_steps or 1))
|
|
data.jpeg_quality = max(1, min(100, int(data.jpeg_quality or 90)))
|
|
data.preview_max_side = 1024 if int(data.preview_max_side or 512) >= 1024 else 512
|
|
data.live_preview_method = str(data.live_preview_method or "server_default")
|
|
if data.live_preview_method == "default":
|
|
data.live_preview_method = "server_default"
|
|
data.image_format = str(data.image_format or "jpg").lower()
|
|
if data.image_format == "jpeg":
|
|
data.image_format = "jpg"
|
|
if data.image_format not in {"jpg", "png", "webp"}:
|
|
data.image_format = "jpg"
|
|
return data
|
|
|
|
|
|
@dataclass
|
|
class LLPSControllerData:
|
|
enabled: bool = True
|
|
live_preview_method: str = "latent2rgb"
|
|
save_also: bool = False
|
|
save_base_path: str = ""
|
|
subfolder: str = ""
|
|
filename_prefix: str = "LLPS"
|
|
image_format: str = "jpg"
|
|
save_every_n_steps: int = 1
|
|
save_metadata_json: bool = True
|
|
jpeg_quality: int = 90
|
|
preview_max_side: int = 512
|
|
controller_node_id: Optional[str] = None
|
|
|
|
@classmethod
|
|
def from_obj(cls, obj: Optional[Dict[str, Any]]) -> "LLPSControllerData":
|
|
if isinstance(obj, cls):
|
|
return obj
|
|
if not isinstance(obj, dict):
|
|
return cls()
|
|
data = cls()
|
|
for key in asdict(data).keys():
|
|
if key in obj:
|
|
setattr(data, key, obj[key])
|
|
data.enabled = _bool(data.enabled)
|
|
data.save_also = _bool(data.save_also)
|
|
data.live_preview_method = str(data.live_preview_method or "server_default")
|
|
if data.live_preview_method == "default":
|
|
data.live_preview_method = "server_default"
|
|
data.image_format = str(data.image_format or "jpg").lower()
|
|
if data.image_format == "jpeg":
|
|
data.image_format = "jpg"
|
|
if data.image_format not in {"jpg", "png", "webp"}:
|
|
data.image_format = "jpg"
|
|
data.save_every_n_steps = max(1, int(data.save_every_n_steps or 1))
|
|
data.jpeg_quality = max(1, min(100, int(data.jpeg_quality or 90)))
|
|
data.preview_max_side = 1024 if int(data.preview_max_side or 512) >= 1024 else 512
|
|
return data
|
|
|
|
def to_legacy_config(self, show_preview: bool, save_preview: bool) -> LLPSConfigData:
|
|
return LLPSConfigData.from_obj(
|
|
{
|
|
"enabled": self.enabled,
|
|
"live_preview_method": self.live_preview_method,
|
|
"show_preview": show_preview,
|
|
"save_preview": save_preview,
|
|
"save_base_path": self.save_base_path,
|
|
"subfolder": self.subfolder,
|
|
"filename_prefix": self.filename_prefix,
|
|
"image_format": self.image_format,
|
|
"save_every_n_steps": self.save_every_n_steps,
|
|
"save_metadata_json": self.save_metadata_json,
|
|
"jpeg_quality": self.jpeg_quality,
|
|
"preview_max_side": self.preview_max_side,
|
|
}
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class LLPSStreamControllerData:
|
|
enabled: bool = True
|
|
target_sampler_id: str = ""
|
|
override_save_also: bool = False
|
|
save_also: bool = False
|
|
override_subfolder: bool = False
|
|
subfolder: str = ""
|
|
override_filename_prefix: bool = False
|
|
filename_prefix: str = ""
|
|
stream_controller_node_id: Optional[str] = None
|
|
|
|
@classmethod
|
|
def from_obj(cls, obj: Optional[Dict[str, Any]]) -> "LLPSStreamControllerData":
|
|
if isinstance(obj, cls):
|
|
return obj
|
|
if not isinstance(obj, dict):
|
|
return cls()
|
|
data = cls()
|
|
for key in asdict(data).keys():
|
|
if key in obj:
|
|
setattr(data, key, obj[key])
|
|
data.enabled = _bool(data.enabled)
|
|
data.target_sampler_id = str(data.target_sampler_id or "").strip()
|
|
data.override_save_also = _bool(data.override_save_also)
|
|
data.save_also = _bool(data.save_also)
|
|
data.override_subfolder = _bool(data.override_subfolder)
|
|
data.subfolder = str(data.subfolder or "").strip()
|
|
data.override_filename_prefix = _bool(data.override_filename_prefix)
|
|
data.filename_prefix = str(data.filename_prefix or "").strip()
|
|
data.stream_controller_node_id = str(data.stream_controller_node_id) if data.stream_controller_node_id is not None else None
|
|
return data
|
|
|
|
|
|
def _preview_method_value(method: Any) -> str:
|
|
if hasattr(method, "value"):
|
|
return str(method.value)
|
|
return str(method)
|
|
|
|
|
|
def _previewer_identity(previewer: Any) -> Dict[str, Any]:
|
|
if previewer is None:
|
|
return {
|
|
"actual_previewer_class": None,
|
|
"actual_previewer_module": None,
|
|
"actual_previewer_id": None,
|
|
}
|
|
return {
|
|
"actual_previewer_class": previewer.__class__.__name__,
|
|
"actual_previewer_module": previewer.__class__.__module__,
|
|
"actual_previewer_id": id(previewer),
|
|
}
|
|
|
|
|
|
def _is_taesd_previewer(previewer: Any) -> bool:
|
|
return previewer is not None and previewer.__class__.__name__ in TAESD_PREVIEWER_CLASSES
|
|
|
|
|
|
def _previewer_fallback_reason(requested_method: str, resolved_method: str, previewer: Any) -> Optional[str]:
|
|
requested = str(requested_method or "server_default")
|
|
resolved = str(resolved_method or "none")
|
|
cls_name = previewer.__class__.__name__ if previewer is not None else None
|
|
|
|
if previewer is None:
|
|
if requested in {"none", "disabled"} or resolved == "none":
|
|
return "preview_disabled"
|
|
return "previewer_unavailable"
|
|
if requested == "server_default":
|
|
return "server_default_used"
|
|
if requested == "auto":
|
|
return "auto_resolved_by_comfy"
|
|
if requested == "taesd" and cls_name not in TAESD_PREVIEWER_CLASSES:
|
|
return "taesd_requested_but_taesd_previewer_unavailable_or_fell_back"
|
|
if requested == "latent2rgb" and cls_name != "Latent2RGBPreviewer":
|
|
return "latent2rgb_requested_but_different_previewer_returned"
|
|
return None
|
|
|
|
|
|
def _taesd_decoder_candidates(decoder_name: Optional[str]) -> list[str]:
|
|
if not decoder_name:
|
|
return []
|
|
try:
|
|
if hasattr(folder_paths, "get_filename_list_"):
|
|
filenames = folder_paths.get_filename_list_("vae_approx")[0]
|
|
else:
|
|
filenames = folder_paths.get_filename_list("vae_approx")
|
|
return [
|
|
filename
|
|
for filename in filenames
|
|
if filename.startswith(decoder_name)
|
|
]
|
|
except Exception as e:
|
|
logging.warning("[LLPS] Could not list TAESD decoder candidates for %s: %s", decoder_name, e)
|
|
return []
|
|
|
|
|
|
def _refresh_vae_approx_cache() -> None:
|
|
try:
|
|
folder_paths.filename_list_cache.pop("vae_approx", None)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
folder_paths.cache_helper.clear()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _ensure_taesd_decoder_available(decoder_name: Optional[str]) -> Dict[str, Any]:
|
|
result = {
|
|
"attempted": False,
|
|
"success": False,
|
|
"decoder_name": decoder_name,
|
|
"url": None,
|
|
"path": None,
|
|
"error": None,
|
|
}
|
|
if not decoder_name:
|
|
result["error"] = "latent_format_has_no_taesd_decoder_name"
|
|
return result
|
|
|
|
candidates = _taesd_decoder_candidates(decoder_name)
|
|
if candidates:
|
|
result["success"] = True
|
|
result["path"] = folder_paths.get_full_path("vae_approx", candidates[0])
|
|
return result
|
|
|
|
url = TAESD_DECODER_URLS.get(decoder_name)
|
|
result["url"] = url
|
|
if not url:
|
|
result["error"] = "no_known_download_url_for_decoder"
|
|
return result
|
|
|
|
folders = folder_paths.get_folder_paths("vae_approx")
|
|
if not folders:
|
|
result["error"] = "vae_approx_folder_not_registered"
|
|
return result
|
|
|
|
target_dir = folders[0]
|
|
target_path = os.path.join(target_dir, f"{decoder_name}.pth")
|
|
result["path"] = target_path
|
|
result["attempted"] = True
|
|
|
|
try:
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
logging.info("[LLPS] Downloading missing TAESD decoder %s to %s", decoder_name, target_path)
|
|
request = urllib.request.Request(url, headers={"User-Agent": "ComfyUI-LLPS/2"})
|
|
tmp_path = f"{target_path}.part"
|
|
with urllib.request.urlopen(request, timeout=60) as response, open(tmp_path, "wb") as f:
|
|
while True:
|
|
chunk = response.read(1024 * 1024)
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
if os.path.getsize(tmp_path) <= 0:
|
|
raise RuntimeError("downloaded decoder file is empty")
|
|
os.replace(tmp_path, target_path)
|
|
_refresh_vae_approx_cache()
|
|
folder_paths.get_filename_list("vae_approx")
|
|
result["success"] = True
|
|
except (OSError, urllib.error.URLError, RuntimeError) as e:
|
|
result["error"] = str(e)
|
|
try:
|
|
if os.path.exists(f"{target_path}.part"):
|
|
os.remove(f"{target_path}.part")
|
|
except OSError:
|
|
pass
|
|
logging.warning("[LLPS] Could not install TAESD decoder %s: %s", decoder_name, e)
|
|
|
|
return result
|
|
|
|
|
|
def _previewer_for_method_with_info(model, method: str, requested_method: Optional[str] = None):
|
|
requested = str(requested_method or method or "server_default")
|
|
source = str(method or "server_default")
|
|
previous = args.preview_method
|
|
previous_method = _preview_method_value(previous)
|
|
info = {
|
|
"requested_live_preview_method": requested,
|
|
"requested_preview_source_method": source,
|
|
"previous_global_preview_method": previous_method,
|
|
"resolved_preview_source_method": None,
|
|
"actual_previewer_class": None,
|
|
"actual_previewer_module": None,
|
|
"actual_previewer_id": None,
|
|
"fallback_reason": None,
|
|
"latent_format_class": None,
|
|
"latent_format_module": None,
|
|
"taesd_decoder_name": None,
|
|
"taesd_decoder_candidates": [],
|
|
"taesd_decoder_install": None,
|
|
"latent_rgb_factors_available": None,
|
|
}
|
|
|
|
latent_format = getattr(getattr(model, "model", None), "latent_format", None)
|
|
if latent_format is not None:
|
|
info["latent_format_class"] = latent_format.__class__.__name__
|
|
info["latent_format_module"] = latent_format.__class__.__module__
|
|
info["taesd_decoder_name"] = getattr(latent_format, "taesd_decoder_name", None)
|
|
info["taesd_decoder_candidates"] = _taesd_decoder_candidates(info["taesd_decoder_name"])
|
|
info["latent_rgb_factors_available"] = getattr(latent_format, "latent_rgb_factors", None) is not None
|
|
|
|
if source in {"none", "disabled"}:
|
|
info["resolved_preview_source_method"] = "none"
|
|
info["fallback_reason"] = _previewer_fallback_reason(requested, "none", None)
|
|
return None, info
|
|
|
|
try:
|
|
if source not in {"server_default", "default", ""}:
|
|
latent_preview.set_preview_method(source)
|
|
resolved = _preview_method_value(args.preview_method)
|
|
previewer = latent_preview.get_previewer(model.load_device, model.model.latent_format)
|
|
info.update(_previewer_identity(previewer))
|
|
info["resolved_preview_source_method"] = resolved
|
|
info["fallback_reason"] = _previewer_fallback_reason(requested, resolved, previewer)
|
|
if source == "taesd" and not _is_taesd_previewer(previewer):
|
|
install_info = _ensure_taesd_decoder_available(info["taesd_decoder_name"])
|
|
info["taesd_decoder_install"] = install_info
|
|
info["taesd_decoder_candidates"] = _taesd_decoder_candidates(info["taesd_decoder_name"])
|
|
if install_info.get("success"):
|
|
previewer = latent_preview.get_previewer(model.load_device, model.model.latent_format)
|
|
info.update(_previewer_identity(previewer))
|
|
info["fallback_reason"] = _previewer_fallback_reason(requested, resolved, previewer)
|
|
return previewer, info
|
|
finally:
|
|
args.preview_method = previous
|
|
|
|
|
|
def _save_base_dir(cfg: LLPSConfigData) -> str:
|
|
base = cfg.save_base_path.strip()
|
|
if not base:
|
|
base = os.path.join(folder_paths.get_output_directory(), "LLPS")
|
|
# Relative paths are placed below ComfyUI output directory for safety/predictability.
|
|
if not os.path.isabs(base):
|
|
base = os.path.join(folder_paths.get_output_directory(), base)
|
|
return base
|
|
|
|
|
|
def _effective_subfolder(cfg: LLPSConfigData, node_label: str, sampler_subfolder: Optional[str] = None) -> str:
|
|
per_sampler = _sanitize_piece(sampler_subfolder or "", "")
|
|
if per_sampler:
|
|
return per_sampler
|
|
controller_subfolder = _sanitize_piece(cfg.subfolder, "")
|
|
if controller_subfolder:
|
|
return controller_subfolder
|
|
return _sanitize_piece(node_label, "LLPS_KSampler")
|
|
|
|
|
|
def _resolve_save_dir(cfg: LLPSConfigData, node_label: str, sampler_subfolder: Optional[str] = None) -> str:
|
|
base = os.path.join(_save_base_dir(cfg), _effective_subfolder(cfg, node_label, sampler_subfolder))
|
|
os.makedirs(base, exist_ok=True)
|
|
return base
|
|
|
|
|
|
def _previewer_for_method(model, method: str):
|
|
"""Create a ComfyUI latent previewer for a specific method, then restore global setting.
|
|
|
|
ComfyUI's latent_preview.get_previewer currently reads args.preview_method. v1 therefore
|
|
changes it only while constructing the previewer, not during the whole sample.
|
|
"""
|
|
previewer, _info = _previewer_for_method_with_info(model, method, method)
|
|
return previewer
|
|
|
|
|
|
def _normalize_preview_image(img: Image.Image, preview_max_side: int) -> Tuple[Image.Image, Dict[str, int]]:
|
|
original_width, original_height = img.size
|
|
max_side = 1024 if int(preview_max_side or 512) >= 1024 else 512
|
|
longest = max(original_width, original_height, 1)
|
|
scale = max_side / longest
|
|
target_width = max(1, int(round(original_width * scale)))
|
|
target_height = max(1, int(round(original_height * scale)))
|
|
|
|
if (target_width, target_height) == img.size:
|
|
out = img.copy()
|
|
else:
|
|
out = img.resize((target_width, target_height), Image.Resampling.LANCZOS)
|
|
|
|
return out, {
|
|
"preview_max_side": max_side,
|
|
"preview_original_width": int(original_width),
|
|
"preview_original_height": int(original_height),
|
|
"preview_target_width": int(target_width),
|
|
"preview_target_height": int(target_height),
|
|
}
|
|
|
|
|
|
def _normalize_preview_tuple(preview_tuple: Tuple[Any, ...], preview_max_side: int) -> Tuple[Optional[Tuple[Any, ...]], Dict[str, int]]:
|
|
if not preview_tuple or len(preview_tuple) < 2 or not isinstance(preview_tuple[1], Image.Image):
|
|
return None, {}
|
|
out, info = _normalize_preview_image(preview_tuple[1], preview_max_side)
|
|
image_type = preview_tuple[0] if len(preview_tuple) > 0 else "JPEG"
|
|
return (image_type, out, None), info
|
|
|
|
|
|
def _preview_data_url(img: Image.Image, image_type: str = "JPEG", quality: int = 90) -> str:
|
|
pil_format = "PNG" if str(image_type).upper() == "PNG" else "JPEG"
|
|
out = img
|
|
if pil_format == "JPEG" and out.mode not in {"RGB", "L"}:
|
|
out = out.convert("RGB")
|
|
buffer = BytesIO()
|
|
kwargs = {"quality": quality} if pil_format == "JPEG" else {}
|
|
out.save(buffer, pil_format, **kwargs)
|
|
mime = "image/png" if pil_format == "PNG" else "image/jpeg"
|
|
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
|
return f"data:{mime};base64,{encoded}"
|
|
|
|
|
|
def _send_llps_preview_event(payload: Dict[str, Any]) -> None:
|
|
try:
|
|
import server as comfy_server
|
|
|
|
instance = comfy_server.PromptServer.instance
|
|
instance.send_sync("llps_preview", payload, instance.client_id)
|
|
except Exception as e:
|
|
logging.debug("[LLPS] Could not send preview event: %s", e)
|
|
|
|
|
|
def _save_preview_image(
|
|
preview_tuple: Tuple[Any, ...],
|
|
save_dir: str,
|
|
cfg: LLPSConfigData,
|
|
node_label: str,
|
|
step: int,
|
|
total_steps: int,
|
|
seed: int,
|
|
method: str,
|
|
temp_run_id: Optional[str] = None,
|
|
sampler_node_id: Optional[str] = None,
|
|
) -> Optional[str]:
|
|
if not preview_tuple or len(preview_tuple) < 2:
|
|
return None
|
|
|
|
img = preview_tuple[1]
|
|
if not isinstance(img, Image.Image):
|
|
return None
|
|
|
|
out = img.copy()
|
|
max_res = preview_tuple[2] if len(preview_tuple) >= 3 else None
|
|
if isinstance(max_res, int) and max_res > 0:
|
|
out.thumbnail((max_res, max_res), Image.Resampling.LANCZOS)
|
|
|
|
ext = cfg.image_format
|
|
pil_format = {"jpg": "JPEG", "png": "PNG", "webp": "WEBP"}[ext]
|
|
prefix = _sanitize_piece(cfg.filename_prefix, "LLPS")
|
|
label = _sanitize_piece(node_label, "KSampler")
|
|
if temp_run_id:
|
|
node_part = _sanitize_piece(sampler_node_id or label, "node")
|
|
run_part = _sanitize_piece(temp_run_id, "run")
|
|
filename = f"tmp_{run_part}_n{node_part}_step{step + 1:04d}of{int(total_steps):04d}.{ext}"
|
|
else:
|
|
filename = f"{prefix}_{label}_step{step + 1:04d}of{int(total_steps):04d}_{method}.{ext}"
|
|
path = os.path.join(save_dir, filename)
|
|
|
|
save_kwargs = {}
|
|
if pil_format == "JPEG":
|
|
if out.mode not in {"RGB", "L"}:
|
|
out = out.convert("RGB")
|
|
save_kwargs.update({"quality": cfg.jpeg_quality, "optimize": True})
|
|
elif pil_format == "WEBP":
|
|
save_kwargs.update({"quality": cfg.jpeg_quality})
|
|
|
|
out.save(path, pil_format, **save_kwargs)
|
|
return path
|
|
|
|
|
|
def _unique_path(path: str) -> str:
|
|
if not os.path.exists(path):
|
|
return path
|
|
root, ext = os.path.splitext(path)
|
|
counter = 2
|
|
while True:
|
|
candidate = f"{root}_{counter}{ext}"
|
|
if not os.path.exists(candidate):
|
|
return candidate
|
|
counter += 1
|
|
|
|
|
|
def _find_upstream_sampler_ids(prompt: Optional[Dict[str, Any]], node_id: Optional[str]) -> list[str]:
|
|
if not isinstance(prompt, dict) or node_id is None:
|
|
return []
|
|
found: list[str] = []
|
|
visited: set[str] = set()
|
|
|
|
def visit_value(value: Any) -> None:
|
|
if _is_prompt_link(value):
|
|
visit_node(str(value[0]))
|
|
|
|
def visit_node(current_id: str) -> None:
|
|
if current_id in visited:
|
|
return
|
|
visited.add(current_id)
|
|
node = prompt.get(current_id)
|
|
if not isinstance(node, dict):
|
|
return
|
|
if node.get("class_type") in LLPS_SAMPLER_TYPES:
|
|
found.append(current_id)
|
|
inputs = node.get("inputs", {})
|
|
if isinstance(inputs, dict):
|
|
for value in inputs.values():
|
|
if isinstance(value, list) and _is_prompt_link(value):
|
|
visit_value(value)
|
|
elif isinstance(value, list):
|
|
for item in value:
|
|
visit_value(item)
|
|
|
|
save_node = prompt.get(str(node_id))
|
|
if isinstance(save_node, dict):
|
|
inputs = save_node.get("inputs", {})
|
|
if isinstance(inputs, dict):
|
|
visit_value(inputs.get("images"))
|
|
return found
|
|
|
|
|
|
def _rename_pending_previews_for_output(
|
|
prompt_id: Optional[str],
|
|
sampler_ids: list[str],
|
|
final_images: list[Dict[str, Any]],
|
|
) -> None:
|
|
if not prompt_id or not sampler_ids or not final_images:
|
|
return
|
|
first_image = final_images[0]
|
|
final_filename = str(first_image.get("filename") or "")
|
|
final_basename = _sanitize_piece(os.path.splitext(os.path.basename(final_filename))[0], "")
|
|
if not final_basename:
|
|
return
|
|
|
|
with _LLPS_RUN_LOCK:
|
|
runs = []
|
|
for sampler_id in sampler_ids:
|
|
runs.extend(_LLPS_PREVIEW_RUNS.get((str(prompt_id), str(sampler_id)), []))
|
|
|
|
for run in runs:
|
|
if run.get("filename_resolution_status") == "renamed_to_final_image":
|
|
continue
|
|
metadata = run.get("metadata") if isinstance(run.get("metadata"), dict) else {}
|
|
prefix = _sanitize_piece(metadata.get("filename_prefix", "LLPS"), "LLPS")
|
|
renamed_files = []
|
|
for entry in run.get("saved_files", []):
|
|
old_path = entry.get("path")
|
|
if not old_path or not os.path.exists(old_path):
|
|
continue
|
|
ext = os.path.splitext(old_path)[1]
|
|
step = int(entry.get("step", 0)) + 1
|
|
total_steps = int(entry.get("total_steps", 0))
|
|
new_name = f"{prefix}_{final_basename}_step{step:04d}of{total_steps:04d}{ext}"
|
|
new_path = _unique_path(os.path.join(os.path.dirname(old_path), new_name))
|
|
try:
|
|
os.replace(old_path, new_path)
|
|
except OSError as e:
|
|
logging.warning("[LLPS] Could not rename preview %s to %s: %s", old_path, new_path, e)
|
|
continue
|
|
entry["path"] = new_path
|
|
entry["filename"] = os.path.basename(new_path)
|
|
renamed_files.append(os.path.basename(new_path))
|
|
|
|
if renamed_files:
|
|
run["renamed_files"] = renamed_files
|
|
run["filename_resolution_status"] = "renamed_to_final_image"
|
|
if isinstance(metadata, dict):
|
|
metadata["filename_resolution_status"] = "renamed_to_final_image"
|
|
metadata["final_output"] = first_image
|
|
metadata["saved_files"] = renamed_files
|
|
metadata["renamed_files"] = renamed_files
|
|
metadata["renamed_at"] = time.strftime("%Y-%m-%d %H:%M:%S")
|
|
metadata_path = run.get("metadata_path")
|
|
if metadata_path:
|
|
try:
|
|
_write_metadata_json(os.path.dirname(metadata_path), metadata, os.path.basename(metadata_path))
|
|
except OSError as e:
|
|
logging.warning("[LLPS] Could not update renamed preview metadata: %s", e)
|
|
_send_llps_preview_event(
|
|
{
|
|
"controller_node_id": run.get("controller_node_id"),
|
|
"sampler_node_id": run.get("sampler_node_id"),
|
|
"sampler_node_type": run.get("sampler_node_type"),
|
|
"sampler_label": run.get("node_label"),
|
|
"prompt_id": str(prompt_id),
|
|
"run_id": run.get("run_id"),
|
|
"stream_controller_node_id": run.get("stream_controller_node_id"),
|
|
"stream_overrides": run.get("stream_overrides", {}),
|
|
"filename_resolution_status": "renamed_to_final_image",
|
|
"final_output": first_image,
|
|
"last_saved_file": renamed_files[-1],
|
|
}
|
|
)
|
|
|
|
|
|
def _make_llps_callback(model, steps: int, cfg: LLPSConfigData, node_label: str, seed: int):
|
|
preview_format_for_comfy = "JPEG"
|
|
pbar = comfy.utils.ProgressBar(steps)
|
|
|
|
if not cfg.enabled or cfg.live_preview_method in {"none", "disabled"}:
|
|
previewer = None
|
|
else:
|
|
previewer = _previewer_for_method(model, cfg.live_preview_method)
|
|
|
|
save_dir = None
|
|
saved_files = []
|
|
if cfg.enabled and cfg.save_preview:
|
|
save_dir = _resolve_save_dir(cfg, node_label)
|
|
|
|
def callback(step, x0, x, total_steps):
|
|
preview_tuple = None
|
|
wants_frame = cfg.enabled and previewer is not None and (cfg.show_preview or cfg.save_preview)
|
|
|
|
if wants_frame:
|
|
try:
|
|
# In some environments the first callback-delivered preview can be a stale warm-up frame.
|
|
# We still show it live (to keep UI behavior natural), but saving logic below skips callback
|
|
# step 0 and explicitly flushes the last callback preview after sampling ends.
|
|
# Do not decode result_samples here: Comfy's previewer expects callback x0-space, while
|
|
# result_samples may already be in the final latent-output space, which can produce
|
|
# oversaturated / CFG-burnt looking TAESD previews.
|
|
preview_tuple = previewer.decode_latent_to_preview_image(preview_format_for_comfy, x0)
|
|
preview_tuple, _resolution_info = _normalize_preview_tuple(preview_tuple, cfg.preview_max_side)
|
|
callback.llps_last_preview_tuple = preview_tuple
|
|
callback.llps_last_callback_step = int(step)
|
|
except Exception as e:
|
|
logging.exception("[LLPS] Failed to decode preview at step %s: %s", step, e)
|
|
preview_tuple = None
|
|
|
|
if cfg.enabled and cfg.show_preview and preview_tuple is not None:
|
|
pbar.update_absolute(step + 1, total_steps, preview_tuple)
|
|
else:
|
|
# Keeps normal progress visible without emitting preview frames.
|
|
pbar.update_absolute(step + 1, total_steps, None)
|
|
|
|
if cfg.enabled and cfg.save_preview and preview_tuple is not None:
|
|
# Empirically, saving callback step 0 can capture a stale warm-up frame in some preview methods
|
|
# (notably TAESD on some setups). We therefore skip callback step 0 for file saving, shift
|
|
# numbering by one, and force-save the last callback preview after sampling completes.
|
|
if step > 0:
|
|
adjusted_step = step - 1
|
|
if (adjusted_step % cfg.save_every_n_steps) == 0:
|
|
try:
|
|
saved = _save_preview_image(
|
|
preview_tuple,
|
|
save_dir,
|
|
cfg,
|
|
node_label,
|
|
int(adjusted_step),
|
|
int(total_steps),
|
|
int(seed),
|
|
cfg.live_preview_method,
|
|
)
|
|
if saved:
|
|
saved_files.append(saved)
|
|
except Exception as e:
|
|
logging.exception("[LLPS] Failed to save preview at step %s: %s", step, e)
|
|
|
|
callback.llps_saved_files = saved_files
|
|
callback.llps_save_dir = save_dir
|
|
callback.llps_config = cfg
|
|
callback.llps_previewer = previewer
|
|
callback.llps_preview_format = preview_format_for_comfy
|
|
callback.llps_last_preview_tuple = None
|
|
callback.llps_last_callback_step = None
|
|
return callback
|
|
|
|
|
|
def _is_prompt_link(value: Any) -> bool:
|
|
return isinstance(value, list) and len(value) == 2 and isinstance(value[0], str)
|
|
|
|
|
|
def _resolve_prompt_value(prompt: Dict[str, Any], value: Any, seen: Optional[set] = None) -> Any:
|
|
if not _is_prompt_link(value):
|
|
return value
|
|
|
|
if seen is None:
|
|
seen = set()
|
|
node_id = str(value[0])
|
|
if node_id in seen:
|
|
return value
|
|
seen.add(node_id)
|
|
|
|
node = prompt.get(node_id)
|
|
if not isinstance(node, dict):
|
|
return value
|
|
inputs = node.get("inputs", {})
|
|
if not isinstance(inputs, dict):
|
|
return value
|
|
|
|
for key in ("value", "boolean", "bool", "enabled", "save_also", "save_preview"):
|
|
if key in inputs:
|
|
return _resolve_prompt_value(prompt, inputs[key], seen)
|
|
|
|
literal_values = [v for v in inputs.values() if not _is_prompt_link(v)]
|
|
bool_values = [v for v in literal_values if isinstance(v, bool)]
|
|
if len(bool_values) == 1:
|
|
return bool_values[0]
|
|
if len(literal_values) == 1:
|
|
return literal_values[0]
|
|
return value
|
|
|
|
|
|
def _prompt_bool(prompt: Dict[str, Any], value: Any, default: bool = False) -> bool:
|
|
resolved = _resolve_prompt_value(prompt, value)
|
|
if _is_prompt_link(resolved) or resolved is None:
|
|
return default
|
|
return _bool(resolved)
|
|
|
|
|
|
def _prompt_text(prompt: Dict[str, Any], value: Any, default: str = "") -> str:
|
|
resolved = _resolve_prompt_value(prompt, value)
|
|
if _is_prompt_link(resolved) or resolved is None:
|
|
return default
|
|
return str(resolved)
|
|
|
|
|
|
def _prompt_int(prompt: Dict[str, Any], value: Any, default: int) -> int:
|
|
resolved = _resolve_prompt_value(prompt, value)
|
|
if _is_prompt_link(resolved) or resolved is None:
|
|
return default
|
|
try:
|
|
return int(resolved)
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def _prompt_node_sort_key(node_id: Any) -> Tuple[int, Any]:
|
|
try:
|
|
return (0, int(node_id))
|
|
except Exception:
|
|
return (1, str(node_id))
|
|
|
|
|
|
def _active_controller_from_prompt(prompt: Optional[Dict[str, Any]]) -> Optional[LLPSControllerData]:
|
|
if not isinstance(prompt, dict):
|
|
return None
|
|
|
|
enabled_controllers = []
|
|
for node_id in sorted(prompt.keys(), key=_prompt_node_sort_key):
|
|
node = prompt.get(node_id)
|
|
if not isinstance(node, dict) or node.get("class_type") != "LLPSController":
|
|
continue
|
|
inputs = node.get("inputs", {})
|
|
if not isinstance(inputs, dict):
|
|
continue
|
|
|
|
enabled = _prompt_bool(prompt, inputs.get("enabled", True), True)
|
|
if not enabled:
|
|
continue
|
|
|
|
save_also_value = inputs.get("save_also", inputs.get("save_preview", False))
|
|
controller = LLPSControllerData.from_obj(
|
|
{
|
|
"enabled": enabled,
|
|
"live_preview_method": _prompt_text(prompt, inputs.get("live_preview_method", "latent2rgb"), "latent2rgb"),
|
|
"save_also": _prompt_bool(prompt, save_also_value, False),
|
|
"save_base_path": _prompt_text(prompt, inputs.get("save_base_path", ""), ""),
|
|
"subfolder": _prompt_text(prompt, inputs.get("subfolder", ""), ""),
|
|
"filename_prefix": _prompt_text(prompt, inputs.get("filename_prefix", "LLPS"), "LLPS"),
|
|
"image_format": _prompt_text(prompt, inputs.get("image_format", "jpg"), "jpg"),
|
|
"save_every_n_steps": _prompt_int(prompt, inputs.get("save_every_n_steps", 1), 1),
|
|
"save_metadata_json": _prompt_bool(prompt, inputs.get("save_metadata_json", True), True),
|
|
"jpeg_quality": _prompt_int(prompt, inputs.get("jpeg_quality", 90), 90),
|
|
"preview_max_side": _prompt_int(prompt, inputs.get("preview_max_side", 512), 512),
|
|
"controller_node_id": str(node_id),
|
|
}
|
|
)
|
|
enabled_controllers.append(controller)
|
|
|
|
if enabled_controllers:
|
|
if len(enabled_controllers) > 1:
|
|
logging.info(
|
|
"[LLPS] Multiple enabled Controllers found; using #%s and ignoring %s",
|
|
enabled_controllers[0].controller_node_id,
|
|
[item.controller_node_id for item in enabled_controllers[1:]],
|
|
)
|
|
return enabled_controllers[0]
|
|
|
|
return None
|
|
|
|
|
|
def _stream_controller_from_prompt(
|
|
prompt: Optional[Dict[str, Any]],
|
|
sampler_node_id: Optional[str],
|
|
) -> Optional[LLPSStreamControllerData]:
|
|
if not isinstance(prompt, dict) or sampler_node_id is None:
|
|
return None
|
|
|
|
matches = []
|
|
for node_id in sorted(prompt.keys(), key=_prompt_node_sort_key):
|
|
node = prompt.get(node_id)
|
|
if not isinstance(node, dict) or node.get("class_type") != "LLPSStreamController":
|
|
continue
|
|
inputs = node.get("inputs", {})
|
|
if not isinstance(inputs, dict):
|
|
continue
|
|
target_sampler_id = _prompt_text(prompt, inputs.get("target_sampler_id", ""), "")
|
|
if str(target_sampler_id).strip() != str(sampler_node_id):
|
|
continue
|
|
matches.append(
|
|
LLPSStreamControllerData.from_obj(
|
|
{
|
|
"enabled": _prompt_bool(prompt, inputs.get("enabled", True), True),
|
|
"target_sampler_id": target_sampler_id,
|
|
"override_save_also": _prompt_bool(prompt, inputs.get("override_save_also", False), False),
|
|
"save_also": _prompt_bool(prompt, inputs.get("save_also", False), False),
|
|
"override_subfolder": _prompt_bool(prompt, inputs.get("override_subfolder", False), False),
|
|
"subfolder": _prompt_text(prompt, inputs.get("subfolder", ""), ""),
|
|
"override_filename_prefix": _prompt_bool(prompt, inputs.get("override_filename_prefix", False), False),
|
|
"filename_prefix": _prompt_text(prompt, inputs.get("filename_prefix", ""), ""),
|
|
"stream_controller_node_id": str(node_id),
|
|
}
|
|
)
|
|
)
|
|
|
|
if matches:
|
|
if len(matches) > 1:
|
|
logging.info(
|
|
"[LLPS] Multiple Stream Controllers target sampler #%s; using #%s and ignoring %s",
|
|
sampler_node_id,
|
|
matches[0].stream_controller_node_id,
|
|
[item.stream_controller_node_id for item in matches[1:]],
|
|
)
|
|
return matches[0]
|
|
|
|
return None
|
|
|
|
|
|
def _controller_with_stream_overrides(
|
|
controller: LLPSControllerData,
|
|
stream: Optional[LLPSStreamControllerData],
|
|
) -> Tuple[LLPSControllerData, Dict[str, Any]]:
|
|
effective = LLPSControllerData.from_obj(asdict(controller))
|
|
info = {
|
|
"stream_controller_node_id": None,
|
|
"stream_enabled": True,
|
|
"stream_overrides": {},
|
|
"save_also_source": "global_controller",
|
|
"subfolder_source": None,
|
|
"filename_prefix_source": "global_controller",
|
|
"stream_subfolder": None,
|
|
}
|
|
if stream is None:
|
|
return effective, info
|
|
|
|
info["stream_controller_node_id"] = stream.stream_controller_node_id
|
|
info["stream_enabled"] = bool(stream.enabled)
|
|
if stream.override_save_also:
|
|
effective.save_also = bool(stream.save_also)
|
|
info["stream_overrides"]["save_also"] = bool(stream.save_also)
|
|
info["save_also_source"] = f"Stream Controller #{stream.stream_controller_node_id}"
|
|
if stream.override_filename_prefix and stream.filename_prefix:
|
|
effective.filename_prefix = stream.filename_prefix
|
|
info["stream_overrides"]["filename_prefix"] = stream.filename_prefix
|
|
info["filename_prefix_source"] = f"Stream Controller #{stream.stream_controller_node_id}"
|
|
if stream.override_subfolder and stream.subfolder:
|
|
info["stream_overrides"]["subfolder"] = stream.subfolder
|
|
info["subfolder_source"] = f"Stream Controller #{stream.stream_controller_node_id}"
|
|
info["stream_subfolder"] = stream.subfolder
|
|
return effective, info
|
|
|
|
|
|
def _write_metadata_json(save_dir: str, metadata: Dict[str, Any], filename: str = "metadata.json") -> None:
|
|
with open(os.path.join(save_dir, filename), "w", encoding="utf-8") as f:
|
|
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
def _make_controller_callback(
|
|
model,
|
|
steps: int,
|
|
controller: LLPSControllerData,
|
|
node_label: str,
|
|
seed: int,
|
|
sampler_node_id: Optional[str],
|
|
sampler_node_type: str,
|
|
sampler_subfolder: Optional[str] = None,
|
|
x0_output_dict: Optional[Dict[str, Any]] = None,
|
|
stream_info: Optional[Dict[str, Any]] = None,
|
|
):
|
|
method = controller.live_preview_method
|
|
show_preview = controller.enabled and method not in {"none", "disabled"}
|
|
save_preview = controller.enabled and controller.save_also
|
|
source_method = method if show_preview else "server_default"
|
|
if show_preview or save_preview:
|
|
previewer, previewer_info = _previewer_for_method_with_info(model, source_method, method)
|
|
else:
|
|
previewer, previewer_info = _previewer_for_method_with_info(model, "none", method)
|
|
cfg = controller.to_legacy_config(show_preview=show_preview, save_preview=save_preview)
|
|
prompt_context = get_executing_context()
|
|
prompt_id = prompt_context.prompt_id if prompt_context is not None else "unknown_prompt"
|
|
run_id = uuid.uuid4().hex[:10]
|
|
effective_subfolder = _effective_subfolder(cfg, node_label, sampler_subfolder)
|
|
stream_info = dict(stream_info or {})
|
|
if not stream_info.get("subfolder_source"):
|
|
if sampler_subfolder:
|
|
stream_info["subfolder_source"] = "sampler_llps_subfolder"
|
|
elif cfg.subfolder:
|
|
stream_info["subfolder_source"] = "global_controller"
|
|
else:
|
|
stream_info["subfolder_source"] = "automatic_sampler_fallback"
|
|
|
|
save_dir = None
|
|
saved_files = []
|
|
saved_file_entries = []
|
|
if save_preview:
|
|
save_dir = _resolve_save_dir(cfg, node_label, sampler_subfolder)
|
|
|
|
pbar = comfy.utils.ProgressBar(steps, node_id=controller.controller_node_id)
|
|
preview_format_for_comfy = "JPEG"
|
|
run_record = {
|
|
"prompt_id": str(prompt_id),
|
|
"run_id": run_id,
|
|
"controller_node_id": controller.controller_node_id,
|
|
"sampler_node_id": str(sampler_node_id) if sampler_node_id is not None else None,
|
|
"sampler_node_type": sampler_node_type,
|
|
"node_label": node_label,
|
|
"save_dir": save_dir,
|
|
"saved_files": saved_file_entries,
|
|
"renamed_files": [],
|
|
"metadata_path": None,
|
|
"metadata": None,
|
|
"filename_resolution_status": "pending_final_image" if save_preview else "not_saving",
|
|
"stream_controller_node_id": stream_info.get("stream_controller_node_id"),
|
|
"stream_overrides": stream_info.get("stream_overrides", {}),
|
|
}
|
|
if save_preview and sampler_node_id is not None:
|
|
with _LLPS_RUN_LOCK:
|
|
_LLPS_PREVIEW_RUNS.setdefault((str(prompt_id), str(sampler_node_id)), []).append(run_record)
|
|
|
|
def callback(step, x0, x, total_steps):
|
|
if x0_output_dict is not None:
|
|
x0_output_dict["x0"] = x0
|
|
|
|
preview_tuple = None
|
|
resolution_info = {}
|
|
if previewer is not None:
|
|
try:
|
|
preview_tuple = previewer.decode_latent_to_preview_image(preview_format_for_comfy, x0)
|
|
preview_tuple, resolution_info = _normalize_preview_tuple(preview_tuple, cfg.preview_max_side)
|
|
callback.llps_last_preview_tuple = preview_tuple
|
|
callback.llps_last_preview_resolution = resolution_info
|
|
if preview_tuple is not None and len(preview_tuple) > 1:
|
|
callback.llps_last_preview_image_id = id(preview_tuple[1])
|
|
callback.llps_last_callback_step = int(step)
|
|
except Exception as e:
|
|
logging.exception("[LLPS] Controller failed to decode preview at step %s: %s", step, e)
|
|
|
|
if show_preview and preview_tuple is not None:
|
|
pbar.update_absolute(step + 1, total_steps, None)
|
|
_send_llps_preview_event(
|
|
{
|
|
"controller_node_id": controller.controller_node_id,
|
|
"sampler_node_id": str(sampler_node_id) if sampler_node_id is not None else None,
|
|
"sampler_node_type": sampler_node_type,
|
|
"sampler_label": node_label,
|
|
"prompt_id": str(prompt_id),
|
|
"run_id": run_id,
|
|
"step": int(step) + 1,
|
|
"total_steps": int(total_steps),
|
|
"method": callback.llps_source_method,
|
|
"previewer_class": previewer_info.get("actual_previewer_class"),
|
|
"fallback_reason": previewer_info.get("fallback_reason"),
|
|
"save_also": bool(save_preview),
|
|
"save_base_path": _save_base_dir(cfg),
|
|
"controller_subfolder": cfg.subfolder,
|
|
"sampler_subfolder": sampler_subfolder or "",
|
|
"effective_subfolder": effective_subfolder,
|
|
"subfolder_source": stream_info.get("subfolder_source"),
|
|
"stream_controller_node_id": stream_info.get("stream_controller_node_id"),
|
|
"stream_overrides": stream_info.get("stream_overrides", {}),
|
|
"save_also_source": stream_info.get("save_also_source", "global_controller"),
|
|
"filename_prefix_source": stream_info.get("filename_prefix_source", "global_controller"),
|
|
"last_saved_file": os.path.basename(saved_files[-1]) if saved_files else None,
|
|
"image": _preview_data_url(preview_tuple[1], preview_tuple[0], cfg.jpeg_quality),
|
|
**resolution_info,
|
|
}
|
|
)
|
|
else:
|
|
pbar.update_absolute(step + 1, total_steps, None)
|
|
|
|
if save_preview and preview_tuple is not None and step > 0:
|
|
adjusted_step = step - 1
|
|
if (adjusted_step % cfg.save_every_n_steps) == 0:
|
|
try:
|
|
saved = _save_preview_image(
|
|
preview_tuple,
|
|
save_dir,
|
|
cfg,
|
|
node_label,
|
|
int(adjusted_step),
|
|
int(total_steps),
|
|
int(seed),
|
|
source_method,
|
|
temp_run_id=run_id,
|
|
sampler_node_id=str(sampler_node_id) if sampler_node_id is not None else None,
|
|
)
|
|
if saved:
|
|
saved_files.append(saved)
|
|
saved_file_entries.append(
|
|
{
|
|
"path": saved,
|
|
"filename": os.path.basename(saved),
|
|
"step": int(adjusted_step),
|
|
"total_steps": int(total_steps),
|
|
}
|
|
)
|
|
except Exception as e:
|
|
logging.exception("[LLPS] Controller failed to save preview at step %s: %s", step, e)
|
|
|
|
callback.llps_saved_files = saved_files
|
|
callback.llps_save_dir = save_dir
|
|
callback.llps_config = cfg
|
|
callback.llps_controller = controller
|
|
callback.llps_source_method = source_method
|
|
callback.llps_previewer_info = previewer_info
|
|
callback.llps_run_record = run_record
|
|
callback.llps_run_id = run_id
|
|
callback.llps_prompt_id = str(prompt_id)
|
|
callback.llps_sampler_node_id = str(sampler_node_id) if sampler_node_id is not None else None
|
|
callback.llps_sampler_node_type = sampler_node_type
|
|
callback.llps_sampler_subfolder = sampler_subfolder or ""
|
|
callback.llps_effective_subfolder = effective_subfolder
|
|
callback.llps_stream_controller_node_id = stream_info.get("stream_controller_node_id")
|
|
callback.llps_stream_overrides = stream_info.get("stream_overrides", {})
|
|
callback.llps_subfolder_source = stream_info.get("subfolder_source")
|
|
callback.llps_save_also_source = stream_info.get("save_also_source", "global_controller")
|
|
callback.llps_filename_prefix_source = stream_info.get("filename_prefix_source", "global_controller")
|
|
callback.llps_last_preview_tuple = None
|
|
callback.llps_last_preview_image_id = None
|
|
callback.llps_last_callback_step = None
|
|
callback.llps_last_preview_resolution = {}
|
|
|
|
logging.info(
|
|
"[LLPS] Controller sampler previewer: controller=%s sampler=%s run=%s requested=%s source=%s resolved=%s class=%s module=%s fallback=%s",
|
|
controller.controller_node_id,
|
|
node_label,
|
|
run_id,
|
|
previewer_info.get("requested_live_preview_method"),
|
|
previewer_info.get("requested_preview_source_method"),
|
|
previewer_info.get("resolved_preview_source_method"),
|
|
previewer_info.get("actual_previewer_class"),
|
|
previewer_info.get("actual_previewer_module"),
|
|
previewer_info.get("fallback_reason"),
|
|
)
|
|
return callback
|
|
|
|
|
|
def _flush_controller_callback(
|
|
callback,
|
|
controller: LLPSControllerData,
|
|
node_label: str,
|
|
seed: int,
|
|
steps: int,
|
|
) -> None:
|
|
if (
|
|
controller.save_also
|
|
and callback.llps_save_dir
|
|
and callback.llps_last_preview_tuple is not None
|
|
):
|
|
try:
|
|
saved = _save_preview_image(
|
|
callback.llps_last_preview_tuple,
|
|
callback.llps_save_dir,
|
|
callback.llps_config,
|
|
node_label,
|
|
int(steps) - 1,
|
|
int(steps),
|
|
int(seed),
|
|
callback.llps_source_method,
|
|
temp_run_id=callback.llps_run_id,
|
|
sampler_node_id=callback.llps_sampler_node_id,
|
|
)
|
|
if saved:
|
|
callback.llps_saved_files.append(saved)
|
|
callback.llps_run_record["saved_files"].append(
|
|
{
|
|
"path": saved,
|
|
"filename": os.path.basename(saved),
|
|
"step": int(steps) - 1,
|
|
"total_steps": int(steps),
|
|
}
|
|
)
|
|
except Exception as e:
|
|
logging.exception("[LLPS] Controller failed to flush final preview image: %s", e)
|
|
|
|
|
|
def _write_controller_metadata(
|
|
callback,
|
|
controller: LLPSControllerData,
|
|
sampler_metadata: Optional[Dict[str, Any]] = None,
|
|
) -> None:
|
|
if not (controller.save_also and callback.llps_config.save_metadata_json and callback.llps_save_dir):
|
|
return
|
|
try:
|
|
previewer_info = dict(callback.llps_previewer_info)
|
|
previewer_info["last_preview_image_id"] = callback.llps_last_preview_image_id
|
|
resolution_info = dict(callback.llps_last_preview_resolution or {})
|
|
metadata_filename = f"metadata_{callback.llps_run_id}.json"
|
|
metadata = {
|
|
"node": "LLPS Controller",
|
|
"prompt_id": callback.llps_prompt_id,
|
|
"run_id": callback.llps_run_id,
|
|
"controller_node_id": controller.controller_node_id,
|
|
"stream_controller_node_id": callback.llps_stream_controller_node_id,
|
|
"sampler_node_id": callback.llps_sampler_node_id,
|
|
"sampler_node_type": callback.llps_sampler_node_type,
|
|
"live_preview_method": controller.live_preview_method,
|
|
"preview_source_method": callback.llps_source_method,
|
|
"previewer_info": previewer_info,
|
|
"save_also": bool(controller.save_also),
|
|
"filename_prefix": callback.llps_config.filename_prefix,
|
|
"save_base_path": _save_base_dir(callback.llps_config),
|
|
"controller_subfolder": callback.llps_config.subfolder,
|
|
"sampler_subfolder": callback.llps_sampler_subfolder,
|
|
"effective_subfolder": callback.llps_effective_subfolder,
|
|
"subfolder_source": callback.llps_subfolder_source,
|
|
"stream_overrides": callback.llps_stream_overrides,
|
|
"save_also_source": callback.llps_save_also_source,
|
|
"filename_prefix_source": callback.llps_filename_prefix_source,
|
|
"save_dir": callback.llps_save_dir,
|
|
"filename_resolution_status": callback.llps_run_record.get("filename_resolution_status"),
|
|
"final_output": None,
|
|
"saved_count": len(callback.llps_saved_files),
|
|
"saved_files": [os.path.basename(p) for p in callback.llps_saved_files],
|
|
"renamed_files": [],
|
|
"created_at": time.strftime("%Y-%m-%d %H:%M:%S"),
|
|
**resolution_info,
|
|
}
|
|
if sampler_metadata:
|
|
metadata.update(sampler_metadata)
|
|
callback.llps_run_record["metadata"] = metadata
|
|
callback.llps_run_record["metadata_path"] = os.path.join(callback.llps_save_dir, metadata_filename)
|
|
_write_metadata_json(
|
|
callback.llps_save_dir,
|
|
metadata,
|
|
metadata_filename,
|
|
)
|
|
except Exception as e:
|
|
logging.exception("[LLPS] Controller failed to write metadata.json: %s", e)
|
|
|
|
|
|
_ORIGINAL_COMMON_KSAMPLER = comfy_nodes.common_ksampler
|
|
_ORIGINAL_SAVE_IMAGE_INPUTS = comfy_nodes.SaveImage.INPUT_TYPES
|
|
_ORIGINAL_SAVE_IMAGE_SAVE_IMAGES = comfy_nodes.SaveImage.save_images
|
|
|
|
|
|
def _llps_common_ksampler(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent,
|
|
denoise=1.0,
|
|
disable_noise=False,
|
|
start_step=None,
|
|
last_step=None,
|
|
force_full_denoise=False,
|
|
prompt=None,
|
|
unique_id=None,
|
|
llps_subfolder="",
|
|
):
|
|
controller = _active_controller_from_prompt(prompt)
|
|
if controller is None:
|
|
return _ORIGINAL_COMMON_KSAMPLER(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent,
|
|
denoise=denoise,
|
|
disable_noise=disable_noise,
|
|
start_step=start_step,
|
|
last_step=last_step,
|
|
force_full_denoise=force_full_denoise,
|
|
)
|
|
|
|
sampler_node_id = str(unique_id) if unique_id is not None else None
|
|
stream_controller = _stream_controller_from_prompt(prompt, sampler_node_id)
|
|
if stream_controller is not None and not stream_controller.enabled:
|
|
return _ORIGINAL_COMMON_KSAMPLER(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent,
|
|
denoise=denoise,
|
|
disable_noise=disable_noise,
|
|
start_step=start_step,
|
|
last_step=last_step,
|
|
force_full_denoise=force_full_denoise,
|
|
)
|
|
controller, stream_info = _controller_with_stream_overrides(controller, stream_controller)
|
|
effective_sampler_subfolder = stream_info.get("stream_subfolder") or llps_subfolder
|
|
|
|
latent_image = latent["samples"]
|
|
latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image, latent.get("downscale_ratio_spacial", None))
|
|
|
|
if disable_noise:
|
|
noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
|
|
else:
|
|
batch_inds = latent["batch_index"] if "batch_index" in latent else None
|
|
noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
|
|
|
|
noise_mask = latent.get("noise_mask", None)
|
|
node_type = "KSampler"
|
|
if isinstance(prompt, dict) and unique_id is not None:
|
|
node_type = prompt.get(str(unique_id), {}).get("class_type", node_type)
|
|
node_label = f"{node_type}_{unique_id or 'node'}"
|
|
callback = _make_controller_callback(
|
|
model,
|
|
int(steps),
|
|
controller,
|
|
node_label,
|
|
int(seed),
|
|
sampler_node_id,
|
|
node_type,
|
|
effective_sampler_subfolder,
|
|
stream_info=stream_info,
|
|
)
|
|
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
|
|
|
samples = comfy.sample.sample(
|
|
model,
|
|
noise,
|
|
int(steps),
|
|
float(cfg),
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise=float(denoise),
|
|
disable_noise=disable_noise,
|
|
start_step=start_step,
|
|
last_step=last_step,
|
|
force_full_denoise=force_full_denoise,
|
|
noise_mask=noise_mask,
|
|
callback=callback,
|
|
disable_pbar=disable_pbar,
|
|
seed=int(seed),
|
|
)
|
|
|
|
_flush_controller_callback(callback, controller, node_label, int(seed), int(steps))
|
|
_write_controller_metadata(
|
|
callback,
|
|
controller,
|
|
{
|
|
"seed": int(seed),
|
|
"steps": int(steps),
|
|
"cfg": float(cfg),
|
|
"sampler_name": sampler_name,
|
|
"scheduler": scheduler,
|
|
"denoise": float(denoise),
|
|
},
|
|
)
|
|
|
|
out = latent.copy()
|
|
out.pop("downscale_ratio_spacial", None)
|
|
out["samples"] = samples
|
|
return (out,)
|
|
|
|
|
|
def _with_llps_hidden(input_types: Dict[str, Any]) -> Dict[str, Any]:
|
|
patched = dict(input_types)
|
|
for section in ("required", "optional", "hidden"):
|
|
if section in patched:
|
|
patched[section] = dict(patched[section])
|
|
optional = patched.setdefault("optional", {})
|
|
optional.setdefault(
|
|
"llps_subfolder",
|
|
("STRING", {"default": "", "multiline": False, "tooltip": "Optional LLPS per-sampler preview save subfolder."}),
|
|
)
|
|
hidden = patched.setdefault("hidden", {})
|
|
hidden["prompt"] = "PROMPT"
|
|
hidden["unique_id"] = "UNIQUE_ID"
|
|
return patched
|
|
|
|
|
|
def _with_save_image_unique_id(input_types: Dict[str, Any]) -> Dict[str, Any]:
|
|
patched = dict(input_types)
|
|
for section in ("required", "optional", "hidden"):
|
|
if section in patched:
|
|
patched[section] = dict(patched[section])
|
|
hidden = patched.setdefault("hidden", {})
|
|
hidden["unique_id"] = "UNIQUE_ID"
|
|
return patched
|
|
|
|
|
|
def _with_llps_v3_hidden(schema):
|
|
hidden = list(getattr(schema, "hidden", None) or [])
|
|
for item in (comfy_custom_sampler_nodes.io.Hidden.prompt, comfy_custom_sampler_nodes.io.Hidden.unique_id):
|
|
if item not in hidden:
|
|
hidden.append(item)
|
|
schema.hidden = hidden
|
|
return schema
|
|
|
|
|
|
def _v3_hidden(cls, name: str, default: Any = None) -> Any:
|
|
hidden = getattr(cls, "hidden", None)
|
|
value = getattr(hidden, name, None) if hidden is not None else None
|
|
return default if value is None else value
|
|
|
|
|
|
def _object_label(obj: Any) -> str:
|
|
return str(getattr(obj, "name", None) or getattr(getattr(obj, "__class__", None), "__name__", None) or obj)
|
|
|
|
|
|
def _patch_builtin_samplers() -> None:
|
|
original_ksampler_inputs = comfy_nodes.KSampler.INPUT_TYPES
|
|
original_advanced_inputs = comfy_nodes.KSamplerAdvanced.INPUT_TYPES
|
|
original_custom_schema = comfy_custom_sampler_nodes.SamplerCustom.define_schema
|
|
original_custom_execute = comfy_custom_sampler_nodes.SamplerCustom.execute
|
|
original_custom_advanced_schema = comfy_custom_sampler_nodes.SamplerCustomAdvanced.define_schema
|
|
original_custom_advanced_execute = comfy_custom_sampler_nodes.SamplerCustomAdvanced.execute
|
|
|
|
@classmethod
|
|
def ksampler_input_types(cls):
|
|
return _with_llps_hidden(original_ksampler_inputs())
|
|
|
|
@classmethod
|
|
def advanced_input_types(cls):
|
|
return _with_llps_hidden(original_advanced_inputs())
|
|
|
|
def ksampler_sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, llps_subfolder="", prompt=None, unique_id=None):
|
|
return _llps_common_ksampler(
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise=denoise,
|
|
prompt=prompt,
|
|
unique_id=unique_id,
|
|
llps_subfolder=llps_subfolder,
|
|
)
|
|
|
|
def advanced_sample(
|
|
self,
|
|
model,
|
|
add_noise,
|
|
noise_seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
start_at_step,
|
|
end_at_step,
|
|
return_with_leftover_noise,
|
|
denoise=1.0,
|
|
llps_subfolder="",
|
|
prompt=None,
|
|
unique_id=None,
|
|
):
|
|
force_full_denoise = return_with_leftover_noise != "enable"
|
|
disable_noise = add_noise == "disable"
|
|
return _llps_common_ksampler(
|
|
model,
|
|
noise_seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise=denoise,
|
|
disable_noise=disable_noise,
|
|
start_step=start_at_step,
|
|
last_step=end_at_step,
|
|
force_full_denoise=force_full_denoise,
|
|
prompt=prompt,
|
|
unique_id=unique_id,
|
|
llps_subfolder=llps_subfolder,
|
|
)
|
|
|
|
@classmethod
|
|
def custom_schema(cls):
|
|
return _with_llps_v3_hidden(original_custom_schema())
|
|
|
|
@classmethod
|
|
def custom_advanced_schema(cls):
|
|
return _with_llps_v3_hidden(original_custom_advanced_schema())
|
|
|
|
@classmethod
|
|
def custom_execute(cls, model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image):
|
|
prompt = _v3_hidden(cls, "prompt")
|
|
unique_id = _v3_hidden(cls, "unique_id")
|
|
controller = _active_controller_from_prompt(prompt)
|
|
if controller is None:
|
|
return original_custom_execute(model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image)
|
|
sampler_node_id = str(unique_id) if unique_id is not None else None
|
|
stream_controller = _stream_controller_from_prompt(prompt, sampler_node_id)
|
|
if stream_controller is not None and not stream_controller.enabled:
|
|
return original_custom_execute(model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image)
|
|
controller, stream_info = _controller_with_stream_overrides(controller, stream_controller)
|
|
effective_sampler_subfolder = stream_info.get("stream_subfolder")
|
|
|
|
latent = latent_image
|
|
latent_image = latent["samples"]
|
|
latent = latent.copy()
|
|
latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image, latent.get("downscale_ratio_spacial", None))
|
|
latent["samples"] = latent_image
|
|
|
|
if not add_noise:
|
|
noise = comfy_custom_sampler_nodes.Noise_EmptyNoise().generate_noise(latent)
|
|
else:
|
|
noise = comfy_custom_sampler_nodes.Noise_RandomNoise(noise_seed).generate_noise(latent)
|
|
|
|
noise_mask = latent.get("noise_mask", None)
|
|
total_steps = max(0, int(sigmas.shape[-1]) - 1)
|
|
node_type = "SamplerCustom"
|
|
if isinstance(prompt, dict) and unique_id is not None:
|
|
node_type = prompt.get(str(unique_id), {}).get("class_type", node_type)
|
|
node_label = f"{node_type}_{unique_id or 'node'}"
|
|
x0_output = {}
|
|
callback = _make_controller_callback(
|
|
model,
|
|
total_steps,
|
|
controller,
|
|
node_label,
|
|
int(noise_seed),
|
|
sampler_node_id,
|
|
node_type,
|
|
effective_sampler_subfolder,
|
|
x0_output_dict=x0_output,
|
|
stream_info=stream_info,
|
|
)
|
|
|
|
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
|
samples = comfy.sample.sample_custom(
|
|
model,
|
|
noise,
|
|
float(cfg),
|
|
sampler,
|
|
sigmas,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
noise_mask=noise_mask,
|
|
callback=callback,
|
|
disable_pbar=disable_pbar,
|
|
seed=int(noise_seed),
|
|
)
|
|
|
|
_flush_controller_callback(callback, controller, node_label, int(noise_seed), total_steps)
|
|
_write_controller_metadata(
|
|
callback,
|
|
controller,
|
|
{
|
|
"seed": int(noise_seed),
|
|
"steps": int(total_steps),
|
|
"cfg": float(cfg),
|
|
"sampler_name": _object_label(sampler),
|
|
"scheduler": "custom_sigmas",
|
|
"sigmas_count": int(sigmas.shape[-1]),
|
|
"denoise": None,
|
|
"hook_controlled": True,
|
|
},
|
|
)
|
|
|
|
out = latent.copy()
|
|
out.pop("downscale_ratio_spacial", None)
|
|
out["samples"] = samples
|
|
if "x0" in x0_output:
|
|
x0_out = model.model.process_latent_out(x0_output["x0"].cpu())
|
|
if samples.is_nested:
|
|
latent_shapes = [x.shape for x in samples.unbind()]
|
|
x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes))
|
|
out_denoised = latent.copy()
|
|
out_denoised["samples"] = x0_out
|
|
else:
|
|
out_denoised = out
|
|
return comfy_custom_sampler_nodes.io.NodeOutput(out, out_denoised)
|
|
|
|
@classmethod
|
|
def custom_advanced_execute(cls, noise, guider, sampler, sigmas, latent_image):
|
|
prompt = _v3_hidden(cls, "prompt")
|
|
unique_id = _v3_hidden(cls, "unique_id")
|
|
controller = _active_controller_from_prompt(prompt)
|
|
if controller is None:
|
|
return original_custom_advanced_execute(noise, guider, sampler, sigmas, latent_image)
|
|
sampler_node_id = str(unique_id) if unique_id is not None else None
|
|
stream_controller = _stream_controller_from_prompt(prompt, sampler_node_id)
|
|
if stream_controller is not None and not stream_controller.enabled:
|
|
return original_custom_advanced_execute(noise, guider, sampler, sigmas, latent_image)
|
|
controller, stream_info = _controller_with_stream_overrides(controller, stream_controller)
|
|
effective_sampler_subfolder = stream_info.get("stream_subfolder")
|
|
|
|
latent = latent_image
|
|
latent_image = latent["samples"]
|
|
latent = latent.copy()
|
|
latent_image = comfy.sample.fix_empty_latent_channels(guider.model_patcher, latent_image, latent.get("downscale_ratio_spacial", None))
|
|
latent["samples"] = latent_image
|
|
|
|
noise_mask = latent.get("noise_mask", None)
|
|
total_steps = max(0, int(sigmas.shape[-1]) - 1)
|
|
seed = int(getattr(noise, "seed", 0) or 0)
|
|
node_type = "SamplerCustomAdvanced"
|
|
if isinstance(prompt, dict) and unique_id is not None:
|
|
node_type = prompt.get(str(unique_id), {}).get("class_type", node_type)
|
|
node_label = f"{node_type}_{unique_id or 'node'}"
|
|
x0_output = {}
|
|
callback = _make_controller_callback(
|
|
guider.model_patcher,
|
|
total_steps,
|
|
controller,
|
|
node_label,
|
|
seed,
|
|
sampler_node_id,
|
|
node_type,
|
|
effective_sampler_subfolder,
|
|
x0_output_dict=x0_output,
|
|
stream_info=stream_info,
|
|
)
|
|
|
|
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
|
samples = guider.sample(
|
|
noise.generate_noise(latent),
|
|
latent_image,
|
|
sampler,
|
|
sigmas,
|
|
denoise_mask=noise_mask,
|
|
callback=callback,
|
|
disable_pbar=disable_pbar,
|
|
seed=seed,
|
|
)
|
|
samples = samples.to(comfy.model_management.intermediate_device())
|
|
|
|
_flush_controller_callback(callback, controller, node_label, seed, total_steps)
|
|
_write_controller_metadata(
|
|
callback,
|
|
controller,
|
|
{
|
|
"seed": seed,
|
|
"steps": int(total_steps),
|
|
"sampler_name": _object_label(sampler),
|
|
"scheduler": "custom_sigmas",
|
|
"sigmas_count": int(sigmas.shape[-1]),
|
|
"guider_class": _object_label(guider),
|
|
"hook_controlled": True,
|
|
},
|
|
)
|
|
|
|
out = latent.copy()
|
|
out.pop("downscale_ratio_spacial", None)
|
|
out["samples"] = samples
|
|
if "x0" in x0_output:
|
|
x0_out = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())
|
|
if samples.is_nested:
|
|
latent_shapes = [x.shape for x in samples.unbind()]
|
|
x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes))
|
|
out_denoised = latent.copy()
|
|
out_denoised["samples"] = x0_out
|
|
else:
|
|
out_denoised = out
|
|
return comfy_custom_sampler_nodes.io.NodeOutput(out, out_denoised)
|
|
|
|
if not getattr(comfy_nodes.KSampler, "_llps_controller_patch", False):
|
|
comfy_nodes.KSampler.INPUT_TYPES = ksampler_input_types
|
|
comfy_nodes.KSampler.sample = ksampler_sample
|
|
comfy_nodes.KSampler._llps_controller_patch = True
|
|
if not getattr(comfy_nodes.KSamplerAdvanced, "_llps_controller_patch", False):
|
|
comfy_nodes.KSamplerAdvanced.INPUT_TYPES = advanced_input_types
|
|
comfy_nodes.KSamplerAdvanced.sample = advanced_sample
|
|
comfy_nodes.KSamplerAdvanced._llps_controller_patch = True
|
|
if not getattr(comfy_custom_sampler_nodes.SamplerCustom, "_llps_controller_patch", False):
|
|
comfy_custom_sampler_nodes.SamplerCustom.define_schema = custom_schema
|
|
comfy_custom_sampler_nodes.SamplerCustom.execute = custom_execute
|
|
comfy_custom_sampler_nodes.SamplerCustom.sample = custom_execute
|
|
comfy_custom_sampler_nodes.SamplerCustom._llps_controller_patch = True
|
|
if not getattr(comfy_custom_sampler_nodes.SamplerCustomAdvanced, "_llps_controller_patch", False):
|
|
comfy_custom_sampler_nodes.SamplerCustomAdvanced.define_schema = custom_advanced_schema
|
|
comfy_custom_sampler_nodes.SamplerCustomAdvanced.execute = custom_advanced_execute
|
|
comfy_custom_sampler_nodes.SamplerCustomAdvanced.sample = custom_advanced_execute
|
|
comfy_custom_sampler_nodes.SamplerCustomAdvanced._llps_controller_patch = True
|
|
|
|
@classmethod
|
|
def save_image_input_types(cls):
|
|
return _with_save_image_unique_id(_ORIGINAL_SAVE_IMAGE_INPUTS())
|
|
|
|
def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, unique_id=None):
|
|
result = _ORIGINAL_SAVE_IMAGE_SAVE_IMAGES(self, images, filename_prefix, prompt, extra_pnginfo)
|
|
if getattr(self, "type", "output") == "output" and isinstance(result, dict):
|
|
output_images = result.get("ui", {}).get("images", [])
|
|
context = get_executing_context()
|
|
prompt_id = context.prompt_id if context is not None else None
|
|
sampler_ids = _find_upstream_sampler_ids(prompt, str(unique_id) if unique_id is not None else None)
|
|
_rename_pending_previews_for_output(prompt_id, sampler_ids, output_images)
|
|
return result
|
|
|
|
if not getattr(comfy_nodes.SaveImage, "_llps_controller_patch", False):
|
|
comfy_nodes.SaveImage.INPUT_TYPES = save_image_input_types
|
|
comfy_nodes.SaveImage.save_images = save_images
|
|
comfy_nodes.SaveImage._llps_controller_patch = True
|
|
|
|
|
|
_patch_builtin_samplers()
|
|
|
|
|
|
|
|
class LLPSConfig:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
"live_preview_method": (["server_default", "none", "auto", "latent2rgb", "taesd"], {"default": "latent2rgb"}),
|
|
"show_preview": ("BOOLEAN", {"default": True}),
|
|
"save_preview": ("BOOLEAN", {"default": False}),
|
|
"save_base_path": ("STRING", {"default": "", "multiline": False}),
|
|
"subfolder": ("STRING", {"default": "", "multiline": False}),
|
|
"filename_prefix": ("STRING", {"default": "LLPS", "multiline": False}),
|
|
"image_format": (["jpg", "png", "webp"], {"default": "jpg"}),
|
|
"save_every_n_steps": ("INT", {"default": 1, "min": 1, "max": 1000, "step": 1}),
|
|
"save_metadata_json": ("BOOLEAN", {"default": True}),
|
|
"jpeg_quality": ("INT", {"default": 90, "min": 1, "max": 100, "step": 1}),
|
|
"preview_max_side": (["512", "1024"], {"default": "512"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LLPS_CONFIG",)
|
|
RETURN_NAMES = ("llps_config",)
|
|
FUNCTION = "make_config"
|
|
CATEGORY = LLPS_LEGACY_CATEGORY
|
|
|
|
def make_config(
|
|
self,
|
|
enabled=True,
|
|
live_preview_method="latent2rgb",
|
|
show_preview=True,
|
|
save_preview=False,
|
|
save_base_path="",
|
|
subfolder="",
|
|
filename_prefix="LLPS",
|
|
image_format="jpg",
|
|
save_every_n_steps=1,
|
|
save_metadata_json=True,
|
|
jpeg_quality=90,
|
|
preview_max_side=512,
|
|
):
|
|
cfg = LLPSConfigData(
|
|
enabled=enabled,
|
|
live_preview_method=live_preview_method,
|
|
show_preview=show_preview,
|
|
save_preview=save_preview,
|
|
save_base_path=save_base_path,
|
|
subfolder=subfolder,
|
|
filename_prefix=filename_prefix,
|
|
image_format=image_format,
|
|
save_every_n_steps=save_every_n_steps,
|
|
save_metadata_json=save_metadata_json,
|
|
jpeg_quality=jpeg_quality,
|
|
preview_max_side=preview_max_side,
|
|
)
|
|
return (asdict(LLPSConfigData.from_obj(asdict(cfg))),)
|
|
|
|
|
|
class LLPSController:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
"live_preview_method": (["server_default", "none", "auto", "latent2rgb", "taesd"], {"default": "latent2rgb"}),
|
|
"save_also": ("BOOLEAN", {"default": False}),
|
|
"save_base_path": ("STRING", {"default": "", "multiline": False}),
|
|
"subfolder": ("STRING", {"default": "", "multiline": False}),
|
|
"filename_prefix": ("STRING", {"default": "LLPS", "multiline": False}),
|
|
"image_format": (["jpg", "png", "webp"], {"default": "jpg"}),
|
|
"save_every_n_steps": ("INT", {"default": 1, "min": 1, "max": 1000, "step": 1}),
|
|
"save_metadata_json": ("BOOLEAN", {"default": True}),
|
|
"jpeg_quality": ("INT", {"default": 90, "min": 1, "max": 100, "step": 1}),
|
|
"preview_max_side": (["512", "1024"], {"default": "512"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "make_controller"
|
|
CATEGORY = LLPS_CATEGORY
|
|
|
|
def make_controller(
|
|
self,
|
|
enabled=True,
|
|
live_preview_method="latent2rgb",
|
|
save_also=False,
|
|
save_base_path="",
|
|
subfolder="",
|
|
filename_prefix="LLPS",
|
|
image_format="jpg",
|
|
save_every_n_steps=1,
|
|
save_metadata_json=True,
|
|
jpeg_quality=90,
|
|
preview_max_side=512,
|
|
):
|
|
cfg = LLPSConfigData(
|
|
enabled=enabled,
|
|
live_preview_method=live_preview_method,
|
|
show_preview=live_preview_method not in {"none", "disabled"},
|
|
save_preview=save_also,
|
|
save_base_path=save_base_path,
|
|
subfolder=subfolder,
|
|
filename_prefix=filename_prefix,
|
|
image_format=image_format,
|
|
save_every_n_steps=save_every_n_steps,
|
|
save_metadata_json=save_metadata_json,
|
|
jpeg_quality=jpeg_quality,
|
|
preview_max_side=preview_max_side,
|
|
)
|
|
controller = asdict(LLPSConfigData.from_obj(asdict(cfg)))
|
|
controller["node"] = "LLPS Controller"
|
|
return ()
|
|
|
|
|
|
class LLPSStreamController:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
"target_sampler_id": ("STRING", {"default": "", "multiline": False, "tooltip": "Sampler node id this Stream Controller applies to."}),
|
|
"override_save_also": ("BOOLEAN", {"default": False}),
|
|
"save_also": ("BOOLEAN", {"default": False}),
|
|
"override_subfolder": ("BOOLEAN", {"default": False}),
|
|
"subfolder": ("STRING", {"default": "", "multiline": False}),
|
|
"override_filename_prefix": ("BOOLEAN", {"default": False}),
|
|
"filename_prefix": ("STRING", {"default": "", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "make_stream_controller"
|
|
CATEGORY = LLPS_CATEGORY
|
|
|
|
def make_stream_controller(
|
|
self,
|
|
enabled=True,
|
|
target_sampler_id="",
|
|
override_save_also=False,
|
|
save_also=False,
|
|
override_subfolder=False,
|
|
subfolder="",
|
|
override_filename_prefix=False,
|
|
filename_prefix="",
|
|
):
|
|
_stream = LLPSStreamControllerData.from_obj(
|
|
{
|
|
"enabled": enabled,
|
|
"target_sampler_id": target_sampler_id,
|
|
"override_save_also": override_save_also,
|
|
"save_also": save_also,
|
|
"override_subfolder": override_subfolder,
|
|
"subfolder": subfolder,
|
|
"override_filename_prefix": override_filename_prefix,
|
|
"filename_prefix": filename_prefix,
|
|
}
|
|
)
|
|
return ()
|
|
|
|
|
|
class LLPSKSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent_image": ("LATENT",),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"llps_config": ("LLPS_CONFIG",),
|
|
},
|
|
"optional": {
|
|
"node_label": ("STRING", {"default": "LLPS_KSampler", "multiline": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("latent",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = LLPS_LEGACY_CATEGORY
|
|
|
|
def sample(
|
|
self,
|
|
model,
|
|
seed,
|
|
steps,
|
|
cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent_image,
|
|
denoise=1.0,
|
|
llps_config=None,
|
|
node_label="LLPS_KSampler",
|
|
):
|
|
cfg_data = LLPSConfigData.from_obj(llps_config)
|
|
samples_tensor = latent_image["samples"]
|
|
batch_inds = latent_image.get("batch_index", None)
|
|
|
|
noise = comfy.sample.prepare_noise(samples_tensor, seed, batch_inds)
|
|
callback = _make_llps_callback(model, int(steps), cfg_data, node_label, int(seed))
|
|
|
|
result_samples = comfy.sample.sample(
|
|
model,
|
|
noise,
|
|
int(steps),
|
|
float(cfg),
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
samples_tensor,
|
|
denoise=float(denoise),
|
|
disable_noise=False,
|
|
start_step=None,
|
|
last_step=None,
|
|
force_full_denoise=True,
|
|
noise_mask=latent_image.get("noise_mask", None),
|
|
callback=callback,
|
|
disable_pbar=False,
|
|
seed=int(seed),
|
|
)
|
|
|
|
# Explicitly flush the last callback preview as the final saved frame.
|
|
# Important: do NOT decode result_samples with the previewer here. Comfy's previewer is built
|
|
# for the sampler callback's x0 tensor. Feeding the final returned latent can place the tensor
|
|
# in the wrong latent space for TAESD/latent2rgb and creates oversaturated final previews.
|
|
if (
|
|
cfg_data.enabled
|
|
and cfg_data.save_preview
|
|
and callback.llps_save_dir
|
|
and callback.llps_last_preview_tuple is not None
|
|
):
|
|
try:
|
|
saved = _save_preview_image(
|
|
callback.llps_last_preview_tuple,
|
|
callback.llps_save_dir,
|
|
cfg_data,
|
|
node_label,
|
|
int(steps) - 1,
|
|
int(steps),
|
|
int(seed),
|
|
cfg_data.live_preview_method,
|
|
)
|
|
if saved:
|
|
callback.llps_saved_files.append(saved)
|
|
except Exception as e:
|
|
logging.exception("[LLPS] Failed to flush final callback preview image: %s", e)
|
|
|
|
out = latent_image.copy()
|
|
out["samples"] = result_samples
|
|
|
|
if cfg_data.save_preview and cfg_data.save_metadata_json and callback.llps_save_dir:
|
|
try:
|
|
metadata = {
|
|
"node": "LLPS KSampler",
|
|
"node_label": node_label,
|
|
"seed": int(seed),
|
|
"steps": int(steps),
|
|
"cfg": float(cfg),
|
|
"sampler_name": sampler_name,
|
|
"scheduler": scheduler,
|
|
"denoise": float(denoise),
|
|
"llps_config": asdict(cfg_data),
|
|
"saved_count": len(callback.llps_saved_files),
|
|
"saved_files": [os.path.basename(p) for p in callback.llps_saved_files],
|
|
"created_at": time.strftime("%Y-%m-%d %H:%M:%S"),
|
|
}
|
|
with open(os.path.join(callback.llps_save_dir, "metadata.json"), "w", encoding="utf-8") as f:
|
|
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
|
except Exception as e:
|
|
logging.exception("[LLPS] Failed to write metadata.json: %s", e)
|
|
|
|
return (out,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LLPSConfig": LLPSConfig,
|
|
"LLPSController": LLPSController,
|
|
"LLPSStreamController": LLPSStreamController,
|
|
"LLPSKSampler": LLPSKSampler,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LLPSConfig": "LLPS Config (Legacy)",
|
|
"LLPSController": "LLPS Controller",
|
|
"LLPSStreamController": "LLPS Stream Controller",
|
|
"LLPSKSampler": "LLPS KSampler (Legacy)",
|
|
}
|
|
|
|
WEB_DIRECTORY = "./web/js"
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|