standalone_app.py

#!/usr/bin/env python3
"""
Standalone OneTrainer UI with LyCORIS support

This script provides a standalone version of the OneTrainer UI that includes
LyCORIS support without requiring problematic dependencies.
"""

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

# Import the standalone LoRA tab
from standalone_lora_tab import LoraTab, Components, ModelType, PeftType, DataType

# Create minimal versions of required enums and classes
class TrainingMethod(enum.Enum):
    FINE_TUNE = "FINE_TUNE"
    LORA = "LORA"
    EMBEDDING = "EMBEDDING"
    
    def __str__(self):
        return self.value

class SimpleStateManager:
    """Simple state manager for the application"""
    
    def __init__(self):
        """Initialize the state manager"""
        self.model_type = ModelType.STABLE_DIFFUSION_15
        self.training_method = TrainingMethod.LORA
        self.peft_type = PeftType.LORA
        self.lora_rank = 32
        self.lora_alpha = 32.0
        self.dropout_probability = 0.0
        self.bundle_embeddings = False
        self.layer_preset = "full"
        self.layer_text = ""
        
    def get_model_type(self):
        """Get the model type"""
        return self.model_type
        
    def get_lora_state(self):
        """Get the LoRA state"""
        return {
            "peft_type": self.peft_type,
            "lora_rank": self.lora_rank,
            "lora_alpha": self.lora_alpha,
            "dropout_probability": self.dropout_probability,
            "bundle_embeddings": self.bundle_embeddings,
            "layer_preset": self.layer_preset,
            "layers": self.layer_text,
        }
        
    def update_from_lora_tab(self, state):
        """Update the state from the LoRA tab"""
        if "peft_type" in state:
            self.peft_type = state["peft_type"]
        if "lora_rank" in state:
            self.lora_rank = state["lora_rank"]
        if "lora_alpha" in state:
            self.lora_alpha = state["lora_alpha"]
        if "dropout_probability" in state:
            self.dropout_probability = state["dropout_probability"]
        if "bundle_embeddings" in state:
            self.bundle_embeddings = state["bundle_embeddings"]
        if "layer_preset" in state:
            self.layer_preset = state["layer_preset"]
        if "layers" in state:
            self.layer_text = state["layers"]


# Tab base class
class TabBase:
    """Base class for tabs"""
    
    def __init__(self, parent, state_manager):
        """Initialize the tab"""
        self.parent = parent
        self.state_manager = state_manager
        self.components = {}
        
    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)


# Model Tab
class ModelTab(TabBase):
    """Model tab implementation"""
    
    def __init__(self, parent, state_manager):
        """Initialize the model tab"""
        super().__init__(parent, state_manager)
        self.create_ui()
        
    def create_ui(self):
        """Create the UI elements"""
        with dpg.group(parent=self.parent):
            # Title
            title = Components.label(dpg.last_item(), "Model Settings")
            self.add_component("title", title)
            
            # Model type selection
            with dpg.child_window(height=-1, width=-1):
                with dpg.group():
                    # Model type
                    with dpg.group(horizontal=True):
                        Components.label(dpg.last_item(), "Model Type:")
                        
                        # Model type options
                        model_types = [
                            ("Stable Diffusion 1.5", ModelType.STABLE_DIFFUSION_15),
                            ("Stable Diffusion XL", ModelType.STABLE_DIFFUSION_XL_10_BASE),
                            ("Stable Diffusion 3", ModelType.STABLE_DIFFUSION_3),
                            ("HiDream", ModelType.HI_DREAM_FULL),
                            ("PixArt Alpha", ModelType.PIXART_ALPHA),
                            ("FLUX", ModelType.FLUX_DEV_1),
                            ("SANA", ModelType.SANA),
                            ("Hunyuan Video", ModelType.HUNYUAN_VIDEO),
                        ]
                        
                        model_combo = Components.options_kv(
                            dpg.last_item(),
                            values=model_types,
                            default_value=self.state_manager.get_model_type(),
                            width=200,
                            on_change=self.on_model_type_change
                        )
                        self.add_component("model_type", model_combo)
                    
                    # Training method
                    with dpg.group(horizontal=True):
                        Components.label(dpg.last_item(), "Training Method:")
                        
                        # Training method options
                        training_methods = [
                            ("LoRA / LyCORIS", TrainingMethod.LORA),
                            ("Fine-Tune", TrainingMethod.FINE_TUNE),
                            ("Embedding", TrainingMethod.EMBEDDING),
                        ]
                        
                        training_combo = Components.options_kv(
                            dpg.last_item(),
                            values=training_methods,
                            default_value=self.state_manager.training_method,
                            width=200,
                            on_change=self.on_training_method_change
                        )
                        self.add_component("training_method", training_combo)
                    
                    # Base model path
                    with dpg.group(horizontal=True):
                        Components.label(dpg.last_item(), "Base Model Path:")
                        model_selector = Components.file_selector(
                            dpg.last_item(),
                            default_path="",
                            width=400
                        )
                        self.add_component("model_path", model_selector)
                        
                    # More settings would go here in a full implementation
                    with dpg.group():
                        Components.label(
                            dpg.last_item(),
                            "This is a simplified model tab. In the full OneTrainer app,\n"
                            "this tab would contain more model configuration options.",
                            wrap=600
                        )
    
    def on_model_type_change(self, model_type):
        """Handle model type change"""
        self.state_manager.model_type = model_type
        print(f"Changed model type to: {model_type}")
        
    def on_training_method_change(self, method):
        """Handle training method change"""
        self.state_manager.training_method = method
        print(f"Changed training method to: {method}")


