Revert to commit a9df757

vms/ui/app_ui.py

@@ -403,7 +403,6 @@ class AppUI:
             ]
         )
         
-        
         # Button update timer for button components (every 1 second)
         button_timer = gr.Timer(value=1)
         button_outputs = [

vms/ui/models/tabs/training_tab.py

@@ -88,8 +88,9 @@ class TrainingTab(BaseTab):
                             gr.Markdown(model.model_display_name or "Unknown")
                         
                         with gr.Column(scale=2, min_width=20):
-                            progress_text = f"Step {model.current_step}/{model.total_steps} ({model.training_progress:.1f}%)"
+                            progress_text = f"Step {model.current_step}/{model.total_steps}"
                             gr.Markdown(progress_text)
+                            gr.Progress(value=model.training_progress/100)
                         
                         with gr.Column(scale=2, min_width=20):
                             with gr.Row():

vms/ui/project/services/training.py

@@ -1823,9 +1823,12 @@ class TrainingService:
         try:
             checkpoints = list(self.app.output_path.glob("finetrainers_step_*"))
             if not checkpoints:
-                return "No checkpoints available"
+                return "📥 Download checkpoints (not available)"
             
-            return f"💽 Download checkpoints"
+            # Get the latest checkpoint by step number
+            latest_checkpoint = max(checkpoints, key=lambda x: int(x.name.split("_")[-1]))
+            step_num = int(latest_checkpoint.name.split("_")[-1])
+            return f"📥 Download checkpoints (step {step_num})"
         except Exception as e:
             logger.warning(f"Error getting checkpoint info for button text: {e}")
-            return "No checkpoints available"
+            return "📥 Download checkpoints (not available)"

vms/ui/project/tabs/manage_tab.py

@@ -25,6 +25,50 @@ class ManageTab(BaseTab):
         self.id = "manage_tab"
         self.title = "5️⃣ Storage"
     
+    def get_download_button_text(self) -> str:
+        """Get the dynamic text for the download button based on current model state"""
+        try:
+            model_info = self.app.training.get_model_output_info()
+            if model_info["path"] and model_info["steps"]:
+                return f"🧠 Download weights ({model_info['steps']} steps)"
+            elif model_info["path"]:
+                return "🧠 Download weights (.safetensors)"
+            else:
+                return "🧠 Download weights (not available)"
+        except Exception as e:
+            logger.warning(f"Error getting model info for button text: {e}")
+            return "🧠 Download weights (.safetensors)"
+
+    def get_checkpoint_button_text(self) -> str:
+        """Get the dynamic text for the download checkpoint button"""
+        try:
+            return self.app.training.get_checkpoint_button_text()
+        except Exception as e:
+            logger.warning(f"Error getting checkpoint button text: {e}")
+            return "📥 Download checkpoints (not available)"
+
+    def update_download_button_text(self) -> gr.update:
+        """Update the download button text"""
+        return gr.update(value=self.get_download_button_text())
+    
+    def update_checkpoint_button_text(self) -> gr.update:
+        """Update the checkpoint button text"""
+        return gr.update(value=self.get_checkpoint_button_text())
+    
+    def update_both_download_buttons(self) -> Tuple[gr.update, gr.update]:
+        """Update both download button texts"""
+        return (
+            gr.update(value=self.get_download_button_text()),
+            gr.update(value=self.get_checkpoint_button_text())
+        )
+    
+    def download_and_update_button(self):
+        """Handle download and return updated button with current text"""
+        # Get the safetensors path for download
+        path = self.app.training.get_model_output_safetensors()
+        # For DownloadButton, we need to return the file path directly for download
+        # The button text will be updated on next render
+        return path
     
     def create(self, parent=None) -> gr.TabItem:
         """Create the Manage tab UI components"""
@@ -46,19 +90,19 @@ class ManageTab(BaseTab):
                             gr.Markdown("📦 Training dataset download disabled for large datasets")
                             
                         self.components["download_model_btn"] = gr.DownloadButton(
-                            "🧠 Download LoRA weights",
+                            self.get_download_button_text(),
                             variant="secondary",
                             size="lg"
                         )
                         
                         self.components["download_checkpoint_btn"] = gr.DownloadButton(
-                            "💽 Download Checkpoints",
+                            self.get_checkpoint_button_text(),
                             variant="secondary",
                             size="lg"
                         )
                         
                         self.components["download_output_btn"] = gr.DownloadButton(
-                            "📁 Download output/ (.zip)",
+                            "📁 Download output directory (.zip)",
                             variant="secondary",
                             size="lg",
                             visible=False

vms/ui/project/tabs/train_tab.py

@@ -494,7 +494,12 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
             save_iterations, repo_id, progress
         )
         
-        return status, logs
+        # Update download button texts
+        manage_tab = self.app.tabs["manage_tab"]
+        download_btn_text = gr.update(value=manage_tab.get_download_button_text())
+        checkpoint_btn_text = gr.update(value=manage_tab.get_checkpoint_button_text())
+        
+        return status, logs, download_btn_text, checkpoint_btn_text
 
     def handle_resume_training(
         self, model_type, model_version, training_type, 
@@ -506,7 +511,10 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
         checkpoints = list(self.app.output_path.glob("finetrainers_step_*"))
 
         if not checkpoints:
-            return "No checkpoints found to resume from", "Please start a new training session instead"
+            manage_tab = self.app.tabs["manage_tab"]
+            download_btn_text = gr.update(value=manage_tab.get_download_button_text())
+            checkpoint_btn_text = gr.update(value=manage_tab.get_checkpoint_button_text())
+            return "No checkpoints found to resume from", "Please start a new training session instead", download_btn_text, checkpoint_btn_text
         
         self.app.training.append_log(f"Resuming training from latest checkpoint")
     
@@ -518,7 +526,12 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
             resume_from_checkpoint="latest"
         )
         
-        return status, logs
+        # Update download button texts
+        manage_tab = self.app.tabs["manage_tab"]
+        download_btn_text = gr.update(value=manage_tab.get_download_button_text())
+        checkpoint_btn_text = gr.update(value=manage_tab.get_checkpoint_button_text())
+        
+        return status, logs, download_btn_text, checkpoint_btn_text
 
     def handle_start_from_lora_training(
         self, model_type, model_version, training_type, 
@@ -529,22 +542,26 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
         # Find the latest LoRA weights
         lora_weights_path = self.app.output_path / "lora_weights"
         
+        manage_tab = self.app.tabs["manage_tab"]
+        download_btn_text = gr.update(value=manage_tab.get_download_button_text())
+        checkpoint_btn_text = gr.update(value=manage_tab.get_checkpoint_button_text())
+        
         if not lora_weights_path.exists():
-            return "No LoRA weights found", "Please train a model first or start a new training session"
+            return "No LoRA weights found", "Please train a model first or start a new training session", download_btn_text, checkpoint_btn_text
         
         # Find the latest LoRA checkpoint directory
         lora_dirs = sorted([d for d in lora_weights_path.iterdir() if d.is_dir()], 
                           key=lambda x: int(x.name), reverse=True)
         
         if not lora_dirs:
-            return "No LoRA weight directories found", "Please train a model first or start a new training session"
+            return "No LoRA weight directories found", "Please train a model first or start a new training session", download_btn_text, checkpoint_btn_text
         
         latest_lora_dir = lora_dirs[0]
         
         # Verify the LoRA weights file exists
         lora_weights_file = latest_lora_dir / "pytorch_lora_weights.safetensors"
         if not lora_weights_file.exists():
-            return f"LoRA weights file not found in {latest_lora_dir}", "Please check your LoRA weights directory"
+            return f"LoRA weights file not found in {latest_lora_dir}", "Please check your LoRA weights directory", download_btn_text, checkpoint_btn_text
         
         # Clear checkpoints to start fresh (but keep LoRA weights)
         for checkpoint in self.app.output_path.glob("finetrainers_step_*"):
@@ -565,7 +582,11 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
             save_iterations, repo_id, progress,
         )
         
-        return status, logs
+        # Update download button texts
+        download_btn_text = gr.update(value=manage_tab.get_download_button_text())
+        checkpoint_btn_text = gr.update(value=manage_tab.get_checkpoint_button_text())
+        
+        return status, logs, download_btn_text, checkpoint_btn_text
 
     def connect_events(self) -> None:
         """Connect event handlers to UI components"""
@@ -748,7 +769,9 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
             ],
             outputs=[
                 self.components["status_box"],
-                self.components["log_box"]
+                self.components["log_box"],
+                self.app.tabs["manage_tab"].components["download_model_btn"],
+                self.app.tabs["manage_tab"].components["download_checkpoint_btn"]
             ]
         )
 
