From 8fbbe8a5ae3047a495fc6ef9094b3e4cedc300e3 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Mon, 27 Apr 2026 11:48:53 +0200 Subject: [PATCH] refactor to use only unifiedefficientloader and fix header loading --- nodes/downloader_utils.py | 9 +- nodes/lora_extract_svd.py | 8 +- nodes/lora_merger.py | 8 +- nodes/lora_resize.py | 10 +- nodes/merger.py | 72 ++++------ nodes/merger_ops.py | 2 +- nodes/merger_utils.py | 291 +------------------------------------- nodes/prunekeys.py | 9 +- nodes/renamekeys.py | 8 +- 9 files changed, 37 insertions(+), 380 deletions(-) diff --git a/nodes/downloader_utils.py b/nodes/downloader_utils.py index a27b072..ba6500a 100644 --- a/nodes/downloader_utils.py +++ b/nodes/downloader_utils.py @@ -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) diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index b528ba6..440fb44 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -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 diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index 66cd886..2d74c60 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -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( diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index b75e0d6..9e6879c 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -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] diff --git a/nodes/merger.py b/nodes/merger.py index bf88066..fd90a77 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -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" diff --git a/nodes/merger_ops.py b/nodes/merger_ops.py index fe19be3..71504d6 100644 --- a/nodes/merger_ops.py +++ b/nodes/merger_ops.py @@ -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.""" diff --git a/nodes/merger_utils.py b/nodes/merger_utils.py index dbe56a9..4725508 100644 --- a/nodes/merger_utils.py +++ b/nodes/merger_utils.py @@ -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(" 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) \ No newline at end of file + return sorted(all_keys) diff --git a/nodes/prunekeys.py b/nodes/prunekeys.py index e0a704e..db398be 100644 --- a/nodes/prunekeys.py +++ b/nodes/prunekeys.py @@ -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, diff --git a/nodes/renamekeys.py b/nodes/renamekeys.py index b764591..e299a77 100644 --- a/nodes/renamekeys.py +++ b/nodes/renamekeys.py @@ -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,