Refactor data_mean and data_std initialization

Refactor data_mean and data_std to use register_buffer and remove nn.Buffer.

Tested with pytorch version: 2.4.0+cu121
xformers version: 0.0.27.post2
Set vram state to: LOW_VRAM
Device: cuda:0 NVIDIA GeForce RTX 3090 : cudaMallocAsync
Enabled pinned memory 122281.0
Using xformers attention
Python version: 3.10.12 (main, Aug 15 2025, 14:32:43) [GCC 11.4.0]
ComfyUI version: 0.3.68
ComfyUI frontend version: 1.28.8

Ovi/vae/vae.py

@@ -75,14 +75,21 @@ class VAE(nn.Module):
         super().__init__()
 
         if data_dim == 80:
-            self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_80D, dtype=torch.float32))
-            self.data_std = nn.Buffer(torch.tensor(DATA_STD_80D, dtype=torch.float32))
+            data_mean = torch.tensor(DATA_MEAN_80D, dtype=torch.float32)
+            data_std  = torch.tensor(DATA_STD_80D, dtype=torch.float32)
         elif data_dim == 128:
-            self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32))
-            self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32))
+            data_mean = torch.tensor(DATA_MEAN_128D, dtype=torch.float32)
+            data_std  = torch.tensor(DATA_STD_128D, dtype=torch.float32)
+        else:
+            raise ValueError(f"Unsupported data_dim={data_dim}, expected 80 or 128")
 
-        self.data_mean = self.data_mean.view(1, -1, 1)
-        self.data_std = self.data_std.view(1, -1, 1)
+        # match old shape: (1, channels, 1)
+        data_mean = data_mean.view(1, -1, 1)
+        data_std  = data_std.view(1, -1, 1)
+
+        # register as buffers so they move with .to(device) / .cuda()
+        self.register_buffer("data_mean", data_mean)
+        self.register_buffer("data_std", data_std)
 
         self.encoder = Encoder1D(
             dim=hidden_dim,