From bc0cff88ce082d6b3d3cb5e7229519f34d389790 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Thu, 13 Aug 2026 23:41:25 +0700 Subject: [PATCH] Normalize fastsafetensors CUDA device --- src/h3_blackwell_runtime/checkpoint.py | 5 ++++- src/h3_blackwell_runtime/vae_decoder.py | 3 ++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/h3_blackwell_runtime/checkpoint.py b/src/h3_blackwell_runtime/checkpoint.py index 044b7a8..4ddd308 100644 --- a/src/h3_blackwell_runtime/checkpoint.py +++ b/src/h3_blackwell_runtime/checkpoint.py @@ -22,11 +22,14 @@ class H3Checkpoint: def _use_fast_safetensors(self) -> bool: return os.getenv("H3_FAST_SAFETENSORS", "").lower() in {"1", "true", "yes", "on"} + def _fast_device(self) -> str: + return "cuda:0" if self.device == "cuda" else self.device + 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: + with fastsafe_open(filenames=[self.path], nogds=True, device=self._fast_device()) as checkpoint: self._no_mmap_tensors = { name: checkpoint.get_tensor(name).clone().detach() for name in checkpoint.get_keys() diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index 62677ff..39dba9c 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/src/h3_blackwell_runtime/vae_decoder.py @@ -168,7 +168,8 @@ class MiniMaxH3VideoVAE(nn.Module): 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: + fast_device = "cuda:0" if str(device) == "cuda" else str(device) + with fastsafe_open(filenames=[str(path)], nogds=True, device=fast_device) as checkpoint: available_weights = { name: checkpoint.get_tensor(name).clone().detach() for name in checkpoint.get_keys()