# Training Tab
class TrainingTab(TabBase):
    """Training tab implementation"""
    
    def __init__(self, parent, state_manager):
        """Initialize the training tab"""
        super().__init__(parent, state_manager)
        self.create_ui()
        
    def create_ui(self):
        """Create the UI elements"""
        with dpg.group(parent=self.parent):
            # Title
            title = Components.label(dpg.last_item(), "Training Settings")
            self.add_component("title", title)
            
            # Training settings
            with dpg.child_window(height=-1, width=-1):
                with dpg.group():
                    # Learning rate
                    with dpg.group(horizontal=True):
                        Components.label(dpg.last_item(), "Learning Rate:")
                        lr_input = Components.input_float(
                            dpg.last_item(),
                            default_value=0.0001,
                            width=200
                        )
                        self.add_component("learning_rate", lr_input)
                    
                    # Batch size
                    with dpg.group(horizontal=True):
                        Components.label(dpg.last_item(), "Batch Size:")
                        batch_input = Components.input_int(
                            dpg.last_item(),
                            default_value=1,
                            width=200
                        )
                        self.add_component("batch_size", batch_input)
                    
                    # Epochs
                    with dpg.group(horizontal=True):
                        Components.label(dpg.last_item(), "Epochs:")
                        epochs_input = Components.input_int(
                            dpg.last_item(),
                            default_value=10,
                            width=200
                        )
                        self.add_component("epochs", epochs_input)
                    
                    # Train button
                    with dpg.group(horizontal=True):
                        train_button = Components.button(
                            dpg.last_item(),
                            text="Start Training",
                            callback=self.on_train_button,
                            width=200
                        )
                        self.add_component("train_button", train_button)
                        
                        # Status text
                        status_text = Components.label(
                            dpg.last_item(),
                            ""
                        )
                        self.add_component("status_text", status_text)
                        
                    # More settings would go here in a full implementation
                    with dpg.group():
                        Components.label(
                            dpg.last_item(),
                            "This is a simplified training tab. In the full OneTrainer app,\n"
                            "this tab would contain more training configuration options.",
                            wrap=600
                        )
    
    def on_train_button(self):
        """Handle train button click"""
        # Get current settings
        learning_rate = dpg.get_value(self.get_component("learning_rate"))
        batch_size = dpg.get_value(self.get_component("batch_size"))
        epochs = dpg.get_value(self.get_component("epochs"))
        
        # Update status
        dpg.set_value(
            self.get_component("status_text"),
            f"Training started with LR={learning_rate}, Batch={batch_size}, Epochs={epochs}"
        )
        
        # In a real implementation, this would start the training process
        print(f"Would start training with: LR={learning_rate}, Batch={batch_size}, Epochs={epochs}")
        print(f"Current model type: {self.state_manager.model_type}")
        print(f"Current training method: {self.state_manager.training_method}")
        print(f"Current PEFT type: {self.state_manager.peft_type}")


