perf(extract): stream adaptive SVD inputs

nodes/adaptive_svd.py

@@ -0,0 +1,57 @@
+"""Bounded progressive partial SVD for adaptive extraction modes."""
+
+import torch
+
+
+ADAPTIVE_PARTIAL_MODES = {"ratio", "quantile", "sv_cumulative", "sv_fro"}
+
+
+def _bounded_rank(singular_values, mode, parameter, output_cap, matrix_fro_sq):
+    capped = singular_values[:output_cap]
+    if capped.numel() == 0 or capped[0] < 1e-8:
+        return 1, True
+
+    if mode == "ratio":
+        rank = int(torch.sum(capped > capped[0] / parameter).item()) if parameter > 0 else output_cap
+        return max(1, min(rank, output_cap)), rank < output_cap
+
+    if mode in {"quantile", "sv_cumulative"}:
+        total = capped.sum()
+        if total < 1e-8:
+            return 1, True
+        rank = int(torch.searchsorted(torch.cumsum(capped, 0) / total, parameter).item()) + 1
+        return max(1, min(rank, output_cap)), rank < output_cap
+
+    target_sq = float(parameter) ** 2 * matrix_fro_sq
+    cumulative = torch.cumsum(capped.square(), 0)
+    reached = bool(cumulative[-1] >= target_sq)
+    if reached:
+        rank = int(torch.searchsorted(cumulative, target_sq).item()) + 1
+        return max(1, min(rank, output_cap)), rank < output_cap
+    return output_cap, False
+
+
+def adaptive_partial_svd(weight, mode, parameter, max_rank, probe_offset, niter):
+    """Return a bounded top spectrum and selected output rank."""
+    max_possible = min(weight.shape)
+    output_cap = max(1, min(int(max_rank), max_possible))
+    offset = max(1, int(probe_offset))
+    initial_probe = min(max_possible, output_cap + offset)
+    probe_limit = min(max_possible, max(initial_probe, output_cap * 2))
+    matrix_fro_sq = float(torch.sum(weight.square()).item()) if mode == "sv_fro" else 0.0
+
+    probe_rank = initial_probe
+    while True:
+        if probe_rank >= max_possible:
+            U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
+        else:
+            U, S, V = torch.svd_lowrank(weight, q=probe_rank, niter=niter)
+            Vh = V.T
+
+        rank, resolved_before_cap = _bounded_rank(
+            S, mode, parameter, output_cap, matrix_fro_sq
+        )
+        near_probe_tail = rank >= max(1, probe_rank - offset)
+        if resolved_before_cap or not near_probe_tail or probe_rank >= probe_limit:
+            return U[:, :rank], S[:rank], Vh[:rank, :], rank
+        probe_rank = probe_limit

nodes/dora_extract_wd.py

@@ -15,10 +15,12 @@ from tqdm import tqdm
 from comfy_api.latest import io
 from .device_utils import (
     estimate_model_size, prepare_for_large_operation,
-    cleanup_after_operation, get_device_capabilities
+    cleanup_after_operation
 )
 
-from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
+from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
+from .adaptive_svd import ADAPTIVE_PARTIAL_MODES, adaptive_partial_svd
+from .extraction_stream import dora_difference, paired_async_tensors, retry_cuda_oom_on_cpu
 
 
 
@@ -272,12 +274,16 @@ def _svd_extract_linear(
             target_rank = min(target_rank, max_rank)
         return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter)
 
-    if mode in {"sv_knee", "sv_cumulative_knee"} and max_rank is not None:
+    if mode in ADAPTIVE_PARTIAL_MODES and max_rank is not None:
+        U, S, Vh, rank = adaptive_partial_svd(
+            weight, mode, mode_param, max_rank, knee_probe_offset, niter
+        )
+    elif mode in {"sv_knee", "sv_cumulative_knee"} and max_rank is not None:
         U, S, Vh, rank = _svd_extract_knee_lowrank(
             weight, mode, max_rank, knee_probe_offset, niter
         )
     else:
-        # Other adaptive modes require the full spectrum under their current contract.
+        # Uncapped modes require the full spectrum.
         try:
             U, S, Vh = linalg.svd(weight, full_matrices=False)
         except Exception as e:
@@ -532,91 +538,53 @@ def extract_dora_from_files(
     try:
         keys_a = set(handler_a.keys())
         keys_b = set(handler_b.keys())
-        weight_keys = [k for k in keys_a if k.endswith(".weight")]
-        pbar = comfy.utils.ProgressBar(len(weight_keys))
+        weight_keys = sorted(k for k in keys_a if k.endswith(".weight"))
+        work_units = [
+            (key, key, key if key in keys_b else None)
+            for key in weight_keys
+            if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)
+            and not (key not in keys_b and mismatch_mode == "skip")
+        ]
+        pbar = comfy.utils.ProgressBar(len(work_units))
         stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
 
-        def _process_layer(key):
+        def _process_layer(key, cpu_a, cpu_b):
             lora_name = _format_lora_key(key)
+            weight_diff = None
+            layer_device = device
 
             if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
                 return "skipped", None
 
             # Load tensors with pinned memory for CUDA
-            use_pinned = device == 'cuda'
-
             if key not in keys_b:
                 if mismatch_mode == "skip":
                     return "skipped", None
                 if mismatch_mode == "error":
                     raise ValueError(f"Key {key} not found in model B")
 
-                cpu_a = handler_a.get_tensor(key)
-                if use_pinned:
-                    weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
-                else:
-                    weight_diff = cpu_a.to(device=device, dtype=torch.float32)
-                del cpu_a
-
-                # Without base, DoRA fallback just treats tensor_a as diff
-                dora_scale = None
+                weight_diff, dora_scale, layer_device = dora_difference(
+                    cpu_a, None, device
+                )
             else:
-                cpu_a = handler_a.get_tensor(key)
-                cpu_b = handler_b.get_tensor(key)
-
-                if use_pinned:
-                    tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
-                    tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
-                else:
-                    tensor_a = cpu_a.to(device=device, dtype=torch.float32)
-                    tensor_b = cpu_b.to(device=device, dtype=torch.float32)
-                del cpu_a, cpu_b
-
-                if tensor_a.shape != tensor_b.shape:
+                if cpu_a.shape != cpu_b.shape:
                     if mismatch_mode == "skip":
-                        del tensor_a, tensor_b
                         return "skipped", None
                     if mismatch_mode == "error":
-                        raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
+                        raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}")
 
-                    weight_diff = tensor_a
-                    dora_scale = None
-                    del tensor_b
+                    weight_diff, dora_scale, layer_device = dora_difference(
+                        cpu_a, None, device
+                    )
                 else:
-                    # 1-D weights (e.g. x_pad_token, norm .weight) have no
-                    # weight-decomposition direction and are skipped below
-                    # anyway; bail out before the norm, which would index
-                    # dim=1 and raise on a 1-D tensor.
-                    if tensor_a.ndim < 2:
-                        del tensor_a, tensor_b
+                    if cpu_a.ndim < 2:
                         return "skipped", None
+                    weight_diff, dora_scale, layer_device = dora_difference(
+                        cpu_a, cpu_b, device
+                    )
 
-                    # Weight Decomposition calculation
-                    is_conv_dim = tensor_a.ndim == 4
-                    
-                    if is_conv_dim:
-                        out_ch = tensor_a.shape[0]
-                        flat_a = tensor_a.reshape(out_ch, -1)
-                        flat_b = tensor_b.reshape(out_ch, -1)
-                        norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
-                        norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
-                    else:
-                        norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
-                        norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
-                    
-                    # Compute directional diff: W_dir = W_ft * (norm_base / norm_ft)
-                    # Use a small epsilon to avoid division by zero
-                    eps = 1e-8
-                    norm_ratio = norm_b / (norm_a + eps)
-                    if is_conv_dim:
-                        norm_ratio = norm_ratio.view(-1, 1, 1, 1)
-                        dora_scale = norm_a.view(-1, 1, 1, 1)
-                    else:
-                        dora_scale = norm_a
-                        
-                    w_dir = tensor_a * norm_ratio
-                    weight_diff = w_dir - tensor_b
-                    del tensor_a, tensor_b, w_dir
+            if weight_diff is None:
+                raise RuntimeError(f"No extraction tensor was produced for {key}")
 
             # Skip small differences
             if min_diff > 0 and weight_diff.abs().max() < min_diff:
