makign some fixes

app.py

@@ -36,7 +36,7 @@ from splitting_service import SplittingService
 from import_service import ImportService
 from config import (
     STORAGE_PATH, VIDEOS_TO_SPLIT_PATH, STAGING_PATH,
-    TRAINING_PATH, TRAINING_VIDEOS_PATH, MODEL_PATH, OUTPUT_PATH, DEFAULT_CAPTIONING_BOT_INSTRUCTIONS,
+    TRAINING_PATH, LOG_FILE_PATH, TRAINING_VIDEOS_PATH, MODEL_PATH, OUTPUT_PATH, DEFAULT_CAPTIONING_BOT_INSTRUCTIONS,
     DEFAULT_PROMPT_PREFIX, HF_API_TOKEN, ASK_USER_TO_DUPLICATE_SPACE, MODEL_TYPES, TRAINING_BUCKETS
 )
 from utils import make_archive, count_media_files, format_media_title, is_image_file, is_video_file, validate_model_repo, format_time
@@ -134,6 +134,9 @@ class VideoTrainerUI:
                 self.splitter.processing = False
                 status_messages["splitting"] = "Scene detection stopped"
             
+            if LOG_FILE_PATH.exists():
+                LOG_FILE_PATH.unlink()
+            
             # Clear all data directories
             for path in [VIDEOS_TO_SPLIT_PATH, STAGING_PATH, TRAINING_VIDEOS_PATH, TRAINING_PATH,
                         MODEL_PATH, OUTPUT_PATH]:
@@ -258,15 +261,11 @@ class VideoTrainerUI:
         # Only return name and status columns for display
         return [[file[0], file[1]] for file in files]
     
-    def update_training_buttons(self, training_state: Dict[str, Any]) -> Dict:
+    def update_training_buttons(self, status: str) -> Dict:
         """Update training control buttons based on state"""
-        #print("update_training_buttons: training_state = ", training_state)
-        is_training = training_state["status"] in ["training", "initializing"]
-        if  training_state["message"] == "No training in progress":
-            is_training = False
-        is_paused = training_state["status"] == "paused"
-        is_completed = training_state["status"] in ["completed", "error", "stopped"]
-        #print(f"update_training_buttons: is_training = {is_training}, is_paused = {is_paused}, is_completed = {is_completed}")
+        is_training = status in ["training", "initializing"]
+        is_paused = status == "paused"
+        is_completed = status in ["completed", "error", "stopped"]
         return {
             "start_btn": gr.Button(
                 interactive=not is_training and not is_paused,
@@ -283,32 +282,20 @@ class VideoTrainerUI:
             )
         }
     
-    def handle_training_complete(self):
-        """Handle training completion"""
-        # Reset button states
-        return self.update_training_buttons({
-            "status": "completed",
-            "progress": "100%",
-            "current_step": 0,
-            "total_steps": 0
-        })
-    
     def handle_pause_resume(self):
+        status, _, _ = self.get_latest_status_message_and_logs()
 
-        status = self.trainer.get_status()
-        print("handle_pause_resume: status = ", status)
-        if status["status"] == "paused":
-            result = self.trainer.resume_training()
-            new_state = {"status": "training"}
+        if status == "paused":
+            self.trainer.resume_training()
         else:
-            result = self.trainer.pause_training()
-            new_state = {"status": "paused"}
-        return (
-            *result,
-            *self.update_training_buttons(new_state).values()
-        )
+            self.trainer.pause_training()
+
+        return self.get_latest_status_message_logs_and_button_labels()
+
+    def handle_stop(self):
+        self.trainer.stop_training()
+        return self.get_latest_status_message_logs_and_button_labels()
 
-    
     def handle_training_dataset_select(self, evt: gr.SelectData) -> Tuple[Optional[str], Optional[str], Optional[str]]:
         """Handle selection of both video clips and images"""
         try:
@@ -623,15 +610,10 @@ class VideoTrainerUI:
             return f"Error during scene detection: {str(e)}"
 
 
-    def refresh_training_status_and_logs(self):
-        """Refresh all dynamic lists and training state"""
-        status = self.trainer.get_status()
+    def get_latest_status_message_and_logs(self) -> Tuple[str, str, str]:
+        state = self.trainer.get_status()
         logs = self.trainer.get_logs()
 
