pollockjj-ComfyUI-MultiGPU/device_utils.py

"""
Device detection, management, and inspection utilities for ComfyUI-MultiGPU.
Single source of truth for all device enumeration, compatibility checks, and state inspection.
Handles all device types supported by ComfyUI core.
"""

import torch
import logging
import hashlib
import psutil
import comfy.model_management as mm
import gc
from datetime import datetime, timezone
import server
import weakref
import platform
import ctypes
import sys
import comfy.model_patcher

# DisTorch stores for pruning/diagnostics
from .distorch_2 import (
    safetensor_allocation_store,
    safetensor_settings_store,
    create_safetensor_model_hash,
)

# Optional DisTorch v1 store support
try:
    from .distorch import (
        model_allocation_store,
        create_model_hash,
    )
except Exception:
    model_allocation_store = {}
    create_model_hash = None

logger = logging.getLogger("MultiGPU")

# Module-level cache for device list (populated once on first call)
_DEVICE_LIST_CACHE = None

# ==========================================================================================
# Executor Cache Management and CPU Monitoring (Phases 1, 2, 3)
# ==========================================================================================

# Configuration for CPU Monitoring (Phase 2)
CPU_MEMORY_THRESHOLD_PERCENT = 85.0
# Hysteresis: Only trigger again if usage increased by this amount since the last reset.
CPU_RESET_HYSTERESIS_PERCENT = 5.0
_last_cpu_usage_at_reset = 0.0

def clear_memory_snapshot_history():
    """Clears the stored memory snapshot history. (Phase 3)"""
    # Logging integration
    multigpu_memory_log("mem_mgmt", "pre-history-clear")

    # Snapshot globals exist in this module; operate safely in case of reload
    if '_MEM_SNAPSHOT_LAST' in globals():
        globals()['_MEM_SNAPSHOT_LAST'].clear()
    if '_MEM_SNAPSHOT_SERIES' in globals():
        globals()['_MEM_SNAPSHOT_SERIES'].clear()
    logger.debug("[MultiGPU_Memory_Management] Memory snapshot history cleared.")

    # Logging integration
    multigpu_memory_log("mem_mgmt", "post-history-clear")

def trigger_executor_cache_reset(reason="policy", force=False):
    """
    (Phase 1/2 Core) Triggers PromptExecutor.reset() by setting the 'free_memory' flag.
    Releases CPU-side references held by execution caches.
    """
    global _last_cpu_usage_at_reset

    # Ensure PromptServer singleton is available
    if server.PromptServer.instance is None:
        logger.debug("[MultiGPU_Memory_Management] PromptServer instance not yet initialized.")
        return

    prompt_server = server.PromptServer.instance

    # Stability guard: Avoid during active execution unless forced
    if prompt_server.prompt_queue.currently_running and not force:
        logger.debug(f"[MultiGPU_Memory_Management] Skipping Executor Cache Reset during active prompt execution (Reason: {reason}).")
        return

    multigpu_memory_log("executor_reset", f"pre-trigger ({reason})")
    logger.info(f"[MultiGPU_Memory_Management] Triggering PromptExecutor cache reset (e.reset()). Reason: {reason}")

    # Diagnostics and store pruning prior to reset
    analyze_cpu_memory_leaks(force=force)
    prune_distorch_stores()

    # Phase 3: Clear internal snapshot history as the context is resetting
    clear_memory_snapshot_history()

    # Set the flag on the prompt queue (ComfyUI core mechanism)
    prompt_server.prompt_queue.set_flag("free_memory", True)
    logger.debug("[MultiGPU_Memory_Management] 'free_memory' flag set.")

    # Update usage baseline for hysteresis
    vm = psutil.virtual_memory()
    _last_cpu_usage_at_reset = vm.percent

    # Attempt to return freed memory to OS
    try_malloc_trim()

    multigpu_memory_log("executor_reset", f"post-trigger ({reason})")


def _cpu_used_bytes():
    try:
        vm = psutil.virtual_memory()
        return vm.used
    except Exception:
        return 0


def force_full_system_cleanup(reason="manual", force=True):
    """
    Mirror ComfyUI-Manager 'Free model and node cache' semantics:
    - Only set unload_models=True and free_memory=True flags on the PromptQueue
    - The prompt worker (main.py) performs unload/reset/GC
    """
    pre_cpu = _cpu_used_bytes()
    pre_models = len(getattr(mm, "current_loaded_models", []))

    multigpu_memory_log("full_cleanup", f"start:{reason}")
    logger.mgpu_mm_log(f"[ManagerMatch] Requesting flags-only cleanup (reason={reason}) | pre_models={pre_models}, cpu_used_gib={pre_cpu/(1024**3):.2f}")

    try:
        if server.PromptServer.instance is not None:
            pq = server.PromptServer.instance.prompt_queue
            # Respect currently_running unless forced
            if (not pq.currently_running) or force:
                pq.set_flag("unload_models", True)
                pq.set_flag("free_memory", True)
                logger.mgpu_mm_log("[ManagerMatch] Flags set: unload_models=True, free_memory=True")
            else:
                logger.mgpu_mm_log("[ManagerMatch] Skipped setting flags due to active execution and force=False")
    except Exception as e:
        logger.mgpu_mm_log(f"[ManagerMatch] Failed to set flags: {e}")

    post_cpu = _cpu_used_bytes()
    post_models = len(getattr(mm, "current_loaded_models", []))
    delta_cpu_mb = (post_cpu - pre_cpu) / (1024**2)

    multigpu_memory_log("full_cleanup", f"requested:{reason}")
    summary = (
        f"[ManagerMatch] Flags-only cleanup requested (reason={reason}) | "
        f"models {pre_models}->{post_models} (no immediate unload), cpu_delta_mb={delta_cpu_mb:.2f}"
    )
    logger.mgpu_mm_log(summary)
    return summary

