fixed ui persistence issues

vms/ui/app_ui.py

@@ -32,7 +32,7 @@ from vms.ui.project.services import (
     TrainingService, CaptioningService, SplittingService, ImportingService, PreviewingService
 )
 from vms.ui.project.tabs import (
-    ImportTab, SplitTab, CaptionTab, TrainTab, PreviewTab, ManageTab
+    ImportTab, CaptionTab, TrainTab, PreviewTab, ManageTab
 )
 
 from vms.ui.monitoring.services import (
@@ -164,7 +164,6 @@ class AppUI:
                                 
                                 # 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)
@@ -213,7 +212,6 @@ class AppUI:
             app.load(
                 fn=self.initialize_app_state,
                 outputs=[
-                    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"],
@@ -273,7 +271,6 @@ class AppUI:
         dataset_timer.tick(
             fn=self.refresh_dataset,
             outputs=[
-                self.project_tabs["split_tab"].components["video_list"],
                 self.project_tabs["caption_tab"].components["training_dataset"]
             ]
         )
@@ -283,7 +280,6 @@ class AppUI:
         titles_timer.tick(
             fn=self.update_titles,
             outputs=[
-                self.project_tabs["split_tab"].components["split_title"],
                 self.project_tabs["caption_tab"].components["caption_title"],
                 self.project_tabs["train_tab"].components["train_title"]
             ]
@@ -292,7 +288,6 @@ class AppUI:
     def initialize_app_state(self):
         """Initialize all app state in one function to ensure correct output count"""
         # Get dataset info
-        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
@@ -381,17 +376,40 @@ class AppUI:
         
         # Get model_version value
         model_version_val = ""
+
         # First get the internal model type for the currently selected model
         model_internal_type = MODEL_TYPES.get(model_type_val)
+        logger.info(f"Initializing model version for model_type: {model_type_val} (internal: {model_internal_type})")
+
         if model_internal_type and model_internal_type in MODEL_VERSIONS:
-            # If there's a saved model_version and it's valid for this model type
-            if "model_version" in ui_state and ui_state["model_version"] in MODEL_VERSIONS.get(model_internal_type, {}):
+            # Get available versions for this model type as simple strings
+            available_model_versions = list(MODEL_VERSIONS.get(model_internal_type, {}).keys())
+            
+            # Log for debugging
+            logger.info(f"Available versions: {available_model_versions}")
+            
+            # Set model_version_val to saved value if valid, otherwise first available
+            if "model_version" in ui_state and ui_state["model_version"] in available_model_versions:
                 model_version_val = ui_state["model_version"]
-            else:
-                # Otherwise use the first available version
-                versions = list(MODEL_VERSIONS.get(model_internal_type, {}).keys())
-                if versions:
-                    model_version_val = versions[0]
+                logger.info(f"Using saved model version: {model_version_val}")
+            elif available_model_versions:
+                model_version_val = available_model_versions[0]
+                logger.info(f"Using first available model version: {model_version_val}")
+            
+            # IMPORTANT: Update the dropdown choices directly in the UI component
+            # This is essential to avoid the error when loading the UI
+            try:
+                self.project_tabs["train_tab"].components["model_version"].choices = available_model_versions
+                logger.info(f"Updated model_version dropdown choices: {len(available_model_versions)} options")
+            except Exception as e:
+                logger.error(f"Error updating model_version dropdown: {str(e)}")
+        else:
+            logger.warning(f"No versions available for model type: {model_type_val}")
+            # Set empty choices to avoid errors
+            try:
+                self.project_tabs["train_tab"].components["model_version"].choices = []
+            except Exception as e:
+                logger.error(f"Error setting empty model_version choices: {str(e)}")
                     
         # Ensure training_type is a valid display name
         training_type_val = ui_state.get("training_type", list(TRAINING_TYPES.keys())[0])
@@ -444,7 +462,6 @@ class AppUI:
         
         # Return all values in the exact order expected by outputs
         return (
-            video_list, 
             training_dataset,
             start_btn, 
             stop_btn, 
@@ -464,7 +481,7 @@ class AppUI:
             precomputation_items_val,
             lr_warmup_steps_val
         )
-
+    
     def initialize_ui_from_state(self):
         """Initialize UI components from saved state"""
         ui_state = self.load_ui_values()
@@ -558,12 +575,6 @@ class AppUI:
         Returns:
             Dict of Gradio updates
         """
-        # Count files for splitting
-        split_videos, _, split_size = count_media_files(VIDEOS_TO_SPLIT_PATH)
-        split_title = format_media_title(
-            "split", split_videos, 0, split_size
-        )
-        
         # Count files for captioning
         caption_videos, caption_images, caption_size = count_media_files(STAGING_PATH)
         caption_title = format_media_title(
@@ -577,17 +588,14 @@ class AppUI:
         )
         
         return (
-            gr.Markdown(value=split_title),
             gr.Markdown(value=caption_title),
             gr.Markdown(value=f"{train_title} available for training")
         )
     
     def refresh_dataset(self):
         """Refresh all dynamic lists and training state"""
-        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,
             training_dataset
         )

vms/ui/project/tabs/__init__.py

@@ -3,7 +3,6 @@ Tab components for the "project" view
 """
 
 from .import_tab import ImportTab
-from .split_tab import SplitTab
 from .caption_tab import CaptionTab
 from .train_tab import TrainTab
 from .preview_tab import PreviewTab
@@ -11,7 +10,6 @@ from .manage_tab import ManageTab
 
 __all__ = [
     'ImportTab',
-    'SplitTab',
     'CaptionTab',
     'TrainTab',
     'PreviewTab',

vms/ui/project/tabs/caption_tab.py

@@ -21,7 +21,7 @@ class CaptionTab(BaseTab):
     def __init__(self, app_state):
         super().__init__(app_state)
         self.id = "caption_tab"
-        self.title = "3️⃣  Caption"
+        self.title = "2️⃣ Caption"
         self._should_stop_captioning = False
     
     def create(self, parent=None) -> gr.TabItem:

vms/ui/project/tabs/import_tab/hub_tab.py

@@ -124,6 +124,14 @@ class HubTab(BaseTab):
             ]
         )
         
+        # Check if we have access to project_tabs_component
+        if hasattr(self.app, "project_tabs_component"):
+            tabs_component = self.app.project_tabs_component
+        else:
+            # Fallback to prevent errors
+            logger.warning("project_tabs_component not found in app, using None for tab switching")
+            tabs_component = None
+        
         # Download videos button
         self.components["download_videos_btn"].click(
             fn=self.set_file_type_and_return,
@@ -142,6 +150,17 @@ class HubTab(BaseTab):
                 self.components["download_webdataset_btn"],
                 self.components["download_in_progress"]
             ]
+        ).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=[
+                tabs_component,
+                self.components["status_output"]
+            ]
         )
         
         # Download WebDataset button
@@ -162,6 +181,17 @@ class HubTab(BaseTab):
                 self.components["download_webdataset_btn"],
                 self.components["download_in_progress"]
             ]