-        status_update = status["message"]
-
-        # print(f"refresh_training_status_and_logs: ", status)
-
         # Parse new log lines
         if logs:
             last_state = None
@@ -639,42 +621,28 @@ class VideoTrainerUI:
                 state_update = self.log_parser.parse_line(line)
                 if state_update:
                     last_state = state_update
-                    print("last_state = ", last_state)
             
             if last_state:
                 ui_updates = self.update_training_ui(last_state)
-                status_update = ui_updates.get("status_box", status["message"])
-        
-        return (status_update, logs)
-
-    def refresh_training_status(self):
-        """Refresh training status and update UI"""
-        status, logs = self.refresh_training_status_and_logs()
+                state["message"] = ui_updates.get("status_box", state["message"])
         
         # Parse status for training state
-        is_completed = "completed" in status.lower() or "100.0%" in status
-        current_state = {
-            "status": "completed" if is_completed else "training",
-            "message": status
-        }
-        
-        #print("refresh_training_status: current_state = ", current_state)
-        
-        if is_completed:
-            button_updates = self.handle_training_complete()
-            return (
-                status,
-                logs,
-                *button_updates.values()
-            )
-        
-        # Update based on current training state
-        button_updates = self.update_training_buttons(current_state)
+        if "completed" in state["message"].lower():
+            state["status"] = "completed"
+
+        return (state["status"], state["message"], logs)
+
+    def get_latest_status_message_logs_and_button_labels(self) -> Tuple[str, str, Any, Any, Any]:
+        status, message, logs = self.get_latest_status_message_and_logs()
         return (
-            status,
+            message,
             logs,
-            *button_updates.values()
+            *self.update_training_buttons(status).values()
         )
+
+    def get_latest_button_labels(self) -> Tuple[Any, Any, Any]:
+        status, message, logs = self.get_latest_status_message_and_logs()
+        return self.update_training_buttons(status).values()
     
     def refresh_dataset(self):
         """Refresh all dynamic lists and training state"""
@@ -1141,22 +1109,18 @@ class VideoTrainerUI:
                 ],
                 outputs=[status_box, log_box]
             ).success(
-                fn=lambda: self.update_training_buttons(),
-                outputs=[start_btn, stop_btn, pause_resume_btn]
+                fn=self.get_latest_status_message_logs_and_button_labels,
+                outputs=[status_box, log_box, start_btn, stop_btn, pause_resume_btn]
             )
 
-
             pause_resume_btn.click(
                 fn=self.handle_pause_resume,
                 outputs=[status_box, log_box, start_btn, stop_btn, pause_resume_btn]
             )
 
             stop_btn.click(
-                fn=self.trainer.stop_training,
-                outputs=[status_box, log_box]
-            ).success(
-                fn=self.handle_training_complete,
-                outputs=[start_btn, stop_btn, pause_resume_btn]
+                fn=self.handle_stop,
+                outputs=[status_box, log_box, start_btn, stop_btn, pause_resume_btn]
             )
 
             def handle_global_stop():
@@ -1218,12 +1182,12 @@ class VideoTrainerUI:
             timer = gr.Timer(value=1)
             timer.tick(
                 fn=lambda: (
-                    self.refresh_training_status()
+                    self.get_latest_status_message_logs_and_button_labels()
                 ),
                 outputs=[
                     status_box,
                     log_box,
-                     start_btn,
+                    start_btn,
                     stop_btn,
                     pause_resume_btn
                 ]
@@ -1239,7 +1203,7 @@ class VideoTrainerUI:
                 ]
             )
 
