h3-blackwell-runtime/src/h3_blackwell_runtime/checkpoint.py

"""Lazy loading for the current Comfy-format H3 safetensors checkpoint."""

import os
from pathlib import Path

import torch

from .nvfp4 import Nvfp4Linear, load_nvfp4_linear


class H3Checkpoint:
    """Load individual tensors/modules without materializing the whole checkpoint."""

    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 _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._fast_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:
            value = checkpoint.get_tensor(name)
        return value.to(dtype=dtype) if dtype is not None else value

    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:
            available = set(checkpoint.keys())
            for suffix in names:
                name = f"{prefix}.{suffix}"
                if name in available:
                    tensors[name] = checkpoint.get_tensor(name)
        return load_nvfp4_linear(tensors, prefix, output_dtype=output_dtype)