+        ).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=[
+                tabs_component,
+                self.components["status_output"]
+            ]
         )
 
     def set_file_type_and_return(self):

vms/ui/project/tabs/import_tab/import_tab.py

@@ -6,6 +6,7 @@ import gradio as gr
 import logging
 import asyncio
 import threading
+import shutil
 from pathlib import Path
 from typing import Dict, Any, List, Optional, Tuple
 
@@ -91,27 +92,35 @@ class ImportTab(BaseTab):
     
     def on_import_success(self, enable_splitting, enable_automatic_content_captioning, prompt_prefix):
         """Handle successful import of files"""
-        videos = self.app.tabs["split_tab"].list_unprocessed_videos()
-        
-        # If scene detection isn't already running and there are videos to process,
-        # and auto-splitting is enabled, start the detection
-        if videos and not self.app.splitting.is_processing() and enable_splitting:
-            # Start the scene detection in a separate thread
-            self._start_scene_detection_bg(enable_splitting)
-            msg = "Starting automatic scene detection..."
+        # If splitting is disabled, we need to directly move videos to staging
+        if not enable_splitting:
+            # Copy files without splitting
+            self._start_copy_to_staging_bg()
+            msg = "Copying videos to staging directory without splitting..."
         else:
-            # Just copy files without splitting if auto-split disabled
-            self._start_copy_files_bg(enable_splitting)
-            msg = "Copying videos without splitting..."
-        
-        self.app.tabs["caption_tab"].copy_files_to_training_dir(prompt_prefix)
+            # Start scene detection if not already running and there are videos to process
+            if not self.app.splitting.is_processing():
+                # Start the scene detection in a separate thread
+                self._start_scene_detection_bg(enable_splitting)
+                msg = "Starting automatic scene detection..."
+            else:
+                msg = "Scene detection already running..."
 
