nodes/silveroxides-ComfyUI-ModelUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor to use only unifiedefficientloader and fix header loading

Browse Source
This commit is contained in:
silveroxides
2026-04-27 11:48:53 +02:00
parent 308d839ce3
commit 8fbbe8a5ae
9 changed files with 37 additions and 380 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/downloader_utils.py
+1 -8
View File
@@ -36,12 +36,7 @@ except ImportError:
# 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.")
from unifiedefficientloader import UnifiedSafetensorsLoader
# --- Database Path Logic ---
try:
@@ -416,8 +411,6 @@ def filter_images_by_nsfw(images, threshold_level):
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)
nodes/lora_extract_svd.py
+1 -7
View File
@@ -18,14 +18,8 @@ from .device_utils import (
cleanup_after_operation, get_device_capabilities
)
try:
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
UNIFIED_ENABLED = True
except Exception as e:
UNIFIED_ENABLED = False
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
if not UNIFIED_ENABLED:
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
nodes/lora_merger.py
+1 -7
View File
@@ -20,14 +20,8 @@ from .lora_resize import (
_format_lora_key
)
try:
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
UNIFIED_ENABLED = True
except Exception as e:
UNIFIED_ENABLED = False
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
if not UNIFIED_ENABLED:
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
def merge_multi_loras(
nodes/lora_resize.py
+2 -8
View File
@@ -16,14 +16,8 @@ from safetensors.torch import save_file
from typing import Optional, Dict, Tuple, List
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
try:
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
UNIFIED_ENABLED = True
except Exception as e:
UNIFIED_ENABLED = False
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
if not UNIFIED_ENABLED:
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
# Reuse rank functions from extraction module
@@ -198,7 +192,7 @@ def detect_lora_rank(handler: MemoryEfficientSafeOpen, pairs: Dict) -> Tuple[int
# Get dim from down weight shape
if network_dim is None:
down_key = block_keys["down"]
shape = handler.header[down_key]["shape"]
shape = handler.get_shape(down_key)
# Linear: [rank, in_features] or Conv: [rank, in_ch, k, k]
network_dim = shape[0]
nodes/merger.py
+26 -46
View File
@@ -11,14 +11,8 @@ from .device_utils import (
estimate_model_size, prepare_for_large_operation, cleanup_after_operation
)
try:
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
UNIFIED_ENABLED = True
except Exception as e:
UNIFIED_ENABLED = False
from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
if not UNIFIED_ENABLED:
from .merger_utils import MemoryEfficientSafeOpen, transfer_to_gpu_pinned
def load_documentation_from_file(filename):
"""Loads documentation from a markdown file in the ../docs/ directory."""
@@ -142,9 +136,9 @@ class MergerLogic:
output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
writer = None
if UNIFIED_ENABLED:
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
writer = IncrementalSafetensorsWriter(output_path, metadata=metadata)
writer.__enter__()
batch_buffer = {}
batch_bytes = 0
@@ -170,8 +164,6 @@ class MergerLogic:
# Check discard patterns first - skip entirely
if _matches_any_pattern(key, discard_patterns):
discarded_keys += 1
if not UNIFIED_ENABLED:
pbar.update(1)
continue
# Pre-load Model A's tensor with pinned memory for CUDA
@@ -185,16 +177,13 @@ class MergerLogic:
# Check exclude patterns - use Model A only, no merge
if _matches_any_pattern(key, exclude_patterns):
t = tensor_a.detach().to(save_torch_dtype).cpu()
if UNIFIED_ENABLED:
batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size()
batch_key_count += 1
else:
merged_state_dict[key] = t.clone()
pbar.update(1)
batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size()
batch_key_count += 1
excluded_keys += 1
del tensor_a
if UNIFIED_ENABLED and (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32):
if (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32):
_flush_batch()
continue
@@ -208,9 +197,9 @@ class MergerLogic:
result = recipe.merge()
except MissingTensorError as e:
if mismatch_mode == MissingTensorBehavior.ERROR:
if UNIFIED_ENABLED:
_flush_batch()
writer.__exit__(None, None, None)
_flush_batch()
writer.__exit__(None, None, None)
raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}")
result = None
error_keys.append(key)
@@ -223,11 +212,9 @@ class MergerLogic:
if isinstance(result, dict):
for r_key, r_tensor in result.items():
t = r_tensor.detach().to(save_torch_dtype).cpu()
if UNIFIED_ENABLED:
batch_buffer[r_key] = t
batch_bytes += t.numel() * t.element_size()
else:
merged_state_dict[r_key] = t.clone()
batch_buffer[r_key] = t
batch_bytes += t.numel() * t.element_size()
else:
# Ensure compatibility with Model A's architecture.
# If alignment_mode is 'pad/crop', we crop results that were padded.
@@ -239,16 +226,12 @@ class MergerLogic:
result = result[slices]
t = result.detach().to(save_torch_dtype).cpu()
if UNIFIED_ENABLED:
batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size()
else:
merged_state_dict[key] = t.clone()
batch_buffer[key] = t
batch_bytes += t.numel() * t.element_size()
if UNIFIED_ENABLED:
batch_key_count += 1
else:
pbar.update(1)
batch_key_count += 1
# Clean up references to allow GC immediately.
# Local loop variables must be explicitly deleted to prevent PyTorch from keeping tensors in VRAM.
@@ -258,12 +241,12 @@ class MergerLogic:
del recipe
del result
if UNIFIED_ENABLED and (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32):
if (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32):
_flush_batch()
if UNIFIED_ENABLED:
_flush_batch()
writer.__exit__(None, None, None)
_flush_batch()
writer.__exit__(None, None, None)
# Log summary
if excluded_keys > 0:
@@ -284,12 +267,9 @@ class MergerLogic:
os.makedirs(output_dir, exist_ok=True)
output_filename = recipe_params.get("output_filename")
output_path = os.path.join(output_dir, f"{output_filename}.safetensors")
if not UNIFIED_ENABLED:
save_file(merged_state_dict, output_path, metadata=metadata)
if os.path.exists(output_path):
print(f"[Merger] Output saved to {output_path}")
# Cleanup after heavy operation
if not UNIFIED_ENABLED:
del merged_state_dict
cleanup_after_operation()
return f"{output_filename}.safetensors"
nodes/merger_ops.py
+1 -1
View File
@@ -4,7 +4,7 @@ import torch.nn.functional as F
from enum import Enum
from .merger_utils import transfer_to_gpu_pinned
from unifiedefficientloader import transfer_to_gpu_pinned
class MissingTensorBehavior(Enum):
"""Controls behavior when a tensor key is missing from a model."""
nodes/merger_utils.py
+3 -288
View File
@@ -1,292 +1,7 @@
"""
Memory-efficient safetensors utilities with pinned memory support.
Forces direct file reads to minimize OS page cache usage and avoids VRAM accumulation.
Utilities for model merging operations.
"""
import os
import gc
import json
import struct
import torch
from safetensors import safe_open
from typing import Dict, List, Optional, Any
# 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
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."""
global _pinned_transfer_stats
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)
if not str(device).startswith('cuda'):
if dtype is not None:
return tensor.to(device=device, dtype=dtype)
return tensor.to(device=device)
try:
pinned = tensor.pin_memory()
if dtype is not None:
result = pinned.to(device=device, dtype=dtype, non_blocking=True)
else:
result = pinned.to(device=device, non_blocking=True)
torch.cuda.current_stream().synchronize()
_pinned_transfer_stats["pinned"] += 1
return result
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 streaming and pinned memory support.
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', **kwargs):
# Ignore mmap_mode, low_memory, or any other legacy kwargs. We always stream.
self.filename = filename
self.device = device
self.header, self.header_size = self._read_header()
self.file = open(filename, "rb")
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
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__"]
def metadata(self) -> Dict[str, str]:
"""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]:
keys_with_offsets = []
for key in self.keys():
offset = self.header[key]["data_offsets"][0]
keys_with_offsets.append((key, offset))
keys_with_offsets.sort(key=lambda x: x[1])
return [k for k, _ in keys_with_offsets]
def get_tensor(self, key: str) -> torch.Tensor:
"""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 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:
tensor_bytes = None
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:
cpu_tensor = self.get_tensor(key)
return transfer_to_gpu_pinned(cpu_tensor, device, dtype)
def get_tensors_parallel(
self,
keys: Optional[List[str]] = None,
workers: Optional[int] = None,
sort_by_offset: bool = True
) -> Dict[str, torch.Tensor]:
if keys is None:
keys = self.keys()
if workers is None:
try:
from .device_utils import get_optimal_workers, estimate_model_size
model_size = estimate_model_size(self.filename)
workers = get_optimal_workers(model_size)
except ImportError:
workers = min(4, os.cpu_count() or 4)
if sort_by_offset:
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(
_read_tensor_from_file,
self.filename,
key,
self.header[key],
self.header_size
): key
for key in keys
}
for future in as_completed(futures):
key = futures[future]
try:
tensor = future.result()
if tensor is not None:
result[key] = tensor
except Exception as e:
print(f"[MemoryEfficientSafeOpen] Failed to load {key}: {e}")
return result
def get_tensor_as_dict(self, key: str) -> Dict[str, Any]:
if key not in self.header:
raise KeyError(f"Tensor '{key}' not found")
metadata = self.header[key]
offset_start, offset_end = metadata["data_offsets"]
if offset_start == offset_end:
return {}
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):
with open(self.filename, "rb") as f:
header_size = struct.unpack("<Q", f.read(8))[0]
header_json = f.read(header_size).decode("utf-8")
return json.loads(header_json), header_size
def _deserialize_tensor(self, tensor_bytes, metadata):
dtype_str = metadata["dtype"]
shape = metadata["shape"]
dtype = self._get_torch_dtype(dtype_str)
if tensor_bytes is None:
byte_tensor = torch.empty(0, dtype=torch.uint8)
else:
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 = {
"F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
"I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
"U8": torch.uint8, "BOOL": torch.bool,
}
if hasattr(torch, "float8_e5m2"):
dtype_map["F8_E5M2"] = torch.float8_e5m2
if hasattr(torch, "float8_e4m3fn"):
dtype_map["F8_E4M3"] = torch.float8_e4m3fn
dtype = dtype_map.get(dtype_str)
if dtype is None:
raise ValueError(f"Unsupported dtype: {dtype_str}")
return dtype
@staticmethod
def _convert_float8(byte_tensor, dtype_str, shape):
if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
return byte_tensor.view(torch.float8_e5m2).reshape(shape)
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}")
def _read_tensor_from_file(
filepath: str,
key: str,
metadata: Dict,
header_size: int
) -> Optional[torch.Tensor]:
"""Helper function to read a single tensor from file (for parallel execution).
Each worker opens its own file handle for thread safety.
"""
offset_start, offset_end = metadata["data_offsets"]
dtype_str = metadata["dtype"]
shape = metadata["shape"]
if offset_start == offset_end:
return None
dtype_map = {
"F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
"I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
"U8": torch.uint8, "BOOL": torch.bool,
}
if hasattr(torch, "float8_e5m2"):
dtype_map["F8_E5M2"] = torch.float8_e5m2
if hasattr(torch, "float8_e4m3fn"):
dtype_map["F8_E4M3"] = torch.float8_e4m3fn
dtype = dtype_map.get(dtype_str, torch.float32)
with open(filepath, "rb") as f:
f.seek(header_size + 8 + offset_start)
tensor_bytes = f.read(offset_end - offset_start)
byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
if dtype_str in ["F8_E5M2", "F8_E4M3"]:
if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
return byte_tensor.view(torch.float8_e5m2).reshape(shape)
elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
return byte_tensor.view(dtype).reshape(shape)
from typing import Dict, List
def get_union_keys(*handlers) -> List[str]:
"""Returns the union of all tensor keys across multiple model handlers."""
@@ -305,4 +20,4 @@ def get_intersection_keys(*handlers) -> List[str]:
all_keys = set(handlers[0].keys())
for handler in handlers[1:]:
all_keys &= set(handler.keys())
return sorted(all_keys)
return sorted(all_keys)
nodes/prunekeys.py
+1 -8
View File
@@ -9,14 +9,7 @@ from safetensors import safe_open
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
try:
from unifiedefficientloader import MemoryEfficientSafeOpen
UNIFIED_ENABLED = True
except Exception as e:
UNIFIED_ENABLED = False
if not UNIFIED_ENABLED:
from .merger_utils import MemoryEfficientSafeOpen
from unifiedefficientloader import MemoryEfficientSafeOpen
def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
nodes/renamekeys.py
+1 -7
View File
@@ -8,14 +8,8 @@ from safetensors import safe_open
from .utils import convert_pt_to_safetensors
from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation
try:
from unifiedefficientloader import MemoryEfficientSafeOpen
UNIFIED_ENABLED = True
except Exception as e:
UNIFIED_ENABLED = False
from unifiedefficientloader import MemoryEfficientSafeOpen
if not UNIFIED_ENABLED:
from .merger_utils import MemoryEfficientSafeOpen
def _rename_keys(model_name: str, model_type: str, old_keys_str: str,