Update app.py

app.py

@@ -15,6 +15,7 @@ import torch
 import cv2
 import mediapipe as mp
 import base64
+import gc
 
 
 # Set page config
@@ -350,14 +351,18 @@ st.title("NeuraSense AI")
 
 # Set up DialoGPT model
 @st.cache_resource
-def load_model():
-   tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-   model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
-   return tokenizer, model
-
+def load_tokenizer():
+    return AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
 
+@st.cache_resource
+def load_model():
+    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium", 
+                                                 device_map="auto", 
+                                                 torch_dtype=torch.float16)
+    return model
 
-tokenizer, model = load_model()
+tokenizer = load_tokenizer()
+model = load_model()
 
 
 
@@ -935,3 +940,6 @@ st.write("---")
 st.write("NeuraSense AI: Quantum-Enhanced Sensory Simulation v4.0")
 st.write("Disclaimer: This is an advanced simulation and does not represent current technological capabilities.""")
 
+# After processing
+torch.cuda.empty_cache()
+gc.collect()