test_dpg_lora_tab.py

#!/usr/bin/env python3
"""
Test script to run just the LoRA tab with LyCORIS support.
This is useful for testing the LoRA/LyCORIS functionality independently.
"""

import os
import sys
from enum import Enum
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

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

# Support loading modules
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# Define minimal enum types if OneTrainer modules are not available
try:
    from modules.util.enum.ModelType import ModelType, PeftType
    from modules.util.enum.DataType import DataType
    MOCK_ENUMS = False
except ImportError:
    print("OneTrainer modules not found, using mock enums")
    MOCK_ENUMS = True
    
    # Create mock enums for testing
    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_stable_diffusion(self):
            return self == ModelType.SD
            
        def is_stable_diffusion_xl(self):
            return self == ModelType.SDXL
            
        def is_stable_diffusion_3(self):
            return self == ModelType.SD3
            
        def is_wuerstchen(self):
            return self == ModelType.WUERSTCHEN
            
        def is_pixart(self):
            return self == ModelType.PIXART_ALPHA
            
        def is_flux(self):
            return self == ModelType.FLUX
            
        def is_sana(self):
            return self == ModelType.SANA
            
        def is_hunyuan_video(self):
            return self == ModelType.HUNYUAN_VIDEO
            
        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"
        FLOAT_8 = "float8"
        NFLOAT_4 = "nf4"
        NONE = "none"

# Define minimal config for testing
@dataclass
class TrainConfig:
    """Minimal TrainConfig class for testing"""
    model_type: ModelType = list(ModelType)[0]  # Get the first enum value
    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
    bundle_additional_embeddings: bool = False
    
    # 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 for each model type
DEFAULT_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"],
    "attn-only": ["to_q", "to_k", "to_v", "to_out.0"],
    "sd3-full": ["to_qkv", "to_out.0", "ff.0", "ff.2"],
    "all": [".*"],
}

HIDREAM_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"],
    "extended": ["to_q", "to_k", "to_v", "to_out.0", "ff.net.0", "ff.net.2", "proj_in", "proj_out"],
    "conv-punc": ["conv1", "conv2", "time_emb_proj"],
    "full": [".*"],
}

# Constants for UI
WINDOW_WIDTH = 900
WINDOW_HEIGHT = 700
PRIMARY_COLOR = [0, 119, 200, 255]
SECONDARY_COLOR = [0, 51, 102, 255]
ACCENT_COLOR = [0, 178, 255, 255]