+        # Copy files to training directory
+        self.app.tabs["caption_tab"].copy_files_to_training_dir(prompt_prefix)
+        
         # Start auto-captioning if enabled
         if enable_automatic_content_captioning:
             self._start_captioning_bg(DEFAULT_CAPTIONING_BOT_INSTRUCTIONS, prompt_prefix)
         
-        # Return the correct tuple of values as expected by the UI
-        return gr.update(selected="split_tab"), videos, msg
+        # Check if we have access to project_tabs_component for tab switching
+        if hasattr(self.app, "project_tabs_component") and self.app.project_tabs_component is not None:
+            # Now redirect to the caption tab instead of split tab
+            return gr.update(selected="caption_tab"), msg
+        else:
+            # If no tabs component is available, just return the message
+            logger.warning("Cannot switch tabs - project_tabs_component not available")
+            return None, msg
     
     def _start_scene_detection_bg(self, enable_splitting):
         """Start scene detection in a background thread"""
@@ -120,7 +129,7 @@ class ImportTab(BaseTab):
             asyncio.set_event_loop(loop)
             try:
                 loop.run_until_complete(
-                    self.app.tabs["split_tab"].start_scene_detection(enable_splitting)
+                    self.app.splitting.start_processing(enable_splitting)
                 )
             except Exception as e:
                 logger.error(f"Error in background scene detection: {str(e)}", exc_info=True)
@@ -131,21 +140,48 @@ class ImportTab(BaseTab):
         thread.daemon = True
         thread.start()
     
-    def _start_copy_files_bg(self, enable_splitting):
-        """Start copying files in a background thread"""
+    def _start_copy_to_staging_bg(self):
+        """Start copying files directly to staging directory in a background thread"""
         def run_async_in_thread():
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
             try:
-                async def copy_files():
-                    for video_file in VIDEOS_TO_SPLIT_PATH.glob("*.mp4"):
-                        await self.app.splitting.process_video(video_file, enable_splitting=False)
+                # Copy all videos from videos_to_split to staging without scene detection
+                for video_file in VIDEOS_TO_SPLIT_PATH.glob("*.mp4"):
+                    try:
+                        # Ensure unique filename in staging directory
+                        target_path = STAGING_PATH / video_file.name
+                        counter = 1
+                        
+                        while target_path.exists():
+                            stem = video_file.stem
+                            if "___" in stem:
+                                base_stem = stem.split("___")[0]
+                            else:
+                                base_stem = stem
+                            target_path = STAGING_PATH / f"{base_stem}___{counter}{video_file.suffix}"
+                            counter += 1
+                        
+                        # Copy the video file to staging
+                        shutil.copy2(video_file, target_path)
+                        logger.info(f"Copied video directly to staging: {video_file.name} -> {target_path.name}")
+                        
+                        # Copy caption file if it exists
+                        caption_path = video_file.with_suffix('.txt')
+                        if caption_path.exists():
+                            shutil.copy2(caption_path, target_path.with_suffix('.txt'))
+                            logger.info(f"Copied caption for {video_file.name}")
+                        
+                        # Remove original after successful copy
+                        video_file.unlink()
+                        if caption_path.exists():
+                            caption_path.unlink()
+                            
+                        gr.Info(f"Imported {video_file.name} directly to staging")
+                        
+                    except Exception as e:
+                        logger.error(f"Error copying {video_file.name} to staging: {str(e)}", exc_info=True)
                 
-                loop.run_until_complete(copy_files())
             except Exception as e:
