Update app.py

app.py

@@ -9,7 +9,7 @@ log_path = "train.log"
 model_dir = "trained_model"
 zip_file = "trained_model.zip"
 
-# Try loading model for test tab
+# Try loading the generator model once at the top (for better performance)
 try:
     generator = pipeline("text-generation", model=model_dir, tokenizer="distilgpt2")
 except:
@@ -32,9 +32,9 @@ def start_training():
 
     if os.path.exists(model_dir):
         shutil.make_archive("trained_model", "zip", model_dir)
-        return "✅ Training complete!", os.path.abspath(zip_file)
+        return "✅ Training complete!", zip_file
     else:
-        return "❌ Training failed.", ""
+        return "❌ Training failed.", None
 
 def read_logs():
     if os.path.exists(log_path):
@@ -57,28 +57,46 @@ def generate_response(prompt):
     except Exception as e:
         return f"❌ Error: {e}"
 
+def upload_model(file):
+    global generator
+    if file is None:
+        return "No model uploaded."
+    shutil.unpack_archive(file.name, model_dir)
+    try:
+        generator = pipeline("text-generation", model=model_dir, tokenizer="distilgpt2")
+        return "✅ Model uploaded and loaded successfully!"
+    except Exception as e:
+        return f"❌ Failed to load model: {e}"
+
 # === UI ===
 with gr.Blocks() as app:
-    with gr.Tab("🧠 Train AI"):
-        gr.Markdown("## 📥 Upload your dataset and 🎯 train a Godot AI")
+    with gr.Tab("🧠 Train Python AI"):
+        gr.Markdown("## 📥 Upload your dataset and 🎯 train a *Python AI* model")
         file_input = gr.File(label="Upload JSONL Dataset")
-        upload_btn = gr.Button("Upload")
+        upload_btn = gr.Button("Upload Dataset")
         status_box = gr.Textbox(label="Upload Status")
 
         start_btn = gr.Button("🚀 Start Training")
         log_output = gr.Textbox(label="📜 Training Logs", lines=15)
-        refresh_btn = gr.Button("🔁 Refresh Logs")
-        download_btn = gr.DownloadButton(label="📥 Download Trained Model", visible=False)
+
+        download_btn = gr.File(label="📥 Download Trained Model", visible=False)
 
         upload_btn.click(fn=upload_file, inputs=file_input, outputs=status_box)
         start_btn.click(fn=start_training, outputs=[status_box, download_btn])
-        refresh_btn.click(fn=read_logs, outputs=log_output)
+        start_btn.click(fn=read_logs, outputs=log_output)
 
-    with gr.Tab("🚀 Test AI"):
-        gr.Markdown("## 💡 Try your trained Godot AI below")
+    with gr.Tab("🚀 Test Python AI"):
+        gr.Markdown("## 💡 Try your trained Python AI below")
         prompt_input = gr.Textbox(label="Enter Prompt")
         test_btn = gr.Button("🔍 Test AI")
         response_output = gr.Textbox(label="AI Response", lines=10)
         test_btn.click(fn=generate_response, inputs=prompt_input, outputs=response_output)
 
+    with gr.Tab("📤 Upload Pretrained Model"):
+        gr.Markdown("## 📦 Upload a trained Python AI model (.zip) to test")
+        model_file_input = gr.File(label="Upload Model ZIP")
+        model_upload_btn = gr.Button("Upload Model")
+        model_status = gr.Textbox(label="Model Upload Status")
+        model_upload_btn.click(fn=upload_model, inputs=model_file_input, outputs=model_status)
+
 app.launch()