onetrainer-dpg.py

#!/usr/bin/env python3
"""
OneTrainer DPG UI Launcher (Full Version)

This script provides a comprehensive launcher for the OneTrainer DPG UI,
with proper error handling and fallbacks.
"""

import os
import sys
import time
import importlib
import traceback
from pathlib import Path

# Banner
print("""
╔═══════════════════════════════════════════════════╗
║               ONETRAINER DPG EDITION              ║
║           (GPU Accelerated Dear PyGui UI)         ║
╚═══════════════════════════════════════════════════╝
""")

# Add current directory to path
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

# Check for DearPyGui
try:
    import dearpygui.dearpygui as dpg
    print("✓ DearPyGui is installed")
except ImportError:
    print("✗ DearPyGui is not installed")
    print("Please install it with: pip install dearpygui")
    sys.exit(1)

# Check for OneTrainer modules
try:
    from modules.util.enum.ModelType import ModelType
    print("✓ OneTrainer modules are available")
except ImportError:
    print("✗ OneTrainer modules not found")
    print("Make sure you're running from the OneTrainer directory")
    sys.exit(1)

# Setup the components system
try:
    print("Setting up UI components...")
    # Ensure the necessary directories exist
    os.makedirs(os.path.join(current_dir, "dpg_ui"), exist_ok=True)
    os.makedirs(os.path.join(current_dir, "dpg_ui", "components"), exist_ok=True)
    os.makedirs(os.path.join(current_dir, "dpg_ui", "tabs"), exist_ok=True)
    os.makedirs(os.path.join(current_dir, "dpg_ui", "themes"), exist_ok=True)
    
    # Check if our file_selector component is available
    file_selector_path = os.path.join(current_dir, "dpg_ui", "components", "file_selector.py")
    if not os.path.exists(file_selector_path):
        print("Creating file_selector.py...")
        with open(file_selector_path, 'w') as f:
            f.write('''#!/usr/bin/env python3
"""
File selector component for DPG UI

This module provides a robust file selection component that handles
various edge cases and error conditions properly.
"""

import os
import dearpygui.dearpygui as dpg
from typing import Any, Dict, List, Callable, Optional, Union

def file_selector(
    parent: Union[int, str],
    default_path: str = "",
    directory: str = None,
    extension: str = None,
    extensions: List[str] = None,
    callback: Callable = None,
    width: int = 300,
    tag: str = None,
    hint: str = None,
    is_directory: bool = False
) -> Dict[str, Any]:
    """
    Robust file selector component that handles callbacks properly
    
    Args:
        parent: Parent container ID
        default_path: Default file path to display
        directory: Directory to use as base path
        extension: Single file extension to filter by
        extensions: List of file extensions to filter by
        callback: Function to call when a file is selected
        width: Width of the input field
        tag: Custom tag for the component
        hint: Hint text to display
        is_directory: True to select directories, False to select files
        
    Returns:
        Dictionary with component IDs
    """
    # Generate unique ID if not provided
    if tag is None:
        tag = f"file_selector_{dpg.generate_uuid()}"
        
    # Prepare default values
    default_value = default_path
    
    # If directory is specified and we have a default path, use just the filename
    if directory and default_path:
        if os.path.isabs(default_path):
            default_value = os.path.basename(default_path)
    
    # Create group for the component
    with dpg.group(parent=parent, horizontal=True, tag=tag):
        # Text input field for the path
        input_id = f"{tag}_input"
        dpg.add_input_text(
            hint=hint if hint else "Select a file..." if not is_directory else "Select a directory...",
            default_value=default_value,
            width=width-50,
            tag=input_id
        )
        
        # Browse button
        button_id = f"{tag}_button"
        dpg.add_button(
            label="...",
            width=35,
            height=25,
            tag=button_id
        )
        
        # File dialog for selecting files/directories
        dialog_id = f"{tag}_dialog"
        
        # Create dialog with appropriate settings
        if is_directory:
            # Directory selection dialog
            dpg.add_file_dialog(
                directory_selector=True,
                show=False,
                tag=dialog_id,
                width=700,
                height=400,
                modal=True,
                default_path=directory if directory else "",
                callback=lambda s, a, u: handle_selection(s, a, u, input_id, directory, callback)
            )
        else:
            # File selection dialog
            filter_ext = []
            if extension:
                filter_ext.append(f".{extension}")
            elif extensions:
                filter_ext.extend([f".{ext}" for ext in extensions])
                
            dpg.add_file_dialog(
                directory_selector=False,
                show=False,
                tag=dialog_id,
                width=700,
                height=400,
                modal=True,
                default_path=directory if directory else "",
                callback=lambda s, a, u: handle_selection(s, a, u, input_id, directory, callback)
            )
            
            # Add extension filters if specified
            if filter_ext:
                ext_id = dpg.add_file_extension_info(parent=dialog_id)
                for ext in filter_ext:
                    dpg.add_file_extension(extension=ext, parent=ext_id, custom_text=f"*{ext} Files")
        
        # Set up button callback to show dialog
        dpg.set_item_callback(button_id, lambda: dpg.show_item(dialog_id))
    
    # Return components IDs for future use
    return {
        "group": tag,
        "input": input_id,
        "button": button_id,
        "dialog": dialog_id,
        "get_value": lambda: get_path_value(input_id, directory),
        "set_value": lambda v: set_path_value(input_id, v, directory)
    }

def handle_selection(sender, app_data, user_data, input_id, directory, callback):
    """
    Handle file selection event
    
    This function properly extracts the file path and calls the callback
    with the appropriate arguments.
    
    Args:
        sender: Sender ID
        app_data: Event data (contains file_path_name)
        user_data: User data
        input_id: ID of the input field to update
        directory: Base directory if specified
        callback: Function to call with the selected path
    """
    try:
        # Extract file path from dialog data
        if "file_path_name" in app_data:
            full_path = app_data["file_path_name"]
            
            # If directory is specified, use just the filename
            if directory:
                filename = os.path.basename(full_path)
                path_value = filename
                full_value = os.path.join(directory, filename)
            else:
                path_value = full_path
                full_value = full_path
                
            # Update the input field
            dpg.set_value(input_id, path_value)
            
            # Call callback if provided
            if callback:
                try:
                    # Try with the full path
                    callback(full_value)
                except TypeError:
                    try:
                        # Try with input ID and path
                        callback(input_id, full_value)
                    except Exception as e:
                        print(f"Error in file selection callback: {e}")
                        # Last resort - try with sender and app_data
                        try:
                            callback(sender, app_data)
                        except Exception:
                            print(f"All callback signature attempts failed for {callback.__name__ if hasattr(callback, '__name__') else 'callback'}")
    except Exception as e:
        print(f"Error handling file selection: {e}")

def get_path_value(input_id, directory=None):
    """Get the full path value from the input field"""
    try:
        # Get path from input field
        path = dpg.get_value(input_id)
        
        # Combine with directory if provided
        if directory and path:
            return os.path.join(directory, path)
        
        return path
    except Exception as e:
        print(f"Error getting path value: {e}")
        return ""

def set_path_value(input_id, value, directory=None):
    """Set the path value in the input field"""
    try:
        # If directory is specified, use just the filename
        if directory and value:
            if os.path.isabs(value):
                filename = os.path.basename(value)
                dpg.set_value(input_id, filename)
            else:
                dpg.set_value(input_id, value)
        else:
            dpg.set_value(input_id, value)
    except Exception as e:
        print(f"Error setting path value: {e}")
''')
    
    # Ensure components/__init__.py exists and has correct content
    init_path = os.path.join(current_dir, "dpg_ui", "components", "__init__.py")
    if not os.path.exists(init_path):
        print("Creating components/__init__.py...")
        with open(init_path, 'w') as f:
            f.write('''#!/usr/bin/env python3
"""
Components package for Dear PyGui UI
"""

# Import all components for easy access
from .basic import Components
try:
    from .file_selector import file_selector
    # Add file_selector as a static method to Components
    Components.file_selector = staticmethod(file_selector)
except ImportError:
    print("Warning: file_selector not available")

try:
    from .extensions import extend_components
    # Extend the Components class with additional methods
    extend_components(Components)
except ImportError:
    print("Warning: extensions not available")
''')
    
    # Create basic components
    basic_path = os.path.join(current_dir, "dpg_ui", "components", "basic.py")
    if not os.path.exists(basic_path):
        print("Creating basic components...")
        with open(basic_path, 'w') as f:
            f.write('''#!/usr/bin/env python3
"""
Basic UI components for DPG UI
"""

import dearpygui.dearpygui as dpg
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

class Components:
    """
    Collection of static UI components for building the Dear PyGui UI
    
    This class provides a set of methods for creating common UI components
    with consistent styling and behavior.
    """
    
    @staticmethod
    def label(parent: int, text: str, **kwargs) -> int:
        """Add a label with standard formatting"""
        return dpg.add_text(text, parent=parent, **kwargs)
    
    @staticmethod
    def button(parent: int, label: str, callback: Callable = None, **kwargs) -> int:
        """Add a button with standard styling"""
        button_id = dpg.add_button(label=label, parent=parent, **kwargs)
        if callback:
            dpg.set_item_callback(button_id, callback)
        return button_id
    
    @staticmethod
    def input_text(parent: int, default_value: str = "", **kwargs) -> int:
        """Add a text input field with standard styling"""
        return dpg.add_input_text(default_value=default_value, parent=parent, **kwargs)
    
    @staticmethod
    def input_int(parent: int, default_value: int = 0, **kwargs) -> int:
        """Add an integer input field with standard styling"""
        return dpg.add_input_int(default_value=default_value, parent=parent, **kwargs)
    
    @staticmethod
    def input_float(parent: int, default_value: float = 0.0, **kwargs) -> int:
        """Add a float input field with standard styling"""
        return dpg.add_input_float(default_value=default_value, parent=parent, **kwargs)
    
    @staticmethod
    def combo(parent: int, items: List[str], default_value: str = None, **kwargs) -> int:
        """Add a combo box with standard styling"""
        combo_id = dpg.add_combo(items=items, parent=parent, **kwargs)
        if default_value is not None:
            dpg.set_value(combo_id, default_value)
        return combo_id
    
    @staticmethod
    def checkbox(parent: int, label: str, default_value: bool = False, **kwargs) -> int:
        """Add a checkbox with standard styling"""
        return dpg.add_checkbox(label=label, default_value=default_value, parent=parent, **kwargs)
    
    @staticmethod
    def radio_button(parent: int, items: List[str], default_value: int = 0, **kwargs) -> int:
        """Add radio buttons with standard styling"""
        return dpg.add_radio_button(items=items, default_value=default_value, parent=parent, **kwargs)
    
    @staticmethod
    def slider_float(parent: int, default_value: float = 0.0, min_value: float = 0.0, 
                     max_value: float = 1.0, **kwargs) -> int:
        """Add a float slider with standard styling"""
        return dpg.add_slider_float(default_value=default_value, min_value=min_value, 
                                   max_value=max_value, parent=parent, **kwargs)
    
    @staticmethod
    def slider_int(parent: int, default_value: int = 0, min_value: int = 0, 
                   max_value: int = 100, **kwargs) -> int:
        """Add an integer slider with standard styling"""
        return dpg.add_slider_int(default_value=default_value, min_value=min_value, 
                                 max_value=max_value, parent=parent, **kwargs)
    
    @staticmethod
    def separator(parent: int, **kwargs) -> int:
        """Add a separator with standard styling"""
        return dpg.add_separator(parent=parent, **kwargs)
    
    @staticmethod
    def selectable(parent: int, label: str, default_value: bool = False, **kwargs) -> int:
        """Add a selectable item with standard styling"""
        return dpg.add_selectable(label=label, default_value=default_value, parent=parent, **kwargs)
    
    @staticmethod
    def progress_bar(parent: int, default_value: float = 0.0, **kwargs) -> int:
        """Add a progress bar with standard styling"""
        return dpg.add_progress_bar(default_value=default_value, parent=parent, **kwargs)
    
    @staticmethod
    def image(parent: int, texture_id: int, **kwargs) -> int:
        """Add an image with standard styling"""
        return dpg.add_image(texture_id, parent=parent, **kwargs)
    
    @staticmethod
    def tooltip(parent: int, text: str, **kwargs) -> int:
        """Add a tooltip with standard styling"""
        with dpg.tooltip(parent=parent, **kwargs) as tooltip_id:
            dpg.add_text(text)
        return tooltip_id
    
    @staticmethod
    def tab_bar(parent: int, **kwargs) -> int:
        """Add a tab bar with standard styling"""
        return dpg.add_tab_bar(parent=parent, **kwargs)
    
    @staticmethod
    def tab(parent: int, label: str, **kwargs) -> int:
        """Add a tab with standard styling"""
        return dpg.add_tab(label=label, parent=parent, **kwargs)
    
    @staticmethod
    def group(parent: int, **kwargs) -> int:
        """Add a group with standard styling"""
        return dpg.add_group(parent=parent, **kwargs)
    
    @staticmethod
    def collapsing_header(parent: int, label: str, default_open: bool = False, **kwargs) -> int:
        """Add a collapsing header with standard styling"""
        return dpg.add_collapsing_header(label=label, default_open=default_open, parent=parent, **kwargs)
    
    @staticmethod
    def tree_node(parent: int, label: str, default_open: bool = False, **kwargs) -> int:
        """Add a tree node with standard styling"""
        return dpg.add_tree_node(label=label, default_open=default_open, parent=parent, **kwargs)
    
    @staticmethod
    def popup(parent: int, title: str, **kwargs) -> int:
        """Add a popup with standard styling"""
        return dpg.add_popup(parent=parent, title=title, **kwargs)
    
    @staticmethod
    def table(parent: int, headers: List[str], **kwargs) -> int:
        """Add a table with standard styling"""
        table_id = dpg.add_table(parent=parent, **kwargs)
        for header in headers:
            dpg.add_table_column(label=header, parent=table_id)
        return table_id
    
    @staticmethod
    def file_entry(parent: int, default_path: str = "", width: int = 300, 
                  filter_extensions: List[str] = None, path_modifier: Callable = None, **kwargs) -> Dict:
        """Add a file entry with standard styling"""
        # Check if file_selector is available (it might be added later as a static method)
        if hasattr(Components, "file_selector"):
            extensions = filter_extensions if filter_extensions else []
            return Components.file_selector(
                parent=parent, 
                default_path=default_path,
                width=width,
                extensions=extensions,
                **kwargs
            )
        else:
            # Fallback to simple text input
            from typing import Dict
            input_id = dpg.add_input_text(default_value=default_path, width=width, parent=parent)
            return {
                "input": input_id,
                "get_value": lambda: dpg.get_value(input_id),
                "set_value": lambda v: dpg.set_value(input_id, v)
            }
    
    @staticmethod
    def directory_entry(parent: int, default_path: str = "", width: int = 300, **kwargs) -> Dict:
        """Add a directory entry with standard styling"""
        # Check if file_selector is available (it might be added later as a static method)
        if hasattr(Components, "file_selector"):
            return Components.file_selector(
                parent=parent, 
                default_path=default_path,
                width=width,
                is_directory=True,
                **kwargs
            )
        else:
            # Fallback to simple text input
            input_id = dpg.add_input_text(default_value=default_path, width=width, parent=parent)
            return {
                "input": input_id,
                "get_value": lambda: dpg.get_value(input_id),
                "set_value": lambda v: dpg.set_value(input_id, v)
            }
    
    @staticmethod
    def entry(parent: int, default_value: str = "", **kwargs) -> int:
        """Add a text entry with standard styling (alias for input_text)"""
        return Components.input_text(parent, default_value, **kwargs)
    
    @staticmethod
    def switch(parent: int, default_value: bool = False, **kwargs) -> int:
        """Add a switch control (using checkbox internally)"""
        return dpg.add_checkbox(default_value=default_value, parent=parent, **kwargs)
    
    @staticmethod
    def options(parent: int, values: List[str], default_value: str = None, 
                on_change: Callable = None, **kwargs) -> int:
        """Add an options dropdown with standard styling"""
        combo_id = dpg.add_combo(items=values, default_value=default_value or values[0] if values else None, 
                                 parent=parent, **kwargs)
        if on_change:
            dpg.set_item_callback(combo_id, lambda s, a: on_change(a))
        return combo_id
    
    @staticmethod
    def options_kv(parent: int, values: List[Tuple[str, Any]], default_value: Any = None,
                  on_change: Callable = None, **kwargs) -> int:
        """
        Add a key-value options dropdown
        
        This displays the first item in the tuple as the display value
        and returns the second item as the actual value.
        """
        # Extract display values for the combo
        display_values = [x[0] for x in values]
        
        # Create value map for lookup
        value_map = {str(x[0]): x[1] for x in values}
        
        # Find default display value
        default_display = None
        if default_value is not None:
            for display, value in values:
                if value == default_value:
                    default_display = display
                    break
            
            if default_display is None:
                default_display = display_values[0] if display_values else None
        else:
            default_display = display_values[0] if display_values else None
        
        # Create the combo
        combo_id = dpg.add_combo(
            items=display_values, 
            default_value=default_display,
            parent=parent, 
            **kwargs
        )
        
        # Add callback to convert display value to actual value
        if on_change:
            # Create a wrapper that converts the display value to the actual value
            def wrapper(sender, app_data):
                # Look up the actual value from the map
                actual_value = value_map.get(app_data, app_data)
                # Call the original callback with the actual value
                on_change(actual_value)
                
            dpg.set_item_callback(combo_id, wrapper)
        
        return combo_id
''')
    
    # Import Components to check if it works
    try:
        dpg_components_dir = os.path.join(current_dir, "dpg_ui", "components")
        if dpg_components_dir not in sys.path:
            sys.path.insert(0, dpg_components_dir)
            
        from dpg_ui.components import Components
        print("✓ Components system initialized")
        
        # Check if file_selector is available
        if hasattr(Components, 'file_selector'):
            print("✓ file_selector is available")
        else:
            print("✗ file_selector not available - UI will have limited functionality")
    except ImportError as e:
        print(f"✗ Error importing Components: {e}")
