dearpygui_train_ui.py

#!/usr/bin/env python3
"""
Prototype implementation of OneTrainer UI using Dear PyGui instead of CustomTkinter
This file follows the structure of the original TrainUI.py file but uses Dear PyGui
"""

import json
import os
import threading
import traceback
import webbrowser
from collections.abc import Callable
from pathlib import Path
import sys

# Add current directory to Python path to ensure imports work
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# Dear PyGui
import dearpygui.dearpygui as dpg

# Try importing modules from original codebase
try:
    from modules.util.config.TrainConfig import TrainConfig
    from modules.util.config.SampleConfig import SampleConfig
    from modules.trainer.GenericTrainer import GenericTrainer
    from modules.trainer.CloudTrainer import CloudTrainer
    from modules.util.callbacks.TrainCallbacks import TrainCallbacks
    from modules.util.commands.TrainCommands import TrainCommands
    from modules.util.enum.DataType import DataType
    from modules.util.enum.ImageFormat import ImageFormat
    from modules.util.enum.ModelType import ModelType
    from modules.util.enum.TrainingMethod import TrainingMethod
    from modules.util.torch_util import torch_gc
    from modules.util.TrainProgress import TrainProgress
    from modules.zluda import ZLUDA
    HAS_MODULES = True
except ImportError:
    # Mock for testing without original modules
    print("Warning: Original modules not found. Using mock classes for testing.")
    from dataclasses import dataclass
    
    @dataclass
    class TrainConfig:
        @staticmethod
        def default_values():
            return TrainConfig()
            
    @dataclass
    class UIState:
        def __init__(self, master, config):
            self.master = master
            self.config = config
            
        def get_var(self, var_name):
            return UIVar()
            
    class UIVar:
        def __init__(self):
            self.value = ""
            
        def set(self, value):
            self.value = value
            
        def get(self):
            return self.value
            
        def trace_add(self, *args, **kwargs):
            pass
            
    # Mock classes
    ModelType = type('ModelType', (), {})
    TrainingMethod = type('TrainingMethod', (), {})
    HAS_MODULES = False


class DPGUIState:
    """State management class similar to UIState from the original codebase, but for Dear PyGui"""
    
    def __init__(self, train_config):
        """Initialize state with reference to train_config"""
        self.train_config = train_config
        self.vars = {}
        self.callback_registry = {}
        self.callback_count = 0
        
    def register_var(self, var_name, initial_value=None):
        """Register a variable in the state system"""
        # Create registry entry if needed
        if var_name not in self.vars:
            self.vars[var_name] = {
                "value": initial_value,
                "callbacks": {}
            }
        return var_name
        
    def get_var_value(self, var_name):
        """Get the current value of a variable"""
        if var_name in self.vars:
            return self.vars[var_name]["value"]
        return None
        
    def set_var_value(self, var_name, value, trigger_callbacks=True):
        """Set the value of a variable and trigger callbacks if needed"""
        if var_name not in self.vars:
            self.register_var(var_name, value)
        else:
            self.vars[var_name]["value"] = value
            
        # Trigger callbacks if requested
        if trigger_callbacks:
            for callback_id, callback in self.vars[var_name]["callbacks"].items():
                try:
                    callback(value)
                except Exception as e:
                    print(f"Error in callback for {var_name}: {e}")
                    traceback.print_exc()
    
    def add_var_callback(self, var_name, callback):
        """Add a callback to be triggered when a variable changes"""
        if var_name not in self.vars:
            self.register_var(var_name)
            
        callback_id = self.callback_count
        self.callback_count += 1
        self.vars[var_name]["callbacks"][callback_id] = callback
        return callback_id
        
    def remove_var_callback(self, var_name, callback_id):
        """Remove a callback from a variable"""
        if var_name in self.vars and callback_id in self.vars[var_name]["callbacks"]:
            del self.vars[var_name]["callbacks"][callback_id]
    
    def update_from_config(self, config):
        """Update state values from a config object"""
        # Similar to the update method in the original UIState
        pass