def check_cpu_memory_threshold(threshold_percent=CPU_MEMORY_THRESHOLD_PERCENT):
    """
    (Phase 2) Checks CPU memory usage and triggers a reset if threshold is exceeded (with hysteresis).
    """
    # Ensure PromptServer singleton is available
    if server.PromptServer.instance is None:
        return

    # Stability/optimization: Do not trigger during active execution
    if server.PromptServer.instance.prompt_queue.currently_running:
        return

    vm = psutil.virtual_memory()
    current_usage = vm.percent

    if current_usage > threshold_percent:
        # Hysteresis gating
        if current_usage > (_last_cpu_usage_at_reset + CPU_RESET_HYSTERESIS_PERCENT):
            logger.warning(f"[MultiGPU_Memory_Monitor] CPU usage ({current_usage:.1f}%) exceeds threshold ({threshold_percent:.1f}%) and hysteresis.")
            multigpu_memory_log("cpu_monitor", f"trigger:{current_usage:.1f}pct")
            trigger_executor_cache_reset(reason="cpu_threshold_exceeded", force=False)
        else:
            logger.debug(f"[MultiGPU_Memory_Monitor] CPU usage high ({current_usage:.1f}%) but within hysteresis range. Skipping reset.")
            multigpu_memory_log("cpu_monitor", f"skip_hysteresis:{current_usage:.1f}pct")

def get_device_list():
    """
    Enumerate ALL physically available devices that can store torch tensors.
    This includes all device types supported by ComfyUI core.
    Results are cached after first call since devices don't change during runtime.
    
    Returns a comprehensive list of all available devices across all types:
    - CPU (always available)
    - CUDA devices (NVIDIA GPUs)
    - XPU devices (Intel GPUs)
    - NPU devices (Ascend NPUs from Huawei)
    - MLU devices (Cambricon MLUs)
    - MPS device (Apple Metal)
    - DirectML devices (Windows DirectML)
    - CoreX/IXUCA devices
    """
    global _DEVICE_LIST_CACHE
    
    # Return cached result if already populated
    if _DEVICE_LIST_CACHE is not None:
        return _DEVICE_LIST_CACHE
    
    # First time - do the actual detection
    devs = []
    
    # CPU is always physically present and can store tensors
    devs.append("cpu")
    
    # CUDA devices (NVIDIA GPUs)
    try:
        if hasattr(torch, "cuda") and hasattr(torch.cuda, "is_available") and torch.cuda.is_available():
            device_count = torch.cuda.device_count()
            devs += [f"cuda:{i}" for i in range(device_count)]
            logger.debug(f"[MultiGPU_Device_Utils] Found {device_count} CUDA device(s)")
    except Exception as e:
        logger.debug(f"[MultiGPU_Device_Utils] CUDA detection failed: {e}")
    
    # XPU devices (Intel GPUs)
    try:
        # Try to import intel extension first (may be required for XPU support)
        import intel_extension_for_pytorch as ipex
    except ImportError:
        pass
    try:
        if hasattr(torch, "xpu") and hasattr(torch.xpu, "is_available") and torch.xpu.is_available():
            device_count = torch.xpu.device_count()
            devs += [f"xpu:{i}" for i in range(device_count)]
            logger.debug(f"[MultiGPU_Device_Utils] Found {device_count} XPU device(s)")
    except Exception as e:
        logger.debug(f"[MultiGPU_Device_Utils] XPU detection failed: {e}")
    
    # NPU devices (Ascend NPUs from Huawei)
    try:
        import torch_npu
        if hasattr(torch, "npu") and hasattr(torch.npu, "is_available") and torch.npu.is_available():
            device_count = torch.npu.device_count()
            devs += [f"npu:{i}" for i in range(device_count)]
            logger.debug(f"[MultiGPU_Device_Utils] Found {device_count} NPU device(s)")
    except Exception as e:
        logger.debug(f"[MultiGPU_Device_Utils] NPU detection failed: {e}")
    
    # MLU devices (Cambricon MLUs)
    try:
        import torch_mlu
        if hasattr(torch, "mlu") and hasattr(torch.mlu, "is_available") and torch.mlu.is_available():
            device_count = torch.mlu.device_count()
            devs += [f"mlu:{i}" for i in range(device_count)]
            logger.debug(f"[MultiGPU_Device_Utils] Found {device_count} MLU device(s)")
    except Exception as e:
        logger.debug(f"[MultiGPU_Device_Utils] MLU detection failed: {e}")
    
    # MPS device (Apple Metal - single device only)
    try:
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            devs.append("mps")
            logger.debug("[MultiGPU_Device_Utils] Found MPS device")
    except Exception as e:
        logger.debug(f"[MultiGPU_Device_Utils] MPS detection failed: {e}")
    
    # DirectML devices (Windows DirectML for AMD/Intel/NVIDIA)
    try:
        import torch_directml
        adapter_count = torch_directml.device_count()
        if adapter_count > 0:
            devs += [f"directml:{i}" for i in range(adapter_count)]
            logger.debug(f"[MultiGPU_Device_Utils] Found {adapter_count} DirectML adapter(s)")
    except Exception as e:
        logger.debug(f"[MultiGPU_Device_Utils] DirectML detection failed: {e}")
    
    # IXUCA/CoreX devices (special accelerator)
    try:
        if hasattr(torch, "corex"):
            # CoreX typically exposes single device, but check if there's a count method
            if hasattr(torch.corex, "device_count"):
                device_count = torch.corex.device_count()
                devs += [f"corex:{i}" for i in range(device_count)]
                logger.debug(f"[MultiGPU_Device_Utils] Found {device_count} CoreX device(s)")
            else:
                devs.append("corex:0")
                logger.debug("[MultiGPU_Device_Utils] Found CoreX device")
    except Exception as e:
        logger.debug(f"[MultiGPU_Device_Utils] CoreX detection failed: {e}")
    
    # Cache the result for future calls
    _DEVICE_LIST_CACHE = devs
    
    # Log only once when initially populated
    logger.debug(f"[MultiGPU_Device_Utils] Device list initialized: {devs}")
    
    return devs