except Exception as e:
    print(f"Error setting up components: {e}")
    traceback.print_exc()

# Ensure the window module exists
window_path = os.path.join(current_dir, "dpg_ui", "window.py")
if not os.path.exists(window_path):
    print("Creating window.py...")
    with open(window_path, 'w') as f:
        f.write('''#!/usr/bin/env python3
"""
Window implementation for the Dear PyGui version of OneTrainer
"""

import os
import sys
import dearpygui.dearpygui as dpg
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

# Add the parent directory to sys.path to allow importing from the modules package
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

class OneTrainerWindow:
    """
    Main window implementation for the Dear PyGui version of OneTrainer
    
    This class manages the window layout, tab system, and state management
    for the OneTrainer UI.
    """
    
    def __init__(self, title: str = "OneTrainer", width: int = 1280, height: int = 800):
        """Initialize the window"""
        self.title = title
        self.width = width
        self.height = height
        self.tabs = {}
        self.tab_instances = {}
        self.current_tab = None
        self.top_bar = None
        self.state = {}
        
        # Initialize DPG
        dpg.create_context()
        dpg.create_viewport(title=title, width=width, height=height)
        dpg.setup_dearpygui()
        
        # Create main window
        with dpg.window(tag="main_window", label=title):
            # Main layout with top bar and tab panel
            with dpg.group(horizontal=False):
                # Add top area for the top bar
                self.top_bar_container = dpg.add_group(horizontal=True, tag="top_bar_container")
                
                # Add separator
                dpg.add_separator()
                
                # Add tab bar
                self.tab_bar = dpg.add_tab_bar(tag="main_tab_bar", callback=self.on_tab_change)
                
                # Add status bar at bottom
                dpg.add_separator()
                with dpg.group(horizontal=True, tag="status_bar"):
                    self.status_text = dpg.add_text("Ready", tag="status_text")
        
        # Set primary window
        dpg.set_primary_window("main_window", True)
    
    def show(self):
        """Show the window and start the main loop"""
        dpg.show_viewport()
        dpg.start_dearpygui()
        dpg.destroy_context()
    
    def bind_to_config(self, config: Dict[str, Any]):
        """Bind to a configuration dictionary"""
        self.state["config"] = config
    
    def get_model_type(self):
        """Get the current model type from the state"""
        if "config" in self.state and "model_type" in self.state["config"]:
            from modules.util.enum.ModelType import ModelType
            try:
                return ModelType(self.state["config"]["model_type"])
            except (ValueError, TypeError):
                return None
        return None
    
    def get_lora_state(self):
        """Get the current LoRA state from the configuration"""
        if "config" in self.state:
            # Extract LoRA-related settings from config
            config = self.state["config"]
            lora_state = {}
            
            # Map config keys to lora state
            lora_keys = [
                "lora_rank", "lora_alpha", "dropout_probability", 
                "peft_type", "weight_dtype", "bundle_embeddings",
                "layer_preset", "layers", "decompose", "norm_epsilon", 
                "output_axis", "lycoris_factor", "lycoris_full_matrix",
                "lycoris_bypass_mode", "lora_model_path"
            ]
            
            for key in lora_keys:
                if key in config:
                    lora_state[key] = config[key]
            
            return lora_state
        return {}
    
    def add_tab(self, name: str, label: str, content_factory: Callable[[int, Dict], Any]):
        """
        Add a tab to the tab bar
        
        Args:
            name: Unique name for the tab
            label: Display label for the tab
            content_factory: Function that creates the tab content and returns it
        """
        tab_id = dpg.add_tab(label=label, parent=self.tab_bar, tag=f"tab_{name}")
        self.tabs[name] = tab_id
        
        # Create the content
        content = content_factory(tab_id, self)
        self.tab_instances[name] = content
    
    def select_tab(self, name: str):
        """Select a tab by name"""
        if name in self.tabs:
            dpg.set_value(self.tab_bar, self.tabs[name])
            self.current_tab = name
    
    def on_tab_change(self, sender, value):
        """Handle tab change event"""
        # Find the tab name from the value
        for name, tab_id in self.tabs.items():
            if tab_id == value:
                self.current_tab = name
                break
    
    def set_top_bar(self, content_factory: Callable[[int, Dict], Any]):
        """Set the top bar content"""
        dpg.delete_item(self.top_bar_container, children_only=True)
        self.top_bar = content_factory(self.top_bar_container, self)
    
    def set_status(self, text: str):
        """Set the status bar text"""
        dpg.set_value(self.status_text, text)
''')