class DearPyGuiComponents:
    """Helper class to create UI components consistently"""
    
    @staticmethod
    def label(parent, text, tag=None, tooltip=None, **kwargs):
        """Create a label with optional tooltip"""
        tag = tag or f"label_{dpg.generate_uuid()}"
        dpg.add_text(text, parent=parent, tag=tag, **kwargs)
        
        if tooltip:
            with dpg.tooltip(parent=tag):
                dpg.add_text(tooltip)
                
        return tag
        
    @staticmethod
    def button(parent, text, callback=None, tag=None, tooltip=None, width=120, height=30, **kwargs):
        """Create a button with optional tooltip"""
        tag = tag or f"button_{dpg.generate_uuid()}"
        dpg.add_button(label=text, callback=callback, parent=parent, 
                       tag=tag, width=width, height=height, **kwargs)
        
        if tooltip:
            with dpg.tooltip(parent=tag):
                dpg.add_text(tooltip)
                
        return tag
        
    @staticmethod
    def input_text(parent, default_value="", callback=None, tag=None, tooltip=None, width=200, **kwargs):
        """Create a text input field with optional tooltip"""
        tag = tag or f"input_{dpg.generate_uuid()}"
        dpg.add_input_text(default_value=default_value, callback=callback, 
                          parent=parent, tag=tag, width=width, **kwargs)
        
        if tooltip:
            with dpg.tooltip(parent=tag):
                dpg.add_text(tooltip)
                
        return tag
        
    @staticmethod
    def checkbox(parent, label, default_value=False, callback=None, tag=None, tooltip=None, **kwargs):
        """Create a checkbox with optional tooltip"""
        tag = tag or f"checkbox_{dpg.generate_uuid()}"
        dpg.add_checkbox(label=label, default_value=default_value, 
                        callback=callback, parent=parent, tag=tag, **kwargs)
        
        if tooltip:
            with dpg.tooltip(parent=tag):
                dpg.add_text(tooltip)
                
        return tag
        
    @staticmethod
    def combo(parent, items, default_value=None, callback=None, tag=None, tooltip=None, width=200, **kwargs):
        """Create a dropdown with optional tooltip"""
        tag = tag or f"combo_{dpg.generate_uuid()}"
        
        if default_value is None and items:
            default_value = items[0]
            
        dpg.add_combo(items=items, default_value=default_value, 
                     callback=callback, parent=parent, tag=tag, width=width, **kwargs)
        
        if tooltip:
            with dpg.tooltip(parent=tag):
                dpg.add_text(tooltip)
                
        return tag
        
    @staticmethod
    def slider_float(parent, default_value=0.0, min_value=0.0, max_value=1.0, 
                    callback=None, tag=None, tooltip=None, width=200, **kwargs):
        """Create a floating point slider with optional tooltip"""
        tag = tag or f"slider_{dpg.generate_uuid()}"
        dpg.add_slider_float(default_value=default_value, min_value=min_value, max_value=max_value, 
                            callback=callback, parent=parent, tag=tag, width=width, **kwargs)
        
        if tooltip:
            with dpg.tooltip(parent=tag):
                dpg.add_text(tooltip)
                
        return tag
        
    @staticmethod
    def slider_int(parent, default_value=0, min_value=0, max_value=100, 
                  callback=None, tag=None, tooltip=None, width=200, **kwargs):
        """Create an integer slider with optional tooltip"""
        tag = tag or f"slider_{dpg.generate_uuid()}"
        dpg.add_slider_int(default_value=default_value, min_value=min_value, max_value=max_value, 
                          callback=callback, parent=parent, tag=tag, width=width, **kwargs)
        
        if tooltip:
            with dpg.tooltip(parent=tag):
                dpg.add_text(tooltip)
                
        return tag
        
    @staticmethod
    def file_selector(parent, directory_selector=False, callback=None, 
                     default_path="", tag=None, tooltip=None, width=200, **kwargs):
        """Create a file dialog opener with a text field and button"""
        tag = tag or f"file_selector_{dpg.generate_uuid()}"
        
        # Create a group to hold the components
        group_tag = f"{tag}_group"
        dpg.add_group(horizontal=True, parent=parent, tag=group_tag)
        
        # Text input for the path
        input_tag = f"{tag}_input"
        dpg.add_input_text(default_value=default_path, parent=group_tag, tag=input_tag, width=width-50, **kwargs)
        
        # Create the file dialog
        dialog_tag = f"{tag}_dialog"
        if directory_selector:
            dpg.add_file_dialog(tag=dialog_tag, directory_selector=True, show=False,
                                callback=lambda s, a, u: dpg.set_value(input_tag, a["file_path_name"]))
        else:
            dpg.add_file_dialog(tag=dialog_tag, show=False, 
                                callback=lambda s, a, u: dpg.set_value(input_tag, a["file_path_name"]))
        
        # Button to open dialog  
        button_tag = f"{tag}_button"
        dpg.add_button(label="...", parent=group_tag, tag=button_tag, width=30, height=23,
                       callback=lambda: dpg.show_item(dialog_tag))
        
        if tooltip:
            with dpg.tooltip(parent=input_tag):
                dpg.add_text(tooltip)
                
        return {
            "group": group_tag, 
            "input": input_tag,
            "button": button_tag,
            "dialog": dialog_tag
        }
        
    @staticmethod
    def progress_bar(parent, default_value=0.0, tag=None, tooltip=None, width=200, **kwargs):
        """Create a progress bar with optional tooltip"""
        tag = tag or f"progress_{dpg.generate_uuid()}"
        dpg.add_progress_bar(default_value=default_value, parent=parent, tag=tag, width=width, **kwargs)
        
        if tooltip:
            with dpg.tooltip(parent=tag):
                dpg.add_text(tooltip)
                
        return tag
        
    @staticmethod
    def tab_bar(parent, tag=None, **kwargs):
        """Create a tab bar container"""
        tag = tag or f"tab_bar_{dpg.generate_uuid()}"
        return dpg.add_tab_bar(parent=parent, tag=tag, **kwargs)
        
    @staticmethod
    def tab(parent, label, tag=None, **kwargs):
        """Create a tab within a tab bar"""
        tag = tag or f"tab_{label}_{dpg.generate_uuid()}"
        return dpg.add_tab(label=label, parent=parent, tag=tag, **kwargs)
        
    @staticmethod
    def menu_bar(parent, tag=None, **kwargs):
        """Create a menu bar"""
        tag = tag or f"menu_bar_{dpg.generate_uuid()}"
        return dpg.add_menu_bar(parent=parent, tag=tag, **kwargs)
        
    @staticmethod
    def menu(parent, label, tag=None, **kwargs):
        """Create a menu within a menu bar"""
        tag = tag or f"menu_{dpg.generate_uuid()}"
        return dpg.add_menu(label=label, parent=parent, tag=tag, **kwargs)
        
    @staticmethod
    def menu_item(parent, label, callback=None, tag=None, **kwargs):
        """Create a menu item"""
        tag = tag or f"menu_item_{dpg.generate_uuid()}"
        return dpg.add_menu_item(label=label, callback=callback, parent=parent, tag=tag, **kwargs)


