fix

vms/ui/app_ui.py

@@ -392,7 +392,7 @@ class AppUI:
                 versions = list(MODEL_VERSIONS.get(model_internal_type, {}).keys())
                 if versions:
                     model_version_val = versions[0]
-
+                    
         # Ensure training_type is a valid display name
         training_type_val = ui_state.get("training_type", list(TRAINING_TYPES.keys())[0])
         if training_type_val not in TRAINING_TYPES:

vms/ui/project/services/training.py

@@ -14,6 +14,7 @@ import zipfile
 import logging
 import traceback
 import threading
+import fcntl
 import select
 
 from typing import Any, Optional, Dict, List, Union, Tuple
@@ -63,6 +64,8 @@ class TrainingService:
         self.pid_file = OUTPUT_PATH / "training.pid"
         self.log_file = OUTPUT_PATH / "training.log"
 
+        self.file_lock = threading.Lock()
+
         self.file_handler = None
         self.setup_logging()
         self.ensure_valid_ui_state_file()
@@ -131,67 +134,69 @@ class TrainingService:
         """Save current UI state to file with validation"""
         ui_state_file = OUTPUT_PATH / "ui_state.json"
         
-        # Validate values before saving
-        validated_values = {}
-        default_state = {
-            "model_type": list(MODEL_TYPES.keys())[0],
-            "model_version": "",
-            "training_type": list(TRAINING_TYPES.keys())[0],
-            "lora_rank": DEFAULT_LORA_RANK_STR,
-            "lora_alpha": DEFAULT_LORA_ALPHA_STR, 
-            "train_steps": DEFAULT_NB_TRAINING_STEPS,
-            "batch_size": DEFAULT_BATCH_SIZE,
-            "learning_rate": DEFAULT_LEARNING_RATE,
-            "save_iterations": DEFAULT_SAVE_CHECKPOINT_EVERY_N_STEPS,
-            "training_preset": list(TRAINING_PRESETS.keys())[0],
-            "num_gpus": DEFAULT_NUM_GPUS,
-            "precomputation_items": DEFAULT_PRECOMPUTATION_ITEMS,
-            "lr_warmup_steps": DEFAULT_NB_LR_WARMUP_STEPS
-        }
-        
-        # Copy default values first
-        validated_values = default_state.copy()
-        
-        # Update with provided values, converting types as needed
-        for key, value in values.items():
-            if key in default_state:
-                if key == "train_steps":
-                    try:
-                        validated_values[key] = int(value)
-                    except (ValueError, TypeError):
-                        validated_values[key] = default_state[key]
-                elif key == "batch_size":
-                    try:
-                        validated_values[key] = int(value)
-                    except (ValueError, TypeError):
-                        validated_values[key] = default_state[key]
-                elif key == "learning_rate":
-                    try:
-                        validated_values[key] = float(value)
-                    except (ValueError, TypeError):
+        # Use a lock to prevent concurrent writes
+        with self.file_lock:
+            # Validate values before saving
+            validated_values = {}
+            default_state = {
+                "model_type": list(MODEL_TYPES.keys())[0],
+                "model_version": "",
+                "training_type": list(TRAINING_TYPES.keys())[0],
+                "lora_rank": DEFAULT_LORA_RANK_STR,
+                "lora_alpha": DEFAULT_LORA_ALPHA_STR, 
+                "train_steps": DEFAULT_NB_TRAINING_STEPS,
+                "batch_size": DEFAULT_BATCH_SIZE,
+                "learning_rate": DEFAULT_LEARNING_RATE,
+                "save_iterations": DEFAULT_SAVE_CHECKPOINT_EVERY_N_STEPS,
+                "training_preset": list(TRAINING_PRESETS.keys())[0],
+                "num_gpus": DEFAULT_NUM_GPUS,
+                "precomputation_items": DEFAULT_PRECOMPUTATION_ITEMS,
+                "lr_warmup_steps": DEFAULT_NB_LR_WARMUP_STEPS
+            }
+            
+            # Copy default values first
+            validated_values = default_state.copy()
+            
+            # Update with provided values, converting types as needed
+            for key, value in values.items():
+                if key in default_state:
+                    if key == "train_steps":
+                        try:
+                            validated_values[key] = int(value)
+                        except (ValueError, TypeError):
+                            validated_values[key] = default_state[key]
+                    elif key == "batch_size":
+                        try:
+                            validated_values[key] = int(value)
+                        except (ValueError, TypeError):
+                            validated_values[key] = default_state[key]
+                    elif key == "learning_rate":
+                        try:
+                            validated_values[key] = float(value)
+                        except (ValueError, TypeError):
+                            validated_values[key] = default_state[key]
+                    elif key == "save_iterations":
+                        try:
+                            validated_values[key] = int(value)
+                        except (ValueError, TypeError):
+                            validated_values[key] = default_state[key]
+                    elif key == "lora_rank" and value not in ["16", "32", "64", "128", "256", "512", "1024"]:
                         validated_values[key] = default_state[key]
