Update app.py

app.py

@@ -1,16 +1,37 @@
 import os
 import gradio as gr
+import subprocess
+import zipfile
 
-# Check for trained model zip
-model_zip = "trained_model.zip"
+model_dir = "trained_model"
+zip_path = "trained_model.zip"
+
+def model_exists():
+    return os.path.exists(zip_path)
+
+def train_model():
+    result = subprocess.run(["python", "train.py"], capture_output=True, text=True)
+    if result.returncode == 0 and os.path.exists(zip_path):
+        return "✅ Model trained successfully! Ready for download.", zip_path
+    else:
+        return f"❌ Training failed:\n\n{result.stderr}", None
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 🤖 Python AI Model Trainer")
+    gr.Markdown("## 🧠 Python AI Model Trainer")
 
-    if os.path.exists(model_zip):
-        gr.Markdown("### ✅ Model trained! Click below to download:")
-        gr.File(value=model_zip, label="Download trained model")
+    if model_exists():
+        gr.Markdown("✅ Trained model found. Click below to download:")
+        gr.File(value=zip_path, label="Download Trained Model")
     else:
-        gr.Markdown("### 🚫 No trained model found yet. Run train.py to train it.")
+        gr.Markdown("🚫 No trained model found yet.")
+        output = gr.Textbox(label="Training Log")
+        download = gr.File(visible=False)
+        train_button = gr.Button("🚀 Train Model")
+
+        def on_click_train():
+            message, path = train_model()
+            return message, gr.update(value=path, visible=True) if path else gr.update(visible=False)
+
+        train_button.click(fn=on_click_train, outputs=[output, download])
 
 demo.launch()