daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Normalize fastsafetensors CUDA device

Browse source
This commit is contained in:
Daniel Maddern 2026-08-13 23:41:25 +07:00
parent 8ed9eecf17
commit bc0cff88ce
2 changed files with 6 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
src/h3_blackwell_runtime/checkpoint.py
View file

@ -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()

3
src/h3_blackwell_runtime/vae_decoder.py
View file

@ -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()