diff --git a/__init__.py b/__init__.py index e0e7e5c..c073b45 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,7 @@ import os from .datastasher import CsvWriterNode +from .stringhasher import StringHasher from .dataloader import ImmatureImageDataLoader from .imagecounter import ImmatureImageCounter from .cached_checkpoint import CachedCheckpoint @@ -9,12 +10,15 @@ from .girlselector import RandomCharacterSelector from .schedulerselector import * from .advancedksampler import * from .multimodelcheckpointiterator import * +from .multimodelcheckpointiteratorfirst import MultiModelCheckpointIteratorFirst + from .danbooru import Ranbooru # A dictionary that contains all nodes you want to export with their names NODE_CLASS_MAPPINGS = { "CsvWriterNode": CsvWriterNode, + "StringHasher": StringHasher, "ImmatureImageDataLoader": ImmatureImageDataLoader, "ImmatureImageCounter": ImmatureImageCounter, "CachedCheckpoint": CachedCheckpoint, @@ -27,11 +31,15 @@ NODE_CLASS_MAPPINGS = { "MultiModelCheckpointIterator": MultiModelCheckpointIterator, "MultiModelPromptSaver": MultiModelPromptSaver, "MultiModelPromptSaverIterative": MultiModelPromptSaverIterative, + "MultiModelPromptSaverIterativeFirst": MultiModelCheckpointIteratorFirst, + "MultiModelPromptSaverIterativeFirst": MultiModelCheckpointIteratorFirst, + } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "CsvWriterNode": "CSV Writer", + "StringHasher": "StringHasher", "ImmatureImageDataLoader": "Immature Image Data Loader", "ImmatureImageCounter": "Immature Image Counter", "CachedCheckpoint": "Cached Checkpoint", @@ -42,7 +50,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "MultiModelAdvancedKsampler": "MultiModelAdvancedKsampler", "MultiModelCheckpointIterator": "MultiModelCheckpointIterator", "MultiModelPromptSaver": "MultiModelPromptSaver", - "MultiModelPromptSaverIterative": "MultiModelPromptSaverIterative" + "MultiModelPromptSaverIterative": "MultiModelPromptSaverIterative", + "MultiModelPromptSaverIterativeFirst": "MultiModelPromptSaverIterativeFirst" } __all__ = NODE_CLASS_MAPPINGS diff --git a/advancedksampler.py b/advancedksampler.py index dde908c..f2866c8 100644 --- a/advancedksampler.py +++ b/advancedksampler.py @@ -3,6 +3,7 @@ from PIL import Image, ImageOps, ImageDraw, ImageFont from PIL.PngImagePlugin import PngInfo import numpy as np import torch +import gc from datetime import datetime from itertools import chain @@ -22,7 +23,7 @@ import psutil import re import nodes as native -from .adv_encode import advanced_encode #, advanced_encode_XL +from .adv_encode import advanced_encode import comfy.sd from comfy_extras.nodes_align_your_steps import AlignYourStepsScheduler @@ -58,7 +59,6 @@ from comfy import samplers sys.path.append(custom_nodes_dir) - class LoraTagLoader: def __init__(self): self.loaded_lora = None @@ -77,10 +77,7 @@ class LoraTagLoader: 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) @@ -165,17 +162,20 @@ class AdvancedCLIPTextEncodeWithBreak: 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 MultiModelAdvancedKsampler: + 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 + @classmethod def INPUT_TYPES(s): return { @@ -211,21 +211,39 @@ class MultiModelAdvancedKsampler: } } - RETURN_TYPES = ("LATENT", "STRING") - RETURN_NAMES = ("LATENT", "settings") + RETURN_TYPES = ("LATENT", "STRING", "INT", "STRING", "STRING") + RETURN_NAMES = ("LATENT", "prompt", "seed", "lora", "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 + 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 cleanup of loaded checkpoints") + 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() - comfy.model_management.soft_empty_cache() + + # 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""" @@ -234,9 +252,7 @@ class MultiModelAdvancedKsampler: settings = { "used_model_count": used_model_count, "ckpt_name1": kwargs.get("ckpt_name1", ""), - # Only include ckpt_name2 if used_model_count >= 2 **({"ckpt_name2": kwargs.get("ckpt_name2", "")} if used_model_count >= 2 else {}), - # Only include ckpt_name3 if used_model_count >= 3 **({"ckpt_name3": kwargs.get("ckpt_name3", "")} if used_model_count >= 3 else {}), "positive": kwargs.get("positive", ""), "negative": kwargs.get("negative", ""), @@ -245,19 +261,15 @@ class MultiModelAdvancedKsampler: "rescaled_steps": kwargs.get("rescaled_steps", 8), "rescale_multiplier": kwargs.get("rescale_multiplier", 0.7), "total_steps_original": kwargs.get("total_steps_original", 25), - # Only include second model steps if used_model_count >= 2 **({"total_steps_shift_second": kwargs.get("total_steps_shift_second", 0)} if used_model_count >= 2 else {}), - # Only include third model steps if used_model_count >= 3 **({"total_steps_shift_third": kwargs.get("total_steps_shift_third", 0)} if used_model_count >= 3 else {}), "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), - # Only include second model steps if used_model_count >= 2 **({"steps_shift_second": kwargs.get("steps_shift_second", 0)} if used_model_count >= 2 else {}), **({"steps_end_second": kwargs.get("steps_end_second", 0)} if used_model_count >= 2 else {}), - # Only include third model steps if used_model_count >= 3 **({"steps_shift_third": kwargs.get("steps_shift_third", 0)} if used_model_count >= 3 else {}), "token_normalization": kwargs.get("token_normalization", "none"), "weight_interpretation": kwargs.get("weight_interpretation", "comfy"), @@ -311,6 +323,10 @@ class MultiModelAdvancedKsampler: # 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 @@ -319,6 +335,8 @@ class MultiModelAdvancedKsampler: 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] @@ -337,7 +355,14 @@ class MultiModelAdvancedKsampler: # Apply LoRA if specified if lora_name and lora_name.strip(): - model, clip, _ = LoraTagLoader().load_lora(model, clip, lora_name) + # 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( @@ -358,6 +383,9 @@ class MultiModelAdvancedKsampler: 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) @@ -397,7 +425,7 @@ class MultiModelAdvancedKsampler: "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_seed": (loaded_settings.get("detached_seed", False),), } # Extract all parameters from effective_kwargs @@ -578,9 +606,12 @@ class MultiModelAdvancedKsampler: # Serialize the settings used for this run settings_str = self.serialize_settings(**effective_kwargs) + # Print final memory usage + self.print_memory_usage() + return { "ui": ui_updates, - "result": (samples[0], settings_str), + "result": (samples[0], positive, noise_seed, lora_name, settings_str), } except Exception as e: @@ -600,8 +631,8 @@ class MultiModelAdvancedKsampler: except Exception as e: self.cleanup() raise e - - + + class MultiModelPromptSaver: def __init__(self): self.output_dir = folder_paths.get_output_directory() @@ -786,4 +817,4 @@ class MultiModelPromptSaver: @staticmethod def unpack_singleton(arr: list): - return arr[0] if len(arr) == 1 else arr \ No newline at end of file + return arr[0] if len(arr) == 1 else arr \ No newline at end of file diff --git a/multimodelcheckpointiterator.py b/multimodelcheckpointiterator.py index 0fb70a7..ddf5341 100644 --- a/multimodelcheckpointiterator.py +++ b/multimodelcheckpointiterator.py @@ -59,10 +59,9 @@ from comfy import samplers sys.path.append(custom_nodes_dir) - class LoraTagLoader: def __init__(self): - self.loaded_lora = None + self.loaded_loras = {} # Changed from single loaded_lora to dictionary self.tag_pattern = "\<[0-9a-zA-Z\:\_\-\.\s\/\(\)\\\\]+\>" @classmethod @@ -77,11 +76,14 @@ class LoraTagLoader: CATEGORY = "loaders" - def load_lora(self, model, clip, text): - # print(f"\nLoraTagLoader input text: { text }") + 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) - # print(f"\nfoound lora tags: { founds }") if len(founds) < 1: return (model, clip, text) @@ -89,18 +91,23 @@ class LoraTagLoader: 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: @@ -110,38 +117,44 @@ class LoraTagLoader: wClip = float(pak[3]) except ValueError: continue - if name == None: + + 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 == None: + + 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) - 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 - + 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_lora = (lora_path, lora) - + 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 @@ -205,6 +218,8 @@ class MultiModelCheckpointIterator: "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": { @@ -261,6 +276,8 @@ class MultiModelCheckpointIterator: "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) @@ -281,6 +298,8 @@ class MultiModelCheckpointIterator: 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) @@ -298,11 +317,19 @@ class MultiModelCheckpointIterator: # 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()] - # If we have a list, select the next checkpoint in sequence - if self.ckpt_name2_list and detached_checkpoint == True: - print("detached checkpoint is ON") - 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 + # 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: @@ -312,10 +339,19 @@ class MultiModelCheckpointIterator: # 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()] - # If we have a list, select the next checkpoint in sequence + # Handle random checkpoint selection if self.ckpt_name2_list: - 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 + 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"] = "" @@ -421,6 +457,8 @@ class MultiModelCheckpointIterator: "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 @@ -485,11 +523,15 @@ class MultiModelCheckpointIterator: samples = None try: # First sampling pass - if rescaled_steps > 0: + if rescaled_steps >= 0: # Changed from > 0 to >= 0 to handle 0 case try: - CFGR = comfy_extras.nodes_model_advanced.RescaleCFG() - patched_model = CFGR.patch(pipelines[0]['model'], rescale_multiplier)[0] - print("Attempting rescaled CFG pass") + 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, @@ -516,10 +558,10 @@ class MultiModelCheckpointIterator: # First model's main pass try: - actual_steps = rescaled_steps if rescaled_steps > 0 else 0 + 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", # add_noise + "disable" if rescaled_steps >= 0 else "enable", # Changed condition noise_seed, total_steps_original, starting_cfg, @@ -527,7 +569,7 @@ class MultiModelCheckpointIterator: scheduler, pipelines[0]['positive_conditioning'], pipelines[0]['negative_conditioning'], - samples[0] if rescaled_steps > 0 else latent_image, + samples[0] if rescaled_steps >= 0 else latent_image, # Changed condition actual_steps, steps_end_first, "enable" # return_with_leftover_noise @@ -624,7 +666,6 @@ class MultiModelCheckpointIterator: self.cleanup() raise e - class MultiModelPromptSaverIterative: def __init__(self): diff --git a/multimodelcheckpointiteratorfirst.py b/multimodelcheckpointiteratorfirst.py new file mode 100644 index 0000000..4ea6d43 --- /dev/null +++ b/multimodelcheckpointiteratorfirst.py @@ -0,0 +1,888 @@ +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 MultiModelCheckpointIteratorFirst: + @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}), + "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_name1_index = 0 # Index for cycling through first model checkpoints + self.ckpt_name1_list = [] # List of first 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", ""), # 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() + 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_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 == True: + 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 + else: + # Process the ckpt_name1_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}") + 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_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"] + 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: + try: + 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, + 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 + 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, + 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 \ No newline at end of file