Add fastsafetensors loader option

Dockerfile.spark

@@ -11,6 +11,8 @@ RUN python -m pip install --no-cache-dir --no-deps /tmp/wheels/sageattn3-*.whl \
 
 RUN python -m pip install --no-cache-dir --no-deps comfy-kitchen==0.2.28
 
+RUN python -m pip install --no-cache-dir "fastsafetensors>=0.1.10"
+
 RUN python -m pip install --no-cache-dir --no-deps -e . \
     && python -c "import comfy_kitchen, torch; from sageattn3 import sageattn3_blackwell; assert hasattr(torch.ops.comfy_kitchen, 'rms_rope_split_half_'); print(torch.__version__, torch.version.cuda)"
 

pyproject.toml

@@ -5,6 +5,7 @@ description = "Direct MiniMax H3 Blackwell inference research runtime"
 requires-python = ">=3.12"
 dependencies = [
   "comfy-kitchen==0.2.28",
+  "fastsafetensors>=0.1.10",
   "safetensors>=0.5.0",
   "torch==2.9.1+cu130",
   "transformers>=4.51,<5"

src/h3_blackwell_runtime/checkpoint.py

@@ -19,6 +19,20 @@ class H3Checkpoint:
     def _disable_mmap(self) -> bool:
         return os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}
 
+    def _use_fast_safetensors(self) -> bool:
+        return os.getenv("H3_FAST_SAFETENSORS", "").lower() in {"1", "true", "yes", "on"}
+
+    def _all_tensors_fast_safetensors(self) -> dict[str, torch.Tensor]:
+        if self._no_mmap_tensors is None:
+            from fastsafetensors import fastsafe_open
+
+            with fastsafe_open(filenames=[self.path], nogds=True, device=self.device) as checkpoint:
+                self._no_mmap_tensors = {
+                    name: checkpoint.get_tensor(name).clone().detach()
+                    for name in checkpoint.get_keys()
+                }
+        return self._no_mmap_tensors
+
     def _all_tensors_no_mmap(self) -> dict[str, torch.Tensor]:
         if self._no_mmap_tensors is None:
             from safetensors.torch import load
@@ -29,6 +43,9 @@ class H3Checkpoint:
         return self._no_mmap_tensors
 
     def tensor(self, name: str, *, dtype: torch.dtype | None = None) -> torch.Tensor:
+        if self._use_fast_safetensors():
+            value = self._all_tensors_fast_safetensors()[name]
+            return value.to(dtype=dtype) if dtype is not None else value
         if self._disable_mmap():
             value = self._all_tensors_no_mmap()[name]
             return value.to(dtype=dtype) if dtype is not None else value
@@ -42,6 +59,13 @@ class H3Checkpoint:
     def nvfp4_linear(self, prefix: str, *, output_dtype=torch.bfloat16) -> Nvfp4Linear:
         names = ("comfy_quant", "weight", "weight_scale", "weight_scale_2", "bias", "pre_quant_scale")
         tensors = {}
+        if self._use_fast_safetensors():
+            available_tensors = self._all_tensors_fast_safetensors()
+            for suffix in names:
+                name = f"{prefix}.{suffix}"
+                if name in available_tensors:
+                    tensors[name] = available_tensors[name]
+            return load_nvfp4_linear(tensors, prefix, output_dtype=output_dtype)
         if self._disable_mmap():
             available_tensors = self._all_tensors_no_mmap()
             for suffix in names:

src/h3_blackwell_runtime/vae_decoder.py

@@ -165,7 +165,22 @@ class MiniMaxH3VideoVAE(nn.Module):
     def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True) -> "MiniMaxH3VideoVAE":
         model = cls(device="meta", tiling=tiling)
         expected = model.state_dict()
-        if os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}:
+        if os.getenv("H3_FAST_SAFETENSORS", "").lower() in {"1", "true", "yes", "on"}:
+            from fastsafetensors import fastsafe_open
+
+            with fastsafe_open(filenames=[str(path)], nogds=True, device=str(device)) as checkpoint:
+                available_weights = {
+                    name: checkpoint.get_tensor(name).clone().detach()
+                    for name in checkpoint.get_keys()
+                }
+            available = set(available_weights)
+            missing = sorted(set(expected) - available)
+            shape_errors = [(name, tuple(expected[name].shape), tuple(available_weights[name].shape)) for name in expected if name in available and tuple(expected[name].shape) != tuple(available_weights[name].shape)]
+            if missing or shape_errors:
+                details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else [])
+                raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details))
+            weights = {name: available_weights[name] for name in expected}
+        elif os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}:
             from safetensors.torch import load
 
             with open(path, "rb") as file:

tools/direct_t2v_preview.py

@@ -48,7 +48,7 @@ def report_memory(stage: str) -> None:
         pass
     cuda_alloc = torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0.0
     cuda_reserved = torch.cuda.memory_reserved() / 1024**3 if torch.cuda.is_available() else 0.0
-    print({"stage": stage, "rss_gb": round(rss_kb / 1024**2, 3), "cuda_alloc_gb": round(cuda_alloc, 3), "cuda_reserved_gb": round(cuda_reserved, 3), "disable_mmap": os.getenv("H3_DISABLE_MMAP", "")}, flush=True)
+    print({"stage": stage, "rss_gb": round(rss_kb / 1024**2, 3), "cuda_alloc_gb": round(cuda_alloc, 3), "cuda_reserved_gb": round(cuda_reserved, 3), "fast_safetensors": os.getenv("H3_FAST_SAFETENSORS", ""), "disable_mmap": os.getenv("H3_DISABLE_MMAP", "")}, flush=True)
 
 report_memory("start")
 checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")