Update app.py

app.py

@@ -8,8 +8,9 @@ import os
 from threading import Thread
 
 # Load model directly
+device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer = AutoTokenizer.from_pretrained("Navid-AI/Mulhem-1-Mini", token=os.getenv("HF_TOKEN"))
-model = AutoModelForCausalLM.from_pretrained("Navid-AI/Mulhem-1-Mini", torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="flash_attention_2", token=os.getenv("HF_TOKEN"))
+model = AutoModelForCausalLM.from_pretrained("Navid-AI/Mulhem-1-Mini", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", token=os.getenv("HF_TOKEN")).to(device)
 streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
 
 def respond(
@@ -30,7 +31,7 @@ def respond(
             messages.append({"role": "assistant", "content": val[1]})
 
     messages.append({"role": "user", "content": message})
-    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True, enable_reasoning=enable_reasoning)
+    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True, enable_reasoning=enable_reasoning).to(device)
     generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_tokens, temperature=temperature, top_p=top_p)
     thread = Thread(target=model.generate, kwargs=generation_kwargs)