-                logger.error(f"Error in background file copying: {str(e)}", exc_info=True)
-            finally:
-                loop.close()
+                logger.error(f"Error in background file copying to staging: {str(e)}", exc_info=True)
         
         thread = threading.Thread(target=run_async_in_thread)
         thread.daemon = True
@@ -174,7 +210,7 @@ class ImportTab(BaseTab):
     async def update_titles_after_import(self, enable_splitting, enable_automatic_content_captioning, prompt_prefix):
         """Handle post-import updates including titles"""
         # Call the non-async version since we need to return immediately for the UI
-        tabs, video_list, detect_status = self.on_import_success(
+        tabs, status_msg = self.on_import_success(
             enable_splitting, enable_automatic_content_captioning, prompt_prefix
         )
         
@@ -182,4 +218,4 @@ class ImportTab(BaseTab):
         titles = self.app.update_titles()
         
         # Return all expected outputs
-        return tabs, video_list, detect_status, *titles
+        return tabs, status_msg, *titles

vms/ui/project/tabs/import_tab/upload_tab.py

@@ -70,16 +70,17 @@ class UploadTab(BaseTab):
         )
         
         # 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"):
+        if hasattr(self.app.tabs, "import_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"]
+                if hasattr(self.app, "project_tabs_component"):
+                    tabs_component = self.app.project_tabs_component
+                else:
+                    logger.warning("project_tabs_component not found in app, using None for tab switching")
+                    tabs_component = None
+                    
                 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"]
@@ -94,9 +95,7 @@ class UploadTab(BaseTab):
                     ],
                     outputs=[
                         tabs_component,
-                        video_list,
-                        detect_status,
-                        split_title,
+                        self.components["import_status"],
                         caption_title,
                         train_title
                     ]

vms/ui/project/tabs/import_tab/youtube_tab.py

@@ -74,6 +74,13 @@ class YouTubeTab(BaseTab):
         except (AttributeError, KeyError):
             logger.warning("Could not access custom_prompt_prefix component")
             
+        # Check if we have access to project_tabs_component
+        if hasattr(self.app, "project_tabs_component"):
+            tabs_component = self.app.project_tabs_component
+        else:
+            logger.warning("project_tabs_component not found in app, using None for tab switching")
+            tabs_component = None
+            
         # YouTube download event
         download_event = self.components["youtube_download_btn"].click(
             fn=self.app.importing.download_youtube_video,
@@ -93,9 +100,8 @@ class YouTubeTab(BaseTab):
                         custom_prompt_prefix
                     ],
                     outputs=[
-                        self.app.tabs_component,
-                        self.app.tabs["split_tab"].components["video_list"],
-                        self.app.tabs["split_tab"].components["detect_status"]
+                        tabs_component,
+                        self.components["import_status"]
                     ]
                 )
             except (AttributeError, KeyError) as e:

vms/ui/project/tabs/manage_tab.py

@@ -22,7 +22,7 @@ class ManageTab(BaseTab):
     def __init__(self, app_state):
         super().__init__(app_state)
         self.id = "manage_tab"
-        self.title = "6️⃣  Storage"
+        self.title = "5️⃣ Storage"
     
     def create(self, parent=None) -> gr.TabItem:
         """Create the Manage tab UI components"""
@@ -108,11 +108,9 @@ class ManageTab(BaseTab):
             fn=self.handle_global_stop,
             outputs=[
                 self.components["global_status"],
-                self.app.tabs["split_tab"].components["video_list"],
                 self.app.tabs["caption_tab"].components["training_dataset"],
                 self.app.tabs["train_tab"].components["status_box"],
                 self.app.tabs["train_tab"].components["log_box"],
-                self.app.tabs["split_tab"].components["detect_status"],
                 self.app.tabs["import_tab"].components["import_status"],
                 self.app.tabs["caption_tab"].components["preview_status"]
             ]
@@ -169,11 +167,9 @@ class ManageTab(BaseTab):
         
         return {
             self.components["global_status"]: gr.update(value=full_status, visible=True),
-            self.app.tabs["split_tab"].components["video_list"]: videos,
             self.app.tabs["caption_tab"].components["training_dataset"]: clips,
             self.app.tabs["train_tab"].components["status_box"]: "Training stopped and data cleared",
             self.app.tabs["train_tab"].components["log_box"]: "",
-            self.app.tabs["split_tab"].components["detect_status"]: "Scene detection stopped",
             self.app.tabs["import_tab"].components["import_status"]: "All data cleared",
             self.app.tabs["caption_tab"].components["preview_status"]: "Captioning stopped"
         }

