From e043e9af827f5cfad661324d2129da1acc702cef Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Fri, 14 Aug 2026 00:41:20 +0700 Subject: [PATCH] Use Comfy attention for VAE when available --- src/h3_blackwell_runtime/vae_decoder.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index b4afc8d..81aa42d 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/src/h3_blackwell_runtime/vae_decoder.py @@ -104,8 +104,14 @@ class Attention(nn.Module): except Exception: 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) - output = F.scaled_dot_product_attention(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)) + query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) + 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):