Update app.py

app.py

@@ -6,6 +6,9 @@ import gradio as gr
 from threading import Thread
 from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
 
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
 parser = argparse.ArgumentParser()
 
 if torch.cuda.is_available():
@@ -17,7 +20,8 @@ model_id = "vikhyatk/moondream2"
 revision = "2024-04-02"
 tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
 moondream = AutoModelForCausalLM.from_pretrained(
-    model_id, trust_remote_code=True, revision=revision
+    model_id, trust_remote_code=True, revision=revision,
+    attn_implementation="flash_attention_2"
 ).to(device=device, dtype=dtype)
 moondream.eval()