From 16e3f3034f2bba1fff6c70cbd759339778555cd6 Mon Sep 17 00:00:00 2001 From: rivadart <67722383+rivadart@users.noreply.github.com> Date: Mon, 3 Aug 2026 19:06:12 -0400 Subject: [PATCH] fix(minimax): cast raw parameters to input device in H3 VAEs (#15268) --- comfy/ldm/minimax/audio_vae.py | 9 +++++---- comfy/ldm/minimax/vae.py | 6 +++--- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/comfy/ldm/minimax/audio_vae.py b/comfy/ldm/minimax/audio_vae.py index 033ae6966..a1be63d54 100644 --- a/comfy/ldm/minimax/audio_vae.py +++ b/comfy/ldm/minimax/audio_vae.py @@ -35,7 +35,8 @@ class Snake1d(nn.Module): self.alpha = nn.Parameter(torch.empty(1, channels, 1)) def forward(self, x): - return snake(x, self.alpha, self.alpha) + alpha = comfy.ops.cast_to_input(self.alpha, x) + return snake(x, alpha, alpha) class SnakeBeta(nn.Module): @@ -47,8 +48,8 @@ class SnakeBeta(nn.Module): self.beta = nn.Parameter(torch.empty(in_features)) def forward(self, x): - alpha = torch.exp(self.alpha).view(1, -1, 1) - beta = torch.exp(self.beta).view(1, -1, 1) + alpha = torch.exp(comfy.ops.cast_to_input(self.alpha, x)).view(1, -1, 1) + beta = torch.exp(comfy.ops.cast_to_input(self.beta, x)).view(1, -1, 1) return snake(x, alpha, beta) @@ -239,7 +240,7 @@ class CausalAttention(nn.Module): def forward(self, x): B, N, C = x.shape weight, _, offload_stream = comfy.ops.cast_bias_weight(self.qkv, x, offloadable=True) - qkv = F.linear(x, weight=weight, bias=torch.cat((self.q_bias, self.zero_k_bias, self.v_bias))) + qkv = F.linear(x, weight=weight, bias=comfy.ops.cast_to_input(torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)), x)) comfy.ops.uncast_bias_weight(self.qkv, weight, None, offload_stream) q, k, v = qkv.reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0) diff --git a/comfy/ldm/minimax/vae.py b/comfy/ldm/minimax/vae.py index aeb3421a2..b03bf0c9a 100644 --- a/comfy/ldm/minimax/vae.py +++ b/comfy/ldm/minimax/vae.py @@ -253,8 +253,8 @@ class TransformerBlock(nn.Module): self.scale2 = nn.Parameter(torch.empty(dim)) def forward(self, x, rotary_pos_emb=None): - x = x.addcmul_(self.attn(comfy.rmsnorm.rms_norm(x, self.norm1.weight, self.norm1.eps), rotary_pos_emb), self.scale1) - return x.addcmul_(self.ff(comfy.rmsnorm.rms_norm(x, self.norm2.weight, self.norm2.eps)), self.scale2) + x = x.addcmul_(self.attn(comfy.rmsnorm.rms_norm(x, self.norm1.weight, self.norm1.eps), rotary_pos_emb), comfy.ops.cast_to_input(self.scale1, x)) + return x.addcmul_(self.ff(comfy.rmsnorm.rms_norm(x, self.norm2.weight, self.norm2.eps)), comfy.ops.cast_to_input(self.scale2, x)) class ViT3DDecoder(nn.Module): @@ -289,7 +289,7 @@ class ViT3DDecoder(nn.Module): num_patches = h.shape[1] num_suffix = 1 + self.num_register_tokens - h = torch.cat([h, self.register_tokens.expand(B, -1, -1), torch.zeros_like(h[:, 0:1, :])], dim=1) + h = torch.cat([h, comfy.ops.cast_to_input(self.register_tokens, h).expand(B, -1, -1), torch.zeros_like(h[:, 0:1, :])], dim=1) img_ids = create_token_ids((latent_T, latent_H, latent_W), x.device, x.dtype).expand(B, -1, -1) suffix_ids = torch.zeros((B, num_suffix, 3), device=x.device, dtype=img_ids.dtype)