From 15cdc3eab3f22677245fc81637503fa2e88412d9 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 17 Jun 2024 19:47:34 +0300
Subject: [PATCH] Fix flash_attn fallback

---
 lumina_models/nextdit.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/lumina_models/nextdit.py b/lumina_models/nextdit.py
index 69da3fe..72280dc 100644
--- a/lumina_models/nextdit.py
+++ b/lumina_models/nextdit.py
@@ -26,7 +26,6 @@ import torch.nn.functional as F
 from .components import RMSNorm
 
 import comfy.model_management
-
 import comfy.ops
 ops = comfy.ops.manual_cast
 
@@ -365,7 +364,10 @@ class Attention(nn.Module):
             # end var_len_flash_attn
 
         else:
-            raise Exception("Flash attention (flash_attn) is not available and currently required for Lumina-next -models.")
+            n_rep = self.n_local_heads // self.n_local_kv_heads
+            if n_rep >= 1:
+                xk = xk.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3)
+                xv = xv.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3)
             output = (
                 F.scaled_dot_product_attention(
                     xq.permute(0, 2, 1, 3),