def is_accelerator_available():
    """
    Check if any accelerator device is available.
    Used by patched functions to determine CPU fallback.
    
    Returns True if any GPU/accelerator is available, False otherwise.
    """
    # Check CUDA
    try:
        if torch.cuda.is_available():
            return True
    except:
        pass
    
    # Check XPU (Intel GPU)
    try:
        if hasattr(torch, "xpu") and torch.xpu.is_available():
            return True
    except:
        pass
    
    # Check NPU (Ascend)
    try:
        import torch_npu
        if hasattr(torch, "npu") and torch.npu.is_available():
            return True
    except:
        pass
    
    # Check MLU (Cambricon)
    try:
        import torch_mlu
        if hasattr(torch, "mlu") and torch.mlu.is_available():
            return True
    except:
        pass
    
    # Check MPS (Apple Metal)
    try:
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return True
    except:
        pass
    
    # Check DirectML
    try:
        import torch_directml
        if torch_directml.device_count() > 0:
            return True
    except:
        pass
    
    # Check CoreX/IXUCA
    try:
        if hasattr(torch, "corex"):
            return True
    except:
        pass
    
    return False


def is_device_compatible(device_string):
    """
    Check if a device string represents a valid, available device.
    
    Args:
        device_string: Device identifier like "cuda:0", "cpu", "xpu:1", etc.
    
    Returns:
        True if the device is available, False otherwise.
    """
    available_devices = get_device_list()
    return device_string in available_devices


def get_device_type(device_string):
    """
    Extract the device type from a device string.
    
    Args:
        device_string: Device identifier like "cuda:0", "cpu", "xpu:1", etc.
    
    Returns:
        Device type string (e.g., "cuda", "cpu", "xpu", "npu", "mlu", "mps", "directml", "corex")
    """
    if ":" in device_string:
        return device_string.split(":")[0]
    return device_string


def parse_device_string(device_string):
    """
    Parse a device string into type and index.

    Args:
        device_string: Device identifier like "cuda:0", "cpu", "xpu:1", etc.

    Returns:
        Tuple of (device_type, device_index) where index is None for non-indexed devices
    """
    if ":" in device_string:
        parts = device_string.split(":")
        return parts[0], int(parts[1])
    return device_string, None


def soft_empty_cache_multigpu():
    """
    Replicate ComfyUI's cache clearing but for ALL devices in MultiGPU.
    MultiGPU adaptation of ComfyUI's soft_empty_cache() functionality.
    Uses context managers to ensure the calling thread's device context is restored.
    """
    import gc

    logger.mgpu_mm_log("soft_empty_cache_multigpu: starting GC and multi-device cache clear")
    # Record pre-GC snapshot for general system view
    multigpu_memory_log("general", "pre-soft-empty")

    multigpu_memory_log("general", "pre-gc")
    # Lifecycle status before GC
    log_tracked_modelpatchers_status(tag="pre-gc")
    gc.collect()
    # Lifecycle status after GC
    log_tracked_modelpatchers_status(tag="post-gc")
    multigpu_memory_log("general", "post-gc")
    logger.mgpu_mm_log("soft_empty_cache_multigpu: garbage collection complete")
    # Attempt to release freed heap memory to OS
    try_malloc_trim()

    # Clear cache for ALL devices (not just ComfyUI's single device)
    all_devices = get_device_list()
    logger.mgpu_mm_log(f"soft_empty_cache_multigpu: devices to clear = {all_devices}")
    
    # Check global availability first to avoid unnecessary iteration if backend is missing
    is_cuda_available = hasattr(torch, "cuda") and hasattr(torch.cuda, "is_available") and torch.cuda.is_available()

    for device_str in all_devices:
        if device_str.startswith("cuda:"):
            if is_cuda_available:
                device_idx = int(device_str.split(":")[1])
                # Use context manager for safe switching and automatic restoration
                logger.mgpu_mm_log(f"Clearing CUDA cache on {device_str} (idx={device_idx})")
                multigpu_memory_log("general", f"pre-empty:{device_str}")
                with torch.cuda.device(device_idx):
                    torch.cuda.empty_cache()
                    if hasattr(torch.cuda, "ipc_collect"):
                        torch.cuda.ipc_collect()  # ComfyUI's CUDA optimization
                logger.mgpu_mm_log(f"Cleared CUDA cache (and IPC if available) on {device_str}")
                multigpu_memory_log("general", f"post-empty:{device_str}")

        elif device_str == "mps":
            if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
                logger.mgpu_mm_log("Clearing MPS cache")
                multigpu_memory_log("general", f"pre-empty:{device_str}")
                torch.mps.empty_cache()
                logger.mgpu_mm_log("Cleared MPS cache")
                multigpu_memory_log("general", f"post-empty:{device_str}")

        elif device_str.startswith("xpu:"):
            if hasattr(torch, "xpu") and hasattr(torch.xpu, "empty_cache"):
                logger.mgpu_mm_log(f"Clearing XPU cache on {device_str}")
                multigpu_memory_log("general", f"pre-empty:{device_str}")
                torch.xpu.empty_cache()
                logger.mgpu_mm_log(f"Cleared XPU cache on {device_str}")
                multigpu_memory_log("general", f"post-empty:{device_str}")

        elif device_str.startswith("npu:"):
            if hasattr(torch, "npu") and hasattr(torch.npu, "empty_cache"):
                logger.mgpu_mm_log(f"Clearing NPU cache on {device_str}")
                multigpu_memory_log("general", f"pre-empty:{device_str}")
                torch.npu.empty_cache()
                logger.mgpu_mm_log(f"Cleared NPU cache on {device_str}")
                multigpu_memory_log("general", f"post-empty:{device_str}")

        elif device_str.startswith("mlu:"):
            if hasattr(torch, "mlu") and hasattr(torch.mlu, "empty_cache"):
                logger.mgpu_mm_log(f"Clearing MLU cache on {device_str}")
                multigpu_memory_log("general", f"pre-empty:{device_str}")
                torch.mlu.empty_cache()
                logger.mgpu_mm_log(f"Cleared MLU cache on {device_str}")
                multigpu_memory_log("general", f"post-empty:{device_str}")

        elif device_str.startswith("corex:"):
            if hasattr(torch, "corex") and hasattr(torch.corex, "empty_cache"):
                logger.mgpu_mm_log(f"Clearing CoreX cache on {device_str}")
                multigpu_memory_log("general", f"pre-empty:{device_str}")
                torch.corex.empty_cache()
                logger.mgpu_mm_log(f"Cleared CoreX cache on {device_str}")
                multigpu_memory_log("general", f"post-empty:{device_str}")

    # Record post-GC snapshot for general system view
    multigpu_memory_log("general", "post-soft-empty")



