From 8ed9eecf175b9b6d7c794838fde26d55ed4d5961 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Thu, 13 Aug 2026 23:35:01 +0700 Subject: [PATCH] Add fastsafetensors loader option --- Dockerfile.spark | 2 ++ pyproject.toml | 1 + src/h3_blackwell_runtime/checkpoint.py | 24 ++++++++++++++++++++++++ src/h3_blackwell_runtime/vae_decoder.py | 17 ++++++++++++++++- tools/direct_t2v_preview.py | 2 +- 5 files changed, 44 insertions(+), 2 deletions(-) diff --git a/Dockerfile.spark b/Dockerfile.spark index 9345e90..83bbde1 100644 --- a/Dockerfile.spark +++ b/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)" diff --git a/pyproject.toml b/pyproject.toml index 7a604ad..a26172a 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/h3_blackwell_runtime/checkpoint.py b/src/h3_blackwell_runtime/checkpoint.py index 720cd6f..044b7a8 100644 --- a/src/h3_blackwell_runtime/checkpoint.py +++ b/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: diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index e3a5447..62677ff 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/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: diff --git a/tools/direct_t2v_preview.py b/tools/direct_t2v_preview.py index 2d5b35a..e62b0ca 100644 --- a/tools/direct_t2v_preview.py +++ b/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")