-                elif key == "save_iterations":
-                    try:
-                        validated_values[key] = int(value)
-                    except (ValueError, TypeError):
+                    elif key == "lora_alpha" and value not in ["16", "32", "64", "128", "256", "512", "1024"]:
                         validated_values[key] = default_state[key]
-                elif key == "lora_rank" and value not in ["16", "32", "64", "128", "256", "512", "1024"]:
-                    validated_values[key] = default_state[key]
-                elif key == "lora_alpha" and value not in ["16", "32", "64", "128", "256", "512", "1024"]:
-                    validated_values[key] = default_state[key]
-                else:
-                    validated_values[key] = value
-        
-        try:
-            # First verify we can serialize to JSON
-            json_data = json.dumps(validated_values, indent=2)
+                    else:
+                        validated_values[key] = value
             
-            # Write to the file
-            with open(ui_state_file, 'w') as f:
-                f.write(json_data)
-            logger.debug(f"UI state saved successfully")
-        except Exception as e:
-            logger.error(f"Error saving UI state: {str(e)}")
+            try:
+                # First verify we can serialize to JSON
+                json_data = json.dumps(validated_values, indent=2)
+                
+                # Write to the file
+                with open(ui_state_file, 'w') as f:
+                    f.write(json_data)
+                logger.debug(f"UI state saved successfully")
+            except Exception as e:
+                logger.error(f"Error saving UI state: {str(e)}")
 
     def _backup_and_recreate_ui_state(self, ui_state_file, default_state):
         """Backup the corrupted UI state file and create a new one with defaults"""
@@ -229,130 +234,133 @@ class TrainingService:
             "lr_warmup_steps": DEFAULT_NB_LR_WARMUP_STEPS
         }
         
-        if not ui_state_file.exists():
-            logger.info("UI state file does not exist, using default values")
-            return default_state
-                
-        try:
-            # First check if the file is empty
-            file_size = ui_state_file.stat().st_size
-            if file_size == 0:
-                logger.warning("UI state file exists but is empty, using default values")
+        # Use lock for reading too to avoid reading during a write
+        with self.file_lock:
+
+            if not ui_state_file.exists():
+                logger.info("UI state file does not exist, using default values")
                 return default_state
-                
-            with open(ui_state_file, 'r') as f:
-                file_content = f.read().strip()
-                if not file_content:
-                    logger.warning("UI state file is empty or contains only whitespace, using default values")
-                    return default_state
                     
-                try:
-                    saved_state = json.loads(file_content)
-                except json.JSONDecodeError as e:
-                    logger.error(f"Error parsing UI state JSON: {str(e)}")
-                    # Instead of showing the error, recreate the file with defaults
-                    self._backup_and_recreate_ui_state(ui_state_file, default_state)
+            try:
+                # First check if the file is empty
+                file_size = ui_state_file.stat().st_size
+                if file_size == 0:
+                    logger.warning("UI state file exists but is empty, using default values")
                     return default_state
