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

Use Comfy attention for VAE when available

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 00:41:20 +07:00
parent e96af86e72
commit e043e9af82
1 changed files with 8 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
src/h3_blackwell_runtime/vae_decoder.py
View file

@ -104,8 +104,14 @@ class Attention(nn.Module):
except Exception: 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)
output = F.scaled_dot_product_attention(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)
return self.to_out(output.transpose(1, 2).reshape(batch, sequence, -1).nan_to_num_(0.0)) try:
from comfy.ldm.modules.attention import optimized_attention
output = optimized_attention(query, key, value, self.heads, skip_reshape=True)
except Exception:
output = F.scaled_dot_product_attention(query, key, value).transpose(1, 2)
return self.to_out(output.reshape(batch, sequence, -1).nan_to_num_(0.0))
class TransformerBlock(nn.Module): class TransformerBlock(nn.Module):