onetrainer_dpg_full.py

#!/usr/bin/env python3
"""
Full DPG Implementation of OneTrainer with all key features from the original.
Includes all tabs and full LyCORIS support.
"""

import os
import sys
import json
import threading
import traceback
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__)))

# Try to import OneTrainer modules, otherwise use mock classes
try:
    from modules.util.enum.ModelType import ModelType, PeftType
    from modules.util.enum.DataType import DataType
    from modules.util.enum.ImageFormat import ImageFormat
    from modules.util.enum.TrainingMethod import TrainingMethod
    from modules.util.enum.TimeUnit import TimeUnit
    from modules.util.enum.ConfigPart import ConfigPart
    from modules.util.enum.ModelFormat import ModelFormat
    from modules.util.config.TrainConfig import TrainConfig
    USING_MOCKS = False
except ImportError as e:
    print(f"Error importing OneTrainer modules: {e}")
    print("Using mock classes for demonstration")
    USING_MOCKS = True
    
    # Define mock enum types
    class ModelType(Enum):
        SD = "Stable Diffusion 1.5"
        SDXL = "Stable Diffusion XL"
        SD3 = "Stable Diffusion 3"
        WUERSTCHEN = "Wuerstchen"
        FLUX = "Flux"
        PIXART_ALPHA = "PixArt Alpha"
        SANA = "Sana"
        HUNYUAN_VIDEO = "Hunyuan Video"
        HI_DREAM = "HiDream"
        
        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"
    
    class ImageFormat(Enum):
        PNG = "png"
        JPG = "jpg"
        WEBP = "webp"
    
    class TrainingMethod(Enum):
        FINE_TUNE = "fine_tune"
        EMBEDDING = "embedding"
        LORA = "lora"
    
    class TimeUnit(Enum):
        STEP = "step"
        EPOCH = "epoch"
    
    class ConfigPart(Enum):
        NONE = "none"
        MINIMAL = "minimal"
        FULL = "full"
    
    class ModelFormat(Enum):
        CKPT = "ckpt"
        SAFETENSORS = "safetensors"
        DIFFUSERS = "diffusers"
    
    @dataclass
    class CloudConfig:
        enabled: bool = False
        ssh_host: str = ""
        ssh_port: int = 22
        ssh_user: str = ""
        ssh_key: str = ""
        remote_path: str = ""
        
    @dataclass
    class SecretsConfig:
        huggingface_token: str = ""
        cloud: CloudConfig = field(default_factory=CloudConfig)
    
    @dataclass
    class TrainConfig:
        """Mock TrainConfig for demonstration"""
        # General settings
        model_type: ModelType = ModelType.SD
        training_method: TrainingMethod = TrainingMethod.LORA
        peft_type: PeftType = PeftType.LORA
        weight_dtype: DataType = DataType.FLOAT_32
        output_dtype: DataType = DataType.FLOAT_32
        train_dtype: DataType = DataType.FLOAT_16
        fallback_train_dtype: DataType = DataType.BFLOAT_16
        output_model_format: ModelFormat = ModelFormat.SAFETENSORS
        include_train_config: ConfigPart = ConfigPart.MINIMAL
        
        # Paths
        workspace_dir: str = "./workspace"
        cache_dir: str = "./cache"
        debug_dir: str = "./debug"
        base_model_name: str = ""
        output_model_destination: str = ""
        
        # Flags
        debug_mode: bool = False
        continue_last_backup: bool = False
        only_cache: bool = False
        
        # Tensorboard
        tensorboard: bool = True
        tensorboard_expose: bool = False
        tensorboard_port: int = 6006
        
        # Validation
        validation: bool = False
        validate_after: int = 100
        validate_after_unit: TimeUnit = TimeUnit.STEP
        
        # System
        dataloader_threads: int = 4
        train_device: str = "cuda:0"
        temp_device: str = "cuda:0"
        
        # Data
        aspect_ratio_bucketing: bool = True
        latent_caching: bool = True
        clear_cache_before_training: bool = True
        max_resolution: int = 512
        
        # Training
        batch_size: int = 1
        epochs: int = 10
        gradient_accumulation_steps: int = 1
        mixed_precision: bool = True
        learning_rate: float = 1e-5
        learning_rate_warmup_steps: int = 0
        optimizer_weight_decay: float = 0.01
        optimizer_beta1: float = 0.9
        optimizer_beta2: float = 0.999
        optimizer_epsilon: float = 1e-8
        use_ema: bool = True
        ema_decay: float = 0.9999
        
        # LoRA
        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
        
        # DoRA
        lora_decompose: bool = False
        lora_decompose_norm_epsilon: bool = True
        lora_decompose_output_axis: bool = False
        
        # LyCORIS
        lycoris_factor: int = 4
        lycoris_full_matrix: bool = False
        lycoris_bypass_mode: bool = False
        
        # Sampling
        sample_after: int = 500
        sample_after_unit: TimeUnit = TimeUnit.STEP
        sample_skip_first: int = 0
        sample_image_format: ImageFormat = ImageFormat.PNG
        non_ema_sampling: bool = False
        samples_to_tensorboard: bool = True
        cfg_scale: float = 7.5
        sample_steps: int = 30
        scheduler: str = "euler_a"
        
        # Backup
        backup_after: int = 500
        backup_after_unit: TimeUnit = TimeUnit.STEP
        rolling_backup: bool = False
        rolling_backup_count: int = 5
        backup_before_save: bool = True
        
        # Save
        save_every: int = 500
        save_every_unit: TimeUnit = TimeUnit.STEP
        save_skip_first: int = 0
        save_filename_prefix: str = ""
        
        # Cloud
        cloud: CloudConfig = field(default_factory=CloudConfig)
        secrets: SecretsConfig = field(default_factory=SecretsConfig)
        
        # Model components
        class ModelComponent:
            def __init__(self):
                self.weight_dtype = DataType.NONE
                self.model_name = ""
        
        unet = ModelComponent()
        text_encoder = ModelComponent()
        text_encoder_2 = ModelComponent()
        text_encoder_3 = ModelComponent()
        vae = ModelComponent()
        prior = ModelComponent()
        effnet_encoder = ModelComponent()
        
        @staticmethod
        def default_values():
            """Return a new instance with default values"""
            return TrainConfig()

# Default layer presets for different model types
DEFAULT_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": [".*"],
    "sd3-full": ["to_qkv", "to_out.0", "ff.0", "ff.2"]
}

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

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