-                
-                # Clean up model type if it contains " (LoRA)" suffix
-                if "model_type" in saved_state and " (LoRA)" in saved_state["model_type"]:
-                    saved_state["model_type"] = saved_state["model_type"].replace(" (LoRA)", "")
-                    logger.info(f"Removed (LoRA) suffix from saved model type: {saved_state['model_type']}")
-
-                # Convert numeric values to appropriate types
-                if "train_steps" in saved_state:
-                    try:
-                        saved_state["train_steps"] = int(saved_state["train_steps"])
-                    except (ValueError, TypeError):
-                        saved_state["train_steps"] = default_state["train_steps"]
-                        logger.warning("Invalid train_steps value, using default")
-                        
-                if "batch_size" in saved_state:
-                    try:
-                        saved_state["batch_size"] = int(saved_state["batch_size"])
-                    except (ValueError, TypeError):
-                        saved_state["batch_size"] = default_state["batch_size"]
-                        logger.warning("Invalid batch_size value, using default")
-                        
-                if "learning_rate" in saved_state:
-                    try:
-                        saved_state["learning_rate"] = float(saved_state["learning_rate"])
-                    except (ValueError, TypeError):
-                        saved_state["learning_rate"] = default_state["learning_rate"]
-                        logger.warning("Invalid learning_rate value, using default")
+                    
+                with open(ui_state_file, 'r') as f:
+                    file_content = f.read().strip()
+                    if not file_content:
+                        logger.warning("UI state file is empty or contains only whitespace, using default values")
+                        return default_state
                         
-                if "save_iterations" in saved_state:
                     try:
-                        saved_state["save_iterations"] = int(saved_state["save_iterations"])
-                    except (ValueError, TypeError):
-                        saved_state["save_iterations"] = default_state["save_iterations"]
-                        logger.warning("Invalid save_iterations value, using default")
+                        saved_state = json.loads(file_content)
+                    except json.JSONDecodeError as e:
+                        logger.error(f"Error parsing UI state JSON: {str(e)}")
+                        # Instead of showing the error, recreate the file with defaults
+                        self._backup_and_recreate_ui_state(ui_state_file, default_state)
+                        return default_state
                     
-                # Make sure we have all keys (in case structure changed)
-                merged_state = default_state.copy()
-                merged_state.update({k: v for k, v in saved_state.items() if v is not None})
-                
-                # Validate model_type is in available choices
-                if merged_state["model_type"] not in MODEL_TYPES:
-                    # Try to map from internal name
-                    model_found = False
-                    for display_name, internal_name in MODEL_TYPES.items():
-                        if internal_name == merged_state["model_type"]:
-                            merged_state["model_type"] = display_name
-                            model_found = True
-                            break
-                    # If still not found, use default
-                    if not model_found:
-                        merged_state["model_type"] = default_state["model_type"]
-                        logger.warning(f"Invalid model type in saved state, using default")
+                    # Clean up model type if it contains " (LoRA)" suffix
+                    if "model_type" in saved_state and " (LoRA)" in saved_state["model_type"]:
+                        saved_state["model_type"] = saved_state["model_type"].replace(" (LoRA)", "")
+                        logger.info(f"Removed (LoRA) suffix from saved model type: {saved_state['model_type']}")
+
+                    # Convert numeric values to appropriate types
+                    if "train_steps" in saved_state:
+                        try:
+                            saved_state["train_steps"] = int(saved_state["train_steps"])
+                        except (ValueError, TypeError):
+                            saved_state["train_steps"] = default_state["train_steps"]
+                            logger.warning("Invalid train_steps value, using default")
                             
