Dockerfile.spark

@@ -11,8 +11,6 @@ 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,7 +5,6 @@ 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

@@ -1,6 +1,5 @@
 """Lazy loading for the current Comfy-format H3 safetensors checkpoint."""
 
-import os
 from pathlib import Path
 
 import torch
@@ -14,42 +13,8 @@ class H3Checkpoint:
     def __init__(self, path: str | Path, device: str | torch.device = "cuda"):
         self.path = str(path)
         self.device = str(device)
-        self._no_mmap_tensors: dict[str, torch.Tensor] | None = None
-
-    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
-
-            with open(self.path, "rb") as file:
-                tensors = load(file.read())
-            self._no_mmap_tensors = {name: value.to(self.device) for name, value in tensors.items()}
-        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
-
         from safetensors import safe_open
 
         with safe_open(self.path, framework="pt", device=self.device) as checkpoint:
@@ -59,21 +24,6 @@ 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:
-                name = f"{prefix}.{suffix}"
-                if name in available_tensors:
-                    tensors[name] = available_tensors[name]
-            return load_nvfp4_linear(tensors, prefix, output_dtype=output_dtype)
-
         from safetensors import safe_open
 
         with safe_open(self.path, framework="pt", device=self.device) as checkpoint:

src/h3_blackwell_runtime/vae_decoder.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import math
-import os
 from pathlib import Path
 
 import torch
@@ -165,34 +164,6 @@ 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_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:
-                available_weights = load(file.read())
-            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].to(device) for name in expected}
-        else:
         with safe_open(str(path), framework="pt", device=str(device)) as checkpoint:
             available = set(checkpoint.keys())
             missing = sorted(set(expected) - available)

tools/direct_t2v_preview.py

@@ -1,7 +1,6 @@
 """Generate a minimal direct, video-only H3 T2V preview without ComfyUI."""
 
 import argparse
-import os
 from pathlib import Path
 import subprocess
 import warnings
@@ -30,38 +29,16 @@ parser.add_argument("--steps", type=int, default=12)
 parser.add_argument("--seed", type=int, default=440204)
 parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
 parser.add_argument("--model-timesteps-capture", type=Path, help="Directory containing captured input_XX.pt H3 timesteps for strict parity checks.")
-parser.add_argument("--profile-memory", action="store_true")
 args = parser.parse_args()
 
-
-def report_memory(stage: str) -> None:
-    if not args.profile_memory:
-        return
-    rss_kb = 0
-    try:
-        with open("/proc/self/status", encoding="utf-8") as file:
-            for line in file:
-                if line.startswith("VmRSS:"):
-                    rss_kb = int(line.split()[1])
-                    break
-    except OSError:
-        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), "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")
 conditioner = Qwen3VLPromptConditioner(
     "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors",
     "/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer",
 )
-report_memory("qwen_loaded")
 video, audio, frames = random_av_latents(args.width, args.height, args.frames, args.seed)
 model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend=args.attention).eval()
-report_memory("h3_loaded")
 text = H3TokenRefiner(checkpoint, attention_backend=args.attention)(conditioner(args.prompt))
-report_memory("text_conditioned")
 model_timesteps = None
 if args.model_timesteps_capture is not None:
     model_timesteps = [
@@ -69,14 +46,9 @@ if args.model_timesteps_capture is not None:
         for index in range(args.steps)
     ]
 latent = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps, model_timesteps=model_timesteps)
-report_memory("sampled")
 vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
-report_memory("vae_loaded")
+pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames]
-with torch.inference_mode():
-    pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames]
-report_memory("vae_decoded")
 pixels = ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()
-report_memory("pixels_cpu")
 args.output.parent.mkdir(parents=True, exist_ok=True)
 raw = args.output.with_suffix(".rgb")
 pixels.numpy().tofile(raw)