# Add the minimal UI for fallback
minimal_ui_path = os.path.join(current_dir, "minimal_dpg_ui.py")
if not os.path.exists(minimal_ui_path):
    print("Creating minimal_dpg_ui.py...")
    with open(minimal_ui_path, 'w') as f:
        f.write('''#!/usr/bin/env python3
"""
Minimal Dear PyGui UI for OneTrainer

This module provides a minimal UI implementation for OneTrainer
that works as a fallback when the full UI cannot be loaded.
"""

import os
import sys
import dearpygui.dearpygui as dpg
from typing import Any, Dict, List, Optional, Union

# Add the parent directory to sys.path to allow importing from the modules package
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

from modules.util.enum.ModelType import ModelType, PeftType
from modules.util.config.TrainConfig import TrainConfig
from modules.util.enum.TrainingMethod import TrainingMethod

class MinimalOneTrainerUI:
    """
    Minimal UI implementation for OneTrainer
    
    This class provides a simple UI that works as a fallback when
    the full UI cannot be loaded.
    """
    
    def __init__(self):
        """Initialize the UI"""
        self.tabs = {}
        self.components = {}
        
        # Set up state
        self.train_config = TrainConfig.default_values()
        
    def run(self):
        """Run the UI"""
        # Initialize DPG
        dpg.create_context()
        dpg.create_viewport(title="OneTrainer (Minimal)", width=1000, height=800)
        dpg.setup_dearpygui()
        
        # Create main window
        with dpg.window(tag="main_window", label="OneTrainer (Minimal)"):
            # Tab bar
            tab_bar = dpg.add_tab_bar(tag="main_tab_bar")
            
            # Training tab
            with dpg.tab(label="Training", parent=tab_bar, tag="training_tab"):
                self.create_training_tab("training_tab")
            
            # Model tab
            with dpg.tab(label="Model", parent=tab_bar, tag="model_tab"):
                self.create_model_tab("model_tab")
            
            # LoRA tab
            with dpg.tab(label="LoRA", parent=tab_bar, tag="lora_tab"):
                self.create_lora_tab("lora_tab")
            
            # Concept tab
            with dpg.tab(label="Concept", parent=tab_bar, tag="concept_tab"):
                self.create_concept_tab("concept_tab")
            
            # Status bar
            dpg.add_separator()
            with dpg.group(horizontal=True):
                self.status_text = dpg.add_text("Ready", tag="status_text")
                
                # Add run button
                dpg.add_button(label="Run Training", callback=self.on_run_training)
        
        # Set primary window
        dpg.set_primary_window("main_window", True)
        
        # Show the UI
        dpg.show_viewport()
        dpg.start_dearpygui()
        dpg.destroy_context()
    
    def create_training_tab(self, parent):
        """Create the training tab"""
        with dpg.group(parent=parent, horizontal=False):
            dpg.add_text("Training Settings", color=(255, 255, 0))
            dpg.add_separator()
            
            # Epochs and batch size
            with dpg.group(horizontal=True):
                dpg.add_text("Epochs:", width=150)
                epochs = dpg.add_input_int(default_value=self.train_config.epochs, width=200)
                self.add_component("epochs", epochs)
            
            with dpg.group(horizontal=True):
                dpg.add_text("Batch Size:", width=150)
                batch_size = dpg.add_input_int(default_value=self.train_config.batch_size, width=200)
                self.add_component("batch_size", batch_size)
            
            with dpg.group(horizontal=True):
                dpg.add_text("Learning Rate:", width=150)
                learning_rate = dpg.add_input_float(default_value=self.train_config.learning_rate, 
                                                 width=200, format="%.6f")
                self.add_component("learning_rate", learning_rate)
            
            # Training method
            with dpg.group(horizontal=True):
                dpg.add_text("Training Method:", width=150)
                methods = [method.name for method in TrainingMethod]
                training_method = dpg.add_combo(items=methods, 
                                             default_value=self.train_config.training_method.name, 
                                             width=200)
                self.add_component("training_method", training_method)
    
    def create_model_tab(self, parent):
        """Create the model tab"""
        with dpg.group(parent=parent, horizontal=False):
            dpg.add_text("Model Settings", color=(255, 255, 0))
            dpg.add_separator()
            
            # Model type
            with dpg.group(horizontal=True):
                dpg.add_text("Model Type:", width=150)
                model_types = [model_type.name for model_type in ModelType]
                model_type = dpg.add_combo(items=model_types, 
                                         default_value=self.train_config.model_type.name, 
                                         width=300,
                                         callback=self.on_model_type_change)
                self.add_component("model_type", model_type)
            
            # Model path
            with dpg.group(horizontal=True):
                dpg.add_text("Model Path:", width=150)
                model_path = dpg.add_input_text(default_value=self.train_config.model_path, width=300)
                self.add_component("model_path", model_path)
                dpg.add_button(label="...", width=30, callback=self.on_select_model)
            
            # VAE path
            with dpg.group(horizontal=True):
                dpg.add_text("VAE Path:", width=150)
                vae_path = dpg.add_input_text(default_value=self.train_config.vae_path, width=300)
                self.add_component("vae_path", vae_path)
                dpg.add_button(label="...", width=30, callback=self.on_select_vae)
    
    def create_lora_tab(self, parent):
        """Create the LoRA tab"""
        with dpg.group(parent=parent, horizontal=False):
            dpg.add_text("LoRA Settings", color=(255, 255, 0))
            dpg.add_separator()
            
            # PEFT type selector
            with dpg.group(horizontal=True):
                dpg.add_text("PEFT Type:", width=150,
                             tooltip="The type of parameter-efficient fine-tuning method")
                peft_types = [
                    (PeftType.LORA.name, "LoRA - Standard adaptation"),
                    (PeftType.LOHA.name, "LoHa - Hadamard product"),
                    (PeftType.LOKR.name, "LoKr - Kronecker product"),
                    (PeftType.DIA.name, "DiA - Diagonal adaptation"),
                    (PeftType.IA3.name, "iA3 - Infused adapter by anisotropic adaptation"),
                    (PeftType.DYLORA.name, "DyLoRA - Dynamic LoRA")
                ]
                peft_values = [p[0] for p in peft_types]
                peft_labels = [f"{p[0]} ({p[1]})" for p in peft_types]
                
                peft_type = dpg.add_combo(
                    items=peft_labels,
                    default_value=peft_labels[0],
                    width=300,
                    callback=self.on_peft_type_change
                )
                self.add_component("peft_type", peft_type)
            
            # LoRA rank
            with dpg.group(horizontal=True):
                dpg.add_text("LoRA Rank:", width=150)
                lora_rank = dpg.add_input_int(default_value=self.train_config.lora_rank, width=200)
                self.add_component("lora_rank", lora_rank)
            
            # LoRA alpha
            with dpg.group(horizontal=True):
                dpg.add_text("LoRA Alpha:", width=150)
                lora_alpha = dpg.add_input_float(default_value=self.train_config.lora_alpha, width=200)
                self.add_component("lora_alpha", lora_alpha)
            
            # Warning for HiDream + LoKr
            with dpg.group(tag="lokr_warning", show=False):
                with dpg.drawlist(width=550, height=80):
                    dpg.draw_rectangle(
                        (0, 0), (550, 80),
                        fill=(255, 230, 230, 200),
                        color=(255, 0, 0, 255),
                        thickness=2
                    )
                    dpg.add_text(
                        "⚠️ WARNING FOR HIDREAM MODELS ⚠️\\n"
                        "LoKr is known to cause sampling issues with HiDream models, which may lead to:\\n"
                        "- Infinite sampling loops or freezing\\n"
                        "- Corrupted outputs or crashes\\n"
                        "For HiDream models, please use standard LoRA, LoHa, or other PEFT types.",
                        pos=(10, 10),
                        color=(180, 0, 0, 255),
                        wrap=530
                    )
    
    def create_concept_tab(self, parent):
        """Create the concept tab"""
        with dpg.group(parent=parent, horizontal=False):
            dpg.add_text("Concept Settings", color=(255, 255, 0))
            dpg.add_separator()
            
            # Dataset path
            with dpg.group(horizontal=True):
                dpg.add_text("Dataset Path:", width=150)
                dataset_path = dpg.add_input_text(default_value=self.train_config.dataset_path, width=300)
                self.add_component("dataset_path", dataset_path)
                dpg.add_button(label="...", width=30, callback=self.on_select_dataset)
            
            # Resolution
            with dpg.group(horizontal=True):
                dpg.add_text("Resolution:", width=150)
                resolution = dpg.add_input_int(default_value=self.train_config.resolution, width=200)
                self.add_component("resolution", resolution)
    
    def add_component(self, name, component_id):
        """Add a component to the registry"""
        self.components[name] = component_id
    
    def get_component(self, name):
        """Get a component by name"""
        return self.components.get(name)
    
    def on_model_type_change(self, sender, app_data):
        """Handle model type change"""
        # Check if we need to show the LoKr warning
        try:
            model_type = ModelType(app_data)
            peft_type_combo = self.get_component("peft_type")
            if peft_type_combo:
                peft_type_value = dpg.get_value(peft_type_combo)
                
                # Show warning if HiDream + LoKr
                if model_type == ModelType.HI_DREAM_FULL and "LOKR" in peft_type_value:
                    dpg.configure_item("lokr_warning", show=True)
                else:
                    dpg.configure_item("lokr_warning", show=False)
        except Exception as e:
            print(f"Error in model type change handler: {e}")
    
    def on_peft_type_change(self, sender, app_data):
        """Handle PEFT type change"""
        # Check if we need to show the LoKr warning
        try:
            model_type_combo = self.get_component("model_type")
            if model_type_combo:
                model_type_value = dpg.get_value(model_type_combo)
                
                # Show warning if HiDream + LoKr
                if model_type_value == "HI_DREAM_FULL" and "LOKR" in app_data:
                    dpg.configure_item("lokr_warning", show=True)
                else:
                    dpg.configure_item("lokr_warning", show=False)
        except Exception as e:
            print(f"Error in PEFT type change handler: {e}")
    
    def on_select_model(self):
        """Handle model selection"""
        # Simple placeholder
        dpg.set_value(self.get_component("status_text"), "Model selection dialog would appear here")
    
    def on_select_vae(self):
        """Handle VAE selection"""
        # Simple placeholder
        dpg.set_value(self.get_component("status_text"), "VAE selection dialog would appear here")
    
    def on_select_dataset(self):
        """Handle dataset selection"""
        # Simple placeholder
        dpg.set_value(self.get_component("status_text"), "Dataset selection dialog would appear here")
    
    def on_run_training(self):
        """Handle run training button"""
        # Get all values
        epochs = dpg.get_value(self.get_component("epochs"))
        batch_size = dpg.get_value(self.get_component("batch_size"))
        learning_rate = dpg.get_value(self.get_component("learning_rate"))
        model_type = dpg.get_value(self.get_component("model_type"))
        
        # Display a message
        dpg.set_value(self.get_component("status_text"), 
                   f"Would start training with: {model_type}, {epochs} epochs, {batch_size} batch size")
        
        # Show a temporary message in a modal dialog
        with dpg.window(label="Training", modal=True, width=400, height=200, pos=(100, 100)):
            dpg.add_text("This is a minimal UI. Training not implemented.")
            dpg.add_text(f"Model Type: {model_type}")
            dpg.add_text(f"Epochs: {epochs}")
            dpg.add_text(f"Batch Size: {batch_size}")
            dpg.add_text(f"Learning Rate: {learning_rate}")
            dpg.add_button(label="Close", callback=lambda: dpg.delete_item(dpg.last_container()))
''')