def _bytes_to_gib(b: int) -> float:
    """Convert bytes to GiB as a float."""
    return float(b) / (1024.0 ** 3)


def comfyui_memory_load(tag: str) -> str:
    """
    Returns a single-line, pipe-delimited snapshot of system and device memory usage.

    Format: "tag=<TAG>|cpu=<used_GiB>/<total_GiB>|<device>=<used_GiB>/<total_GiB>|..."
    - CPU values represent system RAM via psutil.
    - Device values represent VRAM via comfy.model_management across all non-CPU devices.
    - Device identifiers use the torch device string from get_device_list() (e.g., 'cuda:0', 'xpu:0', 'mps').
    - Values are in GiB with 2 decimals.
    """
    # CPU RAM
    vm = psutil.virtual_memory()
    cpu_used_gib = _bytes_to_gib(vm.used)
    cpu_total_gib = _bytes_to_gib(vm.total)

    segments = [f"tag={tag}", f"cpu={cpu_used_gib:.2f}/{cpu_total_gib:.2f}"]

    # Enumerate non-CPU devices
    devices = [d for d in get_device_list() if d != "cpu"]

    # Append per-device VRAM used/total
    for dev_str in devices:
        device = torch.device(dev_str)
        total = mm.get_total_memory(device)
        free_info = mm.get_free_memory(device, torch_free_too=True)
        # free_info may be a tuple (system_free, torch_cache_free) or a single value
        if isinstance(free_info, tuple):
            system_free = free_info[0]
        else:
            system_free = free_info
        used = max(0, (total or 0) - (system_free or 0))

        used_gib = _bytes_to_gib(used)
        total_gib = _bytes_to_gib(total or 0)
        if total_gib > 0:
            segments.append(f"{dev_str}={used_gib:.2f}/{total_gib:.2f}")

    return "|".join(segments)


# ==========================================================================================
# Delta-capable memory logging (identifier + tag) with timestamped series
# ==========================================================================================


# Stores the last snapshot per identifier: identifier -> (last_tag, snapshot_map)
# snapshot_map: device_str -> (used_bytes, total_bytes)
_MEM_SNAPSHOT_LAST = {}

# Full chronological series per identifier: identifier -> list[(timestamp, tag, snapshot_map)]
_MEM_SNAPSHOT_SERIES = {}


def _capture_memory_snapshot() -> dict[str, tuple[int, int]]:
    """
    Capture an absolute memory snapshot for CPU and all non-CPU devices.
    Values are returned in bytes (used, total) for each device string key.
    """
    snapshot: dict[str, tuple[int, int]] = {}

    # CPU
    vm = psutil.virtual_memory()
    snapshot["cpu"] = (vm.used, vm.total)

    # Non-CPU devices
    devices = [d for d in get_device_list() if d != "cpu"]
    for dev_str in devices:
        device = torch.device(dev_str)
        total = mm.get_total_memory(device)
        free_info = mm.get_free_memory(device, torch_free_too=True)
        system_free = free_info[0] if isinstance(free_info, tuple) else free_info
        used = max(0, (total or 0) - (system_free or 0))
        snapshot[dev_str] = (used, total or 0)

    return snapshot


def _format_delta_gib(delta_bytes: int) -> str:
    """Format a signed GiB delta with two decimals."""
    gib = _bytes_to_gib(abs(delta_bytes))
    sign = "+" if delta_bytes >= 0 else "-"
    return f"{sign}{gib:.2f}"