@@ -630,17 +598,20 @@ def extract_dora_from_files(
 
             is_conv = weight_diff.ndim == 4
             layer_results = {}
+            svd_weight = weight_diff
 
             try:
                 if is_conv:
-                    result, mode_str = _svd_extract_conv(
-                        weight_diff, mode, conv_param, device, conv_max_rank,
-                        clamp_quantile, knee_probe_offset
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_conv(svd_weight, mode, conv_param, layer_device, conv_max_rank, clamp_quantile, knee_probe_offset),
+                        lambda: _svd_extract_conv(svd_weight.cpu(), mode, conv_param, "cpu", conv_max_rank, clamp_quantile, knee_probe_offset),
+                        layer_device,
                     )
                 else:
-                    result, mode_str = _svd_extract_linear(
-                        weight_diff, mode, linear_param, device, linear_max_rank,
-                        clamp_quantile, svd_niter, knee_probe_offset
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_linear(svd_weight, mode, linear_param, layer_device, linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset),
+                        lambda: _svd_extract_linear(svd_weight.cpu(), mode, linear_param, "cpu", linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset),
+                        layer_device,
                     )
             except Exception as e:
                 # Try chunked extraction for large tensors
@@ -649,7 +620,7 @@ def extract_dora_from_files(
                     if num_chunks > 1:
                         print(f"[DoRA Extract] Chunked: {key} ({num_chunks} chunks)")
                         lora_up, lora_down, rank = _extract_chunked_layer(
-                            weight_diff, num_chunks, mode, linear_param, device,
+                            weight_diff, num_chunks, mode, linear_param, layer_device,
                             linear_max_rank, knee_probe_offset
                         )
                         if lora_up is not None:
@@ -686,8 +657,11 @@ def extract_dora_from_files(
         writer = IncrementalSafetensorsWriter(output_path)
         writer.__enter__()
         try:
-            for key in tqdm(weight_keys, desc="Extracting DoRA", unit="layers"):
-                status, layer_sd = _process_layer(key)
+            stream = paired_async_tensors(
+                handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda")
+            )
+            for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting DoRA", unit="layers"):
+                status, layer_sd = _process_layer(key, cpu_a, cpu_b)
                 stats[status] += 1
                 if layer_sd:
                     writer.write_dict(layer_sd)
@@ -699,6 +673,7 @@ def extract_dora_from_files(
                         torch.cuda.empty_cache()
 
                 pbar.update(1)
+            stats["skipped"] += len(weight_keys) - len(work_units)
         finally:
             writer.__exit__(None, None, None)
 
@@ -780,7 +755,7 @@ def _format_lora_key(key: str) -> str:
 def _get_common_inputs():
     return [
         io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
-        io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
+        io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer; slower but useful under severe VRAM pressure."),
         io.Boolean.Input("chunk_large_layers", default=False,
                         tooltip="Split large fused layers (QKV, MLP) into chunks"),
         io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
@@ -858,6 +833,8 @@ class DoRAExtractRatio(io.ComfyNode):
                               tooltip="Ratio threshold for linear layers (higher = more SVs kept)"),
                 io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1,
                               tooltip="Ratio threshold for conv layers (higher = more SVs kept)"),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096,
+                             tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
                             tooltip="Maximum rank for linear layers"),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
@@ -869,7 +846,7 @@ class DoRAExtractRatio(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -882,7 +859,7 @@ class DoRAExtractRatio(io.ComfyNode):
             device, save_dtype, output_path, linear_max_rank, conv_max_rank,
             clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
             lazy_load=lazy_load, force_clear_cache=force_clear_cache,
-            glob_skip_patterns=glob_skip_patterns
+            glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset
         )
 
         return io.NodeOutput(output_path)
@@ -904,6 +881,8 @@ class DoRAExtractQuantile(io.ComfyNode):
                               tooltip="Target cumulative % for linear layers"),
                 io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
                               tooltip="Target cumulative % for conv layers"),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096,
+                             tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
                             tooltip="Maximum rank for linear layers"),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
@@ -915,7 +894,7 @@ class DoRAExtractQuantile(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -928,7 +907,7 @@ class DoRAExtractQuantile(io.ComfyNode):
             device, save_dtype, output_path, linear_max_rank, conv_max_rank,
             clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
             lazy_load=lazy_load, force_clear_cache=force_clear_cache,
-            glob_skip_patterns=glob_skip_patterns
+            glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset
         )
 
         return io.NodeOutput(output_path)
@@ -998,6 +977,8 @@ class DoRAExtractFrobenius(io.ComfyNode):
                               tooltip="Target Frobenius norm fraction for linear"),
                 io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01,
                               tooltip="Target Frobenius norm fraction for conv"),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096,
+                             tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
                             tooltip="Maximum rank for linear layers"),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
@@ -1009,7 +990,7 @@ class DoRAExtractFrobenius(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -1022,7 +1003,7 @@ class DoRAExtractFrobenius(io.ComfyNode):
             device, save_dtype, output_path, linear_max_rank, conv_max_rank,
             clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
             lazy_load=lazy_load, force_clear_cache=force_clear_cache,
-            glob_skip_patterns=glob_skip_patterns
+            glob_skip_patterns=glob_skip_patterns, knee_probe_offset=probe_offset
         )
 
         return io.NodeOutput(output_path)

nodes/dora_learned_wd.py

