standalone_lora_test.py

#!/usr/bin/env python3
"""
Standalone LoRA tab test with LyCORIS support.
Uses a simpler DPG approach that's compatible with most versions.
"""

import os
import sys
from enum import Enum
from dataclasses import dataclass

try:
    import dearpygui.dearpygui as dpg
except ImportError:
    print("DearPyGui not found. Please install it with: pip install dearpygui")
    sys.exit(1)

# Define minimal enum types
class ModelType(Enum):
    SD = "sd"
    SDXL = "sdxl"
    SD3 = "sd3"
    WUERSTCHEN = "wuerstchen"
    FLUX = "flux"
    PIXART_ALPHA = "pixart_alpha"
    SANA = "sana"
    HUNYUAN_VIDEO = "hunyuan_video"
    HI_DREAM = "hi_dream"
    
    def is_hi_dream(self):
        return self == ModelType.HI_DREAM

class PeftType(Enum):
    LORA = "lora"
    LOHA = "loha"
    LOKR = "lokr"
    DIA = "dia"
    IA3 = "ia3"
    DYLORA = "dylora"

class DataType(Enum):
    FLOAT_32 = "float32"
    BFLOAT_16 = "bfloat16"
    FLOAT_16 = "float16"

@dataclass
class TrainConfig:
    """Minimal TrainConfig class for testing"""
    model_type: ModelType = ModelType.SD
    peft_type: PeftType = PeftType.LORA
    
    # LoRA settings
    lora_model_name: str = ""
    lora_rank: int = 32
    lora_alpha: float = 32.0
    lora_layer_preset: str = "default"
    lora_layers: str = ""
    lora_weight_dtype: DataType = DataType.FLOAT_32
    dropout_probability: float = 0.0
    
    # LoRA DoRA settings
    lora_decompose: bool = False
    lora_decompose_norm_epsilon: bool = True
    lora_decompose_output_axis: bool = False
    
    # LyCORIS settings
    lycoris_factor: int = 4
    lycoris_full_matrix: bool = False
    lycoris_bypass_mode: bool = False

# LoRA layer presets
LAYER_PRESETS = {
    "default": ["to_q", "to_k", "to_v", "to_out.0"],
    "attn-mlp": ["to_q", "to_k", "to_v", "to_out.0", "ff.net.0", "ff.net.2"],
    "full": [".*"],
}

# Constants for UI
WINDOW_WIDTH = 900
WINDOW_HEIGHT = 700
PRIMARY_COLOR = [0, 119, 200, 255]