class LoraTabTest:
    """Test application for the LoRA tab with LyCORIS support"""
    
    def __init__(self):
        # Initialize DPG
        dpg.create_context()
        dpg.create_viewport(title="OneTrainer - LoRA Tab Test", width=WINDOW_WIDTH, height=WINDOW_HEIGHT)
        
        # Initialize config
        self.train_config = TrainConfig()
        
        # Set up UI
        self.setup_themes()
        self.create_main_window()
        
        # Show viewport
        dpg.setup_dearpygui()
        dpg.show_viewport()
    
    def setup_themes(self):
        """Set up custom themes for the application"""
        # Main theme
        with dpg.theme(tag="main_theme"):
            with dpg.theme_component(dpg.mvAll):
                dpg.add_theme_color(dpg.mvThemeCol_WindowBg, [20, 20, 20, 255])
                dpg.add_theme_color(dpg.mvThemeCol_TitleBgActive, PRIMARY_COLOR)
                dpg.add_theme_color(dpg.mvThemeCol_Button, PRIMARY_COLOR)
                dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, ACCENT_COLOR)
                dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, SECONDARY_COLOR)
                dpg.add_theme_color(dpg.mvThemeCol_Text, [240, 240, 240, 255])
                dpg.add_theme_color(dpg.mvThemeCol_FrameBg, [45, 45, 45, 255])
        
        # Button theme
        with dpg.theme(tag="button_theme"):
            with dpg.theme_component(dpg.mvButton):
                dpg.add_theme_color(dpg.mvThemeCol_Button, PRIMARY_COLOR)
                dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, ACCENT_COLOR)
                dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, SECONDARY_COLOR)
                dpg.add_theme_style(dpg.mvStyleVar_FrameRounding, 5)
                dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 6, 6)
    
    def create_main_window(self):
        """Create the main window with LoRA tab"""
        # Create main window
        with dpg.window(tag="main_window", label="OneTrainer - LoRA Tab Test",
                        no_collapse=True, width=WINDOW_WIDTH, height=WINDOW_HEIGHT):
            
            # Apply theme
            dpg.bind_item_theme("main_window", "main_theme")
            
            # Model type selector
            with dpg.group(horizontal=True):
                dpg.add_text("Model Type:")
                model_types = [(mt.name, mt) for mt in ModelType]
                self.add_combo("model_type_combo", model_types, 
                              callback=self.on_model_type_changed,
                              default_value=self.train_config.model_type.name)
            
            # Create LoRA tab
            self.create_lora_tab()
            
            # Add test buttons
            with dpg.group(horizontal=True):
                dpg.add_button(label="Print Config", callback=self.print_config)
                dpg.add_button(label="Reset", callback=self.reset_config)
    
    def create_lora_tab(self):
        """Create LoRA configuration tab with LyCORIS support"""
        with dpg.group(horizontal=False, tag="lora_group"):
            # Title
            dpg.add_text("LoRA / LyCORIS Settings", color=[255, 255, 255, 255])
            dpg.add_separator()
            
            # PEFT type selector with tooltip
            with dpg.group(horizontal=True):
                dpg.add_text("PEFT Type:", tag="peft_type_label")
                self.add_combo("peft_type", 
                              [(pt.name, pt) for pt in PeftType],
                              callback=self.on_peft_type_changed,
                              default_value=self.train_config.peft_type.name)
                
                # Add info tooltip
                with dpg.tooltip(dpg.last_item()):
                    dpg.add_text(
                        "LoRA: Low-Rank Adaptation (standard)\n"
                        "LoHa: Low-Rank Hadamard Product\n"
                        "LoKr: Low-Rank Kronecker Product\n"
                        "DyLoRA: Dynamic Low-Rank Adaptation\n"
                        "DiA: Dynamic Adapters\n"
                        "iA3: Infused Adapter by Inhibiting and Amplifying"
                    )
            
            # Base model path
            with dpg.group(horizontal=True):
                dpg.add_text("LoRA base model:", tag="lora_model_name_label")
                dpg.add_input_text(tag="lora_model_name", default_value=self.train_config.lora_model_name, 
                                 width=-100, callback=lambda s, v: self.set_config("lora_model_name", v))
                dpg.add_button(label="...", callback=lambda: self.file_dialog_callback("lora_model_name"))
            
            # Common settings that apply to most PEFT types
            with dpg.group(horizontal=True):
                dpg.add_text("LoRA rank:", tag="lora_rank_label")
                dpg.add_input_int(tag="lora_rank", default_value=self.train_config.lora_rank,
                                callback=lambda s, v: self.set_config("lora_rank", v), width=100)
                
                # Add info tooltip
                with dpg.tooltip(dpg.last_item()):
                    dpg.add_text(
                        "The rank of the update matrices, expressed as an integer. "
                        "Higher values = larger file size and VRAM usage, but potentially higher quality."
                    )
            
            with dpg.group(horizontal=True):
                dpg.add_text("LoRA alpha:", tag="lora_alpha_label")
                dpg.add_input_float(tag="lora_alpha", default_value=self.train_config.lora_alpha,
                                  callback=lambda s, v: self.set_config("lora_alpha", v), width=100)
                
                # Add info tooltip
                with dpg.tooltip(dpg.last_item()):
                    dpg.add_text(
                        "Scaling factor for the LoRA adapter. Higher values = stronger effect."
                        "Typically set to same value as rank. Commonly used values include 16, 32, 64."
                    )
            
            # LyCORIS-specific settings container (will be shown/hidden based on PEFT type)
            self.lycoris_settings_group = dpg.add_group(tag="lycoris_settings_group", horizontal=False)
            with self.lycoris_settings_group:
                # LoKr factor
                with dpg.group(horizontal=True):
                    dpg.add_text("LoKr Factor:", tag="lycoris_factor_label")
                    dpg.add_input_int(tag="lycoris_factor", default_value=self.train_config.lycoris_factor,
                                   callback=lambda s, v: self.set_config("lycoris_factor", v), width=100)
                    
                    # Add info tooltip
                    with dpg.tooltip(dpg.last_item()):
                        dpg.add_text(
                            "Compression factor for LoKr/LoHa. Lower = less parameters.\n"
                            "Used in conjunction with rank. Total parameters = (width × height) ÷ (factor × rank)"
                        )
                
                # DyLoRA settings
                with dpg.group(horizontal=True):
                    dpg.add_text("Use Full Matrix:", tag="lycoris_full_matrix_label")
                    dpg.add_checkbox(label="Enabled", tag="lycoris_full_matrix", 
                                    default_value=self.train_config.lycoris_full_matrix,
                                    callback=lambda s, v: self.set_config("lycoris_full_matrix", v))
                    
                    # Add info tooltip
                    with dpg.tooltip(dpg.last_item()):
                        dpg.add_text(
                            "DyLoRA: Use full weight matrix (much larger model size but better quality).\n"
                            "This increases the parameter count significantly."
                        )
                
                # Bypass mode (for advanced use)
                with dpg.group(horizontal=True):
                    dpg.add_text("Bypass Mode:", tag="lycoris_bypass_mode_label")
                    dpg.add_checkbox(label="Enabled", tag="lycoris_bypass_mode", 
                                    default_value=self.train_config.lycoris_bypass_mode,
                                    callback=lambda s, v: self.set_config("lycoris_bypass_mode", v))
                    
                    # Add info tooltip
                    with dpg.tooltip(dpg.last_item()):
                        dpg.add_text(
                            "Advanced: Bypass problematic LyCORIS types (like LoKr) for HiDream models.\n"
                            "Only enable if you're having compatibility issues with specific model types."
                        )
            
            # LoRA-specific settings container
            self.lora_settings_group = dpg.add_group(tag="lora_settings_group", horizontal=False)
            with self.lora_settings_group:
                # DoRA (weight decomposition)
                with dpg.group(horizontal=True):
                    dpg.add_text("Weight Decomposition (DoRA):", tag="lora_decompose_label")
                    dpg.add_checkbox(label="Enabled", tag="lora_decompose", 
                                    default_value=self.train_config.lora_decompose,
                                    callback=lambda s, v: self.set_config("lora_decompose", v))
                    
                    # Add info tooltip
                    with dpg.tooltip(dpg.last_item()):
                        dpg.add_text(
                            "Enable weight decomposition (DoRA - Decomposed Rank Adaptation).\n"
                            "This can significantly improve quality but increases training time."
                        )
                
                # DoRA settings
                with dpg.group(horizontal=True, tag="dora_settings"):
                    dpg.add_text("DoRA Settings:")
                    dpg.add_checkbox(label="Use Norm Epsilon", tag="lora_decompose_norm_epsilon", 
                                    default_value=self.train_config.lora_decompose_norm_epsilon,
                                    callback=lambda s, v: self.set_config("lora_decompose_norm_epsilon", v))
                    
                    dpg.add_checkbox(label="Apply on Output Axis", tag="lora_decompose_output_axis", 
                                    default_value=self.train_config.lora_decompose_output_axis,
                                    callback=lambda s, v: self.set_config("lora_decompose_output_axis", v))
            
            # Common settings for all PEFT types
            with dpg.group(horizontal=True):
                dpg.add_text("Dropout Probability:", tag="dropout_probability_label")
                dpg.add_slider_float(tag="dropout_probability", default_value=self.train_config.dropout_probability,
                                   callback=lambda s, v: self.set_config("dropout_probability", v), 
                                   min_value=0.0, max_value=1.0, width=200)
                
                # Add info tooltip
                with dpg.tooltip(dpg.last_item()):
                    dpg.add_text(
                        "Dropout probability for LoRA layers (0.0 = disabled).\n"
                        "Higher values can help prevent overfitting but may require longer training."
                    )
            
            with dpg.group(horizontal=True):
                dpg.add_text("Weight Data Type:", tag="lora_weight_dtype_label")
                self.add_combo("lora_weight_dtype", 
                              [(dt.name, dt) for dt in [DataType.FLOAT_32, DataType.BFLOAT_16]],
                              callback=lambda s, v: self.set_config("lora_weight_dtype", DataType[v]),
                              default_value=self.train_config.lora_weight_dtype.name)
            
            with dpg.group(horizontal=True):
                dpg.add_checkbox(label="Bundle Embeddings", tag="bundle_additional_embeddings", 
                                default_value=self.train_config.bundle_additional_embeddings,
                                callback=lambda s, v: self.set_config("bundle_additional_embeddings", v))
                
                # Add info tooltip
                with dpg.tooltip(dpg.last_item()):
                    dpg.add_text(
                        "Include trained textual embeddings within the LoRA file."
                    )
            
            # Layer presets
            with dpg.group(horizontal=True):
                dpg.add_text("Layer Preset:", tag="lora_layer_preset_label")
                # Get the right preset list based on model type
                self.presets = self.get_lora_presets()
                preset_list = list(self.presets.keys()) if self.presets else []
                preset_list.append("custom")
                self.add_combo("lora_layer_preset", [(p, p) for p in preset_list],
                              callback=self.on_layer_preset_changed,
                              default_value=self.train_config.lora_layer_preset)
                
                # Add info tooltip
                with dpg.tooltip(dpg.last_item()):
                    dpg.add_text(
                        "Presets for which layers to target with the LoRA adapter.\n"
                        "Different models have different recommended layer targeting strategies."
                    )
            
            with dpg.group(horizontal=True):
                dpg.add_text("Custom Layers:", tag="lora_layers_label")
                dpg.add_input_text(tag="lora_layers", default_value=self.train_config.lora_layers,
                                 callback=lambda s, v: self.set_config("lora_layers", v), width=-100)
                
                # Add info tooltip
                with dpg.tooltip(dpg.last_item()):
                    dpg.add_text(
                        "Comma-separated list of layers to apply LoRA to.\n"
                        "Only used when 'custom' preset is selected."
                    )
            
            # Update visibility of PEFT-specific settings based on current type
            self.update_peft_settings_visibility()
    
    # Helper methods for UI components
    
    def add_combo(self, tag, items, callback=None, default_value=None, width=-1):
        """Helper to add a combo box with label/value pairs"""
        labels = [item[0] for item in items]
        values = [item[1] for item in items]
        
        if default_value is not None and isinstance(default_value, str):
            # Find index of the default value
            default_idx = 0
            for i, label in enumerate(labels):
                if label == default_value:
                    default_idx = i
                    break
        else:
            default_idx = 0
        
        def combo_callback(sender, app_data):
            # Call the provided callback with the selected value
            if callback:
                callback(sender, labels[app_data] if isinstance(app_data, int) else app_data)
        
        return dpg.add_combo(
            items=labels,
            default_value=labels[default_idx] if labels else "",
            tag=tag,
            width=width,
            callback=combo_callback
        )
    
    def file_dialog_callback(self, target_tag):
        """Show a file dialog and update the target input field"""
        # For now, we'll just print a message since file dialogs are complex in DPG
        print(f"Would show file dialog for {target_tag}")
    
    def set_config(self, key, value):
        """Set a config value and update UI if needed"""
        # We're using a simplified approach without dot-notation
        if hasattr(self.train_config, key):
            setattr(self.train_config, key, value)
            
            # Update UI if needed based on the changed value
            if key == "model_type":
                self.update_lora_presets()
                self.update_peft_settings_visibility()
            elif key == "peft_type":
                self.update_peft_settings_visibility()
            elif key == "lora_decompose":
                self.update_dora_settings_visibility()
    
    def get_lora_presets(self):
        """Get appropriate LORA presets based on model type"""
        model_type = self.train_config.model_type
        
        # Return the appropriate presets based on model type
        if model_type == ModelType.HI_DREAM:
            return HIDREAM_PRESETS
        else:
            return DEFAULT_PRESETS
    
    def update_lora_presets(self):
        """Update LoRA layer presets based on model type"""
        if not dpg.does_item_exist("lora_layer_preset"):
            return
        
        # Get the right preset list based on model type
        self.presets = self.get_lora_presets()
        preset_list = list(self.presets.keys()) if self.presets else []
        preset_list.append("custom")
        
        # Update combo box
        dpg.configure_item("lora_layer_preset", items=preset_list)
        
        # Set default value
        default_preset = self.train_config.lora_layer_preset or preset_list[0]
        dpg.set_value("lora_layer_preset", default_preset)
    
    def on_peft_type_changed(self, sender, value):
        """Handle PEFT type selection change"""
        try:
            peft_type = PeftType[value]
            self.set_config("peft_type", peft_type)
            self.update_peft_settings_visibility()
        except Exception as e:
            print(f"Error setting PEFT type: {e}")
    
    def update_peft_settings_visibility(self):
        """Update visibility of PEFT-specific settings based on current type"""
        peft_type = self.train_config.peft_type
        
        # Show/hide LyCORIS-specific settings
        if dpg.does_item_exist("lycoris_settings_group"):
            show_lycoris = peft_type in [PeftType.LOHA, PeftType.LOKR, PeftType.DIA, PeftType.IA3, PeftType.DYLORA]
            dpg.configure_item("lycoris_settings_group", show=show_lycoris)
            
            # For LoKr/LoHa, always show factor
            if dpg.does_item_exist("lycoris_factor_label"):
                show_factor = peft_type in [PeftType.LOHA, PeftType.LOKR]
                dpg.configure_item("lycoris_factor_label", show=show_factor)
                dpg.configure_item("lycoris_factor", show=show_factor)
            
            # For DyLoRA, show full matrix option
            if dpg.does_item_exist("lycoris_full_matrix_label"):
                show_full_matrix = peft_type == PeftType.DYLORA
                dpg.configure_item("lycoris_full_matrix_label", show=show_full_matrix)
                dpg.configure_item("lycoris_full_matrix", show=show_full_matrix)
            
            # Only show bypass mode for problematic types with HiDream
            if dpg.does_item_exist("lycoris_bypass_mode_label"):
                problematic_with_hidream = (
                    self.train_config.model_type == ModelType.HI_DREAM and 
                    peft_type in [PeftType.LOKR]
                )
                dpg.configure_item("lycoris_bypass_mode_label", show=problematic_with_hidream)
                dpg.configure_item("lycoris_bypass_mode", show=problematic_with_hidream)
        
        # Show/hide LoRA-specific settings
        if dpg.does_item_exist("lora_settings_group"):
            show_lora = peft_type == PeftType.LORA
            dpg.configure_item("lora_settings_group", show=show_lora)
            self.update_dora_settings_visibility()
    
    def update_dora_settings_visibility(self):
        """Update visibility of DoRA-specific settings"""
        if not dpg.does_item_exist("dora_settings"):
            return
            
        # Only show DoRA settings if DoRA is enabled and PEFT type is LoRA
        show_dora = (
            self.train_config.peft_type == PeftType.LORA and
            self.train_config.lora_decompose
        )
        dpg.configure_item("dora_settings", show=show_dora)
    
    def on_layer_preset_changed(self, sender, value):
        """Handle layer preset selection change"""
        self.set_config("lora_layer_preset", value)
        
        # Update layers based on selected preset if not custom
        if value != "custom" and value in self.presets:
            layers = ", ".join(self.presets[value])
            self.set_config("lora_layers", layers)
            if dpg.does_item_exist("lora_layers"):
                dpg.set_value("lora_layers", layers)
    
    def on_model_type_changed(self, sender, value):
        """Handle model type selection change"""
        try:
            model_type = ModelType[value]
            self.set_config("model_type", model_type)
            self.update_lora_presets()
            self.update_peft_settings_visibility()
        except Exception as e:
            print(f"Error setting model type: {e}")
    
    def print_config(self):
        """Print current config for debugging"""
        print("\n--- Current Config ---")
        for key, value in vars(self.train_config).items():
            print(f"{key}: {value}")
    
    def reset_config(self):
        """Reset config to defaults"""
        self.train_config = TrainConfig()
        
        # Update UI
        if dpg.does_item_exist("peft_type"):
            dpg.set_value("peft_type", self.train_config.peft_type.name)
        
        if dpg.does_item_exist("lora_rank"):
            dpg.set_value("lora_rank", self.train_config.lora_rank)
        
        if dpg.does_item_exist("lora_alpha"):
            dpg.set_value("lora_alpha", self.train_config.lora_alpha)
        
        if dpg.does_item_exist("lycoris_factor"):
            dpg.set_value("lycoris_factor", self.train_config.lycoris_factor)
        
        if dpg.does_item_exist("lycoris_full_matrix"):
            dpg.set_value("lycoris_full_matrix", self.train_config.lycoris_full_matrix)
        
        if dpg.does_item_exist("lycoris_bypass_mode"):
            dpg.set_value("lycoris_bypass_mode", self.train_config.lycoris_bypass_mode)
        
        if dpg.does_item_exist("lora_decompose"):
            dpg.set_value("lora_decompose", self.train_config.lora_decompose)
        
        if dpg.does_item_exist("lora_decompose_norm_epsilon"):
            dpg.set_value("lora_decompose_norm_epsilon", self.train_config.lora_decompose_norm_epsilon)
        
        if dpg.does_item_exist("lora_decompose_output_axis"):
            dpg.set_value("lora_decompose_output_axis", self.train_config.lora_decompose_output_axis)
        
        if dpg.does_item_exist("dropout_probability"):
            dpg.set_value("dropout_probability", self.train_config.dropout_probability)
        
        if dpg.does_item_exist("bundle_additional_embeddings"):
            dpg.set_value("bundle_additional_embeddings", self.train_config.bundle_additional_embeddings)
        
        # Update UI visibility
        self.update_peft_settings_visibility()
        self.update_lora_presets()
    
    def run(self):
        """Run the application main loop"""
        while dpg.is_dearpygui_running():
            # Update logic here
            dpg.render_dearpygui_frame()
        
        # Clean up DPG
        dpg.destroy_context()

if __name__ == "__main__":
    app = LoraTabTest()
    app.run()