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

Load direct VAE weights as float32

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 00:52:57 +07:00
parent e4f145e9d5
commit 9c9e0857c2
1 changed files with 4 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
src/h3_blackwell_runtime/vae_decoder.py
View file

@ -168,7 +168,7 @@ class MiniMaxH3VideoVAE(nn.Module):
self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD, device=device).view(1, 3, 1, 1, 1), persistent=False) self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD, device=device).view(1, 3, 1, 1, 1), persistent=False)
@classmethod @classmethod
def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True) -> "MiniMaxH3VideoVAE": def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True, dtype: torch.dtype = torch.float32) -> "MiniMaxH3VideoVAE":
model = cls(device="meta", tiling=tiling) model = cls(device="meta", tiling=tiling)
expected = model.state_dict() expected = model.state_dict()
if os.getenv("H3_FAST_SAFETENSORS", "").lower() in {"1", "true", "yes", "on"}: if os.getenv("H3_FAST_SAFETENSORS", "").lower() in {"1", "true", "yes", "on"}:
@ -186,7 +186,7 @@ class MiniMaxH3VideoVAE(nn.Module):
if missing or shape_errors: if missing or shape_errors:
details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else []) 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)) raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details))
weights = {name: available_weights[name] for name in expected} weights = {name: available_weights[name].to(device=device, dtype=dtype) for name in expected}
elif os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}: elif os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}:
from safetensors.torch import load from safetensors.torch import load
@ -198,7 +198,7 @@ class MiniMaxH3VideoVAE(nn.Module):
if missing or shape_errors: if missing or shape_errors:
details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else []) 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)) raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details))
weights = {name: available_weights[name].to(device) for name in expected} weights = {name: available_weights[name].to(device=device, dtype=dtype) for name in expected}
else: else:
with safe_open(str(path), framework="pt", device=str(device)) as checkpoint: with safe_open(str(path), framework="pt", device=str(device)) as checkpoint:
available = set(checkpoint.keys()) available = set(checkpoint.keys())
@ -207,7 +207,7 @@ class MiniMaxH3VideoVAE(nn.Module):
if missing or shape_errors: if missing or shape_errors:
details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else []) 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)) raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details))
weights = {name: checkpoint.get_tensor(name) for name in expected} weights = {name: checkpoint.get_tensor(name).to(dtype=dtype) for name in expected}
model.load_state_dict(weights, strict=True, assign=True) model.load_state_dict(weights, strict=True, assign=True)
return model return model