Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from verbatim_llm import TokenSwapProcessor
+import gc
 
 # Predefined model pairs
 MODEL_PAIRS = {
@@ -14,6 +15,32 @@ MODEL_PAIRS = {
 loaded_models = {}
 current_pair = None
 
+def clear_models():
+    global loaded_models, current_pair
+    
+    try:
+        # Clear models from memory
+        if loaded_models:
+            # Move models to CPU if they were on GPU
+            for key, value in loaded_models.items():
+                if hasattr(value, 'to'):
+                    value.to('cpu')
+                del value
+        
+        loaded_models = {}
+        current_pair = None
+        
+        # Force garbage collection
+        gc.collect()
+        
+        # Clear GPU cache if available
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        
+        return "✅ Models cleared from memory"
+    except Exception as e:
+        return f"❌ Error clearing models: {str(e)}"
+
 def load_models(model_pair):
     global loaded_models, current_pair
     
@@ -21,6 +48,10 @@ def load_models(model_pair):
         return "Models already loaded!"
     
     try:
+        # Clear existing models first if switching
+        if loaded_models:
+            clear_models()
+        
         main_model_name, aux_model_name = MODEL_PAIRS[model_pair]
         device = "cuda" if torch.cuda.is_available() else "cpu"
         
@@ -87,7 +118,9 @@ with gr.Blocks(title="Verbatim-LLM Demo") as app:
             value=list(MODEL_PAIRS.keys())[0],
             label="Model Pair"
         )
-        load_btn = gr.Button("Load Models", variant="primary")
+        with gr.Column():
+            load_btn = gr.Button("Load Models", variant="primary")
+            clear_btn = gr.Button("Clear Models", variant="secondary")
     
     status = gr.Textbox(label="Status", interactive=False)
     
@@ -114,6 +147,11 @@ with gr.Blocks(title="Verbatim-LLM Demo") as app:
         outputs=[status]
     )
     
+    clear_btn.click(
+        fn=clear_models,
+        outputs=[status]
+    )
+    
     generate_btn.click(
         fn=lambda p, t: (generate_text(p, t, False), generate_text(p, t, True)),
         inputs=[prompt_box, max_tokens],