Update app.py

app.py

@@ -97,9 +97,50 @@ def t2i(prompt, model, seed, steps, scale, w, h):
         generator=gen
     ).images[0]
 
+# ===== NEW: Function to unload all models =====
+def unload_all_models():
+    global PIPE_STANDARD, PIPE_ANIME, LINEART_DETECTOR, LINEART_ANIME_DETECTOR
+    global CURRENT_T2I_PIPE, CURRENT_T2I_MODEL
+    
+    print("Unloading all models from memory...")
+    
+    # Unload lineart models
+    if PIPE_STANDARD is not None:
+        del PIPE_STANDARD
+        PIPE_STANDARD = None
+    if PIPE_ANIME is not None:
+        del PIPE_ANIME
+        PIPE_ANIME = None
+    if LINEART_DETECTOR is not None:
+        del LINEART_DETECTOR
+        LINEART_DETECTOR = None
+    if LINEART_ANIME_DETECTOR is not None:
+        del LINEART_ANIME_DETECTOR
+        LINEART_ANIME_DETECTOR = None
+    
+    # Unload T2I model
+    if CURRENT_T2I_PIPE is not None:
+        del CURRENT_T2I_PIPE
+        CURRENT_T2I_PIPE = None
+    CURRENT_T2I_MODEL = None
+    
+    # Force garbage collection
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        print(f"GPU memory cleared. Allocated: {torch.cuda.memory_allocated()/1024**3:.2f} GB")
+    
+    return "✅ All models unloaded from memory!"
+
 # ===== Gradio UI (Minimal) =====
 with gr.Blocks() as demo:
     gr.Markdown("# 🎨 Minimal Style2Paints")
+    
+    # Add unload button at the top
+    with gr.Row():
+        unload_btn = gr.Button("🗑️ Unload All Models", variant="stop")
+        status_text = gr.Textbox(label="Status", interactive=False)
+    unload_btn.click(unload_all_models, outputs=status_text)
 
     with gr.Tab("🎨 Colorize"):
         with gr.Row():