-            timer = gr.Timer(value=5)
+            timer = gr.Timer(value=6)
             timer.tick(
                 fn=lambda: self.update_titles(),
                 outputs=[

config.py

@@ -16,7 +16,8 @@ STAGING_PATH = STORAGE_PATH / "staging"                    # This is where files
 TRAINING_PATH = STORAGE_PATH / "training"                  # Folder containing the final training dataset
 TRAINING_VIDEOS_PATH = TRAINING_PATH / "videos"            # Captioned clips ready for training
 MODEL_PATH = STORAGE_PATH / "model"                        # Model checkpoints and files
-OUTPUT_PATH = STORAGE_PATH / "output"                      # Training outputs and logs
+OUTPUT_PATH = STORAGE_PATH / "output"                  # Training outputs and logs
+LOG_FILE_PATH = OUTPUT_PATH / "last_session.log"
 
 # On the production server we can afford to preload the big model
 PRELOAD_CAPTIONING_MODEL = parse_bool_env(os.environ.get('PRELOAD_CAPTIONING_MODEL'))
@@ -66,6 +67,8 @@ TRAINING_HEIGHT = 512 # 32 * 16
 # right now, finetrainers will crash if that happens, so the workaround is to have more buckets in here
         
 TRAINING_BUCKETS = [
+    (1, TRAINING_HEIGHT, TRAINING_WIDTH), #  1
+    (8 + 1, TRAINING_HEIGHT, TRAINING_WIDTH), # 8 + 1
     (8 * 2 + 1, TRAINING_HEIGHT, TRAINING_WIDTH), # 16 + 1
     (8 * 4 + 1, TRAINING_HEIGHT, TRAINING_WIDTH), # 32 + 1
     (8 * 6 + 1, TRAINING_HEIGHT, TRAINING_WIDTH), # 48 + 1

finetrainers/dataset.py

@@ -266,9 +266,9 @@ class ImageOrVideoDatasetWithResizing(ImageOrVideoDataset):
     def _preprocess_video(self, path: Path) -> torch.Tensor:
         video_reader = decord.VideoReader(uri=path.as_posix())
         video_num_frames = len(video_reader)
-        print(f"ImageOrVideoDatasetWithResizing: self.resolution_buckets = ", self.resolution_buckets)
-        print(f"ImageOrVideoDatasetWithResizing: self.max_num_frames = ", self.max_num_frames)
-        print(f"ImageOrVideoDatasetWithResizing: video_num_frames = ", video_num_frames)
+        #print(f"ImageOrVideoDatasetWithResizing: self.resolution_buckets = ", self.resolution_buckets)
+        #print(f"ImageOrVideoDatasetWithResizing: self.max_num_frames = ", self.max_num_frames)
+        #print(f"ImageOrVideoDatasetWithResizing: video_num_frames = ", video_num_frames)
         
         video_buckets = [bucket for bucket in self.resolution_buckets if bucket[0] <= video_num_frames]
         

training_log_parser.py

@@ -66,14 +66,18 @@ class TrainingLogParser:
         """Parse a single log line and update state"""
         try:
             # For debugging
-            logger.info(f"Parsing line: {line[:100]}...")
+            #logger.info(f"Parsing line: {line[:100]}...")
 
             # Training step progress line example:
             # Training steps:   1%|▏         | 1/70 [00:14<16:11, 14.08s/it, grad_norm=0.00789, step_loss=0.555, lr=3e-7]
+            
+            if ("Started training" in line) or (("Starting training" in line):
+                self.state.status = "training"
+            
             if "Training steps:" in line:
                 # Set status to training if we see this
                 self.state.status = "training"
-                print("setting status to 'training'")
+                #print("setting status to 'training'")
                 if not self.state.start_time:
                     self.state.start_time = datetime.now()
 

training_service.py

@@ -19,7 +19,7 @@ import select
 from typing import Any, Optional, Dict, List, Union, Tuple
 
 from huggingface_hub import upload_folder, create_repo
-from config import TrainingConfig, TRAINING_VIDEOS_PATH, STORAGE_PATH, TRAINING_PATH, MODEL_PATH, OUTPUT_PATH, HF_API_TOKEN, MODEL_TYPES
+from config import TrainingConfig, LOG_FILE_PATH, TRAINING_VIDEOS_PATH, STORAGE_PATH, TRAINING_PATH, MODEL_PATH, OUTPUT_PATH, HF_API_TOKEN, MODEL_TYPES
 from utils import make_archive, parse_training_log, is_image_file, is_video_file
 from finetrainers_utils import prepare_finetrainers_dataset, copy_files_to_training_dir
 
@@ -29,7 +29,7 @@ logging.basicConfig(
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
     handlers=[
         logging.StreamHandler(sys.stdout),
-        logging.FileHandler(str(OUTPUT_PATH / 'training_service.log'))
+        logging.FileHandler(str(LOG_FILE_PATH))
     ]
 )
 logger = logging.getLogger(__name__)