silveroxides-ComfyUI-ModelUtils/nodes/merger_utils.py

"""
Memory-efficient safetensors utilities with parallel I/O and pinned memory support.

Based on benchmarking, parallel reading with ThreadPoolExecutor provides 2-4x speedup
for large models compared to sequential reads.

Worker count is automatically optimized based on device capabilities when workers=None.
Pinned memory enables faster CPU→GPU transfers (2-3x improvement).
"""
import os
import mmap
import json
import struct
import torch
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List, Optional, Any, Tuple

# Import ComfyUI's memory pinning functions if available
try:
    from comfy.model_management import pin_memory, unpin_memory, is_device_cuda
    PINNING_AVAILABLE = True
except ImportError:
    PINNING_AVAILABLE = False
    def pin_memory(tensor): return False
    def unpin_memory(tensor): return False
    def is_device_cuda(device): return str(device).startswith('cuda')


class PinnedTensor:
    """Context manager for automatic pinned memory management.

    Usage:
        with PinnedTensor(cpu_tensor) as pinned:
            gpu_tensor = pinned.to('cuda')
        # Memory automatically unpinned after transfer
    """

    def __init__(self, tensor: torch.Tensor, auto_pin: bool = True):
        self.tensor = tensor
        self.is_pinned = False
        self.auto_pin = auto_pin and PINNING_AVAILABLE

    def __enter__(self) -> torch.Tensor:
        if self.auto_pin and self.tensor.device.type == 'cpu':
            self.is_pinned = pin_memory(self.tensor)
        return self.tensor

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.is_pinned:
            unpin_memory(self.tensor)
            self.is_pinned = False


def transfer_to_gpu_pinned(
    tensor: torch.Tensor,
    device: str = 'cuda',
    dtype: Optional[torch.dtype] = None
) -> torch.Tensor:
    """Transfer tensor to GPU using pinned memory for faster transfer.

    Args:
        tensor: CPU tensor to transfer
        device: Target GPU device
        dtype: Optional dtype conversion

    Returns:
        Tensor on GPU device
    """
    if not PINNING_AVAILABLE or not is_device_cuda(torch.device(device)):
        # Fall back to regular transfer
        if dtype:
            return tensor.to(device=device, dtype=dtype)
        return tensor.to(device=device)

    # Pin memory for faster transfer
    was_pinned = pin_memory(tensor)
    try:
        if dtype:
            result = tensor.to(device=device, dtype=dtype, non_blocking=True)
        else:
            result = tensor.to(device=device, non_blocking=True)
        # Sync to ensure transfer is complete before unpinning
        if was_pinned:
            torch.cuda.current_stream().synchronize()
        return result
    finally:
        if was_pinned:
            unpin_memory(tensor)


