Update app.py

app.py

@@ -540,7 +540,264 @@ api.upload_folder(
 - Epochs: {num_epochs}
 - Model saved successfully!
 
-**You can still test it locally or manually upload!**"""# ==================== GRADIO INTERFACE ====================
+**You can still test it locally or manually upload!**"""# ==================== MODEL MANAGEMENT ====================
+
+import zipfile
+import shutil
+from pathlib import Path
+
+class ModelManager:
+    def __init__(self):
+        self.models_dir = Path("./saved_models")
+        self.models_dir.mkdir(exist_ok=True)
+    
+    @handle_errors("model_management")
+    def create_model_zip(self, model_path: str, model_name: str) -> Tuple[str, str]:
+        """Create a downloadable zip of a trained model"""
+        if not os.path.exists(model_path):
+            raise ArchitechError(f"Model path not found: {model_path}")
+        
+        zip_filename = f"{model_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
+        zip_path = os.path.join(self.models_dir, zip_filename)
+        
+        # Create zip file
+        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for root, dirs, files in os.walk(model_path):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, model_path)
+                    zipf.write(file_path, arcname)
+        
+        file_size = os.path.getsize(zip_path) / (1024 * 1024)  # MB
+        
+        return zip_path, f"✅ Created {zip_filename} ({file_size:.2f} MB)"
+    
+    @handle_errors("model_management")
+    def extract_model_zip(self, zip_file, progress=gr.Progress()) -> str:
+        """Extract uploaded model zip"""
+        if zip_file is None:
+            raise ArchitechError("No file uploaded!")
+        
+        progress(0.1, "📦 Extracting model archive...")
+        
+        # Get filename
+        zip_filename = Path(zip_file.name).name
+        model_name = zip_filename.replace('.zip', '')
+        extract_path = os.path.join("./uploaded_models", model_name)
+        
+        os.makedirs(extract_path, exist_ok=True)
+        
+        progress(0.3, "📂 Unpacking files...")
+        
+        # Extract zip
+        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
+            zip_ref.extractall(extract_path)
+        
+        progress(0.7, "🔍 Validating model files...")
+        
+        # Check for required files
+        files = os.listdir(extract_path)
+        has_model = any('pytorch_model' in f or 'model.safetensors' in f for f in files)
+        has_config = 'config.json' in files
+        has_tokenizer = any('tokenizer' in f for f in files)
+        
+        validation_status = []
+        if has_model:
+            validation_status.append("✅ Model weights found")
+        else:
+            validation_status.append("⚠️ Model weights not found")
+        
+        if has_config:
+            validation_status.append("✅ Config file found")
+        else:
+            validation_status.append("⚠️ Config file not found")
+        
+        if has_tokenizer:
+            validation_status.append("✅ Tokenizer found")
+        else:
+            validation_status.append("⚠️ Tokenizer not found")
+        
+        progress(1.0, "✅ Extraction complete!")
+        
+        return f"""🎉 **Model Uploaded Successfully!**
+
+**Extracted to:** `{extract_path}`
+
+**Validation:**
+{chr(10).join(validation_status)}
+
+**Files found:** {len(files)} files
+
+**You can now:**
+1. Use this model in the Test Model tab
+2. Continue training from this checkpoint
+3. Push to HuggingFace Hub
+
+*Model path: `{extract_path}`*"""
+    
+    def list_local_models(self) -> str:
+        """List all locally saved models"""
+        trained_models = []
+        uploaded_models = []
+        
+        # Check trained models
+        if os.path.exists("./"):
+            for item in os.listdir("./"):
+                if item.startswith("trained_") and os.path.isdir(item):
+                    size = sum(
+                        os.path.getsize(os.path.join(dirpath, filename))
+                        for dirpath, dirnames, filenames in os.walk(item)
+                        for filename in filenames
+                    ) / (1024 * 1024)
+                    trained_models.append(f"- `{item}` ({size:.2f} MB)")
+        
+        # Check uploaded models
+        if os.path.exists("./uploaded_models"):
+            for item in os.listdir("./uploaded_models"):
+                path = os.path.join("./uploaded_models", item)
+                if os.path.isdir(path):
+                    size = sum(
+                        os.path.getsize(os.path.join(dirpath, filename))
+                        for dirpath, dirnames, filenames in os.walk(path)
+                        for filename in filenames
+                    ) / (1024 * 1024)
+                    uploaded_models.append(f"- `{item}` ({size:.2f} MB)")
+        
+        result = "## 📦 Local Models\n\n"
+        
+        if trained_models:
+            result += "### Trained Models:\n" + "\n".join(trained_models) + "\n\n"
+        else:
+            result += "### Trained Models:\n*No trained models found*\n\n"
+        
+        if uploaded_models:
+            result += "### Uploaded Models:\n" + "\n".join(uploaded_models) + "\n\n"
+        else:
+            result += "### Uploaded Models:\n*No uploaded models found*\n\n"
+        
+        return result
+    
+    @handle_errors("model_management")
+    def delete_model(self, model_path: str) -> str:
+        """Delete a local model"""
+        if not os.path.exists(model_path):
+            raise ArchitechError(f"Model not found: {model_path}")
+        
+        shutil.rmtree(model_path)
+        return f"✅ Deleted: {model_path}"
+
+model_manager = ModelManager()
+
+# Add this to the Gradio interface creation function
+# Insert this tab after the "Test Model" tab and before "About"
+
+def add_model_management_tab():
+    """Add Model Management tab to Gradio interface"""
+    with gr.Tab("💾 Model Management"):
+        gr.Markdown("""
+        ### Manage Your Models
+        Upload, download, and organize your trained models
+        """)
+        
+        with gr.Row():
+            # Upload Section
+            with gr.Column():
+                gr.Markdown("### 📤 Upload Model")
+                
+                upload_file = gr.File(
+                    label="Upload Model ZIP",
+                    file_types=[".zip"],
+                    type="filepath"
+                )
+                
+                upload_btn = gr.Button("📦 Extract and Save", variant="primary")
+                upload_output = gr.Markdown()
+                
+                upload_btn.click(
+                    fn=model_manager.extract_model_zip,
+                    inputs=[upload_file],
+                    outputs=upload_output
+                )
+            
+            # Download Section
+            with gr.Column():
+                gr.Markdown("### 📥 Download Model")
+                
+                model_path_input = gr.Textbox(
+                    label="Model Path",
+                    placeholder="e.g., ./trained_my-model or ./uploaded_models/my-model",
+                    info="Path to the model directory you want to download"
+                )
+                
+                model_name_input = gr.Textbox(
+                    label="Archive Name",
+                    placeholder="e.g., my-awesome-model",
+                    info="Name for the downloaded zip file"
+                )
+                
+                download_btn = gr.Button("📦 Create ZIP", variant="primary")
+                download_file = gr.File(label="Download")
+                download_output = gr.Markdown()
+                
+                def create_and_return_zip(model_path, model_name):
+                    zip_path, message = model_manager.create_model_zip(model_path, model_name)
+                    return zip_path, message
+                
+                download_btn.click(
+                    fn=create_and_return_zip,
+                    inputs=[model_path_input, model_name_input],
+                    outputs=[download_file, download_output]
+                )
+        
+        gr.Markdown("---")
+        
+        # List and Delete Section
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### 📋 Your Models")
+                
+                refresh_btn = gr.Button("🔄 Refresh List", variant="secondary")
+                models_list = gr.Markdown()
+                
+                refresh_btn.click(
+                    fn=model_manager.list_local_models,
+                    inputs=[],
+                    outputs=models_list
+                )
+                
+                # Auto-load on tab open
+                models_list.value = model_manager.list_local_models()
+            
+            with gr.Column():
+                gr.Markdown("### 🗑️ Delete Model")
+                
+                delete_path = gr.Textbox(
+                    label="Model Path to Delete",
+                    placeholder="e.g., ./trained_my-model"
+                )
+                
+                delete_btn = gr.Button("🗑️ Delete Model", variant="stop")
+                delete_output = gr.Markdown()
+                
+                delete_btn.click(
+                    fn=model_manager.delete_model,
+                    inputs=[delete_path],
+                    outputs=delete_output
+                )
+        
+        gr.Markdown("""
+        ---
+        ### 💡 Tips:
+        - **Upload:** Upload model zips from other systems or backups
+        - **Download:** Create portable archives of your trained models
+        - **Organize:** Keep your workspace tidy by managing local models
+        - **Backup:** Download important models before deleting them
+        
+        *Note: Uploaded/downloaded models persist only during your session unless you have persistent storage configured.*
+        """)
+
+# This function should be called in create_gradio_interface()
+# Add it right before the "About" tab# ==================== GRADIO INTERFACE ====================
 
 def create_gradio_interface():
     agent = ArchitechAgent()
