├── Image_text_pair_sequence_loader.py ├── __init__.py ├── image_pair_sequence_loader.py ├── image_sequence_loader.py ├── image_storage_nodes.py ├── lora_loader_elemental.py ├── lora_mixer_elemental.py ├── lora_selector.py ├── lora_weight_randomizer.py ├── quantized_lora_loader.py ├── random_word_replacer.py ├── text_combiner.py ├── text_file_selector.py └── text_processor.py /Image_text_pair_sequence_loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | import PIL 3 | from PIL import Image 4 | import numpy as np 5 | import torch 6 | import random 7 | 8 | class ImageTextPairSequenceLoader: 9 | @classmethod 10 | def INPUT_TYPES(s): 11 | return { 12 | "required": { 13 | "image_folder_path": ("STRING", {"default": ""}), 14 | "text_folder_path": ("STRING", {"default": ""}), 15 | "reset": ("BOOLEAN", {"default": False}), 16 | "loop_or_reset": ("BOOLEAN", {"default": False}), 17 | "reset_on_error": ("BOOLEAN", {"default": False}), 18 | "exclude_loaded_on_reset": ("BOOLEAN", {"default": False}), 19 | "output_alpha": ("BOOLEAN", {"default": False}), 20 | "include_extension": ("BOOLEAN", {"default": False}), 21 | "start_index": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 22 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 23 | } 24 | } 25 | 26 | RETURN_TYPES = ("IMAGE", "STRING", "INT", "STRING") 27 | RETURN_NAMES = ("image", "text", "index", "filename") 28 | FUNCTION = "run" 29 | CATEGORY = "tksw_node" 30 | 31 | def __init__(self): 32 | self.current_index = 0 33 | self.image_files = [] 34 | self.text_files = [] 35 | self.common_basenames = [] 36 | self.image_file_map = {} 37 | self.text_file_map = {} 38 | self.prev_image_folder_path = "" 39 | self.prev_text_folder_path = "" 40 | self.prev_start_index = 0 41 | 42 | def _load_files(self, image_folder_path, text_folder_path): 43 | try: 44 | self.image_files = sorted([ 45 | f for f in os.listdir(image_folder_path) 46 | if os.path.isfile(os.path.join(image_folder_path, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp')) 47 | ]) 48 | except FileNotFoundError: 49 | print(f"Warning: Image folder not found: {image_folder_path}") 50 | self.image_files = [] 51 | except Exception as e: 52 | print(f"Error reading image folder {image_folder_path}: {e}") 53 | self.image_files = [] 54 | 55 | try: 56 | self.text_files = sorted([ 57 | f for f in os.listdir(text_folder_path) 58 | if os.path.isfile(os.path.join(text_folder_path, f)) and f.lower().endswith('.txt') 59 | ]) 60 | except FileNotFoundError: 61 | print(f"Warning: Text folder not found: {text_folder_path}") 62 | self.text_files = [] 63 | except Exception as e: 64 | print(f"Error reading text folder {text_folder_path}: {e}") 65 | self.text_files = [] 66 | 67 | self.image_file_map = {os.path.splitext(f)[0]: f for f in self.image_files} 68 | self.text_file_map = {os.path.splitext(f)[0]: f for f in self.text_files} 69 | 70 | image_basenames = set(self.image_file_map.keys()) 71 | text_basenames = set(self.text_file_map.keys()) 72 | self.common_basenames = sorted(list(image_basenames & text_basenames)) 73 | 74 | if not self.common_basenames: 75 | print("Warning: No common filenames (excluding extension) found between image and text folders.") 76 | 77 | 78 | def _load_image(self, folder_path, filename, alpha): 79 | image_path = os.path.join(folder_path, filename) 80 | try: 81 | with Image.open(image_path) as img: 82 | if alpha and img.mode != 'RGBA': 83 | img = img.convert("RGBA") 84 | elif not alpha and img.mode != 'RGB': 85 | img = img.convert("RGB") 86 | 87 | output_image = np.array(img).astype(np.float32) / 255.0 88 | output_image = torch.from_numpy(output_image).unsqueeze(0) 89 | return output_image 90 | except FileNotFoundError: 91 | print(f"Warning: Image file not found: {image_path}") 92 | return None 93 | except (PIL.UnidentifiedImageError, OSError) as e: 94 | print(f"Warning: Skipping corrupted or unsupported image file: {filename} ({e})") 95 | return None 96 | except Exception as e: 97 | print(f"Error loading image {filename}: {e}") 98 | return None 99 | 100 | def _load_text(self, folder_path, filename): 101 | text_path = os.path.join(folder_path, filename) 102 | try: 103 | with open(text_path, 'r', encoding='utf-8') as f: 104 | text_content = f.read() 105 | return text_content 106 | except FileNotFoundError: 107 | print(f"Warning: Text file not found: {text_path}") 108 | return None 109 | except Exception as e: 110 | print(f"Error loading text file {filename}: {e}") 111 | return None 112 | 113 | def run(self, image_folder_path, text_folder_path, reset, reset_on_error, seed, loop_or_reset, include_extension, exclude_loaded_on_reset, output_alpha, start_index): 114 | random.seed(seed) 115 | 116 | if not image_folder_path or not text_folder_path: 117 | print("Error: Image folder path and Text folder path must be specified.") 118 | return (None, None, 0, None) 119 | 120 | needs_reload = ( 121 | reset or 122 | not self.common_basenames or 123 | image_folder_path != self.prev_image_folder_path or 124 | text_folder_path != self.prev_text_folder_path or 125 | start_index != self.prev_start_index 126 | ) 127 | 128 | if needs_reload: 129 | self._load_files(image_folder_path, text_folder_path) 130 | self.current_index = 0 131 | self.prev_image_folder_path = image_folder_path 132 | self.prev_text_folder_path = text_folder_path 133 | self.prev_start_index = start_index 134 | 135 | if self.common_basenames and start_index >= len(self.common_basenames): 136 | print(f"Warning: start_index ({start_index}) is out of bounds. Setting to last index ({len(self.common_basenames) - 1}).") 137 | start_index = len(self.common_basenames) - 1 138 | elif not self.common_basenames: 139 | start_index = 0 140 | 141 | if not self.common_basenames: 142 | print("Error: No image/text pairs found.") 143 | return (None, None, self.current_index, None) 144 | 145 | effective_index = self.current_index + start_index 146 | 147 | if effective_index >= len(self.common_basenames): 148 | if loop_or_reset: 149 | print("Looping back to start.") 150 | self._load_files(image_folder_path, text_folder_path) 151 | self.current_index = 0 152 | start_index = 0 153 | self.prev_start_index = 0 154 | effective_index = 0 155 | if not self.common_basenames: 156 | print("Error: No image/text pairs found after reload.") 157 | return (None, None, self.current_index, None) 158 | else: 159 | print("Reached end of sequence.") 160 | return (None, None, self.current_index, None) 161 | 162 | output_image = None 163 | output_text = None 164 | current_basename = None 165 | loaded_successfully = False 166 | 167 | original_start_index = self.current_index 168 | 169 | while effective_index < len(self.common_basenames): 170 | current_basename = self.common_basenames[effective_index] 171 | image_filename = self.image_file_map.get(current_basename) 172 | text_filename = self.text_file_map.get(current_basename) 173 | 174 | if not image_filename or not text_filename: 175 | print(f"Error: Internal inconsistency. Basename '{current_basename}' not found in file maps.") 176 | else: 177 | output_image = self._load_image(image_folder_path, image_filename, output_alpha) 178 | output_text = self._load_text(text_folder_path, text_filename) 179 | 180 | if output_image is not None and output_text is not None: 181 | loaded_successfully = True 182 | break 183 | else: 184 | print(f"Failed to load pair for basename: {current_basename}") 185 | if reset_on_error: 186 | print("Resetting sequence due to error.") 187 | if exclude_loaded_on_reset: 188 | print("Excluding previously loaded files on reset.") 189 | loaded_basenames = set(self.common_basenames[:original_start_index + start_index]) 190 | self._load_files(image_folder_path, text_folder_path) 191 | self.common_basenames = [b for b in self.common_basenames if b not in loaded_basenames] 192 | else: 193 | self._load_files(image_folder_path, text_folder_path) 194 | 195 | self.current_index = 0 196 | start_index = 0 197 | self.prev_start_index = 0 198 | effective_index = 0 199 | 200 | if not self.common_basenames: 201 | print("Error: No image/text pairs remaining after reset.") 202 | return (None, None, 0, None) 203 | continue 204 | 205 | else: 206 | print("Skipping current pair due to error.") 207 | self.current_index += 1 208 | effective_index = self.current_index + start_index 209 | if effective_index >= len(self.common_basenames): 210 | print("Reached end of sequence after skipping error.") 211 | return (None, None, self.current_index, None) 212 | 213 | if not loaded_successfully: 214 | print("Error: Could not successfully load any remaining image/text pair.") 215 | return (None, None, self.current_index, None) 216 | 217 | 218 | output_filename = current_basename 219 | if include_extension: 220 | image_filename = self.image_file_map.get(current_basename) 221 | if image_filename: 222 | output_filename = image_filename 223 | 224 | final_index = effective_index 225 | self.current_index += 1 226 | 227 | return (output_image, output_text, final_index, output_filename) -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | from .image_sequence_loader import ImageSequenceLoader 2 | from .image_pair_sequence_loader import ImagePairSequenceLoader 3 | from .text_combiner import TextCombiner 4 | from .text_processor import TextProcessor 5 | from .lora_loader_elemental import LoraLoaderElemental 6 | from .random_word_replacer import RandomWordReplacer 7 | from .lora_weight_randomizer import LoraWeightRandomizer 8 | from .lora_mixer_elemental import LoraMixerElemental 9 | from .quantized_lora_loader import QuantizedLoraLoader 10 | from .lora_selector import LoraSelector 11 | from .text_file_selector import TextFileSelector 12 | from .Image_text_pair_sequence_loader import ImageTextPairSequenceLoader 13 | from .image_storage_nodes import ( 14 | StoreImageByNumber, 15 | RetrieveImageByNumber, 16 | StoreMultipleImagesByNumber, 17 | RetrieveMultipleImagesByNumber 18 | ) 19 | 20 | # --- 基本的なノードマッピング定義 --- 21 | # これらは常に定義される 22 | NODE_CLASS_MAPPINGS = { 23 | "ImageSequenceLoader": ImageSequenceLoader, 24 | "ImagePairSequenceLoader": ImagePairSequenceLoader, 25 | "TextCombiner": TextCombiner, 26 | "TextProcessor": TextProcessor, 27 | "LoraLoaderElemental": LoraLoaderElemental, 28 | "RandomWordReplacer": RandomWordReplacer, 29 | "LoraWeightRandomizer": LoraWeightRandomizer, 30 | "LoraMixerElemental": LoraMixerElemental, 31 | "QuantizedLoraLoader": QuantizedLoraLoader, 32 | "LoraSelector": LoraSelector, 33 | "TextFileSelector": TextFileSelector, 34 | "ImageTextPairSequenceLoader": ImageTextPairSequenceLoader, 35 | "StoreImageByNumber": StoreImageByNumber, 36 | "RetrieveImageByNumber": RetrieveImageByNumber, 37 | "StoreMultipleImagesByNumber": StoreMultipleImagesByNumber, 38 | "RetrieveMultipleImagesByNumber": RetrieveMultipleImagesByNumber, 39 | 40 | } 41 | 42 | NODE_DISPLAY_NAME_MAPPINGS = { 43 | "ImageSequenceLoader": "Image Sequence Loader", 44 | "ImagePairSequenceLoader": "Image Pair Sequence Loader", 45 | "TextCombiner": "Text Combiner", 46 | "TextProcessor": "Text Processor", 47 | "LoraLoaderElemental": "Lora Loader Elemental", 48 | "RandomWordReplacer": "Random Word Replacer", 49 | "LoraWeightRandomizer": "Lora Weight Randomizer", 50 | "LoraMixerElemental": "Lora Mixer Elemental", 51 | "QuantizedLoraLoader": "Quantized Lora Loader", 52 | "LoraSelector": "Lora Selector", 53 | "TextFileSelector": "Text File Selector", 54 | "ImageTextPairSequenceLoader": "Image TextPair SequenceLoader", 55 | "StoreImageByNumber": "Store Image by Number", 56 | "RetrieveImageByNumber": "Retrieve Image by Number", 57 | "StoreMultipleImagesByNumber": "Store Multiple Images by Number", 58 | "RetrieveMultipleImagesByNumber": "Retrieve Multiple Images by Number", 59 | } 60 | 61 | 62 | __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] -------------------------------------------------------------------------------- /image_pair_sequence_loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | from PIL import Image 3 | import numpy as np 4 | import torch 5 | import random 6 | 7 | class ImagePairSequenceLoader: 8 | @classmethod 9 | def INPUT_TYPES(s): 10 | return { 11 | "required": { 12 | "folder_path_A": ("STRING", {"default": ""}), 13 | "folder_path_B": ("STRING", {"default": ""}), 14 | "reset": ("BOOLEAN", {"default": False}), 15 | "loop_or_reset": ("BOOLEAN", {"default": False}), 16 | "reset_on_error": ("BOOLEAN", {"default": False}), 17 | "exclude_loaded_on_reset": ("BOOLEAN", {"default": False}), 18 | "output_alpha": ("BOOLEAN", {"default": False}), 19 | "include_extension": ("BOOLEAN", {"default": False}), 20 | "start_index": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 21 | "match_extension": ("BOOLEAN", {"default": False}), 22 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 23 | } 24 | } 25 | 26 | RETURN_TYPES = ("IMAGE", "IMAGE", "INT", "STRING") 27 | RETURN_NAMES = ("image_A", "image_B", "index", "filename") 28 | FUNCTION = "run" 29 | CATEGORY = "tksw_node" 30 | 31 | def __init__(self): 32 | self.current_index = 0 33 | self.image_files_A = [] 34 | self.image_files_B = [] 35 | self.common_files = [] 36 | self.prev_folder_path_A = "" 37 | self.prev_folder_path_B = "" 38 | self.prev_start_index = 0 39 | 40 | def _load_image_files(self, folder_path_A, folder_path_B, match_extension): 41 | self.image_files_A = sorted([ 42 | f for f in os.listdir(folder_path_A) 43 | if os.path.isfile(os.path.join(folder_path_A, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp')) 44 | ]) 45 | self.image_files_B = sorted([ 46 | f for f in os.listdir(folder_path_B) 47 | if os.path.isfile(os.path.join(folder_path_B, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp')) 48 | ]) 49 | 50 | if match_extension: 51 | self.common_files = sorted(list(set(self.image_files_A) & set(self.image_files_B))) 52 | else: 53 | filenames_A = [os.path.splitext(f)[0] for f in self.image_files_A] 54 | filenames_B = [os.path.splitext(f)[0] for f in self.image_files_B] 55 | common_filenames = sorted(list(set(filenames_A) & set(filenames_B))) 56 | self.common_files = [f + os.path.splitext(self.image_files_A[filenames_A.index(f)])[1] for f in common_filenames] 57 | 58 | def _load_image(self, folder_path, filename, alpha): 59 | image_path = os.path.join(folder_path, filename) 60 | try: 61 | with Image.open(image_path) as image: 62 | if alpha == False: 63 | image = image.convert("RGB") 64 | output_image = np.array(image).astype(np.float32) / 255.0 65 | output_image = torch.from_numpy(output_image).unsqueeze(0) 66 | return output_image 67 | except (PIL.UnidentifiedImageError, OSError) as e: 68 | print(f"Warning: Skipping corrupted image file: {image_path} ({e})") 69 | return None 70 | 71 | def run(self, folder_path_A, folder_path_B, reset, reset_on_error, seed, loop_or_reset, include_extension, exclude_loaded_on_reset, output_alpha, start_index, match_extension): 72 | random.seed(seed) 73 | 74 | if not folder_path_B: 75 | folder_path_B = folder_path_A 76 | 77 | if reset or not self.common_files or folder_path_A != self.prev_folder_path_A or folder_path_B != self.prev_folder_path_B or start_index != self.prev_start_index: 78 | if folder_path_A == folder_path_B: 79 | self._load_image_files(folder_path_A, folder_path_A, match_extension) 80 | self.common_files = self.image_files_A 81 | else: 82 | self._load_image_files(folder_path_A, folder_path_B, match_extension) 83 | 84 | self.current_index = 0 85 | self.prev_folder_path_A = folder_path_A 86 | self.prev_folder_path_B = folder_path_B 87 | self.prev_start_index = start_index 88 | 89 | if start_index >= len(self.common_files): 90 | start_index = len(self.common_files) - 1 91 | 92 | if not self.common_files: 93 | return (None, None, self.current_index, None) 94 | 95 | if self.current_index + start_index >= len(self.common_files): 96 | if loop_or_reset: 97 | if folder_path_A == folder_path_B: 98 | self._load_image_files(folder_path_A, folder_path_A, match_extension) 99 | self.common_files = self.image_files_A 100 | else: 101 | self._load_image_files(folder_path_A, folder_path_B, match_extension) 102 | self.current_index = 0 103 | else: 104 | self.current_index = 0 105 | 106 | filename = self.common_files[self.current_index + start_index] 107 | output_image_A = self._load_image(folder_path_A, filename, output_alpha) 108 | 109 | if folder_path_A == folder_path_B: 110 | output_image_B = output_image_A 111 | else: 112 | filename_B = os.path.splitext(filename)[0] + os.path.splitext(self.image_files_B[self.image_files_A.index(filename)])[1] if not match_extension else filename # 拡張子が異なる場合のファイル名 113 | output_image_B = self._load_image(folder_path_B, filename_B, output_alpha) 114 | 115 | if not include_extension: 116 | filename = os.path.splitext(filename)[0] 117 | 118 | self.current_index += 1 119 | 120 | return (output_image_A, output_image_B, self.current_index + start_index - 1, filename) 121 | 122 | while self.current_index + start_index < len(self.common_files): 123 | filename = self.common_files[self.current_index + start_index] 124 | output_image_A = self._load_image(folder_path_A, filename, output_alpha) 125 | 126 | if folder_path_A == folder_path_B: 127 | output_image_B = output_image_A 128 | else: 129 | filename_B = os.path.splitext(filename)[0] + os.path.splitext(self.image_files_B[self.image_files_A.index(filename)])[1] if not match_extension else filename 130 | output_image_B = self._load_image(folder_path_B, filename_B, output_alpha) 131 | 132 | 133 | if output_image_A is not None and output_image_B is not None: 134 | break 135 | elif not reset_on_error: 136 | self.current_index += 1 137 | else: 138 | if folder_path_A == folder_path_B: 139 | self._load_image_files(folder_path_A, folder_path_A, match_extension) 140 | self.common_files = self.image_files_A 141 | else: 142 | self._load_image_files(folder_path_A, folder_path_B, match_extension) 143 | 144 | if exclude_loaded_on_reset: 145 | loaded_files = set(self.common_files[:self.current_index]) 146 | self.common_files = [f for f in self.common_files if f not in loaded_files] 147 | self.current_index = 0 148 | 149 | if not include_extension: 150 | filename = os.path.splitext(filename)[0] 151 | 152 | self.current_index += 1 153 | 154 | return (output_image_A, output_image_B, self.current_index + start_index - 1, filename) -------------------------------------------------------------------------------- /image_sequence_loader.py: -------------------------------------------------------------------------------- 1 | import os 2 | from PIL import Image 3 | import numpy as np 4 | import torch 5 | import random 6 | 7 | class ImageSequenceLoader: 8 | @classmethod 9 | def INPUT_TYPES(s): 10 | return { 11 | "required": { 12 | "folder_path": ("STRING", {"default": ""}), 13 | "reset": ("BOOLEAN", {"default": False}), 14 | "loop_or_reset": ("BOOLEAN", {"default": False}), 15 | "reset_on_error": ("BOOLEAN", {"default": False}), 16 | "exclude_loaded_on_reset": ("BOOLEAN", {"default": False}), 17 | "output_alpha": ("BOOLEAN", {"default": False}), 18 | "include_extension": ("BOOLEAN", {"default": False}), 19 | "use_manual_index": ("BOOLEAN", {"default": False}), 20 | "start_index": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 21 | "manual_index": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 22 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 23 | } 24 | } 25 | 26 | RETURN_TYPES = ("IMAGE", "INT", "INT", "STRING") 27 | RETURN_NAMES = ("image", "index", "seed", "filename") 28 | FUNCTION = "run" 29 | CATEGORY = "tksw_node" 30 | 31 | def __init__(self): 32 | self.current_index = 0 33 | self.image_files = [] 34 | self.prev_folder_path = "" 35 | 36 | def _load_image_files(self, folder_path): 37 | self.image_files = sorted([ 38 | f for f in os.listdir(folder_path) 39 | if os.path.isfile(os.path.join(folder_path, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp')) 40 | ]) 41 | 42 | def _load_image(self, folder_path, index, alpha): 43 | if 0 <= index < len(self.image_files): 44 | image_path = os.path.join(folder_path, self.image_files[index]) 45 | filename = self.image_files[index] 46 | try: 47 | with Image.open(image_path) as image: 48 | if alpha == False: 49 | image = image.convert("RGB") 50 | output_image = np.array(image).astype(np.float32) / 255.0 51 | output_image = torch.from_numpy(output_image).unsqueeze(0) 52 | return output_image, filename 53 | except (PIL.UnidentifiedImageError, OSError) as e: 54 | print(f"Warning: Skipping corrupted image file: {image_path} ({e})") 55 | return None, filename 56 | else: 57 | return None, None 58 | 59 | def run(self, folder_path, reset, reset_on_error, seed, loop_or_reset, include_extension, exclude_loaded_on_reset, output_alpha, start_index, use_manual_index, manual_index): 60 | random.seed(seed) 61 | 62 | if reset or not self.image_files or folder_path != self.prev_folder_path: 63 | self._load_image_files(folder_path) 64 | self.current_index = 0 65 | self.prev_folder_path = folder_path 66 | 67 | if not self.image_files: 68 | return (None, self.current_index, seed, None) 69 | 70 | while self.current_index < len(self.image_files): 71 | if use_manual_index == False : 72 | output_image, filename = self._load_image(folder_path, self.current_index + start_index, output_alpha) 73 | else : 74 | output_image, filename = self._load_image(folder_path, manual_index + start_index, output_alpha) 75 | if output_image is not None: 76 | break 77 | elif not reset_on_error: 78 | self.current_index += 1 79 | else: 80 | self._load_image_files(folder_path) 81 | if exclude_loaded_on_reset: 82 | loaded_files = set(self.image_files[:self.current_index]) 83 | self.image_files = [f for f in self.image_files if f not in loaded_files] 84 | self.current_index = 0 85 | 86 | if self.current_index >= len(self.image_files): 87 | if loop_or_reset: 88 | self._load_image_files(folder_path) 89 | self.current_index = 0 90 | else: 91 | self.current_index = 0 92 | 93 | output_image, filename = self._load_image(folder_path, self.current_index, output_alpha) 94 | 95 | if not include_extension: 96 | filename = os.path.splitext(filename)[0] 97 | 98 | self.current_index += 1 99 | 100 | if use_manual_index == False : 101 | return_index = self.current_index + start_index - 1 102 | else : 103 | return_index = manual_index + start_index 104 | 105 | return (output_image, return_index, seed, filename) 106 | -------------------------------------------------------------------------------- /image_storage_nodes.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from typing import Tuple, Dict, Any, List, Optional 3 | 4 | shared_image_pool: Dict[int, torch.Tensor] = {} 5 | 6 | MAX_IMAGE_SLOTS = 5 7 | 8 | class StoreImageByNumber: 9 | 10 | @classmethod 11 | def INPUT_TYPES(cls) -> Dict[str, Any]: 12 | return { 13 | "required": { 14 | "image": ("IMAGE",), 15 | "image_id": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "step": 1}), 16 | "skip_if_exists": ("BOOLEAN", {"default": False, "label_on": "Skip if ID exists", "label_off": "Overwrite if ID exists"}), 17 | } 18 | } 19 | 20 | RETURN_TYPES: Tuple[()] = () 21 | FUNCTION: str = "store_image" 22 | OUTPUT_NODE: bool = True 23 | CATEGORY = "tksw_node" 24 | 25 | @classmethod 26 | def IS_CHANGED(cls, image: torch.Tensor, image_id: int, skip_if_exists: bool) -> float: 27 | return float("NaN") 28 | 29 | def store_image(self, image: torch.Tensor, image_id: int, skip_if_exists: bool) -> Dict[str, Any]: 30 | log_prefix = f"[Store Image (Memory)] Number ID {image_id}" 31 | 32 | if skip_if_exists and image_id in shared_image_pool: 33 | message = f"{log_prefix}: ID already exists and 'skip_if_exists' is True. Skipped." 34 | print(message) 35 | return {"ui": {"text": f"ID {image_id} skipped."}} 36 | 37 | is_overwrite = image_id in shared_image_pool 38 | shared_image_pool[image_id] = image.clone() 39 | 40 | if skip_if_exists: 41 | action_msg = "newly stored" 42 | print(f"{log_prefix}: Image {action_msg}. Shape: {image.shape}, Device: {image.device} (skip_if_exists: True)") 43 | else: 44 | action_msg = "overwritten" if is_overwrite else "newly stored" 45 | print(f"{log_prefix}: Image {action_msg}. Shape: {image.shape}, Device: {image.device} (skip_if_exists: False)") 46 | 47 | return {"ui": {"text": f"ID {image_id} {action_msg}."}} 48 | 49 | class RetrieveImageByNumber: 50 | 51 | @classmethod 52 | def INPUT_TYPES(cls) -> Dict[str, Any]: 53 | return { 54 | "required": { 55 | "image_id": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "step": 1}), 56 | "remove_after_retrieval": ("BOOLEAN", {"default": False, "label_on": "Remove after retrieval", "label_off": "Keep after retrieval"}), 57 | }, 58 | "optional": { 59 | "fallback_image": ("IMAGE",) 60 | } 61 | } 62 | 63 | RETURN_TYPES: Tuple[str, ...] = ("IMAGE",) 64 | RETURN_NAMES: Tuple[str, ...] = ("image",) 65 | FUNCTION: str = "retrieve_image" 66 | CATEGORY = "tksw_node" 67 | 68 | @classmethod 69 | def IS_CHANGED(cls, image_id: int, remove_after_retrieval: bool, fallback_image: Optional[torch.Tensor] = None) -> float: 70 | return float("NaN") 71 | 72 | def retrieve_image(self, image_id: int, remove_after_retrieval: bool, fallback_image: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor]: 73 | retrieved_image: Optional[torch.Tensor] = None 74 | log_prefix = f"[Retrieve Image (Memory)] Number ID {image_id}" 75 | 76 | if image_id in shared_image_pool: 77 | retrieved_image = shared_image_pool[image_id].clone() 78 | print(f"{log_prefix}: Image retrieved. Shape: {retrieved_image.shape}, Device: {retrieved_image.device}") 79 | 80 | if remove_after_retrieval: 81 | del shared_image_pool[image_id] 82 | print(f"{log_prefix}: Image removed from pool (remove_after_retrieval: True).") 83 | else: 84 | print(f"{log_prefix}: Image kept in pool (remove_after_retrieval: False).") 85 | else: 86 | print(f"{log_prefix}: Image not found in memory.") 87 | 88 | if retrieved_image is not None: 89 | return (retrieved_image,) 90 | 91 | if fallback_image is not None: 92 | print(f"{log_prefix}: Using provided fallback image.") 93 | return (fallback_image.clone(),) 94 | else: 95 | print(f"{log_prefix}: No fallback image provided. Outputting default 64x64 black image.") 96 | default_img = torch.zeros((1, 64, 64, 3), dtype=torch.float32, device="cpu") 97 | return (default_img,) 98 | 99 | class StoreMultipleImagesByNumber: 100 | @classmethod 101 | def INPUT_TYPES(cls) -> Dict[str, Any]: 102 | inputs: Dict[str, Any] = { 103 | "required": { 104 | "skip_if_exists": ("BOOLEAN", {"default": False, "label_on": "Skip if ID exists", "label_off": "Overwrite if ID exists"}), 105 | }, 106 | "optional": {} 107 | } 108 | for i in range(1, MAX_IMAGE_SLOTS + 1): 109 | inputs["optional"][f"image_{i}"] = ("IMAGE",) 110 | inputs["optional"][f"image_id_{i}"] = ("INT", {"default": i - 1, "min": 0, "max": 0xffffffffffffffff, "step": 1}) 111 | return inputs 112 | 113 | RETURN_TYPES: Tuple[()] = () 114 | FUNCTION: str = "store_images" 115 | OUTPUT_NODE: bool = True 116 | CATEGORY = "tksw_node" 117 | 118 | @classmethod 119 | def IS_CHANGED(cls, **kwargs: Any) -> float: 120 | return float("NaN") 121 | 122 | def store_images(self, skip_if_exists: bool, **kwargs: Any) -> Dict[str, Any]: 123 | stored_count = 0 124 | skipped_count = 0 125 | processed_ids_info: List[str] = [] 126 | 127 | for i in range(1, MAX_IMAGE_SLOTS + 1): 128 | image_slot_key = f"image_{i}" 129 | image_id_slot_key = f"image_id_{i}" 130 | 131 | image: Optional[torch.Tensor] = kwargs.get(image_slot_key) 132 | image_id: Optional[int] = kwargs.get(image_id_slot_key) 133 | 134 | if image is not None and image_id is not None: 135 | log_prefix = f"[Store Multiple (Memory)] Slot {i} (ID {image_id})" 136 | 137 | if skip_if_exists and image_id in shared_image_pool: 138 | print(f"{log_prefix}: ID exists, 'skip_if_exists' is True. Skipped.") 139 | skipped_count += 1 140 | processed_ids_info.append(f"ID {image_id}(S{i}):skipped") 141 | continue 142 | 143 | is_overwrite = image_id in shared_image_pool 144 | shared_image_pool[image_id] = image.clone() 145 | stored_count += 1 146 | 147 | if skip_if_exists: 148 | action_msg = "newly stored" 149 | processed_ids_info.append(f"ID {image_id}(S{i}):stored") 150 | else: 151 | action_msg = "overwritten" if is_overwrite else "newly stored" 152 | processed_ids_info.append(f"ID {image_id}(S{i}):{'overwritten' if is_overwrite else 'stored'}") 153 | 154 | print(f"{log_prefix}: Image {action_msg}. Shape: {image.shape} (skip_if_exists: {skip_if_exists})") 155 | elif image is None and kwargs.get(image_id_slot_key) is not None: 156 | print(f"[Store Multiple (Memory)] Slot {i} (ID {image_id}): image_id provided but no image. Skipped.") 157 | 158 | ui_summary = f"Stored: {stored_count}, Skipped: {skipped_count}." 159 | if processed_ids_info: 160 | ui_summary += " Details: " + ", ".join(processed_ids_info) 161 | if stored_count == 0 and skipped_count == 0: 162 | ui_summary = "No images/IDs provided to process." 163 | 164 | return {"ui": {"text": ui_summary}} 165 | 166 | 167 | class RetrieveMultipleImagesByNumber: 168 | @classmethod 169 | def INPUT_TYPES(cls) -> Dict[str, Any]: 170 | inputs: Dict[str, Any] = { 171 | "required": { 172 | "remove_after_retrieval": ("BOOLEAN", {"default": False, "label_on": "Remove after retrieval", "label_off": "Keep after retrieval"}), 173 | }, 174 | "optional": { 175 | "fallback_image": ("IMAGE",) 176 | } 177 | } 178 | for i in range(1, MAX_IMAGE_SLOTS + 1): 179 | inputs["optional"][f"image_id_{i}"] = ("INT", {"default": i - 1, "min": 0, "max": 0xffffffffffffffff, "step": 1}) 180 | return inputs 181 | 182 | RETURN_TYPES = tuple(["IMAGE"] * MAX_IMAGE_SLOTS) 183 | RETURN_NAMES = tuple([f"image_{i}" for i in range(1, MAX_IMAGE_SLOTS + 1)]) 184 | FUNCTION = "retrieve_images" 185 | CATEGORY = "tksw_node" 186 | 187 | @classmethod 188 | def IS_CHANGED(cls, **kwargs: Any) -> float: 189 | return float("NaN") 190 | 191 | def retrieve_images(self, remove_after_retrieval: bool, **kwargs: Any) -> Tuple[torch.Tensor, ...]: 192 | outputs: List[torch.Tensor] = [] 193 | node_fallback_image: Optional[torch.Tensor] = kwargs.get("fallback_image") 194 | 195 | for i in range(1, MAX_IMAGE_SLOTS + 1): 196 | image_id_slot_key = f"image_id_{i}" 197 | image_id: int = kwargs[image_id_slot_key] 198 | 199 | retrieved_image: Optional[torch.Tensor] = None 200 | log_prefix = f"[Retrieve Multiple (Memory)] Slot {i} (ID {image_id})" 201 | 202 | if image_id in shared_image_pool: 203 | retrieved_image = shared_image_pool[image_id].clone() 204 | print(f"{log_prefix}: Image retrieved. Shape: {retrieved_image.shape}") 205 | 206 | if remove_after_retrieval: 207 | del shared_image_pool[image_id] 208 | print(f"{log_prefix}: Image removed from pool.") 209 | else: 210 | print(f"{log_prefix}: Image kept in pool.") 211 | else: 212 | print(f"{log_prefix}: Image not found in memory.") 213 | 214 | if retrieved_image is None: 215 | if node_fallback_image is not None: 216 | retrieved_image = node_fallback_image.clone() 217 | print(f"{log_prefix}: Using node fallback image.") 218 | else: 219 | default_img = torch.zeros((1, 64, 64, 3), dtype=torch.float32, device="cpu") 220 | retrieved_image = default_img 221 | print(f"{log_prefix}: No node fallback. Using default black image.") 222 | 223 | outputs.append(retrieved_image) 224 | 225 | return tuple(outputs) 226 | -------------------------------------------------------------------------------- /lora_loader_elemental.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import comfy.utils 3 | import folder_paths 4 | import comfy.sd 5 | from safetensors.torch import safe_open, save_file 6 | import json 7 | import io 8 | import os 9 | import re 10 | 11 | 12 | class LoraLoaderElemental: 13 | @classmethod 14 | def INPUT_TYPES(s): 15 | return { 16 | "required": { 17 | "lora_name": (folder_paths.get_filename_list("loras"),), 18 | "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), 19 | "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), 20 | }, 21 | "optional": { 22 | "model": ("MODEL",), 23 | "clip": ("CLIP",), 24 | "lora_strength_string": ("STRING", {"multiline": True, "default": ""}), 25 | "save_lora": ("BOOLEAN", {"default": False}), 26 | "save_name": ("STRING", {"default": "processed_lora"}), 27 | "remove_unspecified_keys": ("BOOLEAN", {"default": False}), 28 | "remove_zero_strength_keys": ("BOOLEAN", {"default": False}), 29 | "regex_mode": ("BOOLEAN", {"default": False}), 30 | } 31 | } 32 | 33 | @classmethod 34 | def IS_CHANGED(cls, **kwargs): 35 | return float("NaN") 36 | 37 | RETURN_TYPES = ("MODEL", "CLIP", "LORA", "STRING", "STRING") 38 | OPTIONAL_INPUTS = ("MODEL", "CLIP") 39 | RETURN_NAMES = ("model", "clip", "processed_lora", "metadata", "lora_keys") 40 | 41 | FUNCTION = "load_lora" 42 | CATEGORY = "tksw_node" 43 | 44 | def _parse_strength_string(self, strength_string): 45 | lora_strengths = {} 46 | with io.StringIO(strength_string) as f: 47 | for index, line in enumerate(f): 48 | line = line.strip() 49 | if line and "=" in line: 50 | try: 51 | key, value = line.split("=", 1) 52 | key = key.strip() 53 | value = float(value.strip()) 54 | lora_strengths[key] = (value, index) 55 | except ValueError: 56 | print(f"Invalid line in strength string: ") 57 | return lora_strengths 58 | 59 | def _save_processed_lora(self, lora, save_name): 60 | if not save_name.endswith(".safetensors"): 61 | save_name += ".safetensors" 62 | lora_path = os.path.join(folder_paths.get_folder_paths("loras")[0], save_name) 63 | 64 | metadata = lora.pop("metadata", {}) if isinstance(lora, dict) else {} 65 | try: 66 | save_file(lora, lora_path, metadata) 67 | print(f"Processed LoRA saved to: {lora_path}") 68 | except Exception as e: 69 | print(f"Error saving processed LoRA: {e}") 70 | if metadata: 71 | lora["metadata"] = metadata 72 | 73 | def _get_lora_keys_string(self, lora): 74 | if not isinstance(lora, dict): 75 | return "" 76 | 77 | prefixes = sorted(list(set(key.split(".")[0] for key in lora if key != "metadata"))) 78 | return "\n".join(prefixes) 79 | 80 | def load_lora(self, lora_name, strength_model, strength_clip, model=None, clip=None, 81 | lora_strength_string="", save_lora=False, save_name="processed_lora", 82 | remove_unspecified_keys=False, remove_zero_strength_keys=False, regex_mode=False): 83 | 84 | lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) 85 | 86 | if model is None and clip is None: 87 | raise ValueError("Either 'model' or 'clip' must be provided.") 88 | 89 | if strength_model == 0 and strength_clip == 0: 90 | try: 91 | with safe_open(lora_path, framework="pt", device="cpu") as f: 92 | lora_metadata = f.metadata() 93 | lora = comfy.utils.load_torch_file(lora_path, safe_load=True) 94 | except Exception as e: 95 | print(f"Error reading LoRA metadata or loading: {e}") 96 | lora_metadata = {} 97 | lora = {} 98 | lora_keys_string = self._get_lora_keys_string(lora) 99 | return (model, clip, None, json.dumps(lora_metadata, indent=4), lora_keys_string) 100 | 101 | try: 102 | with safe_open(lora_path, framework="pt", device="cpu") as f: 103 | lora_metadata = f.metadata() 104 | lora = comfy.utils.load_torch_file(lora_path, safe_load=True) 105 | 106 | except Exception as e: 107 | print(f"Error loading LoRA file: {e}") 108 | return (model, clip, None, None, "") 109 | 110 | lora_strengths = {} 111 | if lora_strength_string: 112 | lora_strengths = self._parse_strength_string(lora_strength_string) 113 | 114 | extended_lora = {} 115 | for key, value in lora.items(): 116 | if key == "metadata": 117 | extended_lora[key] = {"value": value} 118 | elif key.endswith((".lora_down.weight", ".lora_up.weight")): 119 | extended_lora[key] = {"strength": None, "specified": False} 120 | else: 121 | extended_lora[key] = {"strength": None, "specified": False} 122 | 123 | for strength_key, (strength, index) in lora_strengths.items(): 124 | for lora_key in list(extended_lora.keys()): 125 | if regex_mode: 126 | try: 127 | if re.fullmatch(strength_key, lora_key): 128 | extended_lora[lora_key]["strength"] = strength 129 | extended_lora[lora_key]["specified"] = True 130 | except re.error as e: 131 | print(f"Invalid regular expression '{strength_key}': {e}") 132 | continue 133 | else: 134 | if lora_key.startswith(strength_key): 135 | extended_lora[lora_key]["strength"] = strength 136 | extended_lora[lora_key]["specified"] = True 137 | 138 | new_lora = {} 139 | for key, data in extended_lora.items(): 140 | if key == "metadata": 141 | new_lora[key] = data["value"] 142 | continue 143 | 144 | if key.endswith((".lora_down.weight", ".lora_up.weight")): 145 | if data["specified"]: 146 | if remove_zero_strength_keys and data["strength"] == 0: 147 | continue 148 | new_lora[key] = lora[key] * data["strength"] 149 | 150 | elif not remove_unspecified_keys: 151 | new_lora[key] = lora[key] 152 | else: 153 | if not remove_unspecified_keys: 154 | new_lora[key] = lora[key] 155 | 156 | 157 | model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, new_lora, strength_model, strength_clip) 158 | lora_metadata_string = json.dumps(lora_metadata, indent=4) 159 | 160 | if save_lora: 161 | self._save_processed_lora(new_lora, save_name) 162 | 163 | lora_keys_string = self._get_lora_keys_string(new_lora) 164 | 165 | return (model_lora, clip_lora, new_lora, lora_metadata_string, lora_keys_string) -------------------------------------------------------------------------------- /lora_mixer_elemental.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import comfy.utils 3 | import folder_paths 4 | import comfy.sd 5 | from safetensors.torch import safe_open 6 | import random 7 | import os 8 | import json 9 | 10 | class LoraMixerElemental: 11 | MAX_LORAS = 8 12 | 13 | @classmethod 14 | def INPUT_TYPES(cls): 15 | lora_names = ["None"] + folder_paths.get_filename_list("loras") 16 | input_types = { 17 | "required": { 18 | "model_strength": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), 19 | "clip_strength": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), 20 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 21 | "save_mixed_lora": (["Off", "First Only", "All"], {"default": "Off"}), 22 | "save_name": ("STRING", {"default": "mixed_lora"}), 23 | "key_selection": (["All Available Keys", "Common Keys Only"], {"default": "All Available Keys"}), 24 | "key_strength_randomization": (["Off", "Per Key"], {"default": "Off"}), 25 | "key_strength_min": ("FLOAT", {"default": 0.0, "min": -10.0, "max": 10.0, "step": 0.01}), 26 | "key_strength_max": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), 27 | "multi_mix": (["Off", "On", "Reuse Keys"], {"default": "Off"}), 28 | "num_mix_passes": ("INT", {"default": 2, "min": 2, "max": 10}), 29 | "mix_passes_strength_randomization": (["Off", "On"], {"default": "Off"}), 30 | "mix_pass_strength_min": ("FLOAT", {"default": 0.5, "min": -10.0, "max": 10.0, "step": 0.01}), 31 | "mix_pass_strength_max": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), 32 | }, 33 | "optional": { 34 | "model": ("MODEL",), 35 | "clip": ("CLIP",), 36 | } 37 | } 38 | for i in range(1, cls.MAX_LORAS + 1): 39 | input_types["optional"][f"lora_name_{i}"] = (lora_names, {"default": "None"}) 40 | return input_types 41 | 42 | @classmethod 43 | def IS_CHANGED(cls, **kwargs): 44 | return float("NaN") 45 | 46 | RETURN_TYPES = ("MODEL", "CLIP", "LORA", "STRING", "STRING") 47 | RETURN_NAMES = ("model", "clip", "mixed_lora", "lora_keys", "all_lora_keys") 48 | FUNCTION = "mix_loras" 49 | CATEGORY = "tksw_node" 50 | 51 | def mix_loras(self, model_strength, clip_strength, seed, save_mixed_lora, save_name, key_selection, 52 | key_strength_randomization, key_strength_min, key_strength_max, multi_mix, num_mix_passes, 53 | mix_passes_strength_randomization, mix_pass_strength_min, mix_pass_strength_max, model=None, clip=None, **kwargs): 54 | 55 | if model is None and clip is None: 56 | raise ValueError("Either 'model' or 'clip' must be provided.") 57 | 58 | lora_names = [ 59 | kwargs[key] for key in kwargs 60 | if key.startswith("lora_name_") and kwargs[key] != "None" 61 | ] 62 | if not lora_names: 63 | return (model, clip, None, "", "{}") 64 | 65 | loras = [] 66 | lora_name_map = {} 67 | lora_dims = {} 68 | 69 | for i, lora_name in enumerate(lora_names): 70 | lora_path = folder_paths.get_full_path("loras", lora_name) 71 | try: 72 | lora = comfy.utils.load_torch_file(lora_path, safe_load=True) 73 | loras.append(lora) 74 | lora_name_map[f"lora{i+1}"] = lora_name 75 | lora_dims[f"lora{i+1}"] = {} 76 | for key, value in lora.items(): 77 | if isinstance(value, torch.Tensor): 78 | lora_dims[f"lora{i+1}"][key] = value.shape 79 | except Exception as e: 80 | print(f"Error loading LoRA {lora_name}: {e}") 81 | continue 82 | 83 | if not loras: 84 | return (model, clip, None, "", "{}") 85 | 86 | if key_selection == "Common Keys Only": 87 | common_keys = set(loras[0].keys()) 88 | for lora in loras[1:]: 89 | common_keys.intersection_update(lora.keys()) 90 | filtered_keys = {key for key in common_keys if key.endswith((".lora_down.weight", ".lora_up.weight"))} 91 | elif key_selection == "All Available Keys": 92 | all_keys = set() 93 | for lora in loras: 94 | all_keys.update(lora.keys()) 95 | filtered_keys = {key for key in all_keys if key.endswith((".lora_down.weight", ".lora_up.weight"))} 96 | else: 97 | raise ValueError(f"Invalid key_selection_mode: {key_selection}") 98 | 99 | random.seed(seed) 100 | 101 | current_model = model 102 | current_clip = clip 103 | first_mixed_lora = None 104 | first_lora_keys_string = "" 105 | all_lora_keys = {} 106 | 107 | for mix_num in range(num_mix_passes if multi_mix != "Off" else 1): 108 | mixed_lora = {} 109 | key_source_map = {} 110 | 111 | temp_filtered_keys = filtered_keys if multi_mix == "Reuse Keys" else filtered_keys.copy() 112 | 113 | for down_key in list(temp_filtered_keys): 114 | if not down_key.endswith(".lora_down.weight"): 115 | continue 116 | 117 | up_key = down_key.replace(".lora_down.weight", ".lora_up.weight") 118 | if up_key not in temp_filtered_keys: 119 | continue 120 | 121 | available_down_loras = [] 122 | for i, lora in enumerate(loras): 123 | if down_key in lora: 124 | available_down_loras.append((i, lora)) 125 | if not available_down_loras: 126 | continue 127 | 128 | selected_down_index, selected_down_lora = random.choice(available_down_loras) 129 | 130 | available_up_loras = [] 131 | for i, lora in enumerate(loras): 132 | if up_key in lora: 133 | if lora_dims[f"lora{i+1}"][up_key] == lora_dims[f"lora{selected_down_index+1}"][up_key]: 134 | available_up_loras.append((i, lora)) 135 | 136 | if not available_up_loras: 137 | temp_filtered_keys.discard(down_key) 138 | temp_filtered_keys.discard(up_key) 139 | continue 140 | 141 | selected_up_index, selected_up_lora = random.choice(available_up_loras) 142 | 143 | if key_strength_randomization == "Per Key": 144 | strength = random.uniform(key_strength_min, key_strength_max) 145 | mixed_lora[down_key] = selected_down_lora[down_key] * strength 146 | mixed_lora[up_key] = selected_up_lora[up_key] * strength 147 | else: 148 | mixed_lora[down_key] = selected_down_lora[down_key] 149 | mixed_lora[up_key] = selected_up_lora[up_key] 150 | 151 | key_source_map[down_key] = lora_name_map[f"lora{selected_down_index + 1}"] 152 | key_source_map[up_key] = lora_name_map[f"lora{selected_up_index + 1}"] 153 | 154 | if multi_mix != "Reuse Keys": 155 | temp_filtered_keys.discard(down_key) 156 | temp_filtered_keys.discard(up_key) 157 | 158 | if not mixed_lora: 159 | continue 160 | 161 | if multi_mix != "Off" and mix_passes_strength_randomization == "On": 162 | mix_strength_model = random.uniform(mix_pass_strength_min, mix_pass_strength_max) 163 | mix_strength_clip = random.uniform(mix_pass_strength_min, mix_pass_strength_max) 164 | else: 165 | mix_strength_model = model_strength 166 | mix_strength_clip = clip_strength 167 | 168 | current_model, current_clip = comfy.sd.load_lora_for_models(current_model, current_clip, mixed_lora, mix_strength_model, mix_strength_clip) 169 | 170 | sorted_key_source_pairs = sorted(key_source_map.items(), key=lambda item: item[0]) 171 | all_lora_keys[f"mix_pass_{mix_num + 1}"] = {key: source for key, source in sorted_key_source_pairs} 172 | 173 | 174 | if mix_num == 0: 175 | first_mixed_lora = mixed_lora 176 | first_lora_keys_string = "\n".join([f"{key}:{source}" for key, source in sorted_key_source_pairs]) 177 | 178 | if save_mixed_lora == "All": 179 | current_save_name = f"{save_name}_{mix_num + 1}" 180 | if not current_save_name.endswith(".safetensors"): 181 | current_save_name += ".safetensors" 182 | lora_path = os.path.join(folder_paths.get_folder_paths("loras")[0], current_save_name) 183 | try: 184 | lora_to_save = {k: v for k, v in mixed_lora.items() if isinstance(v, torch.Tensor)} 185 | comfy.utils.save_torch_file(lora_to_save, lora_path) 186 | print(f"Mixed LoRA {mix_num + 1} saved to: {lora_path}") 187 | except Exception as e: 188 | print(f"Error saving mixed LoRA {mix_num + 1}: {e}") 189 | 190 | if save_mixed_lora == "First Only" and first_mixed_lora: 191 | if not save_name.endswith(".safetensors"): 192 | save_name += ".safetensors" 193 | lora_path = os.path.join(folder_paths.get_folder_paths("loras")[0], save_name) 194 | try: 195 | lora_to_save = {k: v for k, v in first_mixed_lora.items() if isinstance(v, torch.Tensor)} 196 | comfy.utils.save_torch_file(lora_to_save,lora_path) 197 | print(f"First mixed LoRA saved to: {lora_path}") 198 | except Exception as e: 199 | print(f"Error saving first mixed LoRA: {e}") 200 | 201 | return (current_model, current_clip, first_mixed_lora, first_lora_keys_string, json.dumps(all_lora_keys, indent=4)) -------------------------------------------------------------------------------- /lora_selector.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import torch 4 | import collections 5 | import comfy.utils 6 | import comfy.sd 7 | from folder_paths import get_filename_list, supported_pt_extensions, get_full_path 8 | 9 | LORA_SLOT_COUNT = 8 10 | LORA_EXTENSIONS = [ext.lower() for ext in supported_pt_extensions] 11 | 12 | class LoraSelector: 13 | def __init__(self): 14 | self.round_robin_index = 0 15 | self.last_candidate_list = None 16 | self.last_lora_folder = "" 17 | self.cached_folder_loras = [] 18 | self.current_lora = None 19 | self.remaining_executions = 0 20 | self.loaded_lora_cache = collections.OrderedDict() 21 | self.lora_size_cache = {} 22 | self.current_cache_size_bytes = 0 23 | 24 | @classmethod 25 | def INPUT_TYPES(cls): 26 | slot_lora_list = [""] + get_filename_list("loras") 27 | inputs = { 28 | "required": { 29 | "strength_model": ("FLOAT", {"default": 1.00, "min": -10.00, "max": 10.00, "step": 0.01}), 30 | "strength_clip": ("FLOAT", {"default": 1.00, "min": -10.00, "max": 10.00, "step": 0.01}), 31 | "mode": (["random", "round-robin"], {"default": "random"}), 32 | "switch_interval": ("INT", {"default": 1, "min": 1, "max": 9999}), 33 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 34 | "reset_state": ("BOOLEAN", {"default": False}), 35 | "cache_limit_gb": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 128.0, "step": 0.1}), 36 | "lora_folder": ("STRING", {"multiline": False, "default": ""}), 37 | **{f"lora_{i}": (slot_lora_list, {"default": ""}) for i in range(LORA_SLOT_COUNT)} 38 | }, 39 | "optional": { 40 | "model": ("MODEL",), 41 | "clip": ("CLIP",), 42 | } 43 | } 44 | return inputs 45 | 46 | RETURN_TYPES = ("MODEL", "CLIP", "STRING") 47 | RETURN_NAMES = ("MODEL", "CLIP", "selected_lora_name") 48 | FUNCTION = "select_and_apply_lora" 49 | CATEGORY = "tksw_node" 50 | 51 | def _scan_lora_folder(self, folder_path): 52 | lora_files = [] 53 | try: 54 | print(f"[LoraSelector] Scanning folder for LoRA files: {folder_path}") 55 | if not os.path.isdir(folder_path): 56 | print(f"[LoraSelector] Warning: Specified LoRA folder does not exist or is not a directory: {folder_path}") 57 | return [] 58 | for filename in os.listdir(folder_path): 59 | file_path = os.path.join(folder_path, filename) 60 | if os.path.isfile(file_path): 61 | if any(filename.lower().endswith(ext) for ext in LORA_EXTENSIONS): 62 | lora_files.append(filename) 63 | print(f"[LoraSelector] Found {len(lora_files)} LoRA files in folder.") 64 | return sorted(lora_files) 65 | except Exception as e: 66 | print(f"[LoraSelector] Error scanning folder '{folder_path}': {e}") 67 | return [] 68 | 69 | def select_and_apply_lora(self, strength_model, strength_clip, mode, switch_interval, seed, reset_state, cache_limit_gb, lora_folder, model=None, clip=None, **kwargs): 70 | current_folder_loras = [] 71 | clean_lora_folder = lora_folder.strip() 72 | if clean_lora_folder: 73 | if clean_lora_folder != self.last_lora_folder: 74 | self.cached_folder_loras = self._scan_lora_folder(clean_lora_folder) 75 | self.last_lora_folder = clean_lora_folder 76 | current_folder_loras = self.cached_folder_loras 77 | else: 78 | if self.last_lora_folder != "": print("[LoraSelector] LoRA folder cleared.") 79 | self.cached_folder_loras = [] 80 | self.last_lora_folder = "" 81 | current_folder_loras = [] 82 | slot_loras = [] 83 | for i in range(LORA_SLOT_COUNT): 84 | lora_name = kwargs.get(f"lora_{i}", "") 85 | if lora_name and lora_name != "": slot_loras.append(lora_name) 86 | candidate_loras = sorted(list(set(slot_loras + current_folder_loras))) 87 | 88 | selected_lora_name = "None" 89 | chosen_lora = None 90 | num_candidates = len(candidate_loras) 91 | 92 | if candidate_loras: 93 | needs_reset = False 94 | reset_reason = "" 95 | if reset_state: 96 | needs_reset = True 97 | reset_reason = "Manual reset requested." 98 | elif self.last_candidate_list is not None and candidate_loras != self.last_candidate_list: 99 | needs_reset = True 100 | reset_reason = "LoRA candidate list changed." 101 | elif self.current_lora is not None and self.current_lora not in candidate_loras: 102 | needs_reset = True 103 | reset_reason = f"Current LoRA '{self.current_lora}' not in candidate list." 104 | 105 | if needs_reset: 106 | print(f"[LoraSelector] State reset triggered: {reset_reason}") 107 | self.current_lora = None 108 | self.remaining_executions = 0 109 | self.round_robin_index = 0 110 | if reset_state: 111 | print("[LoraSelector] Clearing LoRA data cache due to manual reset.") 112 | self.loaded_lora_cache.clear() 113 | self.lora_size_cache.clear() 114 | self.current_cache_size_bytes = 0 115 | self.last_candidate_list = list(candidate_loras) 116 | else: 117 | if self.last_candidate_list is None or candidate_loras != self.last_candidate_list: 118 | self.last_candidate_list = list(candidate_loras) 119 | 120 | current_switch_interval = max(1, switch_interval) 121 | if self.remaining_executions > 0 and self.current_lora is not None: 122 | chosen_lora = self.current_lora 123 | self.remaining_executions -= 1 124 | print(f"[LoraSelector] Continuing LoRA: '{chosen_lora}' (Remaining: {self.remaining_executions} / Interval: {current_switch_interval})") 125 | else: 126 | print(f"[LoraSelector] Switching LoRA (Interval: {current_switch_interval})") 127 | if mode == "round-robin": 128 | current_index = self.round_robin_index % num_candidates 129 | chosen_lora = candidate_loras[current_index] 130 | self.round_robin_index += 1 131 | print(f"[LoraSelector] Mode: round-robin switch (Index: {current_index}, Next RR Base: {self.round_robin_index})") 132 | elif mode == "random": 133 | random.seed(seed) 134 | available_choices = [lora for lora in candidate_loras if lora != self.current_lora] 135 | if not available_choices and num_candidates > 0: available_choices = candidate_loras 136 | chosen_lora = random.choice(available_choices) 137 | print(f"[LoraSelector] Mode: random switch (Seed: {seed})") 138 | else: 139 | print(f"[LoraSelector] Warning: Unknown mode '{mode}'. Falling back to random.") 140 | random.seed(seed) 141 | chosen_lora = random.choice(candidate_loras) 142 | 143 | self.current_lora = chosen_lora 144 | self.remaining_executions = current_switch_interval - 1 145 | print(f"[LoraSelector] Switched to LoRA: '{chosen_lora}' (Use for {current_switch_interval} times, Remaining: {self.remaining_executions})") 146 | 147 | 148 | output_model = model 149 | output_clip = clip 150 | cache_enabled = cache_limit_gb > 0 151 | cache_limit_bytes = int(cache_limit_gb * (1024**3)) if cache_enabled else 0 152 | 153 | if chosen_lora: 154 | selected_lora_name = chosen_lora 155 | if output_model is None and output_clip is None: 156 | print(f"[LoraSelector] Warning: LoRA '{chosen_lora}' selected, but no Model or Clip input provided.") 157 | return (output_model, output_clip, selected_lora_name) 158 | 159 | print(f"[LoraSelector] Applying LoRA: '{chosen_lora}'") 160 | effective_strength_model = strength_model if output_model is not None else 0.0 161 | effective_strength_clip = strength_clip if output_clip is not None else 0.0 162 | print(f"[LoraSelector] Effective Strength - Model: {effective_strength_model if model else 'N/A'}, Clip: {effective_strength_clip if clip else 'N/A'}") 163 | 164 | try: 165 | lora_data = None 166 | cache_hit = False 167 | 168 | if cache_enabled and chosen_lora in self.loaded_lora_cache: 169 | lora_data = self.loaded_lora_cache[chosen_lora] 170 | self.loaded_lora_cache.move_to_end(chosen_lora) 171 | print(f"[LoraSelector] Using cached LoRA data for '{chosen_lora}'. Cache Size: {self.current_cache_size_bytes / (1024**3):.2f}/{cache_limit_gb:.1f} GB") 172 | cache_hit = True 173 | 174 | if not cache_hit: 175 | lora_path = get_full_path("loras", chosen_lora) 176 | if not lora_path: 177 | raise FileNotFoundError(f"LoRA file not found in known paths: {chosen_lora}") 178 | 179 | print(f"[LoraSelector] Loading LoRA data from: {lora_path}") 180 | lora_data = comfy.utils.load_torch_file(lora_path, safe_load=True) 181 | 182 | if cache_enabled and lora_data is not None: 183 | try: 184 | lora_file_size = os.path.getsize(lora_path) 185 | 186 | while self.current_cache_size_bytes + lora_file_size > cache_limit_bytes and self.loaded_lora_cache: 187 | oldest_key, _ = self.loaded_lora_cache.popitem(last=False) 188 | removed_size = self.lora_size_cache.pop(oldest_key, 0) # サイズ情報も削除 189 | self.current_cache_size_bytes -= removed_size 190 | print(f"[LoraSelector] Cache evicted '{oldest_key}' (Size: {removed_size / (1024**2):.1f} MB) to make space. New Size: {self.current_cache_size_bytes / (1024**3):.2f}/{cache_limit_gb:.1f} GB") 191 | 192 | if self.current_cache_size_bytes + lora_file_size <= cache_limit_bytes: 193 | self.loaded_lora_cache[chosen_lora] = lora_data 194 | self.lora_size_cache[chosen_lora] = lora_file_size 195 | self.current_cache_size_bytes += lora_file_size 196 | print(f"[LoraSelector] Cached LoRA data for '{chosen_lora}'. Cache Size: {self.current_cache_size_bytes / (1024**3):.2f}/{cache_limit_gb:.1f} GB. Items: {len(self.loaded_lora_cache)}") 197 | else: 198 | print(f"[LoraSelector] Warning: LoRA '{chosen_lora}' (Size: {lora_file_size / (1024**2):.1f} MB) is too large for cache limit ({cache_limit_gb:.1f} GB). Not caching.") 199 | except FileNotFoundError: 200 | print(f"[LoraSelector] Warning: Could not get file size for '{lora_path}'. Caching without size check might be unreliable.") 201 | if cache_enabled: 202 | self.loaded_lora_cache[chosen_lora] = lora_data 203 | print(f"[LoraSelector] Cached LoRA data for '{chosen_lora}' (size unknown).") 204 | except Exception as e: 205 | print(f"[LoraSelector] Error during cache management for '{chosen_lora}': {e}") 206 | 207 | if lora_data is not None: 208 | applied_model, applied_clip = comfy.sd.load_lora_for_models( 209 | output_model, output_clip, lora_data, 210 | effective_strength_model, effective_strength_clip 211 | ) 212 | output_model = applied_model 213 | output_clip = applied_clip 214 | else: 215 | print(f"[LoraSelector] Warning: LoRA data for '{chosen_lora}' is None. Skipping application.") 216 | 217 | 218 | except Exception as e: 219 | print(f"[LoraSelector] ERROR during LoRA loading or application for '{chosen_lora}': {e}") 220 | selected_lora_name = f"ERROR applying {chosen_lora}" 221 | return (model, clip, selected_lora_name) 222 | else: 223 | print("[LoraSelector] No LoRA was chosen.") 224 | 225 | else: 226 | print("[LoraSelector] No LoRA candidates found. Passing through inputs.") 227 | self.current_lora = None 228 | self.remaining_executions = 0 229 | self.round_robin_index = 0 230 | self.last_candidate_list = None 231 | 232 | return (output_model, output_clip, selected_lora_name) -------------------------------------------------------------------------------- /lora_weight_randomizer.py: -------------------------------------------------------------------------------- 1 | from folder_paths import get_filename_list 2 | from nodes import LoraLoader 3 | import random 4 | import torch 5 | 6 | LORA_COUNT = 8 7 | 8 | class LoraWeightRandomizer: 9 | def __init__(self): 10 | self.loras = [LoraLoader() for _ in range(LORA_COUNT)] 11 | 12 | @classmethod 13 | def INPUT_TYPES(cls): 14 | args = { 15 | "model": ("MODEL",), 16 | "clip": ("CLIP",), 17 | "total_strength": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 10.00, "step": 0.01}), 18 | "max_single_strength": ("FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.01}), 19 | "randomize_total_strength": ("BOOLEAN", {"default": False}), 20 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 21 | } 22 | arg_lora_name = ([""] + get_filename_list("loras"),) 23 | for i in range(LORA_COUNT): 24 | args["{}:lora".format(i)] = arg_lora_name 25 | return {"required": args} 26 | 27 | def apply(self, model, clip, total_strength, max_single_strength, randomize_total_strength, seed, **kwargs): 28 | selected_loras = [] 29 | for i in range(LORA_COUNT): 30 | lora_name = kwargs["{}:lora".format(i)] 31 | if lora_name != "": 32 | selected_loras.append((lora_name, i)) 33 | 34 | if not selected_loras: 35 | return (model, clip, "") 36 | 37 | torch.manual_seed(seed) 38 | random.seed(seed) 39 | 40 | if randomize_total_strength: 41 | total_strength = round(random.uniform(0, total_strength), 2) 42 | 43 | num_selected = len(selected_loras) 44 | strengths = [0.00] * num_selected 45 | remaining_strength = total_strength 46 | 47 | allocation_order = list(range(num_selected)) 48 | random.shuffle(allocation_order) 49 | 50 | for i in allocation_order[:-1]: 51 | max_allowed = min(remaining_strength, max_single_strength) 52 | strength = round(random.uniform(0, max_allowed), 2) 53 | strengths[i] = strength 54 | remaining_strength -= strength 55 | 56 | strengths[allocation_order[-1]] = round(min(remaining_strength, max_single_strength), 2) 57 | 58 | allocated_strength_total = sum(strengths) 59 | if allocated_strength_total < total_strength: 60 | diff = total_strength - allocated_strength_total 61 | num_under_max = sum(1 for s in strengths if s < max_single_strength) 62 | if num_under_max > 0: 63 | increment = diff / num_under_max 64 | for i in range(len(strengths)): 65 | if strengths[i] < max_single_strength: 66 | add = min(increment, max_single_strength - strengths[i]) 67 | strengths[i] = round(strengths[i] + add, 2) 68 | 69 | reordered_strengths = [0.00] * num_selected 70 | for i, original_index in enumerate(allocation_order): 71 | reordered_strengths[original_index] = strengths[i] 72 | strengths = reordered_strengths 73 | 74 | output_text = f"LoraWeightRandomizer Settings:\n" 75 | output_text += f" Total Strength: {total_strength:.2f}\n" 76 | output_text += f" Max Single Strength: {max_single_strength:.2f}\n" 77 | output_text += f" Randomize Total Strength: {randomize_total_strength}\n" 78 | output_text += f" Seed: {seed}\n" 79 | output_text += "LoRA Weights:\n" 80 | 81 | for (lora_name, index), strength in zip(selected_loras, strengths): 82 | output_text += f" - {lora_name}: {strength:.2f}\n" 83 | (model, clip) = self.loras[index].load_lora(model, clip, lora_name, strength, strength) 84 | 85 | 86 | return (model, clip, output_text) 87 | 88 | 89 | RETURN_TYPES = ("MODEL", "CLIP", "STRING") 90 | RETURN_NAMES = ("MODEL", "CLIP", "settings") 91 | FUNCTION = "apply" 92 | CATEGORY = "tksw_node" -------------------------------------------------------------------------------- /quantized_lora_loader.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import comfy.utils 3 | import folder_paths 4 | import comfy.sd 5 | from safetensors.torch import safe_open, save_file 6 | import json 7 | import os 8 | import numpy as np 9 | 10 | class QuantizedLoraLoader: 11 | @classmethod 12 | def INPUT_TYPES(s): 13 | return { 14 | "required": { 15 | "lora_name": (folder_paths.get_filename_list("loras"),), 16 | "quantization_bits": ("INT", {"default": 8, "min": 2, "max": 32, "step": 1}), 17 | "strength_model": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), 18 | "strength_clip": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), 19 | "quantization_iterations": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1}), 20 | "stepwise_quantization": ("BOOLEAN", {"default": False}), 21 | "quantization_step_size": ("INT", {"default": 2, "min": 2, "max": 16, "step": 1}), 22 | "blend_mode": ("BOOLEAN", {"default": False}), 23 | "blend_factor": ("FLOAT", {"default": 0.5, "min": -10.0, "max": 10.0, "step": 0.01}), 24 | }, 25 | "optional": { 26 | "model": ("MODEL",), 27 | "clip": ("CLIP",), 28 | "save_quantized_lora": ("BOOLEAN", {"default": False}), 29 | "save_name": ("STRING", {"default": "quantized_lora"}), 30 | } 31 | } 32 | 33 | @classmethod 34 | def IS_CHANGED(cls, **kwargs): 35 | return float("NaN") 36 | 37 | RETURN_TYPES = ("MODEL", "CLIP", "LORA", "STRING") 38 | RETURN_NAMES = ("model", "clip", "quantized_lora", "metadata") 39 | OPTIONAL_INPUTS = ("MODEL", "CLIP") 40 | 41 | FUNCTION = "load_and_quantize_lora" 42 | CATEGORY = "tksw_node" 43 | 44 | def _quantize_tensor(self, tensor, bits): 45 | original_dtype = tensor.dtype 46 | min_val = tensor.min() 47 | max_val = tensor.max() 48 | 49 | if max_val - min_val == 0: 50 | return torch.zeros_like(tensor, dtype=original_dtype) 51 | 52 | scale = (max_val - min_val) / (2**bits - 1) 53 | zero_point = torch.round(-min_val / scale) 54 | 55 | quantized_tensor = torch.clamp(torch.round(tensor / scale + zero_point), 0, 2**bits - 1) 56 | dequantized_tensor = (quantized_tensor - zero_point) * scale 57 | 58 | return dequantized_tensor.to(original_dtype) 59 | 60 | def _save_processed_lora(self, lora, save_name, quantization_bits, quantization_iterations, stepwise_quantization, quantization_step_size, blend_mode, blend_factor): 61 | if not save_name.endswith(".safetensors"): 62 | save_name += ".safetensors" 63 | base, ext = os.path.splitext(save_name) 64 | save_name = f"{base}_q{quantization_bits}_i{quantization_iterations}_s{stepwise_quantization}_ss{quantization_step_size}_b{blend_mode}_bf{blend_factor}{ext}" 65 | lora_path = os.path.join(folder_paths.get_folder_paths("loras")[0], save_name) 66 | 67 | metadata = lora.pop("metadata", {}) if isinstance(lora, dict) else {} 68 | 69 | if isinstance(metadata, dict): 70 | metadata["quantization_bits"] = str(quantization_bits) 71 | metadata["quantization_iterations"] = str(quantization_iterations) 72 | metadata["stepwise_quantization"] = str(stepwise_quantization) 73 | metadata["quantization_step_size"] = str(quantization_step_size) 74 | metadata["blend_mode"] = str(blend_mode) 75 | metadata["blend_factor"] = str(blend_factor) 76 | 77 | for k, v in metadata.items(): 78 | if not isinstance(v, (str, int, float, list)): 79 | metadata[k] = str(v) 80 | else: 81 | metadata = { 82 | "quantization_bits": str(quantization_bits), 83 | "quantization_iterations": str(quantization_iterations), 84 | "stepwise_quantization": str(stepwise_quantization), 85 | "quantization_step_size": str(quantization_step_size), 86 | "blend_mode": str(blend_mode), 87 | "blend_factor": str(blend_factor) 88 | } 89 | 90 | try: 91 | save_file(lora, lora_path, metadata=metadata) 92 | print(f"Quantized LoRA saved to: {lora_path}") 93 | except Exception as e: 94 | print(f"Error saving quantized LoRA: {e}") 95 | 96 | if metadata: 97 | lora["metadata"] = metadata 98 | 99 | def load_and_quantize_lora(self, lora_name, quantization_bits, strength_model, strength_clip, 100 | model=None, clip=None, save_quantized_lora=False, save_name="quantized_lora", 101 | quantization_iterations=1, stepwise_quantization=False, quantization_step_size=1, 102 | blend_mode=False, blend_factor=0.5): 103 | 104 | lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) 105 | 106 | try: 107 | with safe_open(lora_path, framework="pt", device="cpu") as f: 108 | lora_metadata = f.metadata() 109 | lora = comfy.utils.load_torch_file(lora_path, safe_load=True) 110 | 111 | except Exception as e: 112 | print(f"Error loading LoRA: {e}") 113 | return (model, clip, None, None) 114 | 115 | original_dtype = None 116 | for key, tensor in lora.items(): 117 | if "lora_down" in key or "lora_up" in key: 118 | if original_dtype is None: 119 | original_dtype = tensor.dtype 120 | elif original_dtype == torch.float16 and tensor.dtype == torch.bfloat16: 121 | original_dtype = torch.bfloat16 122 | 123 | original_bits = 32 if original_dtype == torch.float32 else 16 124 | quantized_lora = {} 125 | original_lora = lora.copy() 126 | 127 | 128 | if quantization_bits >= original_bits: 129 | stepwise_quantization = False 130 | 131 | if stepwise_quantization: 132 | for key, tensor in lora.items(): 133 | if key == "metadata": 134 | quantized_lora[key] = tensor 135 | elif "lora_down" in key or "lora_up" in key: 136 | temp_tensor = tensor 137 | current_bits = original_bits 138 | for _ in range(quantization_iterations): 139 | while current_bits > quantization_bits: 140 | current_bits = max(current_bits - quantization_step_size, quantization_bits) 141 | temp_tensor = self._quantize_tensor(temp_tensor, current_bits) 142 | quantized_lora[key] = temp_tensor 143 | else: 144 | quantized_lora[key] = tensor 145 | else: 146 | for key, tensor in lora.items(): 147 | if key == "metadata": 148 | quantized_lora[key] = tensor 149 | elif "lora_down" in key or "lora_up" in key: 150 | temp_tensor = tensor 151 | for _ in range(quantization_iterations): 152 | temp_tensor = self._quantize_tensor(temp_tensor, quantization_bits) 153 | quantized_lora[key] = temp_tensor 154 | else: 155 | quantized_lora[key] = tensor 156 | 157 | if blend_mode: 158 | blended_lora = {} 159 | for key in quantized_lora.keys(): 160 | if key == "metadata": 161 | blended_lora[key] = quantized_lora[key] 162 | elif "lora_down" in key or "lora_up" in key: 163 | blended_lora[key] = (1 - blend_factor) * quantized_lora[key] + blend_factor * original_lora[key] 164 | else: 165 | blended_lora[key] = quantized_lora[key] 166 | quantized_lora = blended_lora 167 | 168 | 169 | if model is not None and clip is not None: 170 | model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, quantized_lora, strength_model, strength_clip) 171 | elif model is not None: 172 | model_lora, clip_lora = comfy.sd.load_lora_for_models(model, None, quantized_lora, strength_model, 0.0) 173 | clip = None 174 | elif clip is not None: 175 | model_lora, clip_lora = comfy.sd.load_lora_for_models(None, clip, quantized_lora, 0.0, strength_clip) 176 | model = None 177 | else: 178 | model_lora, clip_lora = None, None 179 | 180 | if save_quantized_lora: 181 | self._save_processed_lora(lora=quantized_lora, save_name=save_name, quantization_bits=quantization_bits, 182 | quantization_iterations=quantization_iterations, stepwise_quantization=stepwise_quantization, 183 | quantization_step_size=quantization_step_size, 184 | blend_mode=blend_mode, blend_factor=blend_factor) 185 | 186 | lora_metadata_string = json.dumps(lora_metadata, indent=4) if lora_metadata else "{}" 187 | return (model_lora, clip_lora, quantized_lora, lora_metadata_string) -------------------------------------------------------------------------------- /random_word_replacer.py: -------------------------------------------------------------------------------- 1 | import re 2 | import random 3 | import os 4 | 5 | class RandomWordReplacer: 6 | @classmethod 7 | def INPUT_TYPES(s): 8 | return { 9 | "required": { 10 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 11 | }, 12 | "optional": { 13 | "input_text": ("STRING", {"multiline": True, "default": "", "forceInput": True}), 14 | "replace_specs": ("STRING", {"multiline": True, "default": ""}), 15 | "replace_specs_file": ("STRING", {"multiline": False, "default": ""}), 16 | "replace_specs_folder": ("STRING", {"multiline": False, "default": ""}), 17 | } 18 | } 19 | 20 | RETURN_TYPES = ("STRING",) 21 | RETURN_NAMES = ("processed_text",) 22 | FUNCTION = "replace_words" 23 | CATEGORY = "tksw_node" 24 | 25 | def replace_words(self, seed, input_text=None, replace_specs_file="", replace_specs="", replace_specs_folder=""): 26 | if seed: 27 | random.seed(seed) 28 | 29 | if not input_text: 30 | return ("",) 31 | 32 | word_groups = [] 33 | 34 | if replace_specs_folder: 35 | try: 36 | for filename in os.listdir(replace_specs_folder): 37 | if filename.endswith((".txt", ".csv")): 38 | with open(os.path.join(replace_specs_folder, filename), "r", encoding="utf-8") as f: 39 | words = [line.strip() for line in f if line.strip()] 40 | word_groups.append(words) 41 | except FileNotFoundError: 42 | return (f"Error: Folder not found: ", ) 43 | 44 | if replace_specs_file: 45 | try: 46 | with open(replace_specs_file, "r", encoding="utf-8") as f: 47 | for line in f: 48 | if line: 49 | words = [word.strip() for word in line.split(',')] 50 | if len(words) > 1: 51 | word_groups.append(words) 52 | except FileNotFoundError: 53 | return (f"Error: File not found: ", ) 54 | 55 | if replace_specs: 56 | for line in replace_specs.splitlines(): 57 | if line: 58 | words = [word.strip() for word in line.split(',')] 59 | if len(words) > 1: 60 | word_groups.append(words) 61 | 62 | processed_text = "" 63 | for line in input_text.splitlines(): 64 | processed_line = self.process_line(line, word_groups) 65 | processed_text += processed_line + "\n" 66 | return (processed_text.strip(),) 67 | 68 | def process_line(self, line, word_groups): 69 | processed_line = line 70 | for group in word_groups: 71 | for i, word in enumerate(group): 72 | if word: 73 | while word in processed_line: 74 | if len(group) == 1: 75 | replacement = word 76 | else: 77 | replacement = random.choice([w for w in group if w != word]) 78 | processed_line = processed_line.replace(word, replacement, 1) 79 | return processed_line -------------------------------------------------------------------------------- /text_combiner.py: -------------------------------------------------------------------------------- 1 | import re 2 | 3 | class TextCombiner: 4 | def __init__(self): 5 | self.text_log = [] 6 | 7 | @classmethod 8 | def INPUT_TYPES(s): 9 | return { 10 | "required": { 11 | "separator": ("STRING", {"default": ","}), 12 | "remember_log": ("BOOLEAN", {"default": True}), 13 | "max_log": ("INT", {"default": 10, "min": 0, "max": 1000}), 14 | "allow_duplicate_log": ("BOOLEAN", {"default": True}), 15 | "use_regex": ("BOOLEAN", {"default": False}), 16 | }, 17 | "optional": { 18 | "text_1": ("STRING", {"multiline": False, "default": "", "forceInput": True}), 19 | "text_2": ("STRING", {"multiline": False, "default": "", "forceInput": True}), 20 | "text_3": ("STRING", {"multiline": False, "default": "", "forceInput": True}), 21 | "text_4": ("STRING", {"multiline": False, "default": "", "forceInput": True}), 22 | "remove_text": ("STRING", {"multiline": False, "default": "", "forceInput": True}), 23 | }, 24 | } 25 | 26 | RETURN_TYPES = ("STRING", "LIST", "STRING", "STRING", "STRING", "STRING", "STRING") 27 | RETURN_NAMES = ("text", "text_log", "recent_text_1", "recent_text_2", "recent_text_3", "recent_text_4", "oldest_text") 28 | OUTPUT_NODE = True 29 | FUNCTION = "process_text" 30 | CATEGORY = "tksw_node" 31 | 32 | def process_text(self, text_1="", text_2="", text_3="", text_4="", separator=",", remember_log=True, max_log=10, allow_duplicate_log=False, remove_text="", use_regex=False): 33 | texts = [text_1, text_2, text_3, text_4] 34 | 35 | compiled_patterns = [] 36 | if use_regex and remove_text: 37 | remove_patterns = [pattern.strip() for pattern in remove_text.split(",") if pattern.strip()] 38 | for pattern in remove_patterns: 39 | try: 40 | compiled_patterns.append(re.compile(pattern)) 41 | except re.error as e: 42 | print(f"Invalid regex pattern '{pattern}': {e}") 43 | 44 | cleaned_texts = [] 45 | for text in texts: 46 | split_text = text.split(separator) 47 | cleaned_parts = [] 48 | for part in split_text: 49 | if use_regex and compiled_patterns: 50 | cleaned_part = part 51 | for compiled_pattern in compiled_patterns: 52 | cleaned_part = compiled_pattern.sub("", cleaned_part) 53 | else: 54 | remove_words = [word.strip() for word in remove_text.split(",") if word.strip()] 55 | cleaned_part = part 56 | for word in remove_words: 57 | cleaned_part = cleaned_part.replace(word, "") 58 | cleaned_parts.append(cleaned_part) 59 | cleaned_text = separator.join(cleaned_parts) 60 | cleaned_texts.append(cleaned_text) 61 | 62 | combined_text = separator.join([text for text in cleaned_texts if text]) 63 | 64 | pattern = r"(? max_log: 74 | self.text_log.pop(0) 75 | 76 | recent_texts = [""] * 4 77 | for i in range(min(4, len(self.text_log))): 78 | recent_texts[i] = self.text_log[-(i + 1)] 79 | 80 | oldest_text = "" 81 | if self.text_log: 82 | oldest_text = self.text_log[0] 83 | 84 | return (combined_text, list(self.text_log), *recent_texts, oldest_text) 85 | 86 | else: 87 | return (combined_text, [], *[""] * 4, "") 88 | -------------------------------------------------------------------------------- /text_file_selector.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import torch 4 | import hashlib 5 | import codecs 6 | 7 | class TextFileSelector: 8 | def __init__(self): 9 | self.round_robin_index = 0 10 | self.last_folder_path = "" 11 | self.cached_file_list = [] 12 | self.file_content_cache = {} 13 | self.cache_progress_index = 0 14 | self.last_list_hash = None 15 | 16 | @classmethod 17 | def INPUT_TYPES(cls): 18 | return { 19 | "required": { 20 | "folder_path": ("STRING", {"multiline": False, "default": ""}), 21 | "mode": (["random", "round-robin"], {"default": "random"}), 22 | "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), 23 | "reset_state": ("BOOLEAN", {"default": False}), 24 | "cache_chunk_size": ("INT", {"default": 10, "min": 0, "max": 1000}), 25 | "encoding": ("STRING", {"multiline": False, "default": "utf-8"}), 26 | "filename_filter": ("STRING", {"multiline": False, "default": ""}), 27 | } 28 | } 29 | 30 | RETURN_TYPES = ("STRING", "STRING") 31 | RETURN_NAMES = ("text", "filename") 32 | FUNCTION = "select_and_read_file" 33 | CATEGORY = "tksw_node" 34 | 35 | def _scan_folder(self, folder_path): 36 | txt_files = [] 37 | try: 38 | print(f"[TextFileSelector] Scanning folder: {folder_path}") 39 | if not os.path.isdir(folder_path): 40 | print(f"[TextFileSelector] Warning: Folder does not exist or is not a directory: {folder_path}") 41 | return [] 42 | for filename in os.listdir(folder_path): 43 | full_path = os.path.join(folder_path, filename) 44 | if os.path.isfile(full_path) and filename.lower().endswith(".txt"): 45 | txt_files.append(filename) 46 | print(f"[TextFileSelector] Found {len(txt_files)} .txt files.") 47 | return sorted(txt_files) 48 | except Exception as e: 49 | print(f"[TextFileSelector] Error scanning folder '{folder_path}': {e}") 50 | return [] 51 | 52 | def _read_and_cache_file(self, filename, folder_path, encoding): 53 | full_path = os.path.join(folder_path, filename) 54 | if filename in self.file_content_cache: 55 | return self.file_content_cache[filename] 56 | print(f"[TextFileSelector] Reading and caching: {filename} (Encoding: {encoding})") 57 | try: 58 | with codecs.open(full_path, 'r', encoding=encoding, errors='ignore') as f: 59 | content = f.read() 60 | self.file_content_cache[filename] = content 61 | return content 62 | except FileNotFoundError: 63 | print(f"[TextFileSelector] Error: File not found at path '{full_path}'.") 64 | self.file_content_cache[filename] = "" 65 | return "" 66 | except Exception as e: 67 | print(f"[TextFileSelector] Error reading file '{full_path}' with encoding '{encoding}': {e}") 68 | self.file_content_cache[filename] = "" 69 | return "" 70 | 71 | def select_and_read_file(self, folder_path, mode, seed, reset_state, cache_chunk_size, encoding, filename_filter): 72 | clean_folder_path = folder_path.strip() 73 | current_list_hash = self.last_list_hash 74 | needs_full_reset = False 75 | reset_reason = "" 76 | 77 | if clean_folder_path != self.last_folder_path: 78 | print(f"[TextFileSelector] Folder path changed to: '{clean_folder_path}'") 79 | self.last_folder_path = clean_folder_path 80 | if clean_folder_path and os.path.isdir(clean_folder_path): 81 | self.cached_file_list = self._scan_folder(clean_folder_path) 82 | else: 83 | self.cached_file_list = [] 84 | if clean_folder_path: 85 | print(f"[TextFileSelector] Warning: Invalid folder path provided: '{clean_folder_path}'") 86 | 87 | current_list_hash = hashlib.sha256(str(self.cached_file_list).encode()).hexdigest() 88 | if current_list_hash != self.last_list_hash: 89 | needs_full_reset = True 90 | reset_reason = "Folder path or file list content changed." 91 | 92 | if reset_state: 93 | needs_full_reset = True 94 | reset_reason = "Manual state reset requested." 95 | elif not self.cached_file_list and self.last_list_hash is not None: 96 | needs_full_reset = True 97 | reset_reason = "File list became empty." 98 | 99 | if needs_full_reset: 100 | print(f"[TextFileSelector] Full state reset triggered: {reset_reason}") 101 | self.round_robin_index = 0 102 | self.file_content_cache = {} 103 | self.cache_progress_index = 0 104 | self.last_list_hash = current_list_hash 105 | if not self.cached_file_list: 106 | self.last_list_hash = None 107 | 108 | num_files = len(self.cached_file_list) 109 | processed_in_chunk = 0 110 | effective_cache_chunk_size = max(0, cache_chunk_size) 111 | if effective_cache_chunk_size > 0 and num_files > 0: 112 | for i in range(effective_cache_chunk_size): 113 | if num_files == 0: break 114 | idx_to_check = (self.cache_progress_index + i) % num_files 115 | filename_to_check = self.cached_file_list[idx_to_check] 116 | if filename_to_check not in self.file_content_cache: 117 | self._read_and_cache_file(filename_to_check, self.last_folder_path, encoding) 118 | processed_in_chunk += 1 119 | if num_files > 0: 120 | self.cache_progress_index = (self.cache_progress_index + effective_cache_chunk_size) % num_files 121 | 122 | chosen_filename = None 123 | clean_filename_filter = filename_filter.strip() 124 | 125 | if not self.cached_file_list: 126 | print("[TextFileSelector] No text files found or folder path is invalid.") 127 | return ("", "None") 128 | else: 129 | effective_candidates = self.cached_file_list 130 | is_filtered = False 131 | if clean_filename_filter: 132 | print(f"[TextFileSelector] Applying filename filter: '{clean_filename_filter}'") 133 | filtered_list_for_random = [f for f in self.cached_file_list if clean_filename_filter in f] 134 | if not filtered_list_for_random: 135 | print(f"[TextFileSelector] Warning: No files match the filter '{clean_filename_filter}'.") 136 | else: 137 | effective_candidates = filtered_list_for_random 138 | is_filtered = True 139 | 140 | num_effective_candidates = len(effective_candidates) 141 | 142 | if num_effective_candidates == 0 and is_filtered: 143 | print(f"[TextFileSelector] No files available after applying filter '{clean_filename_filter}'.") 144 | return ("", "None") 145 | elif num_effective_candidates == 0 and not is_filtered: 146 | print("[TextFileSelector] No files available.") 147 | return ("", "None") 148 | 149 | 150 | if mode == "random": 151 | random.seed(seed) 152 | chosen_filename = random.choice(effective_candidates) 153 | print(f"[TextFileSelector] Mode: random {'with filter' if is_filtered else ''} (Seed: {seed})") 154 | 155 | elif mode == "round-robin": 156 | num_total_files = len(self.cached_file_list) 157 | found = False 158 | search_start_index = self.round_robin_index % num_total_files 159 | 160 | for i in range(num_total_files): 161 | current_check_index = (search_start_index + i) % num_total_files 162 | filename_to_check = self.cached_file_list[current_check_index] 163 | 164 | filter_match = (not clean_filename_filter) or (clean_filename_filter in filename_to_check) 165 | 166 | if filter_match: 167 | chosen_filename = filename_to_check 168 | found_index = current_check_index 169 | self.round_robin_index = found_index + 1 170 | print(f"[TextFileSelector] Mode: round-robin {'with filter' if clean_filename_filter else ''} (Found at index {found_index}, Next RR Base: {self.round_robin_index})") 171 | found = True 172 | break 173 | 174 | if not found: 175 | if clean_filename_filter: 176 | print(f"[TextFileSelector] Mode: round-robin with filter '{clean_filename_filter}'. No matching file found.") 177 | else: 178 | print("[TextFileSelector] Mode: round-robin. No file found (List might be empty unexpectedly).") 179 | return ("", "None") 180 | 181 | else: 182 | print(f"[TextFileSelector] Warning: Unknown mode '{mode}'. Falling back to random.") 183 | random.seed(seed) 184 | chosen_filename = random.choice(effective_candidates) 185 | 186 | if chosen_filename: 187 | print(f"[TextFileSelector] Selected file: {chosen_filename}") 188 | file_content = self._read_and_cache_file(chosen_filename, self.last_folder_path, encoding) 189 | return (file_content, chosen_filename) 190 | else: 191 | print("[TextFileSelector] Internal error: Could not select a file.") 192 | return ("", "None") -------------------------------------------------------------------------------- /text_processor.py: -------------------------------------------------------------------------------- 1 | import re 2 | 3 | class TextProcessor: 4 | @classmethod 5 | def INPUT_TYPES(s): 6 | return { 7 | "required": { 8 | "segment_separator": ("STRING", {"multiline": False, "default": ","}), 9 | }, 10 | "optional": { 11 | "input_text": ("STRING", {"multiline": True, "default": "", "forceInput": True}), 12 | "remove_patterns": ("STRING", {"multiline": False, "default": ""}), 13 | "replace_specs": ("STRING", {"multiline": True, "default": ""}), 14 | } 15 | } 16 | 17 | RETURN_TYPES = ("STRING",) 18 | RETURN_NAMES = ("processed_text",) 19 | FUNCTION = "process_text" 20 | CATEGORY = "tksw_node" 21 | 22 | def process_text(self, input_text="", remove_patterns="", replace_specs="", segment_separator=","): 23 | print(f"Input Text: {input_text}") 24 | print(f"Remove Patterns: {remove_patterns}") 25 | print(f"Replace Specs: {replace_specs}") 26 | print(f"Segment Separator: {segment_separator}") 27 | 28 | if not input_text: 29 | return "" 30 | 31 | if segment_separator.strip() == "": 32 | segments = [input_text] 33 | else: 34 | segments = input_text.split(segment_separator) 35 | 36 | print(f"Segments: {segments}") 37 | 38 | processed_segments = [] 39 | 40 | for segment in segments: 41 | cleaned_segment = segment.strip() 42 | 43 | print(f"Segment before processing: {cleaned_segment}") 44 | 45 | if remove_patterns: 46 | cleaned_segment = self.apply_remove_patterns(cleaned_segment, remove_patterns) 47 | print(f"Segment after remove: {cleaned_segment}") 48 | 49 | if replace_specs: 50 | cleaned_segment = self.apply_replace_specs(cleaned_segment, replace_specs) 51 | print(f"Segment after replace: {cleaned_segment}") 52 | 53 | processed_segments.append(cleaned_segment) 54 | 55 | print(f"Processed Segments: {processed_segments}") 56 | 57 | if segment_separator.strip() == "": 58 | processed_text = "".join(processed_segments) 59 | else: 60 | processed_text = segment_separator.join(processed_segments) 61 | 62 | processed_text = re.sub(r" +", " ", processed_text) 63 | processed_text = re.sub(r"\s*,\s*", ",", processed_text) 64 | processed_text = re.sub(r",+", ",", processed_text) 65 | processed_text = re.sub(r"^,|,$", "", processed_text) 66 | processed_text = re.sub(r",(?=[^\s])", ", ", processed_text) 67 | processed_text = processed_text.strip() 68 | 69 | print(f"Processed Text: {processed_text}") 70 | return (processed_text,) 71 | 72 | def split_into_segments(self, text, separator): 73 | return [segment.strip() for segment in text.split(separator)] 74 | 75 | def clean_segment(self, segment, remove_patterns, replace_specs): 76 | cleaned_segment = segment 77 | 78 | cleaned_segment = self.apply_remove_patterns(cleaned_segment, remove_patterns) 79 | cleaned_segment = self.apply_replace_specs(cleaned_segment, replace_specs) 80 | 81 | return cleaned_segment 82 | 83 | 84 | def apply_remove_patterns(self, text, remove_patterns): 85 | compiled_patterns = [] 86 | for pattern in [p.strip() for p in remove_patterns.split(",") if p.strip()]: 87 | try: 88 | compiled_patterns.append(re.compile(pattern)) 89 | except re.error as e: 90 | print(f"Invalid remove pattern: {e}") 91 | 92 | for pattern in compiled_patterns: 93 | text = pattern.sub("", text) 94 | return text 95 | 96 | def apply_replace_specs(self, text, replace_specs): 97 | replace_lines = [line.strip() for line in replace_specs.splitlines() if line.strip()] 98 | compiled_replace_specs = [] 99 | for line in replace_lines: 100 | parts = [part.strip() for part in line.split(",")] 101 | if parts: 102 | try: 103 | compiled_replace_specs.append((parts[0], [re.compile(p) for p in parts[1:]])) 104 | except re.error as e: 105 | print(f"Invalid replace pattern: {e}") 106 | 107 | for replacement, patterns in compiled_replace_specs: 108 | for pattern in patterns: 109 | text = pattern.sub(replacement, text) 110 | return text --------------------------------------------------------------------------------