class OneTrainerDPG:
    """Complete OneTrainer DPG Implementation with all tabs"""
    
    def __init__(self):
        """Initialize the application"""
        # Initialize DPG
        dpg.create_context()
        dpg.create_viewport(title="OneTrainer (DPG)", width=WINDOW_WIDTH, height=WINDOW_HEIGHT)
        
        # Initialize state
        self.train_config = TrainConfig() if USING_MOCKS else TrainConfig.default_values()
        
        # Setup UI
        self.setup_ui()
        
        # Set up DPG
        dpg.setup_dearpygui()
        dpg.show_viewport()
    
    def setup_ui(self):
        """Set up the user interface"""
        # Create main window
        self.main_window = dpg.add_window(label="OneTrainer", width=WINDOW_WIDTH, height=WINDOW_HEIGHT,
                                      no_title_bar=True, no_resize=True, no_move=True, no_collapse=True)
        
        # Create main layout
        self.main_layout = dpg.add_group(parent=self.main_window)
        
        # Create top bar
        self.create_top_bar()
        
        # Create tab bar
        self.tab_bar = dpg.add_tab_bar(parent=self.main_layout)
        
        # Create tabs
        self.create_model_tab()
        self.create_lora_tab()
        self.create_concept_tab()
        self.create_training_tab()
        self.create_sampling_tab()
        self.create_cloud_tab()
        
        # Create bottom bar
        self.create_bottom_bar()
    
    def create_top_bar(self):
        """Create top bar with logo and model selection"""
        # Create top bar group
        self.top_bar = dpg.add_group(parent=self.main_layout, horizontal=True)
        
        # Add app title
        dpg.add_text(parent=self.top_bar, default_value="OneTrainer")
        dpg.add_spacer(parent=self.top_bar, width=20)
        
        # Add model type selection
        dpg.add_text(parent=self.top_bar, default_value="Model Type:")
        
        model_types = [mt.name for mt in ModelType]
        self.model_type_combo = dpg.add_combo(
            parent=self.top_bar,
            items=model_types,
            default_value=self.train_config.model_type.name,
            callback=self.on_model_type_changed
        )
        dpg.add_spacer(parent=self.top_bar, width=10)
        
        # Add training method selection
        dpg.add_text(parent=self.top_bar, default_value="Training Method:")
        
        training_methods = [tm.name for tm in TrainingMethod]
        self.training_method_combo = dpg.add_combo(
            parent=self.top_bar,
            items=training_methods,
            default_value=self.train_config.training_method.name,
            callback=self.on_training_method_changed
        )
    
    def create_bottom_bar(self):
        """Create bottom bar with status and controls"""
        # Create bottom bar group
        self.bottom_bar = dpg.add_group(parent=self.main_layout, horizontal=True)
        
        # Status and progress
        self.status_group = dpg.add_group(parent=self.bottom_bar, horizontal=False, width=400)
        
        # Progress bars
        dpg.add_text(parent=self.status_group, default_value="Step Progress:")
        self.step_progress = dpg.add_progress_bar(parent=self.status_group, default_value=0, width=-1)
        
        dpg.add_text(parent=self.status_group, default_value="Epoch Progress:")
        self.epoch_progress = dpg.add_progress_bar(parent=self.status_group, default_value=0, width=-1)
        
        # Status text
        self.status_text_group = dpg.add_group(parent=self.bottom_bar, horizontal=True)
        dpg.add_text(parent=self.status_text_group, default_value="Status:")
        self.status_text = dpg.add_text(parent=self.status_text_group, default_value="Ready")
        
        # Control buttons
        self.control_group = dpg.add_group(parent=self.bottom_bar, horizontal=True)
        self.train_button = dpg.add_button(parent=self.control_group, label="Start Training", callback=self.start_training)
        self.export_button = dpg.add_button(parent=self.control_group, label="Export", callback=self.export_model)
    
    def export_model(self):
        """Export the trained model"""
        print("Would export model")
        dpg.set_value(self.status_text, "Exporting model...")
    
    def create_model_tab(self):
        """Create model configuration tab"""
        # Create main tab
        self.model_tab = dpg.add_tab(parent=self.tab_bar, label="Model")
        
        # Create layout group
        self.model_group = dpg.add_group(parent=self.model_tab)
        
        # Base model
        base_model_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=base_model_group, default_value="Base Model:")
        
        self.base_model_input = dpg.add_input_text(
            parent=base_model_group,
            default_value=self.train_config.base_model_name,
            callback=lambda s, a: setattr(self.train_config, "base_model_name", a),
            width=-100
        )
        
        self.base_model_browse = dpg.add_button(
            parent=base_model_group,
            label="...",
            callback=lambda: self.browse_file("base_model_input")
        )
        
        # Weight data type
        weight_dtype_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=weight_dtype_group, default_value="Weight Data Type:")
        
        dtype_options = [dt.name for dt in DataType if dt != DataType.NONE]
        self.weight_dtype_combo = dpg.add_combo(
            parent=weight_dtype_group,
            items=dtype_options,
            default_value=self.train_config.weight_dtype.name,
            callback=lambda s, a: self.set_config_enum("weight_dtype", DataType, a)
        )
        
        # Model-specific settings placeholder
        dpg.add_separator(parent=self.model_group)
        dpg.add_text(parent=self.model_group, default_value="Model-Specific Settings")
        self.model_specific_group = dpg.add_group(parent=self.model_group)
        
        # Output settings
        dpg.add_separator(parent=self.model_group)
        dpg.add_text(parent=self.model_group, default_value="Output Settings")
        
        # Output destination
        output_dest_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=output_dest_group, default_value="Output Destination:")
        
        self.output_dest_input = dpg.add_input_text(
            parent=output_dest_group,
            default_value=self.train_config.output_model_destination,
            callback=lambda s, a: setattr(self.train_config, "output_model_destination", a),
            width=-100
        )
        
        self.output_dest_browse = dpg.add_button(
            parent=output_dest_group,
            label="...",
            callback=lambda: self.browse_file("output_dest_input", is_save=True)
        )
        
        # Output data type
        output_dtype_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=output_dtype_group, default_value="Output Data Type:")
        
        self.output_dtype_combo = dpg.add_combo(
            parent=output_dtype_group,
            items=dtype_options,
            default_value=self.train_config.output_dtype.name,
            callback=lambda s, a: self.set_config_enum("output_dtype", DataType, a)
        )
        
        # Output format
        output_format_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=output_format_group, default_value="Output Format:")
        
        format_options = [fmt.name for fmt in ModelFormat]
        self.output_format_combo = dpg.add_combo(
            parent=output_format_group,
            items=format_options,
            default_value=self.train_config.output_model_format.name,
            callback=lambda s, a: self.set_config_enum("output_model_format", ModelFormat, a)
        )
        
        # Include config
        include_config_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=include_config_group, default_value="Include Train Config:")
        
        config_part_options = [cp.name for cp in ConfigPart]
        self.include_config_combo = dpg.add_combo(
            parent=include_config_group,
            items=config_part_options,
            default_value=self.train_config.include_train_config.name,
            callback=lambda s, a: self.set_config_enum("include_train_config", ConfigPart, a)
        )
        
        # Update model-specific UI based on current model type
        self.update_model_specific_ui()
    
    def update_model_specific_ui(self):
        """Update model-specific UI elements based on model type"""
        # Clear existing elements
        if dpg.does_item_exist(self.model_specific_group):
            for child in dpg.get_item_children(self.model_specific_group, slot=1):
                dpg.delete_item(child)
        
        # Add elements based on model type
        model_type = self.train_config.model_type
        
        # This is a simplified version - in a real implementation, add all the model-specific controls
        info_text = f"Model-specific settings for {model_type.name} would be shown here."
        dpg.add_text(parent=self.model_specific_group, default_value=info_text)
    
    def create_lora_tab(self):
        """Create LoRA tab with full LyCORIS support"""
        # Create main tab
        self.lora_tab = dpg.add_tab(parent=self.tab_bar, label="LoRA")
        
        # Create layout group
        self.lora_group = dpg.add_group(parent=self.lora_tab)
        
        # Title
        dpg.add_text(parent=self.lora_group, default_value="LoRA / LyCORIS Settings", color=[255, 255, 255, 255])
        dpg.add_separator(parent=self.lora_group)
        
        # PEFT type selector with tooltip
        peft_type_group = dpg.add_group(parent=self.lora_group, horizontal=True)
        dpg.add_text(parent=peft_type_group, default_value="PEFT Type:")
        
        peft_types = [pt.name for pt in PeftType]
        self.peft_type_combo = dpg.add_combo(
            parent=peft_type_group,
            items=peft_types,
            default_value=self.train_config.peft_type.name,
            callback=self.on_peft_type_changed
        )
        
        # Add tooltip for PEFT types
        self.peft_tooltip = dpg.add_text(
            parent=dpg.add_tooltip(self.peft_type_combo),
            default_value=(
                "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"
            )
        )
        
        # LoRA base model
        lora_model_group = dpg.add_group(parent=self.lora_group, horizontal=True)
        dpg.add_text(parent=lora_model_group, default_value="LoRA base model:")
        
        self.lora_model_input = dpg.add_input_text(
            parent=lora_model_group,
            default_value=self.train_config.lora_model_name,
            callback=lambda s, a: setattr(self.train_config, "lora_model_name", a),
            width=-100
        )
        
        self.lora_model_browse = dpg.add_button(
            parent=lora_model_group,
            label="...",
            callback=lambda: self.browse_file("lora_model_input")
        )
        
        # Common settings
        # LoRA rank
        lora_rank_group = dpg.add_group(parent=self.lora_group, horizontal=True)
        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),
            min_value=1,
            max_value=1024,
            step=1
        )
        
        # LoRA alpha
        lora_alpha_group = dpg.add_group(parent=self.lora_group, horizontal=True)
        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),
            min_value=0.1,
            max_value=1024.0,
            step=0.1
        )
        
        # LyCORIS-specific settings
        dpg.add_separator(parent=self.lora_group)
        dpg.add_text(parent=self.lora_group, default_value="LyCORIS Settings")
        self.lycoris_group = dpg.add_group(parent=self.lora_group)
        
        # LoKr/LoHa factor
        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),
            min_value=1,
            max_value=16,
            step=1
        )
        
        factor_tooltip = dpg.add_tooltip(self.lycoris_factor_input)
        dpg.add_text(
            parent=factor_tooltip,
            default_value=(
                "Compression factor for LoKr/LoHa. Lower = less parameters.\n"
                "Used with rank to determine total parameters = (width × height) ÷ (factor × rank)"
            )
        )
        
        # DyLoRA 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)
        )
        
        full_matrix_tooltip = dpg.add_tooltip(self.lycoris_full_matrix_check)
        dpg.add_text(
            parent=full_matrix_tooltip,
            default_value=(
                "DyLoRA: Use full weight matrix (larger model but better quality)\n"
                "This increases the parameter count significantly"
            )
        )
        
        # Bypass mode toggle
        bypass_group = dpg.add_group(parent=self.lycoris_group, horizontal=True)
        dpg.add_text(parent=bypass_group, default_value="Bypass Mode (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)
        )
        
        bypass_tooltip = dpg.add_tooltip(self.lycoris_bypass_check)
        dpg.add_text(
            parent=bypass_tooltip,
            default_value=(
                "Bypass problematic LyCORIS types (like LoKr) for HiDream models\n"
                "Only needed if you have compatibility issues"
            )
        )
        
        # LoRA-specific settings (DoRA)
        dpg.add_separator(parent=self.lora_group)
        dpg.add_text(parent=self.lora_group, default_value="LoRA Settings")
        self.lora_specific_group = dpg.add_group(parent=self.lora_group)
        
        # DoRA toggle
        dora_toggle_group = dpg.add_group(parent=self.lora_specific_group, horizontal=True)
        dpg.add_text(parent=dora_toggle_group, default_value="Weight Decomposition (DoRA):")
        
        self.lora_decompose_check = dpg.add_checkbox(
            parent=dora_toggle_group,
            default_value=self.train_config.lora_decompose,
            callback=self.on_dora_toggle
        )
        
        dora_tooltip = dpg.add_tooltip(self.lora_decompose_check)
        dpg.add_text(
            parent=dora_tooltip,
            default_value=(
                "Enable weight decomposition (DoRA - Decomposed Rank Adaptation)\n"
                "This improves quality but increases training time"
            )
        )
        
        # DoRA advanced settings
        self.dora_group = dpg.add_group(parent=self.lora_specific_group)
        
        # Norm epsilon
        norm_group = dpg.add_group(parent=self.dora_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_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)
        )
        
        # Common settings for all PEFT types
        dpg.add_separator(parent=self.lora_group)
        dpg.add_text(parent=self.lora_group, default_value="Common Settings")
        
        # Dropout
        dropout_group = dpg.add_group(parent=self.lora_group, 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,
            callback=lambda s, a: setattr(self.train_config, "dropout_probability", a),
            min_value=0.0,
            max_value=1.0,
            width=200
        )
        
        dropout_tooltip = dpg.add_tooltip(self.dropout_slider)
        dpg.add_text(
            parent=dropout_tooltip,
            default_value=(
                "Dropout probability for LoRA layers (0.0 = disabled)\n"
                "Higher values can help prevent overfitting"
            )
        )
        
        # Weight Data Type
        weight_dtype_group = dpg.add_group(parent=self.lora_group, horizontal=True)
        dpg.add_text(parent=weight_dtype_group, default_value="Weight Data Type:")
        
        lora_dtype_options = [dt.name for dt in [DataType.FLOAT_32, DataType.BFLOAT_16]]
        self.lora_weight_dtype_combo = dpg.add_combo(
            parent=weight_dtype_group,
            items=lora_dtype_options,
            default_value=self.train_config.lora_weight_dtype.name,
            callback=lambda s, a: self.set_config_enum("lora_weight_dtype", DataType, a)
        )
        
        # Bundle embeddings
        bundle_group = dpg.add_group(parent=self.lora_group, horizontal=True)
        
        self.bundle_embeddings_check = dpg.add_checkbox(
            parent=bundle_group,
            label="Bundle Embeddings",
            default_value=self.train_config.bundle_additional_embeddings,
            callback=lambda s, a: setattr(self.train_config, "bundle_additional_embeddings", a)
        )
        
        bundle_tooltip = dpg.add_tooltip(self.bundle_embeddings_check)
        dpg.add_text(
            parent=bundle_tooltip,
            default_value="Include trained textual embeddings within the LoRA file"
        )
        
        # Layer Presets
        dpg.add_separator(parent=self.lora_group)
        dpg.add_text(parent=self.lora_group, default_value="Layer Settings")
        
        # Preset selector
        preset_group = dpg.add_group(parent=self.lora_group, horizontal=True)
        dpg.add_text(parent=preset_group, default_value="Layer Preset:")
        
        # Get presets based on model type
        self.current_presets = self.get_lora_presets(self.train_config.model_type)
        preset_list = list(self.current_presets.keys()) + ["custom"]
        
        self.layer_preset_combo = dpg.add_combo(
            parent=preset_group,
            items=preset_list,
            default_value=self.train_config.lora_layer_preset,
            callback=self.on_layer_preset_changed
        )
        
        preset_tooltip = dpg.add_tooltip(self.layer_preset_combo)
        dpg.add_text(
            parent=preset_tooltip,
            default_value=(
                "Presets for which layers to target with the LoRA adapter\n"
                "Different models have different recommended targeting strategies"
            )
        )
        
        # Custom layers input
        layers_group = dpg.add_group(parent=self.lora_group, 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),
            width=-100
        )
        
        layers_tooltip = dpg.add_tooltip(self.layers_input)
        dpg.add_text(
            parent=layers_tooltip,
            default_value=(
                "Comma-separated list of layers to apply LoRA to\n"
                "Only used when 'custom' preset is selected"
            )
        )
        
        # Update visibility based on current settings
        self.update_lora_ui_visibility()
    
    def update_lora_ui_visibility(self):
        """Update visibility of LoRA UI elements based on settings"""
        # Show/hide LyCORIS settings based on PEFT type
        peft_type = self.train_config.peft_type
        show_lycoris = peft_type in [PeftType.LOHA, PeftType.LOKR, PeftType.DIA, PeftType.IA3, PeftType.DYLORA]
        dpg.configure_item(self.lycoris_group, show=show_lycoris)
        
        # Show specific LyCORIS settings based on type
        if show_lycoris:
            # Show factor for LoHa/LoKr
            show_factor = peft_type in [PeftType.LOHA, PeftType.LOKR]
            for child in dpg.get_item_children(self.lycoris_group, slot=1):
                if "Factor" in dpg.get_item_label(child):
                    dpg.configure_item(child, show=show_factor)
            
            # Show full matrix option for DyLoRA
            show_full_matrix = peft_type == PeftType.DYLORA
            for child in dpg.get_item_children(self.lycoris_group, slot=1):
                if "Full Matrix" in dpg.get_item_label(child):
                    dpg.configure_item(child, show=show_full_matrix)
            
            # Show bypass mode for problematic types with HiDream
            show_bypass = (
                self.train_config.model_type.is_hi_dream() and 
                peft_type in [PeftType.LOKR]
            )
            for child in dpg.get_item_children(self.lycoris_group, slot=1):
                if "Bypass Mode" in dpg.get_item_label(child):
                    dpg.configure_item(child, show=show_bypass)
        
        # Show/hide LoRA-specific settings
        show_lora = peft_type == PeftType.LORA
        dpg.configure_item(self.lora_specific_group, show=show_lora)
        
        # Show/hide DoRA settings
        if show_lora:
            show_dora = self.train_config.lora_decompose
            dpg.configure_item(self.dora_group, show=show_dora)
    
    def get_lora_presets(self, model_type):
        """Get layer presets based on model type"""
        if model_type.is_hi_dream():
            return HIDREAM_LAYER_PRESETS
        else:
            return DEFAULT_LAYER_PRESETS
    
    def on_layer_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 self.current_presets:
            layers = ", ".join(self.current_presets[app_data])
            self.train_config.lora_layers = layers
            dpg.set_value(self.layers_input, layers)
    
    def on_peft_type_changed(self, sender, app_data):
        """Handle PEFT type change"""
        try:
            self.train_config.peft_type = PeftType[app_data]
            self.update_lora_ui_visibility()
        except Exception as e:
            print(f"Error changing PEFT type: {e}")
    
    def on_dora_toggle(self, sender, app_data):
        """Handle DoRA toggle"""
        self.train_config.lora_decompose = app_data
        self.update_lora_ui_visibility()
    
    def create_concept_tab(self):
        """Create concept configuration tab"""
        # Create main tab
        self.concept_tab = dpg.add_tab(parent=self.tab_bar, label="Concepts")
        
        # Create layout group
        self.concept_group = dpg.add_group(parent=self.concept_tab)
        
        # Title
        dpg.add_text(parent=self.concept_group, default_value="Training Concepts", color=[255, 255, 255, 255])
        dpg.add_separator(parent=self.concept_group)
        
        # Add concept list
        dpg.add_text(parent=self.concept_group, default_value="Concepts:")
        
        # Create concept table
        self.concept_table = dpg.add_table(
            parent=self.concept_group,
            header_row=True,
            resizable=True,
            policy=dpg.mvTable_SizingStretchProp,
            borders_innerH=True,
            borders_outerH=True,
            borders_innerV=True,
            borders_outerV=True
        )
        
        # Add table columns
        dpg.add_table_column(parent=self.concept_table, label="Name")
        dpg.add_table_column(parent=self.concept_table, label="Images")
        dpg.add_table_column(parent=self.concept_table, label="Weight")
        dpg.add_table_column(parent=self.concept_table, label="Actions")
        
        # Add sample row
        with dpg.table_row(parent=self.concept_table):
            dpg.add_text(default_value="Example Concept")
            dpg.add_text(default_value="0 images")
            dpg.add_text(default_value="1.0")
            with dpg.group(horizontal=True):
                dpg.add_button(label="Edit", callback=lambda: None)
                dpg.add_button(label="Delete", callback=lambda: None)
        
        # Concept controls
        concept_controls = dpg.add_group(parent=self.concept_group, horizontal=True)
        dpg.add_button(parent=concept_controls, label="Add Concept", callback=self.add_concept)
        dpg.add_button(parent=concept_controls, label="Refresh", callback=self.refresh_concepts)
        
        # Concept settings
        dpg.add_separator(parent=self.concept_group)
        dpg.add_text(parent=self.concept_group, default_value="Concept Settings")
        
        # Resolution
        res_group = dpg.add_group(parent=self.concept_group, horizontal=True)
        dpg.add_text(parent=res_group, default_value="Max Resolution:")
        
        self.max_res_input = dpg.add_input_int(
            parent=res_group,
            default_value=512,  # Default to 512
            callback=lambda s, a: setattr(self.train_config, "max_resolution", a),
            min_value=128,
            max_value=4096,
            step=64
        )
        
        # Resolution type
        res_type_group = dpg.add_group(parent=self.concept_group, horizontal=True)
        dpg.add_text(parent=res_type_group, default_value="Resolution Type:")
        
        self.res_type_combo = dpg.add_combo(
            parent=res_type_group,
            items=["square", "original"],
            default_value="square",
            callback=lambda s, a: setattr(self.train_config, "resolution_type", a)
        )
    
    def add_concept(self):
        """Add a new concept"""
        print("Would add a new concept")
    
    def refresh_concepts(self):
        """Refresh concept list"""
        print("Would refresh concepts")
    
    def create_training_tab(self):
        """Create training configuration tab"""
        # Create main tab
        self.training_tab = dpg.add_tab(parent=self.tab_bar, label="Training")
        
        # Create layout group
        self.training_group = dpg.add_group(parent=self.training_tab)
        
        # Title
        dpg.add_text(parent=self.training_group, default_value="Training Settings", color=[255, 255, 255, 255])
        dpg.add_separator(parent=self.training_group)
        
        # Basic training parameters
        basic_group = dpg.add_group(parent=self.training_group)
        dpg.add_text(parent=basic_group, default_value="Basic Parameters")
        
        # Batch size
        batch_group = dpg.add_group(parent=basic_group, horizontal=True)
        dpg.add_text(parent=batch_group, default_value="Batch Size:")
        
        self.batch_size_input = dpg.add_input_int(
            parent=batch_group,
            default_value=self.train_config.batch_size,
            callback=lambda s, a: setattr(self.train_config, "batch_size", a),
            min_value=1,
            max_value=128,
            step=1
        )
        
        # Epochs
        epochs_group = dpg.add_group(parent=basic_group, horizontal=True)
        dpg.add_text(parent=epochs_group, default_value="Epochs:")
        
        self.epochs_input = dpg.add_input_int(
            parent=epochs_group,
            default_value=self.train_config.epochs,
            callback=lambda s, a: setattr(self.train_config, "epochs", a),
            min_value=1,
            max_value=1000,
            step=1
        )
        
        # Gradient accumulation
        grad_accum_group = dpg.add_group(parent=basic_group, horizontal=True)
        dpg.add_text(parent=grad_accum_group, default_value="Gradient Accumulation Steps:")
        
        self.grad_accum_input = dpg.add_input_int(
            parent=grad_accum_group,
            default_value=self.train_config.gradient_accumulation_steps,
            callback=lambda s, a: setattr(self.train_config, "gradient_accumulation_steps", a),
            min_value=1,
            max_value=64,
            step=1
        )
        
        # Mixed precision
        mixed_prec_group = dpg.add_group(parent=basic_group, horizontal=True)
        
        # Handle mixed precision settings - in OneTrainer, mixed precision is controlled by train_dtype
        # We'll map the checkbox to setting train_dtype to either FLOAT_16 (True) or FLOAT_32 (False)
        is_mixed = getattr(self.train_config, "train_dtype", DataType.FLOAT_16) in [DataType.FLOAT_16, DataType.BFLOAT_16]
        
        def toggle_mixed_precision(sender, data):
            if data:  # If checked, use FLOAT_16
                setattr(self.train_config, "train_dtype", DataType.FLOAT_16)
            else:  # If unchecked, use FLOAT_32
                setattr(self.train_config, "train_dtype", DataType.FLOAT_32)
        
        self.mixed_prec_check = dpg.add_checkbox(
            parent=mixed_prec_group,
            label="Mixed Precision",
            default_value=is_mixed,
            callback=toggle_mixed_precision
        )
        
        # Learning rate settings
        dpg.add_separator(parent=self.training_group)
        dpg.add_text(parent=self.training_group, default_value="Learning Rate Settings")
        
        lr_group = dpg.add_group(parent=self.training_group)
        
        # Learning rate
        lr_input_group = dpg.add_group(parent=lr_group, horizontal=True)
        dpg.add_text(parent=lr_input_group, default_value="Learning Rate:")
        
        self.lr_input = dpg.add_input_float(
            parent=lr_input_group,
            default_value=self.train_config.learning_rate,
            callback=lambda s, a: setattr(self.train_config, "learning_rate", a),
            format="%.8f",
            min_value=1e-8,
            max_value=1.0,
            step=0.00001
        )
        
        # Scheduler
        scheduler_group = dpg.add_group(parent=lr_group, horizontal=True)
        dpg.add_text(parent=scheduler_group, default_value="Learning Rate Scheduler:")
        
        scheduler_options = ["constant", "linear", "cosine", "cosine_with_restarts", "polynomial"]
        self.lr_scheduler_combo = dpg.add_combo(
            parent=scheduler_group,
            items=scheduler_options,
            default_value="constant",
            callback=lambda s, a: setattr(self.train_config, "lr_scheduler", a)
        )
        
        # Warmup steps
        warmup_group = dpg.add_group(parent=lr_group, horizontal=True)
        dpg.add_text(parent=warmup_group, default_value="Warmup Steps:")
        
        self.warmup_input = dpg.add_input_int(
            parent=warmup_group,
            default_value=getattr(self.train_config, "learning_rate_warmup_steps", 200),
            callback=lambda s, a: setattr(self.train_config, "learning_rate_warmup_steps", a),
            min_value=0,
            max_value=10000,
            step=10
        )
        
        # Optimizer settings
        dpg.add_separator(parent=self.training_group)
        dpg.add_text(parent=self.training_group, default_value="Optimizer Settings")
        
        opt_group = dpg.add_group(parent=self.training_group)
        
        # Optimizer
        opt_type_group = dpg.add_group(parent=opt_group, horizontal=True)
        dpg.add_text(parent=opt_type_group, default_value="Optimizer:")
        
        optimizer_options = ["adamw", "adam", "sgd", "lion", "adamw8bit", "lion8bit", "came"]
        self.optimizer_combo = dpg.add_combo(
            parent=opt_type_group,
            items=optimizer_options,
            default_value="adamw",
            callback=lambda s, a: setattr(self.train_config, "optimizer", a)
        )
        
        # Weight decay
        weight_decay_group = dpg.add_group(parent=opt_group, horizontal=True)
        dpg.add_text(parent=weight_decay_group, default_value="Weight Decay:")
        
        self.weight_decay_input = dpg.add_input_float(
            parent=weight_decay_group,
            default_value=getattr(self.train_config, "optimizer_weight_decay", 0.01),
            callback=lambda s, a: setattr(self.train_config, "optimizer_weight_decay", a),
            format="%.4f",
            min_value=0.0,
            max_value=1.0,
            step=0.001
        )
        
        # Beta1
        beta1_group = dpg.add_group(parent=opt_group, horizontal=True)
        dpg.add_text(parent=beta1_group, default_value="Adam Beta1:")
        
        self.beta1_input = dpg.add_input_float(
            parent=beta1_group,
            default_value=getattr(self.train_config, "optimizer_beta1", 0.9),
            callback=lambda s, a: setattr(self.train_config, "optimizer_beta1", a),
            format="%.4f",
            min_value=0.0,
            max_value=1.0,
            step=0.001
        )
        
        # Beta2
        beta2_group = dpg.add_group(parent=opt_group, horizontal=True)
        dpg.add_text(parent=beta2_group, default_value="Adam Beta2:")
        
        self.beta2_input = dpg.add_input_float(
            parent=beta2_group,
            default_value=getattr(self.train_config, "optimizer_beta2", 0.999),
            callback=lambda s, a: setattr(self.train_config, "optimizer_beta2", a),
            format="%.6f",
            min_value=0.0,
            max_value=1.0,
            step=0.0001
        )
        
        # Epsilon
        epsilon_group = dpg.add_group(parent=opt_group, horizontal=True)
        dpg.add_text(parent=epsilon_group, default_value="Adam Epsilon:")
        
        self.epsilon_input = dpg.add_input_float(
            parent=epsilon_group,
            default_value=getattr(self.train_config, "optimizer_epsilon", 1e-8),
            callback=lambda s, a: setattr(self.train_config, "optimizer_epsilon", a),
            format="%.8f",
            min_value=1e-10,
            max_value=1e-3,
            step=1e-9
        )
        
        # EMA settings
        dpg.add_separator(parent=self.training_group)
        dpg.add_text(parent=self.training_group, default_value="EMA Settings")
        
        ema_group = dpg.add_group(parent=self.training_group)
        
        # EMA checkbox
        ema_toggle_group = dpg.add_group(parent=ema_group, horizontal=True)
        
        self.ema_check = dpg.add_checkbox(
            parent=ema_toggle_group,
            label="Use EMA",
            default_value=self.train_config.use_ema,
            callback=lambda s, a: setattr(self.train_config, "use_ema", a)
        )
        
        # EMA decay
        ema_decay_group = dpg.add_group(parent=ema_group, horizontal=True)
        dpg.add_text(parent=ema_decay_group, default_value="EMA Decay:")
        
        self.ema_decay_input = dpg.add_input_float(
            parent=ema_decay_group,
            default_value=self.train_config.ema_decay,
            callback=lambda s, a: setattr(self.train_config, "ema_decay", a),
            format="%.6f",
            min_value=0.9,
            max_value=0.99999,
            step=0.0001
        )
    
    def create_sampling_tab(self):
        """Create sampling configuration tab"""
        # Create main tab
        self.sampling_tab = dpg.add_tab(parent=self.tab_bar, label="Sampling")
        
        # Create layout group
        self.sampling_group = dpg.add_group(parent=self.sampling_tab)
        
        # Title
        dpg.add_text(parent=self.sampling_group, default_value="Sampling Settings", color=[255, 255, 255, 255])
        dpg.add_separator(parent=self.sampling_group)
        
        # Sample timing
        dpg.add_text(parent=self.sampling_group, default_value="Sampling Schedule")
        
        # After
        after_group = dpg.add_group(parent=self.sampling_group, horizontal=True)
        dpg.add_text(parent=after_group, default_value="Sample After:")
        
        self.sample_after_input = dpg.add_input_int(
            parent=after_group,
            default_value=self.train_config.sample_after,
            callback=lambda s, a: setattr(self.train_config, "sample_after", a),
            min_value=1,
            max_value=10000,
            step=10
        )
        
        time_units = [tu.name for tu in TimeUnit]
        self.sample_after_unit_combo = dpg.add_combo(
            parent=after_group,
            items=time_units,
            default_value=self.train_config.sample_after_unit.name,
            callback=lambda s, a: self.set_config_enum("sample_after_unit", TimeUnit, a)
        )
        
        # Skip first
        skip_group = dpg.add_group(parent=self.sampling_group, horizontal=True)
        dpg.add_text(parent=skip_group, default_value="Skip First:")
        
        self.sample_skip_input = dpg.add_input_int(
            parent=skip_group,
            default_value=self.train_config.sample_skip_first,
            callback=lambda s, a: setattr(self.train_config, "sample_skip_first", a),
            min_value=0,
            max_value=1000,
            step=1
        )
        
        # Image format
        format_group = dpg.add_group(parent=self.sampling_group, horizontal=True)
        dpg.add_text(parent=format_group, default_value="Image Format:")
        
        image_formats = [fmt.name for fmt in ImageFormat]
        self.sample_format_combo = dpg.add_combo(
            parent=format_group,
            items=image_formats,
            default_value=self.train_config.sample_image_format.name,
            callback=lambda s, a: self.set_config_enum("sample_image_format", ImageFormat, a)
        )
        
        # Sample controls
        sample_controls = dpg.add_group(parent=self.sampling_group, horizontal=True)
        dpg.add_button(parent=sample_controls, label="Sample Now", callback=self.sample_now)
        dpg.add_button(parent=sample_controls, label="Open Sampler", callback=self.open_sampler)
        
        # Settings checkboxes
        checkboxes_group = dpg.add_group(parent=self.sampling_group)
        
        self.non_ema_check = dpg.add_checkbox(
            parent=checkboxes_group,
            label="Non-EMA Sampling",
            default_value=self.train_config.non_ema_sampling,
            callback=lambda s, a: setattr(self.train_config, "non_ema_sampling", a)
        )
        
        self.tb_samples_check = dpg.add_checkbox(
            parent=checkboxes_group,
            label="Samples to Tensorboard",
            default_value=self.train_config.samples_to_tensorboard,
            callback=lambda s, a: setattr(self.train_config, "samples_to_tensorboard", a)
        )
        
        # Sample parameters
        dpg.add_separator(parent=self.sampling_group)
        dpg.add_text(parent=self.sampling_group, default_value="Sampling Parameters")
        
        # CFG Scale
        cfg_group = dpg.add_group(parent=self.sampling_group, horizontal=True)
        dpg.add_text(parent=cfg_group, default_value="CFG Scale:")
        
        self.cfg_scale_slider = dpg.add_slider_float(
            parent=cfg_group,
            default_value=self.train_config.cfg_scale,
            callback=lambda s, a: setattr(self.train_config, "cfg_scale", a),
            min_value=1.0,
            max_value=20.0,
            width=200
        )
        
        # Steps
        steps_group = dpg.add_group(parent=self.sampling_group, horizontal=True)
        dpg.add_text(parent=steps_group, default_value="Steps:")
        
        self.sample_steps_slider = dpg.add_slider_int(
            parent=steps_group,
            default_value=self.train_config.sample_steps,
            callback=lambda s, a: setattr(self.train_config, "sample_steps", a),
            min_value=1,
            max_value=150,
            width=200
        )
        
        # Scheduler
        scheduler_group = dpg.add_group(parent=self.sampling_group, horizontal=True)
        dpg.add_text(parent=scheduler_group, default_value="Scheduler:")
        
        scheduler_options = ["ddim", "ddpm", "k_lms", "pndm", "euler", "euler_a", "dpmpp_2m"]
        self.scheduler_combo = dpg.add_combo(
            parent=scheduler_group,
            items=scheduler_options,
            default_value=self.train_config.scheduler,
            callback=lambda s, a: setattr(self.train_config, "scheduler", a)
        )
        
        # Sample prompts table
        dpg.add_separator(parent=self.sampling_group)
        dpg.add_text(parent=self.sampling_group, default_value="Sample Prompts")
        
        # Create sample table
        self.sample_table = dpg.add_table(
            parent=self.sampling_group,
            header_row=True,
            resizable=True,
            policy=dpg.mvTable_SizingStretchProp,
            borders_innerH=True,
            borders_outerH=True,
            borders_innerV=True,
            borders_outerV=True
        )
        
        # Add table columns
        dpg.add_table_column(parent=self.sample_table, label="Name")
        dpg.add_table_column(parent=self.sample_table, label="Prompt")
        dpg.add_table_column(parent=self.sample_table, label="Negative Prompt")
        dpg.add_table_column(parent=self.sample_table, label="Actions")
        
        # Add sample row
        with dpg.table_row(parent=self.sample_table):
            dpg.add_text(default_value="Sample 1")
            dpg.add_text(default_value="a photo of [concept]")
            dpg.add_text(default_value="blurry, bad quality")
            with dpg.group(horizontal=True):
                dpg.add_button(label="Edit", callback=lambda: None)
                dpg.add_button(label="Delete", callback=lambda: None)
        
        # Sample controls
        sample_controls = dpg.add_group(parent=self.sampling_group, horizontal=True)
        dpg.add_button(parent=sample_controls, label="Add Sample", callback=self.add_sample)
    
    def sample_now(self):
        """Run sampling now"""
        print("Would run sampling")
        dpg.set_value(self.status_text, "Sampling...")
    
    def open_sampler(self):
        """Open the standalone sampler"""
        print("Would open standalone sampler")
    
    def add_sample(self):
        """Add a new sample prompt"""
        print("Would add a new sample")
    
    def create_cloud_tab(self):
        """Create cloud configuration tab"""
        # Create main tab
        self.cloud_tab = dpg.add_tab(parent=self.tab_bar, label="Cloud")
        
        # Create layout group
        self.cloud_group = dpg.add_group(parent=self.cloud_tab)
        
        # Title
        dpg.add_text(parent=self.cloud_group, default_value="Cloud Training Settings", color=[255, 255, 255, 255])
        dpg.add_separator(parent=self.cloud_group)
        
        # Enable cloud training
        enable_group = dpg.add_group(parent=self.cloud_group, horizontal=True)
        
        self.cloud_enabled_check = dpg.add_checkbox(
            parent=enable_group,
            label="Enable Cloud Training",
            default_value=self.train_config.cloud.enabled if hasattr(self.train_config, "cloud") and hasattr(self.train_config.cloud, "enabled") else False,
            callback=lambda s, a: setattr(self.train_config.cloud, "enabled", a) if hasattr(self.train_config, "cloud") else None
        )
        
        # Cloud type
        type_group = dpg.add_group(parent=self.cloud_group, horizontal=True)
        dpg.add_text(parent=type_group, default_value="Cloud Type:")
        
        cloud_types = ["linux", "runpod"]
        self.cloud_type_combo = dpg.add_combo(
            parent=type_group,
            items=cloud_types,
            default_value="linux",
            callback=lambda s, a: setattr(self.train_config.cloud, "type", a) if hasattr(self.train_config, "cloud") else None
        )
        
        # SSH settings
        dpg.add_separator(parent=self.cloud_group)
        dpg.add_text(parent=self.cloud_group, default_value="SSH Connection Settings")
        
        # Host
        host_group = dpg.add_group(parent=self.cloud_group, horizontal=True)
        dpg.add_text(parent=host_group, default_value="SSH Host:")
        
        self.ssh_host_input = dpg.add_input_text(
            parent=host_group,
            default_value=self.train_config.cloud.ssh_host if hasattr(self.train_config, "cloud") and hasattr(self.train_config.cloud, "ssh_host") else "",
            callback=lambda s, a: setattr(self.train_config.cloud, "ssh_host", a) if hasattr(self.train_config, "cloud") else None,
            width=-100
        )
        
        # Port
        port_group = dpg.add_group(parent=self.cloud_group, horizontal=True)
        dpg.add_text(parent=port_group, default_value="SSH Port:")
        
        self.ssh_port_input = dpg.add_input_int(
            parent=port_group,
            default_value=self.train_config.cloud.ssh_port if hasattr(self.train_config, "cloud") and hasattr(self.train_config.cloud, "ssh_port") else 22,
            callback=lambda s, a: setattr(self.train_config.cloud, "ssh_port", a) if hasattr(self.train_config, "cloud") else None,
            min_value=1,
            max_value=65535
        )
        
        # User
        user_group = dpg.add_group(parent=self.cloud_group, horizontal=True)
        dpg.add_text(parent=user_group, default_value="SSH User:")
        
        self.ssh_user_input = dpg.add_input_text(
            parent=user_group,
            default_value=self.train_config.cloud.ssh_user if hasattr(self.train_config, "cloud") and hasattr(self.train_config.cloud, "ssh_user") else "",
            callback=lambda s, a: setattr(self.train_config.cloud, "ssh_user", a) if hasattr(self.train_config, "cloud") else None,
            width=-100
        )
        
        # Key
        key_group = dpg.add_group(parent=self.cloud_group, horizontal=True)
        dpg.add_text(parent=key_group, default_value="SSH Key:")
        
        self.ssh_key_input = dpg.add_input_text(
            parent=key_group,
            default_value=self.train_config.cloud.ssh_key if hasattr(self.train_config, "cloud") and hasattr(self.train_config.cloud, "ssh_key") else "",
            callback=lambda s, a: setattr(self.train_config.cloud, "ssh_key", a) if hasattr(self.train_config, "cloud") else None,
            width=-100
        )
        
        self.ssh_key_browse = dpg.add_button(
            parent=key_group,
            label="...",
            callback=lambda: self.browse_file("ssh_key_input")
        )
        
        # Password
        password_group = dpg.add_group(parent=self.cloud_group, horizontal=True)
        dpg.add_text(parent=password_group, default_value="SSH Password:")
        
        self.ssh_password_input = dpg.add_input_text(
            parent=password_group,
            default_value="",
            password=True,
            callback=lambda s, a: setattr(self.train_config.secrets.cloud, "ssh_password", a) if hasattr(self.train_config, "secrets") and hasattr(self.train_config.secrets, "cloud") else None,
            width=-100
        )
        
        # Remote path
        remote_path_group = dpg.add_group(parent=self.cloud_group, horizontal=True)
        dpg.add_text(parent=remote_path_group, default_value="Remote Path:")
        
        self.remote_path_input = dpg.add_input_text(
            parent=remote_path_group,
            default_value=self.train_config.cloud.remote_path if hasattr(self.train_config, "cloud") and hasattr(self.train_config.cloud, "remote_path") else "",
            callback=lambda s, a: setattr(self.train_config.cloud, "remote_path", a) if hasattr(self.train_config, "cloud") else None,
            width=-100
        )
        
        # Connection test
        test_group = dpg.add_group(parent=self.cloud_group, horizontal=True)
        dpg.add_button(parent=test_group, label="Test Connection", callback=self.test_connection)
        self.connection_status = dpg.add_text(parent=test_group, default_value="")
    
    def test_connection(self):
        """Test cloud connection"""
        print("Would test cloud connection")
        dpg.set_value(self.connection_status, "Testing...")
    
    # General helper methods
    
    def set_config_enum(self, attr_name, enum_class, value):
        """Set an enum value in the config"""
        try:
            enum_value = enum_class[value]
            setattr(self.train_config, attr_name, enum_value)
        except Exception as e:
            print(f"Error setting {attr_name} to {value}: {e}")
    
    def browse_file(self, target_input, is_directory=False, is_save=False):
        """Show a file browser and update the target input"""
        # This is a simple version without a real file dialog
        # In a real implementation, this would use a proper file dialog
        print(f"Would show file browser for {target_input} (dir={is_directory}, save={is_save})")
    
    def on_model_type_changed(self, sender, app_data):
        """Handle model type change"""
        try:
            self.train_config.model_type = ModelType[app_data]
            self.update_model_specific_ui()
            
            # Update LoRA presets if on LoRA tab
            if hasattr(self, "layer_preset_combo"):
                self.current_presets = self.get_lora_presets(self.train_config.model_type)
                preset_list = list(self.current_presets.keys()) + ["custom"]
                dpg.configure_item(self.layer_preset_combo, items=preset_list)
                if preset_list:
                    dpg.set_value(self.layer_preset_combo, preset_list[0])
                    self.on_layer_preset_changed(None, preset_list[0])
            
            # Update LyCORIS UI if on LoRA tab and using a LyCORIS type
            if hasattr(self, "lycoris_group"):
                self.update_lora_ui_visibility()
        except Exception as e:
            print(f"Error changing model type: {e}")
    
    def on_training_method_changed(self, sender, app_data):
        """Handle training method change"""
        try:
            self.train_config.training_method = TrainingMethod[app_data]
        except Exception as e:
            print(f"Error changing training method: {e}")
    
    def on_tab_changed(self, sender, app_data):
        """Handle tab change"""
        print(f"Tab changed to {app_data}")
    
    def start_training(self):
        """Start the training process"""
        print("Would start training process")
        dpg.set_value(self.status_text, "Training...")
        
        # In a real implementation, this would start the training process
        # Possibly in a separate thread
    
    def run(self):
        """Run the application"""
        while dpg.is_dearpygui_running():
            # Update logic here
            dpg.render_dearpygui_frame()
        
        # Clean up
        dpg.destroy_context()

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