diff --git a/nodes/consensus_merger.py b/nodes/consensus_merger.py index 5251c55..c87fed9 100644 --- a/nodes/consensus_merger.py +++ b/nodes/consensus_merger.py @@ -4,8 +4,7 @@ from __future__ import annotations import logging import os -import uuid -from contextlib import contextmanager +from contextlib import closing from dataclasses import dataclass from typing import Iterable @@ -15,17 +14,15 @@ import torch import torch.nn.functional as F from comfy_api.latest import io from tqdm import tqdm -from unifiedefficientloader import ( - IncrementalSafetensorsWriter, - MemoryEfficientSafeOpen, - transfer_to_gpu_pinned, -) +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned from .device_utils import ( cleanup_after_operation, estimate_model_size, prepare_for_large_operation, ) +from .uel_io import atomic_uel_writer, stream_work_units +from .lora_alpha import normalize_lora_pair from .lora_resize import ( canonical_lora_key, layer_has_companions, @@ -41,7 +38,6 @@ from .merger import ( ) from .quantization_guard import ( inspect_low_bit_input, - layer_has_low_bit, write_preserved_tensor, ) @@ -593,25 +589,6 @@ def _merge_lora_pair_to_cpu( return execute("cpu") -@contextmanager -def _atomic_output_writer(output_path: str, metadata: dict): - """Publish a completed safetensors file atomically and remove failed partials.""" - directory = os.path.dirname(output_path) or "." - basename = os.path.basename(output_path) - temporary_path = os.path.join(directory, f".{basename}.{uuid.uuid4().hex}.tmp") - try: - with IncrementalSafetensorsWriter(temporary_path, metadata=metadata) as writer: - yield writer - os.replace(temporary_path, output_path) - logging.info("[CWB Merge] Published '%s'.", output_path) - except BaseException: - try: - os.remove(temporary_path) - except FileNotFoundError: - pass - raise - - def _clear_previous_layer(params: dict) -> None: if not params["force_clear_cache"]: return @@ -762,11 +739,24 @@ class ConsensusMergerLogic: secondary_only_copied = 0 secondary_only_merged = 0 - writer_context = _atomic_output_writer(output_path, primary.metadata()) - writer = writer_context.__enter__() - try: - with torch.no_grad(): - for key in tqdm(keys, desc="CWB merging tensors", unit="tensors"): + handler_map = dict(enumerate(handlers)) + work_units = [] + for key in keys: + source_indices = [i for i, handler in enumerate(handlers) if key in handler.keys()] + discarded = _matches_any_pattern(key, discard, glob_mode=glob_mode) + guarded = any(key in low_bit_sets[i] for i in source_indices) + entries = {} if discarded or guarded else { + i: [key] for i in source_indices + } + work_units.append((key, entries)) + + with atomic_uel_writer(output_path, primary.metadata()) as writer, torch.no_grad(), closing( + stream_work_units( + handler_map, work_units, + pin_memory=str(params["process_device"]).startswith("cuda"), + ) + ) as streamed: + for key, loaded in tqdm(streamed, total=len(work_units), desc="CWB merging tensors", unit="tensors"): _clear_previous_layer(params) if _matches_any_pattern(key, discard, glob_mode=glob_mode): pbar.update(1) @@ -778,14 +768,23 @@ class ConsensusMergerLogic: guarded = any(key in low_bit_sets[i] for i in source_indices) excluded = _matches_any_pattern(key, exclude, glob_mode=glob_mode) if guarded or excluded: - write_preserved_tensor(writer, key, handlers[preserve_index]) + write_preserved_tensor( + writer, + key, + handlers[preserve_index], + force_raw=guarded, + tensor=loaded.get((preserve_index, key)), + ) if secondary_only: secondary_only_copied += 1 pbar.update(1) continue if secondary_only and len(source_indices) == 1: - write_preserved_tensor(writer, key, handlers[preserve_index]) + write_preserved_tensor( + writer, key, handlers[preserve_index], + tensor=loaded[(preserve_index, key)], + ) secondary_only_copied += 1 pbar.update(1) continue @@ -797,7 +796,10 @@ class ConsensusMergerLogic: key, preserve_index + 1, ) - write_preserved_tensor(writer, key, handlers[preserve_index]) + write_preserved_tensor( + writer, key, handlers[preserve_index], + tensor=loaded[(preserve_index, key)], + ) if secondary_only: secondary_only_copied += 1 pbar.update(1) @@ -807,11 +809,13 @@ class ConsensusMergerLogic: if mismatch_mode == "error": raise ValueError(f"Tensor '{key}' is missing from a CWB input.") if mismatch_mode == "skip": - write_preserved_tensor(writer, key, primary) + write_preserved_tensor( + writer, key, primary, tensor=loaded[(0, key)] + ) pbar.update(1) continue - raw = {i: handlers[i].get_tensor(key) for i in source_indices} + raw = {i: loaded[(i, key)] for i in source_indices} reference_index = 0 if embedding_union: reference_source = max( @@ -864,7 +868,10 @@ class ConsensusMergerLogic: actual_dtypes.append(handler.get_dtype(key)) if preserve_for_mismatch: - write_preserved_tensor(writer, key, handlers[preserve_index]) + write_preserved_tensor( + writer, key, handlers[preserve_index], + tensor=loaded[(preserve_index, key)], + ) if secondary_only: secondary_only_copied += 1 pbar.update(1) @@ -886,16 +893,11 @@ class ConsensusMergerLogic: allow_similarity_alignment=embedding_union, operation_label=key, ) - writer.write(key, merged) + writer.write_batch([(key, merged)]) del merged if secondary_only: secondary_only_merged += 1 pbar.update(1) - except BaseException as exc: - writer_context.__exit__(type(exc), exc, exc.__traceback__) - raise - else: - writer_context.__exit__(None, None, None) if secondary_only_copied or secondary_only_merged: logging.info( "[CWB Merge] Secondary-only tensors: %d copied, %d merged.", @@ -932,11 +934,19 @@ class ConsensusMergerLogic: secondary_only_copied = 0 secondary_only_merged = 0 pbar = comfy.utils.ProgressBar(len(logical_cores) + len(passthrough_keys)) + current_loaded = {} def preserve_keys(writer, keys: Iterable[str], source_index: int = 0): for key in keys: if key not in written: - write_preserved_tensor(writer, key, handlers[source_index]) + guarded = key in low_bit_sets[source_index] + write_preserved_tensor( + writer, + key, + handlers[source_index], + force_raw=guarded, + tensor=current_loaded.get((source_index, key)), + ) written.add(key) def preserve_roles( @@ -952,12 +962,30 @@ class ConsensusMergerLogic: continue output_key = canonical_lora_key(block_name, role) recognized_sources.add(source_key) + if role == "alpha": + written.add(output_key) + continue if output_key not in written: + tensor = current_loaded.get((source_index, source_key)) + if role == "up" and "alpha" in keys: + alpha_key = keys["alpha"] + alpha = current_loaded.get((source_index, alpha_key)) + down = current_loaded.get((source_index, keys.get("down"))) + if tensor is None or alpha is None or down is None: + raise ValueError( + f"Cannot normalize alpha for preserved LoRA layer '{block_name}'." + ) + _, tensor = normalize_lora_pair( + down, tensor, alpha, layer=block_name + ) + current_loaded[(source_index, source_key)] = tensor write_preserved_tensor( writer, source_key, handlers[source_index], output_key, + force_raw=source_key in low_bit_sets[source_index], + tensor=tensor, ) written.add(output_key) return recognized_sources @@ -984,18 +1012,105 @@ class ConsensusMergerLogic: source_index, ) - writer_context = _atomic_output_writer(output_path, handlers[0].metadata()) - writer = writer_context.__enter__() - try: - with torch.no_grad(): - for core in tqdm( - logical_cores, desc="CWB merging LoRA layers", unit="layers" + def core_matches(core): + result = [] + for index, ((pairs, _), logical_map) in enumerate(zip(parsed, logical_maps)): + block = logical_map.get(core) + result.append((index, block, pairs.get(block) if block else None)) + return result + + def has_non_alpha_low_bit(keys, low_bit_keys): + return any( + key in low_bit_keys + for role, key in layer_tensor_keys(keys).items() + if role != "alpha" + ) + + def reject_unnormalizable_alpha(matches): + for index, block, keys in matches: + if ( + keys is not None + and "alpha" in keys + and has_non_alpha_low_bit(keys, low_bit_sets[index]) ): + raise ValueError( + f"Cannot alpha-normalize low-bit LoRA factors for '{block}'." + ) + + work_units = [] + for core in logical_cores: + matches = core_matches(core) + reject_unnormalizable_alpha(matches) + layer_keys = { + index: [ + key for key in handlers[index].keys() + if key in keys.values() or key.startswith(f"{block}.") + ] + for index, block, keys in matches + if block is not None and keys is not None + } + flat_keys = [key for keys in layer_keys.values() for key in keys] + discarded = any( + _matches_any_pattern(key, discard, glob_mode=glob_mode) + for key in flat_keys + ) + guarded = any( + keys is not None and has_non_alpha_low_bit(keys, low_bit_sets[index]) + for index, _, keys in matches + ) + work_units.append((("core", core), {} if discarded or guarded else layer_keys)) + for key in sorted(passthrough_keys): + source_index = next( + index for index, (_, passthrough) in enumerate(parsed) if key in passthrough + ) + entries = {} if key in low_bit_sets[source_index] else {source_index: [key]} + work_units.append((("passthrough", key), entries)) + + output_metadata = (handlers[0].metadata() or {}).copy() + output_metadata["alpha_normalized"] = "true" + output_metadata["alpha_normalization"] = ( + "lora_up := lora_up * (alpha / rank); alpha tensors removed" + ) + with atomic_uel_writer(output_path, output_metadata) as writer, torch.no_grad(), closing( + stream_work_units( + dict(enumerate(handlers)), work_units, + pin_memory=str(params["process_device"]).startswith("cuda"), + ) + ) as streamed: + progress = tqdm( + streamed, + total=len(work_units), + desc="CWB merging LoRA layers", + unit="layers", + ) + for (unit_kind, unit_value), current_loaded in progress: + if unit_kind == "passthrough": + key = unit_value _clear_previous_layer(params) - matches = [] - for index, ((pairs, _), logical_map) in enumerate(zip(parsed, logical_maps)): - block = logical_map.get(core) - matches.append((index, block, pairs.get(block) if block else None)) + if key in written: + continue + if _matches_any_pattern(key, discard, glob_mode=glob_mode): + pbar.update(1) + continue + source_index = next( + index for index, (_, passthrough) in enumerate(parsed) + if key in passthrough + ) + preserve_keys(writer, [key], source_index) + if source_index > 0: + secondary_only_copied += 1 + pbar.update(1) + continue + + core = unit_value + if unit_kind == "core": + _clear_previous_layer(params) + pair_sources = [] + compatible = [] + direct = [] + down = up = alpha = tensor = anchor_tensor = None + anchor_down_template = anchor_up_template = None + matches = core_matches(core) anchor_index, anchor_block, anchor_keys = next( (index, block, keys) for index, block, keys in matches @@ -1015,7 +1130,7 @@ class ConsensusMergerLogic: pbar.update(1) continue guarded = any( - keys is not None and layer_has_low_bit(keys, low_bit_sets[index]) + keys is not None and has_non_alpha_low_bit(keys, low_bit_sets[index]) for index, _, keys in matches ) excluded = any( @@ -1060,8 +1175,8 @@ class ConsensusMergerLogic: if "down" in anchor_keys and "up" in anchor_keys: pair_sources = [] pair_failed = False - anchor_down_template = handlers[anchor_index].get_tensor(anchor_keys["down"]) - anchor_up_template = handlers[anchor_index].get_tensor(anchor_keys["up"]) + anchor_down_template = current_loaded[(anchor_index, anchor_keys["down"])] + anchor_up_template = current_loaded[(anchor_index, anchor_keys["up"])] for index, _, keys in matches: if not keys: if secondary_only: @@ -1091,8 +1206,8 @@ class ConsensusMergerLogic: 1.0, )) continue - down = handlers[index].get_tensor(keys["down"]) - up = handlers[index].get_tensor(keys["up"]) + down = current_loaded[(index, keys["down"])] + up = current_loaded[(index, keys["up"])] if down.ndim < 2 or up.ndim < 2 or down.shape[0] != up.shape[1]: if mismatch_mode == "error" or index == anchor_index: raise ValueError(f"Invalid LoRA rank dimensions for '{anchor_block}'.") @@ -1106,15 +1221,16 @@ class ConsensusMergerLogic: 1.0, )) continue - scale = 1.0 - if "alpha" in keys: - alpha = handlers[index].get_tensor(keys["alpha"]) - if alpha.numel() != 1: - raise ValueError( - f"LoRA alpha for '{anchor_block}' must be scalar." - ) - scale = float(alpha.reshape(-1)[0].item()) / down.shape[0] - pair_sources.append((index, down, up, scale)) + alpha = ( + current_loaded[(index, keys["alpha"])] + if "alpha" in keys else None + ) + down, up = normalize_lora_pair( + down, up, alpha, layer=anchor_block + ) + current_loaded[(index, keys["down"])] = down + current_loaded[(index, keys["up"])] = up + pair_sources.append((index, down, up, 1.0)) if pair_failed: preserve_roles( writer, @@ -1182,10 +1298,6 @@ class ConsensusMergerLogic: canonical_lora_key(anchor_block, "up"), }) if "alpha" in anchor_keys: - writer.write( - canonical_lora_key(anchor_block, "alpha"), - torch.tensor(float(max_rank), dtype=target_dtype), - ) written.add(canonical_lora_key(anchor_block, "alpha")) group_merged = True @@ -1194,7 +1306,7 @@ class ConsensusMergerLogic: continue anchor_key = anchor_keys[kind] output_key = canonical_lora_key(anchor_block, kind) - anchor_tensor = handlers[anchor_index].get_tensor(anchor_key) + anchor_tensor = current_loaded[(anchor_index, anchor_key)] if anchor_tensor.ndim == 1 and not include_1d: if output_key not in written: write_preserved_tensor( @@ -1229,7 +1341,7 @@ class ConsensusMergerLogic: break direct.append(torch.zeros_like(anchor_tensor, dtype=torch.float32)) continue - tensor = handlers[index].get_tensor(keys[kind]) + tensor = current_loaded[(index, keys[kind])] if tensor.shape != anchor_tensor.shape: if mismatch_mode == "error": raise ValueError(f"Direct LoRA shape mismatch for '{anchor_block}.{kind}'.") @@ -1285,28 +1397,14 @@ class ConsensusMergerLogic: secondary_only_merged += 1 else: secondary_only_copied += 1 + pair_sources.clear() + compatible.clear() + direct.clear() + current_loaded.clear() + down = up = alpha = tensor = anchor_tensor = None + anchor_down_template = anchor_up_template = None pbar.update(1) - for key in sorted(passthrough_keys): - _clear_previous_layer(params) - if key in written: - continue - if _matches_any_pattern(key, discard, glob_mode=glob_mode): - pbar.update(1) - continue - source_index = next( - index for index, (_, passthrough) in enumerate(parsed) - if key in passthrough - ) - preserve_keys(writer, [key], source_index) - if source_index > 0: - secondary_only_copied += 1 - pbar.update(1) - except BaseException as exc: - writer_context.__exit__(type(exc), exc, exc.__traceback__) - raise - else: - writer_context.__exit__(None, None, None) if preserved_companion_groups: logging.warning( "[CWB LoRA Merge] Preserved %d companion-bearing group(s)", diff --git a/nodes/dora_extract_wd.py b/nodes/dora_extract_wd.py index 15f3d69..518925d 100644 --- a/nodes/dora_extract_wd.py +++ b/nodes/dora_extract_wd.py @@ -18,7 +18,8 @@ from .device_utils import ( cleanup_after_operation ) -from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen +from .uel_io import atomic_uel_writer from .adaptive_svd import ADAPTIVE_PARTIAL_MODES, adaptive_partial_svd from .extraction_stream import dora_difference, paired_async_tensors, retry_cuda_oom_on_cpu @@ -654,8 +655,8 @@ def extract_dora_from_files( del weight_diff return status, layer_results - writer = IncrementalSafetensorsWriter(output_path) - writer.__enter__() + writer_context = atomic_uel_writer(output_path) + writer = writer_context.__enter__() try: stream = paired_async_tensors( handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda") @@ -674,8 +675,11 @@ def extract_dora_from_files( pbar.update(1) stats["skipped"] += len(weight_keys) - len(work_units) - finally: - writer.__exit__(None, None, None) + except BaseException as exc: + writer_context.__exit__(type(exc), exc, exc.__traceback__) + raise + else: + writer_context.__exit__(None, None, None) print(f"[DoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, " f"{stats['full']} full, {stats['skipped']} skipped") diff --git a/nodes/dora_learned_wd.py b/nodes/dora_learned_wd.py index 83d8b54..82992b2 100644 --- a/nodes/dora_learned_wd.py +++ b/nodes/dora_learned_wd.py @@ -16,7 +16,8 @@ from .device_utils import ( cleanup_after_operation ) -from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen +from .uel_io import atomic_uel_writer from .extraction_stream import dora_difference, paired_async_tensors, retry_cuda_oom_on_cpu # Import the existing SV ranking and chunking logic so we don't have to duplicate it @@ -509,8 +510,8 @@ def extract_dora_learned_from_files( del weight_diff, tensor_base, tensor_ft return status, layer_results - writer = IncrementalSafetensorsWriter(output_path) - writer.__enter__() + writer_context = atomic_uel_writer(output_path) + writer = writer_context.__enter__() try: stream = paired_async_tensors( handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda") @@ -529,8 +530,11 @@ def extract_dora_learned_from_files( pbar.update(1) stats["skipped"] += len(weight_keys) - len(work_units) - finally: - writer.__exit__(None, None, None) + except BaseException as exc: + writer_context.__exit__(type(exc), exc, exc.__traceback__) + raise + else: + writer_context.__exit__(None, None, None) print(f"[Learned DoRA] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, " f"{stats['full']} full, {stats['skipped']} skipped") diff --git a/nodes/dtype_conversion.py b/nodes/dtype_conversion.py index 6d1ff40..a9bd6fc 100644 --- a/nodes/dtype_conversion.py +++ b/nodes/dtype_conversion.py @@ -4,21 +4,17 @@ from __future__ import annotations import os import re -import uuid from collections import Counter -from contextlib import contextmanager import comfy.utils import folder_paths import torch from comfy_api.latest import io from tqdm import tqdm -from unifiedefficientloader import ( - IncrementalSafetensorsWriter, - MemoryEfficientSafeOpen, -) +from unifiedefficientloader import MemoryEfficientSafeOpen from .device_utils import cleanup_after_operation +from .uel_io import atomic_uel_writer TARGET_DTYPES = { @@ -58,27 +54,6 @@ def _output_path(output_filename: str) -> str: return os.path.join(output_dir, f"{name}.safetensors") -@contextmanager -def _atomic_writer(output_path: str, metadata: dict): - directory = os.path.dirname(output_path) or "." - basename = os.path.basename(output_path) - temporary_path = os.path.join( - directory, f".{basename}.{uuid.uuid4().hex}.tmp" - ) - try: - with IncrementalSafetensorsWriter( - temporary_path, metadata=metadata, max_workers=1 - ) as writer: - yield writer - os.replace(temporary_path, output_path) - except BaseException: - try: - os.remove(temporary_path) - except FileNotFoundError: - pass - raise - - def convert_diffusion_model_dtype( model_name: str, target_dtype: str, @@ -106,7 +81,7 @@ def convert_diffusion_model_dtype( try: with MemoryEfficientSafeOpen(source_path, low_memory=True) as loader: metadata = (loader.metadata() or {}).copy() - with _atomic_writer(output_path, metadata) as writer: + with atomic_uel_writer(output_path, metadata) as writer: keys = list(loader.keys()) progress = comfy.utils.ProgressBar(len(keys)) stream = loader.async_stream( diff --git a/nodes/extraction_stream.py b/nodes/extraction_stream.py index 0d4f9e3..4c6eca0 100644 --- a/nodes/extraction_stream.py +++ b/nodes/extraction_stream.py @@ -1,9 +1,9 @@ """Bounded paired UEL streaming helpers for model extraction nodes.""" -from collections import deque - import torch +from .uel_io import AsyncTensorCursor + def is_cuda_oom(error): return isinstance(error, torch.cuda.OutOfMemoryError) or ( isinstance(error, RuntimeError) @@ -109,57 +109,6 @@ def dora_difference(cpu_a, cpu_b, device, *, keep_sources=False): return (*compute(cpu_a.float(), cpu_b_float), "cpu") -class AsyncTensorCursor: - """Consume one ordered, bounded UEL stream and release each consumed key.""" - - def __init__(self, handler, keys, *, pin_memory: bool): - self.handler = handler - self.expected = tuple(keys) - self.position = 0 - self.pending = deque() - self.stream = None - if self.expected: - self.stream = handler.async_stream( - list(self.expected), - batch_size=1, - prefetch_batches=1, - pin_memory=pin_memory, - ) - - def take(self, expected_key): - if self.position >= len(self.expected): - raise RuntimeError(f"Unexpected tensor request: {expected_key}") - planned_key = self.expected[self.position] - if planned_key != expected_key: - raise RuntimeError( - f"UEL stream order mismatch: expected {planned_key}, requested {expected_key}" - ) - if not self.pending: - self.pending.extend(next(self.stream)) - key, tensor = self.pending.popleft() - if key != expected_key: - raise RuntimeError(f"UEL yielded {key} while {expected_key} was expected") - self.position += 1 - return tensor - - def release(self, key): - self.handler.mark_processed(key) - - def finish(self): - if self.position != len(self.expected) or self.pending: - raise RuntimeError("UEL stream ended before all planned tensors were consumed") - if self.stream is not None: - try: - next(self.stream) - except StopIteration: - return - raise RuntimeError("UEL stream yielded unplanned tensors") - - def close(self): - if self.stream is not None: - self.stream.close() - - def paired_async_tensors(handler_a, handler_b, work_units, *, pin_memory: bool): """Yield ordered A/B source tensors and release them after each work unit.""" keys_a = [unit[1] for unit in work_units if unit[1] is not None] diff --git a/nodes/lora_alpha.py b/nodes/lora_alpha.py new file mode 100644 index 0000000..aa170f7 --- /dev/null +++ b/nodes/lora_alpha.py @@ -0,0 +1,31 @@ +import torch + + +def lora_alpha_scale(alpha: torch.Tensor | None, rank: int, *, layer: str) -> float: + """Return the effective LoRA multiplier represented by alpha and rank.""" + + if rank <= 0: + raise ValueError(f"LoRA rank for '{layer}' must be positive.") + if alpha is None: + return 1.0 + if alpha.numel() != 1: + raise ValueError(f"LoRA alpha for '{layer}' must be scalar.") + return float(alpha.reshape(-1)[0].item()) / rank + + +def normalize_lora_pair( + down: torch.Tensor, + up: torch.Tensor, + alpha: torch.Tensor | None, + *, + layer: str, +) -> tuple[torch.Tensor, torch.Tensor]: + """Materialize an alpha-free factor pair before rank alignment or merging.""" + + if down.ndim < 2 or up.ndim < 2: + raise ValueError(f"LoRA factors for '{layer}' must have at least two dimensions.") + rank = int(down.shape[0]) + if int(up.shape[1]) != rank: + raise ValueError(f"LoRA factor rank mismatch for '{layer}'.") + scale = lora_alpha_scale(alpha, rank, layer=layer) + return down, up if scale == 1.0 else up * scale diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index f22842e..ba4409d 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -18,7 +18,8 @@ from .device_utils import ( cleanup_after_operation ) -from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen +from .uel_io import atomic_uel_writer from .quantization_guard import inspect_low_bit_input from .adaptive_svd import ADAPTIVE_PARTIAL_MODES, adaptive_partial_svd from .extraction_stream import lora_difference, paired_async_tensors, retry_cuda_oom_on_cpu @@ -679,8 +680,8 @@ def extract_lora_from_files( del weight_diff return status, layer_results - writer = IncrementalSafetensorsWriter(output_path) - writer.__enter__() + writer_context = atomic_uel_writer(output_path) + writer = writer_context.__enter__() try: stream = paired_async_tensors( handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda") @@ -699,8 +700,11 @@ def extract_lora_from_files( pbar.update(1) stats["skipped"] += len(extraction_keys) - len(work_units) - finally: - writer.__exit__(None, None, None) + except BaseException as exc: + writer_context.__exit__(type(exc), exc, exc.__traceback__) + raise + else: + writer_context.__exit__(None, None, None) print(f"[LoRA Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, " f"{stats['full']} full, {stats['skipped']} skipped") diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index ccbd9b4..ae3503e 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -4,7 +4,7 @@ Resolves different naming conventions and ranks via concatenation. """ import os import gc -from collections import Counter, deque +from collections import Counter import torch import folder_paths import comfy.utils @@ -22,8 +22,10 @@ from .lora_resize import ( validate_canonical_blocks, ) from .quantization_guard import inspect_low_bit_input, layer_has_low_bit, write_preserved_tensor +from .uel_io import AsyncTensorCursor as _AsyncTensorCursor, atomic_uel_writer +from .lora_alpha import normalize_lora_pair -from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned BASE_PREFIXES = ["model.diffusion_model.", "diffusion_model.", "transformer.", "model.", "net."] @@ -217,6 +219,13 @@ def _build_merge_work_units( if guarded or raw: idx, block_name = min(indices, key=lambda item: item[0]) block_keys = lora_infos[idx]["pairs"][block_name] + if raw and "alpha" in block_keys and any( + block_keys.get(role) in lora_infos[idx]["low_bit_keys"] + for role in ("down", "up") + ): + raise ValueError( + f"Cannot alpha-normalize low-bit LoRA factors for '{block_name}'." + ) unit = { "kind": "raw_copy" if raw else "copy", "core": core, @@ -321,58 +330,6 @@ def _build_merge_work_units( return units, stream_keys -class _AsyncTensorCursor: - def __init__(self, handler, keys, pin_memory): - self.handler = handler - self.expected = tuple(keys) - self.position = 0 - self.pending = deque() - self.stream = None - if self.expected: - self.stream = handler.async_stream( - list(self.expected), - batch_size=1, - prefetch_batches=1, - pin_memory=pin_memory, - ) - - def take(self, expected_key): - if self.position >= len(self.expected): - raise RuntimeError(f"Unexpected tensor request: {expected_key}") - planned_key = self.expected[self.position] - if planned_key != expected_key: - raise RuntimeError( - f"UEL stream order mismatch: expected {planned_key}, requested {expected_key}" - ) - if not self.pending: - self.pending.extend(next(self.stream)) - key, tensor = self.pending.popleft() - if key != expected_key: - raise RuntimeError( - f"UEL yielded {key} while {expected_key} was expected" - ) - self.position += 1 - return tensor - - def release(self, key): - self.handler.mark_processed(key) - - def finish(self): - if self.position != len(self.expected) or self.pending: - raise RuntimeError("UEL stream ended before all planned tensors were consumed") - if self.stream is not None: - try: - next(self.stream) - except StopIteration: - pass - else: - raise RuntimeError("UEL stream yielded unplanned tensors") - - def close(self): - if self.stream is not None: - self.stream.close() - - def _is_cuda_oom_error(error): return isinstance(error, torch.cuda.OutOfMemoryError) or ( isinstance(error, RuntimeError) @@ -481,12 +438,6 @@ def _enhanced_dare_ties_merge( return _ties_result(values) -def _source_pair_scale(block_keys, loaded, idx, source_rank): - if "alpha" not in block_keys: - return 1.0 - return float(loaded[(idx, block_keys["alpha"])].item()) / source_rank - - def _process_merge_unit_on_device( unit, loaded, @@ -513,13 +464,22 @@ def _process_merge_unit_on_device( f"LoRA rank mismatch for {block_keys['down']} and {block_keys['up']}" ) maximum_rank = max(maximum_rank, source_rank) - scale = _source_pair_scale( - block_keys, loaded, idx, source_rank + alpha = ( + loaded[(idx, block_keys["alpha"])] + if "alpha" in block_keys else None ) + down_cpu, up_cpu = normalize_lora_pair( + down_cpu, + up_cpu, + alpha, + layer=block_name, + ) + loaded[(idx, block_keys["down"])] = down_cpu + loaded[(idx, block_keys["up"])] = up_cpu pair_contributions.append( ( _to_processing_device(down_cpu, process_device), - _to_processing_device(up_cpu, process_device) * scale, + _to_processing_device(up_cpu, process_device), float(info["weight"]), ) ) @@ -800,15 +760,17 @@ def _run_multi_lora_merge( metadata = { "ss_training_comment": f"Merged {len(lora_paths)} LoRAs via {strategy}", "ss_network_module": "networks.lora", + "alpha_normalized": "true", + "alpha_normalization": ( + "lora_up := lora_up * (alpha / rank); alpha tensors removed" + ), } output_ranks = Counter() fallback_count = 0 tensor_count = 0 pbar = comfy.utils.ProgressBar(len(units)) - with IncrementalSafetensorsWriter( - output_path, metadata=metadata, max_workers=1 - ) as writer: + with atomic_uel_writer(output_path, metadata) as writer: with torch.no_grad(): for work_index, unit in enumerate( tqdm( diff --git a/nodes/lora_rename.py b/nodes/lora_rename.py index 5535142..fb26275 100644 --- a/nodes/lora_rename.py +++ b/nodes/lora_rename.py @@ -3,11 +3,11 @@ import folder_paths import comfy.utils from tqdm import tqdm from comfy_api.latest import io -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 -from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen +from .uel_io import atomic_uel_writer REVERSE_RENAME_DICT = { "transformer_blocks.": "blocks.", @@ -56,21 +56,20 @@ def _convert_diffusers_to_non_diffusers_anima_lora(lora_name: str, output_filena model_size_gb = estimate_model_size(model_path_to_load) prepare_for_large_operation(model_size_gb * 1.2) - # Get metadata from safe_open (header only, no tensor loading) - with safe_open(model_path_to_load, framework="pt", device="cpu") as f: - metadata = f.metadata() or {} - model_dir = folder_paths.get_folder_paths(model_type)[-1] output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors") # Stream tensors, rename on the fly, write immediately - writer = IncrementalSafetensorsWriter(output_path, metadata=metadata) - writer.__enter__() - try: - with MemoryEfficientSafeOpen(model_path_to_load) as handler: + with MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) as handler: + metadata = handler.metadata() or {} + with atomic_uel_writer(output_path, metadata) as writer: original_keys = handler.keys() pbar = comfy.utils.ProgressBar(len(original_keys)) - for key in tqdm(original_keys, desc="Renaming Anima LoRA keys", unit="keys"): + stream = handler.async_stream( + original_keys, batch_size=1, prefetch_batches=1, pin_memory=False + ) + for batch in tqdm(stream, total=len(original_keys), desc="Renaming Anima LoRA keys", unit="keys"): + key, tensor = batch[0] new_key = key if key.startswith("text_conditioner."): @@ -81,10 +80,9 @@ def _convert_diffusers_to_non_diffusers_anima_lora(lora_name: str, output_filena new_key = new_key.replace(old_part, new_part) new_key = f"diffusion_model.{new_key}" - writer.write(new_key, handler.get_tensor(key).contiguous()) + writer.write_batch([(new_key, tensor.contiguous())]) + handler.mark_processed(key) pbar.update(1) - finally: - writer.__exit__(None, None, None) # Cleanup after operation cleanup_after_operation() diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index d39d5be..dbf5a12 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -16,7 +16,9 @@ from tqdm import tqdm from comfy_api.latest import io from .device_utils import estimate_model_size, prepare_for_large_operation, cleanup_after_operation -from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned +from .uel_io import atomic_uel_writer, stream_work_units +from .lora_alpha import normalize_lora_pair from .quantization_guard import inspect_low_bit_input, layer_has_low_bit, write_preserved_tensor from typing import Optional, Dict, Tuple, List @@ -24,14 +26,9 @@ from typing import Optional, Dict, Tuple, List # Reuse rank functions from extraction module from .lora_extract_svd import ( - _index_sv_ratio, _index_sv_cumulative, _index_sv_fro, - _svd_extract_linear, - _svd_extract_conv, _compile_patterns, - _matches_any_pattern, - _format_lora_key, ) @@ -333,16 +330,7 @@ def detect_lora_rank( # Linear: [rank, in_features] or Conv: [rank, in_ch, k, k] network_dim = shape[0] - # Get alpha if present - if ( - network_alpha is None - and "alpha" in block_keys - and block_keys["alpha"] not in (low_bit_keys or set()) - ): - alpha_tensor = handler.get_tensor(block_keys["alpha"]) - network_alpha = float(alpha_tensor.item()) - - if network_dim is not None and network_alpha is not None: + if network_dim is not None: break # Default alpha to dim if not found @@ -710,10 +698,8 @@ def resize_lora_file( total_units = len(raw_units) + len(stream_units) pbar = comfy.utils.ProgressBar(total_units) - writer = IncrementalSafetensorsWriter( - output_path, metadata=metadata, max_workers=1 - ) - writer.__enter__() + writer_context = atomic_uel_writer(output_path, metadata) + writer = writer_context.__enter__() stream = None try: with torch.no_grad(): @@ -873,12 +859,19 @@ def resize_lora_file( pass else: raise RuntimeError("UEL resize stream returned extra tensors") - finally: + except BaseException as exc: if stream is not None: close_stream = getattr(stream, "close", None) if close_stream is not None: close_stream() - writer.__exit__(None, None, None) + writer_context.__exit__(type(exc), exc, exc.__traceback__) + raise + else: + if stream is not None: + close_stream = getattr(stream, "close", None) + if close_stream is not None: + close_stream() + writer_context.__exit__(None, None, None) if verbose and fro_list: avg_fro = sum(fro_list) / len(fro_list) @@ -1329,11 +1322,55 @@ def merge_loras_to_model( if verbose: print(f"[LoRA Merge To Model] Processing {len(base_keys)} base model keys...") - writer = IncrementalSafetensorsWriter(output_path, metadata=base_metadata) - writer.__enter__() + handler_map = {-1: base_handler} + handler_map.update({info["index"]: info["handler"] for info in lora_infos}) + work_units = [] + for planned_base_key in base_keys: + planned_core = planned_base_key + for prefix in BASE_PREFIXES: + if planned_core.startswith(prefix): + planned_core = planned_core[len(prefix):] + break + planned_direct = direct_lookup.get(planned_core, []) + if not planned_direct: + planned_direct = direct_lookup.get(planned_core.replace(".", "_"), []) + planned_low_rank = [] + if planned_base_key.endswith(".weight"): + core = extract_core_layer_base(planned_base_key) + planned_low_rank = lora_lookup.get(core.replace(".", "_"), []) + skipped = any(pattern.search(planned_base_key) for pattern in skip_patterns) + entries = {} + if planned_base_key not in base_low_bit_keys and not skipped: + entries[-1] = [planned_base_key] + for info, _, _, direct_key, block_keys in planned_direct: + if not layer_has_low_bit(block_keys, info["low_bit_keys"]): + entries.setdefault(info["index"], []).append(direct_key) + for info, _, block_keys in planned_low_rank: + if not layer_has_low_bit(block_keys, info["low_bit_keys"]): + entries.setdefault(info["index"], []).extend( + layer_tensor_keys(block_keys).values() + ) + entries = { + source: list(dict.fromkeys(keys)) for source, keys in entries.items() + } + work_units.append((planned_base_key, entries)) + + streamed_units = stream_work_units( + handler_map, + work_units, + pin_memory=str(device).startswith("cuda"), + ) + + writer_context = atomic_uel_writer(output_path, base_metadata) + writer = writer_context.__enter__() try: with torch.no_grad(): - for base_key in tqdm(base_keys, desc="Merging to model", unit="keys"): + for base_key, loaded in tqdm( + streamed_units, + total=len(work_units), + desc="Merging to model", + unit="keys", + ): core_with_suffix = base_key for prefix in BASE_PREFIXES: if core_with_suffix.startswith(prefix): @@ -1402,7 +1439,7 @@ def merge_loras_to_model( continue # Load base weight only after guarded and skipped keys are classified. - cpu_base = base_handler.get_tensor(base_key) + cpu_base = loaded[(-1, base_key)] if direct_contributions or low_rank_contributions: if device == 'cuda': @@ -1423,7 +1460,7 @@ def merge_loras_to_model( + low_bit_description(info, block_keys), ) continue - cpu_patch = info["handler"].get_tensor(direct_key) + cpu_patch = loaded[(info["index"], direct_key)] is_additive = direct_role != "set_weight" if is_additive and cpu_patch.ndim == 1 and not include_1d_diffs: stats["disabled_1d"] += 1 @@ -1493,12 +1530,23 @@ def merge_loras_to_model( info["handler"].get_dtype(key) for key in tensor_keys.values() ] tensors = { - key: info["handler"].get_tensor(key) + key: loaded[(info["index"], key)] for key in tensor_keys.values() } - alpha = None + down_key = block_keys["down"] + up_key = block_keys["up"] + alpha_tensor = tensors.get(block_keys.get("alpha")) + tensors[down_key], tensors[up_key] = normalize_lora_pair( + tensors[down_key], + tensors[up_key], + alpha_tensor, + layer=block_name, + ) + loaded[(info["index"], down_key)] = tensors[down_key] + loaded[(info["index"], up_key)] = tensors[up_key] if "alpha" in block_keys: - alpha = float(tensors[block_keys["alpha"]].item()) + tensors.pop(block_keys["alpha"], None) + alpha = None dora_scale = tensors.get(block_keys.get("dora_scale")) try: adapter = LoRAAdapter.load( @@ -1634,8 +1682,13 @@ def merge_loras_to_model( torch.cuda.empty_cache() pbar.update(1) - finally: - writer.__exit__(None, None, None) + except BaseException as exc: + streamed_units.close() + writer_context.__exit__(type(exc), exc, exc.__traceback__) + raise + else: + streamed_units.close() + writer_context.__exit__(None, None, None) adapter_outcomes = { "APPLIED": [], diff --git a/nodes/merger.py b/nodes/merger.py index aa0992c..346fc62 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -1,6 +1,7 @@ import fnmatch import os import re +from contextlib import closing import torch import folder_paths import comfy.utils @@ -11,9 +12,11 @@ from .device_utils import ( estimate_model_size, prepare_for_large_operation, cleanup_after_operation ) -from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned from .quantization_guard import inspect_low_bit_input, write_preserved_tensor -from .lora_resize import is_direct_diff_key +from .lora_resize import is_direct_diff_key, parse_lora_layers +from .uel_io import atomic_uel_writer, stream_work_units +from .lora_alpha import lora_alpha_scale def load_documentation_from_file(filename): @@ -91,9 +94,6 @@ class MergerLogic: process_device = recipe_params.get('device', 'cpu') print(f"[Merger] Preparing memory for {total_size_gb:.2f}GB merge operation...") prepare_for_large_operation(total_size_gb * 1.2, torch.device(process_device)) - flush_threshold_bytes = int(total_size_gb * 1024**3 / 16) - else: - flush_threshold_bytes = 256 * 1024 * 1024 # 256MB fallback (min 8GB RAM assumed) handlers = {} for name in model_names.values(): @@ -142,7 +142,6 @@ class MergerLogic: if extra: print(f"[Merger] {name} has {len(extra)} extra keys not in Model A (ignored)") - merged_state_dict = {} pbar = comfy.utils.ProgressBar(len(all_keys)) save_dtype = recipe_params.pop('save_dtype') save_torch_dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}.get(save_dtype) @@ -181,47 +180,101 @@ class MergerLogic: output_filename = recipe_params.get("output_filename") output_path = os.path.join(output_dir, f"{output_filename}.safetensors") - writer = None + alpha_normalization = {} + alpha_keys = set() + if model_type == "loras": + for name, handler in handlers.items(): + pairs, _ = parse_lora_layers(handler.keys()) + mapping = {} + for block_keys in pairs.values(): + if all(role in block_keys for role in ("down", "up", "alpha")): + up_key = block_keys["up"] + alpha_key = block_keys["alpha"] + mapping[up_key] = ( + alpha_key, + int(handler.get_shape(block_keys["down"])[0]), + block_keys["down"], + ) + alpha_keys.add(alpha_key) + alpha_normalization[name] = mapping + for up_key, (_, _, down_key) in mapping.items(): + if up_key in low_bit_keys_by_name[name] or down_key in low_bit_keys_by_name[name]: + raise ValueError( + f"Cannot alpha-normalize low-bit LoRA factors for '{up_key}'." + ) - writer = IncrementalSafetensorsWriter(output_path, metadata=metadata) - writer.__enter__() + work_units = [] + for key in all_keys: + preserve_low_bit = model_type == "loras" and any( + key in low_bit_keys for low_bit_keys in low_bit_keys_by_name.values() + ) + discarded = _matches_any_pattern(key, discard_patterns, glob_mode=glob_mode) + entries = {} + if key in alpha_keys: + discarded = True + if not preserve_low_bit and not discarded: + preserve_a = ( + model_type == "loras" + and len(primary_handler.get_shape(key)) == 1 + and not include_1d_diffs + ) or _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode) + requested_names = [primary_model_name] + if not preserve_a: + requested_names.extend( + model_names.get(f"model_{label.lower()}") + for label in calc_mode_class.models_used + if label != "A" + ) + for name in dict.fromkeys(requested_names): + if name in handlers and key in handlers[name].keys(): + entries[name] = [key] + alpha_info = alpha_normalization.get(name, {}).get(key) + if alpha_info is not None: + entries[name].append(alpha_info[0]) + work_units.append((key, entries)) - batch_buffer = {} - batch_bytes = 0 - batch_key_count = 0 - - def _flush_batch(): - nonlocal batch_buffer, batch_bytes, batch_key_count - if not batch_buffer: - return - writer.write_dict(batch_buffer) - pbar.update(batch_key_count) - batch_buffer = {} - batch_bytes = 0 - batch_key_count = 0 - if recipe_params.get('force_clear_cache', True): - import gc - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - with torch.no_grad(): - for key in tqdm(all_keys, desc="Merging layers", unit="layers"): + output_metadata = (metadata or {}).copy() + if model_type == "loras": + output_metadata["alpha_normalized"] = "true" + output_metadata["alpha_normalization"] = ( + "lora_up := lora_up * (alpha / rank); alpha tensors removed" + ) + with atomic_uel_writer(output_path, output_metadata) as writer, torch.no_grad(), closing( + stream_work_units( + handlers, work_units, + pin_memory=str(process_device).startswith("cuda"), + ) + ) as streamed: + for key, loaded in tqdm(streamed, total=len(work_units), desc="Merging layers", unit="layers"): preserve_low_bit = model_type == "loras" and any( key in low_bit_keys for low_bit_keys in low_bit_keys_by_name.values() ) + if key in alpha_keys: + discarded_keys += 1 + pbar.update(1) + continue if preserve_low_bit: - _flush_batch() - write_preserved_tensor(writer, key, primary_handler) + write_preserved_tensor(writer, key, primary_handler, force_raw=True) pbar.update(1) continue # Check discard patterns first - skip entirely if _matches_any_pattern(key, discard_patterns, glob_mode=glob_mode): discarded_keys += 1 + pbar.update(1) continue + for name, mapping in alpha_normalization.items(): + alpha_info = mapping.get(key) + if alpha_info is None or (name, key) not in loaded: + continue + alpha_key, source_rank, _ = alpha_info + alpha = loaded[(name, alpha_key)] + scale = lora_alpha_scale(alpha, source_rank, layer=key) + if scale != 1.0: + loaded[(name, key)] = loaded[(name, key)] * scale + # Pre-load Model A's tensor with pinned memory for CUDA - cpu_tensor = primary_handler.get_tensor(key) + cpu_tensor = loaded[(primary_model_name, key)] # Determine original dtypes across all active models for this key original_dtypes = [primary_handler.get_dtype(key)] @@ -248,29 +301,24 @@ class MergerLogic: if preserve_model_a_1d or _matches_any_pattern(key, exclude_patterns, glob_mode=glob_mode): t = tensor_a.detach().to(target_dtype).cpu() - batch_buffer[key] = t - batch_bytes += t.numel() * t.element_size() - batch_key_count += 1 + writer.write_batch([(key, t)]) + pbar.update(1) if not preserve_model_a_1d: excluded_keys += 1 del tensor_a - if (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32): - _flush_batch() 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 + recipe_params['preloaded_tensors'] = loaded try: recipe = calc_mode_class.create_recipe(key=key, **recipe_params) result = recipe.merge() except MissingTensorError as e: if mismatch_mode == MissingTensorBehavior.ERROR: - - _flush_batch() - writer.__exit__(None, None, None) raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}") result = None error_keys.append(key) @@ -281,11 +329,10 @@ class MergerLogic: skipped_keys += 1 if isinstance(result, dict): + outputs = [] for r_key, r_tensor in result.items(): t = r_tensor.detach().to(target_dtype).cpu() - - batch_buffer[r_key] = t - batch_bytes += t.numel() * t.element_size() + outputs.append((r_key, t)) else: # Ensure compatibility with Model A's architecture. # If alignment_mode is 'pad/crop', we crop results that were padded. @@ -298,27 +345,20 @@ class MergerLogic: t = result.detach().to(target_dtype).cpu() - batch_buffer[key] = t - batch_bytes += t.numel() * t.element_size() + outputs = [(key, t)] - - batch_key_count += 1 + writer.write_batch(outputs) + pbar.update(1) # 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 recipe_params['preloaded_tensors'] del tensor_a del recipe del result - if (batch_bytes >= flush_threshold_bytes or batch_key_count >= 32): - _flush_batch() - - - _flush_batch() - writer.__exit__(None, None, None) - # Log summary if excluded_keys > 0: print(f"[Merger] Excluded {excluded_keys} keys from merge (kept Model A only)") diff --git a/nodes/merger_ops.py b/nodes/merger_ops.py index 1973224..91d8c1f 100644 --- a/nodes/merger_ops.py +++ b/nodes/merger_ops.py @@ -146,12 +146,14 @@ class Operation: class LoadTensor(Operation): def __init__(self, key, model_name, handlers, device, dtype, + preloaded_tensors=None, on_missing=MissingTensorBehavior.ERROR, fallback_shape=None, fallback_dtype=None): super().__init__(key) self.model_name = model_name self.handlers = handlers self.device = device self.dtype = dtype + self.preloaded_tensors = preloaded_tensors or {} self.on_missing = on_missing self.fallback_shape = fallback_shape self.fallback_dtype = fallback_dtype @@ -160,8 +162,8 @@ class LoadTensor(Operation): if self.model_name not in self.handlers: raise ValueError(f"Model '{self.model_name}' is required for this mode but was not provided.") - handler = self.handlers[self.model_name] - if self.key not in handler.keys(): + source = (self.model_name, self.key) + if source not in self.preloaded_tensors: if self.on_missing == MissingTensorBehavior.ERROR: raise MissingTensorError( f"Key '{self.key}' not found in model '{self.model_name}'" @@ -176,7 +178,7 @@ 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) - cpu_tensor = handler.get_tensor(self.key) + cpu_tensor = self.preloaded_tensors[source] if self.device == 'cuda': self._tensor = transfer_to_gpu_pinned(cpu_tensor, self.device, self.dtype) else: @@ -239,7 +241,7 @@ class SVD(Operation): a = a[slices] b = b[slices] - diff, weights, conv2d = a - b, {}, len(a.size()) == 4 + diff, conv2d = a - b, len(a.size()) == 4 kernel_size = None if not conv2d else a.size()[2:4] conv2d_3x3 = conv2d and kernel_size != (1, 1) rank = self.alpha if not conv2d_3x3 or self.beta is None else self.beta @@ -454,6 +456,7 @@ class CalcMode: """Helper to build loader args for secondary models (B, C) with fallback support.""" base_args = { "handlers": kwargs['handlers'], + "preloaded_tensors": kwargs.get('preloaded_tensors', {}), "device": kwargs['device'], "dtype": kwargs['dtype'], "on_missing": mismatch_mode, diff --git a/nodes/metakeys.py b/nodes/metakeys.py index a02c3eb..8c8fbb7 100644 --- a/nodes/metakeys.py +++ b/nodes/metakeys.py @@ -1,11 +1,9 @@ import os -import json -import struct import folder_paths import comfy.utils from tqdm import tqdm from comfy_api.latest import io -from safetensors import safe_open +from unifiedefficientloader import MemoryEfficientSafeOpen from .utils import convert_pt_to_safetensors @@ -21,25 +19,15 @@ def _get_metakeys(model_name: str, model_type: str) -> tuple[str, str]: return f"Conversion failed: {error_message}", "" model_path = temp_safe_path - # Read header only - no tensor loading needed - with open(model_path, "rb") as f: - header_size = struct.unpack(" bool: return any( block_keys.get(name) in low_bit_keys for name in ( - "down", "up", "alpha", "mid", "reshape", "dora_scale", + "down", "up", "mid", "reshape", "dora_scale", "diff", "diff_b", "w_norm", "b_norm", "set_weight", ) ) @@ -101,13 +101,30 @@ def write_preserved_tensor( output_key: str | None = None, *, force_raw: bool = False, + tensor=None, ) -> None: """Copy one tensor's original safetensors bytes and dtype into an active writer.""" destination = output_key or key header = _header(handler) entry = header[key] if not force_raw and entry["dtype"] in UPSTREAM_WRITER_STORAGE_CODES: - writer.write(destination, handler.get_tensor(key).cpu().contiguous()) + if tensor is not None: + writer.write(destination, tensor.cpu().contiguous()) + return + stream = handler.async_stream( + [key], batch_size=1, prefetch_batches=1, pin_memory=False + ) + yielded = False + try: + batch = next(stream) + if len(batch) != 1 or batch[0][0] != key: + raise RuntimeError(f"UEL yielded an unexpected preserved tensor for {key}") + yielded = True + writer.write(destination, batch[0][1].cpu().contiguous()) + finally: + if yielded: + handler.mark_processed(key) + stream.close() return source_start, source_end = entry["data_offsets"] diff --git a/nodes/renamekeys.py b/nodes/renamekeys.py index ece6323..79f4e44 100644 --- a/nodes/renamekeys.py +++ b/nodes/renamekeys.py @@ -3,11 +3,11 @@ import folder_paths import comfy.utils from tqdm import tqdm from comfy_api.latest import io -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 -from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen +from .uel_io import atomic_uel_writer @@ -30,10 +30,6 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str, model_size_gb = estimate_model_size(model_path_to_load) prepare_for_large_operation(model_size_gb * 1.2) - # Get metadata from safe_open (header only, no tensor loading) - with safe_open(model_path_to_load, framework="pt", device="cpu") as f: - metadata = f.metadata() or {} - old_keys = [key.strip() for key in old_keys_str.strip().split('\n') if key.strip()] new_keys = [key.strip() for key in new_keys_str.strip().split('\n') if key.strip()] @@ -47,18 +43,20 @@ def _rename_keys(model_name: str, model_type: str, old_keys_str: str, output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors") # Stream tensors, rename on the fly, write immediately - writer = IncrementalSafetensorsWriter(output_path, metadata=metadata) - writer.__enter__() - try: - with MemoryEfficientSafeOpen(model_path_to_load) as handler: + with MemoryEfficientSafeOpen(model_path_to_load, low_memory=True) as handler: + metadata = handler.metadata() or {} + with atomic_uel_writer(output_path, metadata) as writer: original_keys = handler.keys() pbar = comfy.utils.ProgressBar(len(original_keys)) - for key in tqdm(original_keys, desc="Renaming keys", unit="keys"): + stream = handler.async_stream( + original_keys, batch_size=1, prefetch_batches=1, pin_memory=False + ) + for batch in tqdm(stream, total=len(original_keys), desc="Renaming keys", unit="keys"): + key, tensor = batch[0] new_key = key_map.get(key, key) - writer.write(new_key, handler.get_tensor(key).contiguous()) + writer.write_batch([(new_key, tensor.contiguous())]) + handler.mark_processed(key) pbar.update(1) - finally: - writer.__exit__(None, None, None) # Cleanup after operation cleanup_after_operation() diff --git a/nodes/text_encoder_extract.py b/nodes/text_encoder_extract.py index 99d8def..053c3b3 100644 --- a/nodes/text_encoder_extract.py +++ b/nodes/text_encoder_extract.py @@ -15,7 +15,8 @@ from .device_utils import ( cleanup_after_operation ) -from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter +from unifiedefficientloader import MemoryEfficientSafeOpen +from .uel_io import atomic_uel_writer from .extraction_stream import dora_difference, lora_difference, paired_async_tensors, retry_cuda_oom_on_cpu # Import SVD rank utilities from standard lora extract file to avoid duplication @@ -231,8 +232,8 @@ def extract_te_from_files( del weight_diff return status, layer_results - writer = IncrementalSafetensorsWriter(output_path) - writer.__enter__() + writer_context = atomic_uel_writer(output_path) + writer = writer_context.__enter__() try: stream = paired_async_tensors( handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda") @@ -251,8 +252,11 @@ def extract_te_from_files( pbar.update(1) stats["skipped"] += len(weight_keys) - len(work_units) - finally: - writer.__exit__(None, None, None) + except BaseException as exc: + writer_context.__exit__(type(exc), exc, exc.__traceback__) + raise + else: + writer_context.__exit__(None, None, None) print(f"[TE Extract] Done: {stats['extracted']} extracted, {stats['chunked']} chunked, " f"{stats['full']} full, {stats['skipped']} skipped") diff --git a/nodes/uel_io.py b/nodes/uel_io.py new file mode 100644 index 0000000..52c1e4a --- /dev/null +++ b/nodes/uel_io.py @@ -0,0 +1,115 @@ +import os +import uuid +from collections import deque +from contextlib import contextmanager + +from unifiedefficientloader import IncrementalSafetensorsWriter + + +class AsyncTensorCursor: + """Ordered, bounded tensor access over one UEL asynchronous stream.""" + + def __init__(self, handler, keys, *, pin_memory: bool): + self.handler = handler + self.expected = tuple(keys) + self.position = 0 + self.pending = deque() + self.stream = None + if self.expected: + self.stream = handler.async_stream( + list(self.expected), + batch_size=1, + prefetch_batches=1, + pin_memory=pin_memory, + ) + + def take(self, expected_key): + if self.position >= len(self.expected): + raise RuntimeError(f"Unexpected tensor request: {expected_key}") + planned = self.expected[self.position] + if planned != expected_key: + raise RuntimeError( + f"UEL stream order mismatch: expected {planned}, requested {expected_key}" + ) + if not self.pending: + self.pending.extend(next(self.stream)) + key, tensor = self.pending.popleft() + if key != expected_key: + raise RuntimeError(f"UEL yielded {key} while {expected_key} was expected") + self.position += 1 + return tensor + + def release(self, key): + self.handler.mark_processed(key) + + def finish(self): + if self.position != len(self.expected) or self.pending: + raise RuntimeError("UEL stream ended before all planned tensors were consumed") + if self.stream is not None: + try: + next(self.stream) + except StopIteration: + return + raise RuntimeError("UEL stream yielded unplanned tensors") + + def close(self): + if self.stream is not None: + close = getattr(self.stream, "close", None) + if close is not None: + close() + self.stream = None + self.pending.clear() + + +@contextmanager +def atomic_uel_writer(output_path: str, metadata: dict | None = None): + """Write incrementally and expose the destination only after finalization.""" + + directory = os.path.dirname(output_path) or "." + os.makedirs(directory, exist_ok=True) + basename = os.path.basename(output_path) + temporary = os.path.join(directory, f".{basename}.{uuid.uuid4().hex}.tmp") + try: + with IncrementalSafetensorsWriter( + temporary, + metadata=metadata or {}, + max_workers=1, + ) as writer: + yield writer + os.replace(temporary, output_path) + finally: + if os.path.exists(temporary): + os.remove(temporary) + + +def stream_work_units(handlers, work_units, *, pin_memory: bool): + """Yield bounded logical units from one ordered stream per input handler.""" + + planned = {source: [] for source in handlers} + for _, entries in work_units: + for source, keys in entries.items(): + planned[source].extend(keys) + cursors = { + source: AsyncTensorCursor(handler, planned[source], pin_memory=pin_memory) + for source, handler in handlers.items() + } + try: + for logical_key, entries in work_units: + loaded = {} + consumed = [] + try: + for source, keys in entries.items(): + cursor = cursors[source] + for key in keys: + loaded[(source, key)] = cursor.take(key) + consumed.append((source, key)) + yield logical_key, loaded + finally: + loaded.clear() + for source, key in consumed: + cursors[source].release(key) + for cursor in cursors.values(): + cursor.finish() + finally: + for cursor in cursors.values(): + cursor.close() diff --git a/nodes/utils.py b/nodes/utils.py index 23dd765..44edf99 100644 --- a/nodes/utils.py +++ b/nodes/utils.py @@ -1,7 +1,6 @@ import torch -from safetensors import safe_open -from safetensors.torch import save_file -from comfy.model_management import get_torch_device + +from .uel_io import atomic_uel_writer def convert_pt_to_safetensors(pt_path, safe_path): try: @@ -12,23 +11,13 @@ def convert_pt_to_safetensors(pt_path, safe_path): state_dict = model["model"] else: state_dict = model - state_dict = {k.replace("state_dict.", ""): v for k, v in state_dict.items()} metadata = {"format": "pt"} - save_file(state_dict, safe_path, metadata) + with atomic_uel_writer(safe_path, metadata) as writer: + for key in list(state_dict.keys()): + tensor = state_dict.pop(key) + writer.write_batch( + [(key.replace("state_dict.", ""), tensor.contiguous())] + ) return True, "" except Exception as e: return False, str(e) - -def load_metadata_from_safetensors(safetensors_file: str) -> dict: - try: - with safe_open(safetensors_file, framework="pt", device="cpu") as f: - model_weights = {} - for key in f.keys(): - model_weights[key] = f.get_tensor(key) - metadata = f.metadata() - if metadata is None: - metadata = {} - return model_weights, metadata - except Exception as e: - print(f"Error loading safetensors file: {e}") - return {}, {} diff --git a/tests/test_consensus_merger.py b/tests/test_consensus_merger.py index af50ccc..f191d32 100644 --- a/tests/test_consensus_merger.py +++ b/tests/test_consensus_merger.py @@ -584,7 +584,7 @@ def test_lora_cross_format_companion_group_preserves_model_a_canonically( ) -def test_lora_secondary_only_pair_is_preserved_without_scaling( +def test_lora_secondary_only_pair_is_alpha_normalized( monkeypatch, tmp_path, cwb ): a = tmp_path / "lora_secondary_a.safetensors" @@ -595,7 +595,7 @@ def test_lora_secondary_only_pair_is_preserved_without_scaling( save_file({ "transformer.extra.lora_down.weight": torch.tensor([[2.0, 3.0]]), "transformer.extra.lora_up.weight": torch.tensor([[4.0], [5.0]]), - "transformer.extra.alpha": torch.tensor(1.0), + "transformer.extra.alpha": torch.tensor(0.5), }, str(b)) _patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)}) @@ -617,10 +617,10 @@ def test_lora_secondary_only_pair_is_preserved_without_scaling( ) torch.testing.assert_close( tensors["diffusion_model.extra.lora_B.weight"], - torch.tensor([[4.0], [5.0]]), + torch.tensor([[2.0], [2.5]]), ) assert tensors["diffusion_model.extra.lora_A.weight"].dtype == torch.float32 - assert tensors["diffusion_model.extra.alpha"].item() == pytest.approx(1.0) + assert "diffusion_model.extra.alpha" not in tensors def test_lora_three_input_merges_b_and_c_group_without_a(monkeypatch, tmp_path, cwb): @@ -720,9 +720,8 @@ def test_lora_peft_prefix_matching_and_existing_alpha(monkeypatch, tmp_path, cwb assert set(tensors) == { "diffusion_model.foo.lora_A.weight", "diffusion_model.foo.lora_B.weight", - "diffusion_model.foo.alpha", } - assert tensors["diffusion_model.foo.alpha"].item() == pytest.approx(2.0) + assert "diffusion_model.foo.alpha" not in tensors def test_lora_mochi_inputs_emit_preferred_canonical_output( @@ -747,7 +746,6 @@ def test_lora_mochi_inputs_emit_preferred_canonical_output( assert set(tensors) == { "diffusion_model.foo.lora_A.weight", "diffusion_model.foo.lora_B.weight", - "diffusion_model.foo.alpha", } @@ -841,7 +839,28 @@ def test_isolated_low_bit_preserves_complete_lora_layer(monkeypatch, tmp_path, c torch.testing.assert_close( tensors["diffusion_model.foo.lora_B.weight"], torch.tensor([[3.0], [4.0]]) ) - assert tensors["diffusion_model.foo.alpha"].item() == pytest.approx(1.0) + assert "diffusion_model.foo.alpha" not in tensors torch.testing.assert_close( tensors["diffusion_model.foo.dora_scale"], torch.tensor([5.0, 6.0]) ) + + +def test_alpha_bearing_low_bit_lora_factors_are_rejected(monkeypatch, tmp_path, cwb): + a = tmp_path / "alpha_low_bit_a.safetensors" + b = tmp_path / "alpha_low_bit_b.safetensors" + save_file({ + "diffusion_model.foo.lora_A.weight": torch.tensor([[1.0, 2.0]]), + "diffusion_model.foo.lora_B.weight": torch.tensor([[3.0], [4.0]]), + "diffusion_model.foo.alpha": torch.tensor(1.0), + }, str(a)) + save_file({ + "diffusion_model.foo.lora_A.weight": torch.tensor([[1, 2]], dtype=torch.uint8), + "diffusion_model.foo.lora_B.weight": torch.tensor([[3.0], [4.0]]), + "diffusion_model.foo.alpha": torch.tensor(1.0), + }, str(b)) + _patch_io(monkeypatch, cwb, tmp_path, {"a": str(a), "b": str(b)}) + + with pytest.raises(ValueError, match="Cannot alpha-normalize low-bit LoRA factors"): + cwb.ConsensusMergerLogic.execute( + ["a", "b"], "loras", _params("alpha_low_bit"), lora_mode=True + ) diff --git a/tests/test_lora_diff_dtype_contract.py b/tests/test_lora_diff_dtype_contract.py index c797cd4..a273335 100644 --- a/tests/test_lora_diff_dtype_contract.py +++ b/tests/test_lora_diff_dtype_contract.py @@ -646,6 +646,27 @@ def test_standard_multi_merge_emits_weighted_direct_union(monkeypatch, tmp_path, assert "diffusion_model.lowrank.lora_A.weight" in default_tensors +def test_multi_merge_rejects_alpha_bearing_low_bit_factors( + monkeypatch, tmp_path, modules +): + _, merger = modules + _patch_output(monkeypatch, merger, tmp_path) + source = tmp_path / "alpha_low_bit.safetensors" + save_file({ + "diffusion_model.foo.lora_down.weight": torch.tensor( + [[1, 2]], dtype=torch.uint8 + ), + "diffusion_model.foo.lora_up.weight": torch.tensor([[3.0], [4.0]]), + "diffusion_model.foo.alpha": torch.tensor(1.0), + }, str(source)) + + with pytest.raises(ValueError, match="Cannot alpha-normalize low-bit LoRA factors"): + merger.merge_multi_loras( + [str(source)], [1.0], "concatenate", "cpu", torch.float16, + "alpha_low_bit", verbose=False, + ) + + def test_multi_merge_rejects_direct_shape_mismatch(monkeypatch, tmp_path, modules): _, merger = modules _patch_output(monkeypatch, merger, tmp_path) @@ -896,3 +917,46 @@ def test_generic_two_and_three_mergers_enforce_direct_dtype_contract( three_enabled["diffusion_model.norm.diff"], torch.tensor([2.25, 3.5]) ) assert three_enabled["diffusion_model.norm.diff"].dtype == torch.float32 + + +def test_generic_lora_merge_normalizes_alpha_before_factor_merge( + monkeypatch, tmp_path, generic_modules +): + generic, operations = generic_modules + paths = {} + for name, tensors in { + "a": { + "diffusion_model.layer.lora_A.weight": torch.ones((1, 2)), + "diffusion_model.layer.lora_B.weight": torch.full((2, 1), 2.0), + "diffusion_model.layer.alpha": torch.tensor(0.5), + }, + "b": { + "diffusion_model.layer.lora_A.weight": torch.ones((1, 2)), + "diffusion_model.layer.lora_B.weight": torch.full((2, 1), 3.0), + }, + }.items(): + path = tmp_path / f"alpha_{name}.safetensors" + save_file(tensors, str(path)) + paths[name] = str(path) + + monkeypatch.setattr(generic.folder_paths, "get_full_path", lambda _, name: paths.get(name)) + monkeypatch.setattr(generic.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(generic, "prepare_for_large_operation", lambda *args, **kwargs: None) + monkeypatch.setattr(generic, "cleanup_after_operation", lambda: None) + + params = _generic_params("alpha_normalized", override_dtype=False) + params.update({"model_a": "a", "model_b": "b"}) + generic.MergerLogic.execute_merge( + {"model_a": "a", "model_b": "b"}, + "Weight-Sum", + operations.TWO_MODEL_MODES, + params, + "loras", + ) + + tensors = load_file(str(tmp_path / "loras" / "alpha_normalized.safetensors")) + torch.testing.assert_close( + tensors["diffusion_model.layer.lora_B.weight"], + torch.full((2, 1), 2.0), + ) + assert not any(key.endswith(".alpha") for key in tensors) diff --git a/tests/test_lora_quantization_guard.py b/tests/test_lora_quantization_guard.py index 68f222c..fdbb15b 100644 --- a/tests/test_lora_quantization_guard.py +++ b/tests/test_lora_quantization_guard.py @@ -9,6 +9,7 @@ from pathlib import Path import pytest import torch from safetensors.torch import load_file, save_file +from unifiedefficientloader import IncrementalSafetensorsWriter REPO_ROOT = Path(__file__).resolve().parents[1] @@ -53,7 +54,7 @@ def _inspect_file(path, modules, label="input"): def _round_trip_preserved(path, output, modules): handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True) try: - with modules["resize"].IncrementalSafetensorsWriter(str(output)) as writer: + with IncrementalSafetensorsWriter(str(output)) as writer: modules["guard"].write_preserved_tensor(writer, "low", handler) finally: handler.__exit__(None, None, None) diff --git a/tests/test_uel_repository_contract.py b/tests/test_uel_repository_contract.py new file mode 100644 index 0000000..9da91f2 --- /dev/null +++ b/tests/test_uel_repository_contract.py @@ -0,0 +1,114 @@ +import importlib.util +from pathlib import Path + +import pytest +import torch +from unifiedefficientloader import IncrementalSafetensorsWriter + + +REPO_ROOT = Path(__file__).resolve().parents[1] + + +def _load_uel_io(): + spec = importlib.util.spec_from_file_location( + "modelutils_uel_io_test", REPO_ROOT / "nodes" / "uel_io.py" + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class _Handler: + def __init__(self, values): + self.values = values + self.calls = [] + self.processed = [] + self.closed = 0 + + def async_stream(self, keys, **kwargs): + self.calls.append((list(keys), kwargs)) + + def generate(): + try: + for key in keys: + yield [(key, self.values[key])] + finally: + self.closed += 1 + + return generate() + + def mark_processed(self, key): + self.processed.append(key) + + +def test_production_nodes_have_no_synchronous_safetensors_io(): + violations = [] + for path in sorted((REPO_ROOT / "nodes").glob("*.py")): + source = path.read_text(encoding="utf-8") + if ".get_tensor(" in source: + violations.append(f"{path.name}: get_tensor") + if "from safetensors" in source or "import safetensors" in source: + violations.append(f"{path.name}: direct safetensors import") + assert violations == [] + + +def test_work_units_use_bounded_streams_and_release_every_yield(): + uel = _load_uel_io() + first = _Handler({"a": torch.tensor([1.0]), "b": torch.tensor([2.0])}) + second = _Handler({"x": torch.tensor([3.0])}) + units = [("one", {0: ["a"], 1: ["x"]}), ("two", {0: ["b"]})] + + observed = [] + for logical, loaded in uel.stream_work_units( + {0: first, 1: second}, units, pin_memory=True + ): + observed.append((logical, sorted(loaded))) + + assert observed == [ + ("one", [(0, "a"), (1, "x")]), + ("two", [(0, "b")]), + ] + assert first.processed == ["a", "b"] + assert second.processed == ["x"] + assert first.closed == second.closed == 1 + for handler in (first, second): + assert handler.calls[0][1] == { + "batch_size": 1, + "prefetch_batches": 1, + "pin_memory": True, + } + + +def test_work_unit_failure_releases_consumed_sources_and_closes_streams(): + uel = _load_uel_io() + handler = _Handler({"a": torch.tensor([1.0]), "b": torch.tensor([2.0])}) + stream = uel.stream_work_units( + {0: handler}, [("one", {0: ["a"]}), ("two", {0: ["b"]})], + pin_memory=False, + ) + + with pytest.raises(RuntimeError, match="consumer failed"): + try: + for _, _ in stream: + raise RuntimeError("consumer failed") + finally: + stream.close() + + assert handler.processed == ["a"] + assert handler.closed == 1 + + +def test_atomic_writer_preserves_destination_on_failure(tmp_path): + uel = _load_uel_io() + output = tmp_path / "output.safetensors" + with IncrementalSafetensorsWriter(str(output), max_workers=1) as writer: + writer.write("old", torch.tensor([1.0])) + original = output.read_bytes() + + with pytest.raises(RuntimeError, match="planned failure"): + with uel.atomic_uel_writer(str(output)) as writer: + writer.write("new", torch.tensor([2.0])) + raise RuntimeError("planned failure") + + assert output.read_bytes() == original + assert list(tmp_path.glob(".output.safetensors.*.tmp")) == []