zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-05 10:00:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(minimax): cast raw parameters to input device in H3 VAEs (#15268)

Browse Source
This commit is contained in:
rivadart
2026-08-03 19:06:12 -04:00
committed by GitHub
parent e377e26304
commit 16e3f3034f
2 changed files with 8 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
comfy/ldm/minimax/audio_vae.py
View File

@@ -35,7 +35,8 @@ class Snake1d(nn.Module):
self.alpha = nn.Parameter(torch.empty(1, channels, 1)) self.alpha = nn.Parameter(torch.empty(1, channels, 1))
def forward(self, x): 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): class SnakeBeta(nn.Module):
@@ -47,8 +48,8 @@ class SnakeBeta(nn.Module):
self.beta = nn.Parameter(torch.empty(in_features)) self.beta = nn.Parameter(torch.empty(in_features))
def forward(self, x): def forward(self, x):
alpha = torch.exp(self.alpha).view(1, -1, 1) alpha = torch.exp(comfy.ops.cast_to_input(self.alpha, x)).view(1, -1, 1)
beta = torch.exp(self.beta).view(1, -1, 1) beta = torch.exp(comfy.ops.cast_to_input(self.beta, x)).view(1, -1, 1)
return snake(x, alpha, beta) return snake(x, alpha, beta)
@@ -239,7 +240,7 @@ class CausalAttention(nn.Module):
def forward(self, x): def forward(self, x):
B, N, C = x.shape B, N, C = x.shape
weight, _, offload_stream = comfy.ops.cast_bias_weight(self.qkv, x, offloadable=True) 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) 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) q, k, v = qkv.reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)

6
comfy/ldm/minimax/vae.py
View File

@@ -253,8 +253,8 @@ class TransformerBlock(nn.Module):
self.scale2 = nn.Parameter(torch.empty(dim)) self.scale2 = nn.Parameter(torch.empty(dim))
def forward(self, x, rotary_pos_emb=None): 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) 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)), self.scale2) 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): class ViT3DDecoder(nn.Module):
@@ -289,7 +289,7 @@ class ViT3DDecoder(nn.Module):
num_patches = h.shape[1] num_patches = h.shape[1]
num_suffix = 1 + self.num_register_tokens 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) 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) suffix_ids = torch.zeros((B, num_suffix, 3), device=x.device, dtype=img_ids.dtype)