36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
|
"""Lazy loading for the current Comfy-format H3 safetensors checkpoint."""
|
||
|
|
|
||
|
|
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)
|
||
|
|
|
||
|
|
def tensor(self, name: str, *, dtype: torch.dtype | None = None) -> torch.Tensor:
|
||
|
|
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")
|
||
|
|
tensors = {}
|
||
|
|
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)
|