# Concept Tab
class ConceptTab(TabBase):
    """Concept tab implementation"""
    
    def __init__(self, parent, state_manager):
        """Initialize the concept tab"""
        super().__init__(parent, state_manager)
        self.create_ui()
        
    def create_ui(self):
        """Create the UI elements"""
        with dpg.group(parent=self.parent):
            # Title
            title = Components.label(dpg.last_item(), "Concept Settings")
            self.add_component("title", title)
            
            # Concept settings
            with dpg.child_window(height=-1, width=-1):
                with dpg.group():
                    # Add concept button
                    with dpg.group(horizontal=True):
                        add_button = Components.button(
                            dpg.last_item(),
                            text="Add Concept",
                            callback=self.on_add_concept,
                            width=200
                        )
                        self.add_component("add_button", add_button)
                    
                    # Concept list
                    with dpg.group():
                        Components.label(
                            dpg.last_item(),
                            "This is a simplified concept tab. In the full OneTrainer app,\n"
                            "this tab would contain a list of concepts for training.",
                            wrap=600
                        )
    
    def on_add_concept(self):
        """Handle add concept button click"""
        print("Would add a new concept")


# Sampling Tab
class SamplingTab(TabBase):
    """Sampling tab implementation"""
    
    def __init__(self, parent, state_manager):
        """Initialize the sampling tab"""
        super().__init__(parent, state_manager)
        self.create_ui()
        
    def create_ui(self):
        """Create the UI elements"""
        with dpg.group(parent=self.parent):
            # Title
            title = Components.label(dpg.last_item(), "Sampling Settings")
            self.add_component("title", title)
            
            # Sampling settings
            with dpg.child_window(height=-1, width=-1):
                with dpg.group():
                    # Prompt
                    with dpg.group():
                        Components.label(dpg.last_item(), "Prompt:")
                        prompt_input = Components.entry(
                            dpg.last_item(),
                            default_value="",
                            width=400
                        )
                        self.add_component("prompt", prompt_input)
                    
                    # Negative prompt
                    with dpg.group():
                        Components.label(dpg.last_item(), "Negative Prompt:")
                        neg_prompt_input = Components.entry(
                            dpg.last_item(),
                            default_value="",
                            width=400
                        )
                        self.add_component("negative_prompt", neg_prompt_input)
                    
                    # CFG Scale
                    with dpg.group(horizontal=True):
                        Components.label(dpg.last_item(), "CFG Scale:")
                        cfg_input = Components.input_float(
                            dpg.last_item(),
                            default_value=7.0,
                            width=200
                        )
                        self.add_component("cfg_scale", cfg_input)
                    
                    # Steps
                    with dpg.group(horizontal=True):
                        Components.label(dpg.last_item(), "Steps:")
                        steps_input = Components.input_int(
                            dpg.last_item(),
                            default_value=20,
                            width=200
                        )
                        self.add_component("steps", steps_input)
                    
                    # Sample button
                    with dpg.group(horizontal=True):
                        sample_button = Components.button(
                            dpg.last_item(),
                            text="Generate Sample",
                            callback=self.on_sample_button,
                            width=200
                        )
                        self.add_component("sample_button", sample_button)
                        
                    # Image display
                    with dpg.group():
                        Components.label(
                            dpg.last_item(),
                            "This is a simplified sampling tab. In the full OneTrainer app,\n"
                            "this tab would display generated images here.",
                            wrap=600
                        )
    
    def on_sample_button(self):
        """Handle sample button click"""
        # Get current settings
        prompt = dpg.get_value(self.get_component("prompt"))
        negative_prompt = dpg.get_value(self.get_component("negative_prompt"))
        cfg_scale = dpg.get_value(self.get_component("cfg_scale"))
        steps = dpg.get_value(self.get_component("steps"))
        
        # In a real implementation, this would generate a sample image
        print(f"Would generate sample with:")
        print(f"  Prompt: {prompt}")
        print(f"  Negative Prompt: {negative_prompt}")
        print(f"  CFG Scale: {cfg_scale}")
        print(f"  Steps: {steps}")
        print(f"  Model Type: {self.state_manager.model_type}")