@@ -768,7 +791,9 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
             ],
             outputs=[
                 self.components["status_box"],
-                self.components["log_box"]
+                self.components["log_box"],
+                self.app.tabs["manage_tab"].components["download_model_btn"],
+                self.app.tabs["manage_tab"].components["download_checkpoint_btn"]
             ]
         )
 
@@ -788,7 +813,9 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
             ],
             outputs=[
                 self.components["status_box"],
-                self.components["log_box"]
+                self.components["log_box"],
+                self.app.tabs["manage_tab"].components["download_model_btn"],
+                self.app.tabs["manage_tab"].components["download_checkpoint_btn"]
             ]
         )
         
@@ -804,7 +831,9 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
                 self.components["current_task_box"],
                 self.components["start_btn"],
                 self.components["stop_btn"],
-                third_btn
+                third_btn,
+                self.app.tabs["manage_tab"].components["download_model_btn"],
+                self.app.tabs["manage_tab"].components["download_checkpoint_btn"]
             ]
         )
 
@@ -816,7 +845,9 @@ For image-to-video tasks, 'index' (usually with index 0) is most common as it co
                 self.components["current_task_box"],
                 self.components["start_btn"],
                 self.components["stop_btn"],
-                third_btn
+                third_btn,
+                self.app.tabs["manage_tab"].components["download_model_btn"],
+                self.app.tabs["manage_tab"].components["download_checkpoint_btn"]
             ]
         )
 
@@ -1209,7 +1240,12 @@ Full finetune mode trains all parameters of the model, requiring more VRAM but p
             variant="stop"
         )
         
-        return start_btn, resume_btn, stop_btn, delete_checkpoints_btn
+        # Update download button texts
+        manage_tab = self.app.tabs["manage_tab"]
+        download_btn_text = gr.update(value=manage_tab.get_download_button_text())
+        checkpoint_btn_text = gr.update(value=manage_tab.get_checkpoint_button_text())
+        
+        return start_btn, resume_btn, stop_btn, delete_checkpoints_btn, download_btn_text, checkpoint_btn_text
             
     def update_training_ui(self, training_state: Dict[str, Any]):
         """Update UI components based on training state"""