# Create dark theme module
theme_dir = os.path.join(current_dir, "dpg_ui", "themes")
os.makedirs(theme_dir, exist_ok=True)
theme_init_path = os.path.join(theme_dir, "__init__.py")
if not os.path.exists(theme_init_path):
    with open(theme_init_path, 'w') as f:
        f.write('# Theme package\n')

dark_theme_path = os.path.join(theme_dir, "dark_theme.py")
if not os.path.exists(dark_theme_path):
    print("Creating dark theme...")
    with open(dark_theme_path, 'w') as f:
        f.write('''#!/usr/bin/env python3
"""
Dark theme implementation for Dear PyGui UI
"""

import dearpygui.dearpygui as dpg

def apply_dark_theme():
    """Apply a dark theme to the UI"""
    with dpg.theme() as global_theme:
        with dpg.theme_component(dpg.mvAll):
            # Text
            dpg.add_theme_color(dpg.mvThemeCol_Text, (255, 255, 255, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TextDisabled, (128, 128, 128, 255))
            
            # Window
            dpg.add_theme_color(dpg.mvThemeCol_WindowBg, (25, 25, 25, 255))
            dpg.add_theme_color(dpg.mvThemeCol_ChildBg, (25, 25, 25, 255))
            dpg.add_theme_color(dpg.mvThemeCol_PopupBg, (30, 30, 30, 255))
            dpg.add_theme_color(dpg.mvThemeCol_Border, (110, 110, 128, 255))
            dpg.add_theme_color(dpg.mvThemeCol_BorderShadow, (0, 0, 0, 0))
            
            # Frame
            dpg.add_theme_color(dpg.mvThemeCol_FrameBg, (41, 74, 122, 138))
            dpg.add_theme_color(dpg.mvThemeCol_FrameBgHovered, (66, 150, 250, 102))
            dpg.add_theme_color(dpg.mvThemeCol_FrameBgActive, (66, 150, 250, 171))
            
            # Tabs
            dpg.add_theme_color(dpg.mvThemeCol_Tab, (46, 89, 148, 220))
            dpg.add_theme_color(dpg.mvThemeCol_TabHovered, (66, 150, 250, 204))
            dpg.add_theme_color(dpg.mvThemeCol_TabActive, (51, 105, 173, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TabUnfocused, (33, 59, 98, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TabUnfocusedActive, (35, 93, 179, 255))
            
            # Title
            dpg.add_theme_color(dpg.mvThemeCol_TitleBg, (35, 35, 35, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TitleBgActive, (41, 74, 122, 255))
            dpg.add_theme_color(dpg.mvThemeCol_TitleBgCollapsed, (35, 35, 35, 128))
            
            # Buttons
            dpg.add_theme_color(dpg.mvThemeCol_Button, (66, 102, 148, 255))
            dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, (66, 150, 250, 255))
            dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, (15, 135, 250, 255))
            
            # Headers
            dpg.add_theme_color(dpg.mvThemeCol_Header, (66, 102, 148, 255))
            dpg.add_theme_color(dpg.mvThemeCol_HeaderHovered, (66, 150, 250, 204))
            dpg.add_theme_color(dpg.mvThemeCol_HeaderActive, (66, 150, 250, 255))
            
            # Sliders & scrollbars
            dpg.add_theme_color(dpg.mvThemeCol_SliderGrab, (66, 150, 250, 204))
            dpg.add_theme_color(dpg.mvThemeCol_SliderGrabActive, (15, 135, 250, 255))
            dpg.add_theme_color(dpg.mvThemeCol_ScrollbarGrab, (66, 102, 148, 255))
            dpg.add_theme_color(dpg.mvThemeCol_ScrollbarGrabHovered, (66, 150, 250, 204))
            dpg.add_theme_color(dpg.mvThemeCol_ScrollbarGrabActive, (15, 135, 250, 255))
            
            # Styles
            dpg.add_theme_style(dpg.mvStyleVar_WindowRounding, 5)
            dpg.add_theme_style(dpg.mvStyleVar_FrameRounding, 3)
            dpg.add_theme_style(dpg.mvStyleVar_PopupRounding, 3)
            dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 6, 6)
            dpg.add_theme_style(dpg.mvStyleVar_ItemSpacing, 6, 6)
    
    # Apply the theme
    dpg.bind_theme(global_theme)
''')

