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

fix(io): stream UEL and normalize LoRA alpha

Browse Source
This commit is contained in:
silveroxides
2026-08-14 08:43:53 +02:00
parent ed8c331ce2
commit f93ca54c81
23 changed files with 883 additions and 456 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/consensus_merger.py
+193 -95
View File
@@ -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)",
nodes/dora_extract_wd.py
+9 -5
View File
@@ -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")
nodes/dora_learned_wd.py
+9 -5
View File
@@ -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")
nodes/dtype_conversion.py
+3 -28
View File
@@ -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(
nodes/extraction_stream.py
+2 -53
View File
@@ -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]
nodes/lora_alpha.py
+31
View File
@@ -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
nodes/lora_extract_svd.py
+9 -5
View File
@@ -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")
nodes/lora_merger.py
+28 -66
View File
@@ -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(
nodes/lora_rename.py
+12 -14
View File
@@ -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()
nodes/lora_resize.py
+85 -32
View File
@@ -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": [],
nodes/merger.py
+95 -55
View File
@@ -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)")
nodes/merger_ops.py
+7 -4
View File
@@ -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,
nodes/metakeys.py
+10 -22
View File
@@ -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("<Q", f.read(8))[0]
header_json = f.read(header_size).decode("utf-8")
header = json.loads(header_json)
# Extract metadata and layer shapes from header
metadata = header.pop("__metadata__", {})
metadata_str = str(metadata)
layer_shapes = ""
keys = [k for k in header.keys()]
pbar = comfy.utils.ProgressBar(len(keys))
for layer_name in tqdm(keys, desc="Reading layer shapes", unit="layers"):
info = header[layer_name]
shape = str(info.get("shape", []))
name_shape = f"{layer_name}, {shape}\n"
layer_shapes += name_shape
pbar.update(1)
with MemoryEfficientSafeOpen(model_path, low_memory=True) as handler:
metadata_str = str(handler.metadata() or {})
keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(keys))
lines = []
for layer_name in tqdm(keys, desc="Reading layer shapes", unit="layers"):
lines.append(f"{layer_name}, {list(handler.get_shape(layer_name))}\n")
pbar.update(1)
layer_shapes = "".join(lines)
return metadata_str, layer_shapes
nodes/prunekeys.py
+20 -23
View File
@@ -4,11 +4,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
def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
@@ -30,10 +30,6 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_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 {}
patterns = [p.strip() for p in keys_to_prune_str.strip().split('\n') if p.strip()]
if not patterns:
@@ -44,26 +40,27 @@ def _prune_keys(model_name: str, model_type: str, keys_to_prune_str: str,
output_path = os.path.join(model_dir, f"{output_filename.strip()}.safetensors")
# Stream tensors, filter 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:
all_keys = handler.keys()
pbar = comfy.utils.ProgressBar(len(all_keys))
for key in tqdm(all_keys, desc="Pruning keys", unit="keys"):
is_match = False
if use_regex:
if any(re.search(pattern, key) for pattern in patterns):
is_match = True
else:
if any(pattern in key for pattern in patterns):
is_match = True
if not is_match:
writer.write(key, handler.get_tensor(key).contiguous())
kept_keys = [
key for key in all_keys
if not (
any(re.search(pattern, key) for pattern in patterns)
if use_regex else any(pattern in key for pattern in patterns)
)
]
stream = handler.async_stream(
kept_keys, batch_size=1, prefetch_batches=1, pin_memory=False
)
for batch in tqdm(stream, total=len(kept_keys), desc="Pruning keys", unit="keys"):
key, tensor = batch[0]
writer.write_batch([(key, tensor.contiguous())])
handler.mark_processed(key)
pbar.update(1)
finally:
writer.__exit__(None, None, None)
pbar.update(len(all_keys) - len(kept_keys))
# Cleanup after operation
cleanup_after_operation()
nodes/quantization_guard.py
+19 -2
View File
@@ -88,7 +88,7 @@ def layer_has_low_bit(block_keys: dict, low_bit_keys: set[str]) -> 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"]
nodes/renamekeys.py
+12 -14
View File
@@ -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()
nodes/text_encoder_extract.py
+9 -5
View File
@@ -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")
nodes/uel_io.py
+115
View File
@@ -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()
nodes/utils.py
+8 -19
View File
@@ -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 {}, {}
tests/test_consensus_merger.py
+27 -8
View File
@@ -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
)
tests/test_lora_diff_dtype_contract.py
+64
View File
@@ -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)
tests/test_lora_quantization_guard.py
+2 -1
View File
@@ -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)
tests/test_uel_repository_contract.py
+114
View File
@@ -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")) == []