flashattn

app.py

@@ -1,3 +1,5 @@
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', shell=True)
 import argparse
 import spaces
 from visualcloze import VisualClozeModel

flux/math.py

@@ -2,8 +2,8 @@ from einops import rearrange
 import torch
 from torch import Tensor
 import torch.nn.functional as F
-# from flash_attn import flash_attn_varlen_func
-# from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
+from flash_attn import flash_attn_varlen_func
+from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
 
 
 def _upad_input(query_layer, key_layer, value_layer, query_mask, key_mask, query_length):