# Run the UI
def run_ui():
    """Run the UI with comprehensive error handling"""
    try:
        # Clean up any existing DPG context
        try:
            # Different ways to check if DPG is running (depends on version)
            dpg_running = False
            try:
                if hasattr(dpg, '_dpg_internal') and hasattr(dpg._dpg_internal, 'is_dearpygui_running'):
                    dpg_running = dpg._dpg_internal.is_dearpygui_running()
                elif hasattr(dpg, 'is_dearpygui_running'):
                    dpg_running = dpg.is_dearpygui_running()
            except:
                pass
                
            if dpg_running:
                dpg.stop_dearpygui()
        except:
            pass
            
        try:
            dpg.destroy_context()
        except:
            pass
        
        # Try the enhanced app
        try:
            print("\nLaunching enhanced DPG UI...")
            from dpg_ui.app_enhanced import OneTrainerApp
            app = OneTrainerApp()
            app.run()
            return True
        except Exception as e:
            print(f"Enhanced UI failed: {e}")
            traceback.print_exc()
            
            # Clean up DPG context
            try:
                dpg.destroy_context()
            except:
                pass
            
            # Try standard app
            try:
                print("\nTrying standard DPG UI...")
                from dpg_ui.app import OneTrainerApp as StandardApp
                app = StandardApp()
                app.run()
                return True
            except Exception as e:
                print(f"Standard UI failed: {e}")
                traceback.print_exc()
                
                # Clean up DPG context
                try:
                    # Stop DPG if running
                    try:
                        if hasattr(dpg, 'is_dearpygui_running') and dpg.is_dearpygui_running():
                            dpg.stop_dearpygui()
                    except:
                        pass
                    
                    dpg.destroy_context()
                except:
                    pass
                
                # Try simplified UI
                try:
                    print("\nTrying minimal UI...")
                    from minimal_dpg_ui import MinimalOneTrainerUI
                    app = MinimalOneTrainerUI()
                    app.run()
                    return True
                except Exception as e:
                    print(f"Minimal UI failed: {e}")
                    traceback.print_exc()
                    
                    # Clean up DPG context
                    try:
                        # Stop DPG if running
                        try:
                            if hasattr(dpg, 'is_dearpygui_running') and dpg.is_dearpygui_running():
                                dpg.stop_dearpygui()
                        except:
                            pass
                        
                        dpg.destroy_context()
                    except:
                        pass
                    
                    # Last resort - create an ultra minimal UI on the fly
                    print("\nCreating emergency UI...")
                    return create_emergency_ui()
    except Exception as e:
        print(f"Critical error in UI launcher: {e}")
        traceback.print_exc()
        return False

