From aa9f4749587c0f8a5041a56bcc4e4a07ca76c4f0 Mon Sep 17 00:00:00 2001 From: jamesjjcondon <33615047+jamesjjcondon@users.noreply.github.com> Date: Fri, 14 Nov 2025 11:57:26 +1030 Subject: [PATCH] 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 | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/Ovi/vae/vae.py b/Ovi/vae/vae.py index d1b9254..0e0d207 100644 --- a/Ovi/vae/vae.py +++ b/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,