Update app.py

app.py

@@ -1,6 +1,11 @@
+import subprocess
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, LlamaForCausalLM
+import spaces
+
+# Install flash-attn with specific environment variable
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 # Initialize model and tokenizer
 model_id = 'akjindal53244/Llama-3.1-Storm-8B'
@@ -29,7 +34,7 @@ def generate_response(message, history):
     messages.append({"role": "user", "content": message})
 
     prompt = format_prompt(messages)
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
     generated_ids = model.generate(input_ids, max_new_tokens=256, temperature=0.7, do_sample=True, eos_token_id=tokenizer.eos_token_id)
     response = tokenizer.decode(generated_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
     return response.strip()