def memory_print_summary(log: logging.Logger = logger):
    """
    Print the entire run as absolute actuals with timestamps for each identifier.
    One line per recorded snapshot in insertion order.
    Format:
      YYYY-MM-DDTHH:MM:SS.mmmZ identifier tag | cpu=U/T | cuda:0=U/T | ...
      (GiB values, two decimals)
    """
    from . import logger
    
    # Stable identifier order for readability
    for identifier in sorted(_MEM_SNAPSHOT_SERIES.keys()):
        series = _MEM_SNAPSHOT_SERIES[identifier]
        if not series:
            continue
        logger.mgpu_mm_log(f"=== memory summary: {identifier} ===")
        for ts, tag, snap in series:
            # Build device list (cpu first, then sorted devices)
            parts = []
            # CPU
            cpu_used, cpu_total = snap.get("cpu", (0, 0))
            parts.append(f"cpu={_bytes_to_gib(cpu_used):.2f}/{_bytes_to_gib(cpu_total):.2f}")
            # Non-CPU (sorted)
            devs = sorted([k for k in snap.keys() if k != "cpu"])
            for dev in devs:
                used, total = snap[dev]
                parts.append(f"{dev}={_bytes_to_gib(used):.2f}/{_bytes_to_gib(total):.2f}")
            ts_str = ts.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
            logger.mgpu_mm_log(f"{ts_str} {identifier} {tag} | " + " | ".join(parts))


def multigpu_memory_log(identifier: str, tag: str, log: logging.Logger = logger):
    """
    Record a timestamped memory snapshot for the given identifier and tag.
    - INFO: per-device deltas vs. previous snapshot for the same identifier (GiB, signed, no totals).
    - DEBUG: absolute snapshot string via comfyui_memory_load(tag) prefixed by identifier.
    - Special identifier: 'print_summary' will dump the entire series as actuals with timestamps.
    """
    from . import logger as mgpu_logger
    
    if identifier == "print_summary":
        memory_print_summary(log=log)
        return

    # Capture current snapshot and timestamp
    ts = datetime.now(timezone.utc)
    curr = _capture_memory_snapshot()

    # Append to full series
    series = _MEM_SNAPSHOT_SERIES.get(identifier)
    if series is None:
        series = []
        _MEM_SNAPSHOT_SERIES[identifier] = series
    series.append((ts, tag, curr))

    # Compute and log delta vs last
    if identifier in _MEM_SNAPSHOT_LAST:
        prev_tag, prev = _MEM_SNAPSHOT_LAST[identifier]
        # Union of device keys
        keys = set(prev.keys()) | set(curr.keys())
        # Stable order: cpu first, then sorted devices
        ordered = ["cpu"] + sorted([k for k in keys if k != "cpu"])
        parts = []
        for k in ordered:
            p_used, _p_tot = prev.get(k, (0, curr.get(k, (0, 0))[1]))
            c_used, _c_tot = curr.get(k, (0, prev.get(k, (0, 0))[1]))
            delta = c_used - p_used
            parts.append(f"{k}={_format_delta_gib(delta)}")
        logger.mgpu_mm_log(f"{identifier} {tag} - {prev_tag}: " + " | ".join(parts))
    else:
        # Baseline vs zero
        keys = set(curr.keys())
        ordered = ["cpu"] + sorted([k for k in keys if k != "cpu"])
        parts = []
        for k in ordered:
            c_used, _c_tot = curr.get(k, (0, 0))
            parts.append(f"{k}=+{_bytes_to_gib(c_used):.2f}")
        logger.mgpu_mm_log(f"{identifier} {tag} - <baseline>: " + " | ".join(parts))

    # Update last snapshot
    _MEM_SNAPSHOT_LAST[identifier] = (tag, curr)


# ==========================================================================================
# Lifecycle Tracking and Leak Analysis Utilities
# ==========================================================================================

# Track ModelPatcher lifecycle to correlate with CPU RAM trends
if '_MGPU_TRACKED_MODELPATCHERS' not in globals():
    _MGPU_TRACKED_MODELPATCHERS = weakref.WeakSet()

def track_modelpatcher(model_patcher):
    """Registers a ModelPatcher instance for lifecycle tracking."""
    try:
        if isinstance(model_patcher, comfy.model_patcher.ModelPatcher):
            if model_patcher not in _MGPU_TRACKED_MODELPATCHERS:
                _MGPU_TRACKED_MODELPATCHERS.add(model_patcher)
                logger.debug(f"[MultiGPU_Lifecycle] Tracking ModelPatcher {id(model_patcher)} (tracked={len(_MGPU_TRACKED_MODELPATCHERS)})")
    except Exception as e:
        logger.debug(f"[MultiGPU_Lifecycle] track_modelpatcher error: {e}")

def log_tracked_modelpatchers_status(tag="checkpoint"):
    """Logs count and estimated CPU RAM for tracked ModelPatchers."""
    alive_count = len(_MGPU_TRACKED_MODELPATCHERS)
    total_cpu_memory_mb = 0.0
    for patcher in list(_MGPU_TRACKED_MODELPATCHERS):
        try:
            if hasattr(patcher, "model") and patcher.model is not None:
                for param in patcher.model.parameters():
                    if getattr(param, "device", torch.device("cpu")).type == "cpu":
                        total_cpu_memory_mb += (param.nelement() * param.element_size()) / (1024.0 * 1024.0)
        except Exception:
            continue
    logger.warning(f"[MultiGPU_Lifecycle] [{tag}] Tracked ModelPatchers={alive_count}, approx CPU RAM={total_cpu_memory_mb:.2f} MB")

