simple_dpg_ui.py

#!/usr/bin/env python3
"""
Simple Dear PyGui UI for OneTrainer with LyCORIS support

This is a simplified version that avoids potential issues with more complex DPG features.
"""

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

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

# Import OneTrainer modules
from modules.util.enum.ModelType import ModelType, PeftType
from modules.util.config.TrainConfig import TrainConfig

class SimpleDpgUI:
    """Simplified DPG UI for OneTrainer with LyCORIS support"""
    
    def __init__(self):
        """Initialize the UI"""
        self.components = {}
        self.config = TrainConfig.default_values()
        self.model_type = ModelType.STABLE_DIFFUSION_15
        self.peft_type = PeftType.LORA
    
    def run(self):
        """Run the UI"""
        # Create context
        dpg.create_context()
        
        # Create viewport
        dpg.create_viewport(title="OneTrainer (Simple DPG)", width=1000, height=800)
        
        # Create main window
        with dpg.window(tag="main_window", label="OneTrainer with LyCORIS support"):
            with dpg.tab_bar(tag="tab_bar"):
                # Create LoRA tab
                with dpg.tab(label="LoRA / LyCORIS", tag="lora_tab"):
                    self.create_lora_tab()
                
                # Create Model tab
                with dpg.tab(label="Model", tag="model_tab"):
                    self.create_model_tab()
                
                # Create Training tab
                with dpg.tab(label="Training", tag="training_tab"):
                    self.create_training_tab()
            
            # Add status bar
            with dpg.group(horizontal=True):
                self.status_text = dpg.add_text("Ready", tag="status_text")
        
        # Setup and start DPG
        dpg.setup_dearpygui()
        dpg.show_viewport()
        dpg.set_primary_window("main_window", True)
        dpg.start_dearpygui()
        dpg.destroy_context()
    
    def create_model_tab(self):
        """Create the model tab"""
        with dpg.group():
            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)
                
                # Get model types
                model_types = [m.name for m in ModelType]
                
                # Create selector
                model_type_combo = dpg.add_combo(
                    items=model_types,
                    default_value=self.model_type.name,
                    callback=self.on_model_type_change,
                    width=300
                )
                self.add_component("model_type", model_type_combo)
            
            # Model path
            with dpg.group(horizontal=True):
                dpg.add_text("Model Path:", width=150)
                
                # Create path field
                model_path = dpg.add_input_text(
                    default_value=self.config.model_path or "",
                    width=300
                )
                self.add_component("model_path", model_path)
                
                # Browse button
                dpg.add_button(
                    label="...",
                    callback=lambda: self.set_status("Model path selection would appear here")
                )
    
    def create_lora_tab(self):
        """Create the LoRA tab with LyCORIS support"""
        with dpg.group():
            dpg.add_text("LoRA / LyCORIS 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")
                
                # Define PEFT types
                peft_types = [
                    ("LoRA", PeftType.LORA),
                    ("LoHa", PeftType.LOHA),
                    ("LoKr", PeftType.LOKR),
                    ("DiA", PeftType.DIA),
                    ("iA3", PeftType.IA3),
                    ("DyLoRA", PeftType.DYLORA)
                ]
                
                # Extract labels
                peft_labels = [p[0] for p in peft_types]
                
                # Create selector
                peft_combo = dpg.add_combo(
                    items=peft_labels,
                    default_value=peft_labels[0],
                    callback=self.on_peft_type_change,
                    width=300
                )
                self.add_component("peft_type", peft_combo)
            
            # LoRA rank
            with dpg.group(horizontal=True, tag="rank_group"):
                dpg.add_text("LoRA Rank:", width=150)
                rank_input = dpg.add_input_int(
                    default_value=32,
                    width=200
                )
                self.add_component("lora_rank", rank_input)
            
            # LoRA alpha
            with dpg.group(horizontal=True, tag="alpha_group"):
                dpg.add_text("LoRA Alpha:", width=150)
                alpha_input = dpg.add_input_float(
                    default_value=32.0,
                    width=200
                )
                self.add_component("lora_alpha", alpha_input)
            
            # LoKr factor
            with dpg.group(horizontal=True, tag="lokr_group", show=False):
                dpg.add_text("LoKr Factor:", width=150)
                factor_input = dpg.add_input_int(
                    default_value=4,
                    width=200
                )
                self.add_component("lokr_factor", factor_input)
            
            # Warning message for HiDream
            with dpg.group(tag="warning_group", show=False):
                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.",
                    color=(255, 0, 0),
                    wrap=580
                )
    
    def create_training_tab(self):
        """Create the training tab"""
        with dpg.group():
            dpg.add_text("Training Settings", color=(255, 255, 0))
            dpg.add_separator()
            
            # Epochs
            with dpg.group(horizontal=True):
                dpg.add_text("Epochs:", width=150)
                epochs_input = dpg.add_input_int(
                    default_value=self.config.epochs,
                    width=200
                )
                self.add_component("epochs", epochs_input)
            
            # Batch size
            with dpg.group(horizontal=True):
                dpg.add_text("Batch Size:", width=150)
                batch_size_input = dpg.add_input_int(
                    default_value=self.config.batch_size,
                    width=200
                )
                self.add_component("batch_size", batch_size_input)
            
            # Learning rate
            with dpg.group(horizontal=True):
                dpg.add_text("Learning Rate:", width=150)
                learning_rate_input = dpg.add_input_float(
                    default_value=self.config.learning_rate,
                    format="%.6f",
                    width=200
                )
                self.add_component("learning_rate", learning_rate_input)
            
            # Start training button
            dpg.add_button(
                label="Start Training",
                callback=self.on_start_training
            )
    
    def on_model_type_change(self, sender, app_data):
        """Handle model type change"""
        try:
            self.model_type = ModelType(app_data)
            self.update_warnings()
        except:
            pass
    
    def on_peft_type_change(self, sender, app_data):
        """Handle PEFT type change"""
        # Map from display name to enum
        peft_map = {
            "LoRA": PeftType.LORA,
            "LoHa": PeftType.LOHA,
            "LoKr": PeftType.LOKR,
            "DiA": PeftType.DIA,
            "iA3": PeftType.IA3,
            "DyLoRA": PeftType.DYLORA
        }
        
        # Update PEFT type
        self.peft_type = peft_map.get(app_data, PeftType.LORA)
        
        # Show/hide LoKr factor
        dpg.configure_item("lokr_group", show=(self.peft_type == PeftType.LOKR))
        
        # Update warnings
        self.update_warnings()
    
    def update_warnings(self):
        """Update warnings based on model type and PEFT type"""
        # Show warning for HiDream + LoKr combination
        show_warning = self.model_type == ModelType.HI_DREAM_FULL and self.peft_type == PeftType.LOKR
        dpg.configure_item("warning_group", show=show_warning)
    
    def on_start_training(self):
        """Handle start training button"""
        # Get values
        epochs = dpg.get_value(self.get_component("epochs"))
        batch_size = dpg.get_value(self.get_component("batch_size"))
        
        # Update status
        self.set_status(f"Would start training with {epochs} epochs, batch size {batch_size}")
    
    def add_component(self, name: str, component_id: int):
        """Add a component to the registry"""
        self.components[name] = component_id
    
    def get_component(self, name: str) -> int:
        """Get a component by name"""
        return self.components.get(name)
    
    def set_status(self, text: str):
        """Set the status text"""
        dpg.set_value(self.status_text, text)

if __name__ == "__main__":
    try:
        print("Starting Simple OneTrainer DPG UI...")
        ui = SimpleDpgUI()
        ui.run()
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()