Update app.py

app.py

@@ -2,10 +2,6 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import subprocess
-
-# Install flash-attn
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 # Load the model and tokenizer
 model_name = "akjindal53244/Llama-3.1-Storm-8B"
@@ -13,7 +9,6 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.bfloat16,
-    use_flash_attention_2=True,
     device_map="auto"
 )
 
@@ -25,7 +20,7 @@ def generate_text(prompt, max_length, temperature):
     ]
     formatted_prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
     
-    inputs = tokenizer(formatted_prompt, return_tensors="pt").to("cuda")
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
     
     outputs = model.generate(
         **inputs,