class MemoryEfficientSafeOpen:
    """Memory-efficient safetensors file reader with mmap, parallel I/O, and pinned memory support.

    Features:
    - mmap mode: Zero-copy tensor access via memory-mapped file
    - low_memory mode: Direct file reads to minimize OS page cache usage
    - Parallel loading: Multi-threaded tensor reads for 2-4x speedup
    - Sorted batch reads: Keys sorted by file offset for sequential I/O
    - Auto-optimized workers: Adjusts parallelism based on device capabilities
    - Pinned memory: Faster CPU→GPU transfers when transferring to CUDA

    Args:
        filename: Path to safetensors file
        device: Target device (default 'cpu')
        mmap_mode: Use memory-mapped file for zero-copy (default True)
        low_memory: Use direct file reads to minimize memory (overrides mmap_mode)
    """

    def __init__(self, filename: str, device: str = 'cpu', mmap_mode: bool = True, low_memory: bool = False):
        self.filename = filename
        self.device = device
        # low_memory mode forces mmap off to avoid OS page caching
        self.low_memory = low_memory
        self.mmap_mode = mmap_mode and not low_memory
        self.header, self.header_size = self._read_header()
        self.file = open(filename, "rb")
        self.mmap_obj = None

        if self.mmap_mode:
            self.mmap_obj = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_READ)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.mmap_obj:
            self.mmap_obj.close()
        self.file.close()


    def keys(self) -> List[str]:
        """Return all tensor keys (excluding metadata)."""
        return [k for k in self.header.keys() if k != "__metadata__"]

    def metadata(self) -> Dict[str, str]:
        """Return file metadata."""
        return self.header.get("__metadata__", {})

    def keys_sorted_by_offset(self) -> List[str]:
        """Return keys sorted by file offset for optimal sequential I/O."""
        keys_with_offsets = []
        for key in self.keys():
            offset = self.header[key]["data_offsets"][0]
            keys_with_offsets.append((key, offset))
        keys_with_offsets.sort(key=lambda x: x[1])
        return [k for k, _ in keys_with_offsets]

    def get_tensor(self, key: str) -> torch.Tensor:
        """Load a single tensor by key."""
        if key not in self.header:
            raise KeyError(f"Tensor '{key}' not found in the file")

        metadata = self.header[key]
        offset_start, offset_end = metadata["data_offsets"]

        if self.mmap_mode and self.mmap_obj:
            if offset_start != offset_end:
                start = self.header_size + 8 + offset_start
                end = self.header_size + 8 + offset_end
                tensor_bytes = memoryview(self.mmap_obj)[start:end]
            else:
                tensor_bytes = None
        else:
            # Non-mmap mode: use pre-allocated bytearray with readinto for minimal copies
            tensor_bytes = None
            if offset_start != offset_end:
                self.file.seek(self.header_size + 8 + offset_start)
                # Pre-allocate writable buffer and read directly into it
                tensor_bytes = bytearray(offset_end - offset_start)
                self.file.readinto(tensor_bytes)

        return self._deserialize_tensor(tensor_bytes, metadata)


    def get_tensor_to_gpu(
        self,
        key: str,
        device: str = 'cuda',
        dtype: Optional[torch.dtype] = None
    ) -> torch.Tensor:
        """Load a tensor and transfer to GPU using pinned memory for faster transfer.

        Args:
            key: Tensor key to load
            device: Target GPU device (default 'cuda')
            dtype: Optional dtype conversion

        Returns:
            Tensor on GPU device
        """
        cpu_tensor = self.get_tensor(key)
        return transfer_to_gpu_pinned(cpu_tensor, device, dtype)

    def get_tensors_parallel(
        self,
        keys: Optional[List[str]] = None,
        workers: Optional[int] = None,
        sort_by_offset: bool = True
    ) -> Dict[str, torch.Tensor]:
        """Load multiple tensors in parallel using ThreadPoolExecutor.

        Args:
            keys: List of keys to load (default: all keys)
            workers: Number of worker threads (default: auto-calculated based on device)
            sort_by_offset: Sort keys by file offset before loading (default: True)

        Returns:
            Dict mapping key -> tensor
        """
        if keys is None:
            keys = self.keys()

        # Auto-calculate workers based on device capabilities
        if workers is None:
            try:
                from .device_utils import get_optimal_workers, estimate_model_size
                model_size = estimate_model_size(self.filename)
                workers = get_optimal_workers(model_size)
            except ImportError:
                workers = min(4, os.cpu_count() or 4)

        if sort_by_offset:
            # Sort by offset for better I/O patterns
            keys_with_offsets = [(k, self.header[k]["data_offsets"][0]) for k in keys]
            keys_with_offsets.sort(key=lambda x: x[1])
            keys = [k for k, _ in keys_with_offsets]

        result = {}

        with ThreadPoolExecutor(max_workers=workers) as executor:
            futures = {
                executor.submit(
                    _read_tensor_from_file,
                    self.filename,
                    key,
                    self.header[key],
                    self.header_size
                ): key
                for key in keys
            }

            for future in as_completed(futures):
                key = futures[future]
                try:
                    tensor = future.result()
                    if tensor is not None:
                        result[key] = tensor
                except Exception as e:
                    print(f"[MemoryEfficientSafeOpen] Failed to load {key}: {e}")

        return result

    def get_tensor_as_dict(self, key: str) -> Dict[str, Any]:
        """Load a tensor and decode as JSON dict (for config tensors)."""
        if key not in self.header:
            raise KeyError(f"Tensor '{key}' not found")

        metadata = self.header[key]
        offset_start, offset_end = metadata["data_offsets"]

        if offset_start == offset_end:
            return {}

        if self.mmap_mode and self.mmap_obj:
            start = self.header_size + 8 + offset_start
            end = self.header_size + 8 + offset_end
            tensor_bytes = bytes(self.mmap_obj[start:end])
        else:
            self.file.seek(self.header_size + 8 + offset_start)
            tensor_bytes = self.file.read(offset_end - offset_start)

        return json.loads(tensor_bytes.decode("utf-8"))

    def _read_header(self):
        with open(self.filename, "rb") as f:
            header_size = struct.unpack("<Q", f.read(8))[0]
            header_json = f.read(header_size).decode("utf-8")
        return json.loads(header_json), header_size

    def _deserialize_tensor(self, tensor_bytes, metadata):
        dtype_str = metadata["dtype"]
        shape = metadata["shape"]
        dtype = self._get_torch_dtype(dtype_str)

        if tensor_bytes is None:
            byte_tensor = torch.empty(0, dtype=torch.uint8)
        else:
            # Avoid extra copy if already a bytearray (low_memory mode)
            if isinstance(tensor_bytes, bytearray):
                byte_tensor = torch.frombuffer(tensor_bytes, dtype=torch.uint8)
            else:
                # mmap memoryview needs copy to create writable tensor
                byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)

        if dtype_str in ["F8_E5M2", "F8_E4M3"]:
            return self._convert_float8(byte_tensor, dtype_str, shape)

        return byte_tensor.view(dtype).reshape(shape)


    @staticmethod
    def _get_torch_dtype(dtype_str):
        dtype_map = {
            "F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
            "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
            "U8": torch.uint8, "BOOL": torch.bool,
        }
        if hasattr(torch, "float8_e5m2"):
            dtype_map["F8_E5M2"] = torch.float8_e5m2
        if hasattr(torch, "float8_e4m3fn"):
            dtype_map["F8_E4M3"] = torch.float8_e4m3fn

        dtype = dtype_map.get(dtype_str)
        if dtype is None:
            raise ValueError(f"Unsupported dtype: {dtype_str}")
        return dtype

    @staticmethod
    def _convert_float8(byte_tensor, dtype_str, shape):
        if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
            return byte_tensor.view(torch.float8_e5m2).reshape(shape)
        elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
            return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)
        else:
            raise ValueError(f"Unsupported float8 type: {dtype_str}. Your PyTorch version may be too old.")