-                # Validate model_version is appropriate for model_type
-                if "model_type" in merged_state and "model_version" in merged_state:
-                    model_internal_type = MODEL_TYPES.get(merged_state["model_type"])
-                    if model_internal_type:
-                        valid_versions = MODEL_VERSIONS.get(model_internal_type, {}).keys()
-                        if merged_state["model_version"] not in valid_versions:
-                            # Set to default for this model type
-                            from vms.ui.project.tabs.train_tab import TrainTab
-                            train_tab = TrainTab(None)  # Temporary instance just for the helper method
-                            merged_state["model_version"] = train_tab.get_default_model_version(saved_state["model_type"])
-                            logger.warning(f"Invalid model version for {merged_state['model_type']}, using default")
-                
-                # Validate training_type is in available choices
-                if merged_state["training_type"] not in TRAINING_TYPES:
-                    # Try to map from internal name
-                    training_found = False
-                    for display_name, internal_name in TRAINING_TYPES.items():
-                        if internal_name == merged_state["training_type"]:
-                            merged_state["training_type"] = display_name
-                            training_found = True
-                            break
-                    # If still not found, use default
-                    if not training_found:
-                        merged_state["training_type"] = default_state["training_type"]
-                        logger.warning(f"Invalid training type in saved state, using default")
-                
-                # Validate training_preset is in available choices
-                if merged_state["training_preset"] not in TRAINING_PRESETS:
-                    merged_state["training_preset"] = default_state["training_preset"]
-                    logger.warning(f"Invalid training preset in saved state, using default")
+                    if "batch_size" in saved_state:
+                        try:
+                            saved_state["batch_size"] = int(saved_state["batch_size"])
+                        except (ValueError, TypeError):
+                            saved_state["batch_size"] = default_state["batch_size"]
+                            logger.warning("Invalid batch_size value, using default")
+                            
+                    if "learning_rate" in saved_state:
+                        try:
+                            saved_state["learning_rate"] = float(saved_state["learning_rate"])
+                        except (ValueError, TypeError):
+                            saved_state["learning_rate"] = default_state["learning_rate"]
+                            logger.warning("Invalid learning_rate value, using default")
+                            
+                    if "save_iterations" in saved_state:
+                        try:
+                            saved_state["save_iterations"] = int(saved_state["save_iterations"])
+                        except (ValueError, TypeError):
+                            saved_state["save_iterations"] = default_state["save_iterations"]
+                            logger.warning("Invalid save_iterations value, using default")
+                        
+                    # Make sure we have all keys (in case structure changed)
+                    merged_state = default_state.copy()
+                    merged_state.update({k: v for k, v in saved_state.items() if v is not None})
                     
-                # Validate lora_rank is in allowed values
-                if merged_state.get("lora_rank") not in ["16", "32", "64", "128", "256", "512", "1024"]:
-                    merged_state["lora_rank"] = default_state["lora_rank"]
-                    logger.warning(f"Invalid lora_rank in saved state, using default")
+                    # Validate model_type is in available choices
+                    if merged_state["model_type"] not in MODEL_TYPES:
+                        # Try to map from internal name
+                        model_found = False
+                        for display_name, internal_name in MODEL_TYPES.items():
+                            if internal_name == merged_state["model_type"]:
+                                merged_state["model_type"] = display_name
+                                model_found = True
+                                break
+                        # If still not found, use default
+                        if not model_found:
+                            merged_state["model_type"] = default_state["model_type"]
+                            logger.warning(f"Invalid model type in saved state, using default")
+                                
+                    # Validate model_version is appropriate for model_type
+                    if "model_type" in merged_state and "model_version" in merged_state:
+                        model_internal_type = MODEL_TYPES.get(merged_state["model_type"])
+                        if model_internal_type:
+                            valid_versions = MODEL_VERSIONS.get(model_internal_type, {}).keys()
+                            if merged_state["model_version"] not in valid_versions:
+                                # Set to default for this model type
+                                from vms.ui.project.tabs.train_tab import TrainTab
+                                train_tab = TrainTab(None)  # Temporary instance just for the helper method
+                                merged_state["model_version"] = train_tab.get_default_model_version(saved_state["model_type"])
+                                logger.warning(f"Invalid model version for {merged_state['model_type']}, using default")
                     
