Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -10,14 +11,6 @@ tokenizer = AutoTokenizer.from_pretrained(
     trust_remote_code=True
 )
 
-print("Loading model...")
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    trust_remote_code=True
-)
-
 SYSTEM_PROMPT = """You are CodeSentinel.
 
 You are an expert AI software engineer.
@@ -34,7 +27,33 @@ Capabilities:
 Never generate malware or unsafe code.
 """
 
+# Global model variable
+model = None
+
+
+def load_model():
+    """Load the model only once after GPU allocation."""
+    global model
+
+    if model is None:
+        print("Loading model...")
+
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True,
+        )
+
+        model.eval()
+
+    return model
+
+
+@spaces.GPU
 def generate(message, history):
+    model = load_model()
+
     messages = [
         {
             "role": "system",
@@ -78,6 +97,7 @@ def generate(message, history):
 
     return response.strip()
 
+
 demo = gr.ChatInterface(
     fn=generate,
     title="🛡️ CodeSentinel",
@@ -91,6 +111,7 @@ demo = gr.ChatInterface(
     ],
 )
 
+
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",