917 lines
40 KiB
Python
917 lines
40 KiB
Python
from torch import Tensor
|
|
from PIL import Image, ImageOps, ImageDraw, ImageFont
|
|
from PIL.PngImagePlugin import PngInfo
|
|
import numpy as np
|
|
import torch
|
|
|
|
from datetime import datetime
|
|
from itertools import chain
|
|
from comfy.cli_args import args
|
|
import folder_paths
|
|
import comfy.model_management
|
|
|
|
import ast
|
|
from pathlib import Path
|
|
from importlib import import_module
|
|
import os
|
|
import sys
|
|
import copy
|
|
import subprocess
|
|
import json
|
|
import psutil
|
|
import hashlib
|
|
|
|
import re
|
|
import nodes as native
|
|
from .adv_encode import advanced_encode #, advanced_encode_XL
|
|
|
|
import comfy.sd
|
|
from comfy_extras.nodes_align_your_steps import AlignYourStepsScheduler
|
|
from comfy_extras.nodes_gits import GITSScheduler
|
|
import comfy_extras.nodes_model_advanced
|
|
|
|
SUPPORTED_FORMATS = [".png", ".jpg", ".jpeg", ".webp"]
|
|
|
|
# Get the absolute path of various directories
|
|
my_dir = os.path.dirname(os.path.abspath(__file__))
|
|
custom_nodes_dir = os.path.abspath(os.path.join(my_dir, '..'))
|
|
comfy_dir = os.path.abspath(os.path.join(my_dir, '..', '..'))
|
|
|
|
# Construct the path to the font file
|
|
font_path = os.path.join(my_dir, 'arial.ttf')
|
|
|
|
# Append comfy_dir to sys.path & import files
|
|
sys.path.append(comfy_dir)
|
|
from nodes import LatentUpscaleBy, KSampler, KSamplerAdvanced, VAEDecode, VAEDecodeTiled, VAEEncode, VAEEncodeTiled, \
|
|
ImageScaleBy, CLIPSetLastLayer, CLIPTextEncode, ControlNetLoader, ControlNetApply, ControlNetApplyAdvanced, \
|
|
PreviewImage, MAX_RESOLUTION
|
|
from comfy_extras.nodes_upscale_model import UpscaleModelLoader, ImageUpscaleWithModel
|
|
from comfy_extras.nodes_clip_sdxl import CLIPTextEncodeSDXL, CLIPTextEncodeSDXLRefiner
|
|
import comfy.sample
|
|
import comfy.samplers
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import comfy.latent_formats
|
|
sys.path.remove(comfy_dir)
|
|
|
|
from comfy import samplers
|
|
# Append custom_nodes_dir to sys.path
|
|
sys.path.append(custom_nodes_dir)
|
|
|
|
|
|
class LoraTagLoader:
|
|
def __init__(self):
|
|
self.loaded_loras = {} # Changed from single loaded_lora to dictionary
|
|
self.tag_pattern = "\<[0-9a-zA-Z\:\_\-\.\s\/\(\)\\\\]+\>"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "model": ("MODEL",),
|
|
"clip": ("CLIP", ),
|
|
"text": ("STRING", {"multiline": True}),
|
|
}}
|
|
RETURN_TYPES = ("MODEL", "CLIP", "STRING")
|
|
RETURN_NAMES = ("MODEL", "CLIP", "STRING")
|
|
FUNCTION = "load_lora"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def cleanup(self):
|
|
"""Clean up all loaded LoRA weights"""
|
|
print("Performing LoRA cleanup")
|
|
self.loaded_loras = {}
|
|
comfy.model_management.soft_empty_cache()
|
|
|
|
def load_lora(self, model, clip, text):
|
|
founds = re.findall(self.tag_pattern, text)
|
|
|
|
if len(founds) < 1:
|
|
return (model, clip, text)
|
|
|
|
model_lora = model
|
|
clip_lora = clip
|
|
|
|
# Track which LoRAs we're using this run
|
|
current_loras = set()
|
|
lora_files = folder_paths.get_filename_list("loras")
|
|
|
|
for f in founds:
|
|
tag = f[1:-1]
|
|
pak = tag.split(":")
|
|
type = pak[0]
|
|
if type != 'lora':
|
|
continue
|
|
|
|
name = None
|
|
if len(pak) > 1 and len(pak[1]) > 0:
|
|
name = pak[1]
|
|
else:
|
|
continue
|
|
|
|
wModel = wClip = 0
|
|
try:
|
|
if len(pak) > 2 and len(pak[2]) > 0:
|
|
wModel = float(pak[2])
|
|
wClip = wModel
|
|
if len(pak) > 3 and len(pak[3]) > 0:
|
|
wClip = float(pak[3])
|
|
except ValueError:
|
|
continue
|
|
|
|
if name is None:
|
|
continue
|
|
|
|
lora_name = None
|
|
for lora_file in lora_files:
|
|
if Path(lora_file).name.startswith(name) or lora_file.startswith(name):
|
|
lora_name = lora_file
|
|
break
|
|
|
|
if lora_name is None:
|
|
print(f"bypassed lora tag: { (type, name, wModel, wClip) } >> { lora_name }")
|
|
continue
|
|
|
|
print(f"detected lora tag: { (type, name, wModel, wClip) } >> { lora_name }")
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
current_loras.add(lora_path)
|
|
|
|
# Check if we already have this LoRA loaded
|
|
lora = self.loaded_loras.get(lora_path)
|
|
|
|
if lora is None:
|
|
# Load new LoRA
|
|
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
self.loaded_loras[lora_path] = lora
|
|
|
|
# Apply LoRA to model
|
|
model_lora, clip_lora = comfy.sd.load_lora_for_models(model_lora, clip_lora, lora, wModel, wClip)
|
|
|
|
# Clean up any LoRAs not used in this run
|
|
to_remove = [path for path in self.loaded_loras if path not in current_loras]
|
|
for path in to_remove:
|
|
print(f"Purging unused LoRA: {path}")
|
|
del self.loaded_loras[path]
|
|
|
|
plain_prompt = re.sub(self.tag_pattern, "", text)
|
|
return (model_lora, clip_lora, plain_prompt)
|
|
|
|
class AdvancedCLIPTextEncodeWithBreak:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"text": ("STRING", {"multiline": True}),
|
|
"clip": ("CLIP", ),
|
|
"token_normalization": (["none", "mean", "length", "length+mean"],),
|
|
"weight_interpretation": (["comfy", "A1111", "compel", "comfy++" ,"down_weight"],),
|
|
}}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "encode"
|
|
|
|
CATEGORY = "conditioning/advanced"
|
|
|
|
def _encode(self, clip, text, token_normalization, weight_interpretation):
|
|
embeddings_final, pooled = advanced_encode(clip, text, token_normalization, weight_interpretation, w_max=1.0)
|
|
return ([[embeddings_final, {"pooled_output": pooled}]], )
|
|
|
|
def encode(self, clip, text, token_normalization, weight_interpretation):
|
|
prompts = re.split(r"\s*\bBREAK\b\s*", text)
|
|
# encode first prompt fragment
|
|
prompt = prompts.pop(0)
|
|
# print(f"prompt: {prompt}")
|
|
out = self._encode(clip, prompt, token_normalization, weight_interpretation)
|
|
# encode and concatenate the rest of the prompt
|
|
for prompt in prompts:
|
|
# print(f"prompt: {prompt}")
|
|
cond_to = self._encode(clip, prompt, token_normalization, weight_interpretation)
|
|
out = native.ConditioningConcat.concat(self, cond_to[0], out[0])
|
|
return out
|
|
|
|
class MultiModelCheckpointIterator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"used_model_count": ("INT",{"default": 2, "min": 1, "max": 3},),
|
|
"ckpt_name1": (folder_paths.get_filename_list("checkpoints"),),
|
|
"ckpt_name2_list": ("STRING", {"multiline": True, "default": ""}),
|
|
"ckpt_name3": (folder_paths.get_filename_list("checkpoints"),),
|
|
"positive": ("STRING", {"multiline": False}),
|
|
"negative": ("STRING", {"multiline": False}),
|
|
"lora_name": ("STRING", {"default": ""}),
|
|
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"rescaled_steps": ("INT",{"default": 8, "min": 0, "max": 100}),
|
|
"rescale_multiplier": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1, "step": 0.05}),
|
|
"total_steps_original": ("INT",{"default": 25, "min": 1, "max": 100}),
|
|
"total_steps_shift_second": ("INT",{"default": 0, "min": -50, "max": 100}),
|
|
"total_steps_shift_third": ("INT",{"default": 0, "min": -50, "max": 100}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
|
|
"starting_cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1}),
|
|
"cfg_shift": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 100.0, "step": 0.1}),
|
|
"steps_end_first": ("INT", {"default": 15, "min": 0, "max": 100}),
|
|
"steps_shift_second": ("INT", {"default": 0, "min": -20, "max": 100}),
|
|
"steps_end_second": ("INT", {"default": 0, "min": 0, "max": 100}),
|
|
"steps_shift_third": ("INT", {"default": 0, "min": -20, "max": 100}),
|
|
"token_normalization": (["none", "mean", "length", "length+mean"],),
|
|
"weight_interpretation": (["comfy", "A1111", "compel", "comfy++" ,"down_weight"],),
|
|
"detached_seed": ("BOOLEAN", {"default": False}),
|
|
"detached_checkpoint": ("BOOLEAN", {"default": False}),
|
|
"reset_rng": ("BOOLEAN", {"default": False}), # New toggle
|
|
"random_checkpoint": ("BOOLEAN", {"default": False}), # New toggle
|
|
"latent_image": ("LATENT",),
|
|
},
|
|
"optional": {
|
|
"load_settings": ("STRING", {"multiline": True, "default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "STRING")
|
|
RETURN_NAMES = ("LATENT", "settings")
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling"
|
|
|
|
def __init__(self):
|
|
self.loaded_checkpoints = {} # Dictionary to store loaded checkpoints by name
|
|
self.current_checkpoints = set() # Track currently requested checkpoints
|
|
self.ckpt_name2_index = 0 # Index for cycling through second model checkpoints
|
|
self.ckpt_name2_list = [] # List of second model checkpoints
|
|
|
|
def cleanup(self):
|
|
"""Force cleanup of all resources"""
|
|
print("Performing cleanup of loaded checkpoints")
|
|
self.loaded_checkpoints = {}
|
|
self.current_checkpoints = set()
|
|
comfy.model_management.soft_empty_cache()
|
|
|
|
def serialize_settings(self, **kwargs):
|
|
"""Serialize all input parameters into a human-readable string"""
|
|
used_model_count = kwargs.get("used_model_count", 2)
|
|
|
|
settings = {
|
|
"used_model_count": used_model_count,
|
|
"ckpt_name1": kwargs.get("ckpt_name1", ""),
|
|
"ckpt_name2": kwargs.get("ckpt_name2", ""), # Store the currently used ckpt_name2
|
|
"ckpt_name2_list": kwargs.get("ckpt_name2_list", ""), # Store the full list
|
|
"ckpt_name3": kwargs.get("ckpt_name3", ""),
|
|
"positive": kwargs.get("positive", ""),
|
|
"negative": kwargs.get("negative", ""),
|
|
"lora_name": kwargs.get("lora_name", ""),
|
|
"noise_seed": kwargs.get("noise_seed", 0),
|
|
"rescaled_steps": kwargs.get("rescaled_steps", 8),
|
|
"rescale_multiplier": kwargs.get("rescale_multiplier", 0.7),
|
|
"total_steps_original": kwargs.get("total_steps_original", 25),
|
|
"total_steps_shift_second": kwargs.get("total_steps_shift_second", 0),
|
|
"total_steps_shift_third": kwargs.get("total_steps_shift_third", 0),
|
|
"sampler_name": kwargs.get("sampler_name", "euler"),
|
|
"scheduler": kwargs.get("scheduler", "normal"),
|
|
"starting_cfg": kwargs.get("starting_cfg", 8.0),
|
|
"cfg_shift": kwargs.get("cfg_shift", 0.0),
|
|
"steps_end_first": kwargs.get("steps_end_first", 15),
|
|
"steps_shift_second": kwargs.get("steps_shift_second", 0),
|
|
"steps_end_second": kwargs.get("steps_end_second", 0),
|
|
"steps_shift_third": kwargs.get("steps_shift_third", 0),
|
|
"token_normalization": kwargs.get("token_normalization", "none"),
|
|
"weight_interpretation": kwargs.get("weight_interpretation", "comfy"),
|
|
"detached_seed": kwargs.get("detached_seed", False),
|
|
"detached_checkpoint": kwargs.get("detached_checkpoint", False),
|
|
"reset_rng": kwargs.get("reset_rng", False), # Include new toggle
|
|
"random_checkpoint": kwargs.get("random_checkpoint", False), # Include new toggle
|
|
}
|
|
|
|
return json.dumps(settings, indent=2)
|
|
|
|
def deserialize_settings(self, settings_str):
|
|
"""Deserialize settings string back to a dictionary"""
|
|
if not settings_str.strip():
|
|
return None
|
|
|
|
try:
|
|
return json.loads(settings_str)
|
|
except json.JSONDecodeError as e:
|
|
print(f"Error decoding settings: {e}")
|
|
return None
|
|
|
|
def get_effective_settings(self, **kwargs):
|
|
"""Get the effective settings, either from load_settings or direct inputs"""
|
|
load_settings = kwargs.get("load_settings", "")
|
|
detached_seed = kwargs.get("detached_seed", False)
|
|
detached_checkpoint = kwargs.get("detached_checkpoint", False)
|
|
reset_rng = kwargs.get("reset_rng", False)
|
|
random_checkpoint = kwargs.get("random_checkpoint", False)
|
|
|
|
if load_settings and load_settings.strip():
|
|
loaded = self.deserialize_settings(load_settings)
|
|
if loaded:
|
|
# Create a new kwargs with loaded settings, but keep original latent_image
|
|
effective_kwargs = loaded.copy()
|
|
effective_kwargs["latent_image"] = kwargs["latent_image"]
|
|
|
|
# If detached_seed is True, override the seed from direct input
|
|
if detached_seed:
|
|
effective_kwargs["noise_seed"] = kwargs["noise_seed"]
|
|
|
|
ckpt_name2_list = kwargs.get("ckpt_name2_list", "")
|
|
if ckpt_name2_list:
|
|
# Split the list by line breaks and clean up each entry
|
|
self.ckpt_name2_list = [name.strip() for name in ckpt_name2_list.split('\n') if name.strip()]
|
|
|
|
# Handle random checkpoint selection
|
|
if self.ckpt_name2_list and detached_checkpoint:
|
|
if reset_rng:
|
|
self.ckpt_name2_index = 0 # Reset counter if requested
|
|
|
|
if random_checkpoint:
|
|
# Select a random checkpoint from the list
|
|
import random
|
|
effective_kwargs["ckpt_name2"] = random.choice(self.ckpt_name2_list)
|
|
else:
|
|
# Sequential selection
|
|
effective_kwargs["ckpt_name2"] = self.ckpt_name2_list[self.ckpt_name2_index % len(self.ckpt_name2_list)]
|
|
self.ckpt_name2_index += 1 # Increment for next run
|
|
|
|
return effective_kwargs
|
|
else:
|
|
# Process the ckpt_name2_list when not loading settings
|
|
ckpt_name2_list = kwargs.get("ckpt_name2_list", "")
|
|
if ckpt_name2_list:
|
|
# Split the list by line breaks and clean up each entry
|
|
self.ckpt_name2_list = [name.strip() for name in ckpt_name2_list.split('\n') if name.strip()]
|
|
|
|
# Handle random checkpoint selection
|
|
if self.ckpt_name2_list:
|
|
if reset_rng:
|
|
self.ckpt_name2_index = 0 # Reset counter if requested
|
|
|
|
if random_checkpoint:
|
|
# Select a random checkpoint from the list
|
|
import random
|
|
kwargs["ckpt_name2"] = random.choice(self.ckpt_name2_list)
|
|
else:
|
|
# Sequential selection
|
|
kwargs["ckpt_name2"] = self.ckpt_name2_list[self.ckpt_name2_index % len(self.ckpt_name2_list)]
|
|
self.ckpt_name2_index += 1 # Increment for next run
|
|
else:
|
|
kwargs["ckpt_name2"] = ""
|
|
|
|
return kwargs
|
|
|
|
def purge_unused_checkpoints(self, requested_checkpoints):
|
|
"""Purges checkpoints that aren't in the current request"""
|
|
# Convert to set for faster lookups
|
|
requested_set = set(requested_checkpoints)
|
|
|
|
# Find checkpoints to remove
|
|
to_remove = [name for name in self.loaded_checkpoints
|
|
if name not in requested_set]
|
|
|
|
# Purge unused checkpoints
|
|
for name in to_remove:
|
|
print(f"Purging unused checkpoint: {name}")
|
|
del self.loaded_checkpoints[name]
|
|
|
|
# Update current checkpoints
|
|
self.current_checkpoints = requested_set
|
|
|
|
def load_model_pipeline(self, ckpt_name, lora_name, positive, negative,
|
|
token_normalization, weight_interpretation):
|
|
"""Helper function to load a complete model pipeline"""
|
|
# Check if we already have this checkpoint loaded
|
|
if ckpt_name in self.loaded_checkpoints:
|
|
model, clip, vae = self.loaded_checkpoints[ckpt_name]
|
|
print(f"Using cached checkpoint: {ckpt_name}")
|
|
else:
|
|
# Load base checkpoint if not found in cache
|
|
print(f"Loading new checkpoint: {ckpt_name}")
|
|
out = comfy.sd.load_checkpoint_guess_config(
|
|
folder_paths.get_full_path("checkpoints", ckpt_name),
|
|
output_vae=True,
|
|
output_clip=True,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings")
|
|
)
|
|
model, clip, vae = out[0], out[1], out[2]
|
|
self.loaded_checkpoints[ckpt_name] = (model, clip, vae)
|
|
|
|
# Apply LoRA if specified
|
|
if lora_name and lora_name.strip():
|
|
model, clip, _ = LoraTagLoader().load_lora(model, clip, lora_name)
|
|
|
|
# Generate conditioning
|
|
positive_conditioning = AdvancedCLIPTextEncodeWithBreak().encode(
|
|
clip, positive, token_normalization, weight_interpretation
|
|
)
|
|
|
|
negative_conditioning = AdvancedCLIPTextEncodeWithBreak().encode(
|
|
clip, negative, token_normalization, weight_interpretation
|
|
)
|
|
|
|
return {
|
|
'model': model,
|
|
'clip': clip,
|
|
'vae': vae,
|
|
'positive_conditioning': positive_conditioning[0],
|
|
'negative_conditioning': negative_conditioning[0]
|
|
}
|
|
|
|
def sample(self, **kwargs):
|
|
try:
|
|
# Get effective settings (either from load_settings or direct inputs)
|
|
effective_kwargs = self.get_effective_settings(**kwargs)
|
|
if effective_kwargs is None:
|
|
effective_kwargs = kwargs
|
|
|
|
load_settings = kwargs.get("load_settings", "")
|
|
|
|
# Initialize UI update dictionary
|
|
ui_updates = {}
|
|
|
|
# If we loaded settings from the string, populate ui_updates
|
|
if load_settings and load_settings.strip():
|
|
loaded_settings = self.deserialize_settings(load_settings)
|
|
if loaded_settings:
|
|
ui_updates = {
|
|
"used_model_count": (loaded_settings["used_model_count"],),
|
|
"ckpt_name1": (loaded_settings["ckpt_name1"],),
|
|
"ckpt_name2": (loaded_settings.get("ckpt_name2", ""),),
|
|
"ckpt_name2_list": (loaded_settings.get("ckpt_name2_list", ""),),
|
|
"ckpt_name3": (loaded_settings.get("ckpt_name3", ""),),
|
|
"positive": (loaded_settings["positive"],),
|
|
"negative": (loaded_settings["negative"],),
|
|
"lora_name": (loaded_settings["lora_name"],),
|
|
"noise_seed": (loaded_settings["noise_seed"],),
|
|
"rescaled_steps": (loaded_settings["rescaled_steps"],),
|
|
"rescale_multiplier": (loaded_settings["rescale_multiplier"],),
|
|
"total_steps_original": (loaded_settings["total_steps_original"],),
|
|
"total_steps_shift_second": (loaded_settings.get("total_steps_shift_second", 0),),
|
|
"total_steps_shift_third": (loaded_settings.get("total_steps_shift_third", 0),),
|
|
"sampler_name": (loaded_settings["sampler_name"],),
|
|
"scheduler": (loaded_settings["scheduler"],),
|
|
"starting_cfg": (loaded_settings["starting_cfg"],),
|
|
"cfg_shift": (loaded_settings["cfg_shift"],),
|
|
"steps_end_first": (loaded_settings["steps_end_first"],),
|
|
"steps_shift_second": (loaded_settings.get("steps_shift_second", 0),),
|
|
"steps_end_second": (loaded_settings.get("steps_end_second", 0),),
|
|
"steps_shift_third": (loaded_settings.get("steps_shift_third", 0),),
|
|
"token_normalization": (loaded_settings["token_normalization"],),
|
|
"weight_interpretation": (loaded_settings["weight_interpretation"],),
|
|
"detached_seed": (loaded_settings.get("detached_seed", False),),
|
|
"detached_checkpoint": (loaded_settings.get("detached_checkpoint", False),),
|
|
"reset_rng": (loaded_settings.get("reset_rng", False),), # Update UI
|
|
"random_checkpoint": (loaded_settings.get("random_checkpoint", False),), # Update UI
|
|
}
|
|
|
|
# Extract all parameters from effective_kwargs
|
|
used_model_count = effective_kwargs["used_model_count"]
|
|
ckpt_name1 = effective_kwargs["ckpt_name1"]
|
|
ckpt_name2 = effective_kwargs.get("ckpt_name2", "")
|
|
ckpt_name3 = effective_kwargs.get("ckpt_name3", "")
|
|
positive = effective_kwargs["positive"]
|
|
negative = effective_kwargs["negative"]
|
|
lora_name = effective_kwargs["lora_name"]
|
|
noise_seed = effective_kwargs["noise_seed"]
|
|
rescaled_steps = effective_kwargs["rescaled_steps"]
|
|
rescale_multiplier = effective_kwargs["rescale_multiplier"]
|
|
total_steps_original = effective_kwargs["total_steps_original"]
|
|
total_steps_shift_second = effective_kwargs.get("total_steps_shift_second", 0)
|
|
total_steps_shift_third = effective_kwargs.get("total_steps_shift_third", 0)
|
|
sampler_name = effective_kwargs["sampler_name"]
|
|
scheduler = effective_kwargs["scheduler"]
|
|
starting_cfg = effective_kwargs["starting_cfg"]
|
|
cfg_shift = effective_kwargs["cfg_shift"]
|
|
steps_end_first = effective_kwargs["steps_end_first"]
|
|
steps_shift_second = effective_kwargs.get("steps_shift_second", 0)
|
|
steps_end_second = effective_kwargs.get("steps_end_second", 0)
|
|
steps_shift_third = effective_kwargs.get("steps_shift_third", 0)
|
|
token_normalization = effective_kwargs["token_normalization"]
|
|
weight_interpretation = effective_kwargs["weight_interpretation"]
|
|
latent_image = effective_kwargs["latent_image"]
|
|
|
|
# Determine which checkpoints are being requested
|
|
requested_checkpoints = [ckpt_name1]
|
|
if used_model_count >= 2 and ckpt_name2:
|
|
requested_checkpoints.append(ckpt_name2)
|
|
if used_model_count >= 3 and ckpt_name3:
|
|
requested_checkpoints.append(ckpt_name3)
|
|
|
|
# Purge checkpoints that aren't being used this run
|
|
self.purge_unused_checkpoints(requested_checkpoints)
|
|
|
|
# Load all required pipelines
|
|
pipelines = []
|
|
try:
|
|
pipelines.append(self.load_model_pipeline(ckpt_name1, lora_name, positive, negative,
|
|
token_normalization, weight_interpretation))
|
|
if comfy.model_management.processing_interrupted():
|
|
raise RuntimeError("Processing interrupted by user")
|
|
|
|
if used_model_count >= 2 and ckpt_name2:
|
|
pipelines.append(self.load_model_pipeline(ckpt_name2, lora_name, positive, negative,
|
|
token_normalization, weight_interpretation))
|
|
if comfy.model_management.processing_interrupted():
|
|
raise RuntimeError("Processing interrupted by user")
|
|
|
|
if used_model_count >= 3 and ckpt_name3:
|
|
pipelines.append(self.load_model_pipeline(ckpt_name3, lora_name, positive, negative,
|
|
token_normalization, weight_interpretation))
|
|
if comfy.model_management.processing_interrupted():
|
|
raise RuntimeError("Processing interrupted by user")
|
|
except Exception as e:
|
|
self.cleanup()
|
|
raise e
|
|
|
|
samples = None
|
|
try:
|
|
# First sampling pass
|
|
if rescaled_steps >= 0: # Changed from > 0 to >= 0 to handle 0 case
|
|
try:
|
|
if rescaled_steps > 0: # Only apply RescaleCFG if steps > 0
|
|
CFGR = comfy_extras.nodes_model_advanced.RescaleCFG()
|
|
patched_model = CFGR.patch(pipelines[0]['model'], rescale_multiplier)[0]
|
|
print("Attempting rescaled CFG pass")
|
|
else:
|
|
patched_model = pipelines[0]['model']
|
|
print("Skipping rescaled CFG pass (0 steps)")
|
|
|
|
samples = KSamplerAdvanced().sample(
|
|
patched_model,
|
|
"enable", # add_noise
|
|
noise_seed,
|
|
total_steps_original,
|
|
starting_cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
pipelines[0]['positive_conditioning'],
|
|
pipelines[0]['negative_conditioning'],
|
|
latent_image,
|
|
0,
|
|
rescaled_steps,
|
|
"enable" # return_with_leftover_noise
|
|
)
|
|
print(f"Rescaled CFG pass completed, latent shape: {samples[0]['samples'].shape}")
|
|
except Exception as e:
|
|
print(f"Error in rescaled CFG pass: {str(e)}")
|
|
raise e
|
|
|
|
if comfy.model_management.processing_interrupted():
|
|
raise RuntimeError("Processing interrupted by user")
|
|
|
|
# First model's main pass
|
|
try:
|
|
actual_steps = rescaled_steps if rescaled_steps >= 0 else 0 # Changed condition
|
|
samples = KSamplerAdvanced().sample(
|
|
pipelines[0]['model'],
|
|
"disable" if rescaled_steps >= 0 else "enable", # Changed condition
|
|
noise_seed,
|
|
total_steps_original,
|
|
starting_cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
pipelines[0]['positive_conditioning'],
|
|
pipelines[0]['negative_conditioning'],
|
|
samples[0] if rescaled_steps >= 0 else latent_image, # Changed condition
|
|
actual_steps,
|
|
steps_end_first,
|
|
"enable" # return_with_leftover_noise
|
|
)
|
|
print(f"First model pass completed, latent shape: {samples[0]['samples'].shape}")
|
|
except Exception as e:
|
|
print(f"Error in first model pass: {str(e)}")
|
|
raise e
|
|
|
|
if comfy.model_management.processing_interrupted():
|
|
raise RuntimeError("Processing interrupted by user")
|
|
|
|
# Second model pass if enabled
|
|
if used_model_count >= 2 and ckpt_name2:
|
|
try:
|
|
end_steps = steps_end_second if steps_end_second > 0 else total_steps_original + total_steps_shift_second
|
|
new_cfg = starting_cfg + cfg_shift
|
|
model2_start_step = steps_end_first + steps_shift_second
|
|
new_total_steps = total_steps_original + total_steps_shift_second
|
|
|
|
samples = KSamplerAdvanced().sample(
|
|
pipelines[1]['model'],
|
|
"disable",
|
|
noise_seed,
|
|
new_total_steps,
|
|
new_cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
pipelines[1]['positive_conditioning'],
|
|
pipelines[1]['negative_conditioning'],
|
|
samples[0],
|
|
model2_start_step,
|
|
end_steps,
|
|
"enable"
|
|
)
|
|
print(f"Second model pass completed, latent shape: {samples[0]['samples'].shape}")
|
|
except Exception as e:
|
|
print(f"Error in second model pass: {str(e)}")
|
|
raise e
|
|
|
|
if comfy.model_management.processing_interrupted():
|
|
raise RuntimeError("Processing interrupted by user")
|
|
|
|
# Third model pass if enabled
|
|
if used_model_count >= 3 and ckpt_name3:
|
|
try:
|
|
new_total_steps = total_steps_original + total_steps_shift_third + total_steps_shift_second
|
|
new_cfg = starting_cfg + cfg_shift
|
|
model3_start_step = steps_end_second + steps_shift_third
|
|
|
|
samples = KSamplerAdvanced().sample(
|
|
pipelines[2]['model'],
|
|
"disable",
|
|
noise_seed,
|
|
new_total_steps,
|
|
new_cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
pipelines[2]['positive_conditioning'],
|
|
pipelines[2]['negative_conditioning'],
|
|
samples[0],
|
|
model3_start_step,
|
|
new_total_steps,
|
|
"enable"
|
|
)
|
|
print(f"Third model pass completed, latent shape: {samples[0]['samples'].shape}")
|
|
except Exception as e:
|
|
print(f"Error in third model pass: {str(e)}")
|
|
raise e
|
|
|
|
# Serialize the settings used for this run
|
|
settings_str = self.serialize_settings(**effective_kwargs)
|
|
|
|
return {
|
|
"ui": ui_updates,
|
|
"result": (samples[0], settings_str),
|
|
}
|
|
|
|
except Exception as e:
|
|
self.cleanup()
|
|
if str(e) == "Processing interrupted by user":
|
|
print("Processing was interrupted, cleaned up resources")
|
|
return {
|
|
"ui": ui_updates,
|
|
"result": (latent_image, ""),
|
|
}
|
|
else:
|
|
print(f"Error during sampling: {str(e)}")
|
|
raise e
|
|
finally:
|
|
if samples is None or comfy.model_management.processing_interrupted():
|
|
self.cleanup()
|
|
except Exception as e:
|
|
self.cleanup()
|
|
raise e
|
|
|
|
|
|
class MultiModelPromptSaverIterative:
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
self.prefix_append = ""
|
|
self.ckpt_cache = {} # Cache for checkpoint names
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"settings": ("STRING", {"multiline": True}),
|
|
},
|
|
"optional": {
|
|
"filename": (
|
|
"STRING",
|
|
{"default": "ComfyUI_%time_%seed_%counter_%ckpt", "multiline": False},
|
|
),
|
|
"path": ("STRING", {"default": "%date/%ckpt", "multiline": False}),
|
|
"extension": (["png", "jpg", "jpeg", "webp"],),
|
|
"lossless_webp": ("BOOLEAN", {"default": True}),
|
|
"jpg_webp_quality": ("INT", {"default": 100, "min": 1, "max": 100}),
|
|
"date_format": (
|
|
"STRING",
|
|
{"default": "%Y-%m-%d", "multiline": False},
|
|
),
|
|
"time_format": (
|
|
"STRING",
|
|
{"default": "%H%M%S", "multiline": False},
|
|
),
|
|
"save_metadata_file": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("FILENAME", "FILE_PATH", "METADATA")
|
|
FUNCTION = "save_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "SD Prompt Reader"
|
|
|
|
def save_images(
|
|
self,
|
|
images,
|
|
settings: str,
|
|
filename: str = "ComfyUI_%time_%seed_%counter_%ckpt",
|
|
path: str = "%date/%ckpt",
|
|
extension: str = "png",
|
|
lossless_webp: bool = True,
|
|
jpg_webp_quality: int = 100,
|
|
date_format: str = "%Y-%m-%d",
|
|
time_format: str = "%H%M%S",
|
|
save_metadata_file: bool = False,
|
|
prompt=None,
|
|
extra_pnginfo=None,
|
|
):
|
|
# Get or generate cache key for settings
|
|
settings_hash = hashlib.md5(settings.encode()).hexdigest()
|
|
|
|
# Check cache first
|
|
if settings_hash in self.ckpt_cache:
|
|
ckpt_name = self.ckpt_cache[settings_hash]
|
|
else:
|
|
ckpt_name = self.extract_checkpoint_name(settings)
|
|
self.ckpt_cache[settings_hash] = ckpt_name
|
|
|
|
# Prepare variable substitutions
|
|
variable_map = {
|
|
"%date": self.get_time(date_format),
|
|
"%time": self.get_time(time_format),
|
|
"%extension": extension,
|
|
"%quality": jpg_webp_quality,
|
|
"%ckpt": ckpt_name or "unknown_model",
|
|
}
|
|
|
|
# Get output paths
|
|
full_output_folder = folder_paths.get_output_directory()
|
|
subfolder = self.get_path(path, variable_map)
|
|
output_folder = Path(full_output_folder) / subfolder
|
|
output_folder.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Process images
|
|
results = []
|
|
files = []
|
|
file_paths = []
|
|
|
|
for image in images:
|
|
# Update counter for each image
|
|
counter = self.get_counter(output_folder)
|
|
variable_map["%counter"] = f"{counter:05}"
|
|
|
|
# Generate filename
|
|
stem = self.get_path(filename, variable_map)
|
|
file = self.get_unique_filename(stem, extension, output_folder)
|
|
file_path = output_folder / file
|
|
|
|
# Convert and save image
|
|
img = self.tensor_to_image(image)
|
|
self.save_image_with_metadata(
|
|
img,
|
|
file_path,
|
|
extension,
|
|
settings,
|
|
prompt,
|
|
extra_pnginfo,
|
|
lossless_webp,
|
|
jpg_webp_quality
|
|
)
|
|
|
|
# Save metadata file if requested
|
|
if save_metadata_file:
|
|
self.save_metadata_file(file_path, settings)
|
|
|
|
# Store results
|
|
results.append({
|
|
"filename": file.name,
|
|
"subfolder": str(subfolder),
|
|
"type": self.type
|
|
})
|
|
files.append(str(file))
|
|
file_paths.append(str(file_path))
|
|
|
|
return {
|
|
"ui": {"images": results},
|
|
"result": (
|
|
self.unpack_singleton(files),
|
|
self.unpack_singleton(file_paths),
|
|
settings, # Return original settings as metadata
|
|
),
|
|
}
|
|
|
|
def extract_checkpoint_name(self, settings_str: str) -> str:
|
|
"""Robust checkpoint name extraction from settings string"""
|
|
if not settings_str.strip():
|
|
return ""
|
|
|
|
try:
|
|
# First try direct JSON parse
|
|
try:
|
|
settings = json.loads(settings_str)
|
|
except json.JSONDecodeError:
|
|
# Handle cases where settings might be wrapped in metadata
|
|
if settings_str.startswith('{"parameters":'):
|
|
settings = json.loads(settings_str)
|
|
settings_str = settings.get("parameters", "")
|
|
settings = json.loads(settings_str)
|
|
else:
|
|
return ""
|
|
|
|
# Try to find checkpoint names in various locations
|
|
ckpt_name = (
|
|
settings.get("ckpt_name2") or
|
|
settings.get("ckpt_name1") or
|
|
""
|
|
)
|
|
|
|
# Clean up the name
|
|
if ckpt_name:
|
|
ckpt_name = Path(ckpt_name).stem
|
|
ckpt_name = re.sub(r'[\\/:*?"<>|\s]', '_', ckpt_name)
|
|
ckpt_name = re.sub(r'_{2,}', '_', ckpt_name).strip('_')
|
|
return ckpt_name[:64] # Limit length
|
|
|
|
return ""
|
|
except Exception as e:
|
|
print(f"Error extracting checkpoint: {str(e)}")
|
|
return ""
|
|
|
|
def tensor_to_image(self, tensor):
|
|
"""Convert torch tensor to PIL Image"""
|
|
i = 255.0 * tensor.cpu().numpy()
|
|
return Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
|
|
def save_image_with_metadata(
|
|
self,
|
|
img,
|
|
file_path,
|
|
extension,
|
|
settings,
|
|
prompt,
|
|
extra_pnginfo,
|
|
lossless_webp,
|
|
quality
|
|
):
|
|
"""Save image with appropriate metadata"""
|
|
if extension == "png":
|
|
metadata = PngInfo()
|
|
if not args.disable_metadata:
|
|
metadata.add_text("parameters", settings)
|
|
if prompt is not None:
|
|
metadata.add_text("prompt", json.dumps(prompt))
|
|
if extra_pnginfo is not None:
|
|
for x in extra_pnginfo:
|
|
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
|
img.save(file_path, pnginfo=metadata, compress_level=4)
|
|
else:
|
|
img.save(file_path, quality=quality, lossless=lossless_webp)
|
|
if not args.disable_metadata:
|
|
metadata = piexif.dump({
|
|
"Exif": {
|
|
piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(
|
|
settings, encoding="unicode"
|
|
)
|
|
},
|
|
})
|
|
piexif.insert(metadata, str(file_path))
|
|
|
|
def save_metadata_file(self, file_path, settings):
|
|
"""Save metadata to separate text file"""
|
|
with open(file_path.with_suffix(".txt"), "w", encoding="utf-8") as f:
|
|
f.write(settings)
|
|
|
|
@staticmethod
|
|
def get_counter(directory: Path):
|
|
"""Get next counter value for directory"""
|
|
existing = list(directory.glob("*"))
|
|
return len(existing) + 1
|
|
|
|
@staticmethod
|
|
def get_path(name, variable_map):
|
|
"""Apply variable substitutions to path/filename"""
|
|
for var, val in variable_map.items():
|
|
name = name.replace(var, str(val))
|
|
return Path(name)
|
|
|
|
@staticmethod
|
|
def get_time(time_format):
|
|
"""Get formatted time string"""
|
|
try:
|
|
return datetime.now().strftime(time_format)
|
|
except:
|
|
return ""
|
|
|
|
@staticmethod
|
|
def get_unique_filename(stem: Path, extension: str, output_folder: Path):
|
|
"""Generate unique filename by appending index if needed"""
|
|
file = stem.with_suffix(f".{extension}")
|
|
index = 0
|
|
while (output_folder / file).exists():
|
|
index += 1
|
|
file = stem.with_name(f"{stem.name}_{index}").with_suffix(f".{extension}")
|
|
return file
|
|
|
|
@staticmethod
|
|
def unpack_singleton(arr: list):
|
|
"""Return single item or full list"""
|
|
return arr[0] if len(arr) == 1 else arr |