class StandaloneLoraTest:
    """Test application for LoRA tab with LyCORIS support"""
    
    def __init__(self):
        self.train_config = TrainConfig()
        self.setup_ui()
    
    def setup_ui(self):
        """Set up the user interface"""
        dpg.create_context()
        
        # Create viewport
        dpg.create_viewport(title="OneTrainer - LoRA Test", width=WINDOW_WIDTH, height=WINDOW_HEIGHT)
        
        # Create window
        self.window = dpg.add_window(label="LoRA/LyCORIS Settings", width=WINDOW_WIDTH, height=WINDOW_HEIGHT)
        
        # Add UI elements
        self.add_title()
        self.add_model_type_selector()
        self.add_peft_type_selector()
        self.add_lora_settings()
        self.add_lycoris_settings()
        self.add_dora_settings()
        self.add_common_settings()
        self.add_layer_preset_settings()
        self.add_debug_buttons()
        
        # Show the UI
        dpg.setup_dearpygui()
        dpg.show_viewport()
    
    def add_title(self):
        """Add title to the UI"""
        dpg.add_text(parent=self.window, default_value="OneTrainer - LoRA/LyCORIS Settings")
        dpg.add_separator(parent=self.window)
    
    def add_model_type_selector(self):
        """Add model type selector"""
        # Create horizontal group for model type
        model_group = dpg.add_group(parent=self.window, horizontal=True)
        dpg.add_text(parent=model_group, default_value="Model Type:")
        
        # Create combo box
        model_types = [mt.name for mt in ModelType]
        self.model_type_combo = dpg.add_combo(
            parent=model_group,
            items=model_types,
            default_value=self.train_config.model_type.name,
            callback=self.on_model_type_changed
        )
    
    def add_peft_type_selector(self):
        """Add PEFT type selector"""
        # Create horizontal group for PEFT type
        peft_group = dpg.add_group(parent=self.window, horizontal=True)
        dpg.add_text(parent=peft_group, default_value="PEFT Type:")
        
        # Create combo box
        peft_types = [pt.name for pt in PeftType]
        self.peft_type_combo = dpg.add_combo(
            parent=peft_group,
            items=peft_types,
            default_value=self.train_config.peft_type.name,
            callback=self.on_peft_type_changed
        )
    
    def add_lora_settings(self):
        """Add general LoRA settings"""
        # Create groups
        lora_base_group = dpg.add_group(parent=self.window, horizontal=True)
        lora_rank_group = dpg.add_group(parent=self.window, horizontal=True)
        lora_alpha_group = dpg.add_group(parent=self.window, horizontal=True)
        
        # Add UI elements
        dpg.add_text(parent=lora_base_group, default_value="LoRA base model:")
        self.lora_model_input = dpg.add_input_text(
            parent=lora_base_group,
            default_value=self.train_config.lora_model_name,
            callback=lambda s, a: setattr(self.train_config, "lora_model_name", a)
        )
        
        dpg.add_text(parent=lora_rank_group, default_value="LoRA rank:")
        self.lora_rank_input = dpg.add_input_int(
            parent=lora_rank_group,
            default_value=self.train_config.lora_rank,
            callback=lambda s, a: setattr(self.train_config, "lora_rank", a)
        )
        
        dpg.add_text(parent=lora_alpha_group, default_value="LoRA alpha:")
        self.lora_alpha_input = dpg.add_input_float(
            parent=lora_alpha_group,
            default_value=self.train_config.lora_alpha,
            callback=lambda s, a: setattr(self.train_config, "lora_alpha", a)
        )
    
    def add_lycoris_settings(self):
        """Add LyCORIS-specific settings"""
        # Create parent group
        self.lycoris_group = dpg.add_group(parent=self.window)
        dpg.add_text(parent=self.lycoris_group, default_value="LyCORIS Settings")
        
        # Factor settings
        factor_group = dpg.add_group(parent=self.lycoris_group, horizontal=True)
        dpg.add_text(parent=factor_group, default_value="LoKr/LoHa Factor:")
        self.lycoris_factor_input = dpg.add_input_int(
            parent=factor_group,
            default_value=self.train_config.lycoris_factor,
            callback=lambda s, a: setattr(self.train_config, "lycoris_factor", a)
        )
        
        # Full matrix toggle
        full_matrix_group = dpg.add_group(parent=self.lycoris_group, horizontal=True)
        dpg.add_text(parent=full_matrix_group, default_value="Use Full Matrix (DyLoRA):")
        self.lycoris_full_matrix_check = dpg.add_checkbox(
            parent=full_matrix_group,
            default_value=self.train_config.lycoris_full_matrix,
            callback=lambda s, a: setattr(self.train_config, "lycoris_full_matrix", a)
        )
        
        # Bypass mode toggle
        bypass_group = dpg.add_group(parent=self.lycoris_group, horizontal=True)
        dpg.add_text(parent=bypass_group, default_value="Bypass Mode (for HiDream + LoKr):")
        self.lycoris_bypass_check = dpg.add_checkbox(
            parent=bypass_group,
            default_value=self.train_config.lycoris_bypass_mode,
            callback=lambda s, a: setattr(self.train_config, "lycoris_bypass_mode", a)
        )
    
    def add_dora_settings(self):
        """Add DoRA-specific settings"""
        # Create parent group
        self.dora_group = dpg.add_group(parent=self.window)
        dpg.add_text(parent=self.dora_group, default_value="DoRA Settings (LoRA only)")
        
        # Enable DoRA
        dora_enable_group = dpg.add_group(parent=self.dora_group, horizontal=True)
        dpg.add_text(parent=dora_enable_group, default_value="Enable Weight Decomposition:")
        self.lora_decompose_check = dpg.add_checkbox(
            parent=dora_enable_group,
            default_value=self.train_config.lora_decompose,
            callback=self.on_dora_toggle
        )
        
        # Advanced settings
        self.dora_advanced_group = dpg.add_group(parent=self.dora_group)
        
        # Norm epsilon
        norm_group = dpg.add_group(parent=self.dora_advanced_group, horizontal=True)
        dpg.add_text(parent=norm_group, default_value="Use Norm Epsilon:")
        self.norm_epsilon_check = dpg.add_checkbox(
            parent=norm_group,
            default_value=self.train_config.lora_decompose_norm_epsilon,
            callback=lambda s, a: setattr(self.train_config, "lora_decompose_norm_epsilon", a)
        )
        
        # Output axis
        output_axis_group = dpg.add_group(parent=self.dora_advanced_group, horizontal=True)
        dpg.add_text(parent=output_axis_group, default_value="Apply on Output Axis:")
        self.output_axis_check = dpg.add_checkbox(
            parent=output_axis_group,
            default_value=self.train_config.lora_decompose_output_axis,
            callback=lambda s, a: setattr(self.train_config, "lora_decompose_output_axis", a)
        )
    
    def add_common_settings(self):
        """Add common settings for all PEFT types"""
        # Dropout
        dropout_group = dpg.add_group(parent=self.window, horizontal=True)
        dpg.add_text(parent=dropout_group, default_value="Dropout Probability:")
        self.dropout_slider = dpg.add_slider_float(
            parent=dropout_group,
            default_value=self.train_config.dropout_probability,
            min_value=0.0,
            max_value=1.0,
            callback=lambda s, a: setattr(self.train_config, "dropout_probability", a)
        )
        
        # Data type
        dtype_group = dpg.add_group(parent=self.window, horizontal=True)
        dpg.add_text(parent=dtype_group, default_value="Weight Data Type:")
        dtype_options = [dt.name for dt in DataType]
        self.dtype_combo = dpg.add_combo(
            parent=dtype_group,
            items=dtype_options,
            default_value=self.train_config.lora_weight_dtype.name,
            callback=self.on_dtype_changed
        )
    
    def add_layer_preset_settings(self):
        """Add layer preset settings"""
        # Layer preset
        preset_group = dpg.add_group(parent=self.window, horizontal=True)
        dpg.add_text(parent=preset_group, default_value="Layer Preset:")
        
        preset_options = list(LAYER_PRESETS.keys()) + ["custom"]
        self.preset_combo = dpg.add_combo(
            parent=preset_group,
            items=preset_options,
            default_value=self.train_config.lora_layer_preset,
            callback=self.on_preset_changed
        )
        
        # Custom layers
        layers_group = dpg.add_group(parent=self.window, horizontal=True)
        dpg.add_text(parent=layers_group, default_value="Custom Layers:")
        self.layers_input = dpg.add_input_text(
            parent=layers_group,
            default_value=self.train_config.lora_layers,
            callback=lambda s, a: setattr(self.train_config, "lora_layers", a)
        )
    
    def add_debug_buttons(self):
        """Add debug buttons"""
        debug_group = dpg.add_group(parent=self.window, horizontal=True)
        dpg.add_button(
            parent=debug_group,
            label="Print Config",
            callback=self.print_config
        )
        dpg.add_button(
            parent=debug_group,
            label="Reset",
            callback=self.reset_config
        )
    
    def on_model_type_changed(self, sender, app_data):
        """Handle model type change"""
        try:
            self.train_config.model_type = ModelType[app_data]
            self.update_ui_visibility()
        except Exception as e:
            print(f"Error changing model type: {e}")
    
    def on_peft_type_changed(self, sender, app_data):
        """Handle PEFT type change"""
        try:
            self.train_config.peft_type = PeftType[app_data]
            self.update_ui_visibility()
        except Exception as e:
            print(f"Error changing PEFT type: {e}")
    
    def on_dtype_changed(self, sender, app_data):
        """Handle data type change"""
        try:
            self.train_config.lora_weight_dtype = DataType[app_data]
        except Exception as e:
            print(f"Error changing data type: {e}")
    
    def on_dora_toggle(self, sender, app_data):
        """Handle DoRA toggle"""
        self.train_config.lora_decompose = app_data
        self.update_ui_visibility()
    
    def on_preset_changed(self, sender, app_data):
        """Handle layer preset change"""
        self.train_config.lora_layer_preset = app_data
        
        # If not custom, set the layers based on the preset
        if app_data != "custom" and app_data in LAYER_PRESETS:
            layers = ", ".join(LAYER_PRESETS[app_data])
            self.train_config.lora_layers = layers
            dpg.set_value(self.layers_input, layers)
    
    def update_ui_visibility(self):
        """Update UI element visibility based on current config"""
        # Show/hide LyCORIS settings
        show_lycoris = self.train_config.peft_type in [PeftType.LOHA, PeftType.LOKR, PeftType.DIA, PeftType.IA3, PeftType.DYLORA]
        dpg.configure_item(self.lycoris_group, show=show_lycoris)
        
        # Show/hide DoRA settings
        show_dora = self.train_config.peft_type == PeftType.LORA
        dpg.configure_item(self.dora_group, show=show_dora)
        
        # Show/hide DoRA advanced settings
        show_dora_advanced = show_dora and self.train_config.lora_decompose
        dpg.configure_item(self.dora_advanced_group, show=show_dora_advanced)
    
    def print_config(self):
        """Print current config to console"""
        print("\n--- Current Config ---")
        print(f"Model Type: {self.train_config.model_type}")
        print(f"PEFT Type: {self.train_config.peft_type}")
        print(f"LoRA Rank: {self.train_config.lora_rank}")
        print(f"LoRA Alpha: {self.train_config.lora_alpha}")
        print(f"Dropout: {self.train_config.dropout_probability}")
        print(f"Layer Preset: {self.train_config.lora_layer_preset}")
        print(f"Layers: {self.train_config.lora_layers}")
        
        if self.train_config.peft_type == PeftType.LORA:
            print(f"DoRA Enabled: {self.train_config.lora_decompose}")
            if self.train_config.lora_decompose:
                print(f"DoRA Norm Epsilon: {self.train_config.lora_decompose_norm_epsilon}")
                print(f"DoRA Output Axis: {self.train_config.lora_decompose_output_axis}")
        
        if self.train_config.peft_type in [PeftType.LOHA, PeftType.LOKR, PeftType.DIA, PeftType.IA3, PeftType.DYLORA]:
            print(f"LyCORIS Factor: {self.train_config.lycoris_factor}")
            print(f"LyCORIS Full Matrix: {self.train_config.lycoris_full_matrix}")
            print(f"LyCORIS Bypass Mode: {self.train_config.lycoris_bypass_mode}")
    
    def reset_config(self):
        """Reset config to defaults"""
        self.train_config = TrainConfig()
        
        # Update UI
        dpg.set_value(self.model_type_combo, self.train_config.model_type.name)
        dpg.set_value(self.peft_type_combo, self.train_config.peft_type.name)
        dpg.set_value(self.lora_model_input, self.train_config.lora_model_name)
        dpg.set_value(self.lora_rank_input, self.train_config.lora_rank)
        dpg.set_value(self.lora_alpha_input, self.train_config.lora_alpha)
        dpg.set_value(self.dropout_slider, self.train_config.dropout_probability)
        dpg.set_value(self.dtype_combo, self.train_config.lora_weight_dtype.name)
        dpg.set_value(self.preset_combo, self.train_config.lora_layer_preset)
        dpg.set_value(self.layers_input, self.train_config.lora_layers)
        dpg.set_value(self.lora_decompose_check, self.train_config.lora_decompose)
        dpg.set_value(self.norm_epsilon_check, self.train_config.lora_decompose_norm_epsilon)
        dpg.set_value(self.output_axis_check, self.train_config.lora_decompose_output_axis)
        dpg.set_value(self.lycoris_factor_input, self.train_config.lycoris_factor)
        dpg.set_value(self.lycoris_full_matrix_check, self.train_config.lycoris_full_matrix)
        dpg.set_value(self.lycoris_bypass_check, self.train_config.lycoris_bypass_mode)
        
        # Update visibility
        self.update_ui_visibility()
    
    def run(self):
        """Run the application"""
        # Initial UI visibility update
        self.update_ui_visibility()
        
        # Start main loop
        while dpg.is_dearpygui_running():
            dpg.render_dearpygui_frame()
        
        # Cleanup
        dpg.destroy_context()

if __name__ == "__main__":
    print("Starting OneTrainer LoRA/LyCORIS test...")
    app = StandaloneLoraTest()
    app.run()