Omni: Account whole LoRA weights in ComfyUI memory budgets

omni/ComfyUI-OmniXPU/README.md

@@ -18,6 +18,7 @@ No workflow or model-pipeline replacement is required.
 |---|---|
 | Kitchen XPU backend | INT8/QTensor operations, FP8 QDQ and stochastic rounding, SVDQuant, AdaLN, four RoPE APIs, and ConvRot |
 | ComfyUI adapter | Attention routing, LayerNorm/RMSNorm class integration, the remaining FP8 model/factory bridge, and fused Lumina/Z-Image INT8 FFN wiring |
+| Memory adapter | Cached whole-LoRA model budgets plus optional DynamicVRAM per-layer XPU staging measurements |
 | Legacy fix | Global `F.interpolate` and `torch.median`/`torch.nanmedian` workarounds; disabled by default |
 
 RoPE, generic INT8 linear dispatch, and the old FP8 negative-zero wrapper are
@@ -45,8 +46,15 @@ OMNIXPU_ATTENTION=0         # Disable the attention adapter
 OMNIXPU_NORM=0              # Disable the norm adapter
 OMNIXPU_FP8_GEMM=0          # Disable the temporary FP8 model/factory adapter
 OMNIXPU_INT8_FFN=0          # Disable fused Lumina/Z-Image INT8 FFN wiring
+OMNIXPU_DYNAMIC_VRAM_BOUNDARY_TRIM=0  # Disable Windows XPU model-boundary trim
+OMNIXPU_LORA_MEMORY=0       # Disable cached whole-LoRA budgets and staging logs
 ```
 
+On Windows XPU, the boundary trim turns an unmet DynamicVRAM minimum-memory
+budget into an explicit partial VBAR reclaim before model loading. It preserves
+loaded models and is enabled by default; the environment variable above is the
+A/B-test escape hatch.
+
 Validated sub-routes can be disabled independently:
 
 ```bash
@@ -91,6 +99,18 @@ Dispatch decisions and fallback reasons:
 OMNIXPU_DEBUG_VERBOSE=1 python main.py
 ```
 
+LoRA weights are measured once when the LoRA node executes. Unique tensor sizes
+are cached in a `ModelPatcher` attachment, inherited by clones, accumulated for
+stacked LoRAs, and added to both `memory_required` and an explicitly supplied
+`minimum_memory_required`. The base model's `model_size()` semantics stay
+unchanged. Model loads read the cached attachment instead of rescanning patches.
+DynamicVRAM layer scanning is disabled by default. To diagnose every LoRA
+staging operation, including its XPU state and any failure, enable:
+
+```bash
+OMNIXPU_LORA_MEMORY_TRACE=1 python main.py
+```
+
 Set tracing variables before startup. The **OmniXPU Status** node reports:
 
 - GPU and `omni_xpu_kernel` capabilities;

omni/ComfyUI-OmniXPU/adapters/dynamic_vram.py

@@ -0,0 +1,145 @@
+"""Reclaim DynamicVRAM residency at Windows XPU model boundaries."""
+
+from __future__ import annotations
+
+import functools
+import logging
+import sys
+
+
+log = logging.getLogger("ComfyUI-OmniXPU")
+
+_PATCH_MARKER = "__omnixpu_dynamic_vram_boundary_original__"
+
+
+def _load_argument(args, kwargs, name, position, default):
+    if name in kwargs:
+        return kwargs[name]
+    if len(args) > position:
+        return args[position]
+    return default
+
+
+def _expanded_models(models):
+    pending = list(models)
+    expanded = []
+    seen = set()
+    while pending:
+        model = pending.pop(0)
+        identity = id(model)
+        if identity in seen:
+            continue
+        seen.add(identity)
+        expanded.append(model)
+        pending.extend(model.model_patches_models())
+    return expanded
+
+
+def _minimum_memory_target(model_management, args, kwargs):
+    memory_required = _load_argument(args, kwargs, "memory_required", 0, 0)
+    minimum_required = _load_argument(
+        args, kwargs, "minimum_memory_required", 2, None
+    )
+    reserved = model_management.extra_reserved_memory()
+    inference = model_management.minimum_inference_memory()
+    if minimum_required is None:
+        return max(inference, memory_required + reserved)
+    return max(inference, minimum_required + reserved)
+
+
+def _trim_pass(model_management, device, target, requested, include_requested):
+    reclaimed = 0
+    trimmed = 0
+    for loaded in reversed(model_management.current_loaded_models):
+        model = loaded.model
+        if loaded.device != device or loaded.is_dead() or not model.is_dynamic():
+            continue
+        is_requested = id(model) in requested
+        if is_requested != include_requested:
+            continue
+
+        shortfall = target - model_management.get_free_memory(device)
+        if shortfall <= 0:
+            break
+        if model.loaded_size() <= 0:
+            continue
+
+        freed = int(model.partially_unload(model.offload_device, shortfall))
+        if freed > 0:
+            reclaimed += freed
+            trimmed += 1
+    return reclaimed, trimmed
+
+
+def _trim_dynamic_boundary(model_management, models, target):
+    if not models or not all(model.is_dynamic() for model in models):
+        return
+
+    requested = {id(model) for model in models}
+    devices = {model.load_device for model in models}
+    for device in devices:
+        if getattr(device, "type", None) != "xpu":
+            continue
+
+        free_before = int(model_management.get_free_memory(device))
+        if free_before >= target:
+            continue
+
+        inactive_bytes, inactive_models = _trim_pass(
+            model_management, device, target, requested, False
+        )
+        active_bytes, active_models = _trim_pass(
+            model_management, device, target, requested, True
+        )
+        free_after = int(model_management.get_free_memory(device))
+        log.info(
+            "[OmniXPU] DynamicVRAM boundary trim: device=%s target=%.1fMiB "
+            "free=%.1f->%.1fMiB reclaimed=%.1fMiB models=%d inactive/%d active",
+            device,
+            target / (1024 ** 2),
+            free_before / (1024 ** 2),
+            free_after / (1024 ** 2),
+            (inactive_bytes + active_bytes) / (1024 ** 2),
+            inactive_models,
+            active_models,
+        )
+
+
+def _patch_model_loader(model_management):
+    original = model_management.load_models_gpu
+    if hasattr(original, _PATCH_MARKER):
+        return
+
+    @functools.wraps(original)
+    def boundary_trimmed(models, *args, **kwargs):
+        expanded = _expanded_models(models)
+        target = _minimum_memory_target(model_management, args, kwargs)
+        _trim_dynamic_boundary(model_management, expanded, target)
+        return original(models, *args, **kwargs)
+
+    setattr(boundary_trimmed, _PATCH_MARKER, original)
+    model_management.load_models_gpu = boundary_trimmed
+
+
+def apply():
+    if sys.platform != "win32":
+        return False, "Windows only"
+
+    import comfy.model_management
+
+    required = (
+        "current_loaded_models",
+        "extra_reserved_memory",
+        "get_free_memory",
+        "load_models_gpu",
+        "minimum_inference_memory",
+    )
+    missing = [name for name in required if not hasattr(comfy.model_management, name)]
+    if missing:
+        return False, "missing ComfyUI hooks: " + ", ".join(missing)
+
+    _patch_model_loader(comfy.model_management)
+    return True, "Windows XPU DynamicVRAM boundary reclaim enabled"
+
+
+__all__ = ["apply"]