@@ -5,12 +5,8 @@ Extracts Weight-Decomposed Low-Rank Adapters using gradient descent optimization
 to minimize the reconstruction MSE loss, matching the true non-linear DoRA formulation.
 Uses analytical SVD purely as the initialization seed.
 """
-import fnmatch
-import os
-import re
 import math
 import torch
-import torch.linalg as linalg
 import folder_paths
 import comfy.utils
 from tqdm import tqdm
@@ -20,11 +16,12 @@ from .device_utils import (
     cleanup_after_operation
 )
 
-from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
+from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
+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
 from .dora_extract_wd import (
-    _compute_rank, _svd_extract_linear_lowrank, _svd_extract_linear,
+    _svd_extract_linear,
     _svd_extract_conv, _detect_fused_layer, _compile_patterns,
     _matches_any_pattern, _build_lora_output_path, _get_model_inputs,
     _format_lora_key, _get_common_inputs
@@ -378,82 +375,54 @@ def extract_dora_learned_from_files(
     try:
         keys_a = set(handler_a.keys())
         keys_b = set(handler_b.keys())
-        weight_keys = [k for k in keys_a if k.endswith(".weight")]
-        pbar = comfy.utils.ProgressBar(len(weight_keys))
+        weight_keys = sorted(k for k in keys_a if k.endswith(".weight"))
+        work_units = [
+            (key, key, key if key in keys_b else None)
+            for key in weight_keys
+            if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)
+            and not (key not in keys_b and mismatch_mode == "skip")
+        ]
+        pbar = comfy.utils.ProgressBar(len(work_units))
         stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
 
-        def _process_layer(key):
+        def _process_layer(key, cpu_a, cpu_b):
             lora_name = _format_lora_key(key)
+            weight_diff = None
+            tensor_base = None
+            tensor_ft = None
+            layer_device = device
 
             if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
                 return "skipped", None
 
-            use_pinned = device == 'cuda'
-
             if key not in keys_b:
                 if mismatch_mode == "skip":
                     return "skipped", None
                 if mismatch_mode == "error":
                     raise ValueError(f"Key {key} not found in model B")
 
-                # Cannot extract true DoRA without base. Fallback to diff.
-                cpu_a = handler_a.get_tensor(key)
-                tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32) if use_pinned else cpu_a.to(device=device, dtype=torch.float32)
-                del cpu_a
-                weight_diff = tensor_a
-                tensor_base = torch.zeros_like(tensor_a)
-                tensor_ft = tensor_a
-                dora_scale = None
+                weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = (
+                    dora_difference(cpu_a, None, device, keep_sources=True)
+                )
             else:
-                cpu_a = handler_a.get_tensor(key)
-                cpu_b = handler_b.get_tensor(key)
-
-                if use_pinned:
-                    tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
-                    tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
-                else:
-                    tensor_a = cpu_a.to(device=device, dtype=torch.float32)
-                    tensor_b = cpu_b.to(device=device, dtype=torch.float32)
-                del cpu_a, cpu_b
-
-                if tensor_a.shape != tensor_b.shape:
+                if cpu_a.shape != cpu_b.shape:
                     if mismatch_mode == "skip":
-                        del tensor_a, tensor_b
                         return "skipped", None
                     if mismatch_mode == "error":
-                        raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
+                        raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}")
 
-                    weight_diff = tensor_a
-                    tensor_base = torch.zeros_like(tensor_a)
-                    tensor_ft = tensor_a
-                    dora_scale = None
-                    del tensor_b
+                    weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = (
+                        dora_difference(cpu_a, None, device, keep_sources=True)
+                    )
                 else:
-                    tensor_base = tensor_b
-                    tensor_ft = tensor_a
+                    if cpu_a.ndim < 2:
+                        return "skipped", None
+                    weight_diff, dora_scale, tensor_base, tensor_ft, layer_device = (
+                        dora_difference(cpu_a, cpu_b, device, keep_sources=True)
+                    )
 
-                    is_conv_dim = tensor_a.ndim == 4
-                    if is_conv_dim:
-                        out_ch = tensor_a.shape[0]
-                        flat_a = tensor_a.reshape(out_ch, -1)
-                        flat_b = tensor_b.reshape(out_ch, -1)
-                        norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
-                        norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
-                    else:
-                        norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
-                        norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
-
-                    eps = 1e-8
-                    norm_ratio = norm_b / (norm_a + eps)
-                    if is_conv_dim:
-                        norm_ratio = norm_ratio.view(-1, 1, 1, 1)
-                        dora_scale = norm_a.view(-1, 1, 1, 1)
-                    else:
-                        dora_scale = norm_a
-
-                    w_dir = tensor_a * norm_ratio
-                    weight_diff = w_dir - tensor_b
-                    del w_dir
+            if weight_diff is None or tensor_base is None or tensor_ft is None:
+                raise RuntimeError(f"No extraction tensors were produced for {key}")
 
             if min_diff > 0 and weight_diff.abs().max() < min_diff:
                 del weight_diff, tensor_base, tensor_ft
@@ -465,28 +434,37 @@ def extract_dora_learned_from_files(
 
             is_conv = weight_diff.ndim == 4
             layer_results = {}
+            svd_weight = weight_diff
+            optimize_base = tensor_base
+            optimize_ft = tensor_ft
 
             try:
                 # 1. Run Analytical SVD (Initialization)
                 if is_conv:
-                    result, mode_str = _svd_extract_conv(
-                        weight_diff, mode, conv_param, device, conv_max_rank,
-                        clamp_quantile, knee_probe_offset
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_conv(svd_weight, mode, conv_param, layer_device, conv_max_rank, clamp_quantile, knee_probe_offset),
+                        lambda: _svd_extract_conv(svd_weight.cpu(), mode, conv_param, "cpu", conv_max_rank, clamp_quantile, knee_probe_offset),
+                        layer_device,
                     )
                 else:
-                    result, mode_str = _svd_extract_linear(
-                        weight_diff, mode, linear_param, device, linear_max_rank,
-                        clamp_quantile, svd_niter, knee_probe_offset
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_linear(svd_weight, mode, linear_param, layer_device, linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset),
+                        lambda: _svd_extract_linear(svd_weight.cpu(), mode, linear_param, "cpu", linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset),
+                        layer_device,
                     )
 
                 # 2. Run Optimization (Learned Rounding)
                 if mode_str != "full" and dora_scale is not None:
                     lora_down, lora_up, _ = result
-                    lora_down, lora_up = _optimize_dora(
-                        tensor_base, tensor_ft, lora_down, lora_up, optimize_iters, learning_rate,
-                        optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
-                        early_stop_loss, early_stop_stall, early_stop_lr,
-                        layer_name=lora_name
+                    optimize_args = (
+                        optimize_iters, learning_rate, optimizer_type, lr_schedule,
+                        lr_patience, lr_factor, lr_cooldown, early_stop_loss,
+                        early_stop_stall, early_stop_lr,
+                    )
+                    lora_down, lora_up = retry_cuda_oom_on_cpu(
+                        lambda: _optimize_dora(optimize_base, optimize_ft, lora_down, lora_up, *optimize_args, layer_name=lora_name),
+                        lambda: _optimize_dora(optimize_base.cpu(), optimize_ft.cpu(), lora_down.cpu(), lora_up.cpu(), *optimize_args, layer_name=lora_name),
+                        layer_device,
                     )
                     result = (lora_down, lora_up, _)
 
@@ -496,7 +474,7 @@ def extract_dora_learned_from_files(
                     if num_chunks > 1:
                         print(f"[Learned DoRA] Chunked: {key} ({num_chunks} chunks)")
                         lora_up, lora_down, rank = _extract_chunked_learned_layer(
-                            tensor_base, tensor_ft, num_chunks, mode, linear_param, device, linear_max_rank, optimize_iters, learning_rate,
+                            tensor_base, tensor_ft, num_chunks, mode, linear_param, layer_device, linear_max_rank, optimize_iters, learning_rate,
                             optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown,
                             early_stop_loss, early_stop_stall, early_stop_lr,
                             layer_name=lora_name,
@@ -534,8 +512,11 @@ def extract_dora_learned_from_files(
         writer = IncrementalSafetensorsWriter(output_path)
         writer.__enter__()
         try:
-            for key in tqdm(weight_keys, desc="Extracting Learned DoRA", unit="layers"):
-                status, layer_sd = _process_layer(key)
+            stream = paired_async_tensors(
+                handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda")
+            )
+            for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting Learned DoRA", unit="layers"):
+                status, layer_sd = _process_layer(key, cpu_a, cpu_b)
                 stats[status] += 1
                 if layer_sd:
                     writer.write_dict(layer_sd)
@@ -547,6 +528,7 @@ def extract_dora_learned_from_files(
                         torch.cuda.empty_cache()
 
                 pbar.update(1)
+            stats["skipped"] += len(weight_keys) - len(work_units)
         finally:
             writer.__exit__(None, None, None)
 
@@ -642,6 +624,7 @@ class DoRALearnedExtractRatio(io.ComfyNode):
                 *_get_model_inputs(),
                 io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for linear layers; higher values retain more singular values."),
                 io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
                 *_get_learned_inputs(),
@@ -652,7 +635,7 @@ class DoRALearnedExtractRatio(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank,
                 optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
                 output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -664,7 +647,8 @@ class DoRALearnedExtractRatio(io.ComfyNode):
             model_a_path, model_b_path, "ratio", linear_ratio, conv_ratio,
             device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
             clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
-            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
+            knee_probe_offset=probe_offset
         )
 
         return io.NodeOutput(output_path)
@@ -685,6 +669,7 @@ class DoRALearnedExtractQuantile(io.ComfyNode):
                 *_get_model_inputs(),
                 io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."),
                 io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
                 *_get_learned_inputs(),
@@ -695,7 +680,7 @@ class DoRALearnedExtractQuantile(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank,
                 optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
                 output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -707,7 +692,8 @@ class DoRALearnedExtractQuantile(io.ComfyNode):
             model_a_path, model_b_path, "quantile", linear_quantile, conv_quantile,
             device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
             clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
-            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
+            knee_probe_offset=probe_offset
         )
 
         return io.NodeOutput(output_path)
@@ -773,6 +759,7 @@ class DoRALearnedExtractFrobenius(io.ComfyNode):
                 *_get_model_inputs(),
                 io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."),
                 io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
                 *_get_learned_inputs(),
@@ -783,7 +770,7 @@ class DoRALearnedExtractFrobenius(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank,
                 optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode,
                 output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -795,7 +782,8 @@ class DoRALearnedExtractFrobenius(io.ComfyNode):
             model_a_path, model_b_path, "sv_fro", linear_target, conv_target,
             device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank,
             clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
-            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
+            knee_probe_offset=probe_offset
         )
 
         return io.NodeOutput(output_path)

nodes/extraction_stream.py

@@ -0,0 +1,185 @@
+"""Bounded paired UEL streaming helpers for model extraction nodes."""
+
+from collections import deque
+
+import torch
+
+def is_cuda_oom(error):
+    return isinstance(error, torch.cuda.OutOfMemoryError) or (
+        isinstance(error, RuntimeError)
+        and "cuda" in str(error).lower()
+        and "out of memory" in str(error).lower()
+    )
+
+
+def release_cuda_oom():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def move_source_tensors(cpu_a, cpu_b, device, dtype=torch.float32):
+    """Move one source pair, retrying a confirmed CUDA OOM on CPU."""
+    use_cuda = str(device).startswith("cuda")
+    try:
+        if use_cuda:
+            tensor_a = cpu_a.to(device=device, dtype=dtype, non_blocking=True)
+            tensor_b = cpu_b.to(device=device, dtype=dtype, non_blocking=True) if cpu_b is not None else None
+        else:
+            tensor_a = cpu_a.to(device=device, dtype=dtype)
+            tensor_b = cpu_b.to(device=device, dtype=dtype) if cpu_b is not None else None
+        return tensor_a, tensor_b, device
+    except BaseException as error:
+        if not use_cuda or not is_cuda_oom(error):
+            raise
+        release_cuda_oom()
+        return (
+            cpu_a.to(device="cpu", dtype=dtype),
+            cpu_b.to(device="cpu", dtype=dtype) if cpu_b is not None else None,
+            "cpu",
+        )
+
+
+def retry_cuda_oom_on_cpu(cuda_call, cpu_call, device):
+    """Retry one compute operation on CPU only for a confirmed CUDA OOM."""
+    try:
+        return cuda_call()
+    except BaseException as error:
+        if not str(device).startswith("cuda") or not is_cuda_oom(error):
+            raise
+        release_cuda_oom()
+        return cpu_call()
+
+
+def lora_difference(cpu_a, cpu_b, device):
+    """Construct A minus B with transfer and arithmetic CUDA OOM fallback."""
+    tensor_a, tensor_b, actual_device = move_source_tensors(
+        cpu_a, cpu_b, device
+    )
+    if tensor_b is None:
+        return tensor_a, actual_device
+    try:
+        return tensor_a - tensor_b, actual_device
+    except BaseException as error:
+        if not str(actual_device).startswith("cuda") or not is_cuda_oom(error):
+            raise
+        del tensor_a, tensor_b
+        release_cuda_oom()
+        return cpu_a.float() - cpu_b.float(), "cpu"
+
+
+def dora_difference(cpu_a, cpu_b, device, *, keep_sources=False):
+    """Construct a directional DoRA delta with complete CUDA OOM fallback."""
+    tensor_a, tensor_b, actual_device = move_source_tensors(
+        cpu_a, cpu_b, device
+    )
+
+    def compute(a, b):
+        if b is None:
+            base = torch.zeros_like(a)
+            result = (a, None, base, a)
+            return result if keep_sources else result[:2]
+        if a.ndim < 2:
+            result = (None, None, b, a)
+            return result if keep_sources else result[:2]
+        is_conv = a.ndim == 4
+        if is_conv:
+            out_channels = a.shape[0]
+            norm_a = torch.linalg.norm(a.reshape(out_channels, -1), dim=1, keepdim=True)
+            norm_b = torch.linalg.norm(b.reshape(out_channels, -1), dim=1, keepdim=True)
+        else:
+            norm_a = torch.linalg.norm(a, dim=1, keepdim=True)
+            norm_b = torch.linalg.norm(b, dim=1, keepdim=True)
+        ratio = norm_b / (norm_a + 1e-8)
+        if is_conv:
+            ratio = ratio.view(-1, 1, 1, 1)
+            scale = norm_a.view(-1, 1, 1, 1)
+        else:
+            scale = norm_a
+        result = (a * ratio - b, scale, b, a)
+        return result if keep_sources else result[:2]
+
+    try:
+        return (*compute(tensor_a, tensor_b), actual_device)
+    except BaseException as error:
+        if not str(actual_device).startswith("cuda") or not is_cuda_oom(error):
+            raise
+        del tensor_a, tensor_b
+        release_cuda_oom()
+        cpu_b_float = cpu_b.float() if cpu_b is not None else None
+        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]
+    keys_b = [unit[2] for unit in work_units if unit[2] is not None]
+    cursor_a = AsyncTensorCursor(handler_a, keys_a, pin_memory=pin_memory)
+    cursor_b = AsyncTensorCursor(handler_b, keys_b, pin_memory=pin_memory)
+    try:
+        for logical_key, key_a, key_b in work_units:
+            tensor_a = cursor_a.take(key_a) if key_a is not None else None
+            tensor_b = cursor_b.take(key_b) if key_b is not None else None
+            try:
+                yield logical_key, tensor_a, tensor_b
+            finally:
+                del tensor_a, tensor_b
+                if key_a is not None:
+                    cursor_a.release(key_a)
+                if key_b is not None:
+                    cursor_b.release(key_b)
+        cursor_a.finish()
+        cursor_b.finish()
+    finally:
+        cursor_a.close()
+        cursor_b.close()

nodes/lora_extract_svd.py

@@ -15,11 +15,13 @@ from tqdm import tqdm
 from comfy_api.latest import io
 from .device_utils import (
     estimate_model_size, prepare_for_large_operation,
-    cleanup_after_operation, get_device_capabilities
+    cleanup_after_operation
 )
 
-from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
+from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
 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
 
 
 
@@ -273,12 +275,16 @@ def _svd_extract_linear(
             target_rank = min(target_rank, max_rank)
         return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter)
 
-    if mode in {"sv_knee", "sv_cumulative_knee"} and max_rank is not None:
+    if mode in ADAPTIVE_PARTIAL_MODES and max_rank is not None:
+        U, S, Vh, rank = adaptive_partial_svd(
+            weight, mode, mode_param, max_rank, knee_probe_offset, niter
+        )
+    elif mode in {"sv_knee", "sv_cumulative_knee"} and max_rank is not None:
         U, S, Vh, rank = _svd_extract_knee_lowrank(
             weight, mode, max_rank, knee_probe_offset, niter
         )
     else:
-        # Other adaptive modes require the full spectrum under their current contract.
+        # Uncapped modes require the full spectrum.
         try:
             U, S, Vh = linalg.svd(weight, full_matrices=False)
         except Exception as e:
@@ -549,11 +555,20 @@ def extract_lora_from_files(
                 and (len(handler_a.get_shape(k)) == 1 or k.endswith(direct_diff_suffixes))
             ))
         ]
-        pbar = comfy.utils.ProgressBar(len(extraction_keys))
+        extraction_keys.sort()
+        work_units = [
+            (key, key, key if key in keys_b else None)
+            for key in extraction_keys
+            if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)
+            and not (key not in keys_b and mismatch_mode == "skip")
+        ]
+        pbar = comfy.utils.ProgressBar(len(work_units))
         stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
 
-        def _process_layer(key):
+        def _process_layer(key, cpu_a, cpu_b):
             lora_name = _format_lora_key(key)
+            weight_diff = None
+            layer_device = device
 
             if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
                 return "skipped", None
@@ -566,50 +581,32 @@ def extract_lora_from_files(
             )
 
             # Load tensors with pinned memory for CUDA
-            use_pinned = device == 'cuda'
-
             if key not in keys_b:
                 if mismatch_mode == "skip":
                     return "skipped", None
                 if mismatch_mode == "error":
                     raise ValueError(f"Key {key} not found in model B")
 
-                cpu_a = handler_a.get_tensor(key)
-                if use_pinned:
-                    weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
-                else:
-                    weight_diff = cpu_a.to(device=device, dtype=torch.float32)
-                del cpu_a
+                weight_diff, layer_device = lora_difference(cpu_a, None, device)
 
                 if mismatch_mode == "zeros":
                     # For zeros mode, we treat the missing base as a zeroed tensor of same shape
                     # which is already captured by weight_diff = tensor_a
                     pass
             else:
-                cpu_a = handler_a.get_tensor(key)
-                cpu_b = handler_b.get_tensor(key)
-
-                if use_pinned:
-                    tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
-                    tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
-                else:
-                    tensor_a = cpu_a.to(device=device, dtype=torch.float32)
-                    tensor_b = cpu_b.to(device=device, dtype=torch.float32)
-                del cpu_a, cpu_b
-
-                if tensor_a.shape != tensor_b.shape:
+                if cpu_a.shape != cpu_b.shape:
                     if mismatch_mode == "skip":
-                        del tensor_a, tensor_b
                         return "skipped", None
                     if mismatch_mode == "error":
-                        raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
+                        raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}")
 
                     # For zeros/fallback, use tensor_a as the difference
-                    weight_diff = tensor_a
-                    del tensor_b
+                    weight_diff, layer_device = lora_difference(cpu_a, None, device)
                 else:
-                    weight_diff = tensor_a - tensor_b
-                    del tensor_a, tensor_b
+                    weight_diff, layer_device = lora_difference(cpu_a, cpu_b, device)
+
+            if weight_diff is None:
+                raise RuntimeError(f"No extraction tensor was produced for {key}")
 
             # Skip small differences
             if min_diff > 0 and weight_diff.abs().max() < min_diff:
@@ -634,17 +631,20 @@ def extract_lora_from_files(
 
             is_conv = weight_diff.ndim == 4
             layer_results = {}
+            svd_weight = weight_diff
 
             try:
                 if is_conv:
-                    result, mode_str = _svd_extract_conv(
-                        weight_diff, mode, conv_param, device, conv_max_rank,
-                        clamp_quantile, knee_probe_offset
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_conv(svd_weight, mode, conv_param, layer_device, conv_max_rank, clamp_quantile, knee_probe_offset),
+                        lambda: _svd_extract_conv(svd_weight.cpu(), mode, conv_param, "cpu", conv_max_rank, clamp_quantile, knee_probe_offset),
+                        layer_device,
                     )
                 else:
-                    result, mode_str = _svd_extract_linear(
-                        weight_diff, mode, linear_param, device, linear_max_rank,
-                        clamp_quantile, svd_niter, knee_probe_offset
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_linear(svd_weight, mode, linear_param, layer_device, linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset),
+                        lambda: _svd_extract_linear(svd_weight.cpu(), mode, linear_param, "cpu", linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset),
+                        layer_device,
                     )
             except Exception as e:
                 # Try chunked extraction for large tensors
@@ -653,7 +653,7 @@ def extract_lora_from_files(
                     if num_chunks > 1:
                         print(f"[LoRA Extract] Chunked: {key} ({num_chunks} chunks)")
                         lora_up, lora_down, rank = _extract_chunked_layer(
-                            weight_diff, num_chunks, mode, linear_param, device,
+                            weight_diff, num_chunks, mode, linear_param, layer_device,
                             linear_max_rank, knee_probe_offset
                         )
                         if lora_up is not None:
@@ -682,8 +682,11 @@ def extract_lora_from_files(
         writer = IncrementalSafetensorsWriter(output_path)
         writer.__enter__()
         try:
-            for key in tqdm(extraction_keys, desc="Extracting LoRA", unit="layers"):
-                status, layer_sd = _process_layer(key)
+            stream = paired_async_tensors(
+                handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda")
+            )
+            for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting LoRA", unit="layers"):
+                status, layer_sd = _process_layer(key, cpu_a, cpu_b)
                 stats[status] += 1
                 if layer_sd:
                     writer.write_dict(layer_sd)
@@ -695,6 +698,7 @@ def extract_lora_from_files(
                         torch.cuda.empty_cache()
 
                 pbar.update(1)
+            stats["skipped"] += len(extraction_keys) - len(work_units)
         finally:
             writer.__exit__(None, None, None)
 
@@ -783,7 +787,7 @@ def _format_direct_diff_key(source_key: str, lora_name: str) -> str:
 def _get_common_inputs():
     return [
         io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
-        io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
+        io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer; slower but useful under severe VRAM pressure."),
         io.Boolean.Input("chunk_large_layers", default=False,
                         tooltip="Split large fused layers (QKV, MLP) into chunks"),
         io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
@@ -865,6 +869,8 @@ class LoRAExtractRatio(io.ComfyNode):
                               tooltip="Ratio threshold for linear layers (higher = more SVs kept)"),
                 io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1,
                               tooltip="Ratio threshold for conv layers (higher = more SVs kept)"),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096,
+                             tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
                             tooltip="Maximum rank for linear layers"),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
@@ -876,7 +882,7 @@ class LoRAExtractRatio(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache,
                 include_1d_diffs=False) -> io.NodeOutput:
@@ -890,7 +896,8 @@ class LoRAExtractRatio(io.ComfyNode):
             device, save_dtype, output_path, linear_max_rank, conv_max_rank,
             clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
             lazy_load=lazy_load, force_clear_cache=force_clear_cache,
-            glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
+            glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs,
+            knee_probe_offset=probe_offset,
         )
 
         return io.NodeOutput(output_path)
@@ -912,6 +919,8 @@ class LoRAExtractQuantile(io.ComfyNode):
                               tooltip="Target cumulative % for linear layers"),
                 io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01,
                               tooltip="Target cumulative % for conv layers"),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096,
+                             tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
                             tooltip="Maximum rank for linear layers"),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
@@ -923,7 +932,7 @@ class LoRAExtractQuantile(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache,
                 include_1d_diffs=False) -> io.NodeOutput:
@@ -937,7 +946,8 @@ class LoRAExtractQuantile(io.ComfyNode):
             device, save_dtype, output_path, linear_max_rank, conv_max_rank,
             clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
             lazy_load=lazy_load, force_clear_cache=force_clear_cache,
-            glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
+            glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs,
+            knee_probe_offset=probe_offset,
         )
 
         return io.NodeOutput(output_path)
@@ -1008,6 +1018,8 @@ class LoRAExtractFrobenius(io.ComfyNode):
                               tooltip="Target Frobenius norm fraction for linear"),
                 io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01,
                               tooltip="Target Frobenius norm fraction for conv"),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096,
+                             tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384,
                             tooltip="Maximum rank for linear layers"),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384,
@@ -1019,7 +1031,7 @@ class LoRAExtractFrobenius(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache,
                 include_1d_diffs=False) -> io.NodeOutput:
@@ -1033,7 +1045,8 @@ class LoRAExtractFrobenius(io.ComfyNode):
             device, save_dtype, output_path, linear_max_rank, conv_max_rank,
             clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers,
             lazy_load=lazy_load, force_clear_cache=force_clear_cache,
-            glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs
+            glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs,
+            knee_probe_offset=probe_offset,
         )
 
         return io.NodeOutput(output_path)

nodes/text_encoder_extract.py

@@ -4,11 +4,8 @@ Text Encoder LoRA and DoRA Extraction.
 Specifically targets extracting text encoder deltas as LoRA/DoRA with custom mapping rules
 for layer name prefixes and PEFT weight formatting.
 """
