Upload 150 files

source/unet_hacked.py

@@ -527,7 +527,14 @@ class MemoryEfficientCrossAttention(nn.Module):
         )
 
         # actually compute the attention, what we cannot get enough of
-        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
+        try:
+            out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
+        except (NotImplementedError, RuntimeError):
+            # Fallback to standard attention for CPU or unsupported configs
+            scale = self.dim_head ** -0.5
+            attn_weights = torch.bmm(q * scale, k.transpose(-2, -1))
+            attn_weights = torch.softmax(attn_weights, dim=-1)
+            out = torch.bmm(attn_weights, v)
 
         if exists(mask):
             raise NotImplementedError

source/vae_hacked.py

@@ -260,7 +260,14 @@ class MemoryEfficientAttnBlock(nn.Module):
             .contiguous(),
             (q, k, v),
         )
-        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
+        try:
+            out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op)
+        except (NotImplementedError, RuntimeError):
+            # Fallback to standard attention for CPU or unsupported configs
+            scale = C ** -0.5
+            attn_weights = torch.bmm(q * scale, k.transpose(-2, -1))
+            attn_weights = torch.softmax(attn_weights, dim=-1)
+            out = torch.bmm(attn_weights, v)
 
         out = (
             out.unsqueeze(0)