def analyze_cpu_memory_leaks(force=False):
    """Diagnostic: scan referrers of tracked ModelPatchers when memory is high."""
    try:
        vm = psutil.virtual_memory()
        patchers = list(_MGPU_TRACKED_MODELPATCHERS)
        if not force and len(patchers) <= 5 and vm.percent <= 80.0:
            logger.debug(f"[MultiGPU_Leak_Analyzer] Skipping analysis. Patcher count ({len(patchers)}) and memory usage ({vm.percent:.1f}%) normal.")
            return
        logger.warning(f"[MultiGPU_Leak_Analyzer] High pressure (patchers={len(patchers)}, cpu_mem={vm.percent:.1f}%). Inspecting up to 5 referrer sets.")
        for i, patcher in enumerate(patchers[:5]):
                try:
                    referrers = gc.get_referrers(patcher)
                    logger.warning(f"[MultiGPU_Leak_Analyzer] Patcher #{i} id={id(patcher)} referrers={len(referrers)}")
                    for j, ref in enumerate(referrers[:10]):
                        rtype = type(ref).__name__
                        rmod = getattr(type(ref), "__module__", "unknown")
                        if isinstance(ref, dict):
                            logger.warning(f"  Ref {j}: dict(len={len(ref)}) mod={rmod}")
                        elif isinstance(ref, list):
                            logger.warning(f"  Ref {j}: list(len={len(ref)}) mod={rmod}")
                        else:
                            logger.warning(f"  Ref {j}: {rtype} mod={rmod}")
                except Exception:
                    logger.warning("[MultiGPU_Leak_Analyzer] Failed to inspect referrers for a patcher.")
    except Exception as e:
        logger.debug(f"[MultiGPU_Leak_Analyzer] analyze error: {e}")

def try_malloc_trim():
    """Attempt to return freed heap memory to OS (Linux/glibc)."""
    try:
        if platform.system() == "Linux":
            libc = ctypes.CDLL("libc.so.6")
            if hasattr(libc, "malloc_trim"):
                logger.info("[MultiGPU_Memory_Management] malloc_trim(0) begin")
                multigpu_memory_log("mem_mgmt", "pre-malloc-trim")
                res = libc.malloc_trim(0)
                multigpu_memory_log("mem_mgmt", "post-malloc-trim")
                if res == 1:
                    logger.info("[MultiGPU_Memory_Management] malloc_trim(0) released memory")
                else:
                    logger.debug("[MultiGPU_Memory_Management] malloc_trim(0) no release")
    except Exception as e:
        logger.debug(f"[MultiGPU_Memory_Management] malloc_trim error: {e}")

def prune_distorch_stores():
    """Prune stale allocation/settings entries not tied to active models."""
    try:
        multigpu_memory_log("distorch_prune", "start")
        active_hashes_v2 = set()
        active_hashes_v1 = set()
        for lm in getattr(mm, "current_loaded_models", []):
            mp = getattr(lm, "model", None)
            if mp is not None:
                try:
                    h2 = create_safetensor_model_hash(mp, "prune_check_v2")
                    active_hashes_v2.add(h2)
                except Exception:
                    pass
                if create_model_hash is not None:
                    try:
                        h1 = create_model_hash(mp, "prune_check_v1")
                        active_hashes_v1.add(h1)
                    except Exception:
                        pass

        # V1
        if isinstance(model_allocation_store, dict) and active_hashes_v1:
            stale = set(model_allocation_store.keys()) - active_hashes_v1
            if stale:
                logger.info(f"[MultiGPU_Memory_Management] Pruning {len(stale)} DisTorch V1 entries")
                for k in stale:
                    model_allocation_store.pop(k, None)

        # V2
        for store, name in ((safetensor_allocation_store, "allocation"), (safetensor_settings_store, "settings")):
            try:
                if isinstance(store, dict):
                    stale2 = set(store.keys()) - active_hashes_v2
                    if stale2:
                        logger.info(f"[MultiGPU_Memory_Management] Pruning {len(stale2)} V2 {name} entries")
                        for k in stale2:
                            store.pop(k, None)
            except Exception:
                pass
        multigpu_memory_log("distorch_prune", "end")
    except Exception as e:
        logger.debug(f"[MultiGPU_Memory_Management] prune_distorch_stores error: {e}")


# ==========================================================================================
# Model Management Inspection Utilities (End-to-End Tracking)
# ==========================================================================================

def create_model_identifier(model_patcher):
    """Creates a concise, unique identifier for a model patcher based on type and size."""
    if not model_patcher or not model_patcher.model:
        return "N/A (Detached)"

    model = model_patcher.model
    model_type = type(model).__name__

    # Try the fast path first (using size calculated by ModelPatcher)
    try:
        model_size = model_patcher.model_size()
    except Exception:
        model_size = 0

    # If the fast path fails or returns 0, perform a safe deep inspection
    if model_size == 0:
        try:
            # Safely inspect parameters without triggering hooks/loads
            with model_patcher.use_ejected(skip_and_inject_on_exit_only=True):
                 # We must iterate parameters() AND buffers() as both consume memory
                 params = list(model.parameters()) + list(model.buffers())
                 # Use data_ptr to handle potential weight tying/shared tensors correctly
                 seen_tensors = set()
                 for p in params:
                     if p.data_ptr() not in seen_tensors:
                        model_size += p.numel() * p.element_size()
                        seen_tensors.add(p.data_ptr())
        except Exception as e:
            logger.debug(f"[MultiGPU_Inspection] Error during safe size calculation for identifier: {e}")
            return f"{model_type} (ID_Err)"

    # Create a hash based on type and calculated size
    identifier = f"{model_type}_{model_size}"
    model_hash = hashlib.sha256(identifier.encode()).hexdigest()
    return f"{model_type} ({model_hash[:8]})"


