onetrainer_full.py

#!/usr/bin/env python3
"""
Complete OneTrainer DPG implementation with all tabs and functionality.
"""

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

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

# Define basic 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_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"

# Default layer presets
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"],
    "full": [".*"]
}

HIDREAM_PRESETS = {
    "default": ["to_q", "to_k", "to_v", "to_out.0"],
    "extended": ["to_q", "to_k", "to_v", "to_out.0", "ff.net.0", "ff.net.2", "proj_in", "proj_out"],
    "full": [".*"]
}

@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 ModelComponentConfig:
    weight_dtype: DataType = DataType.NONE
    model_name: str = ""

@dataclass
class TrainConfig:
    # 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
    
    # 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_from_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
    warmup_steps: int = 0
    weight_decay: float = 0.01
    adam_beta1: float = 0.9
    adam_beta2: float = 0.999
    adam_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
    unet: ModelComponentConfig = field(default_factory=ModelComponentConfig)
    text_encoder: ModelComponentConfig = field(default_factory=ModelComponentConfig)
    text_encoder_2: ModelComponentConfig = field(default_factory=ModelComponentConfig)
    vae: ModelComponentConfig = field(default_factory=ModelComponentConfig)
    prior: ModelComponentConfig = field(default_factory=ModelComponentConfig)

# Sample concepts for demonstration
SAMPLE_CONCEPTS = [
    {"name": "Portrait", "images_count": 24, "weight": 1.0},
    {"name": "Landscape", "images_count": 32, "weight": 1.0},
    {"name": "Concept3", "images_count": 18, "weight": 0.8}
]

# Sample prompts for demonstration
SAMPLE_PROMPTS = [
    {"name": "Sample 1", "prompt": "a photo of [concept]", "negative_prompt": "blurry, bad quality"},
    {"name": "Sample 2", "prompt": "a detailed painting of [concept]", "negative_prompt": "ugly, deformed"}
]


