From 9c9e0857c290df66ff129caf8c4cc236629e3e30 Mon Sep 17 00:00:00 2001
From: Daniel Maddern <daniel.maddern@illuvium.io>
Date: Fri, 14 Aug 2026 00:52:57 +0700
Subject: [PATCH] Load direct VAE weights as float32

---
 src/h3_blackwell_runtime/vae_decoder.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py
index ecaced5..bed1238 100644
--- a/src/h3_blackwell_runtime/vae_decoder.py
+++ b/src/h3_blackwell_runtime/vae_decoder.py
@@ -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)
 
     @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)
         expected = model.state_dict()
         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:
                 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))
-            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"}:
             from safetensors.torch import load
 
@@ -198,7 +198,7 @@ class MiniMaxH3VideoVAE(nn.Module):
             if missing or shape_errors:
                 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))
-            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:
             with safe_open(str(path), framework="pt", device=str(device)) as checkpoint:
                 available = set(checkpoint.keys())
@@ -207,7 +207,7 @@ class MiniMaxH3VideoVAE(nn.Module):
                 if missing or shape_errors:
                     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))
-                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)
         return model