def analyze_tensor_locations(model_patcher):
    """
    Analyzes the physical device placement of model tensors (parameters and buffers).
    This provides the Ground Truth location of the data, handling shared weights correctly.
    """
    device_summary = {}
    seen_tensors = set()
    total_memory = 0

    if not model_patcher or not model_patcher.model:
        return {"error": "Model not available"}, 0

    model = model_patcher.model

    # Crucial: Use the ejector to ensure we can access the model weights safely
    # without interfering with injections, hooks, or triggering unintended loads (like in standard LowVRAM mode).
    try:
        with model_patcher.use_ejected(skip_and_inject_on_exit_only=True):
            # Helper to process tensors (parameters or buffers)
            def process_tensor(tensor):
                nonlocal total_memory
                # Use data_ptr() for unique identification of the underlying memory
                if tensor.data_ptr() in seen_tensors:
                    return
                seen_tensors.add(tensor.data_ptr())

                if tensor.numel() > 0:
                    tensor_mem = tensor.numel() * tensor.element_size()
                    total_memory += tensor_mem

                    if hasattr(tensor, 'device'):
                        device = str(tensor.device)
                    else:
                        # Handle cases like NF4 quantization or other custom tensors
                        device = "Unknown/Managed"

                    if device not in device_summary:
                        device_summary[device] = {'tensors': 0, 'memory': 0}

                    device_summary[device]['tensors'] += 1
                    device_summary[device]['memory'] += tensor_mem

            # Iterate over all parameters (weights, biases)
            for param in model.parameters():
                process_tensor(param)

            # Iterate over all buffers (like batch norm running stats)
            for buffer in model.buffers():
                process_tensor(buffer)

    except Exception as e:
        logger.error(f"[MultiGPU_Inspection] Error during tensor location analysis: {e}")
        return {"error": str(e)}, 0

    return device_summary, total_memory