def _read_tensor_from_file(
    filepath: str,
    key: str,
    metadata: Dict,
    header_size: int
) -> Optional[torch.Tensor]:
    """Helper function to read a single tensor from file (for parallel execution).

    Each worker opens its own file handle for thread safety.
    """
    offset_start, offset_end = metadata["data_offsets"]
    dtype_str = metadata["dtype"]
    shape = metadata["shape"]

    if offset_start == offset_end:
        return None

    dtype_map = {
        "F64": torch.float64, "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
        "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
        "U8": torch.uint8, "BOOL": torch.bool,
    }
    if hasattr(torch, "float8_e5m2"):
        dtype_map["F8_E5M2"] = torch.float8_e5m2
    if hasattr(torch, "float8_e4m3fn"):
        dtype_map["F8_E4M3"] = torch.float8_e4m3fn

    dtype = dtype_map.get(dtype_str, torch.float32)

    with open(filepath, "rb") as f:
        f.seek(header_size + 8 + offset_start)
        tensor_bytes = f.read(offset_end - offset_start)

    byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)

    if dtype_str in ["F8_E5M2", "F8_E4M3"]:
        if dtype_str == "F8_E5M2" and hasattr(torch, "float8_e5m2"):
            return byte_tensor.view(torch.float8_e5m2).reshape(shape)
        elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"):
            return byte_tensor.view(torch.float8_e4m3fn).reshape(shape)

    return byte_tensor.view(dtype).reshape(shape)


def get_union_keys(*handlers) -> List[str]:
    """Returns the union of all tensor keys across multiple model handlers."""
    all_keys = set()
    for handler in handlers:
        if handler is not None:
            all_keys.update(handler.keys())
    return sorted(all_keys)


def get_intersection_keys(*handlers) -> List[str]:
    """Returns only keys present in ALL model handlers."""
    handlers = [h for h in handlers if h is not None]
    if not handlers:
        return []
    all_keys = set(handlers[0].keys())
    for handler in handlers[1:]:
        all_keys &= set(handler.keys())
    return sorted(all_keys)