vms/ui/project/tabs/preview_tab.py

@@ -23,7 +23,7 @@ class PreviewTab(BaseTab):
     def __init__(self, app_state):
         super().__init__(app_state)
         self.id = "preview_tab"
-        self.title = "5️⃣  Preview"
+        self.title = "4️⃣ Preview"
          
     def create(self, parent=None) -> gr.TabItem:
         """Create the Preview tab UI components"""
@@ -193,26 +193,31 @@ class PreviewTab(BaseTab):
         """Get model version choices based on model type"""
         # Convert UI display name to internal name
         internal_type = MODEL_TYPES.get(model_type)
-        if not internal_type:
-            return []
-            
-        # Get versions from preview service
-        versions = self.app.previewing.get_model_versions(internal_type)
-        if not versions:
+        if not internal_type or internal_type not in MODEL_VERSIONS:
+            logger.warning(f"No model versions found for {model_type} (internal type: {internal_type})")
             return []
             
-        # Format choices with display name and description
-        choices = []
-        for model_id, info in versions.items():
-            choices.append(f"{model_id} - {info.get('name', '')}")
+        # Return just the model IDs as a list of simple strings
+        version_ids = list(MODEL_VERSIONS.get(internal_type, {}).keys())
+        logger.info(f"Found {len(version_ids)} versions for {model_type}: {version_ids}")
+        return version_ids
             
-        return choices
-    
     def get_default_model_version(self, model_type: str) -> str:
-        """Get default model version for the model type"""
-        choices = self.get_model_version_choices(model_type)
-        if choices:
-            return choices[0]
+        """Get default model version for the given model type"""
+        # Convert UI display name to internal name
+        internal_type = MODEL_TYPES.get(model_type)
+        logger.debug(f"get_default_model_version({model_type}) = {internal_type}")
+        
+        if not internal_type or internal_type not in MODEL_VERSIONS:
+            logger.warning(f"No valid model versions found for {model_type}")
+            return ""
+            
+        # Get the first version available for this model type
+        versions = list(MODEL_VERSIONS.get(internal_type, {}).keys())
+        if versions:
+            default_version = versions[0]
+            logger.debug(f"Default version for {model_type}: {default_version}")
+            return default_version
         return ""
 
     def get_default_model_type(self) -> str:

vms/ui/project/tabs/split_tab.py

@@ -1,81 +0,0 @@
-"""
-Split tab for Video Model Studio UI
-"""
-
-import gradio as gr
-import logging
-from typing import Dict, Any, List, Optional
-
-from vms.utils import BaseTab
-
-logger = logging.getLogger(__name__)
-
-class SplitTab(BaseTab):
-    """Split tab for scene detection and video splitting"""
-    
-    def __init__(self, app_state):
-        super().__init__(app_state)
-        self.id = "split_tab"
-        self.title = "2️⃣  Split"
-    
-    def create(self, parent=None) -> gr.TabItem:
-        """Create the Split tab UI components"""
-        with gr.TabItem(self.title, id=self.id) as tab:
-            with gr.Row():
-                self.components["split_title"] = gr.Markdown("## Splitting of 0 videos (0 bytes)")
-            
-            with gr.Row():
-                with gr.Column():
-                    self.components["detect_btn"] = gr.Button("Split videos into single-camera shots", variant="primary")
-                    self.components["detect_status"] = gr.Textbox(label="Status", interactive=False)
-
-                with gr.Column():
-                    self.components["video_list"] = gr.Dataframe(
-                        headers=["name", "status"],
-                        label="Videos to split (note: Nvidia A100 cannot split videos encoded in AV1)",
-                        interactive=False,
-                        wrap=True
-                    )
-                    
-        return tab
-    
-    def connect_events(self) -> None:
-        """Connect event handlers to UI components"""
-        # Scene detection button event
-        self.components["detect_btn"].click(
-            fn=self.start_scene_detection,
-            inputs=[self.app.tabs["import_tab"].components["enable_automatic_video_split"]],
-            outputs=[self.components["detect_status"]]
-        )
-        
-    def refresh(self) -> Dict[str, Any]:
-        """Refresh the video list with current data"""
-        videos = self.list_unprocessed_videos()
-        return {
-            "video_list": videos
-        }
-    
-    def list_unprocessed_videos(self) -> gr.Dataframe:
-        """Update list of unprocessed videos"""
-        videos = self.app.splitting.list_unprocessed_videos()
-        # videos is already in [[name, status]] format from splitting_service
-        return gr.Dataframe(
-            headers=["name", "status"],
-            value=videos,
-            interactive=False
-        )
-
-    async def start_scene_detection(self, enable_splitting: bool) -> str:
-        """Start background scene detection process
-        
-        Args:
-            enable_splitting: Whether to split videos into scenes
-        """
-        if self.app.splitting.is_processing():
-            return "Scene detection already running"
-            
-        try:
-            await self.app.splitting.start_processing(enable_splitting)
-            return "Scene detection completed"
-        except Exception as e:
-            return f"Error during scene detection: {str(e)}"