class OneTrainerFull:
    """Complete OneTrainer DPG Implementation with all tabs"""
    
    def __init__(self):
        # Initialize DPG
        dpg.create_context()
        dpg.create_viewport(title="OneTrainer", width=1000, height=700)
        
        # Initialize state
        self.train_config = TrainConfig()
        self.current_presets = DEFAULT_PRESETS
        
        # Setup UI
        self.setup_ui()
        
        # Setup DPG
        dpg.setup_dearpygui()
        dpg.show_viewport()
    
    def setup_ui(self):
        """Set up the user interface"""
        # Create main window
        self.window = dpg.add_window(label="OneTrainer", width=1000, height=700)
        
        # Create main layout
        self.main_layout = dpg.add_group(parent=self.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.model_tab = dpg.add_tab(parent=self.tab_bar, label="Model")
        self.lora_tab = dpg.add_tab(parent=self.tab_bar, label="LoRA")
        self.concept_tab = dpg.add_tab(parent=self.tab_bar, label="Concepts")
        self.training_tab = dpg.add_tab(parent=self.tab_bar, label="Training")
        self.sampling_tab = dpg.add_tab(parent=self.tab_bar, label="Sampling")
        self.cloud_tab = dpg.add_tab(parent=self.tab_bar, label="Cloud")
        
        # Create tab content
        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 model type and training method selection"""
        self.top_bar = dpg.add_group(parent=self.main_layout, horizontal=True)
        
        dpg.add_text(parent=self.top_bar, default_value="OneTrainer")
        dpg.add_spacer(parent=self.top_bar, width=20)
        
        # Model type selector
        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,
            width=150
        )
        
        dpg.add_spacer(parent=self.top_bar, width=10)
        
        # Training method selector
        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,
            width=150
        )
        
        dpg.add_spacer(parent=self.top_bar, width=10)
        
        # Preset selector
        dpg.add_text(parent=self.top_bar, default_value="Preset:")
        
        self.preset_files = self.get_preset_files()
        preset_options = ["Custom"] + [os.path.basename(p) for p in self.preset_files]
        self.preset_combo = dpg.add_combo(
            parent=self.top_bar,
            items=preset_options,
            default_value="Custom",
            callback=self.on_preset_selected,
            width=150
        )
        
        # Save preset button
        self.save_preset_button = dpg.add_button(
            parent=self.top_bar,
            label="Save",
            callback=self.save_preset
        )
    
    def create_bottom_bar(self):
        """Create bottom bar with status and controls"""
        self.bottom_bar = dpg.add_group(parent=self.main_layout, horizontal=True)
        
        # Progress bar
        self.progress_group = dpg.add_group(parent=self.bottom_bar, horizontal=False, width=300)
        
        dpg.add_text(parent=self.progress_group, default_value="Progress:")
        self.progress_bar = dpg.add_progress_bar(parent=self.progress_group, default_value=0, width=-1)
        
        # Status text
        dpg.add_text(parent=self.bottom_bar, default_value="Status:")
        self.status_text = dpg.add_text(parent=self.bottom_bar, default_value="Ready")
        
        dpg.add_spacer(parent=self.bottom_bar, width=20)
        
        # Training controls
        self.train_button = dpg.add_button(
            parent=self.bottom_bar,
            label="Start Training",
            callback=self.start_training
        )
        
        self.export_button = dpg.add_button(
            parent=self.bottom_bar,
            label="Export",
            callback=self.export_model
        )
    
    def create_model_tab(self):
        """Create model configuration tab"""
        self.model_group = dpg.add_group(parent=self.model_tab)
        
        # Basic model settings
        dpg.add_text(parent=self.model_group, default_value="Model Settings")
        dpg.add_separator(parent=self.model_group)
        
        # 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
        )
        
        dpg.add_button(
            parent=base_model_group,
            label="...",
            callback=lambda: self.browse_file("base_model_input")
        )
        
        # Weight data type
        weight_type_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=weight_type_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_type_group,
            items=dtype_options,
            default_value=self.train_config.weight_dtype.name,
            callback=lambda s, a: setattr(self.train_config, "weight_dtype", DataType[a]),
            width=150
        )
        
        # Model component settings
        dpg.add_separator(parent=self.model_group)
        dpg.add_text(parent=self.model_group, default_value="Model Component Settings")
        
        # Container for model-specific settings
        self.model_specific_group = dpg.add_group(parent=self.model_group)
        self.update_model_specific_ui()
        
        # Output settings
        dpg.add_separator(parent=self.model_group)
        dpg.add_text(parent=self.model_group, default_value="Output Settings")
        
        # Output path
        output_path_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=output_path_group, default_value="Output Path:")
        
        self.output_path_input = dpg.add_input_text(
            parent=output_path_group,
            default_value=self.train_config.output_model_destination,
            callback=lambda s, a: setattr(self.train_config, "output_model_destination", a),
            width=-100
        )
        
        dpg.add_button(
            parent=output_path_group,
            label="...",
            callback=lambda: self.browse_file("output_path_input", is_save=True)
        )
        
        # Output data type
        output_type_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=output_type_group, default_value="Output Data Type:")
        
        self.output_dtype_combo = dpg.add_combo(
            parent=output_type_group,
            items=dtype_options,
            default_value=self.train_config.output_dtype.name,
            callback=lambda s, a: setattr(self.train_config, "output_dtype", DataType[a]),
            width=150
        )
        
        # System settings
        dpg.add_separator(parent=self.model_group)
        dpg.add_text(parent=self.model_group, default_value="System Settings")
        
        # Workspace directory
        workspace_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=workspace_group, default_value="Workspace Directory:")
        
        self.workspace_input = dpg.add_input_text(
            parent=workspace_group,
            default_value=self.train_config.workspace_dir,
            callback=lambda s, a: setattr(self.train_config, "workspace_dir", a),
            width=-100
        )
        
        dpg.add_button(
            parent=workspace_group,
            label="...",
            callback=lambda: self.browse_file("workspace_input", is_directory=True)
        )
        
        # Cache directory
        cache_group = dpg.add_group(parent=self.model_group, horizontal=True)
        dpg.add_text(parent=cache_group, default_value="Cache Directory:")
        
        self.cache_input = dpg.add_input_text(
            parent=cache_group,
            default_value=self.train_config.cache_dir,
            callback=lambda s, a: setattr(self.train_config, "cache_dir", a),
            width=-100
        )
        
        dpg.add_button(
            parent=cache_group,
            label="...",
            callback=lambda: self.browse_file("cache_input", is_directory=True)
        )
        
        # Options
        options_group = dpg.add_group(parent=self.model_group)
        
        self.debug_check = dpg.add_checkbox(
            parent=options_group,
            label="Debug Mode",
            default_value=self.train_config.debug_mode,
            callback=lambda s, a: setattr(self.train_config, "debug_mode", a)
        )
        
        self.continue_check = dpg.add_checkbox(
            parent=options_group,
            label="Continue From Backup",
            default_value=self.train_config.continue_from_backup,
            callback=lambda s, a: setattr(self.train_config, "continue_from_backup", a)
        )
        
        self.cache_only_check = dpg.add_checkbox(
            parent=options_group,
            label="Only Cache (No Training)",
            default_value=self.train_config.only_cache,
            callback=lambda s, a: setattr(self.train_config, "only_cache", a)
        )
    
    def create_lora_tab(self):
        """Create LoRA configuration tab with LyCORIS support"""
        self.lora_group = dpg.add_group(parent=self.lora_tab)
        
        # Title
        dpg.add_text(parent=self.lora_group, default_value="LoRA / LyCORIS Settings")
        dpg.add_separator(parent=self.lora_group)
        
        # PEFT type selector
        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,
            width=150
        )
        
        # Add info tooltip
        with dpg.tooltip(self.peft_type_combo):
            dpg.add_text(
                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"
                )
            )
        
        # Base model
        base_model_group = dpg.add_group(parent=self.lora_group, horizontal=True)
        dpg.add_text(parent=base_model_group, default_value="LoRA Base Model:")
        
        self.lora_model_input = dpg.add_input_text(
            parent=base_model_group,
            default_value=self.train_config.lora_model_name,
            callback=lambda s, a: setattr(self.train_config, "lora_model_name", a),
            width=-100
        )
        
        dpg.add_button(
            parent=base_model_group,
            label="...",
            callback=lambda: self.browse_file("lora_model_input")
        )
        
        # Common settings
        common_group = dpg.add_group(parent=self.lora_group)
        
        # LoRA rank
        rank_group = dpg.add_group(parent=common_group, horizontal=True)
        dpg.add_text(parent=rank_group, default_value="LoRA Rank:")
        
        self.lora_rank_input = dpg.add_input_int(
            parent=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,
            width=100
        )
        
        # Add info tooltip
        with dpg.tooltip(self.lora_rank_input):
            dpg.add_text(
                default_value=(
                    "The rank of the update matrices. "
                    "Higher values = larger model size but potentially higher quality."
                )
            )
        
        # LoRA alpha
        alpha_group = dpg.add_group(parent=common_group, horizontal=True)
        dpg.add_text(parent=alpha_group, default_value="LoRA Alpha:")
        
        self.lora_alpha_input = dpg.add_input_float(
            parent=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,
            width=100
        )
        
        # Add info tooltip
        with dpg.tooltip(self.lora_alpha_input):
            dpg.add_text(
                default_value=(
                    "Scaling factor for the LoRA adapter. "
                    "Typically set to same value as rank. Common values: 16, 32, 64."
                )
            )
        
        # Dropout
        dropout_group = dpg.add_group(parent=common_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
        )
        
        # Add info tooltip
        with dpg.tooltip(self.dropout_slider):
            dpg.add_text(
                default_value=(
                    "Dropout probability for LoRA layers (0.0 = disabled).\n"
                    "Higher values help prevent overfitting but require longer training."
                )
            )
        
        # Bundle embeddings
        bundle_group = dpg.add_group(parent=common_group, horizontal=True)
        
        self.bundle_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)
        )
        
        # Weight data type
        dtype_group = dpg.add_group(parent=common_group, horizontal=True)
        dpg.add_text(parent=dtype_group, default_value="Weight Data Type:")
        
        lora_dtype_options = [dt.name for dt in [DataType.FLOAT_32, DataType.BFLOAT_16]]
        self.lora_dtype_combo = dpg.add_combo(
            parent=dtype_group,
            items=lora_dtype_options,
            default_value=self.train_config.lora_weight_dtype.name,
            callback=lambda s, a: setattr(self.train_config, "lora_weight_dtype", DataType[a]),
            width=150
        )
        
        # 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,
            width=100
        )
        
        # Add info tooltip
        with dpg.tooltip(self.lycoris_factor_input):
            dpg.add_text(
                default_value=(
                    "Compression factor for LoKr/LoHa. Lower = less parameters.\n"
                    "Used with rank to determine total parameters."
                )
            )
        
        # DyLoRA full matrix
        full_matrix_group = dpg.add_group(parent=self.lycoris_group, horizontal=True)
        
        self.full_matrix_check = dpg.add_checkbox(
            parent=full_matrix_group,
            label="Use Full Matrix (DyLoRA)",
            default_value=self.train_config.lycoris_full_matrix,
            callback=lambda s, a: setattr(self.train_config, "lycoris_full_matrix", a)
        )
        
        # Bypass mode
        bypass_group = dpg.add_group(parent=self.lycoris_group, horizontal=True)
        
        self.bypass_check = dpg.add_checkbox(
            parent=bypass_group,
            label="Bypass Mode (for HiDream + LoKr)",
            default_value=self.train_config.lycoris_bypass_mode,
            callback=lambda s, a: setattr(self.train_config, "lycoris_bypass_mode", a)
        )
        
        # DoRA settings
        dpg.add_separator(parent=self.lora_group)
        dpg.add_text(parent=self.lora_group, default_value="DoRA Settings (LoRA only)")
        
        self.dora_group = dpg.add_group(parent=self.lora_group)
        
        # Enable DoRA
        dora_enable_group = dpg.add_group(parent=self.dora_group, horizontal=True)
        
        self.dora_check = dpg.add_checkbox(
            parent=dora_enable_group,
            label="Enable Weight Decomposition (DoRA)",
            default_value=self.train_config.lora_decompose,
            callback=self.on_dora_toggle
        )
        
        # DoRA 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)
        
        self.norm_check = dpg.add_checkbox(
            parent=norm_group,
            label="Use Norm Epsilon",
            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_group = dpg.add_group(parent=self.dora_advanced_group, horizontal=True)
        
        self.output_check = dpg.add_checkbox(
            parent=output_group,
            label="Apply on Output Axis",
            default_value=self.train_config.lora_decompose_output_axis,
            callback=lambda s, a: setattr(self.train_config, "lora_decompose_output_axis", a)
        )
        
        # Layer settings
        dpg.add_separator(parent=self.lora_group)
        dpg.add_text(parent=self.lora_group, default_value="Layer Settings")
        
        # Layer preset
        preset_group = dpg.add_group(parent=self.lora_group, horizontal=True)
        dpg.add_text(parent=preset_group, default_value="Layer Preset:")
        
        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,
            width=150
        )
        
        # Custom layers
        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
        )
        
        # Update visibility
        self.update_lora_ui_visibility()
    
    def create_concept_tab(self):
        """Create concept configuration tab"""
        self.concept_group = dpg.add_group(parent=self.concept_tab)
        
        # Title
        dpg.add_text(parent=self.concept_group, default_value="Training Concepts")
        dpg.add_separator(parent=self.concept_group)
        
        # Create concepts table
        concept_table_group = dpg.add_group(parent=self.concept_group)
        dpg.add_text(parent=concept_table_group, default_value="Concepts:")
        
        self.concept_table = dpg.add_table(
            parent=concept_table_group,
            header_row=True,
            resizable=True,
            borders_innerH=True,
            borders_outerH=True,
            borders_innerV=True,
            borders_outerV=True,
            policy=dpg.mvTable_SizingStretchProp
        )
        
        # 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")
        
        # Populate with sample data
        self.populate_concept_table()
        
        # 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")
        
        # Max 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=self.train_config.max_resolution,
            callback=lambda s, a: setattr(self.train_config, "max_resolution", a),
            min_value=128,
            max_value=2048,
            step=64,
            width=100
        )
        
        # Data settings
        data_group = dpg.add_group(parent=self.concept_group)
        
        self.aspect_ratio_check = dpg.add_checkbox(
            parent=data_group,
            label="Aspect Ratio Bucketing",
            default_value=self.train_config.aspect_ratio_bucketing,
            callback=lambda s, a: setattr(self.train_config, "aspect_ratio_bucketing", a)
        )
        
        self.latent_caching_check = dpg.add_checkbox(
            parent=data_group,
            label="Latent Caching",
            default_value=self.train_config.latent_caching,
            callback=lambda s, a: setattr(self.train_config, "latent_caching", a)
        )
        
        self.clear_cache_check = dpg.add_checkbox(
            parent=data_group,
            label="Clear Cache Before Training",
            default_value=self.train_config.clear_cache_before_training,
            callback=lambda s, a: setattr(self.train_config, "clear_cache_before_training", a)
        )
    
    def create_training_tab(self):
        """Create training configuration tab"""
        self.training_group = dpg.add_group(parent=self.training_tab)
        
        # Title
        dpg.add_text(parent=self.training_group, default_value="Training Settings")
        dpg.add_separator(parent=self.training_group)
        
        # Basic settings
        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,
            width=100
        )
        
        # 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,
            width=100
        )
        
        # 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,
            width=100
        )
        
        # Mixed precision
        mixed_prec_group = dpg.add_group(parent=basic_group, horizontal=True)
        
        self.mixed_prec_check = dpg.add_checkbox(
            parent=mixed_prec_group,
            label="Mixed Precision",
            default_value=self.train_config.mixed_precision,
            callback=lambda s, a: setattr(self.train_config, "mixed_precision", a)
        )
        
        # 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_value_group = dpg.add_group(parent=lr_group, horizontal=True)
        dpg.add_text(parent=lr_value_group, default_value="Learning Rate:")
        
        self.lr_input = dpg.add_input_float(
            parent=lr_value_group,
            default_value=self.train_config.learning_rate,
            callback=lambda s, a: setattr(self.train_config, "learning_rate", a),
            format="%.8f",
            width=150
        )
        
        # Scheduler
        scheduler_group = dpg.add_group(parent=lr_group, horizontal=True)
        dpg.add_text(parent=scheduler_group, default_value="LR Scheduler:")
        
        scheduler_options = ["constant", "linear", "cosine", "cosine_with_restarts", "polynomial"]
        self.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),
            width=150
        )
        
        # 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=self.train_config.warmup_steps,
            callback=lambda s, a: setattr(self.train_config, "warmup_steps", a),
            min_value=0,
            max_value=10000,
            width=100
        )
        
        # 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 type
        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),
            width=150
        )
        
        # 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=self.train_config.weight_decay,
            callback=lambda s, a: setattr(self.train_config, "weight_decay", a),
            format="%.4f",
            width=100
        )
        
        # Adam parameters
        adam_betas_group = dpg.add_group(parent=opt_group, horizontal=True)
        dpg.add_text(parent=adam_betas_group, default_value="Adam Betas:")
        
        self.beta1_input = dpg.add_input_float(
            parent=adam_betas_group,
            default_value=self.train_config.adam_beta1,
            callback=lambda s, a: setattr(self.train_config, "adam_beta1", a),
            format="%.3f",
            width=100
        )
        
        self.beta2_input = dpg.add_input_float(
            parent=adam_betas_group,
            default_value=self.train_config.adam_beta2,
            callback=lambda s, a: setattr(self.train_config, "adam_beta2", a),
            format="%.6f",
            width=100
        )
        
        # Adam epsilon
        adam_eps_group = dpg.add_group(parent=opt_group, horizontal=True)
        dpg.add_text(parent=adam_eps_group, default_value="Adam Epsilon:")
        
        self.epsilon_input = dpg.add_input_float(
            parent=adam_eps_group,
            default_value=self.train_config.adam_epsilon,
            callback=lambda s, a: setattr(self.train_config, "adam_epsilon", a),
            format="%.8f",
            width=150
        )
        
        # 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)
        
        # Use EMA
        use_ema_group = dpg.add_group(parent=ema_group, horizontal=True)
        
        self.use_ema_check = dpg.add_checkbox(
            parent=use_ema_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",
            width=150
        )
        
        # Save/backup settings
        dpg.add_separator(parent=self.training_group)
        dpg.add_text(parent=self.training_group, default_value="Save & Backup Settings")
        
        save_group = dpg.add_group(parent=self.training_group)
        
        # Save every
        save_every_group = dpg.add_group(parent=save_group, horizontal=True)
        dpg.add_text(parent=save_every_group, default_value="Save Every:")
        
        self.save_every_input = dpg.add_input_int(
            parent=save_every_group,
            default_value=self.train_config.save_every,
            callback=lambda s, a: setattr(self.train_config, "save_every", a),
            width=100
        )
        
        time_units = [tu.name for tu in TimeUnit]
        self.save_unit_combo = dpg.add_combo(
            parent=save_every_group,
            items=time_units,
            default_value=self.train_config.save_every_unit.name,
            callback=lambda s, a: setattr(self.train_config, "save_every_unit", TimeUnit[a]),
            width=100
        )
        
        # Backup every
        backup_every_group = dpg.add_group(parent=save_group, horizontal=True)
        dpg.add_text(parent=backup_every_group, default_value="Backup Every:")
        
        self.backup_every_input = dpg.add_input_int(
            parent=backup_every_group,
            default_value=self.train_config.backup_after,
            callback=lambda s, a: setattr(self.train_config, "backup_after", a),
            width=100
        )
        
        self.backup_unit_combo = dpg.add_combo(
            parent=backup_every_group,
            items=time_units,
            default_value=self.train_config.backup_after_unit.name,
            callback=lambda s, a: setattr(self.train_config, "backup_after_unit", TimeUnit[a]),
            width=100
        )
        
        # Backup options
        backup_options_group = dpg.add_group(parent=save_group)
        
        self.rolling_backup_check = dpg.add_checkbox(
            parent=backup_options_group,
            label="Rolling Backup",
            default_value=self.train_config.rolling_backup,
            callback=lambda s, a: setattr(self.train_config, "rolling_backup", a)
        )
        
        self.backup_before_save_check = dpg.add_checkbox(
            parent=backup_options_group,
            label="Backup Before Save",
            default_value=self.train_config.backup_before_save,
            callback=lambda s, a: setattr(self.train_config, "backup_before_save", a)
        )
    
    def create_sampling_tab(self):
        """Create sampling configuration tab"""
        self.sampling_group = dpg.add_group(parent=self.sampling_tab)
        
        # Title
        dpg.add_text(parent=self.sampling_group, default_value="Sampling Settings")
        dpg.add_separator(parent=self.sampling_group)
        
        # Sampling schedule
        dpg.add_text(parent=self.sampling_group, default_value="Sampling Schedule")
        
        # Sample 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),
            width=100
        )
        
        time_units = [tu.name for tu in TimeUnit]
        self.sample_unit_combo = dpg.add_combo(
            parent=after_group,
            items=time_units,
            default_value=self.train_config.sample_after_unit.name,
            callback=lambda s, a: setattr(self.train_config, "sample_after_unit", TimeUnit[a]),
            width=100
        )
        
        # Skip first
        skip_group = dpg.add_group(parent=self.sampling_group, horizontal=True)
        dpg.add_text(parent=skip_group, default_value="Skip First:")
        
        self.skip_first_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),
            width=100
        )
        
        # Image format
        format_group = dpg.add_group(parent=self.sampling_group, horizontal=True)
        dpg.add_text(parent=format_group, default_value="Image Format:")
        
        format_options = [fmt.name for fmt in ImageFormat]
        self.format_combo = dpg.add_combo(
            parent=format_group,
            items=format_options,
            default_value=self.train_config.sample_image_format.name,
            callback=lambda s, a: setattr(self.train_config, "sample_image_format", ImageFormat[a]),
            width=100
        )
        
        # Sample options
        options_group = dpg.add_group(parent=self.sampling_group)
        
        self.non_ema_check = dpg.add_checkbox(
            parent=options_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.samples_to_tb_check = dpg.add_checkbox(
            parent=options_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 controls
        controls_group = dpg.add_group(parent=self.sampling_group, horizontal=True)
        
        dpg.add_button(
            parent=controls_group,
            label="Sample Now",
            callback=self.sample_now
        )
        
        dpg.add_button(
            parent=controls_group,
            label="Open Sampler",
            callback=self.open_sampler
        )
        
        # Sampling parameters
        dpg.add_separator(parent=self.sampling_group)
        dpg.add_text(parent=self.sampling_group, default_value="Sampling Parameters")
        
        params_group = dpg.add_group(parent=self.sampling_group)
        
        # CFG scale
        cfg_group = dpg.add_group(parent=params_group, horizontal=True)
        dpg.add_text(parent=cfg_group, default_value="CFG Scale:")
        
        self.cfg_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=params_group, horizontal=True)
        dpg.add_text(parent=steps_group, default_value="Steps:")
        
        self.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
        sched_group = dpg.add_group(parent=params_group, horizontal=True)
        dpg.add_text(parent=sched_group, default_value="Scheduler:")
        
        scheduler_options = ["ddim", "ddpm", "k_lms", "pndm", "euler", "euler_a", "dpmpp_2m"]
        self.sched_combo = dpg.add_combo(
            parent=sched_group,
            items=scheduler_options,
            default_value=self.train_config.scheduler,
            callback=lambda s, a: setattr(self.train_config, "scheduler", a),
            width=150
        )
        
        # Sample prompts
        dpg.add_separator(parent=self.sampling_group)
        dpg.add_text(parent=self.sampling_group, default_value="Sample Prompts")
        
        # Prompt table
        self.sample_table = dpg.add_table(
            parent=self.sampling_group,
            header_row=True,
            resizable=True,
            borders_innerH=True,
            borders_outerH=True,
            borders_innerV=True,
            borders_outerV=True,
            policy=dpg.mvTable_SizingStretchProp
        )
        
        # 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")
        
        # Populate with sample data
        self.populate_sample_table()
        
        # Add prompt button
        prompt_controls = dpg.add_group(parent=self.sampling_group, horizontal=True)
        
        dpg.add_button(
            parent=prompt_controls,
            label="Add Sample Prompt",
            callback=self.add_sample_prompt
        )
    
    def create_cloud_tab(self):
        """Create cloud configuration tab"""
        self.cloud_group = dpg.add_group(parent=self.cloud_tab)
        
        # Title
        dpg.add_text(parent=self.cloud_group, default_value="Cloud Training Settings")
        dpg.add_separator(parent=self.cloud_group)
        
        # Enable cloud
        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,
            callback=lambda s, a: setattr(self.train_config.cloud, "enabled", a)
        )
        
        # 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),
            width=150
        )
        
        # SSH connection 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,
            callback=lambda s, a: setattr(self.train_config.cloud, "ssh_host", a),
            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,
            callback=lambda s, a: setattr(self.train_config.cloud, "ssh_port", a),
            width=100
        )
        
        # 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,
            callback=lambda s, a: setattr(self.train_config.cloud, "ssh_user", a),
            width=-100
        )
        
        # SSH 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,
            callback=lambda s, a: setattr(self.train_config.cloud, "ssh_key", a),
            width=-100
        )
        
        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),
            width=-100
        )
        
        # Remote path
        path_group = dpg.add_group(parent=self.cloud_group, horizontal=True)
        dpg.add_text(parent=path_group, default_value="Remote Path:")
        
        self.remote_path_input = dpg.add_input_text(
            parent=path_group,
            default_value=self.train_config.cloud.remote_path,
            callback=lambda s, a: setattr(self.train_config.cloud, "remote_path", a),
            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=""
        )
    
    # UI update methods
    
    def update_model_specific_ui(self):
        """Update model-specific UI based on selected model type"""
        # Clear previous UI
        if self.model_specific_group:
            # Get all children
            children = dpg.get_item_children(self.model_specific_group, slot=1)
            
            # Delete all children
            if children:
                for child in children:
                    dpg.delete_item(child)
        
        # Add model-specific UI based on model type
        model_type = self.train_config.model_type
        
        # For demonstration, just show text describing what would be here
        dpg.add_text(
            parent=self.model_specific_group,
            default_value=f"Model-specific settings for {model_type.name}"
        )
        
        # In a real implementation, this would add appropriate UI elements
        # based on the selected model type (SD, SDXL, etc.)
    
    def update_lora_ui_visibility(self):
        """Update LoRA tab UI visibility based on settings"""
        # LyCORIS settings visibility
        is_lycoris = self.train_config.peft_type in [
            PeftType.LOHA, PeftType.LOKR, PeftType.DIA, 
            PeftType.IA3, PeftType.DYLORA
        ]
        dpg.configure_item(self.lycoris_group, show=is_lycoris)
        
        # Factor setting only for LoHa/LoKr
        is_factor_type = self.train_config.peft_type in [PeftType.LOHA, PeftType.LOKR]
        dpg.configure_item(self.lycoris_factor_input, show=is_factor_type)
        for item in dpg.get_item_children(dpg.get_item_parent(self.lycoris_factor_input), slot=1):
            if dpg.get_item_type(item) == "mvText":
                dpg.configure_item(item, show=is_factor_type)
        
        # Full matrix only for DyLoRA
        is_dylora = self.train_config.peft_type == PeftType.DYLORA
        dpg.configure_item(self.full_matrix_check, show=is_dylora)
        dpg.configure_item(dpg.get_item_parent(self.full_matrix_check), show=is_dylora)
        
        # Bypass mode only for HiDream + LoKr
        is_bypass_needed = (
            self.train_config.model_type.is_hi_dream() and 
            self.train_config.peft_type == PeftType.LOKR
        )
        dpg.configure_item(self.bypass_check, show=is_bypass_needed)
        dpg.configure_item(dpg.get_item_parent(self.bypass_check), show=is_bypass_needed)
        
        # DoRA settings only for LoRA
        is_lora = self.train_config.peft_type == PeftType.LORA
        dpg.configure_item(self.dora_group, show=is_lora)
        
        # DoRA advanced settings only when DoRA is enabled
        use_dora = self.train_config.lora_decompose
        dpg.configure_item(self.dora_advanced_group, show=is_lora and use_dora)
    
    def populate_concept_table(self):
        """Populate concept table with sample data"""
        # Clear existing rows (except header)
        for row in dpg.get_item_children(self.concept_table, slot=1):
            dpg.delete_item(row)
        
        # Add sample concepts
        for concept in SAMPLE_CONCEPTS:
            with dpg.table_row(parent=self.concept_table):
                dpg.add_text(default_value=concept["name"])
                dpg.add_text(default_value=f"{concept['images_count']} images")
                dpg.add_text(default_value=str(concept["weight"]))
                
                with dpg.group(horizontal=True):
                    dpg.add_button(label="Edit", callback=lambda: None)
                    dpg.add_button(label="Delete", callback=lambda: None)
    
    def populate_sample_table(self):
        """Populate sample prompt table with sample data"""
        # Clear existing rows (except header)
        for row in dpg.get_item_children(self.sample_table, slot=1):
            dpg.delete_item(row)
        
        # Add sample prompts
        for prompt in SAMPLE_PROMPTS:
            with dpg.table_row(parent=self.sample_table):
                dpg.add_text(default_value=prompt["name"])
                dpg.add_text(default_value=prompt["prompt"])
                dpg.add_text(default_value=prompt["negative_prompt"])
                
                with dpg.group(horizontal=True):
                    dpg.add_button(label="Edit", callback=lambda: None)
                    dpg.add_button(label="Delete", callback=lambda: None)
    
    # Event handlers
    
    def on_model_type_changed(self, sender, app_data):
        """Handle model type change"""
        try:
            # Update model type
            self.train_config.model_type = ModelType[app_data]
            
            # Update UI based on new model type
            self.update_model_specific_ui()
            
            # Update LoRA presets
            self.current_presets = self.get_lora_presets()
            preset_list = list(self.current_presets.keys()) + ["custom"]
            
            # Update preset combo
            if dpg.does_item_exist(self.layer_preset_combo):
                dpg.configure_item(self.layer_preset_combo, items=preset_list)
                
                # Set first preset as default if current is not in list
                if dpg.get_value(self.layer_preset_combo) not in preset_list:
                    dpg.set_value(self.layer_preset_combo, preset_list[0])
                    self.on_layer_preset_changed(None, preset_list[0])
            
            # Update LoRA UI visibility
            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_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 on_layer_preset_changed(self, sender, app_data):
        """Handle layer preset change"""
        try:
            self.train_config.lora_layer_preset = app_data
            
            # If not custom, update layers based on selected 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
                
                # Update input field
                if dpg.does_item_exist(self.layers_input):
                    dpg.set_value(self.layers_input, layers)
        except Exception as e:
            print(f"Error changing layer preset: {e}")
    
    def on_preset_selected(self, sender, app_data):
        """Handle preset selection"""
        if app_data == "Custom":
            # Custom setting, no preset to load
            return
        
        # Find preset file
        for preset_file in self.preset_files:
            if os.path.basename(preset_file) == app_data:
                self.load_preset(preset_file)
                break
    
    # Action methods
    
    def get_lora_presets(self):
        """Get LoRA layer presets based on model type"""
        if self.train_config.model_type.is_hi_dream():
            return HIDREAM_PRESETS
        else:
            return DEFAULT_PRESETS
    
    def get_preset_files(self):
        """Get list of available preset files"""
        preset_dir = "training_presets"
        preset_files = []
        
        if os.path.exists(preset_dir):
            preset_files = [
                os.path.join(preset_dir, f) for f in os.listdir(preset_dir)
                if f.endswith(".json")
            ]
        
        return preset_files
    
    def load_preset(self, preset_path):
        """Load training preset from file"""
        if not os.path.exists(preset_path):
            print(f"Preset file not found: {preset_path}")
            return
        
        try:
            with open(preset_path, "r") as f:
                preset_data = json.load(f)
                
            # In a real implementation, this would properly load
            # the preset data into the train_config object
            print(f"Would load preset: {preset_path}")
            dpg.set_value(self.status_text, f"Loaded preset: {os.path.basename(preset_path)}")
        except Exception as e:
            print(f"Error loading preset: {e}")
    
    def save_preset(self):
        """Save current settings as a preset"""
        # In a real implementation, this would save current settings
        # to a JSON file in the training_presets directory
        print("Would save preset")
        dpg.set_value(self.status_text, "Saved preset")
    
    def browse_file(self, target_input, is_directory=False, is_save=False):
        """Show file browser dialog"""
        # In a real implementation, this would show a file browser
        # and update the target input field
        print(f"Would show file browser for {target_input}")
    
    def add_concept(self):
        """Add a new concept"""
        # In a real implementation, this would show a dialog to add a concept
        print("Would add a new concept")
        dpg.set_value(self.status_text, "Adding concept...")
    
    def refresh_concepts(self):
        """Refresh concepts list"""
        print("Would refresh concepts")
        dpg.set_value(self.status_text, "Refreshing concepts...")
        self.populate_concept_table()
    
    def add_sample_prompt(self):
        """Add a new sample prompt"""
        # In a real implementation, this would show a dialog to add a prompt
        print("Would add a new sample prompt")
        dpg.set_value(self.status_text, "Adding sample prompt...")
    
    def sample_now(self):
        """Run sampling now"""
        print("Would run sampling")
        dpg.set_value(self.status_text, "Sampling...")
    
    def open_sampler(self):
        """Open standalone sampler"""
        print("Would open standalone sampler")
        dpg.set_value(self.status_text, "Opening sampler...")
    
    def start_training(self):
        """Start training process"""
        print("Would start training")
        dpg.set_value(self.status_text, "Training started")
        dpg.set_value(self.progress_bar, 0.1)  # Show some progress
    
    def export_model(self):
        """Export trained model"""
        print("Would export model")
        dpg.set_value(self.status_text, "Exporting model...")
    
    def test_connection(self):
        """Test cloud connection"""
        print("Would test cloud connection")
        dpg.set_value(self.connection_status, "Connection successful")
    
    def run(self):
        """Run the application"""
        while dpg.is_dearpygui_running():
            # Update logic here (would be used for progress updates, etc.)
            dpg.render_dearpygui_frame()
        
        # Cleanup
        dpg.destroy_context()

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