Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
-import torch
 import spaces
+import torch
 
 from transformers import (
     AutoTokenizer,
@@ -10,60 +10,74 @@ from transformers import (
 from peft import PeftModel
 
 
-# ==========================================
+# ============================================================
 # MODEL
-# ==========================================
+# ============================================================
 
 BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
 
 ADAPTER_MODEL = "dd253B/DhanushAI-0.5B"
 
 
-# ==========================================
-# LOAD TOKENIZER
-# ==========================================
+# ============================================================
+# GLOBAL VARIABLES
+# ============================================================
 
-print("Loading tokenizer...")
+tokenizer = None
+model = None
 
-tokenizer = AutoTokenizer.from_pretrained(
-    BASE_MODEL
-)
 
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+# ============================================================
+# LOAD MODEL
+# ============================================================
 
+def load_model():
 
-# ==========================================
-# LOAD MODEL
-# ==========================================
+    global tokenizer
+    global model
 
-print("Loading base model...")
+    if model is not None:
+        return
 
-model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
+    print("Loading tokenizer...")
 
-print("Loading DhanushAI adapter...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        BASE_MODEL
+    )
 
-model = PeftModel.from_pretrained(
-    model,
-    ADAPTER_MODEL
-)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
 
-model.eval()
+    print("Loading base model...")
 
-print("DhanushAI loaded!")
+    base_model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        torch_dtype=torch.float16
+    )
 
+    print("Loading DhanushAI adapter...")
 
-# ==========================================
-# GPU FUNCTION
-# ==========================================
+    model = PeftModel.from_pretrained(
+        base_model,
+        ADAPTER_MODEL
+    )
+
+    model = model.to("cuda")
+
+    model.eval()
+
+    print("DhanushAI loaded successfully!")
+
+
+# ============================================================
+# CHAT
+# ============================================================
 
 @spaces.GPU
 def chat(message):
 
+    load_model()
+
     if not message or not message.strip():
         return "Please enter a message."
 
@@ -78,9 +92,8 @@ Assistant:"""
         return_tensors="pt"
     )
 
-    # Move inputs to the same device as model
     inputs = {
-        key: value.to(model.device)
+        key: value.to("cuda")
         for key, value in inputs.items()
     }
 
@@ -119,11 +132,12 @@ Assistant:"""
     return answer
 
 
-# ==========================================
+# ============================================================
 # GRADIO
-# ==========================================
+# ============================================================
 
 demo = gr.Interface(
+
     fn=chat,
 
     inputs=gr.Textbox(