-import fnmatch
 import os
-import re
 import torch
-import torch.linalg as linalg
 import folder_paths
 import comfy.utils
 from tqdm import tqdm
@@ -18,11 +15,11 @@ from .device_utils import (
     cleanup_after_operation
 )
 
-from unifiedefficientloader import MemoryEfficientSafeOpen, transfer_to_gpu_pinned, IncrementalSafetensorsWriter
+from unifiedefficientloader import MemoryEfficientSafeOpen, IncrementalSafetensorsWriter
+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
 from .lora_extract_svd import (
-    _compute_rank,
     _svd_extract_linear,
     _svd_extract_conv,
     _extract_chunked_layer,
@@ -106,91 +103,67 @@ def extract_te_from_files(
     try:
         keys_a = set(handler_a.keys())
         keys_b = set(handler_b.keys())
-        weight_keys = [k for k in keys_a if k.endswith(".weight")]
-        pbar = comfy.utils.ProgressBar(len(weight_keys))
+        weight_keys = sorted(k for k in keys_a if k.endswith(".weight"))
+        work_units = [
+            (key, key, key if key in keys_b else None)
+            for key in weight_keys
+            if not _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns)
+            and not (key not in keys_b and mismatch_mode == "skip")
+        ]
+        pbar = comfy.utils.ProgressBar(len(work_units))
         stats = {"extracted": 0, "full": 0, "skipped": 0, "chunked": 0}
 
-        def _process_layer(key):
+        def _process_layer(key, cpu_a, cpu_b):
             lora_name = _format_te_lora_key(key)
+            weight_diff = None
+            layer_device = device
 
             if _matches_any_pattern(key, skip_patterns, glob_mode=glob_skip_patterns):
                 return "skipped", None
 
             # Load tensors with pinned memory for CUDA
-            use_pinned = device == 'cuda'
-
             if key not in keys_b:
                 if mismatch_mode == "skip":
                     return "skipped", None
                 if mismatch_mode == "error":
                     raise ValueError(f"Key {key} not found in model B")
 
-                cpu_a = handler_a.get_tensor(key)
-                if use_pinned:
-                    weight_diff = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
+                if is_dora:
+                    weight_diff, dora_scale, layer_device = dora_difference(
+                        cpu_a, None, device
+                    )
                 else:
-                    weight_diff = cpu_a.to(device=device, dtype=torch.float32)
-                del cpu_a
-
-                dora_scale = None
+                    weight_diff, layer_device = lora_difference(cpu_a, None, device)
+                    dora_scale = None
             else:
-                cpu_a = handler_a.get_tensor(key)
-                cpu_b = handler_b.get_tensor(key)
-
-                if use_pinned:
-                    tensor_a = transfer_to_gpu_pinned(cpu_a, device, torch.float32)
-                    tensor_b = transfer_to_gpu_pinned(cpu_b, device, torch.float32)
-                else:
-                    tensor_a = cpu_a.to(device=device, dtype=torch.float32)
-                    tensor_b = cpu_b.to(device=device, dtype=torch.float32)
-                del cpu_a, cpu_b
-
-                if tensor_a.shape != tensor_b.shape:
+                if cpu_a.shape != cpu_b.shape:
                     if mismatch_mode == "skip":
-                        del tensor_a, tensor_b
                         return "skipped", None
                     if mismatch_mode == "error":
-                        raise ValueError(f"Shape mismatch for {key}: {tensor_a.shape} vs {tensor_b.shape}")
-
-                    weight_diff = tensor_a
-                    dora_scale = None
-                    del tensor_b
-                else:
-                    # Skip 1D weights
-                    if tensor_a.ndim < 2:
-                        del tensor_a, tensor_b
-                        return "skipped", None
+                        raise ValueError(f"Shape mismatch for {key}: {cpu_a.shape} vs {cpu_b.shape}")
 
                     if is_dora:
-                        # Weight Decomposition calculation
-                        is_conv_dim = tensor_a.ndim == 4
-
-                        if is_conv_dim:
-                            out_ch = tensor_a.shape[0]
-                            flat_a = tensor_a.reshape(out_ch, -1)
-                            flat_b = tensor_b.reshape(out_ch, -1)
-                            norm_a = torch.linalg.norm(flat_a, dim=1, keepdim=True)
-                            norm_b = torch.linalg.norm(flat_b, dim=1, keepdim=True)
-                        else:
-                            norm_a = torch.linalg.norm(tensor_a, dim=1, keepdim=True)
-                            norm_b = torch.linalg.norm(tensor_b, dim=1, keepdim=True)
-
-                        # Compute directional diff
-                        eps = 1e-8
-                        norm_ratio = norm_b / (norm_a + eps)
-                        if is_conv_dim:
-                            norm_ratio = norm_ratio.view(-1, 1, 1, 1)
-                            dora_scale = norm_a.view(-1, 1, 1, 1)
-                        else:
-                            dora_scale = norm_a
-
-                        w_dir = tensor_a * norm_ratio
-                        weight_diff = w_dir - tensor_b
-                        del tensor_a, tensor_b, w_dir
+                        weight_diff, dora_scale, layer_device = dora_difference(
+                            cpu_a, None, device
+                        )
                     else:
-                        weight_diff = tensor_a - tensor_b
+                        weight_diff, layer_device = lora_difference(cpu_a, None, device)
                         dora_scale = None
-                        del tensor_a, tensor_b
+                else:
+                    if cpu_a.ndim < 2:
+                        return "skipped", None
+                    if is_dora:
+                        weight_diff, dora_scale, layer_device = dora_difference(
+                            cpu_a, cpu_b, device
+                        )
+                    else:
+                        weight_diff, layer_device = lora_difference(
+                            cpu_a, cpu_b, device
+                        )
+                        dora_scale = None
+
+            if weight_diff is None:
+                raise RuntimeError(f"No extraction tensor was produced for {key}")
 
             # Skip small differences
             if min_diff > 0 and weight_diff.abs().max() < min_diff:
@@ -204,17 +177,20 @@ def extract_te_from_files(
 
             is_conv = weight_diff.ndim == 4
             layer_results = {}
+            svd_weight = weight_diff
 
             try:
                 if is_conv:
-                    result, mode_str = _svd_extract_conv(
-                        weight_diff, mode, conv_param, device, conv_max_rank,
-                        clamp_quantile, knee_probe_offset
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_conv(svd_weight, mode, conv_param, layer_device, conv_max_rank, clamp_quantile, knee_probe_offset),
+                        lambda: _svd_extract_conv(svd_weight.cpu(), mode, conv_param, "cpu", conv_max_rank, clamp_quantile, knee_probe_offset),
+                        layer_device,
                     )
                 else:
-                    result, mode_str = _svd_extract_linear(
-                        weight_diff, mode, linear_param, device, linear_max_rank,
-                        clamp_quantile, svd_niter, knee_probe_offset
+                    result, mode_str = retry_cuda_oom_on_cpu(
+                        lambda: _svd_extract_linear(svd_weight, mode, linear_param, layer_device, linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset),
+                        lambda: _svd_extract_linear(svd_weight.cpu(), mode, linear_param, "cpu", linear_max_rank, clamp_quantile, svd_niter, knee_probe_offset),
+                        layer_device,
                     )
             except Exception as e:
                 # Try chunked extraction for large tensors
@@ -223,7 +199,7 @@ def extract_te_from_files(
                     if num_chunks > 1:
                         print(f"[TE Extract] Chunked: {key} ({num_chunks} chunks)")
                         lora_up, lora_down, rank = _extract_chunked_layer(
-                            weight_diff, num_chunks, mode, linear_param, device,
+                            weight_diff, num_chunks, mode, linear_param, layer_device,
                             linear_max_rank, knee_probe_offset
                         )
                         if lora_up is not None:
@@ -258,8 +234,11 @@ def extract_te_from_files(
         writer = IncrementalSafetensorsWriter(output_path)
         writer.__enter__()
         try:
-            for key in tqdm(weight_keys, desc="Extracting TE Layers", unit="layers"):
-                status, layer_sd = _process_layer(key)
+            stream = paired_async_tensors(
+                handler_a, handler_b, work_units, pin_memory=str(device).startswith("cuda")
+            )
+            for key, cpu_a, cpu_b in tqdm(stream, total=len(work_units), desc="Extracting TE Layers", unit="layers"):
+                status, layer_sd = _process_layer(key, cpu_a, cpu_b)
                 stats[status] += 1
                 if layer_sd:
                     writer.write_dict(layer_sd)
@@ -271,6 +250,7 @@ def extract_te_from_files(
                         torch.cuda.empty_cache()
 
                 pbar.update(1)
+            stats["skipped"] += len(weight_keys) - len(work_units)
         finally:
             writer.__exit__(None, None, None)
 
@@ -302,7 +282,7 @@ def _get_te_model_inputs():
 def _get_common_inputs():
     return [
         io.Boolean.Input("lazy_load", default=True, tooltip="Low memory mode: load tensors from disk on demand"),
-        io.Boolean.Input("force_clear_cache", default=True, tooltip="Clear CUDA cache after each layer"),
+        io.Boolean.Input("force_clear_cache", default=False, tooltip="Clear CUDA cache after each layer; slower but useful under severe VRAM pressure."),
         io.Boolean.Input("chunk_large_layers", default=False,
                         tooltip="Split large fused layers (QKV, MLP) into chunks"),
         io.Float.Input("clamp_quantile", default=0.99, min=0.5, max=1.0, step=0.01,
@@ -376,6 +356,7 @@ class TextEncoderLoRAExtractRatio(io.ComfyNode):
                 *_get_te_model_inputs(),
                 io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for linear layers; higher values retain more singular values."),
                 io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
                 *_get_common_inputs(),
@@ -385,7 +366,7 @@ class TextEncoderLoRAExtractRatio(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -399,7 +380,8 @@ class TextEncoderLoRAExtractRatio(io.ComfyNode):
             linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
             clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
             mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
-            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
+            knee_probe_offset=probe_offset
         )
         return io.NodeOutput(output_path)
 
@@ -416,6 +398,7 @@ class TextEncoderLoRAExtractQuantile(io.ComfyNode):
                 *_get_te_model_inputs(),
                 io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."),
                 io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
                 *_get_common_inputs(),
@@ -425,7 +408,7 @@ class TextEncoderLoRAExtractQuantile(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -439,7 +422,8 @@ class TextEncoderLoRAExtractQuantile(io.ComfyNode):
             linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
             clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
             mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
-            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
+            knee_probe_offset=probe_offset
         )
         return io.NodeOutput(output_path)
 
@@ -498,6 +482,7 @@ class TextEncoderLoRAExtractFrobenius(io.ComfyNode):
                 *_get_te_model_inputs(),
                 io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."),
                 io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
                 *_get_common_inputs(),
@@ -507,7 +492,7 @@ class TextEncoderLoRAExtractFrobenius(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -521,7 +506,8 @@ class TextEncoderLoRAExtractFrobenius(io.ComfyNode):
             linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
             clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
             mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
-            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
+            knee_probe_offset=probe_offset
         )
         return io.NodeOutput(output_path)
 
@@ -581,6 +567,7 @@ class TextEncoderDoRAExtractRatio(io.ComfyNode):
                 *_get_te_model_inputs(),
                 io.Float.Input("linear_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for linear layers; higher values retain more singular values."),
                 io.Float.Input("conv_ratio", default=2.0, min=1.0, max=100.0, step=0.1, tooltip="Ratio threshold for convolution layers; higher values retain more singular values."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
                 *_get_common_inputs(),
@@ -590,7 +577,7 @@ class TextEncoderDoRAExtractRatio(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_ratio, conv_ratio, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -604,7 +591,8 @@ class TextEncoderDoRAExtractRatio(io.ComfyNode):
             linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
             clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
             mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
-            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
+            knee_probe_offset=probe_offset
         )
         return io.NodeOutput(output_path)
 
@@ -621,6 +609,7 @@ class TextEncoderDoRAExtractQuantile(io.ComfyNode):
                 *_get_te_model_inputs(),
                 io.Float.Input("linear_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for linear layers."),
                 io.Float.Input("conv_quantile", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target cumulative singular-value fraction for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
                 *_get_common_inputs(),
@@ -630,7 +619,7 @@ class TextEncoderDoRAExtractQuantile(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_quantile, conv_quantile, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -644,7 +633,8 @@ class TextEncoderDoRAExtractQuantile(io.ComfyNode):
             linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
             clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
             mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
-            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
+            knee_probe_offset=probe_offset
         )
         return io.NodeOutput(output_path)
 
@@ -703,6 +693,7 @@ class TextEncoderDoRAExtractFrobenius(io.ComfyNode):
                 *_get_te_model_inputs(),
                 io.Float.Input("linear_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for linear layers."),
                 io.Float.Input("conv_target", default=0.9, min=0.0, max=1.0, step=0.01, tooltip="Target Frobenius-norm fraction retained for convolution layers."),
+                io.Int.Input("probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values sampled beyond Max Rank for a reliable bounded rank decision."),
                 io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."),
                 io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."),
                 *_get_common_inputs(),
@@ -712,7 +703,7 @@ class TextEncoderDoRAExtractFrobenius(io.ComfyNode):
         )
 
     @classmethod
-    def execute(cls, model_a, model_b, linear_target, conv_target, linear_max_rank, conv_max_rank,
+    def execute(cls, model_a, model_b, linear_target, conv_target, probe_offset, linear_max_rank, conv_max_rank,
                 chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename,
                 save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput:
 
@@ -726,6 +717,7 @@ class TextEncoderDoRAExtractFrobenius(io.ComfyNode):
             linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank,
             clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns,
             mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers,
-            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns
+            lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns,
+            knee_probe_offset=probe_offset
         )
         return io.NodeOutput(output_path)

tests/test_adaptive_extraction_pipeline.py

@@ -0,0 +1,219 @@
+import importlib
+import sys
+import types
+from collections import Counter
+from pathlib import Path
+
+import pytest
+import torch
+from unifiedefficientloader import IncrementalSafetensorsWriter, MemoryEfficientSafeOpen
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+PACKAGE_NAME = "modelutils_adaptive_extraction_tests"
+
+
+def _load_module(name):
+    if PACKAGE_NAME not in sys.modules:
+        package = types.ModuleType(PACKAGE_NAME)
+        package.__path__ = [str(REPO_ROOT)]
+        sys.modules[PACKAGE_NAME] = package
+    return importlib.import_module(f"{PACKAGE_NAME}.{name}")
+
+
+@pytest.fixture(scope="module")
+def adaptive():
+    return _load_module("nodes.adaptive_svd")
+
+
+def _exact_partial(weight, q, niter):
+    del niter
+    U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
+    return U[:, :q], S[:q], Vh[:q, :].T
+
+
+@pytest.mark.parametrize(
+    ("mode", "parameter", "expected"),
+    [("ratio", 2.0, 2), ("quantile", 0.7, 2), ("sv_fro", 0.8, 1)],
+)
+def test_adaptive_partial_rank_matches_reference(monkeypatch, adaptive, mode, parameter, expected):
+    monkeypatch.setattr(adaptive.torch, "svd_lowrank", _exact_partial)
+    weight = torch.diag(torch.tensor([8.0, 5.0, 2.0, 1.0, 0.5, 0.25]))
+    _, _, _, rank = adaptive.adaptive_partial_svd(
+        weight, mode, parameter, max_rank=4, probe_offset=1, niter=2
+    )
+    assert rank == expected
+
+
+def test_unresolved_ratio_expands_and_caps(monkeypatch, adaptive):
+    probes = []
+
+    def tracked(weight, q, niter):
+        probes.append(q)
+        return _exact_partial(weight, q, niter)
+
+    monkeypatch.setattr(adaptive.torch, "svd_lowrank", tracked)
+    weight = torch.diag(torch.linspace(10.0, 8.0, 12))
+    _, _, _, rank = adaptive.adaptive_partial_svd(
+        weight, "ratio", 100.0, max_rank=4, probe_offset=2, niter=2
+    )
+    assert probes == [6, 8]
+    assert rank == 4
+
+
+def test_frobenius_uses_energy_outside_output_cap(monkeypatch, adaptive):
+    monkeypatch.setattr(adaptive.torch, "svd_lowrank", _exact_partial)
+    weight = torch.eye(10)
+    _, _, _, rank = adaptive.adaptive_partial_svd(
+        weight, "sv_fro", 0.9, max_rank=3, probe_offset=2, niter=2
+    )
+    assert rank == 3
+
+
+class _FakeStream:
+    def __init__(self, batches):
+        self._batches = iter(batches)
+        self.closed = False
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return next(self._batches)
+
+    def close(self):
+        self.closed = True
+
+
+class _FakeHandler:
+    def __init__(self, values):
+        self.values = values
+        self.calls = []
+        self.marked = []
+        self.streams = []
+
+    def async_stream(self, keys, **kwargs):
+        self.calls.append((tuple(keys), kwargs))
+        stream = _FakeStream([[(key, self.values[key])] for key in keys])
+        self.streams.append(stream)
+        return stream
+
+    def mark_processed(self, key):
+        self.marked.append(key)
+
+
+def test_paired_stream_is_bounded_ordered_and_releases_once():
+    streaming = _load_module("nodes.extraction_stream")
+    handler_a = _FakeHandler({"a1": torch.tensor(1), "a2": torch.tensor(2)})
+    handler_b = _FakeHandler({"b1": torch.tensor(3)})
+    units = [("one", "a1", "b1"), ("two", "a2", None)]
+
+    observed = []
+    for key, tensor_a, tensor_b in streaming.paired_async_tensors(
+        handler_a, handler_b, units, pin_memory=True
+    ):
+        observed.append((key, tensor_a.item(), None if tensor_b is None else tensor_b.item()))
+
+    assert observed == [("one", 1, 3), ("two", 2, None)]
+    assert handler_a.calls == [(('a1', 'a2'), {"batch_size": 1, "prefetch_batches": 1, "pin_memory": True})]
+    assert handler_b.calls == [(('b1',), {"batch_size": 1, "prefetch_batches": 1, "pin_memory": True})]
+    assert Counter(handler_a.marked) == Counter({"a1": 1, "a2": 1})
+    assert Counter(handler_b.marked) == Counter({"b1": 1})
+    assert all(stream.closed for stream in handler_a.streams + handler_b.streams)
+
+
+def test_cuda_oom_compute_retries_on_cpu(monkeypatch):
+    streaming = _load_module("nodes.extraction_stream")
+    monkeypatch.setattr(streaming, "release_cuda_oom", lambda: None)
+    calls = []
+
+    def fail_cuda():
+        calls.append("cuda")
+        raise RuntimeError("CUDA out of memory")
+
+    def run_cpu():
+        calls.append("cpu")
+        return 7
+
+    assert streaming.retry_cuda_oom_on_cpu(fail_cuda, run_cpu, "cuda:0") == 7
+    assert calls == ["cuda", "cpu"]
+
+
+def test_adaptive_controls_cover_all_public_families():
+    modules_and_nodes = [
+        (_load_module("nodes.lora_extract_svd"), prefix)
+        for prefix in ("LoRAExtractRatio", "LoRAExtractQuantile", "LoRAExtractFrobenius")
+    ]
+    modules_and_nodes += [
+        (_load_module("nodes.dora_extract_wd"), prefix)
+        for prefix in ("DoRAExtractRatio", "DoRAExtractQuantile", "DoRAExtractFrobenius")
+    ]
+    modules_and_nodes += [
+        (_load_module("nodes.dora_learned_wd"), prefix)
+        for prefix in ("DoRALearnedExtractRatio", "DoRALearnedExtractQuantile", "DoRALearnedExtractFrobenius")
+    ]
+    modules_and_nodes += [
+        (_load_module("nodes.text_encoder_extract"), prefix)
+        for prefix in (
+            "TextEncoderLoRAExtractRatio", "TextEncoderLoRAExtractQuantile",
+            "TextEncoderLoRAExtractFrobenius", "TextEncoderDoRAExtractRatio",
+            "TextEncoderDoRAExtractQuantile", "TextEncoderDoRAExtractFrobenius",
+        )
+    ]
+
+    for module, node_name in modules_and_nodes:
+        names = [item.id for item in getattr(module, node_name).define_schema().inputs]
+        probe_index = names.index("probe_offset")
+        assert names[probe_index + 1] == "linear_max_rank"
+
+
+def _write_uel(path, tensors):
+    with IncrementalSafetensorsWriter(str(path), max_workers=1) as writer:
+        writer.write_dict(tensors)
+
+
+def test_lora_extraction_uses_only_async_uel_and_releases_sources(monkeypatch, tmp_path):
+    extraction = _load_module("nodes.lora_extract_svd")
+    model_a = tmp_path / "a.safetensors"
+    model_b = tmp_path / "b.safetensors"
+    output = tmp_path / "out.safetensors"
+    key = "model.diffusion_model.layer.weight"
+    _write_uel(model_a, {key: torch.eye(3)})
+    _write_uel(model_b, {key: torch.zeros(3, 3)})
+    monkeypatch.setattr(extraction, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(extraction, "cleanup_after_operation", lambda: None)
+    original_async = extraction.MemoryEfficientSafeOpen.async_stream
+    original_mark = extraction.MemoryEfficientSafeOpen.mark_processed
+    async_calls = []
+    marked = []
+
+    def track_async(self, keys, **kwargs):
+        async_calls.append((self.filename, tuple(keys), kwargs))
+        yield from original_async(self, keys, **kwargs)
+
+    def track_mark(self, source_key):
+        marked.append((self.filename, source_key))
+        return original_mark(self, source_key)
+
+    monkeypatch.setattr(extraction.MemoryEfficientSafeOpen, "async_stream", track_async)
+    monkeypatch.setattr(extraction.MemoryEfficientSafeOpen, "mark_processed", track_mark)
+    monkeypatch.setattr(
+        extraction.MemoryEfficientSafeOpen,
+        "get_tensor",
+        lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("synchronous tensor load")),
+    )
+
+    extraction.extract_lora_from_files(
+        str(model_a), str(model_b), "fixed", 1, 1, "cpu", "fp32", str(output),
+        linear_max_rank=1, conv_max_rank=1, force_clear_cache=False,
+    )
+
+    assert output.is_file()
+    assert len(async_calls) == 2
+    assert all(options == {"batch_size": 1, "prefetch_batches": 1, "pin_memory": False} for _, _, options in async_calls)
+    assert Counter(marked) == Counter((filename, key) for filename, keys, _ in async_calls for key in keys)
+    with MemoryEfficientSafeOpen(str(output), low_memory=True) as result:
+        assert sorted(result.keys()) == [
+            "diffusion_model.layer.lora_A.weight",
+            "diffusion_model.layer.lora_B.weight",
+        ]