738 lines
35 KiB
Python
738 lines
35 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 gc
|
|
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_lora = None
|
|
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 load_lora(self, model, clip, text):
|
|
# print(f"\nLoraTagLoader input text: { text }")
|
|
|
|
founds = re.findall(self.tag_pattern, text)
|
|
# print(f"\nfoound lora tags: { founds }")
|
|
|
|
if len(founds) < 1:
|
|
return (model, clip, text)
|
|
|
|
model_lora = model
|
|
clip_lora = clip
|
|
|
|
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 == 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 == 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)
|
|
lora = None
|
|
if self.loaded_lora is not None:
|
|
if self.loaded_lora[0] == lora_path:
|
|
lora = self.loaded_lora[1]
|
|
else:
|
|
temp = self.loaded_lora
|
|
self.loaded_lora = None
|
|
del temp
|
|
|
|
if lora is None:
|
|
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
self.loaded_lora = (lora_path, lora)
|
|
|
|
model_lora, clip_lora = comfy.sd.load_lora_for_models(model_lora, clip_lora, lora, wModel, wClip)
|
|
|
|
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 MultiModelCheckpointIteratorFirst:
|
|
def __init__(self):
|
|
self.loaded_checkpoints = {} # Dictionary to store loaded checkpoints by name
|
|
self.loaded_loras = {} # Dictionary to store loaded LoRAs by path
|
|
self.current_checkpoints = set() # Track currently requested checkpoints
|
|
self.ckpt_name1_index = 0 # Index for cycling through first model checkpoints
|
|
self.ckpt_name1_list = [] # List of first model checkpoints
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"used_model_count": ("INT",{"default": 2, "min": 1, "max": 3},),
|
|
"ckpt_name1_list": ("STRING", {"multiline": True, "default": ""}),
|
|
"ckpt_name2": (folder_paths.get_filename_list("checkpoints"),),
|
|
"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}),
|
|
"stop_at_step": ("INT", {"default": 0, "min": 0, "max": 100}),
|
|
"latent_image": ("LATENT",),
|
|
},
|
|
"optional": {
|
|
"load_settings": ("STRING", {"multiline": True, "default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "STRING", "INT", "STRING", "STRING", "STRING") # Added 6th output
|
|
RETURN_NAMES = ("LATENT", "prompt", "seed", "lora", "settings", "used_ckpt_name") # New output
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling"
|
|
|
|
def print_memory_usage(self):
|
|
process = psutil.Process(os.getpid())
|
|
print(f"Memory usage: {process.memory_info().rss / 1024 / 1024:.2f} MB")
|
|
print(f"Loaded checkpoints: {len(self.loaded_checkpoints)}")
|
|
print(f"Loaded LoRAs: {len(self.loaded_loras)}")
|
|
|
|
def cleanup(self):
|
|
"""Force cleanup of all resources"""
|
|
print("Performing full cleanup of loaded models and LoRAs")
|
|
|
|
# Clear checkpoints
|
|
for name, (model, clip, vae) in self.loaded_checkpoints.items():
|
|
del model
|
|
del clip
|
|
del vae
|
|
self.loaded_checkpoints = {}
|
|
|
|
# Clear LoRAs
|
|
for path, lora in self.loaded_loras.items():
|
|
del lora
|
|
self.loaded_loras = {}
|
|
|
|
self.current_checkpoints = set()
|
|
self.ckpt_name1_index = 0
|
|
self.ckpt_name1_list = []
|
|
|
|
# Force garbage collection
|
|
gc.collect()
|
|
comfy.model_management.soft_empty_cache()
|
|
self.print_memory_usage()
|
|
|
|
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", ""), # Store the currently used ckpt_name1
|
|
"ckpt_name1_list": kwargs.get("ckpt_name1_list", ""), # Store the full list
|
|
"ckpt_name2": kwargs.get("ckpt_name2", ""),
|
|
"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),
|
|
}
|
|
|
|
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)
|
|
|
|
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()
|
|
if "latent_image" in kwargs:
|
|
effective_kwargs["latent_image"] = kwargs["latent_image"]
|
|
|
|
# If detached_seed is True, override the seed from direct input
|
|
if detached_seed and "noise_seed" in kwargs:
|
|
effective_kwargs["noise_seed"] = kwargs["noise_seed"]
|
|
|
|
# Process checkpoint list if provided
|
|
ckpt_name1_list = kwargs.get("ckpt_name1_list", "")
|
|
if ckpt_name1_list:
|
|
# Split the list by line breaks and clean up each entry
|
|
self.ckpt_name1_list = [name.strip() for name in ckpt_name1_list.split('\n') if name.strip()]
|
|
|
|
# If we have a list, select the next checkpoint in sequence
|
|
if self.ckpt_name1_list and detached_checkpoint:
|
|
print("detached checkpoint is ON")
|
|
effective_kwargs["ckpt_name1"] = self.ckpt_name1_list[self.ckpt_name1_index % len(self.ckpt_name1_list)]
|
|
self.ckpt_name1_index += 1 # Increment for next run
|
|
|
|
return effective_kwargs
|
|
|
|
# Process checkpoint list when not loading settings
|
|
ckpt_name1_list = kwargs.get("ckpt_name1_list", "")
|
|
if ckpt_name1_list:
|
|
# Split the list by line breaks and clean up each entry
|
|
self.ckpt_name1_list = [name.strip() for name in ckpt_name1_list.split('\n') if name.strip()]
|
|
|
|
# If we have a list, select the next checkpoint in sequence
|
|
if self.ckpt_name1_list:
|
|
kwargs["ckpt_name1"] = self.ckpt_name1_list[self.ckpt_name1_index % len(self.ckpt_name1_list)]
|
|
self.ckpt_name1_index += 1 # Increment for next run
|
|
else:
|
|
kwargs["ckpt_name1"] = ""
|
|
|
|
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}")
|
|
model, clip, vae = self.loaded_checkpoints[name]
|
|
del model
|
|
del clip
|
|
del vae
|
|
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"""
|
|
self.print_memory_usage()
|
|
|
|
# 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():
|
|
# Create a new LoraTagLoader instance for each call
|
|
lora_loader = LoraTagLoader()
|
|
model, clip, _ = lora_loader.load_lora(model, clip, lora_name)
|
|
|
|
# Store the loaded LoRA reference if the loader has one
|
|
if lora_loader.loaded_lora is not None:
|
|
lora_path, lora = lora_loader.loaded_lora
|
|
self.loaded_loras[lora_path] = lora
|
|
|
|
# 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:
|
|
# Print initial memory usage
|
|
self.print_memory_usage()
|
|
|
|
# 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.copy()
|
|
|
|
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_name1_list": (loaded_settings.get("ckpt_name1_list", ""),),
|
|
"ckpt_name2": (loaded_settings.get("ckpt_name2", ""),),
|
|
"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),),
|
|
}
|
|
|
|
# Extract all parameters from effective_kwargs
|
|
used_model_count = effective_kwargs["used_model_count"]
|
|
ckpt_name1 = effective_kwargs["ckpt_name1"]
|
|
used_ckpt_name = effective_kwargs["ckpt_name1"] # The actual checkpoint used
|
|
used_ckpt_short = os.path.splitext(os.path.basename(used_ckpt_name))[0] # Cleaned name only
|
|
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"]
|
|
stop_at_step = kwargs.get("stop_at_step", 0) # Get from direct kwargs, not effective_kwargs
|
|
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:
|
|
try:
|
|
# Adjust rescaled_steps if stop_at_step is set
|
|
adjusted_rescaled_steps = min(rescaled_steps, stop_at_step) if stop_at_step > 0 else rescaled_steps
|
|
|
|
CFGR = comfy_extras.nodes_model_advanced.RescaleCFG()
|
|
patched_model = CFGR.patch(pipelines[0]['model'], rescale_multiplier)[0]
|
|
print("Attempting rescaled CFG pass")
|
|
|
|
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,
|
|
adjusted_rescaled_steps,
|
|
"enable" # return_with_leftover_noise
|
|
)
|
|
print(f"Rescaled CFG pass completed, latent shape: {samples[0]['samples'].shape}")
|
|
|
|
# Check if we should stop after rescaled steps
|
|
if stop_at_step > 0 and adjusted_rescaled_steps >= stop_at_step:
|
|
print(f"Stopping early at step {adjusted_rescaled_steps} due to stop_at_step setting")
|
|
settings_str = self.serialize_settings(**effective_kwargs)
|
|
return {
|
|
"ui": ui_updates,
|
|
"result": (samples[0], positive, noise_seed, lora_name, settings_str),
|
|
}
|
|
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
|
|
|
|
# Adjust steps_end_first if stop_at_step is set
|
|
adjusted_steps_end_first = steps_end_first
|
|
if stop_at_step > 0:
|
|
adjusted_steps_end_first = min(steps_end_first, stop_at_step)
|
|
|
|
samples = KSamplerAdvanced().sample(
|
|
pipelines[0]['model'],
|
|
"disable" if rescaled_steps > 0 else "enable", # add_noise
|
|
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,
|
|
actual_steps,
|
|
adjusted_steps_end_first,
|
|
"enable" # return_with_leftover_noise
|
|
)
|
|
print(f"First model pass completed, latent shape: {samples[0]['samples'].shape}")
|
|
|
|
# Check if we should stop after first model
|
|
if stop_at_step > 0 and adjusted_steps_end_first >= stop_at_step:
|
|
print(f"Stopping early at step {adjusted_steps_end_first} due to stop_at_step setting")
|
|
settings_str = self.serialize_settings(**effective_kwargs)
|
|
return {
|
|
"ui": ui_updates,
|
|
"result": (samples[0], positive, noise_seed, lora_name, settings_str),
|
|
}
|
|
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
|
|
|
|
# Adjust end_steps if stop_at_step is set
|
|
if stop_at_step > 0:
|
|
end_steps = min(end_steps, stop_at_step)
|
|
|
|
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}")
|
|
|
|
# Check if we should stop after second model
|
|
if stop_at_step > 0 and end_steps >= stop_at_step:
|
|
print(f"Stopping early at step {end_steps} due to stop_at_step setting")
|
|
settings_str = self.serialize_settings(**effective_kwargs)
|
|
return {
|
|
"ui": ui_updates,
|
|
"result": (samples[0], positive, noise_seed, lora_name, settings_str),
|
|
}
|
|
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
|
|
|
|
# Adjust new_total_steps if stop_at_step is set
|
|
adjusted_new_total_steps = new_total_steps
|
|
if stop_at_step > 0:
|
|
adjusted_new_total_steps = min(new_total_steps, stop_at_step)
|
|
|
|
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,
|
|
adjusted_new_total_steps,
|
|
"enable"
|
|
)
|
|
print(f"Third model pass completed, latent shape: {samples[0]['samples'].shape}")
|
|
|
|
# Check if we should stop after third model
|
|
if stop_at_step > 0 and adjusted_new_total_steps >= stop_at_step:
|
|
print(f"Stopping early at step {adjusted_new_total_steps} due to stop_at_step setting")
|
|
settings_str = self.serialize_settings(**effective_kwargs)
|
|
return {
|
|
"ui": ui_updates,
|
|
"result": (samples[0], positive, noise_seed, lora_name, settings_str),
|
|
}
|
|
except Exception as e:
|
|
print(f"Error in third model pass: {str(e)}")
|
|
raise e
|
|
|
|
# Serialize the settings used for this run (without stop_at_step)
|
|
settings_str = self.serialize_settings(**effective_kwargs)
|
|
|
|
# Print final memory usage
|
|
self.print_memory_usage()
|
|
|
|
return {
|
|
"ui": ui_updates,
|
|
"result": (
|
|
samples[0],
|
|
positive,
|
|
noise_seed,
|
|
lora_name,
|
|
settings_str,
|
|
used_ckpt_short # New: exposes the actual checkpoint used
|
|
),
|
|
}
|
|
|
|
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, "", 0, "", ""),
|
|
}
|
|
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 |