working on a fix

vms/services/trainer.py

@@ -353,7 +353,7 @@ class TrainingService:
         resume_from_checkpoint: Optional[str] = None,
     ) -> Tuple[str, str]:
         """Start training with finetrainers"""
-            
+        
         self.clear_logs()
 
         if not model_type:
@@ -365,22 +365,31 @@ class TrainingService:
         is_resuming = resume_from_checkpoint is not None
         log_prefix = "Resuming" if is_resuming else "Initializing"
         logger.info(f"{log_prefix} training with model_type={model_type}")
-        self.append_log(f"{log_prefix} training with model_type={model_type}")
-        
-        if is_resuming:
-            self.append_log(f"Resuming from checkpoint: {resume_from_checkpoint}")
         
         try:
-            # Get absolute paths
-            current_dir = Path(__file__).parent.absolute()
-            train_script = current_dir.parent / "train.py"
-
+            # Get absolute paths - FIXED to look in project root instead of within vms directory
+            current_dir = Path(__file__).parent.parent.parent.absolute()  # Go up to project root
+            train_script = current_dir / "train.py"
             
             if not train_script.exists():
-                error_msg = f"Training script not found at {train_script}"
-                logger.error(error_msg)
-                return error_msg, "Training script not found"
+                # Try alternative locations
+                alt_locations = [
+                    current_dir.parent / "train.py",  # One level up from project root
+                    Path("/home/user/app/train.py"),  # Absolute path
+                    Path("train.py")  # Current working directory
+                ]
+                
+                for alt_path in alt_locations:
+                    if alt_path.exists():
+                        train_script = alt_path
+                        logger.info(f"Found train.py at alternative location: {train_script}")
+                        break
                 
+                if not train_script.exists():
+                    error_msg = f"Training script not found at {train_script} or any alternative locations"
+                    logger.error(error_msg)
+                    return error_msg, "Training script not found"
+                    
             # Log paths for debugging
             logger.info("Current working directory: %s", current_dir)
             logger.info("Training script path: %s", train_script)

vms/tabs/train_tab.py

@@ -91,20 +91,35 @@ class TrainTab(BaseTab):
                 
                 with gr.Column():
                     with gr.Row():
+                        # Check for existing checkpoints to determine button text
+                        has_checkpoints = len(list(OUTPUT_PATH.glob("checkpoint-*"))) > 0
+                        start_text = "Continue Training" if has_checkpoints else "Start Training"
+                        
                         self.components["start_btn"] = gr.Button(
-                            "Start Training",
+                            start_text,
                             variant="primary",
                             interactive=not ASK_USER_TO_DUPLICATE_SPACE
                         )
+                        
+                        # Just use stop and pause buttons for now to ensure compatibility
+                        self.components["stop_btn"] = gr.Button(
+                            "Stop at Last Checkpoint",
+                            variant="primary",
+                            interactive=False
+                        )
+                        
                         self.components["pause_resume_btn"] = gr.Button(
                             "Resume Training",
                             variant="secondary",
-                            interactive=False
+                            interactive=False,
+                            visible=False
                         )
-                        self.components["stop_btn"] = gr.Button(
-                            "Stop Training",
+                        
+                        # Add delete checkpoints button - THIS IS THE KEY FIX
+                        self.components["delete_checkpoints_btn"] = gr.Button(
+                            "Delete All Checkpoints",
                             variant="stop",
-                            interactive=False
+                            interactive=True
                         )
 
                     with gr.Row():
@@ -468,31 +483,56 @@ class TrainTab(BaseTab):
 
         return (state["status"], state["message"], logs)
 
-    def get_latest_status_message_logs_and_button_labels(self) -> Tuple[str, str, Any, Any, Any]:
+    def get_latest_status_message_logs_and_button_labels(self) -> Tuple:
+        """Get latest status message, logs and button states"""
         status, message, logs = self.get_latest_status_message_and_logs()
-        return (
-            message,
-            logs,
-            *self.update_training_buttons(status).values()
-        )
-
-    def update_training_buttons(self, status: str) -> Dict:
+        
+        # Add checkpoints detection
+        has_checkpoints = len(list(OUTPUT_PATH.glob("checkpoint-*"))) > 0
+        
+        button_updates = self.update_training_buttons(status, has_checkpoints).values()
+        
+        # Return in order expected by timer
+        return (message, logs, *button_updates)
+    
+    def update_training_buttons(self, status: str, has_checkpoints: bool = None) -> Dict:
         """Update training control buttons based on state"""
+        if has_checkpoints is None:
+            has_checkpoints = len(list(OUTPUT_PATH.glob("checkpoint-*"))) > 0
+            
         is_training = status in ["training", "initializing"]
-        is_paused = status == "paused"
         is_completed = status in ["completed", "error", "stopped"]
-        return {
+        
+        start_text = "Continue Training" if has_checkpoints else "Start Training"
+        
+        # Only include buttons that we know exist in components
+        result = {
             "start_btn": gr.Button(
-                interactive=not is_training and not is_paused,
+                value=start_text,
+                interactive=not is_training,
                 variant="primary" if not is_training else "secondary",
             ),
             "stop_btn": gr.Button(
-                interactive=is_training or is_paused,
+                value="Stop at Last Checkpoint",
+                interactive=is_training,
+                variant="primary" if is_training else "secondary",
+            )
+        }
+        
+        # Add delete_checkpoints_btn only if it exists in components
+        if "delete_checkpoints_btn" in self.components:
+            result["delete_checkpoints_btn"] = gr.Button(
+                value="Delete All Checkpoints",
+                interactive=has_checkpoints and not is_training,
                 variant="stop",
-            ),
-            "pause_resume_btn": gr.Button(
-                value="Resume Training" if is_paused else "Pause Training",
-                interactive=(is_training or is_paused) and not is_completed,
+            )
+        else:
+            # Add pause_resume_btn as fallback
+            result["pause_resume_btn"] = gr.Button(
+                value="Resume Training" if status == "paused" else "Pause Training",
+                interactive=(is_training or status == "paused") and not is_completed,
                 variant="secondary",
+                visible=False
             )
-        }
+        
+        return result

vms/ui/video_trainer_ui.py

@@ -100,15 +100,25 @@ class VideoTrainerUI:
         """Add auto-refresh timers to the UI"""
         # Status update timer (every 1 second)
         status_timer = gr.Timer(value=1)
+        
+        # Use a safer approach - check if the component exists before using it
+        outputs = [
+            self.tabs["train_tab"].components["status_box"],
+            self.tabs["train_tab"].components["log_box"],
+            self.tabs["train_tab"].components["start_btn"],
+            self.tabs["train_tab"].components["stop_btn"]
+        ]
+        
+        # Add delete_checkpoints_btn only if it exists
+        if "delete_checkpoints_btn" in self.tabs["train_tab"].components:
+            outputs.append(self.tabs["train_tab"].components["delete_checkpoints_btn"])
+        else:
+            # Add pause_resume_btn as fallback
+            outputs.append(self.tabs["train_tab"].components["pause_resume_btn"])
+        
         status_timer.tick(
             fn=self.tabs["train_tab"].get_latest_status_message_logs_and_button_labels,
-            outputs=[
-                self.tabs["train_tab"].components["status_box"],
-                self.tabs["train_tab"].components["log_box"],
-                self.tabs["train_tab"].components["start_btn"],
-                self.tabs["train_tab"].components["stop_btn"],
-                self.tabs["train_tab"].components["delete_checkpoints_btn"]  # Replace pause_resume_btn
-            ]
+            outputs=outputs
         )
         
         # Dataset refresh timer (every 5 seconds)