-                # Validate lora_alpha is in allowed values
-                if merged_state.get("lora_alpha") not in ["16", "32", "64", "128", "256", "512", "1024"]:
-                    merged_state["lora_alpha"] = default_state["lora_alpha"]
-                    logger.warning(f"Invalid lora_alpha in saved state, using default")
+                    # Validate training_type is in available choices
+                    if merged_state["training_type"] not in TRAINING_TYPES:
+                        # Try to map from internal name
+                        training_found = False
+                        for display_name, internal_name in TRAINING_TYPES.items():
+                            if internal_name == merged_state["training_type"]:
+                                merged_state["training_type"] = display_name
+                                training_found = True
+                                break
+                        # If still not found, use default
+                        if not training_found:
+                            merged_state["training_type"] = default_state["training_type"]
+                            logger.warning(f"Invalid training type in saved state, using default")
                     
-                return merged_state
-        except Exception as e:
-            logger.error(f"Error loading UI state: {str(e)}")
-            # If anything goes wrong, backup and recreate
-            self._backup_and_recreate_ui_state(ui_state_file, default_state)
-            return default_state
+                    # Validate training_preset is in available choices
+                    if merged_state["training_preset"] not in TRAINING_PRESETS:
+                        merged_state["training_preset"] = default_state["training_preset"]
+                        logger.warning(f"Invalid training preset in saved state, using default")
+                        
+                    # Validate lora_rank is in allowed values
+                    if merged_state.get("lora_rank") not in ["16", "32", "64", "128", "256", "512", "1024"]:
+                        merged_state["lora_rank"] = default_state["lora_rank"]
+                        logger.warning(f"Invalid lora_rank in saved state, using default")
+                        
+                    # Validate lora_alpha is in allowed values
+                    if merged_state.get("lora_alpha") not in ["16", "32", "64", "128", "256", "512", "1024"]:
+                        merged_state["lora_alpha"] = default_state["lora_alpha"]
+                        logger.warning(f"Invalid lora_alpha in saved state, using default")
+                        
+                    return merged_state
+            except Exception as e:
+                logger.error(f"Error loading UI state: {str(e)}")
+                # If anything goes wrong, backup and recreate
+                self._backup_and_recreate_ui_state(ui_state_file, default_state)
+                return default_state
 
     def ensure_valid_ui_state_file(self):
         """Ensure UI state file exists and is valid JSON"""

vms/ui/project/tabs/preview_tab.py