def create_emergency_ui():
    """Create an emergency ultra-minimal UI when everything else fails"""
    try:
        # Initialize DPG
        dpg.create_context()
        dpg.create_viewport(title="OneTrainer Emergency UI", width=800, height=600)
        dpg.setup_dearpygui()
        
        # Create main window with very basic UI
        with dpg.window(tag="emergency_window"):
            dpg.add_text("OneTrainer Emergency UI", color=(255, 255, 0))
            dpg.add_text("The regular UI implementations failed to load.")
            dpg.add_separator()
            
            # Model type selection
            dpg.add_text("Model Type:", color=(220, 220, 220))
            model_types = [model_type.name for model_type in ModelType]
            dpg.add_combo(items=model_types, default_value="SD", width=300)
            
            # Training method selection
            dpg.add_text("Training Method:", color=(220, 220, 220))
            from modules.util.enum.TrainingMethod import TrainingMethod
            training_methods = [method.name for method in TrainingMethod]
            dpg.add_combo(items=training_methods, default_value="LORA", width=300)
            
            # PEFT Type selection with warning for HiDream + LoKr
            dpg.add_text("PEFT Type:", color=(220, 220, 220))
            peft_types = [peft_type.name for peft_type in modules.util.enum.ModelType.PeftType]
            peft_combo = dpg.add_combo(items=peft_types, default_value="LORA", width=300, callback=lambda s, a: update_warning(a))
            
            # Warning text for HiDream + LoKr
            warning_text = dpg.add_text("", color=(255, 100, 100), wrap=700, show=False)
            
            def update_warning(peft_type):
                """Update warning visibility based on selections"""
                model_type_value = dpg.get_value(model_types_combo)
                if model_type_value == "HI_DREAM_FULL" and peft_type == "LOKR":
                    dpg.set_value(warning_text, "⚠️ WARNING: LoKr is known to cause sampling issues with HiDream models!")
                    dpg.configure_item(warning_text, show=True)
                else:
                    dpg.configure_item(warning_text, show=False)
            
            # Model type combo with callback
            model_types_combo = dpg.add_combo(
                items=model_types, 
                default_value="STABLE_DIFFUSION_15", 
                width=300,
                before=peft_combo,
                callback=lambda s, a: update_warning(dpg.get_value(peft_combo))
            )
            
            # Exit button
            dpg.add_separator()
            dpg.add_button(label="Exit", callback=lambda: dpg.stop_dearpygui())
        
        # Show viewport and run
        dpg.set_primary_window("emergency_window", True)
        dpg.show_viewport()
        dpg.start_dearpygui()
        return True
    except Exception as e:
        print(f"Emergency UI failed: {e}")
        traceback.print_exc()
        return False

# Main entry point
if __name__ == "__main__":
    try:
        success = run_ui()
        if not success:
            print("\nAll UI implementations failed. Please check error messages above.")
            print("Try running OneTrainer with a different UI using ./start-ui.sh")
    except KeyboardInterrupt:
        print("\nUser interrupted execution.")
    except Exception as e:
        print(f"\nUnexpected error: {e}")
        traceback.print_exc()
    finally:
        # Make sure DPG is cleaned up
        try:
            # Different ways to check if DPG is running (depends on version)
            dpg_running = False
            try:
                if hasattr(dpg, '_dpg_internal') and hasattr(dpg._dpg_internal, 'is_dearpygui_running'):
                    dpg_running = dpg._dpg_internal.is_dearpygui_running()
                elif hasattr(dpg, 'is_dearpygui_running'):
                    dpg_running = dpg.is_dearpygui_running()
            except:
                pass
                
            if dpg_running:
                dpg.stop_dearpygui()
                
            dpg.destroy_context()
        except:
            pass