From df423cb9fa4914b120697daa1d814a5145b582c4 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 18 Jun 2024 11:07:15 +0300
Subject: [PATCH] Update nextdit.py

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

diff --git a/lumina_models/nextdit.py b/lumina_models/nextdit.py
index 72280dc..84b422c 100644
--- a/lumina_models/nextdit.py
+++ b/lumina_models/nextdit.py
@@ -17,8 +17,10 @@ try:
     from flash_attn import flash_attn_varlen_func
     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
     FLASH_ATTN_AVAILABLE = True
+    print("Flash Attention is available")
 except:
     FLASH_ATTN_AVAILABLE = False
+    print("LuminaWrapper: WARNING! Flash Attention is not available, using much slower torch SDP attention")
 import torch
 import torch.nn as nn
 import torch.nn.functional as F