# Create window management
class OneTrainerWindow:
    """Main window class"""
    
    def __init__(self, title="OneTrainer", width=1280, height=800):
        """Initialize the window"""
        self.title = title
        self.width = width
        self.height = height
        self.state = SimpleStateManager()
        self.tabs = {}
        
        # Setup DPG
        dpg.create_context()
        dpg.create_viewport(title=title, width=width, height=height)
        self.apply_dark_theme()
        
        # Create window
        with dpg.window(label=title, tag="main_window", width=width, height=height):
            # Create header
            self.create_header()
            
            # Create tab bar
            self.tab_bar = dpg.add_tab_bar(tag="main_tab_bar")
            
            # Create tabs
            self.add_tab("model", "Model", ModelTab)
            self.add_tab("training", "Training", TrainingTab)
            self.add_tab("lora", "LoRA", LoraTab)
            self.add_tab("concept", "Concept", ConceptTab)
            self.add_tab("sampling", "Sampling", SamplingTab)
            
            # Show the first tab
            dpg.set_value(self.tab_bar, "model_tab")
        
        # Setup and show viewport
        dpg.setup_dearpygui()
        
    def create_header(self):
        """Create the header"""
        with dpg.group(horizontal=True):
            # Logo text
            dpg.add_text("OneTrainer", color=(255, 255, 255, 255))
            
            # Version
            dpg.add_text("v0.1.0", color=(180, 180, 180, 255))
            
    def add_tab(self, name, label, tab_class):
        """Add a tab to the window"""
        # Create the tab
        tab = dpg.add_tab(label=label, tag=f"{name}_tab", parent=self.tab_bar)
        
        # Create the tab content
        tab_instance = tab_class(tab, self.state)
        self.tabs[name] = tab_instance
        
        return tab
        
    def apply_dark_theme(self):
        """Apply a dark theme to the UI"""
        with dpg.theme() as global_theme:
            with dpg.theme_component(dpg.mvAll):
                # Background and text colors
                dpg.add_theme_color(dpg.mvThemeCol_WindowBg, (30, 30, 30, 255))
                dpg.add_theme_color(dpg.mvThemeCol_Text, (240, 240, 240, 255))
                
                # Button colors
                dpg.add_theme_color(dpg.mvThemeCol_Button, (41, 83, 154, 255))
                dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, (61, 103, 174, 255))
                dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, (21, 63, 134, 255))
                
                # Header colors
                dpg.add_theme_color(dpg.mvThemeCol_Header, (41, 83, 154, 255))
                dpg.add_theme_color(dpg.mvThemeCol_HeaderHovered, (61, 103, 174, 255))
                dpg.add_theme_color(dpg.mvThemeCol_HeaderActive, (21, 63, 134, 255))
                
                # Tab colors
                dpg.add_theme_color(dpg.mvThemeCol_Tab, (41, 83, 154, 255))
                dpg.add_theme_color(dpg.mvThemeCol_TabHovered, (61, 103, 174, 255))
                dpg.add_theme_color(dpg.mvThemeCol_TabActive, (21, 63, 134, 255))
                
                # Frame colors
                dpg.add_theme_color(dpg.mvThemeCol_FrameBg, (45, 45, 45, 255))
                dpg.add_theme_color(dpg.mvThemeCol_FrameBgHovered, (55, 55, 55, 255))
                dpg.add_theme_color(dpg.mvThemeCol_FrameBgActive, (65, 65, 65, 255))
                
                # Add spacing (if available in this DPG version)
                try:
                    dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 6, 6)
                    dpg.add_theme_style(dpg.mvStyleVar_ItemSpacing, 6, 6)
                    dpg.add_theme_style(dpg.mvStyleVar_WindowPadding, 10, 10)
                except (AttributeError, RuntimeError):
                    # Ignore styling if not available in this version
                    pass
                
                # Round corners (if available in this DPG version)
                try:
                    dpg.add_theme_style(dpg.mvStyleVar_FrameRounding, 4)
                    dpg.add_theme_style(dpg.mvStyleVar_TabRounding, 4)
                except (AttributeError, RuntimeError):
                    # Ignore styling if not available in this version
                    pass
        
        # Apply the theme
        dpg.bind_theme(global_theme)
        
    def show(self):
        """Show the window and start the main loop"""
        dpg.show_viewport()
        dpg.start_dearpygui()
        dpg.destroy_context()


# Main application
def main():
    """Main entry point"""
    print("Starting OneTrainer standalone app...")
    
    # Create and show window
    window = OneTrainerWindow()
    window.show()
    
    print("OneTrainer app closed")

if __name__ == "__main__":
    main()