Compare commits

...
Author SHA1 Message Date
silveroxides 0a82e071b8 fix NodeOutput issue 2026-04-08 17:50:10 +02:00
silveroxides dc43040562 Rename downloader nodes and add nodes for displaying previews, metadata of model and workflows extracted from images/videos 2026-04-08 16:46:35 +02:00
silveroxides 39380f533d Remove subprocess, use av. More info in status output from Downloader nodes 2026-04-08 16:28:45 +02:00
silveroxides 2ac0d46935 Add progress bar 2026-04-08 15:43:05 +02:00
silveroxides 378c327406 Add subdirectory selector 2026-04-08 15:31:42 +02:00
silveroxides 4076c14cb5 Add requirements.txt entries 2026-04-08 15:16:16 +02:00
silveroxides c4c967172b Add Model preview/example image/video downloading, model hashing, preview/example workflow extraction from images and videos followed by compression of images to save storage, saving of hash database and much more 2026-04-08 15:04:24 +02:00
Silver 095ac89164 Bump version to 0.2.9 2026-03-07 19:10:13 +01:00
Silver 2d477767d6 Merge pull request #9 from silveroxides/fix/loading-fixes
Fix models getting stuck in memory when merging
2026-03-07 19:09:32 +01:00
silveroxides 4d49477f58 Fix models getting stuck in memory when merging 2026-03-07 16:32:43 +01:00
Silver 6c4d505d84 Bump version from 0.2.7 to 0.2.8 2026-02-24 17:41:25 +01:00
Silver 1e53a87427 Merge pull request #8 from silveroxides/feature/expand-lora-functionality
expand on functionality and reduce restrictions when extracting LoRA
2026-02-24 17:40:47 +01:00
silveroxides b497022c3c expand on functionality and reduce restrictions when extracting LoRA 2026-02-24 17:39:48 +01:00
Silver 5df8c6711a Merge pull request #7 from silveroxides/fix/loading-imporovements
improve memory efficiency
2026-02-12 08:07:39 +01:00
Silver f4cecdc49b Merge branch 'main' into fix/loading-imporovements 2026-02-12 08:07:31 +01:00
Silver 8b34d57003 Merge pull request #6 from silveroxides/patch/lora-changes
changes to LoRAs
2026-02-12 08:05:19 +01:00
silveroxides 0d011d1232 improve memory efficiency 2026-02-07 16:01:00 +01:00
12 changed files with 1673 additions and 362 deletions
+14
View File
@@ -29,6 +29,14 @@ from .nodes.lora_resize import (
LoRAResizeFrobenius, LoRAResizeCumulative,
LoRAMergeToModel
)
from .nodes.downloader_nodes import (
CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader
)
from .nodes.model_info_nodes import (
CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader,
VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader
)
class ModelUtilsExtension(ComfyExtension):
@@ -60,6 +68,12 @@ class ModelUtilsExtension(ComfyExtension):
LoRAMultiMerge, LoRAMultiMergeDARE,
# LoRA Merge To Model
LoRAMergeToModel,
# Downloaders
CheckpointInfoMetaDownloader, DiffusionModelInfoMetaDownloader, LoRAInfoMetaDownloader, EmbeddingInfoMetaDownloader,
VAEInfoMetaDownloader, ControlNetInfoMetaDownloader, ManualPathInfoMetaDownloader,
# Info Loaders
CheckpointInfoLoader, LoRAInfoLoader, EmbeddingInfoLoader,
VAEInfoLoader, ControlNetInfoLoader, DiffusionModelInfoLoader,
]
+341
View File
@@ -0,0 +1,341 @@
import os
import folder_paths
from comfy_api.latest import io
from .downloader_utils import scan_and_process
def get_subdirectories(folder_name: str) -> list[str]:
"""Helper to get a list of subdirectories for a given model folder name."""
paths = folder_paths.get_folder_paths(folder_name)
if not paths:
return ["/"]
subdirs = set()
for base_path in paths:
if not os.path.exists(base_path):
continue
for root, dirs, _ in os.walk(base_path):
# Ignore hidden directories directly during walk
dirs[:] = [d for d in dirs if not d.startswith('.')]
rel_path = os.path.relpath(root, base_path)
if rel_path == ".":
subdirs.add("/")
else:
# Store with forward slashes for consistency
subdirs.add("/" + rel_path.replace("\\", "/"))
return sorted(list(subdirs))
class BaseInfoMetaDownloaderNode(io.ComfyNode):
CATEGORY = "ModelUtils/Downloader"
OUTPUT_NODE = True
@classmethod
def get_common_inputs(cls):
return {
"recursive": io.Boolean.Input("recursive", default=True, tooltip="Whether to scan subdirectories recursively."),
"nsfw_level": io.Combo.Input("nsfw_level", options=["None", "Soft", "Mature", "X", "XXX", "All"], default="All", tooltip="Filter models based on NSFW level in metadata."),
"max_examples": io.Int.Input("max_examples", default=0, min=0, max=100, tooltip="Maximum number of examples to process."),
"threads": io.Int.Input("threads", default=4, min=1, max=32, tooltip="Number of threads to use for processing."),
"api_key": io.String.Input("api_key", default="", tooltip="API key for accessing the model repository.")
}
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name=cls.__name__.replace("Downloader", " Downloader"),
category=cls.CATEGORY,
inputs=list(cls.get_common_inputs().values()),
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
description=("Scans specified directories for models, extracts metadata, generates previews, and optionally compresses files.")
)
@classmethod
def execute(cls, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
scan_dirs = cls.get_scan_dirs(**kwargs)
if not scan_dirs:
return io.NodeOutput("No directories to scan.")
# Call the scan_and_process function
stats = scan_and_process(
scan_dirs=scan_dirs,
recursive=recursive,
nsfw_level=nsfw_level,
max_examples=max_examples,
api_key=api_key if api_key else None,
threads=threads
)
mb_saved = stats.get("space_saved", 0) / (1024 * 1024)
output_msg = (
f"Scan Complete for {len(scan_dirs)} directory/directories.\n"
f"Total Models Processed: {stats.get('total_processed', 0)}\n"
f"New Models Hashed: {stats.get('new_hashes', 0)}\n"
f"Images Processed: {stats.get('images_processed', 0)}\n"
f"Workflows Extracted: {stats.get('workflows_extracted', 0)}\n"
f"Space Saved via Compression: {mb_saved:.2f} MB"
)
return io.NodeOutput(output_msg)
@classmethod
def get_scan_dirs(cls, **kwargs) -> list[str]:
return []
class CheckpointInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("checkpoints"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="Checkpoint Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("checkpoints")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class DiffusionModelInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("diffusion_models"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="Diffusion Model Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("diffusion_models")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class LoRAInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("loras"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="LoRA Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("loras")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class EmbeddingInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("embeddings"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="Embedding Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("embeddings")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class VAEInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("vae"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="VAE Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("vae")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class ControlNetInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("controlnet"), default="/"))
return io.Schema(
node_id=cls.__name__,
display_name="ControlNet Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs):
return super().execute(
recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples,
threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
base_paths = folder_paths.get_folder_paths("controlnet")
if not base_paths:
return []
sub = kwargs.get("subdirectory", "/")
if sub == "/":
return base_paths
target_paths = []
for bp in base_paths:
target_path = os.path.join(bp, sub.lstrip("/"))
if os.path.exists(target_path):
target_paths.append(target_path)
return target_paths
class ManualPathInfoMetaDownloader(BaseInfoMetaDownloaderNode):
@classmethod
def define_schema(cls):
inputs = list(cls.get_common_inputs().values())
inputs.append(io.String.Input("scan_dir", default=""))
return io.Schema(
node_id=cls.__name__,
display_name="Manual Path Downloader",
category=cls.CATEGORY,
inputs=inputs,
outputs=[io.String.Output(display_name="status")],
is_output_node=True,
)
@classmethod
def execute(cls, recursive, nsfw_level, max_examples, threads, api_key, scan_dir, **kwargs):
return super().execute(
recursive=recursive,
nsfw_level=nsfw_level,
max_examples=max_examples,
threads=threads,
api_key=api_key,
scan_dir=scan_dir,
**kwargs
)
@classmethod
def get_scan_dirs(cls, **kwargs):
scan_dir = kwargs.get("scan_dir", "")
if scan_dir.strip():
return [scan_dir.strip()]
return []
+702
View File
@@ -0,0 +1,702 @@
import os
import sys
import hashlib
import json
import sqlite3
import requests
import glob
from pathlib import Path
import threading
import concurrent.futures
import codecs
import shutil
try:
from PIL import Image
HAS_PIL = True
except ImportError:
HAS_PIL = False
print("Warning: Pillow not found. Image workflow extraction disabled.")
try:
from mutagen.mp4 import MP4
HAS_MUTAGEN = True
except ImportError:
HAS_MUTAGEN = False
print("Warning: mutagen not found. MP4 tag extraction will rely entirely on av fallback.")
try:
import av
HAS_AV = True
except ImportError:
HAS_AV = False
print("Warning: av (PyAV) not found. MP4 metadata extraction fallback disabled.")
# Add the reference module to the Python path
# Assuming downloader_utils.py is inside nodes/, the root is one level up
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.join(root_dir, 'unifiedefficientloader-reference'))
try:
from unifiedefficientloader import UnifiedSafetensorsLoader
HAS_UNIFIED_LOADER = True
except ImportError:
HAS_UNIFIED_LOADER = False
print("Warning: unifiedefficientloader not found. Safetensors metadata extraction disabled.")
# --- Database Path Logic ---
try:
import folder_paths
base_user_dir = folder_paths.get_public_user_directory("default")
if not base_user_dir:
base_user_dir = root_dir
except (ImportError, AttributeError):
base_user_dir = root_dir
PLUGIN_DATA_DIR = os.path.join(base_user_dir, "ComfyUI-ModelUtils")
os.makedirs(PLUGIN_DATA_DIR, exist_ok=True)
DEFAULT_DB_PATH = os.path.join(PLUGIN_DATA_DIR, "cache.db")
# --- Workflow Extraction Logic ---
def unescape_and_parse_nested_json(json_string):
if not isinstance(json_string, str): return None
stripped = json_string.strip()
if not (stripped.startswith('{') or stripped.startswith('[')): return None
try: return json.loads(stripped)
except json.JSONDecodeError: pass
try:
first_parse = json.loads(stripped)
if isinstance(first_parse, str): return json.loads(first_parse)
except json.JSONDecodeError: pass
try:
decoded = codecs.decode(stripped, "unicode_escape")
return json.loads(decoded)
except (json.JSONDecodeError, UnicodeDecodeError, ValueError): pass
try:
unescaped = stripped.replace('\\\\"', '\\"').replace('\\"', '"').replace("\\\\", "\\")
unescaped = unescaped.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t")
return json.loads(unescaped)
except json.JSONDecodeError: pass
return None
def get_workflow_blob(metadata):
for key in ['workflow', 'prompt']:
val = metadata.get(key)
if val:
parsed = unescape_and_parse_nested_json(val)
if parsed: return parsed
for val in metadata.values():
parsed = unescape_and_parse_nested_json(val)
if isinstance(parsed, dict) and (any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed):
return parsed
return None
def process_image_workflow_and_compress(filepath_str):
if not HAS_PIL: return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
p = Path(filepath_str)
if p.suffix.lower() not in ['.png', '.jpg', '.jpeg']: return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
workflow = None
actual_format = None
stats = {"images_processed": 1, "workflows_extracted": 0, "space_saved": 0}
try:
with Image.open(p) as img:
actual_format = img.format
# CivitAI sometimes serves PNGs with .jpeg extension in URL
if actual_format == 'PNG':
workflow = get_workflow_blob(img.info)
if workflow:
# Save workflow
json_path = p.with_suffix('.json')
if not json_path.exists():
with open(json_path, 'w', encoding='utf-8') as f:
json.dump(workflow, f, indent=2, ensure_ascii=False)
print(f"Extracted workflow to {json_path}")
stats["workflows_extracted"] = 1
# Compress image to WebP
webp_path = p.with_suffix('.webp')
if not webp_path.exists():
orig_size = os.path.getsize(p)
img.save(webp_path, 'WEBP', quality=85, method=4)
new_size = os.path.getsize(webp_path)
stats["space_saved"] = max(0, orig_size - new_size)
print(f"Compressed {p.name} (was actually {actual_format}) to {webp_path.name}")
# Delete original if we created a webp replacement
if workflow and p.with_suffix('.webp').exists() and actual_format == 'PNG':
os.remove(p)
print(f"Deleted original uncompressed image: {p.name}")
except Exception as e:
print(f"Failed to process image {p.name}: {e}")
return stats
# --- Video Workflow Logic ---
def get_tag_value(mp4, keys):
for key in keys:
if key in mp4.tags: return mp4.tags[key][0] if mp4.tags[key] else None
for key in keys:
search_suffix = key.split(':')[-1].lower()
for tag_key, val in mp4.tags.items():
if tag_key.lower().endswith(f":{search_suffix}"): return val[0] if val else None
return None
def extract_metadata_av(mp4_path):
if not HAS_AV: return None
try:
with av.open(str(mp4_path)) as container:
return {k.lower(): v for k, v in container.metadata.items()}
except Exception as e:
print(f"av extraction error for {mp4_path}: {e}")
return None
def process_video_workflow(filepath_str):
p = Path(filepath_str)
if p.suffix.lower() != '.mp4': return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
json_path = p.with_suffix('.json')
if json_path.exists(): return {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
workflow_data = None
mp4 = None
stats = {"images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
if HAS_MUTAGEN:
try: mp4 = MP4(p)
except: pass
if mp4:
for tag_keys in [['----:QuickTime:Workflow', 'workflow'], ['----:QuickTime:Prompt', 'prompt'], ['©cmt']]:
val = get_tag_value(mp4, tag_keys)
if val:
if isinstance(val, bytes): val = val.decode('utf-8', errors='ignore')
try:
parsed = json.loads(val)
if isinstance(parsed, dict) and ('prompt' in parsed or any(k.isdigit() for k in parsed.keys()) or 'nodes' in parsed):
if 'prompt' in parsed and isinstance(parsed['prompt'], str):
try: workflow_data = json.loads(parsed['prompt'])
except: workflow_data = parsed
else:
workflow_data = parsed
break
except json.JSONDecodeError:
# Try more aggressive parsing if it's a nested string
parsed = unescape_and_parse_nested_json(val)
if parsed:
workflow_data = parsed
break
if not workflow_data:
av_data = extract_metadata_av(p)
if av_data:
for key in ['workflow', 'prompt', 'comment', 'description']:
if key in av_data:
val = av_data[key]
try:
parsed = json.loads(val)
if isinstance(parsed, dict):
if 'prompt' in parsed and isinstance(parsed['prompt'], str):
try: workflow_data = json.loads(parsed['prompt'])
except: workflow_data = parsed
else:
workflow_data = parsed
break
except json.JSONDecodeError:
parsed = unescape_and_parse_nested_json(val)
if parsed:
workflow_data = parsed
break
if workflow_data:
try:
with open(json_path, 'w', encoding='utf-8') as f:
json.dump(workflow_data, f, indent=2, ensure_ascii=False)
print(f"Extracted workflow to {json_path.name}")
stats["workflows_extracted"] = 1
except Exception as e:
print(f"Error saving workflow for {p.name}: {e}")
return stats
# --- Constants for Civitai NSFW levels ---
NSFW_LEVELS = {
"None": 1,
"Soft": 2,
"Mature": 4,
"X": 8,
"XXX": 16,
"All": 32 # Anything
}
class HashCache:
def __init__(self, db_path=None):
self.db_path = db_path if db_path else DEFAULT_DB_PATH
self.is_sqlite = self.db_path.endswith('.db') or self.db_path.endswith('.sqlite')
self.lock = threading.Lock()
if self.is_sqlite:
self._init_sqlite()
else:
self.cache = self._load_json()
def _init_sqlite(self):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS file_hashes (
filepath TEXT PRIMARY KEY,
mtime REAL,
size INTEGER,
hash TEXT
)
''')
conn.commit()
def _load_json(self):
if os.path.exists(self.db_path):
with open(self.db_path, "r", encoding="utf-8") as f:
return json.load(f)
return {}
def _save_json(self):
with open(self.db_path, "w", encoding="utf-8") as f:
json.dump(self.cache, f, indent=4)
def _get_sqlite_hash(self, filepath, mtime, size):
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute("SELECT mtime, size, hash FROM file_hashes WHERE filepath=?", (filepath,))
row = cursor.fetchone()
if row:
db_mtime, db_size, db_hash = row
if db_mtime == mtime and db_size == size:
return db_hash
return None
def _set_sqlite_hash(self, filepath, mtime, size, file_hash):
with sqlite3.connect(self.db_path, timeout=10.0) as conn:
cursor = conn.cursor()
cursor.execute('''
INSERT OR REPLACE INTO file_hashes (filepath, mtime, size, hash)
VALUES (?, ?, ?, ?)
''', (filepath, mtime, size, file_hash))
conn.commit()
def get_hash(self, filepath):
if not os.path.exists(filepath):
return None
mtime = os.path.getmtime(filepath)
size = os.path.getsize(filepath)
key = filepath
with self.lock:
if self.is_sqlite:
cached_hash = self._get_sqlite_hash(key, mtime, size)
if cached_hash: return cached_hash
else:
if key in self.cache:
entry = self.cache[key]
if entry["mtime"] == mtime and entry["size"] == size:
return entry["hash"]
# Calculate new hash (outside the lock to allow parallel hashing)
print(f"Calculating SHA256 for: {filepath}")
sha256_hash = self.calculate_sha256(filepath)
with self.lock:
if self.is_sqlite:
self._set_sqlite_hash(key, mtime, size, sha256_hash)
else:
self.cache[key] = {
"mtime": mtime,
"size": size,
"hash": sha256_hash
}
self._save_json()
return sha256_hash
@staticmethod
def calculate_sha256(filename):
sha256_hash = hashlib.sha256()
with open(filename, "rb") as f:
# Read and update hash string value in blocks of 4K
for byte_block in iter(lambda: f.read(4096), b""):
sha256_hash.update(byte_block)
return sha256_hash.hexdigest()
class CivitaiAPI:
def __init__(self, api_key=None):
self.api_key = api_key
self.base_url = "https://civitai.com/api/v1"
def get_headers(self):
headers = {}
if self.api_key:
headers["Authorization"] = f"Bearer {self.api_key}"
return headers
def get_model_info_by_hash(self, file_hash):
url = f"{self.base_url}/model-versions/by-hash/{file_hash}"
response = requests.get(url, headers=self.get_headers())
if response.status_code == 200:
return response.json()
elif response.status_code == 404:
return None
else:
response.raise_for_status()
class Downloader:
@staticmethod
def download_file(url, target_path, headers=None):
if headers is None:
headers = {}
temp_path = target_path + ".downloading"
downloaded_size = 0
if os.path.exists(temp_path):
downloaded_size = os.path.getsize(temp_path)
if downloaded_size > 0:
headers["Range"] = f"bytes={downloaded_size}-"
response = requests.get(url, headers=headers, stream=True)
if response.status_code == 416: # Range not satisfiable, file probably changed or completed
os.remove(temp_path)
downloaded_size = 0
if "Range" in headers:
del headers["Range"]
response = requests.get(url, headers=headers, stream=True)
if response.status_code not in [200, 206]:
print(f"Failed to download from {url}, status code: {response.status_code}")
return False
mode = "ab" if response.status_code == 206 else "wb"
total_size = int(response.headers.get('content-length', 0)) + downloaded_size
print(f"Downloading: {target_path} ({downloaded_size}/{total_size} bytes)")
with open(temp_path, mode) as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
os.rename(temp_path, target_path)
return True
def get_potential_preview_files(base_path):
exts = ["png", "jpg", "jpeg", "webp", "mp4"]
files = []
for ext in exts:
files.append(f"{base_path}.preview.{ext}")
files.append(f"{base_path}.{ext}")
return files
def preview_exists(base_path):
previews = get_potential_preview_files(base_path)
for p in previews:
if os.path.exists(p):
return True
return False
def next_example_image_path(base_path, ext):
i = 0
while len(glob.glob(f"{base_path}.example.{i}.*")) > 0:
i += 1
return f"{base_path}.example.{i}.{ext}"
def filter_images_by_nsfw(images, threshold_level):
valid_images = []
for img in images:
rating = img.get("nsfwLevel", 1)
if rating <= threshold_level:
valid_images.append(img)
return valid_images
def extract_safetensors_metadata(filepath):
if not HAS_UNIFIED_LOADER:
return None
try:
# Use low_memory=True to only read the header and metadata without loading tensors
loader = UnifiedSafetensorsLoader(filepath, low_memory=True)
metadata = loader.metadata()
loader.close()
return metadata if metadata else None
except Exception as e:
print(f"Failed to read safetensors metadata from {filepath} using unifiedefficientloader: {e}")
return None
def process_file(filepath, cache, api, nsfw_threshold, max_examples):
filepath_str = str(filepath)
base_path = os.path.splitext(filepath_str)[0]
is_safetensors = filepath_str.lower().endswith(".safetensors")
stats = {"new_hashes": 0, "images_processed": 0, "workflows_extracted": 0, "space_saved": 0}
# Check if all targets already exist
info_path = f"{base_path}.civitai.info"
st_metadata_path = f"{base_path}.metadata.json"
has_info = os.path.exists(info_path)
has_preview = preview_exists(base_path)
has_st_metadata = not is_safetensors or os.path.exists(st_metadata_path)
# Process existing images/videos for workflow extraction & compression BEFORE we possibly skip the file
for p in get_potential_preview_files(base_path):
if os.path.exists(p):
ext_lower = p.lower()
if ext_lower.endswith(('.png', '.jpg', '.jpeg')):
res = process_image_workflow_and_compress(p)
for k in stats: stats[k] += res.get(k, 0)
elif ext_lower.endswith('.mp4'):
res = process_video_workflow(p)
for k in stats: stats[k] += res.get(k, 0)
for ex in glob.glob(f"{base_path}.example.*"):
ext_lower = ex.lower()
if ext_lower.endswith(('.png', '.jpg', '.jpeg')):
res = process_image_workflow_and_compress(ex)
for k in stats: stats[k] += res.get(k, 0)
elif ext_lower.endswith('.mp4'):
res = process_video_workflow(ex)
for k in stats: stats[k] += res.get(k, 0)
# If everything exists and we don't want examples or already have enough examples, we might skip API logic
if has_info and has_preview and has_st_metadata and max_examples == 0:
print(f"Skipping fully populated model: {filepath_str}")
return stats
print(f"Processing model: {filepath_str}")
# Safetensors built-in metadata extraction
if is_safetensors and not has_st_metadata:
metadata = extract_safetensors_metadata(filepath_str)
if metadata:
with open(st_metadata_path, "w", encoding="utf-8") as f:
json.dump(metadata, f, indent=4)
print(f"Saved safetensors metadata to {st_metadata_path}")
mtime = os.path.getmtime(filepath_str)
size = os.path.getsize(filepath_str)
# We want to know if it's a NEW hash or cached
is_cached = False
with cache.lock:
if cache.is_sqlite:
cached_hash = cache._get_sqlite_hash(filepath_str, mtime, size)
if cached_hash: is_cached = True
else:
if filepath_str in cache.cache:
entry = cache.cache[filepath_str]
if entry["mtime"] == mtime and entry["size"] == size:
is_cached = True
file_hash = cache.get_hash(filepath_str)
if not file_hash:
return stats
if not is_cached:
stats["new_hashes"] += 1
try:
model_info = api.get_model_info_by_hash(file_hash)
except Exception as e:
print(f"API Error for {filepath_str}: {e}")
return stats
if not model_info:
print(f"Model not found on Civitai: {filepath_str}")
return stats
# 1. Save Info
if not has_info:
with open(info_path, "w", encoding="utf-8") as f:
json.dump(model_info, f, indent=4)
print(f"Saved metadata to {info_path}")
images = model_info.get("images", [])
valid_images = filter_images_by_nsfw(images, nsfw_threshold)
if not valid_images:
print(f"No valid images found matching NSFW criteria for {filepath_str}")
return stats
# 2. Save Preview
if not has_preview and len(valid_images) > 0:
preview_img = valid_images[0]
url = preview_img.get("url")
ext = "jpeg" # Default fallback
if ".png" in url.lower(): ext = "png"
elif ".mp4" in url.lower(): ext = "mp4"
elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg"
preview_path = f"{base_path}.preview.{ext}"
if Downloader.download_file(url, preview_path):
if ext in ("png", "jpg", "jpeg"):
res = process_image_workflow_and_compress(preview_path)
for k in stats: stats[k] += res.get(k, 0)
elif ext == "mp4":
res = process_video_workflow(preview_path)
for k in stats: stats[k] += res.get(k, 0)
# 3. Save Examples
if max_examples > 0:
existing_examples_count = len(glob.glob(f"{base_path}.example.*"))
# We divide by 2 roughly if there are json files, let's just count unique example indices
existing_indices = set()
for f in glob.glob(f"{base_path}.example.*"):
parts = f.split('.example.')
if len(parts) > 1:
idx = parts[1].split('.')[0]
if idx.isdigit(): existing_indices.add(idx)
needed_examples = max_examples - len(existing_indices)
if needed_examples > 0:
example_images = valid_images[1:max_examples+1]
# Skip images we probably already downloaded by skipping the first N
example_images_to_dl = example_images[len(existing_indices):]
for img in example_images_to_dl:
url = img.get("url")
ext = "jpeg"
if ".png" in url.lower(): ext = "png"
elif ".mp4" in url.lower(): ext = "mp4"
elif ".jpg" in url.lower() or ".jpeg" in url.lower(): ext = "jpeg"
example_path = next_example_image_path(base_path, ext)
if Downloader.download_file(url, example_path):
if ext in ("png", "jpg", "jpeg"):
res = process_image_workflow_and_compress(example_path)
for k in stats: stats[k] += res.get(k, 0)
elif ext == "mp4":
res = process_video_workflow(example_path)
for k in stats: stats[k] += res.get(k, 0)
return stats
import numpy as np
from PIL import Image, ImageOps
import torch
def load_image_tensor(path):
if not os.path.exists(path):
return None
try:
img = Image.open(path)
img = ImageOps.exif_transpose(img)
image = img.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
image = torch.from_numpy(image)[None,]
return image
except Exception as e:
print(f"Error loading image {path}: {e}")
return None
def get_model_workflows(base_path):
workflows = []
# Main workflow
main_json = f"{base_path}.json"
if os.path.exists(main_json):
workflows.append(main_json)
# Example workflows
# They are named like base_path.example.N.json
# or base_path.preview.json (if extracted from preview)
preview_json = f"{base_path}.preview.json"
if os.path.exists(preview_json) and preview_json not in workflows:
workflows.append(preview_json)
example_jsons = glob.glob(f"{base_path}.example.*.json")
# Sort them by index
example_jsons.sort()
for ej in example_jsons:
if ej not in workflows:
workflows.append(ej)
return workflows
def get_model_metadata_file(base_path):
# Prefer civitai.info if available
civitai_info = f"{base_path}.civitai.info"
if os.path.exists(civitai_info):
return civitai_info
# Fallback to metadata.json (safetensors metadata)
metadata_json = f"{base_path}.metadata.json"
if os.path.exists(metadata_json):
return metadata_json
return None
def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", max_examples=0, api_key=None, threads=4):
"""
Scans directories for models and fetches their metadata/previews.
Adapted for direct Python API calls rather than argparse CLI.
"""
db_path = db_path if db_path else DEFAULT_DB_PATH
cache = HashCache(db_path)
api = CivitaiAPI(api_key)
model_exts = [".safetensors", ".ckpt", ".pt"]
files_to_process = []
if isinstance(scan_dirs, str):
scan_dirs = [scan_dirs]
for scan_dir in scan_dirs:
path = Path(scan_dir)
if recursive:
for ext in model_exts:
# rglob includes all directories, so we must filter out hidden ones
for f in path.rglob(f"*{ext}"):
if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'):
files_to_process.append(f)
else:
for ext in model_exts:
for f in path.glob(f"*{ext}"):
if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'):
files_to_process.append(f)
nsfw_threshold = NSFW_LEVELS.get(nsfw_level, NSFW_LEVELS["All"])
print(f"Found {len(files_to_process)} models to process.")
global_stats = {
"total_processed": len(files_to_process),
"new_hashes": 0,
"images_processed": 0,
"workflows_extracted": 0,
"space_saved": 0
}
# Import and use ComfyUI's ProgressBar
try:
import comfy.utils
pbar = comfy.utils.ProgressBar(len(files_to_process))
except ImportError:
pbar = None
if threads > 1:
print(f"Starting {threads} threads...")
with concurrent.futures.ThreadPoolExecutor(max_workers=threads) as executor:
futures = [executor.submit(process_file, fp, cache, api, nsfw_threshold, max_examples) for fp in files_to_process]
for future in concurrent.futures.as_completed(futures):
try:
res = future.result()
if res:
for k in global_stats: global_stats[k] += res.get(k, 0)
except Exception as e:
print(f"Error in thread: {e}")
if pbar:
pbar.update(1)
else:
for filepath in files_to_process:
res = process_file(filepath, cache, api, nsfw_threshold, max_examples)
if res:
for k in global_stats: global_stats[k] += res.get(k, 0)
if pbar:
pbar.update(1)
return global_stats
+102 -80
View File
@@ -46,14 +46,14 @@ def _index_sv_ratio(S: torch.Tensor, ratio: float) -> int:
def _index_sv_cumulative(S: torch.Tensor, target: float, max_rank: int = None) -> int:
"""Cumulative mode - keep enough SVs to reach target % of total.
Calculates relative to max_rank if provided, otherwise relative to full.
"""
if max_rank is not None and max_rank < len(S):
total = torch.sum(S[:max_rank])
else:
total = torch.sum(S)
if total < 1e-8:
return 1
cumsum = torch.cumsum(S, dim=0) / total
@@ -63,7 +63,7 @@ def _index_sv_cumulative(S: torch.Tensor, target: float, max_rank: int = None) -
def _index_sv_fro(S: torch.Tensor, target: float, max_rank: int = None) -> int:
"""Frobenius norm mode - preserve target fraction of Frobenius norm.
Calculates relative to max_rank if provided, otherwise relative to full.
This means "retain target% of what's achievable within max_rank".
"""
@@ -75,10 +75,10 @@ def _index_sv_fro(S: torch.Tensor, target: float, max_rank: int = None) -> int:
else:
S_sq = S.pow(2)
total_sq = torch.sum(S_sq)
if total_sq < 1e-8:
return 1
# Cumsum of all S (not capped) to find where we reach target
cumsum = torch.cumsum(S.pow(2), dim=0) / total_sq
rank = int(torch.searchsorted(cumsum, target ** 2).item()) + 1
@@ -172,27 +172,27 @@ def _svd_extract_linear_lowrank(
) -> tuple:
"""
Low-rank SVD decomposition using svd_lowrank for fixed rank extraction.
Much faster than full SVD when rank << min(m, n) because it only computes
the top-k singular values using randomized algorithms.
Args:
niter: Power iterations for SVD accuracy (higher = more accurate but slower)
Returns:
(lora_down, lora_up, diff), "low rank" or (weight, "full")
"""
weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape
# Clamp rank to valid range
max_possible_rank = min(out_dim, in_dim)
rank = max(1, min(rank, max_possible_rank - 1))
rank = max(1, min(rank, max_possible_rank))
# Check if decomposition is worthwhile
if rank >= max_possible_rank // 2:
if rank >= max_possible_rank:
return weight, "full"
# Use svd_lowrank for faster low-rank approximation
try:
U, S, V = torch.svd_lowrank(weight, q=rank, niter=niter)
@@ -200,7 +200,7 @@ def _svd_extract_linear_lowrank(
Vh = V.T # [rank, in_dim]
except Exception as e:
raise RuntimeError(f"svd_lowrank failed: {e}")
# Clamp outliers
if clamp_quantile < 1.0 and len(S) > 0:
try:
@@ -208,14 +208,14 @@ def _svd_extract_linear_lowrank(
S = S.clamp(max=max_val)
except RuntimeError:
pass
# Construct LoRA matrices
lora_up = U @ torch.diag(S) # [out_dim, rank]
lora_down = Vh # [rank, in_dim]
# Compute reconstruction diff
diff = weight - (lora_up @ lora_down)
return (lora_down, lora_up, diff), "low rank"
@@ -239,7 +239,7 @@ def _svd_extract_linear(
"""
weight = weight.to(device=device, dtype=torch.float32)
out_dim, in_dim = weight.shape
# For fixed rank, use the faster svd_lowrank
if mode == "fixed" and mode_param > 0:
target_rank = int(mode_param)
@@ -257,7 +257,7 @@ def _svd_extract_linear(
rank = _compute_rank(S, mode, mode_param, max_rank)
# Check if decomposition is worthwhile
if rank >= min(out_dim, in_dim) // 2:
if rank >= min(out_dim, in_dim):
return weight, "full"
# Truncate to rank
@@ -303,7 +303,7 @@ def _svd_extract_conv(
# Flatten for SVD
if is_1x1:
mat = weight.squeeze() # [out_ch, in_ch]
mat = weight.view(out_ch, in_ch) # [out_ch, in_ch]
else:
mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k]
@@ -382,8 +382,11 @@ def _extract_chunked_layer(
return None, None, 0
# Combine chunks
# Note: Using the first chunk's down matrix as the basis for all chunks
# is a better approximation than a simple mean if the chunks share the same input space (like QKV).
# Ideally, we would use a more sophisticated joint SVD, but this is a reasonable fallback.
combined_lora_up = torch.cat(all_lora_up, dim=0)
combined_lora_down = torch.stack(all_lora_down, dim=0).mean(dim=0)
combined_lora_down = all_lora_down[0]
rank = combined_lora_down.shape[0]
return combined_lora_up, combined_lora_down, rank
@@ -433,6 +436,8 @@ def extract_lora_from_files(
mismatch_mode: str = "skip",
chunk_large_layers: bool = True,
svd_niter: int = 2,
lazy_load: bool = True,
force_clear_cache: bool = True,
) -> dict[str, torch.Tensor]:
"""
Extract LoRA from difference between two models.
@@ -468,45 +473,46 @@ def extract_lora_from_files(
print(f"[LoRA Extract] Preparing memory for {total_size_gb:.2f}GB operation...")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device)) # 1.5x for SVD headroom
handler_a = MemoryEfficientSafeOpen(model_a_path)
handler_b = MemoryEfficientSafeOpen(model_b_path)
handler_a = MemoryEfficientSafeOpen(model_a_path, low_memory=lazy_load)
handler_b = MemoryEfficientSafeOpen(model_b_path, low_memory=lazy_load)
try:
keys_a = set(handler_a.keys())
keys_b = set(handler_b.keys())
weight_keys = [k for k in keys_a if k.endswith(".weight")]
pbar = comfy.utils.ProgressBar(len(weight_keys))
stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
def _process_layer(key):
lora_name = _format_lora_key(key)
if _matches_any_pattern(key, skip_patterns):
stats["skipped"] += 1
pbar.update(1)
continue
return "skipped", None
# Load tensors with pinned memory for CUDA
use_pinned = device == 'cuda'
if key not in keys_b:
if mismatch_mode == "skip":
stats["skipped"] += 1
pbar.update(1)
continue
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Key {key} not found in model B")
cpu_a = handler_a.get_tensor(key)
if use_pinned:
weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
else:
weight_diff = cpu_a.to(device=device, dtype=torch.float32)
del cpu_a
if mismatch_mode == "zeros":
# For zeros mode, we treat the missing base as a zeroed tensor of same shape
# which is already captured by weight_diff = tensor_a
pass
else:
cpu_a = handler_a.get_tensor(key)
cpu_b = handler_b.get_tensor(key)
if use_pinned:
tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
@@ -517,10 +523,12 @@ def extract_lora_from_files(
if tensor_a.shape != tensor_b.shape:
if mismatch_mode == "skip":
stats["skipped"] += 1
pbar.update(1)
del tensor_a, tensor_b
continue
return "skipped", None
if mismatch_mode == "error":
raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
# For zeros/fallback, use tensor_a as the difference
weight_diff = tensor_a
del tensor_b
else:
@@ -529,19 +537,16 @@ def extract_lora_from_files(
# Skip small differences
if min_diff > 0 and weight_diff.abs().max() < min_diff:
stats["skipped"] += 1
del weight_diff
pbar.update(1)
continue
return "skipped", None
# Skip 1D tensors
if weight_diff.ndim < 2:
stats["skipped"] += 1
del weight_diff
pbar.update(1)
continue
return "skipped", None
is_conv = weight_diff.ndim == 4
layer_results = {}
try:
if is_conv:
@@ -562,32 +567,42 @@ def extract_lora_from_files(
weight_diff, num_chunks, mode, linear_param, device, linear_max_rank
)
if lora_up is not None:
output_sd[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
output_sd[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
output_sd[f"{lora_name}.alpha"] = torch.tensor(rank).to(save_torch_dtype)
stats["chunked"] += 1
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(rank).to(save_torch_dtype)
del weight_diff
pbar.update(1)
continue
return "chunked", layer_results
print(f"[LoRA Extract] Failed: {key}: {e}")
stats["skipped"] += 1
del weight_diff
pbar.update(1)
continue
return "skipped", None
# Store result
if mode_str == "full":
output_sd[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
stats["full"] += 1
layer_results[f"{lora_name}.diff"] = weight_diff.to(save_torch_dtype).cpu().contiguous()
status = "full"
else:
lora_down, lora_up, _ = result
output_sd[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
output_sd[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
output_sd[f"{lora_name}.alpha"] = torch.tensor(lora_down.shape[0]).to(save_torch_dtype)
stats["extracted"] += 1
layer_results[f"{lora_name}.lora_down.weight"] = lora_down.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous()
layer_results[f"{lora_name}.alpha"] = torch.tensor(lora_down.shape[0]).to(save_torch_dtype)
status = "extracted"
del weight_diff
return status, layer_results
for key in tqdm(weight_keys, desc="Extracting LoRA", unit="layers"):
status, layer_sd = _process_layer(key)
stats[status] += 1
if layer_sd:
output_sd.update(layer_sd)
if force_clear_cache:
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
print(f"[LoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, "
@@ -698,7 +713,7 @@ def _format_lora_key(key: str) -> str:
# Handle already correct prefixes
if key.startswith("diffusion_model.") or key.startswith("transformer."):
return key
# Handle known Unet blocks (Diffusers format UNet)
if any(key.startswith(p) for p in ["down_blocks", "up_blocks", "mid_block", "conv_in", "conv_out", "time_embedding", "class_embedding"]):
return f"diffusion_model.{key}"
@@ -709,6 +724,8 @@ def _format_lora_key(key: str) -> str:
def _get_common_inputs():
return [
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
io.Boolean.Input("chunk_large_layers", default=False,
tooltip="Split large fused layers (QKV, MLP) into chunks"),
io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
@@ -736,9 +753,9 @@ class LoRAExtractFixed(io.ComfyNode):
description="Extract LoRA with specified fixed rank for each layer type.",
inputs=[
*_get_model_inputs(),
io.Int.Input("linear_dim", default=64, min=1, max=4096,
io.Int.Input("linear_dim", default=64, min=1, max=16384,
tooltip="Rank for linear/attention layers"),
io.Int.Input("conv_dim", default=32, min=1, max=4096,
io.Int.Input("conv_dim", default=32, min=1, max=16384,
tooltip="Rank for conv layers"),
io.Int.Input("svd_niter", default=2, min=0, max=10,
tooltip="SVD power iterations (higher = more accurate but slower)"),
@@ -751,7 +768,7 @@ class LoRAExtractFixed(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_dim, conv_dim, svd_niter, chunk_large_layers,
clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> io.NodeOutput:
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
@@ -759,7 +776,8 @@ class LoRAExtractFixed(io.ComfyNode):
output_sd = extract_lora_from_files(
model_a_path, model_b_path, "fixed", linear_dim, conv_dim,
device, save_dtype, linear_dim, conv_dim,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, svd_niter,
lazy_load, force_clear_cache
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
@@ -781,9 +799,9 @@ class LoRAExtractRatio(io.ComfyNode):
tooltip="Ratio threshold for linear layers (higher = more SVs kept)"),
io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1,
tooltip="Ratio threshold for conv layers (higher = more SVs kept)"),
io.Int.Input("linear_max_rank", default=128, min=1, max=3072,
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=3072,
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
@@ -794,7 +812,7 @@ class LoRAExtractRatio(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> io.NodeOutput:
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
@@ -802,7 +820,8 @@ class LoRAExtractRatio(io.ComfyNode):
output_sd = extract_lora_from_files(
model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
device, save_dtype, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
@@ -824,9 +843,9 @@ class LoRAExtractQuantile(io.ComfyNode):
tooltip="Target cumulative % for linear layers"),
io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target cumulative % for conv layers"),
io.Int.Input("linear_max_rank", default=128, min=1, max=3072,
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=3072,
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
@@ -837,7 +856,7 @@ class LoRAExtractQuantile(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> io.NodeOutput:
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
@@ -845,7 +864,8 @@ class LoRAExtractQuantile(io.ComfyNode):
output_sd = extract_lora_from_files(
model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
device, save_dtype, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
@@ -865,9 +885,9 @@ class LoRAExtractKnee(io.ComfyNode):
*_get_model_inputs(),
io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"],
default="sv_knee", tooltip="Knee detection method"),
io.Int.Input("linear_max_rank", default=128, min=1, max=3072,
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=3072,
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
@@ -878,7 +898,7 @@ class LoRAExtractKnee(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> io.NodeOutput:
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
@@ -886,7 +906,8 @@ class LoRAExtractKnee(io.ComfyNode):
output_sd = extract_lora_from_files(
model_a_path, model_b_path, knee_method, 0, 0,
device, save_dtype, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
@@ -908,9 +929,9 @@ class LoRAExtractFrobenius(io.ComfyNode):
tooltip="Target Frobenius norm fraction for linear"),
io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01,
tooltip="Target Frobenius norm fraction for conv"),
io.Int.Input("linear_max_rank", default=128, min=1, max=3072,
io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for linear layers"),
io.Int.Input("conv_max_rank", default=128, min=1, max=3072,
io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
tooltip="Maximum rank for conv layers"),
*_get_common_inputs(),
],
@@ -921,7 +942,7 @@ class LoRAExtractFrobenius(io.ComfyNode):
@classmethod
def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
save_dtype, device, skip_patterns) -> io.NodeOutput:
save_dtype, device, skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
model_a_path = folder_paths.get_full_path_or_raise("diffusion_models", model_a)
model_b_path = folder_paths.get_full_path_or_raise("diffusion_models", model_b)
@@ -929,7 +950,8 @@ class LoRAExtractFrobenius(io.ComfyNode):
output_sd = extract_lora_from_files(
model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
device, save_dtype, linear_max_rank, conv_max_rank,
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers
clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
lazy_load=lazy_load, force_clear_cache=force_clear_cache
)
return io.NodeOutput(_save_lora(output_sd, output_filename))
+73 -57
View File
@@ -13,8 +13,8 @@ from typing import List, Dict, Tuple, Optional
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
from .lora_resize import (
detect_lora_format,
extract_lora_pairs,
detect_lora_format,
extract_lora_pairs,
detect_lora_rank,
_compile_patterns,
_matches_any_pattern,
@@ -32,7 +32,7 @@ def merge_multi_loras(
) -> str:
"""
Merge multiple LoRAs into a single LoRA file.
Modes:
- concatenate: Mathematically sound merge by stacking ranks.
New rank = sum(input ranks). Alpha set to new rank.
@@ -45,22 +45,22 @@ def merge_multi_loras(
if verbose:
print(f"[LoRA Multi-Merge] Preparing memory for {total_size_gb:.2f}GB operation...")
print(f"[LoRA Multi-Merge] Mode: {merge_mode}")
prepare_for_large_operation(total_size_gb * 2.5, torch.device(device))
handlers = [MemoryEfficientSafeOpen(p) for p in lora_paths]
handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths]
try:
# 1. Analyze all LoRAs
layer_map = {}
lora_infos = []
for i, handler in enumerate(handlers):
keys = handler.keys()
fmt = detect_lora_format(keys)
pairs = extract_lora_pairs(keys, fmt)
rank, alpha = detect_lora_rank(handler, pairs)
info = {
"name": os.path.basename(lora_paths[i]),
"format": fmt,
@@ -71,10 +71,10 @@ def merge_multi_loras(
"scale": alpha / rank if rank > 0 else 1.0
}
lora_infos.append(info)
if verbose:
print(f"[LoRA Multi-Merge] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}, format={fmt['format']}")
for block_name in pairs.keys():
core = _format_lora_key(block_name)
if core not in layer_map:
@@ -94,37 +94,37 @@ def merge_multi_loras(
ups = []
is_conv = False
max_rank = 0
# First pass: load and determine max rank
for idx in indices:
info = lora_infos[idx]
handler = handlers[idx]
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
block_keys = info["pairs"][orig_block]
if "down" not in block_keys or "up" not in block_keys:
continue
t_down = handler.get_tensor(block_keys["down"])
t_up = handler.get_tensor(block_keys["up"])
if device == 'cuda':
t_down = transfer_to_gpu_pinned(t_down, device, torch.float32)
t_up = transfer_to_gpu_pinned(t_up, device, torch.float32)
else:
t_down = t_down.to(device=device, dtype=torch.float32)
t_up = t_up.to(device=device, dtype=torch.float32)
# Store with original info for padding/weighting
is_conv = len(t_down.shape) == 4
current_rank = t_down.shape[0]
max_rank = max(max_rank, current_rank)
# Apply scale immediately for concatenate mode, or keep for later
if merge_mode == "concatenate":
effective_weight = info["weight"] * info["scale"]
t_up = t_up * effective_weight
downs.append((t_down, info))
ups.append((t_up, info))
@@ -148,10 +148,10 @@ def merge_multi_loras(
target_down_shape[0] = max_rank
target_up_shape = list(ups[0][0].shape)
target_up_shape[1] = max_rank
merged_down = torch.zeros(target_down_shape, device=device, dtype=torch.float32)
merged_up = torch.zeros(target_up_shape, device=device, dtype=torch.float32)
for (t_d, info_d), (t_u, info_u) in zip(downs, ups):
r = t_d.shape[0]
w = info_d["weight"]
@@ -159,7 +159,7 @@ def merge_multi_loras(
# Kohya just uses weights. But to be safe with different alphas,
# we apply scale to the final delta.
# For direct weight merge, we'll just use weights.
if r < max_rank:
# Pad down: [r, in...] -> [max_rank, in...]
pad_d = [0] * (len(t_d.shape) * 2)
@@ -168,15 +168,15 @@ def merge_multi_loras(
# For [r, in], padding is (0,0, 0, max_rank-r)
padding_d = [0, 0] * (len(t_d.shape) - 1) + [0, max_rank - r]
t_d = torch.nn.functional.pad(t_d, tuple(padding_d))
# Pad up: [out, r...] -> [out, max_rank...]
# For [out, r], padding is (0, max_rank-r, 0, 0)
padding_u = [0, max_rank - r] + [0, 0] * (len(t_u.shape) - 1)
t_u = torch.nn.functional.pad(t_u, tuple(padding_u))
merged_down += w * t_d
merged_up += w * t_u
new_rank = max_rank
# Alpha is usually max of alphas or same as rank
new_alpha = float(max_rank)
@@ -185,12 +185,18 @@ def merge_multi_loras(
output_sd[f"{core}.lora_down.weight"] = merged_down.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.alpha"] = torch.tensor(new_alpha, dtype=save_dtype)
# Cleanup
del merged_down, merged_up
for d, _ in downs: del d
for u, _ in ups: del u
downs.clear()
ups.clear()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
# 3. Final Summary
@@ -199,7 +205,7 @@ def merge_multi_loras(
for indices in layer_map.values():
for idx in indices:
matched_stats[idx] += 1
print(f"[LoRA Multi-Merge] --- Merge Summary ---")
for i, info in enumerate(lora_infos):
print(f"[LoRA Multi-Merge] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge")
@@ -210,16 +216,16 @@ def merge_multi_loras(
"ss_training_comment": f"Merged {len(lora_paths)} LoRAs via concatenation",
"ss_network_module": "networks.lora",
}
# 4. Save
output_dir = folder_paths.get_folder_paths("loras")[0]
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, f"{output_filename.strip()}.safetensors")
save_file(output_sd, output_path, metadata)
if verbose:
print(f"[LoRA Multi-Merge] Saved merged LoRA to {output_path}")
return output_path
finally:
@@ -265,7 +271,7 @@ class LoRAMultiMerge(io.ComfyNode):
# LoRA 8
io.Combo.Input("lora_8", options=["None"] + folder_paths.get_filename_list("loras"), default="None"),
io.Float.Input("weight_8", default=1.0, min=-10.0, max=10.0, step=0.01),
io.Combo.Input("merge_mode", options=["concatenate", "weighted_sum"], default="concatenate",
tooltip="concatenate: safe, increases rank. weighted_sum: fixed rank (to max input rank), mathematically lossy but standard in Kohya."),
io.String.Input("output_filename", default="merged_lora"),
@@ -281,25 +287,25 @@ class LoRAMultiMerge(io.ComfyNode):
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
merge_mode, output_filename, save_dtype, device) -> io.NodeOutput:
count = int(lora_count)
names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8]
weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8]
valid_paths = []
valid_weights = []
for i in range(count):
if names[i] and names[i] != "None":
path = folder_paths.get_full_path_or_raise("loras", names[i])
valid_paths.append(path)
valid_weights.append(weights[i])
if not valid_paths:
raise ValueError("No LoRAs selected for merging.")
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = merge_multi_loras(
lora_paths=valid_paths,
lora_weights=valid_weights,
@@ -308,7 +314,7 @@ class LoRAMultiMerge(io.ComfyNode):
save_dtype=dtype,
output_filename=output_filename
)
return io.NodeOutput(path)
@@ -359,25 +365,25 @@ class LoRAMultiMergeDARE(io.ComfyNode):
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
drop_rate, trim_quantile, seed, output_filename, save_dtype, device) -> io.NodeOutput:
count = int(lora_count)
names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8]
weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8]
valid_paths = []
valid_weights = []
for i in range(count):
if names[i] and names[i] != "None":
path = folder_paths.get_full_path_or_raise("loras", names[i])
valid_paths.append(path)
valid_weights.append(weights[i])
if not valid_paths:
raise ValueError("No LoRAs selected for merging.")
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[save_dtype]
path = merge_multi_loras_dare(
lora_paths=valid_paths,
lora_weights=valid_weights,
@@ -388,7 +394,7 @@ class LoRAMultiMergeDARE(io.ComfyNode):
save_dtype=dtype,
output_filename=output_filename
)
return io.NodeOutput(path)
def merge_multi_loras_dare(
@@ -412,8 +418,8 @@ def merge_multi_loras_dare(
print(f"[LoRA Multi-Merge DARE] Drop rate: {drop_rate}, Trim quantile: {trim_quantile}")
prepare_for_large_operation(total_size_gb * 2.5, torch.device(device))
handlers = [MemoryEfficientSafeOpen(p) for p in lora_paths]
handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths]
rng = torch.Generator(device=device).manual_seed(seed)
try:
@@ -433,7 +439,7 @@ def merge_multi_loras_dare(
"weight": lora_weights[i]
}
lora_infos.append(info)
if verbose:
print(f"[LoRA Multi-Merge DARE] LoRA {i+1} ({info['name']}): {len(pairs)} layers, dim={rank}")
@@ -454,16 +460,16 @@ def merge_multi_loras_dare(
downs = []
ups = []
max_rank = 0
for idx in indices:
info = lora_infos[idx]
orig_block = next(b for b in info["pairs"].keys() if _format_lora_key(b) == core)
block_keys = info["pairs"][orig_block]
if "down" not in block_keys or "up" not in block_keys: continue
t_d = handlers[idx].get_tensor(block_keys["down"]).to(device=device, dtype=torch.float32)
t_u = handlers[idx].get_tensor(block_keys["up"]).to(device=device, dtype=torch.float32)
max_rank = max(max_rank, t_d.shape[0])
downs.append((t_d, info["weight"]))
ups.append((t_u, info["weight"]))
@@ -484,13 +490,13 @@ def merge_multi_loras_dare(
padding[rev_dim*2 + 1] = max_rank - t.shape[dim_to_pad]
t = torch.nn.functional.pad(t, tuple(padding))
padded.append(t * w)
# DARE
if drop_rate > 0:
for i in range(len(padded)):
mask = (torch.rand(padded[i].shape, generator=rng, device=device) > drop_rate).float()
padded[i] = (padded[i] * mask) / (1 - drop_rate)
# TIES
# 1. Trim
if trim_quantile > 0:
@@ -500,17 +506,17 @@ def merge_multi_loras_dare(
if k > 0:
threshold = torch.kthvalue(flat, k).values
padded[i] = torch.where(padded[i].abs() < threshold, torch.zeros_like(padded[i]), padded[i])
# 2. Elect & Merge
stacked = torch.stack(padded) # [N, ...]
signs = torch.sign(stacked)
sum_signs = signs.sum(dim=0)
dominant_sign = torch.sign(sum_signs)
# Filter those matching dominant sign
mask = (signs == dominant_sign) & (dominant_sign != 0)
filtered = torch.where(mask, stacked, torch.zeros_like(stacked))
# Average matching signs
count = mask.sum(dim=0)
result = filtered.sum(dim=0) / torch.clamp(count, min=1.0)
@@ -522,7 +528,17 @@ def merge_multi_loras_dare(
output_sd[f"{core}.lora_down.weight"] = merged_down.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous()
output_sd[f"{core}.alpha"] = torch.tensor(float(max_rank), dtype=save_dtype)
del merged_down, merged_up
for d, _ in downs: del d
for u, _ in ups: del u
downs.clear()
ups.clear()
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
# Final Summary
@@ -531,7 +547,7 @@ def merge_multi_loras_dare(
for indices in layer_map.values():
for idx in indices:
matched_stats[idx] += 1
print(f"[LoRA Multi-Merge DARE] --- Merge Summary ---")
for i, info in enumerate(lora_infos):
print(f"[LoRA Multi-Merge DARE] LoRA {i+1}: {matched_stats[i]}/{len(info['pairs'])} layers used in merge")
+35 -4
View File
@@ -461,6 +461,12 @@ def resize_lora_file(
if alpha_suffix:
output_sd[f"{new_block_name}{alpha_suffix}"] = torch.tensor(result["new_alpha"]).to(save_dtype)
del result
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
if verbose and fro_list:
@@ -854,6 +860,10 @@ def merge_loras_to_model(
del cpu_base
stats["copied"] += 1
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
@@ -886,11 +896,11 @@ class LoRAMergeToModel(io.ComfyNode):
node_id="LoRAMergeToModel",
display_name="LoRA Merge To Model",
category="ModelUtils/LoRA/Merge",
description="Merge 1-4 LoRAs into a base model and save the result. Saves to base model directory.",
description="Merge 1-8 LoRAs into a base model and save the result. Saves to base model directory.",
inputs=[
io.Combo.Input("base_model", options=folder_paths.get_filename_list("diffusion_models"),
tooltip="Base model the LoRAs were trained on"),
io.Combo.Input("lora_count", options=["1", "2", "3", "4"], default="2",
io.Combo.Input("lora_count", options=["1", "2", "3", "4", "5", "6", "7", "8"], default="2",
tooltip="Number of LoRAs to merge"),
# LoRA 1
io.Combo.Input("lora_1", options=folder_paths.get_filename_list("loras"),
@@ -912,6 +922,26 @@ class LoRAMergeToModel(io.ComfyNode):
default="None", tooltip="Fourth LoRA"),
io.Float.Input("weight_4", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 4"),
# LoRA 5
io.Combo.Input("lora_5", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Fifth LoRA"),
io.Float.Input("weight_5", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 5"),
# LoRA 6
io.Combo.Input("lora_6", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Sixth LoRA"),
io.Float.Input("weight_6", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 6"),
# LoRA 7
io.Combo.Input("lora_7", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Seventh LoRA"),
io.Float.Input("weight_7", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 7"),
# LoRA 8
io.Combo.Input("lora_8", options=["None"] + folder_paths.get_filename_list("loras"),
default="None", tooltip="Eighth LoRA"),
io.Float.Input("weight_8", default=1.0, min=-10.0, max=10.0, step=0.01,
tooltip="Weight strength for LoRA 8"),
# Settings
io.String.Input("skip_patterns", default="", multiline=True,
tooltip="Regex patterns for layers to skip"),
@@ -926,12 +956,13 @@ class LoRAMergeToModel(io.ComfyNode):
@classmethod
def execute(cls, base_model, lora_count,
lora_1, weight_1, lora_2, weight_2, lora_3, weight_3, lora_4, weight_4,
lora_5, weight_5, lora_6, weight_6, lora_7, weight_7, lora_8, weight_8,
skip_patterns, output_filename, save_dtype, device) -> io.NodeOutput:
# Build LoRA list based on count
count = int(lora_count)
lora_names = [lora_1, lora_2, lora_3, lora_4][:count]
lora_weights = [weight_1, weight_2, weight_3, weight_4][:count]
lora_names = [lora_1, lora_2, lora_3, lora_4, lora_5, lora_6, lora_7, lora_8][:count]
lora_weights = [weight_1, weight_2, weight_3, weight_4, weight_5, weight_6, weight_7, weight_8][:count]
# Filter out "None" entries
valid_loras = [(name, weight) for name, weight in zip(lora_names, lora_weights) if name != "None"]
+124 -71
View File
@@ -83,11 +83,11 @@ class MergerLogic:
path = folder_paths.get_full_path(model_type, name)
if not path:
raise FileNotFoundError(f"Model '{name}' not found.")
handlers[name] = MemoryEfficientSafeOpen(path)
handlers[name] = MemoryEfficientSafeOpen(path, low_memory=True)
primary_handler = handlers[primary_model_name]
all_keys = primary_handler.keys()
metadata = primary_handler.header.get("__metadata__", {})
metadata = primary_handler.metadata()
# Convert mismatch_mode string to enum
mismatch_mode_str = recipe_params.get('mismatch_mode', 'skip')
@@ -124,65 +124,78 @@ class MergerLogic:
discarded_keys = 0
error_keys = []
for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
# Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns):
discarded_keys += 1
with torch.no_grad():
for key in tqdm(all_keys, desc="Merging layers", unit="layers"):
# Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns):
discarded_keys += 1
pbar.update(1)
continue
# Pre-load Model A's tensor with pinned memory for CUDA
cpu_tensor = primary_handler.get_tensor(key)
if process_device == 'cuda':
tensor_a = transfer_to_gpu_pinned(cpu_tensor, process_device, process_dtype)
else:
tensor_a = cpu_tensor.to(device=process_device, dtype=process_dtype)
del cpu_tensor
# Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns):
merged_state_dict[key] = tensor_a.detach().to(save_torch_dtype).cpu().clone()
excluded_keys += 1
pbar.update(1)
continue
# Pass tensor_a metadata to recipes for zeros mode and fallback
recipe_params['_tensor_a'] = tensor_a
recipe_params['_tensor_a_shape'] = tensor_a.shape
recipe_params['_tensor_a_dtype'] = tensor_a.dtype
try:
recipe = calc_mode_class.create_recipe(key=key, **recipe_params)
result = recipe.merge()
except MissingTensorError as e:
if mismatch_mode == MissingTensorBehavior.ERROR:
raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}")
result = None
error_keys.append(key)
# Handle None result (mismatch occurred with skip mode)
if result is None:
result = tensor_a
skipped_keys += 1
if isinstance(result, dict):
for r_key, r_tensor in result.items():
merged_state_dict[r_key] = r_tensor.detach().to(save_torch_dtype).cpu().clone()
else:
# Ensure compatibility with Model A's architecture.
# If alignment_mode is 'pad/crop', we crop results that were padded.
# If alignment_mode is 'interpolate', resizing happened during operators.
if alignment_mode == 'pad/crop':
target_shape = recipe_params['_tensor_a_shape']
if result.shape != target_shape:
slices = tuple(slice(0, min(res_s, tgt_s)) for res_s, tgt_s in zip(result.shape, target_shape))
result = result[slices]
merged_state_dict[key] = result.detach().to(save_torch_dtype).cpu().clone()
# Clean up references to allow GC immediately.
# Local loop variables must be explicitly deleted to prevent PyTorch from keeping tensors in VRAM.
recipe.clean()
del recipe_params['_tensor_a']
del tensor_a
del recipe
del result
if recipe_params.get('force_clear_cache', True):
import gc
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
pbar.update(1)
continue
# Pre-load Model A's tensor with pinned memory for CUDA
cpu_tensor = primary_handler.get_tensor(key)
if process_device == 'cuda':
tensor_a = transfer_to_gpu_pinned(cpu_tensor, process_device, process_dtype)
else:
tensor_a = cpu_tensor.to(device=process_device, dtype=process_dtype)
del cpu_tensor
# Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns):
merged_state_dict[key] = tensor_a.to(save_torch_dtype).cpu().clone()
excluded_keys += 1
pbar.update(1)
continue
# Pass tensor_a metadata to recipes for zeros mode and fallback
recipe_params['_tensor_a'] = tensor_a
recipe_params['_tensor_a_shape'] = tensor_a.shape
recipe_params['_tensor_a_dtype'] = tensor_a.dtype
try:
recipe = calc_mode_class.create_recipe(key=key, **recipe_params)
result = recipe.merge()
except MissingTensorError as e:
if mismatch_mode == MissingTensorBehavior.ERROR:
raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}")
result = None
error_keys.append(key)
# Handle None result (mismatch occurred with skip mode)
if result is None:
result = tensor_a
skipped_keys += 1
if isinstance(result, dict):
for r_key, r_tensor in result.items():
merged_state_dict[r_key] = r_tensor.to(save_torch_dtype).cpu().clone()
else:
# Ensure compatibility with Model A's architecture.
# If alignment_mode is 'pad/crop', we crop results that were padded.
# If alignment_mode is 'interpolate', resizing happened during operators.
if alignment_mode == 'pad/crop':
target_shape = recipe_params['_tensor_a_shape']
if result.shape != target_shape:
slices = tuple(slice(0, min(res_s, tgt_s)) for res_s, tgt_s in zip(result.shape, target_shape))
result = result[slices]
merged_state_dict[key] = result.to(save_torch_dtype).cpu().clone()
# Clean up tensor_a reference to allow GC
del recipe_params['_tensor_a']
pbar.update(1)
# Log summary
if excluded_keys > 0:
@@ -240,6 +253,8 @@ class CheckpointTwoMerger(io.ComfyNode):
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -251,7 +266,8 @@ class CheckpointTwoMerger(io.ComfyNode):
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput:
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -263,6 +279,7 @@ class CheckpointTwoMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -295,6 +312,8 @@ class ModelTwoMerger(io.ComfyNode):
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -306,7 +325,8 @@ class ModelTwoMerger(io.ComfyNode):
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput:
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -318,6 +338,7 @@ class ModelTwoMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -350,6 +371,8 @@ class TextEncoderTwoMerger(io.ComfyNode):
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -361,7 +384,8 @@ class TextEncoderTwoMerger(io.ComfyNode):
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput:
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -373,6 +397,7 @@ class TextEncoderTwoMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -405,6 +430,8 @@ class LoRATwoMerger(io.ComfyNode):
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -416,7 +443,8 @@ class LoRATwoMerger(io.ComfyNode):
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput:
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -428,6 +456,7 @@ class LoRATwoMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -460,6 +489,8 @@ class EmbeddingTwoMerger(io.ComfyNode):
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -471,7 +502,8 @@ class EmbeddingTwoMerger(io.ComfyNode):
def execute(cls, execution_mode: str, model_a: str, model_b: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput:
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_2_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -483,6 +515,7 @@ class EmbeddingTwoMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b}
filename = MergerLogic.execute_merge(model_names, calc_mode, TWO_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -518,6 +551,8 @@ class CheckpointThreeMerger(io.ComfyNode):
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -529,7 +564,8 @@ class CheckpointThreeMerger(io.ComfyNode):
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput:
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -541,6 +577,7 @@ class CheckpointThreeMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -574,6 +611,8 @@ class ModelThreeMerger(io.ComfyNode):
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -585,7 +624,8 @@ class ModelThreeMerger(io.ComfyNode):
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput:
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -597,6 +637,7 @@ class ModelThreeMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -622,7 +663,7 @@ class TextEncoderThreeMerger(io.ComfyNode):
io.Combo.Input("alignment_mode", options=["pad/crop", "interpolate"], default="pad/crop"),
io.Float.Input("alpha", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("beta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Float.Input("gamma", default=0.99, min=0.0, max=1.0, step=0.001),
io.Float.Input("delta", default=0.5, min=-2.0, max=3.0, step=0.01),
io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
io.String.Input("output_filename", default="merged_3_textencoder"),
@@ -630,6 +671,8 @@ class TextEncoderThreeMerger(io.ComfyNode):
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -641,7 +684,8 @@ class TextEncoderThreeMerger(io.ComfyNode):
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput:
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -653,6 +697,7 @@ class TextEncoderThreeMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -686,6 +731,8 @@ class LoRAThreeMerger(io.ComfyNode):
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -697,7 +744,8 @@ class LoRAThreeMerger(io.ComfyNode):
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput:
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -709,6 +757,7 @@ class LoRAThreeMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
@@ -742,6 +791,8 @@ class EmbeddingThreeMerger(io.ComfyNode):
io.Combo.Input("process_device", options=["cuda", "cpu"]),
io.String.Input("exclude_patterns", default="", multiline=True),
io.String.Input("discard_patterns", default="", multiline=True),
io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
],
outputs=[
io.String.Output(display_name="output_filename"),
@@ -753,7 +804,8 @@ class EmbeddingThreeMerger(io.ComfyNode):
def execute(cls, execution_mode: str, model_a: str, model_b: str, model_c: str,
calc_mode: str, mismatch_mode: str, alignment_mode: str, alpha: float, beta: float,
gamma: float, delta: float, seed: int, output_filename: str, save_dtype: str,
process_device: str, exclude_patterns: str, discard_patterns: str) -> io.NodeOutput:
process_device: str, exclude_patterns: str, discard_patterns: str,
lazy_load: bool, force_clear_cache: bool) -> io.NodeOutput:
doc = load_documentation_from_file('merger_3_model_modes.md')
if execution_mode == "DOCUMENTATION ONLY":
return io.NodeOutput("Documentation mode active. No merge performed.", doc)
@@ -765,6 +817,7 @@ class EmbeddingThreeMerger(io.ComfyNode):
"output_filename": output_filename, "save_dtype": save_dtype,
"device": process_device, "dtype": torch.float32,
"exclude_patterns": exclude_patterns, "discard_patterns": discard_patterns,
"lazy_load": lazy_load, "force_clear_cache": force_clear_cache,
}
model_names = {"model_a": model_a, "model_b": model_b, "model_c": model_c}
filename = MergerLogic.execute_merge(model_names, calc_mode, THREE_MODEL_MODES, recipe_params, cls.MODEL_TYPE)
+28 -3
View File
@@ -4,6 +4,8 @@ import torch.nn.functional as F
from enum import Enum
from .merger_utils import transfer_to_gpu_pinned
class MissingTensorBehavior(Enum):
"""Controls behavior when a tensor key is missing from a model."""
ERROR = "error" # Raise exception (strict mode)
@@ -127,12 +129,21 @@ class Operation:
raise NotImplementedError
def recurse(self, operation):
source_tensors = [source_oper.merge() for source_oper in operation.sources]
return operation.oper(*source_tensors)
self._source_tensors = [source_oper.merge() for source_oper in operation.sources]
return operation.oper(*self._source_tensors)
def merge(self):
return self.merge_func(self)
def clean(self):
if hasattr(self, '_source_tensors'):
for t in self._source_tensors:
del t
del self._source_tensors
for source in self.sources:
if hasattr(source, 'clean'):
source.clean()
class LoadTensor(Operation):
def __init__(self, key, model_name, handlers, device, dtype,
on_missing=MissingTensorBehavior.ERROR, fallback_shape=None, fallback_dtype=None):
@@ -165,7 +176,17 @@ class LoadTensor(Operation):
dtype = self.fallback_dtype if self.fallback_dtype else self.dtype
return torch.zeros(self.fallback_shape, device=self.device, dtype=dtype)
return handler.get_tensor(self.key).to(device=self.device, dtype=self.dtype)
cpu_tensor = handler.get_tensor(self.key)
if self.device == 'cuda':
self._tensor = transfer_to_gpu_pinned(cpu_tensor, self.device, self.dtype)
else:
self._tensor = cpu_tensor.to(device=self.device, dtype=self.dtype)
del cpu_tensor
return self._tensor
def clean(self):
if hasattr(self, '_tensor'):
del self._tensor
class Multiply(Operation):
def __init__(self, key, alpha, *sources):
@@ -409,6 +430,10 @@ class PreloadedTensor(Operation):
def merge(self) -> torch.Tensor:
return self.tensor
def clean(self):
if hasattr(self, 'tensor'):
del self.tensor
class WeightSum(CalcMode):
name = 'Weight-Sum'
description = 'A * (1 - α) + B * α. Simple linear interpolation.'
+64 -145
View File
@@ -1,133 +1,89 @@
"""
Memory-efficient safetensors utilities with parallel I/O and pinned memory support.
Based on benchmarking, parallel reading with ThreadPoolExecutor provides 2-4x speedup
for large models compared to sequential reads.
Worker count is automatically optimized based on device capabilities when workers=None.
Pinned memory enables faster CPU→GPU transfers (2-3x improvement).
Memory-efficient safetensors utilities with pinned memory support.
Forces direct file reads to minimize OS page cache usage and avoids VRAM accumulation.
"""
import os
import mmap
import gc
import json
import struct
import torch
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List, Optional, Any, Tuple
from safetensors import safe_open
from typing import Dict, List, Optional, Any
# Import ComfyUI's memory pinning functions if available
try:
from comfy.model_management import pin_memory, unpin_memory, is_device_cuda
PINNING_AVAILABLE = True
except ImportError:
PINNING_AVAILABLE = False
def pin_memory(tensor): return False
def unpin_memory(tensor): return False
def is_device_cuda(device): return str(device).startswith('cuda')
# Module-level configuration for pinned transfer
_verbose = False
_pinned_transfer_stats = {"pinned": 0, "fallback": 0}
def set_verbose(enabled: bool):
"""Enable/disable verbose output for pinned transfers."""
global _verbose
_verbose = enabled
class PinnedTensor:
"""Context manager for automatic pinned memory management.
Usage:
with PinnedTensor(cpu_tensor) as pinned:
gpu_tensor = pinned.to('cuda')
# Memory automatically unpinned after transfer
"""
def __init__(self, tensor: torch.Tensor, auto_pin: bool = True):
self.tensor = tensor
self.is_pinned = False
self.auto_pin = auto_pin and PINNING_AVAILABLE
def __enter__(self) -> torch.Tensor:
if self.auto_pin and self.tensor.device.type == 'cpu':
self.is_pinned = pin_memory(self.tensor)
return self.tensor
def __exit__(self, exc_type, exc_val, exc_tb):
if self.is_pinned:
unpin_memory(self.tensor)
self.is_pinned = False
def get_pinned_transfer_stats():
"""Return pinned transfer statistics."""
return _pinned_transfer_stats.copy()
def reset_pinned_transfer_stats():
"""Reset transfer statistics."""
global _pinned_transfer_stats
_pinned_transfer_stats = {"pinned": 0, "fallback": 0}
def transfer_to_gpu_pinned(
tensor: torch.Tensor,
device: str = 'cuda',
dtype: Optional[torch.dtype] = None
) -> torch.Tensor:
"""Transfer tensor to GPU using pinned memory for faster transfer.
"""Transfer tensor to GPU using pinned memory for faster transfer."""
global _pinned_transfer_stats
Args:
tensor: CPU tensor to transfer
device: Target GPU device
dtype: Optional dtype conversion
if tensor.device.type != 'cpu' or not torch.cuda.is_available():
if dtype is not None:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
Returns:
Tensor on GPU device
"""
if not PINNING_AVAILABLE or not is_device_cuda(torch.device(device)):
# Fall back to regular transfer
if dtype:
if not str(device).startswith('cuda'):
if dtype is not None:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
# Pin memory for faster transfer
was_pinned = pin_memory(tensor)
try:
if dtype:
result = tensor.to(device=device, dtype=dtype, non_blocking=True)
pinned = tensor.pin_memory()
if dtype is not None:
result = pinned.to(device=device, dtype=dtype, non_blocking=True)
else:
result = tensor.to(device=device, non_blocking=True)
# Sync to ensure transfer is complete before unpinning
if was_pinned:
torch.cuda.current_stream().synchronize()
result = pinned.to(device=device, non_blocking=True)
torch.cuda.current_stream().synchronize()
_pinned_transfer_stats["pinned"] += 1
return result
finally:
if was_pinned:
unpin_memory(tensor)
except Exception as e:
_pinned_transfer_stats["fallback"] += 1
if dtype is not None:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
class MemoryEfficientSafeOpen:
"""Memory-efficient safetensors file reader with mmap, parallel I/O, and pinned memory support.
"""Memory-efficient safetensors file reader with streaming and pinned memory support.
Features:
- mmap mode: Zero-copy tensor access via memory-mapped file
- low_memory mode: Direct file reads to minimize OS page cache usage
- Parallel loading: Multi-threaded tensor reads for 2-4x speedup
- Sorted batch reads: Keys sorted by file offset for sequential I/O
- Auto-optimized workers: Adjusts parallelism based on device capabilities
- Pinned memory: Faster CPU→GPU transfers when transferring to CUDA
Args:
filename: Path to safetensors file
device: Target device (default 'cpu')
mmap_mode: Use memory-mapped file for zero-copy (default True)
low_memory: Use direct file reads to minimize memory (overrides mmap_mode)
Operates strictly in streaming mode (low memory). Tensors are read byte-by-byte from
their exact offsets without using safe_open or memory mapping, ensuring tensors do not
remain persistently in OS page caches or RAM.
"""
def __init__(self, filename: str, device: str = 'cpu', mmap_mode: bool = True, low_memory: bool = False):
def __init__(self, filename: str, device: str = 'cpu', **kwargs):
# Ignore mmap_mode, low_memory, or any other legacy kwargs. We always stream.
self.filename = filename
self.device = device
# low_memory mode forces mmap off to avoid OS page caching
self.low_memory = low_memory
self.mmap_mode = mmap_mode and not low_memory
self.header, self.header_size = self._read_header()
self.file = open(filename, "rb")
self.mmap_obj = None
if self.mmap_mode:
self.mmap_obj = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_READ)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.mmap_obj:
self.mmap_obj.close()
self.file.close()
def keys(self) -> List[str]:
"""Return all tensor keys (excluding metadata)."""
return [k for k in self.header.keys() if k != "__metadata__"]
@@ -136,8 +92,17 @@ class MemoryEfficientSafeOpen:
"""Return file metadata."""
return self.header.get("__metadata__", {})
def get_shape(self, key: str) -> tuple:
"""Get tensor shape without loading tensor data."""
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found in file")
return tuple(self.header[key]["shape"])
def get_ndim(self, key: str) -> int:
"""Get tensor ndim without loading tensor data."""
return len(self.get_shape(key))
def keys_sorted_by_offset(self) -> List[str]:
"""Return keys sorted by file offset for optimal sequential I/O."""
keys_with_offsets = []
for key in self.keys():
offset = self.header[key]["data_offsets"][0]
@@ -146,48 +111,28 @@ class MemoryEfficientSafeOpen:
return [k for k, _ in keys_with_offsets]
def get_tensor(self, key: str) -> torch.Tensor:
"""Load a single tensor by key."""
"""Load a single tensor by key directly from file byte offset."""
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found in the file")
metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"]
if self.mmap_mode and self.mmap_obj:
if offset_start != offset_end:
start = self.header_size + 8 + offset_start
end = self.header_size + 8 + offset_end
tensor_bytes = memoryview(self.mmap_obj)[start:end]
else:
tensor_bytes = None
if offset_start != offset_end:
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = bytearray(offset_end - offset_start)
self.file.readinto(tensor_bytes)
else:
# Non-mmap mode: use pre-allocated bytearray with readinto for minimal copies
tensor_bytes = None
if offset_start != offset_end:
self.file.seek(self.header_size + 8 + offset_start)
# Pre-allocate writable buffer and read directly into it
tensor_bytes = bytearray(offset_end - offset_start)
self.file.readinto(tensor_bytes)
return self._deserialize_tensor(tensor_bytes, metadata)
def get_tensor_to_gpu(
self,
key: str,
device: str = 'cuda',
dtype: Optional[torch.dtype] = None
) -> torch.Tensor:
"""Load a tensor and transfer to GPU using pinned memory for faster transfer.
Args:
key: Tensor key to load
device: Target GPU device (default 'cuda')
dtype: Optional dtype conversion
Returns:
Tensor on GPU device
"""
cpu_tensor = self.get_tensor(key)
return transfer_to_gpu_pinned(cpu_tensor, device, dtype)
@@ -197,20 +142,9 @@ class MemoryEfficientSafeOpen:
workers: Optional[int] = None,
sort_by_offset: bool = True
) -> Dict[str, torch.Tensor]:
"""Load multiple tensors in parallel using ThreadPoolExecutor.
Args:
keys: List of keys to load (default: all keys)
workers: Number of worker threads (default: auto-calculated based on device)
sort_by_offset: Sort keys by file offset before loading (default: True)
Returns:
Dict mapping key -> tensor
"""
if keys is None:
keys = self.keys()
# Auto-calculate workers based on device capabilities
if workers is None:
try:
from .device_utils import get_optimal_workers, estimate_model_size
@@ -220,13 +154,11 @@ class MemoryEfficientSafeOpen:
workers = min(4, os.cpu_count() or 4)
if sort_by_offset:
# Sort by offset for better I/O patterns
keys_with_offsets = [(k, self.header[k]["data_offsets"][0]) for k in keys]
keys_with_offsets.sort(key=lambda x: x[1])
keys = [k for k, _ in keys_with_offsets]
result = {}
with ThreadPoolExecutor(max_workers=workers) as executor:
futures = {
executor.submit(
@@ -251,7 +183,6 @@ class MemoryEfficientSafeOpen:
return result
def get_tensor_as_dict(self, key: str) -> Dict[str, Any]:
"""Load a tensor and decode as JSON dict (for config tensors)."""
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found")
@@ -261,14 +192,8 @@ class MemoryEfficientSafeOpen:
if offset_start == offset_end:
return {}
if self.mmap_mode and self.mmap_obj:
start = self.header_size + 8 + offset_start
end = self.header_size + 8 + offset_end
tensor_bytes = bytes(self.mmap_obj[start:end])
else:
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - offset_start)
self.file.seek(self.header_size + 8 + offset_start)
tensor_bytes = self.file.read(offset_end - offset_start)
return json.loads(tensor_bytes.decode("utf-8"))
def _read_header(self):
@@ -285,19 +210,13 @@ class MemoryEfficientSafeOpen:
if tensor_bytes is None:
byte_tensor = torch.empty(0, dtype=torch.uint8)
else:
# Avoid extra copy if already a bytearray (low_memory mode)
if isinstance(tensor_bytes, bytearray):
byte_tensor = torch.frombuffer(tensor_bytes, dtype=torch.uint8)
else:
# mmap memoryview needs copy to create writable tensor
byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
byte_tensor = torch.frombuffer(tensor_bytes, dtype=torch.uint8)
if dtype_str in ["F8_E5M2", "F8_E4M3"]:
return self._convert_float8(byte_tensor, dtype_str, shape)
return byte_tensor.view(dtype).reshape(shape)
@staticmethod
def _get_torch_dtype(dtype_str):
dtype_map = {
@@ -322,7 +241,7 @@ class MemoryEfficientSafeOpen:
elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
else:
raise ValueError(f"Unsupported float8 type: {dtype_str}. Your PyTorch version may be too old.")
raise ValueError(f"Unsupported float8 type: {dtype_str}")
def _read_tensor_from_file(
+184
View File
@@ -0,0 +1,184 @@
import os
import folder_paths
from comfy_api.latest import io
from .downloader_utils import get_potential_preview_files, load_image_tensor, get_model_workflows, get_model_metadata_file
class BaseModelInfoLoader(io.ComfyNode):
CATEGORY = "ModelUtils/Info"
@classmethod
def get_info(cls, category, name, workflow_index):
full_path = folder_paths.get_full_path(category, name)
if not full_path or not os.path.exists(full_path):
return (None, "{}", "{}")
base_path = os.path.splitext(full_path)[0]
# 1. Image Preview
image_tensor = None
previews = get_potential_preview_files(base_path)
for p in previews:
if os.path.exists(p) and not p.lower().endswith('.mp4'):
image_tensor = load_image_tensor(p)
if image_tensor is not None:
break
# 2. Workflow
workflow_json = "{}"
workflows = get_model_workflows(base_path)
if workflows and 0 <= workflow_index < len(workflows):
try:
with open(workflows[workflow_index], 'r', encoding='utf-8') as f:
workflow_json = f.read()
except Exception as e:
print(f"Error reading workflow {workflows[workflow_index]}: {e}")
# 3. Metadata
metadata_json = "{}"
metadata_file = get_model_metadata_file(base_path)
if metadata_file:
try:
with open(metadata_file, 'r', encoding='utf-8') as f:
metadata_json = f.read()
except Exception as e:
print(f"Error reading metadata {metadata_file}: {e}")
return (image_tensor, workflow_json, metadata_json)
class CheckpointInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="Checkpoint Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("checkpoint", options=folder_paths.get_filename_list("checkpoints")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, checkpoint, workflow_index):
res = cls.get_info("checkpoints", checkpoint, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class LoRAInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="LoRA Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("lora", options=folder_paths.get_filename_list("loras")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, lora, workflow_index):
res = cls.get_info("loras", lora, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class EmbeddingInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="Embedding Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("embedding", options=folder_paths.get_filename_list("embeddings")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, embedding, workflow_index):
res = cls.get_info("embeddings", embedding, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class VAEInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="VAE Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("vae", options=folder_paths.get_filename_list("vae")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, vae, workflow_index):
res = cls.get_info("vae", vae, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class ControlNetInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="ControlNet Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("controlnet", options=folder_paths.get_filename_list("controlnet")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, controlnet, workflow_index):
res = cls.get_info("controlnet", controlnet, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
class DiffusionModelInfoLoader(BaseModelInfoLoader):
@classmethod
def define_schema(cls):
return io.Schema(
node_id=cls.__name__,
display_name="Diffusion Model Info Loader",
category=cls.CATEGORY,
inputs=[
io.Combo.Input("diffusion_model", options=folder_paths.get_filename_list("diffusion_models")),
io.Int.Input("workflow_index", default=0, min=0, max=100)
],
outputs=[
io.Image.Output(display_name="preview"),
io.String.Output(display_name="workflow_json"),
io.String.Output(display_name="metadata_json")
]
)
@classmethod
def execute(cls, diffusion_model, workflow_index):
res = cls.get_info("diffusion_models", diffusion_model, workflow_index)
return io.NodeOutput(res[0], res[1], res[2])
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-ModelUtils"
description = "[WIP]Custom nodes for handling, inspecting, modifying and creating various model files."
version = "0.2.7"
version = "0.2.9"
license = { file = "LICENSE" }
[project.urls]
+5 -1
View File
@@ -1 +1,5 @@
# No external dependencies required
requests
Pillow
mutagen
unifiedefficientloader
av