class DearPyGuiTrainUI:
    """Main UI class using Dear PyGui instead of CustomTkinter"""
    
    def __init__(self):
        """Initialize Dear PyGui and create the main window"""
        # Initialize Dear PyGui
        dpg.create_context()
        dpg.create_viewport(title="OneTrainer", width=1280, height=800)
        
        # Set up the DPG configuration
        with dpg.font_registry():
            # Add default font
            default_font = dpg.add_font("resources/fonts/OpenSans-Regular.ttf", 16) if os.path.exists("resources/fonts/OpenSans-Regular.ttf") else None
            if default_font:
                dpg.bind_font(default_font)
            
        with dpg.theme() as global_theme:
            with dpg.theme_component(dpg.mvAll):
                # Theme similar to CustomTkinter dark-blue
                dpg.add_theme_color(dpg.mvThemeCol_WindowBg, [26, 27, 38, 255])  # Dark blue background
                dpg.add_theme_color(dpg.mvThemeCol_TitleBg, [32, 34, 48, 255])  # Title bar
                dpg.add_theme_color(dpg.mvThemeCol_TitleBgActive, [39, 41, 57, 255])  # Active title
                dpg.add_theme_color(dpg.mvThemeCol_Button, [41, 83, 154, 255])  # Buttons (blue)
                dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, [59, 113, 202, 255])  # Buttons hovered
                dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, [66, 127, 227, 255])  # Buttons active
                dpg.add_theme_color(dpg.mvThemeCol_FrameBg, [37, 38, 51, 255])  # Frame background
                dpg.add_theme_color(dpg.mvThemeCol_CheckMark, [41, 83, 154, 255])  # Checkboxes
                dpg.add_theme_color(dpg.mvThemeCol_ScrollbarBg, [30, 31, 41, 255])  # Scrollbar bg
                dpg.add_theme_color(dpg.mvThemeCol_ScrollbarGrab, [41, 83, 154, 255])  # Scrollbar
                dpg.add_theme_color(dpg.mvThemeCol_Header, [41, 83, 154, 255])  # Headers
                
                # Add spacing to match CustomTkinter look
                dpg.add_theme_style(dpg.mvStyleVar_FramePadding, [8, 5])
                dpg.add_theme_style(dpg.mvStyleVar_ItemSpacing, [8, 5])
                dpg.add_theme_style(dpg.mvStyleVar_WindowPadding, [10, 10])
                dpg.add_theme_style(dpg.mvStyleVar_FrameRounding, 4.0)  # Rounded corners
                dpg.add_theme_style(dpg.mvStyleVar_GrabRounding, 4.0)  # Rounded corners
            
        dpg.bind_theme(global_theme)
        
        # Create train configuration if modules are available
        if HAS_MODULES:
            self.train_config = TrainConfig.default_values()
            self.ui_state = DPGUIState(self.train_config)
        else:
            # Mock objects for testing
            self.train_config = TrainConfig()
            self.ui_state = DPGUIState(self.train_config)
        
        # Create windows and UI elements
        self.setup_primary_window()
        
        # Training attributes
        self.training_thread = None
        self.training_callbacks = None
        self.training_commands = None
        
        # Start the viewport
        dpg.setup_dearpygui()
        dpg.show_viewport()
        
    def setup_primary_window(self):
        """Create the main window layout"""
        # Create the main window
        with dpg.window(tag="primary_window", label="OneTrainer", no_title_bar=True, no_resize=False,
                       no_move=False, no_close=False, no_collapse=False, no_scrollbar=True):
            
            # Layout with top bar, main area and bottom bar
            with dpg.group(horizontal=False):
                # Top bar (app title, config selector, model selection)
                self.create_top_bar()
                
                # Main content area with tabs
                self.create_tabs_content()
                
                # Bottom bar (progress bars, status, buttons)
                self.create_bottom_bar()
    
    def create_top_bar(self):
        """Create the top bar with logo, preset selector and model type selection"""
        with dpg.group(horizontal=True, tag="top_bar"):
            # App title and logo (left side)
            with dpg.group(horizontal=True):
                # Logo would go here
                DearPyGuiComponents.label(parent="top_bar", text="OneTrainer", tag="app_title")
            
            # Config/preset selection
            with dpg.group(horizontal=True):
                configs = [("Default Config", "#.json")]  # Would be populated from available configs
                
                # Config dropdown - wide for better preset names visibility
                DearPyGuiComponents.combo(
                    parent="top_bar", 
                    items=[name for name, _ in configs],
                    default_value=configs[0][0] if configs else "",
                    callback=self.load_config,
                    tag="config_selector",
                    width=350
                )
                
                # Save config button
                DearPyGuiComponents.button(
                    parent="top_bar",
                    text="Save Config",
                    callback=self.save_config,
                    tooltip="Save the current configuration as a preset"
                )
                
                # Wiki button
                DearPyGuiComponents.button(
                    parent="top_bar",
                    text="Wiki",
                    callback=lambda: webbrowser.open("https://github.com/Nerogar/OneTrainer/wiki", new=0, autoraise=False),
                    tooltip="Open OneTrainer Wiki"
                )
            
            # Flexible space
            dpg.add_spacer(width=20)
            
            # Model type selection (right side)
            with dpg.group(horizontal=True):
                model_types = [
                    "Stable Diffusion 1.5",
                    "Stable Diffusion 1.5 Inpainting",
                    "Stable Diffusion 2.0",
                    "Stable Diffusion 2.0 Inpainting",
                    "Stable Diffusion 2.1",
                    "Stable Diffusion 3",
                    "Stable Diffusion 3.5",
                    "Stable Diffusion XL 1.0 Base",
                    "Stable Diffusion XL 1.0 Base Inpainting",
                    "Wuerstchen v2",
                    "Stable Cascade",
                    "PixArt Alpha",
                    "PixArt Sigma",
                    "Flux Dev",
                    "Flux Fill Dev",
                    "Sana",
                    "Hunyuan Video",
                    "HiDream Full"
                ]
                
                DearPyGuiComponents.combo(
                    parent="top_bar",
                    items=model_types,
                    default_value=model_types[0],
                    callback=self.change_model_type,
                    tag="model_type_selector",
                    width=250
                )
                
                # Training method selection
                training_methods = ["Fine Tune", "LoRA", "Embedding"]
                DearPyGuiComponents.combo(
                    parent="top_bar",
                    items=training_methods,
                    default_value=training_methods[0],
                    callback=self.change_training_method,
                    tag="training_method_selector",
                    width=120
                )
    
    def create_tabs_content(self):
        """Create the tabbed content area"""
        with dpg.child_window(tag="main_content", height=-40, no_scrollbar=True):
            tab_bar = DearPyGuiComponents.tab_bar(parent="main_content", tag="main_tabs")
            
            # General tab
            self.create_general_tab(DearPyGuiComponents.tab(tab_bar, "General", tag="general_tab"))
            
            # Model tab
            self.create_model_tab(DearPyGuiComponents.tab(tab_bar, "Model", tag="model_tab"))
            
            # Data tab
            self.create_data_tab(DearPyGuiComponents.tab(tab_bar, "Data", tag="data_tab"))
            
            # Concepts tab
            self.create_concepts_tab(DearPyGuiComponents.tab(tab_bar, "Concepts", tag="concepts_tab"))
            
            # Training tab
            self.create_training_tab(DearPyGuiComponents.tab(tab_bar, "Training", tag="training_tab"))
            
            # Sampling tab
            self.create_sampling_tab(DearPyGuiComponents.tab(tab_bar, "Sampling", tag="sampling_tab"))
            
            # Backup tab
            self.create_backup_tab(DearPyGuiComponents.tab(tab_bar, "Backup", tag="backup_tab"))
            
            # Tools tab
            self.create_tools_tab(DearPyGuiComponents.tab(tab_bar, "Tools", tag="tools_tab"))
            
            # Additional embeddings tab
            self.create_additional_embeddings_tab(DearPyGuiComponents.tab(tab_bar, "Additional Embeddings", tag="additional_embeddings_tab"))
            
            # Cloud tab
            self.create_cloud_tab(DearPyGuiComponents.tab(tab_bar, "Cloud", tag="cloud_tab"))
            
            # If training method is LoRA, add LoRA tab
            self.create_lora_tab(DearPyGuiComponents.tab(tab_bar, "LoRA", tag="lora_tab"))
    
    def create_bottom_bar(self):
        """Create the bottom bar with progress indicators and action buttons"""
        with dpg.group(horizontal=True, tag="bottom_bar"):
            # Progress indicators
            with dpg.group(horizontal=False, width=400):
                # Step progress
                DearPyGuiComponents.label(parent="bottom_bar", text="Step Progress", tag="step_progress_label")
                self.step_progress_bar = DearPyGuiComponents.progress_bar(
                    parent="bottom_bar", 
                    default_value=0.0,
                    tag="step_progress_bar",
                    width=390
                )
                
                # Epoch progress
                DearPyGuiComponents.label(parent="bottom_bar", text="Epoch Progress", tag="epoch_progress_label")
                self.epoch_progress_bar = DearPyGuiComponents.progress_bar(
                    parent="bottom_bar", 
                    default_value=0.0,
                    tag="epoch_progress_bar",
                    width=390
                )
            
            # Status label
            self.status_label = DearPyGuiComponents.label(
                parent="bottom_bar", 
                text="Ready", 
                tag="status_label"
            )
            
            # Flexible space
            dpg.add_spacer(width=20)
            
            # Action buttons
            DearPyGuiComponents.button(
                parent="bottom_bar",
                text="Tensorboard",
                callback=self.open_tensorboard,
                tooltip="Open TensorBoard visualization"
            )
            
            self.training_button = DearPyGuiComponents.button(
                parent="bottom_bar",
                text="Start Training",
                callback=self.start_training,
                tooltip="Start or stop the training process"
            )
            
            DearPyGuiComponents.button(
                parent="bottom_bar",
                text="Export",
                callback=self.export_training,
                tooltip="Export the current configuration as a script to run without a UI"
            )
    
    def create_general_tab(self, parent):
        """Create the general tab content"""
        with dpg.group(parent=parent, horizontal=False):
            with dpg.collapsing_header(label="Workspace Settings", default_open=True):
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Workspace Directory",
                        tooltip="The directory where all files of this training run are saved"
                    )
                    workspace_selector = DearPyGuiComponents.file_selector(
                        parent=dpg.last_item(),
                        directory_selector=True,
                        default_path="./workspace",
                        tooltip="The directory where all files of this training run are saved",
                        width=400
                    )
                
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Cache Directory",
                        tooltip="The directory where cached data is saved"
                    )
                    cache_selector = DearPyGuiComponents.file_selector(
                        parent=dpg.last_item(),
                        directory_selector=True,
                        default_path="./cache",
                        tooltip="The directory where cached data is saved",
                        width=400
                    )
            
            with dpg.collapsing_header(label="Training Options", default_open=True):
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.checkbox(
                        parent=dpg.last_item(),
                        label="Continue from last backup",
                        tooltip="Automatically continues training from the last backup saved in <workspace>/backup"
                    )
                
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.checkbox(
                        parent=dpg.last_item(),
                        label="Only Cache",
                        tooltip="Only populate the cache, without any training"
                    )
                    
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.checkbox(
                        parent=dpg.last_item(),
                        label="Debug mode",
                        tooltip="Save debug information during the training into the debug directory"
                    )
                
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Debug Directory",
                        tooltip="The directory where debug data is saved"
                    )
                    debug_selector = DearPyGuiComponents.file_selector(
                        parent=dpg.last_item(),
                        directory_selector=True,
                        default_path="./debug",
                        tooltip="The directory where debug data is saved",
                        width=400
                    )
            
            with dpg.collapsing_header(label="Tensorboard Options", default_open=True):
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.checkbox(
                        parent=dpg.last_item(),
                        label="Tensorboard",
                        tooltip="Starts the Tensorboard Web UI during training"
                    )
                
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.checkbox(
                        parent=dpg.last_item(),
                        label="Expose Tensorboard",
                        tooltip="Exposes Tensorboard Web UI to all network interfaces (makes it accessible from the network)"
                    )
                    
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Tensorboard Port",
                        tooltip="Port to use for Tensorboard link"
                    )
                    DearPyGuiComponents.input_text(
                        parent=dpg.last_item(),
                        default_value="6006",
                        width=100
                    )
            
            with dpg.collapsing_header(label="Device Settings", default_open=True):
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Dataloader Threads",
                        tooltip="Number of threads used for the data loader"
                    )
                    DearPyGuiComponents.input_text(
                        parent=dpg.last_item(),
                        default_value="4",
                        width=100
                    )
                
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Train Device",
                        tooltip="The device used for training. Can be \"cuda\", \"cuda:0\", \"cuda:1\" etc. Default:\"cuda\""
                    )
                    DearPyGuiComponents.input_text(
                        parent=dpg.last_item(),
                        default_value="cuda",
                        width=100
                    )
                    
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Temp Device",
                        tooltip="The device used to temporarily offload models while they are not used. Default:\"cpu\""
                    )
                    DearPyGuiComponents.input_text(
                        parent=dpg.last_item(),
                        default_value="cpu",
                        width=100
                    )
    
    def create_model_tab(self, parent):
        """Create the model tab content"""
        # This is a simplification - would need to be dynamic based on model type
        with dpg.group(parent=parent):
            with dpg.collapsing_header(label="Model Settings", default_open=True):
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Base Model",
                        tooltip="The base model to train on"
                    )
                    base_model_selector = DearPyGuiComponents.file_selector(
                        parent=dpg.last_item(),
                        directory_selector=False,
                        tooltip="The base model to train on",
                        width=400
                    )
                
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(), 
                        text="Output Model",
                        tooltip="The destination to save the output model to"
                    )
                    output_model_selector = DearPyGuiComponents.file_selector(
                        parent=dpg.last_item(),
                        directory_selector=True,
                        tooltip="The destination to save the output model to",
                        width=400
                    )
    
    def create_data_tab(self, parent):
        """Create the data tab content"""
        with dpg.group(parent=parent):
            with dpg.collapsing_header(label="Data Processing Options", default_open=True):
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.checkbox(
                        parent=dpg.last_item(),
                        label="Aspect Ratio Bucketing",
                        tooltip="Aspect ratio bucketing enables training on images with different aspect ratios"
                    )
                
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.checkbox(
                        parent=dpg.last_item(),
                        label="Latent Caching",
                        tooltip="Caching of intermediate training data that can be re-used between epochs"
                    )
                
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.checkbox(
                        parent=dpg.last_item(),
                        label="Clear cache before training",
                        tooltip="Clears the cache directory before starting to train"
                    )
    
    def create_concepts_tab(self, parent):
        """Create the concepts tab content"""
        with dpg.group(parent=parent):
            # This would be much more complex in practice,
            # involving a table for concepts and buttons to add/edit/remove them
            with dpg.group(horizontal=True):
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Add Concept",
                    callback=self.add_concept
                )
                
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Edit Selected",
                    callback=self.edit_concept
                )
                
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Remove Selected",
                    callback=self.remove_concept
                )
            
            # Placeholder for concept list
            with dpg.group():
                DearPyGuiComponents.label(parent=dpg.last_item(), text="No concepts defined")
    
    def create_training_tab(self, parent):
        """Create the training tab content"""
        with dpg.group(parent=parent):
            with dpg.collapsing_header(label="Training Parameters", default_open=True):
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Training Method",
                        tooltip="The method used for training"
                    )
                    DearPyGuiComponents.combo(
                        parent=dpg.last_item(),
                        items=["LORA", "Embedding", "Fine-tune", "Fine-tune VAE"],
                        default_value="LORA",
                        width=200
                    )
                
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Resolution",
                        tooltip="The resolution used for training"
                    )
                    DearPyGuiComponents.input_text(
                        parent=dpg.last_item(),
                        default_value="512x512",
                        width=100
                    )
                    
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Batch Size",
                        tooltip="The batch size used for training"
                    )
                    DearPyGuiComponents.input_text(
                        parent=dpg.last_item(),
                        default_value="1",
                        width=100
                    )
                    
                with dpg.group(horizontal=True):
                    DearPyGuiComponents.label(
                        parent=dpg.last_item(),
                        text="Epochs",
                        tooltip="The number of epochs to train for"
                    )
                    DearPyGuiComponents.input_text(
                        parent=dpg.last_item(),
                        default_value="10",
                        width=100
                    )
    
    def create_sampling_tab(self, parent):
        """Create the sampling tab content"""
        with dpg.group(parent=parent):
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Sample After",
                    tooltip="The interval used when automatically sampling from the model during training"
                )
                DearPyGuiComponents.input_text(
                    parent=dpg.last_item(),
                    default_value="500",
                    width=100
                )
                DearPyGuiComponents.combo(
                    parent=dpg.last_item(),
                    items=["STEPS", "EPOCHS"],
                    default_value="STEPS",
                    width=100
                )
                
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Skip First",
                    tooltip="Start sampling automatically after this interval has elapsed"
                )
                DearPyGuiComponents.input_text(
                    parent=dpg.last_item(),
                    default_value="0",
                    width=50
                )
                
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Sample Now",
                    callback=self.sample_now
                )
                
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Manual Sample",
                    callback=self.open_sample_ui
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.checkbox(
                    parent=dpg.last_item(),
                    label="Non-EMA Sampling",
                    tooltip="Whether to include non-ema sampling when using ema"
                )
                
                DearPyGuiComponents.checkbox(
                    parent=dpg.last_item(),
                    label="Samples to Tensorboard",
                    tooltip="Whether to include sample images in the Tensorboard output"
                )
                
                DearPyGuiComponents.checkbox(
                    parent=dpg.last_item(),
                    label="Sample at Startup", 
                    tooltip="Whether to generate samples when the model is loaded, before training begins"
                )
            
            # Sample configurations would be displayed in a table below
            with dpg.group():
                DearPyGuiComponents.label(parent=dpg.last_item(), text="No sample configurations defined")
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Add Sample Config",
                    callback=self.add_sample_config
                )
    
    def create_backup_tab(self, parent):
        """Create the backup tab content"""
        with dpg.group(parent=parent):
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Backup After",
                    tooltip="The interval used when automatically creating model backups during training"
                )
                DearPyGuiComponents.input_text(
                    parent=dpg.last_item(),
                    default_value="5000",
                    width=100
                )
                DearPyGuiComponents.combo(
                    parent=dpg.last_item(),
                    items=["STEPS", "EPOCHS", "MINUTES", "HOURS"],
                    default_value="STEPS",
                    width=100
                )
                
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Backup Now",
                    callback=self.backup_now
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.checkbox(
                    parent=dpg.last_item(),
                    label="Rolling Backup",
                    tooltip="If rolling backups are enabled, older backups are deleted automatically"
                )
                
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Rolling Backup Count",
                    tooltip="Defines the number of backups to keep if rolling backups are enabled"
                )
                DearPyGuiComponents.input_text(
                    parent=dpg.last_item(),
                    default_value="5",
                    width=50
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.checkbox(
                    parent=dpg.last_item(),
                    label="Backup Before Save",
                    tooltip="Create a full backup before saving the final model"
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Save Every",
                    tooltip="The interval used when automatically saving the model during training"
                )
                DearPyGuiComponents.input_text(
                    parent=dpg.last_item(),
                    default_value="10000",
                    width=100
                )
                DearPyGuiComponents.combo(
                    parent=dpg.last_item(),
                    items=["STEPS", "EPOCHS", "MINUTES", "HOURS"],
                    default_value="STEPS",
                    width=100
                )
                
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Save Now",
                    callback=self.save_now
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Skip First",
                    tooltip="Start saving automatically after this interval has elapsed"
                )
                DearPyGuiComponents.input_text(
                    parent=dpg.last_item(),
                    default_value="0",
                    width=50
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Save Filename Prefix",
                    tooltip="The prefix for filenames used when saving the model during training"
                )
                DearPyGuiComponents.input_text(
                    parent=dpg.last_item(),
                    default_value="model",
                    width=200
                )
    
    def create_tools_tab(self, parent):
        """Create the tools tab content"""
        with dpg.group(parent=parent):
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Dataset Tools",
                    tooltip="Open the captioning tool"
                )
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Open",
                    callback=self.open_dataset_tool
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Video Tools",
                    tooltip="Open the video tools"
                )
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Open",
                    callback=self.open_video_tool
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Convert Model Tools",
                    tooltip="Open the model conversion tool"
                )
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Open",
                    callback=self.open_convert_model_tool
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Sampling Tool",
                    tooltip="Open the model sampling tool"
                )
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Open",
                    callback=self.open_sampling_tool
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Profiling Tool",
                    tooltip="Open the profiling tools"
                )
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Open",
                    callback=self.open_profiling_tool
                )
    
    def create_lora_tab(self, parent):
        """Create the LoRA tab content"""
        with dpg.group(parent=parent):
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="LoRA base model",
                    tooltip="The base LoRA to train on. Leave empty to create a new LoRA"
                )
                lora_model_selector = DearPyGuiComponents.file_selector(
                    parent=dpg.last_item(),
                    directory_selector=False,
                    tooltip="The base LoRA to train on. Leave empty to create a new LoRA",
                    width=400
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="LoRA rank",
                    tooltip="The rank parameter used when creating a new LoRA"
                )
                DearPyGuiComponents.input_text(
                    parent=dpg.last_item(),
                    default_value="16",
                    width=100
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="LoRA alpha",
                    tooltip="The alpha parameter used when creating a new LoRA"
                )
                DearPyGuiComponents.input_text(
                    parent=dpg.last_item(),
                    default_value="16",
                    width=100
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.checkbox(
                    parent=dpg.last_item(),
                    label="Bundle Embeddings",
                    tooltip="Bundles any additional embeddings into the LoRA output file, rather than as separate files"
                )
                
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="PEFT Type",
                    tooltip="Type of Parameter-Efficient Fine Tuning"
                )
                DearPyGuiComponents.combo(
                    parent=dpg.last_item(),
                    items=["LORA", "LOHA", "LOKR", "IA3", "DYLORA"],
                    default_value="LORA",
                    width=100,
                    tooltip="Select the PEFT type for the LoRA adapter"
                )
    
    def create_additional_embeddings_tab(self, parent):
        """Create the additional embeddings tab content"""
        with dpg.group(parent=parent):
            with dpg.group(horizontal=True):
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Add Embedding",
                    callback=self.add_embedding
                )
                
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Edit Selected",
                    callback=self.edit_embedding
                )
                
                DearPyGuiComponents.button(
                    parent=dpg.last_item(),
                    text="Remove Selected",
                    callback=self.remove_embedding
                )
            
            # Placeholder for embeddings list
            with dpg.group():
                DearPyGuiComponents.label(parent=dpg.last_item(), text="No additional embeddings defined")
    
    def create_cloud_tab(self, parent):
        """Create the cloud tab content"""
        with dpg.group(parent=parent):
            with dpg.group(horizontal=True):
                DearPyGuiComponents.checkbox(
                    parent=dpg.last_item(),
                    label="Enable Cloud Training",
                    tooltip="Enable training on a cloud instance"
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="Cloud Type",
                    tooltip="The type of cloud provider to use"
                )
                DearPyGuiComponents.combo(
                    parent=dpg.last_item(),
                    items=["RunPod", "Linux"],
                    default_value="RunPod",
                    width=150
                )
            
            with dpg.group(horizontal=True):
                DearPyGuiComponents.label(
                    parent=dpg.last_item(),
                    text="API Key",
                    tooltip="The API key to use for cloud provider authentication"
                )
                DearPyGuiComponents.input_text(
                    parent=dpg.last_item(),
                    default_value="",
                    width=300,
                    password=True
                )
    
    # Event handler methods
    def change_model_type(self, sender, app_data):
        """Handle changes to the model type dropdown"""
        print(f"Model type changed to: {app_data}")
        # In the actual implementation, this would update UI state and refresh tabs
    
    def change_training_method(self, sender, app_data):
        """Handle changes to the training method dropdown"""
        print(f"Training method changed to: {app_data}")
        # In the actual implementation, this would update UI state and refresh tabs
    
    def load_config(self, sender, app_data):
        """Handle loading a configuration preset"""
        print(f"Loading config: {app_data}")
        # In the actual implementation, this would load a config file
    
    def save_config(self):
        """Handle saving the current configuration as a preset"""
        print("Saving current config")
        # In the actual implementation, this would save the current configuration
    
    def start_training(self):
        """Handle the start/stop training button"""
        if self.training_thread is None:
            print("Starting training")
            dpg.set_item_label(self.training_button, "Stop Training")
            
            # In the actual implementation, this would start a training thread
            # self.training_thread = threading.Thread(target=self.__training_thread_function)
            # self.training_thread.start()
        else:
            print("Stopping training")
            dpg.set_item_label(self.training_button, "Start Training")
            dpg.configure_item(self.status_label, default_value="Stopping...")
            
            # In the actual implementation, this would stop the training thread
            # self.training_commands.stop()
    
    def export_training(self):
        """Handle exporting the current configuration"""
        print("Exporting training configuration")
        # In the actual implementation, this would export the configuration to a file
    
    def open_tensorboard(self):
        """Handle opening TensorBoard"""
        print("Opening TensorBoard")
        # In the actual implementation, this would open a web browser to TensorBoard
    
    def set_step_progress(self, current, maximum):
        """Update the step progress bar"""
        if maximum > 0:
            dpg.set_value(self.step_progress_bar, current / maximum)
            dpg.configure_item(self.step_progress_label, default_value=f"Step Progress: {current}/{maximum}")
    
    def set_epoch_progress(self, current, maximum):
        """Update the epoch progress bar"""
        if maximum > 0:
            dpg.set_value(self.epoch_progress_bar, current / maximum)
            dpg.configure_item(self.epoch_progress_label, default_value=f"Epoch Progress: {current}/{maximum}")
    
    def add_concept(self):
        """Handle adding a new concept"""
        print("Adding concept")
        # In the actual implementation, this would open a dialog to add a concept
    
    def edit_concept(self):
        """Handle editing a concept"""
        print("Editing concept")
        # In the actual implementation, this would open a dialog to edit a concept
    
    def remove_concept(self):
        """Handle removing a concept"""
        print("Removing concept")
        # In the actual implementation, this would remove a concept
    
    def add_embedding(self):
        """Handle adding a new embedding"""
        print("Adding embedding")
        # In the actual implementation, this would open a dialog to add an embedding
    
    def edit_embedding(self):
        """Handle editing an embedding"""
        print("Editing embedding")
        # In the actual implementation, this would open a dialog to edit an embedding
    
    def remove_embedding(self):
        """Handle removing an embedding"""
        print("Removing embedding")
        # In the actual implementation, this would remove an embedding
    
    def add_sample_config(self):
        """Handle adding a new sample configuration"""
        print("Adding sample configuration")
        # In the actual implementation, this would open a dialog to add a sample configuration
    
    def sample_now(self):
        """Handle triggering a sample now"""
        print("Sampling now")
        # In the actual implementation, this would trigger a sample
    
    def backup_now(self):
        """Handle triggering a backup now"""
        print("Backing up now")
        # In the actual implementation, this would trigger a backup
    
    def save_now(self):
        """Handle triggering a save now"""
        print("Saving now")
        # In the actual implementation, this would trigger a save
    
    def open_dataset_tool(self):
        """Handle opening the dataset tool"""
        print("Opening dataset tool")
        # In the actual implementation, this would open the dataset tool window
    
    def open_video_tool(self):
        """Handle opening the video tool"""
        print("Opening video tool")
        # In the actual implementation, this would open the video tool window
    
    def open_convert_model_tool(self):
        """Handle opening the convert model tool"""
        print("Opening convert model tool")
        # In the actual implementation, this would open the convert model tool window
    
    def open_sampling_tool(self):
        """Handle opening the sampling tool"""
        print("Opening sampling tool")
        # In the actual implementation, this would open the sampling tool window
    
    def open_profiling_tool(self):
        """Handle opening the profiling tool"""
        print("Opening profiling tool")
        # In the actual implementation, this would open the profiling tool window
    
    def open_sample_ui(self):
        """Handle opening the manual sample UI"""
        print("Opening manual sample UI")
        # In the actual implementation, this would open the sample UI window
    
    def shutdown(self):
        """Clean up resources and exit DearPyGui"""
        dpg.destroy_context()
    
    def run(self):
        """Start the main loop"""
        while dpg.is_dearpygui_running():
            # Render frame
            dpg.render_dearpygui_frame()
        
        # Clean up when done
        self.shutdown()


if __name__ == "__main__":
    # Create and run the UI
    ui = DearPyGuiTrainUI()
    ui.run()