daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Mirror upstream VAE block math

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 00:50:03 +07:00
parent e043e9af82
commit e4f145e9d5
1 changed files with 6 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
src/h3_blackwell_runtime/vae_decoder.py
View file

@ -67,7 +67,7 @@ class FeedForward(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
gate, value = self.w1(x).chunk(2, dim=-1) gate, value = self.w1(x).chunk(2, dim=-1)
return self.w2(F.silu(gate) * value) return self.w2(F.silu(gate).mul_(value))
def _apply_rope_split_half(x: torch.Tensor, table: torch.Tensor) -> torch.Tensor: def _apply_rope_split_half(x: torch.Tensor, table: torch.Tensor) -> torch.Tensor:
@ -94,14 +94,6 @@ class Attention(nn.Module):
batch, sequence, _ = x.shape batch, sequence, _ = x.shape
qkv = self.to_qkv(x).view(batch, sequence, self.heads, 3 * self.dim_head) qkv = self.to_qkv(x).view(batch, sequence, self.heads, 3 * self.dim_head)
query, key, value = qkv.chunk(3, dim=-1) query, key, value = qkv.chunk(3, dim=-1)
try:
import comfy_kitchen # Registers the standalone CUDA extension operators.
del comfy_kitchen
query, key = query.contiguous(), key.contiguous()
weight = torch.ones(self.dim_head, device=query.device, dtype=query.dtype)
torch.ops.comfy_kitchen.rms_rope_split_half_(query, key, rotary_pos_emb, weight, weight, self.norm_q.eps, rotary_pos_emb.shape[-3] * 2)
except Exception:
query, key = self.norm_q(query), self.norm_k(key) query, key = self.norm_q(query), self.norm_k(key)
query, key = _apply_rope_split_half(query, rotary_pos_emb), _apply_rope_split_half(key, rotary_pos_emb) query, key = _apply_rope_split_half(query, rotary_pos_emb), _apply_rope_split_half(key, rotary_pos_emb)
query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
@ -126,8 +118,8 @@ class TransformerBlock(nn.Module):
self.scale2 = nn.Parameter(torch.empty(dim, device=device)) self.scale2 = nn.Parameter(torch.empty(dim, device=device))
def forward(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor) -> torch.Tensor:
x = x + self.attn(self.norm1(x), rotary_pos_emb) * self.scale1.to(x.dtype) x = x.addcmul_(self.attn(self.norm1(x), rotary_pos_emb), self.scale1.to(x.dtype))
return x + self.ff(self.norm2(x)) * self.scale2.to(x.dtype) return x.addcmul_(self.ff(self.norm2(x)), self.scale2.to(x.dtype))
class ViT3DDecoder(nn.Module): class ViT3DDecoder(nn.Module):
@ -156,7 +148,7 @@ class ViT3DDecoder(nn.Module):
h = block(h, rope) h = block(h, rope)
output = self.proj_out(self.norm_out(h))[:, :patches] output = self.proj_out(self.norm_out(h))[:, :patches]
output = output.view(batch, latent_t, latent_h, latent_w, self.out_channels, self.patch_size_t, self.patch_size, self.patch_size) output = output.view(batch, latent_t, latent_h, latent_w, self.out_channels, self.patch_size_t, self.patch_size, self.patch_size)
return output.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(batch, self.out_channels, latent_t * self.patch_size_t, latent_h * self.patch_size, latent_w * self.patch_size) return output.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous().reshape(batch, self.out_channels, latent_t * self.patch_size_t, latent_h * self.patch_size, latent_w * self.patch_size)
class MiniMaxH3VideoVAE(nn.Module): class MiniMaxH3VideoVAE(nn.Module):