@@ -298,7 +298,7 @@ class PreviewTab(BaseTab):
         # Update model_version choices when model_type changes or tab is selected
         if hasattr(self.app, 'tabs_component') and self.app.tabs_component is not None:
             self.app.tabs_component.select(
-                fn=self.sync_model_type_and_verions,
+                fn=self.sync_model_type_and_versions,
                 inputs=[],
                 outputs=[
                     self.components["model_type"],
@@ -391,7 +391,7 @@ class PreviewTab(BaseTab):
             self.components["conditioning_image"]: gr.Image(visible=show_conditioning_image)
         }
     
-    def sync_model_type_and_verions(self) -> Tuple[str, str]:
+    def sync_model_type_and_versions(self) -> Tuple[str, str]:
         """Sync model type with training tab when preview tab is selected and update model version choices"""
         model_type = self.get_default_model_type()
         model_version = ""
@@ -401,19 +401,15 @@ class PreviewTab(BaseTab):
         preview_state = ui_state.get("preview", {})
         model_version = preview_state.get("model_version", "")
         
+        # If no model version specified or invalid, use default
         if not model_version:
-            # Format it as a display choice
+            # Get the internal model type
             internal_type = MODEL_TYPES.get(model_type)
             if internal_type and internal_type in MODEL_VERSIONS:
-                first_version = next(iter(MODEL_VERSIONS[internal_type].keys()), "")
-                if first_version:
-                    model_version_info = MODEL_VERSIONS[internal_type][first_version]
-                    model_version = f"{first_version} - {model_version_info.get('name', '')}"
+                versions = list(MODEL_VERSIONS[internal_type].keys())
+                if versions:
+                    model_version = versions[0]
         
-        # If we couldn't get it, use default
-        if not model_version:
-            model_version = self.get_default_model_version(model_type)
-            
         return model_type, model_version
     
     def update_resolution(self, preset: str) -> Tuple[int, int, float]:

vms/ui/project/tabs/train_tab.py

@@ -69,7 +69,7 @@ class TrainTab(BaseTab):
                             # 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}")
                             self.components["model_version"] = gr.Dropdown(
                                 choices=default_model_versions,
                                 label="Model Version",
@@ -214,6 +214,37 @@ class TrainTab(BaseTab):
                     
         return tab
     
+    def update_model_type_and_version(self, model_type: str, model_version: str):
+        """Update both model type and version together to keep them in sync"""
+        # Get internal model type
+        internal_type = MODEL_TYPES.get(model_type)
+        
+        # Make sure model_version is valid for this model type
+        if internal_type and internal_type in MODEL_VERSIONS:
+            valid_versions = list(MODEL_VERSIONS[internal_type].keys())
+            if not model_version or model_version not in valid_versions:
+                if valid_versions:
+                    model_version = valid_versions[0]
+        
+        # Update UI state with both values to keep them in sync
+        self.app.update_ui_state(model_type=model_type, model_version=model_version)
+        return None
+
+    def save_model_version(self, model_type: str, model_version: str):
+        """Save model version ensuring it's consistent with model type"""
+        internal_type = MODEL_TYPES.get(model_type)
+        
+        # Verify the model_version is compatible with the current model_type
+        if internal_type and internal_type in MODEL_VERSIONS:
+            valid_versions = MODEL_VERSIONS[internal_type].keys()
+            if model_version not in valid_versions:
+                # Don't save incompatible version
+                return None
+                
+        # Save the model version along with current model type to ensure consistency
+        self.app.update_ui_state(model_type=model_type, model_version=model_version)
+        return None
+
     def connect_events(self) -> None:
         """Connect event handlers to UI components"""
         # Model type change event - Update model version dropdown choices
@@ -222,8 +253,8 @@ class TrainTab(BaseTab):
             inputs=[self.components["model_type"]],
             outputs=[self.components["model_version"]]
         ).then(
-            fn=lambda v: self.app.update_ui_state(model_type=v),
-            inputs=[self.components["model_type"]],
+            fn=self.update_model_type_and_version,  # Add this new function
+            inputs=[self.components["model_type"], self.components["model_version"]],
             outputs=[]
         ).then(
             # Use get_model_info instead of update_model_info
@@ -234,8 +265,8 @@ class TrainTab(BaseTab):
         
         # Model version change event
         self.components["model_version"].change(
-            fn=lambda v: self.app.update_ui_state(model_version=v),
-            inputs=[self.components["model_version"]],
+            fn=self.save_model_version,  # Replace with this new function
+            inputs=[self.components["model_type"], self.components["model_version"]],
             outputs=[]
         )
             
@@ -399,10 +430,13 @@ class TrainTab(BaseTab):
         """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)
-        
+
     def handle_training_start(
         self, preset, model_type, model_version, training_type, 
         lora_rank, lora_alpha, train_steps, batch_size, learning_rate, 
@@ -477,22 +511,24 @@ class TrainTab(BaseTab):
         if not internal_type or internal_type not in MODEL_VERSIONS:
             return []
             
-        # Get versions and return them as choices
-        versions = MODEL_VERSIONS.get(internal_type, {})
-        return list(versions.keys())
+        # Return just the model IDs without formatting
+        return list(MODEL_VERSIONS.get(internal_type, {}).keys())
+
     
     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}")
         if not internal_type or internal_type not in MODEL_VERSIONS:
             return ""
             
         # Get the first version available for this model type
         versions = MODEL_VERSIONS.get(internal_type, {})
         if versions:
-            return next(iter(versions.keys()))
-            
+            model_versions = list(versions.keys())
+            if model_versions:
+                return model_versions[0]
         return ""
    
     def update_model_info(self, model_type: str, training_type: str) -> Dict: