fixes for the dataset importer

app.py

@@ -15,7 +15,7 @@ from vms.config import (
     HF_API_TOKEN
 )
 
-from vms.ui.video_trainer_ui import VideoTrainerUI
+from vms.ui.app_ui import AppUI
 
 # Configure logging
 logger = logging.getLogger(__name__)
@@ -37,7 +37,7 @@ To avoid overpaying for your space, you can configure the auto-sleep settings to
         return app
 
     # Create the main application UI
-    ui = VideoTrainerUI()
+    ui = AppUI()
     app = ui.create_ui()
 
     return app

vms/ui/__init__.py

@@ -1,5 +1,5 @@
-from .video_trainer_ui import VideoTrainerUI
+from .app_ui import AppUI
 
 __all__ = [
-    'VideoTrainerUI',
+    'AppUI',
 ]

vms/ui/{video_trainer_ui.py → app_ui.py}

@@ -5,8 +5,7 @@ import logging
 import asyncio
 from typing import Any, Optional, Dict, List, Union, Tuple
 
-from ..services import TrainingService, CaptioningService, SplittingService, ImportingService, PreviewingService, MonitoringService
-from ..config import (
+from vms.config import (
     STORAGE_PATH, VIDEOS_TO_SPLIT_PATH, STAGING_PATH, OUTPUT_PATH,
     TRAINING_PATH, LOG_FILE_PATH, TRAINING_PRESETS, TRAINING_VIDEOS_PATH, MODEL_PATH, OUTPUT_PATH,
     MODEL_TYPES, SMALL_TRAINING_BUCKETS, TRAINING_TYPES,
@@ -22,13 +21,27 @@ from ..config import (
     DEFAULT_NB_TRAINING_STEPS,
     DEFAULT_NB_LR_WARMUP_STEPS
 )
-from ..utils import (
+from vms.utils import (
     get_recommended_precomputation_items,
     count_media_files,
     format_media_title,
     TrainingLogParser
 )
-from ..tabs import ImportTab, SplitTab, CaptionTab, TrainTab, MonitorTab, PreviewTab, ManageTab
+
+from vms.ui.project.services import (
+    TrainingService, CaptioningService, SplittingService, ImportingService, PreviewingService
+)
+from vms.ui.project.tabs import (
+    ImportTab, SplitTab, CaptionTab, TrainTab, PreviewTab, ManageTab
+)
+
+from vms.ui.monitoring.services import (
+    MonitoringService
+)
+
+from vms.ui.monitoring.tabs import (
+    GeneralTab
+)
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -36,22 +49,23 @@ logger.setLevel(logging.INFO)
 httpx_logger = logging.getLogger('httpx')
 httpx_logger.setLevel(logging.WARN)
 
-class VideoTrainerUI:       
+class AppUI:       
     def __init__(self):
         """Initialize services and tabs"""
-        # Initialize core services
+        # Project view
         self.training = TrainingService(self)
         self.splitting = SplittingService()
         self.importing = ImportingService()
         self.captioning = CaptioningService()
-        self.monitoring = MonitoringService()
         self.previewing = PreviewingService()
 
-        # Start the monitoring service on app creation
+        # Monitoring view
+        self.monitoring = MonitoringService()
         self.monitoring.start_monitoring()
     
         # Recovery status from any interrupted training
         recovery_result = self.training.recover_interrupted_training()
+
         # Add null check for recovery_result
         if recovery_result is None:
             recovery_result = {"status": "unknown", "ui_updates": {}}
@@ -67,9 +81,13 @@ class VideoTrainerUI:
             "recovery_result": recovery_result
         }
         
-        # Initialize tabs dictionary (will be populated in create_ui)
+        # Initialize tabs dictionary
         self.tabs = {}
-        self.tabs_component = None
+        self.project_tabs = {}
+        self.monitor_tabs = {}
+        self.main_tabs = None  # Main tabbed interface
+        self.project_tabs_component = None  # Project sub-tabs
+        self.monitor_tabs_component = None  # Monitor sub-tabs
 
         # Log recovery status
         logger.info(f"Initialization complete. Recovery status: {self.recovery_status}")
@@ -99,89 +117,153 @@ class VideoTrainerUI:
             logger.info(f"Added periodic callback {callback_fn.__name__} with interval {interval}s")
         except Exception as e:
             logger.error(f"Error adding periodic callback: {e}", exc_info=True)
+    
+    def switch_to_tab(self, tab_index: int):
+        """Switch to the specified tab index
+        
+        Args:
+            tab_index: Index of the tab to select (0 for Project, 1 for Monitor)
             
+        Returns:
+            Tab selection dictionary for Gradio
+        """
+        
+        return gr.Tabs(selected=tab_index)
+    
     def create_ui(self):
-        """Create the main Gradio UI"""
-        with gr.Blocks(title="🎥 Video Model Studio") as app:
-            gr.Markdown("# 🎥 Video Model Studio")
+        self.components = {}
+        """Create the main Gradio UI with tabbed navigation"""
+        with gr.Blocks(
+            title="🎞️ Video Model Studio",
+
+            # Let's hack Gradio!
+            css="#component-8 > .tab-wrapper{ display: none; }") as app:
+            self.app = app
             
-            # Create main tabs component
-            with gr.Tabs() as self.tabs_component:
-                # Initialize tab objects
-                self.tabs["import_tab"] = ImportTab(self)
-                self.tabs["split_tab"] = SplitTab(self)
-                self.tabs["caption_tab"] = CaptionTab(self)
-                self.tabs["train_tab"] = TrainTab(self)
-                self.tabs["monitor_tab"] = MonitorTab(self)
-                self.tabs["preview_tab"] = PreviewTab(self)
-                self.tabs["manage_tab"] = ManageTab(self)
-                
-                # Create tab UI components
-                for tab_id, tab_obj in self.tabs.items():
-                    tab_obj.create(self.tabs_component)
             
-            # Connect event handlers
+            # Main container with sidebar and tab area
+            with gr.Row():
+                # Sidebar for navigation
+                with gr.Sidebar(position="left", open=True):
+                    gr.Markdown("# 🎞️ Video Model Studio")
+                    self.components["current_project_btn"] = gr.Button("Current Project", variant="primary")
+                    self.components["system_monitoring_btn"] = gr.Button("System Monitoring")
+
+                # Main content area with tabs
+                with gr.Column():
+                    # Main tabbed interface for switching between Project and Monitor views
+                    with gr.Tabs() as main_tabs:
+                        self.main_tabs = main_tabs
+                        
+                        # Project View Tab
+                        with gr.Tab("📁 Current Project", id=0) as project_view:
+                            # Create project tabs
+                            with gr.Tabs() as project_tabs:
+                                # Store reference to project tabs component
+                                self.project_tabs_component = project_tabs
+                                
+                                # Initialize project tab objects
+                                self.project_tabs["import_tab"] = ImportTab(self)
+                                self.project_tabs["split_tab"] = SplitTab(self)
+                                self.project_tabs["caption_tab"] = CaptionTab(self)
+                                self.project_tabs["train_tab"] = TrainTab(self)
+                                self.project_tabs["preview_tab"] = PreviewTab(self)
+                                self.project_tabs["manage_tab"] = ManageTab(self)
+                                
+                                # Create tab UI components for project
+                                for tab_id, tab_obj in self.project_tabs.items():
+                                    tab_obj.create(project_tabs)
+                        
+                        # Monitoring View Tab
+                        with gr.Tab("📊 System Monitoring", id=1) as monitoring_view:
+                            # Create monitoring tabs
+                            with gr.Tabs() as monitoring_tabs:
+                                # Store reference to monitoring tabs component
+                                self.monitor_tabs_component = monitoring_tabs
+                                
+                                # Initialize monitoring tab objects
+                                self.monitor_tabs["general_tab"] = GeneralTab(self)
+                                
+                                # Create tab UI components for monitoring
+                                for tab_id, tab_obj in self.monitor_tabs.items():
+                                    tab_obj.create(monitoring_tabs)
+            
+            # Combine all tabs into a single dictionary for event handling
+            self.tabs = {**self.project_tabs, **self.monitor_tabs}
+
+            # Connect event handlers for all tabs - this must happen AFTER all tabs are created
             for tab_id, tab_obj in self.tabs.items():
                 tab_obj.connect_events()
             
             # app-level timers for auto-refresh functionality
             self._add_timers()
+
+            # Connect navigation events using tab switching
+            self.components["current_project_btn"].click(
+                fn=lambda: self.switch_to_tab(0),
+                outputs=[self.main_tabs],
+            )
+            
+            self.components["system_monitoring_btn"].click(
+                fn=lambda: self.switch_to_tab(1),
+                outputs=[self.main_tabs],
+            )
             
             # Initialize app state on load
             app.load(
                 fn=self.initialize_app_state,
                 outputs=[
-                    self.tabs["split_tab"].components["video_list"],
-                    self.tabs["caption_tab"].components["training_dataset"],
-                    self.tabs["train_tab"].components["start_btn"],
-                    self.tabs["train_tab"].components["stop_btn"],
-                    self.tabs["train_tab"].components["pause_resume_btn"],
-                    self.tabs["train_tab"].components["training_preset"],
-                    self.tabs["train_tab"].components["model_type"],
-                    self.tabs["train_tab"].components["training_type"],
-                    self.tabs["train_tab"].components["lora_rank"],
-                    self.tabs["train_tab"].components["lora_alpha"],
-                    self.tabs["train_tab"].components["train_steps"],
-                    self.tabs["train_tab"].components["batch_size"],
-                    self.tabs["train_tab"].components["learning_rate"],
-                    self.tabs["train_tab"].components["save_iterations"],
-                    self.tabs["train_tab"].components["current_task_box"],
-                    self.tabs["train_tab"].components["num_gpus"],
-                    self.tabs["train_tab"].components["precomputation_items"],
-                    self.tabs["train_tab"].components["lr_warmup_steps"]
+                    self.project_tabs["split_tab"].components["video_list"],
+                    self.project_tabs["caption_tab"].components["training_dataset"],
+                    self.project_tabs["train_tab"].components["start_btn"],
+                    self.project_tabs["train_tab"].components["stop_btn"],
+                    self.project_tabs["train_tab"].components["pause_resume_btn"],
+                    self.project_tabs["train_tab"].components["training_preset"],
+                    self.project_tabs["train_tab"].components["model_type"],
+                    self.project_tabs["train_tab"].components["training_type"],
+                    self.project_tabs["train_tab"].components["lora_rank"],
+                    self.project_tabs["train_tab"].components["lora_alpha"],
+                    self.project_tabs["train_tab"].components["train_steps"],
+                    self.project_tabs["train_tab"].components["batch_size"],
+                    self.project_tabs["train_tab"].components["learning_rate"],
+                    self.project_tabs["train_tab"].components["save_iterations"],
+                    self.project_tabs["train_tab"].components["current_task_box"],
+                    self.project_tabs["train_tab"].components["num_gpus"],
+                    self.project_tabs["train_tab"].components["precomputation_items"],
+                    self.project_tabs["train_tab"].components["lr_warmup_steps"]
                 ]
             )
-            
+        
         return app
-    
+        
     def _add_timers(self):
         """Add auto-refresh timers to the UI"""
         # Status update timer for text components (every 1 second)
         status_timer = gr.Timer(value=1)
         status_timer.tick(
-            fn=self.tabs["train_tab"].get_status_updates,  # Use a new function that returns appropriate updates
+            fn=self.project_tabs["train_tab"].get_status_updates,  # Use a new function that returns appropriate updates
             outputs=[
-                self.tabs["train_tab"].components["status_box"],
-                self.tabs["train_tab"].components["log_box"],
-                self.tabs["train_tab"].components["current_task_box"] if "current_task_box" in self.tabs["train_tab"].components else None
+                self.project_tabs["train_tab"].components["status_box"],
+                self.project_tabs["train_tab"].components["log_box"],
+                self.project_tabs["train_tab"].components["current_task_box"] if "current_task_box" in self.project_tabs["train_tab"].components else None
             ]
         )
         
         # Button update timer for button components (every 1 second)
         button_timer = gr.Timer(value=1)
         button_outputs = [
-            self.tabs["train_tab"].components["start_btn"],
-            self.tabs["train_tab"].components["stop_btn"]
+            self.project_tabs["train_tab"].components["start_btn"],
+            self.project_tabs["train_tab"].components["stop_btn"]
         ]
         
         # Add delete_checkpoints_btn or pause_resume_btn as the third button
-        if "delete_checkpoints_btn" in self.tabs["train_tab"].components:
-            button_outputs.append(self.tabs["train_tab"].components["delete_checkpoints_btn"])
-        elif "pause_resume_btn" in self.tabs["train_tab"].components:
-            button_outputs.append(self.tabs["train_tab"].components["pause_resume_btn"])
+        if "delete_checkpoints_btn" in self.project_tabs["train_tab"].components:
+            button_outputs.append(self.project_tabs["train_tab"].components["delete_checkpoints_btn"])
+        elif "pause_resume_btn" in self.project_tabs["train_tab"].components:
+            button_outputs.append(self.project_tabs["train_tab"].components["pause_resume_btn"])
         
         button_timer.tick(
-            fn=self.tabs["train_tab"].get_button_updates,  # Use a new function for button-specific updates
+            fn=self.project_tabs["train_tab"].get_button_updates,  # Use a new function for button-specific updates
             outputs=button_outputs
         )
         
@@ -190,8 +272,8 @@ class VideoTrainerUI:
         dataset_timer.tick(
             fn=self.refresh_dataset,
             outputs=[
-                self.tabs["split_tab"].components["video_list"],
-                self.tabs["caption_tab"].components["training_dataset"]
+                self.project_tabs["split_tab"].components["video_list"],
+                self.project_tabs["caption_tab"].components["training_dataset"]
             ]
         )
         
@@ -200,17 +282,17 @@ class VideoTrainerUI:
         titles_timer.tick(
             fn=self.update_titles,
             outputs=[
-                self.tabs["split_tab"].components["split_title"],
-                self.tabs["caption_tab"].components["caption_title"],
-                self.tabs["train_tab"].components["train_title"]
+                self.project_tabs["split_tab"].components["split_title"],
+                self.project_tabs["caption_tab"].components["caption_title"],
+                self.project_tabs["train_tab"].components["train_title"]
             ]
         )
     
     def initialize_app_state(self):
         """Initialize all app state in one function to ensure correct output count"""
         # Get dataset info
-        video_list = self.tabs["split_tab"].list_unprocessed_videos()
-        training_dataset = self.tabs["caption_tab"].list_training_files_to_caption()
+        video_list = self.project_tabs["split_tab"].list_unprocessed_videos()
+        training_dataset = self.project_tabs["caption_tab"].list_training_files_to_caption()
         
         # Get button states based on recovery status
         button_states = self.get_initial_button_states()
@@ -474,8 +556,8 @@ class VideoTrainerUI:
     
     def refresh_dataset(self):
         """Refresh all dynamic lists and training state"""
-        video_list = self.tabs["split_tab"].list_unprocessed_videos()
-        training_dataset = self.tabs["caption_tab"].list_training_files_to_caption()
+        video_list = self.project_tabs["split_tab"].list_unprocessed_videos()
+        training_dataset = self.project_tabs["caption_tab"].list_training_files_to_caption()
 
         return (
             video_list,

vms/ui/monitoring/services/__init__.py

@@ -0,0 +1,5 @@
+from .monitoring import MonitoringService
+
+__all__ = [
+    'MonitoringService',
+]

vms/ui/monitoring/tabs/__init__.py

@@ -0,0 +1,9 @@
+"""
+Tab components for the "monitor" view
+"""
+
+from .general_tab import GeneralTab
+
+__all__ = [
+    'GeneralTab'
+]

vms/{tabs/monitor_tab.py → ui/monitoring/tabs/general_tab.py}

@@ -12,39 +12,20 @@ import psutil
 from typing import Dict, Any, List, Optional, Tuple
 from datetime import datetime, timedelta
 
-from .base_tab import BaseTab
-from ..config import STORAGE_PATH
+from vms.utils.base_tab import BaseTab
+from vms.config import STORAGE_PATH
+from vms.ui.monitoring.utils import get_folder_size, human_readable_size
 
 logger = logging.getLogger(__name__)
 
-def get_folder_size(path):
-    """Calculate the total size of a folder in bytes"""
-    total_size = 0
-    for dirpath, dirnames, filenames in os.walk(path):
-        for filename in filenames:
-            file_path = os.path.join(dirpath, filename)
-            if not os.path.islink(file_path):  # Skip symlinks
-                total_size += os.path.getsize(file_path)
-    return total_size
 
-def human_readable_size(size_bytes):
-    """Convert a size in bytes to a human-readable string"""
-    if size_bytes == 0:
-        return "0 B"
-    size_names = ("B", "KB", "MB", "GB", "TB", "PB")
-    i = 0
-    while size_bytes >= 1024 and i < len(size_names) - 1:
-        size_bytes /= 1024
-        i += 1
-    return f"{size_bytes:.2f} {size_names[i]}"
-
-class MonitorTab(BaseTab):
-    """Monitor tab for system resource monitoring"""
+class GeneralTab(BaseTab):
+    """Monitor tab for general system resource monitoring"""
     
     def __init__(self, app_state):
         super().__init__(app_state)
-        self.id = "monitor_tab"
-        self.title = "5️⃣  Monitor"
+        self.id = "General_tab"
+        self.title = "General stats"
         self.refresh_interval = 8
     
     def create(self, parent=None) -> gr.TabItem:

vms/ui/monitoring/utils/__init__.py

@@ -0,0 +1,7 @@
+from .get_folder_size import get_folder_size
+from .human_readable_size import human_readable_size
+
+__all__ = [
+    'get_folder_size',
+    'human_readable_size',
+]

vms/ui/monitoring/utils/get_folder_size.py

@@ -0,0 +1,22 @@
+import os
+
+def get_folder_size(path):
+    """Calculate the total size of a folder in bytes"""
+    total_size = 0
+    for dirpath, dirnames, filenames in os.walk(path):
+        for filename in filenames:
+            file_path = os.path.join(dirpath, filename)
+            if not os.path.islink(file_path):  # Skip symlinks
+                total_size += os.path.getsize(file_path)
+    return total_size
+
+def human_readable_size(size_bytes):
+    """Convert a size in bytes to a human-readable string"""
+    if size_bytes == 0:
+        return "0 B"
+    size_names = ("B", "KB", "MB", "GB", "TB", "PB")
+    i = 0
+    while size_bytes >= 1024 and i < len(size_names) - 1:
+        size_bytes /= 1024
+        i += 1
+    return f"{size_bytes:.2f} {size_names[i]}"

vms/ui/monitoring/utils/human_readable_size.py

@@ -0,0 +1,10 @@
+def human_readable_size(size_bytes):
+    """Convert a size in bytes to a human-readable string"""
+    if size_bytes == 0:
+        return "0 B"
+    size_names = ("B", "KB", "MB", "GB", "TB", "PB")
+    i = 0
+    while size_bytes >= 1024 and i < len(size_names) - 1:
+        size_bytes /= 1024
+        i += 1
+    return f"{size_bytes:.2f} {size_names[i]}"

vms/{services → ui/project/services}/__init__.py

@@ -1,6 +1,5 @@
 from .captioning import CaptioningProgress, CaptioningService
 from .importing import ImportingService
-from .monitoring import MonitoringService
 from .splitting import SplittingService
 from .previewing import PreviewingService
 from .training import TrainingService
@@ -9,7 +8,6 @@ __all__ = [
     'CaptioningProgress',
     'CaptioningService',
     'ImportingService',
-    'MonitoringService',
     'SplittingService',
     'PreviewingService',
     'TrainingService',

vms/{services → ui/project/services}/captioning.py

@@ -17,8 +17,8 @@ from llava.mm_utils import tokenizer_image_token
 from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
 from llava.conversation import conv_templates, SeparatorStyle
 
-from ..config import TRAINING_VIDEOS_PATH, STAGING_PATH, PRELOAD_CAPTIONING_MODEL, CAPTIONING_MODEL, USE_MOCK_CAPTIONING_MODEL, DEFAULT_CAPTIONING_BOT_INSTRUCTIONS, VIDEOS_TO_SPLIT_PATH, DEFAULT_PROMPT_PREFIX
-from ..utils import extract_scene_info, is_image_file, is_video_file, copy_files_to_training_dir, prepare_finetrainers_dataset
+from vms.config import TRAINING_VIDEOS_PATH, STAGING_PATH, PRELOAD_CAPTIONING_MODEL, CAPTIONING_MODEL, USE_MOCK_CAPTIONING_MODEL, DEFAULT_CAPTIONING_BOT_INSTRUCTIONS, VIDEOS_TO_SPLIT_PATH, DEFAULT_PROMPT_PREFIX
+from vms.utils import extract_scene_info, is_image_file, is_video_file, copy_files_to_training_dir, prepare_finetrainers_dataset
 
 logger = logging.getLogger(__name__)
 

vms/{services → ui/project/services}/importing/hub_dataset.py

@@ -113,7 +113,7 @@ class HubDatasetBrowser:
             dataset_info = self.hf_api.dataset_info(dataset_id)
             
             # Format the information for display
-            info_text = f"## {dataset_info.id}\n\n"
+            info_text = f"### {dataset_info.id}\n\n"
             
             # Add description if available (with safer access)
             card_data = getattr(dataset_info, "card_data", None)
@@ -126,11 +126,11 @@ class HubDatasetBrowser:
                 info_text += f"{description[:500]}{'...' if len(description) > 500 else ''}\n\n"
             
             # Add basic stats (with safer access)
-            downloads = getattr(dataset_info, 'downloads', None)
-            info_text += f"## Downloads: {downloads if downloads is not None else 'N/A'}\n"
+            #downloads = getattr(dataset_info, 'downloads', None)
+            #info_text += f"## Downloads: {downloads if downloads is not None else 'N/A'}\n"
             
-            last_modified = getattr(dataset_info, 'last_modified', None)
-            info_text += f"## Last modified: {last_modified if last_modified is not None else 'N/A'}\n"
+            #last_modified = getattr(dataset_info, 'last_modified', None)
+            #info_text += f"## Last modified: {last_modified if last_modified is not None else 'N/A'}\n"
             
             # Group files by type
             file_groups = {

vms/{services → ui/project/services}/importing/import_service.py

@@ -10,11 +10,12 @@ import gradio as gr
 
 from huggingface_hub import HfApi
 
-from .file_upload import FileUploadHandler
-from .youtube import YouTubeDownloader
-from .hub_dataset import HubDatasetBrowser
 from vms.config import HF_API_TOKEN
 
+from vms.ui.project.services.importing.file_upload import FileUploadHandler
+from vms.ui.project.services.importing.youtube import YouTubeDownloader
+from vms.ui.project.services.importing.hub_dataset import HubDatasetBrowser
+
 logger = logging.getLogger(__name__)
 
 class ImportingService:

vms/{services → ui/project/services}/splitting.py

@@ -12,8 +12,8 @@ import gradio as gr
 from scenedetect import detect, ContentDetector, SceneManager, open_video
 from scenedetect.video_splitter import split_video_ffmpeg
 
-from ..config import TRAINING_PATH, STORAGE_PATH, TRAINING_VIDEOS_PATH, VIDEOS_TO_SPLIT_PATH, STAGING_PATH, DEFAULT_PROMPT_PREFIX
-from ..utils import remove_black_bars, extract_scene_info, is_video_file, is_image_file, add_prefix_to_caption
+from vms.config import TRAINING_PATH, STORAGE_PATH, TRAINING_VIDEOS_PATH, VIDEOS_TO_SPLIT_PATH, STAGING_PATH, DEFAULT_PROMPT_PREFIX
+from vms.utils import remove_black_bars, extract_scene_info, is_video_file, is_image_file, add_prefix_to_caption
 
 logger = logging.getLogger(__name__)
 

vms/{services → ui/project/services}/training.py

@@ -20,7 +20,7 @@ from typing import Any, Optional, Dict, List, Union, Tuple
 
 from huggingface_hub import upload_folder, create_repo
 
-from ..config import (
+from vms.config import (
     TrainingConfig, TRAINING_PRESETS, LOG_FILE_PATH, TRAINING_VIDEOS_PATH, 
     STORAGE_PATH, TRAINING_PATH, MODEL_PATH, OUTPUT_PATH, HF_API_TOKEN, 
     MODEL_TYPES, TRAINING_TYPES,
@@ -39,7 +39,7 @@ from ..config import (
     DEFAULT_NB_TRAINING_STEPS,
     DEFAULT_NB_LR_WARMUP_STEPS
 )
-from ..utils import (
+from vms.utils import (
     get_available_gpu_count,
     make_archive,
     parse_training_log,

vms/{tabs → ui/project/tabs}/__init__.py

@@ -1,23 +1,19 @@
 """
-Tab components for Video Model Studio UI
+Tab components for the "project" view
 """
 
-from .base_tab import BaseTab
 from .import_tab import ImportTab
 from .split_tab import SplitTab
 from .caption_tab import CaptionTab
 from .train_tab import TrainTab
-from .monitor_tab import MonitorTab
 from .preview_tab import PreviewTab
 from .manage_tab import ManageTab
 
 __all__ = [
-    'BaseTab',
     'ImportTab',
     'SplitTab',
     'CaptionTab',
     'TrainTab',
-    'MonitorTab',
     'PreviewTab',
     'ManageTab'
 ]

vms/{tabs → ui/project/tabs}/caption_tab.py

@@ -8,10 +8,10 @@ import asyncio
 import traceback
 from typing import Dict, Any, List, Optional, AsyncGenerator, Tuple
 from pathlib import Path
+import mimetypes
 
-from .base_tab import BaseTab
-from ..config import DEFAULT_CAPTIONING_BOT_INSTRUCTIONS, DEFAULT_PROMPT_PREFIX, STAGING_PATH, TRAINING_VIDEOS_PATH
-from ..utils import is_image_file, is_video_file, copy_files_to_training_dir
+from vms.utils import BaseTab, is_image_file, is_video_file, copy_files_to_training_dir
+from vms.config import DEFAULT_CAPTIONING_BOT_INSTRUCTIONS, DEFAULT_PROMPT_PREFIX, STAGING_PATH, TRAINING_VIDEOS_PATH
 
 logger = logging.getLogger(__name__)
 
@@ -533,9 +533,7 @@ class CaptionTab(BaseTab):
         Returns:
             Dict with preview content for each preview component
         """
-        import mimetypes
-        from ..config import TRAINING_VIDEOS_PATH
-        
+
         if not selected_text or "Caption:" in selected_text:
             return {
                 "video": None,

vms/{tabs → ui/project/tabs}/import_tab/hub_tab.py

@@ -10,7 +10,7 @@ import threading
 from pathlib import Path
 from typing import Dict, Any, List, Optional, Tuple
 
-from ..base_tab import BaseTab
+from vms.utils import BaseTab
 
 logger = logging.getLogger(__name__)
 
@@ -26,22 +26,32 @@ class HubTab(BaseTab):
     def create(self, parent=None) -> gr.Tab:
         """Create the Hub tab UI components"""
         with gr.Tab(self.title, id=self.id) as tab:
+
             with gr.Column():
                 with gr.Row():
-                    gr.Markdown("## Import from Hub datasets")
+                    gr.Markdown("## Import a dataset from Hugging Face")
                 
                 with gr.Row():
-                    gr.Markdown("Search for datasets with videos or WebDataset archives:")
-                
+                    with gr.Column():
+                        with gr.Row():
+                            gr.Markdown("You can use any dataset containing video files (.mp4) with optional captions (same names but in .txt format)")
+                        
+                        with gr.Row():
+                            gr.Markdown("You can also use a dataset containing WebDataset shards (.tar files).")
+
+                    with gr.Column():
+                        self.components["dataset_search"] = gr.Textbox(
+                            label="Search Hugging Face Datasets (MP4, WebDataset)",
+                            placeholder="video datasets eg. cakeify, disney, rickroll.."
+                        )
+                    
                 with gr.Row():
-                    self.components["dataset_search"] = gr.Textbox(
-                        label="Search Hugging Face Datasets (eg. cakeify, disney, rickroll..)",
-                        placeholder="Search for video datasets (eg. cakeify, disney, rickroll..)"
+                    self.components["dataset_search_btn"] = gr.Button(
+                        "Search Datasets",
+                        variant="primary",
+                        #size="md"
                     )
                 
-                with gr.Row():
-                    self.components["dataset_search_btn"] = gr.Button("Search Datasets", variant="primary")
-                
                 # Dataset browser results section
                 with gr.Row(visible=False) as dataset_results_row:
                     self.components["dataset_results_row"] = dataset_results_row
@@ -66,8 +76,6 @@ class HubTab(BaseTab):
                         with gr.Column(visible=False) as files_section:
                             self.components["files_section"] = files_section
                             
-                            gr.Markdown("## Files:")
-                            
                             # Video files row (appears if videos are present)
                             with gr.Row() as video_files_row:
                                 self.components["video_files_row"] = video_files_row

vms/{tabs → ui/project/tabs}/import_tab/import_tab.py

@@ -9,10 +9,10 @@ import threading
 from pathlib import Path
 from typing import Dict, Any, List, Optional, Tuple
 
-from ..base_tab import BaseTab
-from .upload_tab import UploadTab
-from .youtube_tab import YouTubeTab
-from .hub_tab import HubTab
+from vms.utils import BaseTab
+from vms.ui.project.tabs.import_tab.upload_tab import UploadTab
+from vms.ui.project.tabs.import_tab.youtube_tab import YouTubeTab
+from vms.ui.project.tabs.import_tab.hub_tab import HubTab
 
 from vms.config import (
     VIDEOS_TO_SPLIT_PATH, DEFAULT_PROMPT_PREFIX, DEFAULT_CAPTIONING_BOT_INSTRUCTIONS,
@@ -28,7 +28,8 @@ class ImportTab(BaseTab):
         super().__init__(app_state)
         self.id = "import_tab"
         self.title = "1️⃣  Import"
-        # Initialize sub-tabs
+        
+        # Initialize sub-tabs - these should be created first
         self.upload_tab = UploadTab(app_state)
         self.youtube_tab = YouTubeTab(app_state)
         self.hub_tab = HubTab(app_state)
@@ -53,7 +54,11 @@ class ImportTab(BaseTab):
                     visible=True,
                 )
             
-            # Create tabs for different import methods
+            # Create the import status textbox before creating the sub-tabs
+            with gr.Row():
+                self.components["import_status"] = gr.Textbox(label="Status", interactive=False)
+            
+            # Now create tabs for different import methods
             with gr.Tabs() as import_tabs:
                 # Create each sub-tab
                 self.upload_tab.create(import_tabs)
@@ -64,19 +69,20 @@ class ImportTab(BaseTab):
                 self.components["upload_tab"] = self.upload_tab
                 self.components["youtube_tab"] = self.youtube_tab
                 self.components["hub_tab"] = self.hub_tab
-            
-            with gr.Row():
-                self.components["import_status"] = gr.Textbox(label="Status", interactive=False)
 
         return tab
     
     def connect_events(self) -> None:
         """Connect event handlers to UI components"""
-        # Set shared components from parent tab to sub-tabs first
+        # Set shared components from parent tab to sub-tabs before connecting events
         for subtab in [self.upload_tab, self.youtube_tab, self.hub_tab]:
-            subtab.components["import_status"] = self.components["import_status"]
-            subtab.components["enable_automatic_video_split"] = self.components["enable_automatic_video_split"]
-            subtab.components["enable_automatic_content_captioning"] = self.components["enable_automatic_content_captioning"]
+            # Ensure these components exist in the parent before sharing them
+            if "import_status" in self.components:
+                subtab.components["import_status"] = self.components["import_status"]
+            if "enable_automatic_video_split" in self.components:
+                subtab.components["enable_automatic_video_split"] = self.components["enable_automatic_video_split"]
+            if "enable_automatic_content_captioning" in self.components:
+                subtab.components["enable_automatic_content_captioning"] = self.components["enable_automatic_content_captioning"]
             
         # Then connect events for each sub-tab
         self.upload_tab.connect_events()

vms/{tabs → ui/project/tabs}/import_tab/upload_tab.py

@@ -8,7 +8,7 @@ import logging
 from pathlib import Path
 from typing import Dict, Any, Optional
 
-from ..base_tab import BaseTab
+from vms.utils import BaseTab
 
 logger = logging.getLogger(__name__)
 
@@ -19,6 +19,12 @@ class UploadTab(BaseTab):
         super().__init__(app_state)
         self.id = "upload_tab"
         self.title = "Manual Upload"
+        # Initialize the components dictionary with None values for expected shared components
+        if "components" not in self.__dict__:
+            self.components = {}
+        self.components["import_status"] = None
+        self.components["enable_automatic_video_split"] = None
+        self.components["enable_automatic_content_captioning"] = None
     
     def create(self, parent=None) -> gr.Tab:
         """Create the Upload tab UI components"""
@@ -51,24 +57,50 @@ class UploadTab(BaseTab):
     
     def connect_events(self) -> None:
         """Connect event handlers to UI components"""
+        # Check if required shared components exist before connecting events
+        if not self.components.get("import_status"):
+            logger.warning("import_status component is not set in UploadTab")
+            return
+        
         # File upload event
-        self.components["files"].upload(
+        upload_event = self.components["files"].upload(
             fn=lambda x: self.app.importing.process_uploaded_files(x),
             inputs=[self.components["files"]],
-            outputs=[self.components["import_status"]]  # This comes from parent tab
-        ).success(
-            fn=self.app.tabs["import_tab"].update_titles_after_import,
-            inputs=[
-                self.components["enable_automatic_video_split"], 
-                self.components["enable_automatic_content_captioning"], 
-                self.app.tabs["caption_tab"].components["custom_prompt_prefix"]
-            ],
-            outputs=[
-                self.app.tabs_component,  # Main tabs component 
-                self.app.tabs["split_tab"].components["video_list"],
-                self.app.tabs["split_tab"].components["detect_status"],
-                self.app.tabs["split_tab"].components["split_title"],
-                self.app.tabs["caption_tab"].components["caption_title"],
-                self.app.tabs["train_tab"].components["train_title"]
-            ]
-        )
+            outputs=[self.components["import_status"]]
+        )
+        
+        # Only add success handler if all required components exist
+        if hasattr(self.app.tabs, "import_tab") and hasattr(self.app.tabs, "split_tab") and \
+           hasattr(self.app.tabs, "caption_tab") and hasattr(self.app.tabs, "train_tab"):
+            
+            # Get required components for success handler
+            try:
+                # If the components are missing, this will raise an AttributeError
+                tabs_component = self.app.tabs_component
+                video_list = self.app.tabs["split_tab"].components["video_list"]
+                detect_status = self.app.tabs["split_tab"].components["detect_status"]
+                split_title = self.app.tabs["split_tab"].components["split_title"]
+                caption_title = self.app.tabs["caption_tab"].components["caption_title"]
+                train_title = self.app.tabs["train_tab"].components["train_title"]
+                custom_prompt_prefix = self.app.tabs["caption_tab"].components["custom_prompt_prefix"]
+                
+                # Add success handler
+                upload_event.success(
+                    fn=self.app.tabs["import_tab"].update_titles_after_import,
+                    inputs=[
+                        self.components["enable_automatic_video_split"], 
+                        self.components["enable_automatic_content_captioning"], 
+                        custom_prompt_prefix
+                    ],
+                    outputs=[
+                        tabs_component,
+                        video_list,
+                        detect_status,
+                        split_title,
+                        caption_title,
+                        train_title
+                    ]
+                )
+            except (AttributeError, KeyError) as e:
+                logger.error(f"Error connecting event handlers in UploadTab: {str(e)}")
+                # Continue without the success handler

vms/{tabs → ui/project/tabs}/import_tab/youtube_tab.py

@@ -8,7 +8,7 @@ import logging
 from pathlib import Path
 from typing import Dict, Any, Optional
 
-from ..base_tab import BaseTab
+from vms.utils import BaseTab
 
 logger = logging.getLogger(__name__)
 
@@ -19,6 +19,12 @@ class YouTubeTab(BaseTab):
         super().__init__(app_state)
         self.id = "youtube_tab"
         self.title = "Download from YouTube"
+        # Initialize components that will be shared from parent
+        if "components" not in self.__dict__:
+            self.components = {}
+        self.components["import_status"] = None
+        self.components["enable_automatic_video_split"] = None
+        self.components["enable_automatic_content_captioning"] = None
     
     def create(self, parent=None) -> gr.Tab:
         """Create the YouTube tab UI components"""
@@ -47,21 +53,51 @@ class YouTubeTab(BaseTab):
     
     def connect_events(self) -> None:
         """Connect event handlers to UI components"""
+        # Check if required shared components exist before connecting events
+        if not self.components.get("import_status"):
+            logger.warning("import_status component is not set in YouTubeTab")
+            return
+        
+        if not self.components.get("enable_automatic_video_split"):
+            logger.warning("enable_automatic_video_split component is not set in YouTubeTab")
+            return
+            
+        if not self.components.get("enable_automatic_content_captioning"):
+            logger.warning("enable_automatic_content_captioning component is not set in YouTubeTab")
+            return
+            
+        # Only try to access custom_prompt_prefix if the caption_tab exists
+        custom_prompt_prefix = None
+        try:
+            if hasattr(self.app.tabs, "caption_tab") and "custom_prompt_prefix" in self.app.tabs["caption_tab"].components:
+                custom_prompt_prefix = self.app.tabs["caption_tab"].components["custom_prompt_prefix"]
+        except (AttributeError, KeyError):
+            logger.warning("Could not access custom_prompt_prefix component")
+            
         # YouTube download event
-        self.components["youtube_download_btn"].click(
+        download_event = self.components["youtube_download_btn"].click(
             fn=self.app.importing.download_youtube_video,
             inputs=[self.components["youtube_url"]],
-            outputs=[self.components["import_status"]]  # This comes from parent tab
-        ).success(
-            fn=self.app.tabs["import_tab"].on_import_success,
-            inputs=[
-                self.components["enable_automatic_video_split"],
-                self.components["enable_automatic_content_captioning"],
-                self.app.tabs["caption_tab"].components["custom_prompt_prefix"]
-            ],
-            outputs=[
-                self.app.tabs_component,
-                self.app.tabs["split_tab"].components["video_list"],
-                self.app.tabs["split_tab"].components["detect_status"]
-            ]
-        )
+            outputs=[self.components["import_status"]]
+        )
+        
+        # Add success handler if all components exist
+        if hasattr(self.app.tabs, "import_tab") and custom_prompt_prefix is not None:
+            try:
+                # Add the success handler
+                download_event.success(
+                    fn=self.app.tabs["import_tab"].on_import_success,
+                    inputs=[
+                        self.components["enable_automatic_video_split"],
+                        self.components["enable_automatic_content_captioning"],
+                        custom_prompt_prefix
+                    ],
+                    outputs=[
+                        self.app.tabs_component,
+                        self.app.tabs["split_tab"].components["video_list"],
+                        self.app.tabs["split_tab"].components["detect_status"]
+                    ]
+                )
+            except (AttributeError, KeyError) as e:
+                logger.error(f"Error connecting success handler in YouTubeTab: {str(e)}")
+                # Continue without the success handler

vms/{tabs → ui/project/tabs}/manage_tab.py

@@ -8,12 +8,11 @@ import shutil
 from pathlib import Path
 from typing import Dict, Any, List, Optional
 
-from .base_tab import BaseTab
-from ..config import (
+from vms.utils import BaseTab, validate_model_repo
+from vms.config import (
     HF_API_TOKEN, VIDEOS_TO_SPLIT_PATH, STAGING_PATH, TRAINING_VIDEOS_PATH, 
     TRAINING_PATH, MODEL_PATH, OUTPUT_PATH, LOG_FILE_PATH
 )
-from ..utils import validate_model_repo
 
 logger = logging.getLogger(__name__)
 
@@ -23,59 +22,64 @@ class ManageTab(BaseTab):
     def __init__(self, app_state):
         super().__init__(app_state)
         self.id = "manage_tab"
-        self.title = "7️⃣  Manage"
+        self.title = "6️⃣  Storage"
     
     def create(self, parent=None) -> gr.TabItem:
         """Create the Manage tab UI components"""
         with gr.TabItem(self.title, id=self.id) as tab:
-            with gr.Column():
-                with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("## Publishing")
-                        gr.Markdown("You model can be pushed to Hugging Face (this will use HF_API_TOKEN)")
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("## Download your model")
+                    gr.Markdown("There is currently a bug, you might have to click multiple times to trigger a download.")
 
-                        with gr.Row():
-                            with gr.Column():
-                                self.components["repo_id"] = gr.Textbox(
-                                    label="HuggingFace Model Repository",
-                                    placeholder="username/model-name",
-                                    info="The repository will be created if it doesn't exist"
-                                )
-                                self.components["make_public"] = gr.Checkbox(
-                                    label="Check this to make your model public (ie. visible and downloadable by anyone)",
-                                    info="You model is private by default"
-                                )
-                                self.components["push_model_btn"] = gr.Button(
-                                    "Push my model"
-                                )
+                    with gr.Row():
+                        self.components["download_dataset_btn"] = gr.DownloadButton(
+                            "Download training dataset",
+                            variant="secondary",
+                            size="lg"
+                        )
+                        self.components["download_model_btn"] = gr.DownloadButton(
+                            "Download model weights",
+                            variant="secondary",
+                            size="lg"
+                        )
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("## Publish your model")
+                    gr.Markdown("You model can be pushed to Hugging Face (this will use HF_API_TOKEN)")
 
-                with gr.Row():
-                    with gr.Column():
-                        with gr.Row():
-                            with gr.Column():
-                                gr.Markdown("## Storage management")
-                                with gr.Row():
-                                    self.components["download_dataset_btn"] = gr.DownloadButton(
-                                        "Download dataset (click again if DL doesn't start)",
-                                        variant="secondary",
-                                        size="lg"
-                                    )
-                                    self.components["download_model_btn"] = gr.DownloadButton(
-                                        "Download model (click again if DL doesn't start)",
-                                        variant="secondary",
-                                        size="lg"
-                                    )
+            with gr.Row():
+                with gr.Column():
+                    self.components["repo_id"] = gr.Textbox(
+                        label="HuggingFace Model Repository",
+                        placeholder="username/model-name",
+                        info="The repository will be created if it doesn't exist"
+                    )
+                    self.components["make_public"] = gr.Checkbox(
+                        label="Check this to make your model public (ie. visible and downloadable by anyone)",
+                        info="You model is private by default"
+                    )
+                    self.components["push_model_btn"] = gr.Button(
+                        "Push my model"
+                    )
 
-                        with gr.Row():
-                            self.components["global_stop_btn"] = gr.Button(
-                                "Stop everything and delete my data",
-                                variant="stop"
-                            )
-                            self.components["global_status"] = gr.Textbox(
-                                label="Global Status",
-                                interactive=False,
-                                visible=False
-                            )
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("## Delete your model")
+                    gr.Markdown("If something went wrong, you can trigger a full reset (model shutdown + data destruction).")
+                    gr.Markdown("Make sure you have made a backup first.")
+                    gr.Markdown("If you are deleting because of a bug, remember you can use the Developer Mode on HF to inspect the working directory (in /data or .data)")
+
+            with gr.Row():
+                self.components["global_stop_btn"] = gr.Button(
+                    "Stop everything and delete my data",
+                    variant="stop"
+                )
+                self.components["global_status"] = gr.Textbox(
+                    label="Global Status",
+                    interactive=False,
+                    visible=False
+                )
         
         return tab
     

vms/{tabs → ui/project/tabs}/preview_tab.py

@@ -8,8 +8,8 @@ from pathlib import Path
 from typing import Dict, Any, List, Optional, Tuple
 import time
 
-from .base_tab import BaseTab
-from ..config import (
+from vms.utils import BaseTab
+from vms.config import (
     MODEL_TYPES, DEFAULT_PROMPT_PREFIX
 )
 
@@ -21,13 +21,13 @@ class PreviewTab(BaseTab):
     def __init__(self, app_state):
         super().__init__(app_state)
         self.id = "preview_tab"
-        self.title = "6️⃣  Preview"
+        self.title = "5️⃣  Preview"
          
     def create(self, parent=None) -> gr.TabItem:
         """Create the Preview tab UI components"""
         with gr.TabItem(self.title, id=self.id) as tab:
             with gr.Row():
-                gr.Markdown("## Test Your Trained Model")
+                gr.Markdown("## Preview your model")
             
             with gr.Row():
                 with gr.Column(scale=2):

vms/{tabs → ui/project/tabs}/split_tab.py

@@ -6,7 +6,7 @@ import gradio as gr
 import logging
 from typing import Dict, Any, List, Optional
 
-from .base_tab import BaseTab
+from vms.utils import BaseTab
 
 logger = logging.getLogger(__name__)
 

vms/{tabs → ui/project/tabs}/train_tab.py

@@ -8,8 +8,8 @@ import os
 from typing import Dict, Any, List, Optional, Tuple
 from pathlib import Path
 
-from .base_tab import BaseTab
-from ..config import (
+from vms.utils import BaseTab
+from vms.config import (
     TRAINING_PRESETS, OUTPUT_PATH, MODEL_TYPES, ASK_USER_TO_DUPLICATE_SPACE, SMALL_TRAINING_BUCKETS, TRAINING_TYPES,
     DEFAULT_NB_TRAINING_STEPS, DEFAULT_SAVE_CHECKPOINT_EVERY_N_STEPS,
     DEFAULT_BATCH_SIZE, DEFAULT_CAPTION_DROPOUT_P,

vms/utils/__init__.py

@@ -9,6 +9,7 @@ from .finetrainers_utils import prepare_finetrainers_dataset, copy_files_to_trai
 from . import webdataset_handler
 
 from .gpu_detector import get_available_gpu_count, get_gpu_info, get_recommended_precomputation_items
+from .base_tab import BaseTab
 
 __all__ = [
     'validate_model_repo',
@@ -39,5 +40,7 @@ __all__ = [
 
     'get_available_gpu_count',
     'get_gpu_info',
-    'get_recommended_precomputation_items'
+    'get_recommended_precomputation_items',
+
+    'BaseTab',
 ]

vms/{tabs → utils}/base_tab.py

@@ -15,7 +15,7 @@ class BaseTab:
         """Initialize the tab with app state reference
         
         Args:
-            app_state: Reference to main VideoTrainerUI instance
+            app_state: Reference to main AppUI instance
         """
         self.app = app_state
         self.components = {}