vms/ui/project/tabs/train_tab.py

@@ -35,7 +35,7 @@ class TrainTab(BaseTab):
     def __init__(self, app_state):
         super().__init__(app_state)
         self.id = "train_tab"
-        self.title = "4️⃣  Train"
+        self.title = "3️⃣ Train"
     
     def create(self, parent=None) -> gr.TabItem:
         """Create the Train tab UI components"""
@@ -58,23 +58,46 @@ class TrainTab(BaseTab):
                         with gr.Column():
                             # Get the default model type from the first preset
                             default_model_type = list(MODEL_TYPES.keys())[0]
-                            
+
                             self.components["model_type"] = gr.Dropdown(
                                 choices=list(MODEL_TYPES.keys()),
                                 label="Model Type",
                                 value=default_model_type,
                                 interactive=True
                             )
-                            
+
                             # Get model versions for the default model type
                             default_model_versions = self.get_model_version_choices(default_model_type)
                             default_model_version = self.get_default_model_version(default_model_type)
-                            print(f"default_model_version(default_model_type) = {default_model_version}")
+
+                            # Ensure default_model_versions is not empty
+                            if not default_model_versions:
+                                # If no versions found for default model, use a fallback
+                                internal_type = MODEL_TYPES.get(default_model_type)
+                                if internal_type in MODEL_VERSIONS:
+                                    default_model_versions = list(MODEL_VERSIONS[internal_type].keys())
+                                else:
+                                    # Last resort - collect all available versions from all models
+                                    default_model_versions = []
+                                    for model_versions in MODEL_VERSIONS.values():
+                                        default_model_versions.extend(list(model_versions.keys()))
+                                    
+                                # If still empty, provide a placeholder
+                                if not default_model_versions:
+                                    default_model_versions = ["-- No versions available --"]
+                                    
+                                # Set default version to first in list if available
+                                if default_model_versions:
+                                    default_model_version = default_model_versions[0]
+                                else:
+                                    default_model_version = ""
+
                             self.components["model_version"] = gr.Dropdown(
                                 choices=default_model_versions,
                                 label="Model Version",
                                 value=default_model_version,
-                                interactive=True
+                                interactive=True,
+                                allow_custom_value=True  # Add this to avoid errors with custom values
                             )
                             
                             self.components["training_type"] = gr.Dropdown(
@@ -428,15 +451,38 @@ class TrainTab(BaseTab):
     
     def update_model_versions(self, model_type: str) -> Dict:
         """Update model version choices based on selected model type"""
-        model_versions = self.get_model_version_choices(model_type)
-        default_version = self.get_default_model_version(model_type)
-        #print(f"update_model_versions({model_type}): default_version = {default_version}")
-        # Update UI state with proper model_type first (add this line)
-        self.app.update_ui_state(model_type=model_type)
-        
-        # Update the model_version dropdown with new choices and default value
-        return gr.Dropdown(choices=model_versions, value=default_version)
-
+        try:
+            # Get version choices for this model type
+            model_versions = self.get_model_version_choices(model_type)
+            
+            # Get default version
+            default_version = self.get_default_model_version(model_type)
+            logger.info(f"update_model_versions({model_type}): default_version = {default_version}, available versions: {model_versions}")
+            
+            # Update UI state with proper model_type first
+            self.app.update_ui_state(model_type=model_type)
+            
+            # Create a new dropdown with the updated choices
+            if not model_versions:
+                logger.warning(f"No model versions available for {model_type}, using empty list")
+                # Return empty dropdown to avoid errors
+                return gr.Dropdown(choices=[], value=None)
+                
+            # Ensure default_version is in model_versions
+            if default_version not in model_versions and model_versions:
+                default_version = model_versions[0]
+                logger.info(f"Default version not in choices, using first available: {default_version}")
+            
+            # Return the updated dropdown
+            logger.info(f"Returning dropdown with {len(model_versions)} choices")
+            return gr.Dropdown(choices=model_versions, value=default_version)
+        except Exception as e:
+            # Log any exceptions for debugging
+            logger.error(f"Error in update_model_versions: {str(e)}")
+            # Return empty dropdown to avoid errors
+            return gr.Dropdown(choices=[], value=None)
+            
+    
     def handle_training_start(
         self, preset, model_type, model_version, training_type, 
         lora_rank, lora_alpha, train_steps, batch_size, learning_rate, 
@@ -509,26 +555,30 @@ class TrainTab(BaseTab):
         # Convert UI display name to internal name
         internal_type = MODEL_TYPES.get(model_type)
         if not internal_type or internal_type not in MODEL_VERSIONS:
+            logger.warning(f"No model versions found for {model_type} (internal type: {internal_type})")
             return []
             
-        # Return just the model IDs without formatting
-        return list(MODEL_VERSIONS.get(internal_type, {}).keys())
+        # Return just the model IDs as a list of simple strings
+        version_ids = list(MODEL_VERSIONS.get(internal_type, {}).keys())
+        logger.info(f"Found {len(version_ids)} versions for {model_type}: {version_ids}")
+        return version_ids
 
-    
     def get_default_model_version(self, model_type: str) -> str:
         """Get default model version for the given model type"""
         # Convert UI display name to internal name
         internal_type = MODEL_TYPES.get(model_type)
-        #print(f"get_default_model_version({model_type}) = {internal_type}")
+        logger.debug(f"get_default_model_version({model_type}) = {internal_type}")
+        
         if not internal_type or internal_type not in MODEL_VERSIONS:
+            logger.warning(f"No valid model versions found for {model_type}")
             return ""
             
         # Get the first version available for this model type
-        versions = MODEL_VERSIONS.get(internal_type, {})
+        versions = list(MODEL_VERSIONS.get(internal_type, {}).keys())
         if versions:
-            model_versions = list(versions.keys())
-            if model_versions:
-                return model_versions[0]
+            default_version = versions[0]
+            logger.debug(f"Default version for {model_type}: {default_version}")
+            return default_version
         return ""
    
     def update_model_info(self, model_type: str, training_type: str) -> Dict:
@@ -698,7 +748,19 @@ class TrainTab(BaseTab):
         # Get the appropriate model version for the selected model type
         model_versions = self.get_model_version_choices(model_display_name)
         default_model_version = self.get_default_model_version(model_display_name)
-        
+
+        # Create the model version dropdown update
+        model_version_update = gr.Dropdown(choices=model_versions, value=default_model_version)
+
+        # Ensure we have valid choices and values
+        if not model_versions:
+            logger.warning(f"No versions found for {model_display_name}, using empty list")
+            model_versions = []
+            default_model_version = None
+        elif default_model_version not in model_versions and model_versions:
+            default_model_version = model_versions[0]
+            logger.info(f"Reset default version to first available: {default_model_version}")
+
         # Return values in the same order as the output components
         return (
             model_display_name,
@@ -714,7 +776,7 @@ class TrainTab(BaseTab):
             num_gpus_val,
             precomputation_items_val,
             lr_warmup_steps_val,
-            gr.Dropdown(choices=model_versions, value=default_model_version)
+            model_version_update,
         )