unifiedefficientloader 0.5.2__tar.gz → 0.5.4__tar.gz

{unifiedefficientloader-0.5.2 → unifiedefficientloader-0.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unifiedefficientloader
-Version: 0.5.2
+Version: 0.5.4
 Summary: A unified interface for memory efficient per tensor loading of safetensors files as raw bytes from offset, handling CPU/GPU pinned transfers, and converting between tensors and dicts.
 Author: silveroxides
 License: MIT License
@@ -33,7 +33,7 @@ Classifier: Operating System :: POSIX :: Linux
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: comfy-aimdo>=0.4.5
+Requires-Dist: comfy-aimdo>=0.5.1
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
 Provides-Extra: safetensors

{unifiedefficientloader-0.5.2 → unifiedefficientloader-0.5.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "unifiedefficientloader"
-version = "0.5.2"
+version = "0.5.4"
 description = "A unified interface for memory efficient per tensor loading of safetensors files as raw bytes from offset, handling CPU/GPU pinned transfers, and converting between tensors and dicts."
 readme = "README.md"
 authors = [
@@ -19,7 +19,7 @@ classifiers = [
     "Operating System :: POSIX :: Linux",
 ]
 requires-python = ">=3.9"
-dependencies = ["comfy-aimdo>=0.4.5"]
+dependencies = ["comfy-aimdo>=0.5.1"]
 
 [project.optional-dependencies]
 torch = ["torch"]

unifiedefficientloader-0.5.4/tests/test_low_bit_dtypes.py

@@ -0,0 +1,83 @@
+import json
+import struct
+
+import pytest
+import torch
+from safetensors.torch import load_file, save_file
+
+from unifiedefficientloader import IncrementalSafetensorsWriter, UnifiedSafetensorsLoader
+from unifiedefficientloader.tensor_utils import get_dtype_size, st_shape_to_torch_shape, st_to_torch_dtype, torch_shape_to_st_shape, torch_to_st_dtype
+
+LOW_BIT_DTYPES = [
+    ("float8_e4m3fn", "F8_E4M3"),
+    ("float8_e5m2", "F8_E5M2"),
+    ("float8_e8m0fnu", "F8_E8M0"),
+    ("float4_e2m1fn_x2", "F4"),
+]
+
+
+def _dtype_or_skip(name):
+    dtype = getattr(torch, name, None)
+    if dtype is None:
+        pytest.skip(f"Installed PyTorch has no torch.{name}")
+    return dtype
+
+
+def _payload(path, key="low"):
+    with open(path, "rb") as source:
+        header_size = struct.unpack("<Q", source.read(8))[0]
+        header = json.loads(source.read(header_size).decode("utf-8"))
+        start, end = header[key]["data_offsets"]
+        source.seek(8 + header_size + start)
+        return header[key], source.read(end - start)
+
+
+@pytest.mark.parametrize(("torch_name", "storage_code"), LOW_BIT_DTYPES)
+def test_low_bit_dtype_mappings(torch_name, storage_code):
+    dtype = _dtype_or_skip(torch_name)
+    assert torch_to_st_dtype(dtype) == storage_code
+    assert st_to_torch_dtype(storage_code) == dtype
+    assert get_dtype_size(storage_code) == 1
+
+
+def test_packed_fp4_shape_conversion():
+    assert torch_shape_to_st_shape((2, 3, 4), "F4") == [2, 3, 8]
+    assert st_shape_to_torch_shape((2, 3, 8), "F4") == (2, 3, 4)
+    with pytest.raises(ValueError, match="at least one dimension"):
+        torch_shape_to_st_shape((), "F4")
+    with pytest.raises(ValueError, match="Invalid packed F4"):
+        st_shape_to_torch_shape((2, 3, 7), "F4")
+
+
+@pytest.mark.parametrize(("torch_name", "storage_code"), LOW_BIT_DTYPES)
+@pytest.mark.parametrize("write_mode", ["single", "batch"])
+def test_low_bit_streaming_and_writer_round_trip(tmp_path, torch_name, storage_code, write_mode):
+    dtype = _dtype_or_skip(torch_name)
+    source_path = tmp_path / f"source_{storage_code}.safetensors"
+    output_path = tmp_path / f"output_{storage_code}_{write_mode}.safetensors"
+    original = torch.zeros((2, 4), dtype=dtype)
+    try:
+        save_file({"low": original}, str(source_path))
+    except (KeyError, RuntimeError, TypeError, ValueError):
+        pytest.skip(f"Installed safetensors cannot write {dtype}")
+
+    with UnifiedSafetensorsLoader(str(source_path), low_memory=True) as loader:
+        assert loader.get_dtype("low") == dtype
+        assert loader.get_shape("low") == original.shape
+        loaded = loader.get_tensor("low")
+        assert loaded.dtype == dtype
+        assert loaded.shape == original.shape
+        with IncrementalSafetensorsWriter(str(output_path)) as writer:
+            if write_mode == "single":
+                writer.write("low", loaded)
+            else:
+                writer.write_batch([("low", loaded)])
+
+    source_header, source_payload = _payload(source_path)
+    output_header, output_payload = _payload(output_path)
+    assert source_header["dtype"] == output_header["dtype"] == storage_code
+    assert source_header["shape"] == output_header["shape"]
+    assert source_payload == output_payload
+    round_tripped = load_file(str(output_path))["low"]
+    assert round_tripped.dtype == dtype
+    assert round_tripped.shape == original.shape

{unifiedefficientloader-0.5.2 → unifiedefficientloader-0.5.4}/tests/test_mmap.py

@@ -106,6 +106,29 @@ def test_mmap_get_tensor_values(sample_file):
             assert torch.equal(t.clone(), expected), f"{key}: value mismatch"
 
 
+@uel_required
+def test_mmap_exposes_file_handle(sample_file):
+    """The compatibility wrapper exposes comfy-aimdo's file handle API."""
+    path, _ = sample_file
+
+    with UnifiedSafetensorsLoader(path, low_memory=True, use_mmap=True) as loader:
+        file_handle = loader._mmap.get_file_handle()
+
+    assert isinstance(file_handle, int)
+    assert file_handle > 0
+
+
+def test_control_init_accepts_aimdo_051_options():
+    """The control wrapper forwards comfy-aimdo 0.5.1 init options."""
+    from unifiedefficientloader.uel import control
+
+    assert control.init(
+        implementation=None,
+        simple_vram_headroom=None,
+        nvml_pressure=False,
+    )
+
+
 @uel_required
 def test_mmap_get_tensor_mmap_ref_attached(sample_file):
     """

{unifiedefficientloader-0.5.2 → unifiedefficientloader-0.5.4}/unifiedefficientloader/incremental_writer.py

@@ -16,7 +16,7 @@ import ctypes
 from concurrent.futures import ThreadPoolExecutor
 
 from . import logging_utils
-from .tensor_utils import torch_to_st_dtype, get_dtype_size
+from .tensor_utils import get_dtype_size, torch_shape_to_st_shape, torch_to_st_dtype
 logger = logging_utils.get_logger(__name__)
 
 def _ensure_torch():
@@ -211,8 +211,9 @@ class IncrementalSafetensorsWriter:
         prepared = []
         for name, tensor in batch:
             st_dtype = torch_to_st_dtype(tensor.dtype)
-            shape = list(tensor.shape)
-            num_elements = math.prod(shape) if shape else 1
+            tensor_shape = list(tensor.shape)
+            shape = torch_shape_to_st_shape(tensor.shape, st_dtype)
+            num_elements = math.prod(tensor_shape) if tensor_shape else 1
             byte_size = num_elements * get_dtype_size(st_dtype)
             prepared.append((name, st_dtype, shape, byte_size, tensor))
 
@@ -270,9 +271,10 @@ class IncrementalSafetensorsWriter:
 
         # Extract metadata only — no data movement in calling thread
         st_dtype = torch_to_st_dtype(tensor.dtype)
-        shape = list(tensor.shape)
+        tensor_shape = list(tensor.shape)
+        shape = torch_shape_to_st_shape(tensor.shape, st_dtype)
         
-        num_elements = math.prod(shape) if shape else 1
+        num_elements = math.prod(tensor_shape) if tensor_shape else 1
         byte_size = num_elements * get_dtype_size(st_dtype)
         
         # Record manifest entry
@@ -293,4 +295,4 @@ class IncrementalSafetensorsWriter:
 
         self._semaphore.acquire()
         future = self._executor.submit(self._batch_worker, [(absolute_offset, tensor)])
-        self._futures.append(future)
+        self._futures.append(future)

{unifiedefficientloader-0.5.2 → unifiedefficientloader-0.5.4}/unifiedefficientloader/memory_efficient_loader.py

@@ -11,7 +11,7 @@ import struct
 from typing import Dict, Optional, Tuple
 
 from . import logging_utils
-from .tensor_utils import st_to_torch_dtype
+from .tensor_utils import st_shape_to_torch_shape, st_to_torch_dtype
 
 logger = logging_utils.get_logger(__name__)
 
@@ -212,7 +212,8 @@ class UnifiedSafetensorsLoader:
         if self.low_memory:
             if key not in self._header:
                 raise KeyError(f"Tensor '{key}' not found in file")
-            return tuple(self._header[key]["shape"])
+            metadata = self._header[key]
+            return st_shape_to_torch_shape(metadata["shape"], metadata["dtype"])
         else:
             return tuple(self._tensors[key].shape)
 
@@ -265,9 +266,8 @@ class UnifiedSafetensorsLoader:
                         "ignore", message="The given buffer is not writable"
                     )
                     dtype = st_to_torch_dtype(metadata["dtype"])
-                    tensor = torch.frombuffer(tensor_view, dtype=dtype).view(
-                        metadata["shape"]
-                    )
+                    shape = st_shape_to_torch_shape(metadata["shape"], metadata["dtype"])
+                    tensor = torch.frombuffer(tensor_view, dtype=dtype).view(shape)
                     storage = tensor.untyped_storage()
                     setattr(storage, "_uel_mmap_ref", self._mmap)
                     return tensor
@@ -311,7 +311,7 @@ class UnifiedSafetensorsLoader:
         """Deserialize raw bytes into a torch tensor."""
         torch = _ensure_torch()
         dtype_str = metadata["dtype"]
-        shape = metadata["shape"]
+        shape = st_shape_to_torch_shape(metadata["shape"], dtype_str)
         dtype = st_to_torch_dtype(dtype_str)
 
         if tensor_bytes is None:

{unifiedefficientloader-0.5.2 → unifiedefficientloader-0.5.4}/unifiedefficientloader/tensor_utils.py

@@ -76,6 +76,10 @@ def torch_to_st_dtype(dtype) -> str:
         mapping[torch.float8_e5m2] = "F8_E5M2"
     if hasattr(torch, "float8_e4m3fn"):
         mapping[torch.float8_e4m3fn] = "F8_E4M3"
+    if hasattr(torch, "float8_e8m0fnu"):
+        mapping[torch.float8_e8m0fnu] = "F8_E8M0"
+    if hasattr(torch, "float4_e2m1fn_x2"):
+        mapping[torch.float4_e2m1fn_x2] = "F4"
     if hasattr(torch, "uint64"):
         mapping[torch.uint64] = "U64"
     if hasattr(torch, "uint32"):
@@ -104,6 +108,8 @@ def st_to_torch_dtype(dtype_str: str):
         "C64": torch.complex64,
         "F8_E5M2": getattr(torch, "float8_e5m2", None),
         "F8_E4M3": getattr(torch, "float8_e4m3fn", None),
+        "F8_E8M0": getattr(torch, "float8_e8m0fnu", None),
+        "F4": getattr(torch, "float4_e2m1fn_x2", None),
         "U64": getattr(torch, "uint64", None),
         "U32": getattr(torch, "uint32", None),
         "U16": getattr(torch, "uint16", None),
@@ -121,6 +127,28 @@ def get_dtype_size(st_dtype: str) -> int:
         "I64": 8, "I32": 4, "I16": 2, "I8": 1, "U8": 1,
         "U64": 8, "U32": 4, "U16": 2,
         "BOOL": 1, "C64": 8,
-        "F8_E5M2": 1, "F8_E4M3": 1,
+        "F8_E5M2": 1, "F8_E4M3": 1, "F8_E8M0": 1, "F4": 1,
     }
     return sizes[st_dtype]
+
+
+def torch_shape_to_st_shape(shape, st_dtype: str) -> list:
+    """Convert a Torch tensor shape to its safetensors header shape."""
+    shape = list(shape)
+    if st_dtype != "F4":
+        return shape
+    if not shape:
+        raise ValueError("Packed F4 tensors must have at least one dimension")
+    shape[-1] *= 2
+    return shape
+
+
+def st_shape_to_torch_shape(shape, st_dtype: str) -> tuple:
+    """Convert a safetensors header shape to its Torch tensor shape."""
+    shape = list(shape)
+    if st_dtype != "F4":
+        return tuple(shape)
+    if not shape or shape[-1] % 2:
+        raise ValueError(f"Invalid packed F4 safetensors shape: {shape}")
+    shape[-1] //= 2
+    return tuple(shape)

{unifiedefficientloader-0.5.2 → unifiedefficientloader-0.5.4}/unifiedefficientloader/uel/model_mmap.py

@@ -23,6 +23,9 @@ def _setup_argtypes(lib):
     lib.model_mmap_get.argtypes = [ctypes.c_void_p]
     lib.model_mmap_get.restype = ctypes.c_void_p
 
+    lib.model_mmap_get_file_handle.argtypes = [ctypes.c_void_p]
+    lib.model_mmap_get_file_handle.restype = ctypes.c_uint64
+
     lib.model_mmap_bounce.argtypes = [ctypes.c_void_p]
     lib.model_mmap_bounce.restype = ctypes.c_bool
 
@@ -53,6 +56,9 @@ class ModelMMAP:
     def get(self):
         return _ctrl.lib.model_mmap_get(self.state)
 
+    def get_file_handle(self):
+        return int(_ctrl.lib.model_mmap_get_file_handle(self.state))
+
     def bounce(self):
         return bool(_ctrl.lib.model_mmap_bounce(self.state))
 

{unifiedefficientloader-0.5.2 → unifiedefficientloader-0.5.4}/unifiedefficientloader.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unifiedefficientloader
-Version: 0.5.2
+Version: 0.5.4
 Summary: A unified interface for memory efficient per tensor loading of safetensors files as raw bytes from offset, handling CPU/GPU pinned transfers, and converting between tensors and dicts.
 Author: silveroxides
 License: MIT License
@@ -33,7 +33,7 @@ Classifier: Operating System :: POSIX :: Linux
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: comfy-aimdo>=0.4.5
+Requires-Dist: comfy-aimdo>=0.5.1
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
 Provides-Extra: safetensors

{unifiedefficientloader-0.5.2 → unifiedefficientloader-0.5.4}/unifiedefficientloader.egg-info/SOURCES.txt

@@ -5,6 +5,7 @@ setup.py
 tests/test_direct_gpu.py
 tests/test_incremental_writer.py
 tests/test_logging.py
+tests/test_low_bit_dtypes.py
 tests/test_mmap.py
 tests/test_mmap_diagnostics.py
 tests/test_unified_data_loader.py

{unifiedefficientloader-0.5.2 → unifiedefficientloader-0.5.4}/unifiedefficientloader.egg-info/requires.txt

@@ -1,4 +1,4 @@
-comfy-aimdo>=0.4.5
+comfy-aimdo>=0.5.1
 
 [all]
 torch