Update app.py

app.py

@@ -1,13 +1,20 @@
+import os
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
+# ---- Safety: prevent VRAM fragmentation ----
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
 MODEL_ID = "deepseek-ai/DeepSeek-R1"
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print("Using device:", device)
+# ---- HARD GPU CHECK ----
+if not torch.cuda.is_available():
+    raise RuntimeError("❌ GPU not detected. Please enable GPU hardware in HF Space settings.")
+
+print("✅ Using GPU:", torch.cuda.get_device_name(0))
 
-# 4-bit config (24GB VRAM optimized)
+# ---- 4bit quant config (24GB optimized) ----
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.float16,
@@ -15,22 +22,25 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_quant_type="nf4",
 )
 
-# Load tokenizer
+# ---- Load tokenizer ----
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID,
     trust_remote_code=True
 )
 
-# Load model
+# ---- Load model ----
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     quantization_config=bnb_config,
     device_map="auto",
-    trust_remote_code=True
+    trust_remote_code=True,
+    attn_implementation="flash_attention_2"
 )
 
 model.eval()
 
+
+# ---- Chat Function ----
 def chat_fn(message, history):
     messages = []
 
@@ -65,9 +75,10 @@ def chat_fn(message, history):
     return response
 
 
+# ---- Gradio UI ----
 demo = gr.ChatInterface(
     fn=chat_fn,
-    title="DeepSeek-R1 32B (4bit) - 24GB VRAM",
+    title="DeepSeek-R1 32B (4bit) - 24GB GPU",
     chatbot=gr.Chatbot(height=500),
 )