Disable sageattention for SAM3 (Comfy-Org#13529)

Causes Nans

Author: kijai

comfy/ldm/sam3/detector.py

@@ -54,7 +54,7 @@ def forward(self, q_input, k_input=None, v_input=None, mask=None):
         if mask is not None and mask.ndim == 2:
             mask = mask[:, None, None, :]  # [B, T] -> [B, 1, 1, T] for SDPA broadcast
         dtype = q.dtype  # manual_cast may produce mixed dtypes
-        out = optimized_attention(q, k.to(dtype), v.to(dtype), self.num_heads, mask=mask)
+        out = optimized_attention(q, k.to(dtype), v.to(dtype), self.num_heads, mask=mask, low_precision_attention=False)
         return self.out_proj(out)
 
 

comfy/ldm/sam3/sam.py

@@ -40,7 +40,7 @@ def forward(self, q, k, v):
         q = self.q_proj(q)
         k = self.k_proj(k)
         v = self.v_proj(v)
-        return self.out_proj(optimized_attention(q, k, v, self.num_heads))
+        return self.out_proj(optimized_attention(q, k, v, self.num_heads, low_precision_attention=False))
 
 
 class TwoWayAttentionBlock(nn.Module):
@@ -179,7 +179,7 @@ def forward(self, x, freqs_cis=None):
         q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(dim=0)
         if self.use_rope and freqs_cis is not None:
             q, k = apply_rope(q, k, freqs_cis)
-        return self.proj(optimized_attention(q, k, v, self.num_heads, skip_reshape=True))
+        return self.proj(optimized_attention(q, k, v, self.num_heads, skip_reshape=True, low_precision_attention=False))
 
 
 class Block(nn.Module):

comfy/ldm/sam3/tracker.py

@@ -364,7 +364,7 @@ def forward(self, q, k=None, v=None, rope=None, num_k_exclude_rope=0):
         v = self.v_proj(v)
         if rope is not None:
             q, k = apply_rope_memory(q, k, rope, self.num_heads, num_k_exclude_rope)
-        out = optimized_attention(q, k, v, self.num_heads)
+        out = optimized_attention(q, k, v, self.num_heads, low_precision_attention=False)
         return self.out_proj(out)
 
 
@@ -657,7 +657,7 @@ def forward(self, image, x, memory_image, memory, memory_image_pos=None,
         v = self.self_attn_v_proj(normed)
         if rope is not None:
             q, k = apply_rope_memory(q, k, rope, self.num_heads, 0)
-        x = x + self.self_attn_out_proj(optimized_attention(q, k, v, self.num_heads))
+        x = x + self.self_attn_out_proj(optimized_attention(q, k, v, self.num_heads, low_precision_attention=False))
 
         # Decoupled cross-attention: fuse image and memory projections
         normed = self.norm2(x)
@@ -668,7 +668,7 @@ def forward(self, image, x, memory_image, memory, memory_image_pos=None,
         v = self.cross_attn_v_proj(memory)
         if rope is not None:
             q, k = apply_rope_memory(q, k, rope, self.num_heads, num_k_exclude_rope)
-        x = x + self.cross_attn_out_proj(optimized_attention(q, k, v, self.num_heads))
+        x = x + self.cross_attn_out_proj(optimized_attention(q, k, v, self.num_heads, low_precision_attention=False))
 
         # FFN
         x = x + self.linear2(F.gelu(self.linear1(self.norm3(x))))