@@ -640,6 +897,100 @@ def create_gradio_interface():
                     outputs=test_output
                 )
             
+            # Model Management Tab
+            with gr.Tab("💾 Model Management"):
+                gr.Markdown("""
+                ### Manage Your Models
+                Upload, download, and organize your trained models
+                """)
+                
+                with gr.Row():
+                    # Upload Section
+                    with gr.Column():
+                        gr.Markdown("### 📤 Upload Model")
+                        
+                        upload_file = gr.File(
+                            label="Upload Model ZIP",
+                            file_types=[".zip"],
+                            type="filepath"
+                        )
+                        
+                        upload_btn = gr.Button("📦 Extract and Save", variant="primary")
+                        upload_output = gr.Markdown()
+                        
+                        upload_btn.click(
+                            fn=model_manager.extract_model_zip,
+                            inputs=[upload_file],
+                            outputs=upload_output
+                        )
+                    
+                    # Download Section
+                    with gr.Column():
+                        gr.Markdown("### 📥 Download Model")
+                        
+                        model_path_input = gr.Textbox(
+                            label="Model Path",
+                            placeholder="e.g., ./trained_my-model",
+                            info="Path to the model directory"
+                        )
+                        
+                        model_name_input = gr.Textbox(
+                            label="Archive Name",
+                            placeholder="e.g., my-awesome-model",
+                            info="Name for the zip file"
+                        )
+                        
+                        download_btn = gr.Button("📦 Create ZIP", variant="primary")
+                        download_file = gr.File(label="Download")
+                        download_output = gr.Markdown()
+                        
+                        def create_and_return_zip(model_path, model_name):
+                            zip_path, message = model_manager.create_model_zip(model_path, model_name)
+                            return zip_path, message
+                        
+                        download_btn.click(
+                            fn=create_and_return_zip,
+                            inputs=[model_path_input, model_name_input],
+                            outputs=[download_file, download_output]
+                        )
+                
+                gr.Markdown("---")
+                
+                # List Models
+                with gr.Row():
+                    with gr.Column():
+                        gr.Markdown("### 📋 Your Models")
+                        refresh_btn = gr.Button("🔄 Refresh List", variant="secondary")
+                        models_list = gr.Markdown()
+                        
+                        refresh_btn.click(
+                            fn=model_manager.list_local_models,
+                            outputs=models_list
+                        )
+                    
+                    with gr.Column():
+                        gr.Markdown("### 🗑️ Delete Model")
+                        delete_path = gr.Textbox(
+                            label="Model Path to Delete",
+                            placeholder="e.g., ./trained_my-model"
+                        )
+                        delete_btn = gr.Button("🗑️ Delete", variant="stop")
+                        delete_output = gr.Markdown()
+                        
+                        delete_btn.click(
+                            fn=model_manager.delete_model,
+                            inputs=[delete_path],
+                            outputs=delete_output
+                        )
+                
+                gr.Markdown("""
+                ### 💡 Tips:
+                - Upload model zips from backups or other systems
+                - Download models as portable archives
+                - Keep your workspace organized
+                - Always backup before deleting!
+                """)
+            
             # About
             with gr.Tab("ℹ️ About"):
                 gr.Markdown("""