def inspect_model_management_state(context_description=""):
    """
    Provides a detailed, structured overview of the current state of ComfyUI's model management,
    including memory usage across all devices and the status, location, and patching of all loaded models.

    Call this function anywhere in the code to get an immediate snapshot of the system state.
    """

    # Ensure logger configuration (handles calls before full MultiGPU init if needed)
    if not logger.handlers:
        handler = logging.StreamHandler()
        formatter = logging.Formatter('%(message)s')
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        # Default to INFO if log level isn't set by main __init__.py
        if logger.level == logging.NOTSET:
            logger.setLevel(logging.INFO)

    # We inspect the state without forcing GC or cache clearing, which might alter the state we want to observe.

    logger.info("\n" + "=" * 100)
    logger.info(f"  INSPECTION: ComfyUI Model Management State [Context: {context_description}]")
    logger.info("=" * 100)

    # 1. Device Memory Overview
    # Provides context on available resources across the system.
    logger.info("--- [1] System Device Memory Overview (GB) ---")
    # Sys Free: Memory available to the OS. Torch Alloc: Memory reserved by PyTorch (Active + Cache).
    fmt_mem = "{:<12} | {:>10} | {:>10} | {:>10} | {:>15}"
    logger.info(fmt_mem.format("Device", "Total", "Sys Free", "Used", "Torch Alloc"))
    logger.info("-" * 70)

    all_devices = get_device_list()
    # Sort devices for consistent display (CPU last)
    sorted_devices = sorted(all_devices, key=lambda d: (d == 'cpu', d))

    for dev_str in sorted_devices:
        try:
            device = torch.device(dev_str)

            if dev_str == "cpu":
                vm = psutil.virtual_memory()
                mem_total, mem_free_sys, mem_used = vm.total, vm.available, vm.used
                torch_alloc = 0 # Difficult to track accurately for CPU globally
            else:
                # Use ComfyUI's management functions which account for different backends (CUDA, XPU, etc.)
                mem_total = mm.get_total_memory(device)

                # get_free_memory returns (system_free, torch_cache_free)
                free_info = mm.get_free_memory(device, torch_free_too=True)
                if isinstance(free_info, tuple):
                     mem_free_sys = free_info[0]
                else:
                     mem_free_sys = free_info # Fallback for backends that return single value (like MPS)

                mem_used = mem_total - mem_free_sys

                # Determine Torch Allocation (Reserved memory) - Specific checks for known backends
                torch_alloc = 0
                if device.type == 'cuda' and hasattr(torch.cuda, 'memory_stats'):
                    stats = torch.cuda.memory_stats(device)
                    torch_alloc = stats.get('reserved_bytes.all.current', 0)
                elif device.type == 'xpu' and hasattr(torch, 'xpu') and hasattr(torch.xpu, 'memory_stats'):
                    stats = torch.xpu.memory_stats(device)
                    torch_alloc = stats.get('reserved_bytes.all.current', 0)
                elif device.type == 'npu' and hasattr(torch, 'npu') and hasattr(torch.npu, 'memory_stats'):
                     stats = torch.npu.memory_stats(device)
                     torch_alloc = stats.get('reserved_bytes.all.current', 0)
                elif device.type == 'mlu' and hasattr(torch, 'mlu') and hasattr(torch.mlu, 'memory_stats'):
                     stats = torch.mlu.memory_stats(device)
                     torch_alloc = stats.get('reserved_bytes.all.current', 0)
                # MPS, DirectML, CoreX do not always expose detailed reserved memory stats easily.

            logger.info(fmt_mem.format(
                dev_str,
                f"{mem_total / (1024**3):.2f}",
                f"{mem_free_sys / (1024**3):.2f}",
                f"{mem_used / (1024**3):.2f}",
                f"{torch_alloc / (1024**3):.2f}"
            ))
        except Exception as e:
            logger.debug(f"Could not retrieve memory stats for {dev_str}: {e}")

    logger.info("-" * 70)

    # 2. Loaded Models Inspection (Logical and Physical View)
    # mm.current_loaded_models holds the list of models ComfyUI is managing.
    loaded_models = mm.current_loaded_models
    logger.info(f"\n--- [2] Loaded Models Inspection (Count: {len(loaded_models)}) ---")

    if not loaded_models:
        logger.info("No models currently managed by comfy.model_management.")
        logger.info("=" * 100)
        return

    for i, lm in enumerate(loaded_models):
        logger.info(f"\nModel {i+1}/{len(loaded_models)}:")

        # Check lifecycle status
        mp = lm.model # weakref call to ModelPatcher
        if mp is None:
            # ModelPatcher is gone. Check if the underlying model is still alive (potential leak)
            if lm.is_dead() and lm.real_model() is not None:
                 logger.warning(f"  [!] Status: LEAK DETECTED (Patcher GC'd, but underlying model {lm.real_model().__class__.__name__} persists)")
            else:
                 logger.info(f"  Status: Cleaned Up (Patcher and Model GC'd)")
            continue

        model_id = create_model_identifier(mp)
        logger.info(f"  Identifier: {model_id}")
        logger.info(f"  Status: {'Active (In Use)' if lm.currently_used else 'Idle (Cache)'}")

        # A. Logical View (What ComfyUI intends/tracks)
        logger.info("  [A] Logical View (ComfyUI Tracking):")

        # Devices: Target (Compute) vs Offload (Storage)
        logger.info(f"    Devices: Target={lm.device} | Offload={mp.offload_device} | Current (Model.device)={mp.current_loaded_device()}")

        # Memory Footprint
        mem_total = lm.model_memory()
        mem_loaded = lm.model_loaded_memory()
        mem_offloaded = lm.model_offloaded_memory()
        logger.info(f"    Memory (MB): Total={mem_total/(1024**2):.2f} | Loaded (on Target)={mem_loaded/(1024**2):.2f} | Offloaded={mem_offloaded/(1024**2):.2f}")

        # Management Mode (LowVRAM/DisTorch)
        # model_lowvram indicates if ComfyUI is managing this model partially
        is_lowvram = getattr(mp.model, 'model_lowvram', False)
        lowvram_patches_pending = mp.lowvram_patch_counter()
        logger.info(f"    Mode: {'Partial Load (LowVRAM/DisTorch)' if is_lowvram else 'Full Load'}")
        if is_lowvram:
            # This indicates how many weights are being managed by the partial loading system
            logger.info(f"    Weights Managed by LowVRAM/DisTorch System: {lowvram_patches_pending}")

        # Patching (LoRAs, etc.) - Tracking Attach/Detach
        num_weight_patches = len(mp.patches)
        # Check the UUID applied to the actual weights vs the UUID defined in the patcher
        current_weight_uuid = getattr(mp.model, 'current_weight_patches_uuid', None)
        weights_synced = (mp.patches_uuid == current_weight_uuid) and (current_weight_uuid is not None)

        if num_weight_patches > 0:
            status = 'Applied & Synced' if weights_synced else 'Pending/Mismatch (Re-patch needed)'
            logger.info(f"    Patches: {num_weight_patches} weight patches defined | Status: {status}")
            logger.info(f"    UUIDs: Defined={str(mp.patches_uuid)[:8]}... | Applied={str(current_weight_uuid)[:8] if current_weight_uuid else 'None'}...")

        # B. Physical View (Ground Truth Tensor Locations)
        logger.info("  [B] Physical View (Ground Truth Tensor Locations):")
        device_summary, calculated_total_mem = analyze_tensor_locations(mp)

        if "error" in device_summary:
            logger.error(f"    Analysis Error: {device_summary['error']}")
            continue

        if not device_summary:
            logger.info("    No tensors found (e.g., fully offloaded CLIP or utility object).")
        else:
            # Sort devices (CPU last)
            sorted_devices = sorted(device_summary.keys(), key=lambda d: (d.startswith("cpu"), d))
            fmt_loc = "    {:<15} | Tensors: {:>6} | Memory (MB): {:>10.2f} | Percent: {:>6.1f}%"
            for device in sorted_devices:
                data = device_summary[device]
                percent = (data['memory'] / calculated_total_mem) * 100 if calculated_total_mem > 0 else 0
                logger.info(fmt_loc.format(device, data['tensors'], data['memory']/(1024**2), percent))

            # Verification Check
            if abs(calculated_total_mem - mem_total) > (1024*1024): # Allow 1MB difference
                logger.warning(f"    [!] Verification WARNING: Physical memory ({calculated_total_mem/(1024**2):.2f}MB) differs from logical memory ({mem_total/(1024**2):.2f}MB).")

        logger.info("-" * 100)

    logger.info("End of Inspection")
    logger.info("=" * 100)