enhanced_onetrainer_ui.py

#!/usr/bin/env python3
"""
Enhanced OneTrainer UI using DearPyGui

This is a more complete version of the OneTrainer UI that incorporates
all major functionality while avoiding segmentation faults and crashes.
"""

import os
import sys
import dearpygui.dearpygui as dpg
from pathlib import Path
import json
import traceback
import glob
import subprocess
import threading
import time
from typing import Dict, Any, List, Optional, Union, Callable, Tuple

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

# Import OneTrainer modules with error handling
try:
    from modules.util.enum.ModelType import ModelType
    from modules.util.enum.TrainingMethod import TrainingMethod
    from modules.util.enum.LossWeight import LossWeight
    from modules.util.enum.Optimizer import Optimizer
    from modules.util.enum.NoiseScheduler import NoiseScheduler
    from modules.util.enum.LearningRateScheduler import LearningRateScheduler
    from modules.util.config.TrainConfig import TrainConfig
    from modules.util.config.ConceptConfig import ConceptConfig
    HAS_MODULES = True
except ImportError as e:
    print(f"Warning: Could not import OneTrainer modules: {e}")
    HAS_MODULES = False

# Set up global config
CONFIG = {}
PRESETS_DIR = "training_presets"
THEMES = {
    "dark_blue": {
        "primary": (41, 83, 154, 255),  # Blue
        "secondary": (63, 119, 183, 255),  # Lighter blue
        "accent": (255, 255, 0, 255),  # Yellow
        "background": (30, 30, 30, 255),  # Dark gray
        "text": (255, 255, 255, 255),  # White
        "text_secondary": (200, 200, 200, 255),  # Light gray
        "disabled": (120, 120, 120, 255)  # Gray
    }
}

class EnhancedOneTrainerUI:
    """
    Enhanced OneTrainer UI using DearPyGui

    This class implements a comprehensive UI with all major tabs and functionality.
    """

    def __init__(self):
        """Initialize the UI"""
        print("Initializing Enhanced OneTrainer UI...")

        # Initialize DearPyGui
        dpg.create_context()
        dpg.create_viewport(title="OneTrainer", width=1280, height=800)
        dpg.setup_dearpygui()

        # Create state
        self.config = {}
        self.training_thread = None
        self.is_training = False
        self.presets = []
        self.preset_paths = {}
        self.progress = {"step": 0, "total_steps": 1000, "epoch": 0, "total_epochs": 10}

        # Load config
        self.load_config()

        # Load presets
        self.load_presets()

        # Create UI
        self.create_ui()

        # Apply theme
        self.apply_theme()
        
    def load_config(self):
        """Load configuration from file"""
        # Try to load config from default location if it exists
        config_file = "train_config.json"
        if os.path.exists(config_file):
            try:
                with open(config_file, "r") as f:
                    self.config = json.load(f)
                    print(f"Loaded config from {config_file}")
            except Exception as e:
                print(f"Error loading config: {e}")
                self.config = {}
        else:
            # Create default config
            self.config = {
                "model_type": "STABLE_DIFFUSION_15",
                "training_method": "LORA",
                "learning_rate": 0.0001,
                "batch_size": 1,
                "max_train_steps": 1000,
                "save_every_n_steps": 100,
                "lora_rank": 32,
                "lora_alpha": 32.0,
                "optimizer": "AdamW8bit",
                "scheduler": "cosine"
            }

    def load_presets(self):
        """Load presets from directory"""
        self.presets = ["-- Select Preset --"]
        self.preset_paths = {}

        if not os.path.exists(PRESETS_DIR):
            os.makedirs(PRESETS_DIR, exist_ok=True)
            return

        try:
            for preset_file in glob.glob(os.path.join(PRESETS_DIR, "*.json")):
                preset_name = os.path.basename(preset_file).replace(".json", "")
                self.presets.append(preset_name)
                self.preset_paths[preset_name] = preset_file
        except Exception as e:
            print(f"Error loading presets: {e}")

    def save_config(self, path=None):
        """Save configuration to file"""
        if path is None:
            path = "train_config.json"

        try:
            with open(path, "w") as f:
                json.dump(self.config, f, indent=2)
                print(f"Saved config to {path}")
        except Exception as e:
            print(f"Error saving config: {e}")

    def apply_theme(self):
        """Apply theme to UI elements"""
        theme = THEMES["dark_blue"]

        with dpg.theme() as global_theme:
            with dpg.theme_component(dpg.mvAll):
                # Background colors
                dpg.add_theme_color(dpg.mvThemeCol_WindowBg, theme["background"])
                dpg.add_theme_color(dpg.mvThemeCol_ChildBg, (25, 25, 25, 255))
                dpg.add_theme_color(dpg.mvThemeCol_PopupBg, (35, 35, 35, 255))
                dpg.add_theme_color(dpg.mvThemeCol_Border, (60, 60, 60, 255))

                # Control colors
                dpg.add_theme_color(dpg.mvThemeCol_Button, theme["primary"])
                dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, theme["secondary"])
                dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, theme["secondary"])

                # Widget background
                dpg.add_theme_color(dpg.mvThemeCol_FrameBg, (45, 45, 45, 255))
                dpg.add_theme_color(dpg.mvThemeCol_FrameBgHovered, (55, 55, 55, 255))
                dpg.add_theme_color(dpg.mvThemeCol_FrameBgActive, (65, 65, 65, 255))

                # Text
                dpg.add_theme_color(dpg.mvThemeCol_Text, theme["text"])
                dpg.add_theme_color(dpg.mvThemeCol_TextDisabled, theme["disabled"])

                # Tab colors
                dpg.add_theme_color(dpg.mvThemeCol_Tab, (35, 35, 35, 255))
                dpg.add_theme_color(dpg.mvThemeCol_TabHovered, theme["secondary"])
                dpg.add_theme_color(dpg.mvThemeCol_TabActive, theme["primary"])

                # Style settings
                dpg.add_theme_style(dpg.mvStyleVar_WindowRounding, 5)
                dpg.add_theme_style(dpg.mvStyleVar_FrameRounding, 3)
                dpg.add_theme_style(dpg.mvStyleVar_FramePadding, 6, 4)
                dpg.add_theme_style(dpg.mvStyleVar_ItemSpacing, 8, 6)

        dpg.bind_theme(global_theme)

    def create_ui(self):
        """Create the UI elements"""
        # Create primary window
        with dpg.window(tag="main_window", label="OneTrainer"):
            # Create top bar
            self.create_top_bar()

            # Create tab bar
            tab_bar = dpg.add_tab_bar(tag="main_tab_bar")

            # Create tabs
            model_tab = dpg.add_tab(label="Model", parent=tab_bar)
            self.create_model_tab(model_tab)

            training_tab = dpg.add_tab(label="Training", parent=tab_bar)
            self.create_training_tab(training_tab)

            lora_tab = dpg.add_tab(label="LoRA", parent=tab_bar)
            self.create_lora_tab(lora_tab)

            concept_tab = dpg.add_tab(label="Concept", parent=tab_bar)
            self.create_concept_tab(concept_tab)

            sampling_tab = dpg.add_tab(label="Sampling", parent=tab_bar)
            self.create_sampling_tab(sampling_tab)

            embeddings_tab = dpg.add_tab(label="Additional Embeddings", parent=tab_bar)
            self.create_embeddings_tab(embeddings_tab)

            cloud_tab = dpg.add_tab(label="Cloud", parent=tab_bar)
            self.create_cloud_tab(cloud_tab)

            # Status bar at the bottom
            self.create_status_bar()

            # Set primary window
            dpg.set_primary_window("main_window", True)

    def create_top_bar(self):
        """Create the top bar"""
        with dpg.group(horizontal=True, tag="top_bar"):
            # Logo placeholder
            with dpg.drawlist(width=40, height=40):
                dpg.draw_rectangle((0, 0), (40, 40), color=THEMES["dark_blue"]["primary"],
                                  fill=THEMES["dark_blue"]["primary"])

            # Title
            dpg.add_text("OneTrainer", color=THEMES["dark_blue"]["accent"], tag="app_title")
            dpg.add_spacer(width=20)

            # Model type selector
            dpg.add_text("Model Type:", color=THEMES["dark_blue"]["text_secondary"])

            # Get model types if available
            model_types = []
            if HAS_MODULES:
                for model_type in ModelType:
                    model_types.append(model_type.name)
            else:
                model_types = ["STABLE_DIFFUSION_15", "STABLE_DIFFUSION_XL", "PIXART_ALPHA"]

            self.model_type_combo = dpg.add_combo(
                items=model_types,
                default_value=self.config.get("model_type", model_types[0] if model_types else ""),
                callback=self.on_model_change,
                width=200
            )

            dpg.add_spacer(width=20)

            # Training method selector
            dpg.add_text("Training Method:", color=THEMES["dark_blue"]["text_secondary"])

            # Get training methods if available
            training_methods = []
            if HAS_MODULES:
                for method in TrainingMethod:
                    training_methods.append(method.name)
            else:
                training_methods = ["LORA", "EMBEDDING", "FINETUNE"]

            self.training_method_combo = dpg.add_combo(
                items=training_methods,
                default_value=self.config.get("training_method", training_methods[0] if training_methods else ""),
                callback=self.on_training_method_change,
                width=150
            )

            dpg.add_spacer(width=20)

            # Preset selector
            dpg.add_text("Preset:", color=THEMES["dark_blue"]["text_secondary"])
            with dpg.group(horizontal=True):
                self.preset_combo = dpg.add_combo(
                    items=self.presets,
                    default_value=self.presets[0] if self.presets else "",
                    width=200
                )

                # Add load button
                dpg.add_button(
                    label="Load",
                    callback=self.on_load_preset,
                    width=50
                )

                # Add save button
                dpg.add_button(
                    label="Save",
                    callback=self.on_save_preset,
                    width=50
                )

            dpg.add_spacer(width=20)

            # Config buttons
            dpg.add_button(
                label="Load Config",
                callback=self.on_load_config_button,
                width=100
            )

            dpg.add_button(
                label="Save Config",
                callback=self.on_save_config_button,
                width=100
            )

    def create_status_bar(self):
        """Create the status bar at the bottom"""
        with dpg.group(horizontal=True, tag="status_bar"):
            # Status label
            dpg.add_text("Ready", tag="status_text")

            dpg.add_spacer(width=20)

            # Progress bar
            dpg.add_text("Progress:")
            self.progress_bar = dpg.add_progress_bar(default_value=0, width=200, tag="progress_bar")
            self.progress_text = dpg.add_text("0/0", tag="progress_text")
    
    def create_model_tab(self, parent):
        """Create the model tab"""
        with dpg.group(parent=parent):
            dpg.add_text("Model Settings", color=THEMES["dark_blue"]["accent"])
            dpg.add_separator()

            # Create a section for base model settings
            with dpg.collapsing_header(label="Base Model", default_open=True):
                # Model path
                dpg.add_text("Model Path:")
                with dpg.group(horizontal=True):
                    self.model_path_input = dpg.add_input_text(
                        default_value=self.config.get("model_path", ""),
                        width=500,
                        callback=lambda s, a: self.update_config("model_path", a)
                    )
                    browse_button = dpg.add_button(label="...", width=30)

                    # Set up browse button callback
                    def browse_model():
                        with dpg.file_dialog(
                            label="Select Model File",
                            callback=lambda s, a: self._handle_model_file_selected(a),
                            width=700,
                            height=400
                        ):
                            pass

                    dpg.set_item_callback(browse_button, browse_model)

                # Model format
                dpg.add_text("Model Format:")
                self.model_format_combo = dpg.add_combo(
                    items=["diffusers", "safetensors", "ckpt", "auto"],
                    default_value=self.config.get("model_format", "diffusers"),
                    callback=lambda s, a: self.update_config("model_format", a),
                    width=300
                )

                # Precision
                dpg.add_text("Precision:")
                precision_combo = dpg.add_combo(
                    items=["fp16", "fp32", "bf16"],
                    default_value=self.config.get("precision", "fp16"),
                    callback=lambda s, a: self.update_config("precision", a),
                    width=300
                )

                # VRAM options
                with dpg.collapsing_header(label="VRAM Options", default_open=False):
                    vram_options = [
                        ("Enable xformers", "use_xformers", self.config.get("use_xformers", True)),
                        ("Enable attention slicing", "attention_slicing", self.config.get("attention_slicing", False)),
                        ("Enable gradient checkpointing", "gradient_checkpointing", self.config.get("gradient_checkpointing", True)),
                        ("Enable memory efficient attention", "memory_efficient_attention", self.config.get("memory_efficient_attention", True)),
                        ("Use FP8 kernels (requires xformers)", "use_fp8_kernels", self.config.get("use_fp8_kernels", False)),
                        ("Off-load CPU", "offload_cpu", self.config.get("offload_cpu", False))
                    ]

                    for label, key, default_value in vram_options:
                        dpg.add_checkbox(
                            label=label,
                            default_value=default_value,
                            callback=lambda s, a, k=key: self.update_config(k, a)
                        )

                # Device selection
                dpg.add_text("Device:")
                device_combo = dpg.add_combo(
                    items=["cuda", "cpu"] + (["mps"] if sys.platform == "darwin" else []),
                    default_value=self.config.get("device", "cuda"),
                    callback=lambda s, a: self.update_config("device", a),
                    width=300
                )

            # VAE section
            with dpg.collapsing_header(label="VAE Settings", default_open=False):
                # VAE path
                dpg.add_text("VAE Path (optional):")
                with dpg.group(horizontal=True):
                    self.vae_path_input = dpg.add_input_text(
                        default_value=self.config.get("vae_path", ""),
                        width=500,
                        callback=lambda s, a: self.update_config("vae_path", a)
                    )
                    browse_button = dpg.add_button(label="...", width=30)

                    # Set up browse button callback
                    def browse_vae():
                        with dpg.file_dialog(
                            label="Select VAE File",
                            callback=lambda s, a: self._handle_vae_file_selected(a),
                            width=700,
                            height=400
                        ):
                            pass

                    dpg.set_item_callback(browse_button, browse_vae)

                # Scale factor for VAE
                dpg.add_text("VAE Scale Factor:")
                vae_scale = dpg.add_slider_float(
                    default_value=self.config.get("vae_scale_factor", 0.18215),
                    min_value=0.0,
                    max_value=1.0,
                    format="%.5f",
                    callback=lambda s, a: self.update_config("vae_scale_factor", a),
                    width=300
                )

            # Model variants section
            with dpg.collapsing_header(label="Model Variants", default_open=False):
                dpg.add_text("This section is for special model variants that require additional configuration.")

                # Inpainting checkbox
                dpg.add_checkbox(
                    label="Inpainting model",
                    default_value=self.config.get("is_inpainting", False),
                    callback=lambda s, a: self.update_config("is_inpainting", a)
                )

                # SD v2 depth checkbox
                dpg.add_checkbox(
                    label="SD v2 Depth model",
                    default_value=self.config.get("is_depth", False),
                    callback=lambda s, a: self.update_config("is_depth", a)
                )

                # SDXL refiner checkbox
                dpg.add_checkbox(
                    label="SDXL Refiner model",
                    default_value=self.config.get("is_refiner", False),
                    callback=lambda s, a: self.update_config("is_refiner", a)
                )

    def _handle_model_file_selected(self, app_data):
        """Handle model file selection from file dialog"""
        if "file_path_name" in app_data:
            path = app_data["file_path_name"]
            dpg.set_value(self.model_path_input, path)
            self.update_config("model_path", path)

            # Auto-detect model format from extension
            if path.endswith(".safetensors"):
                dpg.set_value(self.model_format_combo, "safetensors")
                self.update_config("model_format", "safetensors")
            elif path.endswith(".ckpt"):
                dpg.set_value(self.model_format_combo, "ckpt")
                self.update_config("model_format", "ckpt")
            elif os.path.isdir(path):
                dpg.set_value(self.model_format_combo, "diffusers")
                self.update_config("model_format", "diffusers")

    def _handle_vae_file_selected(self, app_data):
        """Handle VAE file selection from file dialog"""
        if "file_path_name" in app_data:
            path = app_data["file_path_name"]
            dpg.set_value(self.vae_path_input, path)
            self.update_config("vae_path", path)
    
    def create_training_tab(self, parent):
        """Create the training tab"""
        with dpg.group(parent=parent):
            dpg.add_text("Training Settings", color=THEMES["dark_blue"]["accent"])
            dpg.add_separator()

            # Training parameters section
            with dpg.collapsing_header(label="Training Parameters", default_open=True):
                # Learning rate
                dpg.add_text("Learning Rate:")
                lr_input = dpg.add_input_float(
                    default_value=self.config.get("learning_rate", 0.0001),
                    format="%.6f",
                    callback=lambda s, a: self.update_config("learning_rate", a),
                    width=300
                )

                # Batch size
                dpg.add_text("Batch Size:")
                batch_input = dpg.add_input_int(
                    default_value=self.config.get("batch_size", 1),
                    callback=lambda s, a: self.update_config("batch_size", a),
                    width=300
                )

                # Max train steps
                dpg.add_text("Max Train Steps:")
                steps_input = dpg.add_input_int(
                    default_value=self.config.get("max_train_steps", 1000),
                    callback=lambda s, a: self.update_config("max_train_steps", a),
                    width=300
                )

                # Gradient accumulation steps
                dpg.add_text("Gradient Accumulation Steps:")
                grad_accum_input = dpg.add_input_int(
                    default_value=self.config.get("gradient_accumulation_steps", 1),
                    callback=lambda s, a: self.update_config("gradient_accumulation_steps", a),
                    width=300
                )

                # Mixed precision
                dpg.add_text("Mixed Precision:")
                mixed_precision_combo = dpg.add_combo(
                    items=["no", "fp16", "bf16"],
                    default_value=self.config.get("mixed_precision", "fp16"),
                    callback=lambda s, a: self.update_config("mixed_precision", a),
                    width=300
                )

                # Save checkpoints
                dpg.add_text("Save Every N Steps:")
                save_steps_input = dpg.add_input_int(
                    default_value=self.config.get("save_every_n_steps", 100),
                    callback=lambda s, a: self.update_config("save_every_n_steps", a),
                    width=300
                )

            # Optimizer section
            with dpg.collapsing_header(label="Optimizer Settings", default_open=True):
                # Optimizer type
                dpg.add_text("Optimizer:")

                # Get optimizer types
                optimizer_types = []
                if HAS_MODULES:
                    for optimizer in Optimizer:
                        optimizer_types.append(optimizer.name)
                else:
                    optimizer_types = ["AdamW", "AdamW8bit", "Lion", "Lion8bit", "SGD"]

                optimizer_combo = dpg.add_combo(
                    items=optimizer_types,
                    default_value=self.config.get("optimizer", "AdamW8bit"),
                    callback=lambda s, a: self.update_config("optimizer", a),
                    width=300
                )

                # Weight decay
                dpg.add_text("Weight Decay:")
                weight_decay_input = dpg.add_input_float(
                    default_value=self.config.get("weight_decay", 0.01),
                    format="%.4f",
                    callback=lambda s, a: self.update_config("weight_decay", a),
                    width=300
                )

                # Max gradient norm
                dpg.add_text("Max Gradient Norm:")
                max_grad_norm_input = dpg.add_input_float(
                    default_value=self.config.get("max_grad_norm", 1.0),
                    format="%.2f",
                    callback=lambda s, a: self.update_config("max_grad_norm", a),
                    width=300
                )

            # Scheduler section
            with dpg.collapsing_header(label="Learning Rate Scheduler", default_open=True):
                # Scheduler type
                dpg.add_text("Scheduler:")

                # Get scheduler types
                scheduler_types = []
                if HAS_MODULES:
                    for scheduler in LearningRateScheduler:
                        scheduler_types.append(scheduler.name)
                else:
                    scheduler_types = ["constant", "cosine", "linear", "polynomial"]

                scheduler_combo = dpg.add_combo(
                    items=scheduler_types,
                    default_value=self.config.get("scheduler", "cosine"),
                    callback=lambda s, a: self.update_config("scheduler", a),
                    width=300
                )

                # Warmup steps
                dpg.add_text("Warmup Steps:")
                warmup_steps_input = dpg.add_input_int(
                    default_value=self.config.get("warmup_steps", 100),
                    callback=lambda s, a: self.update_config("warmup_steps", a),
                    width=300
                )

                # Warmup ratio
                dpg.add_text("Warmup Ratio:")
                warmup_ratio_input = dpg.add_input_float(
                    default_value=self.config.get("warmup_ratio", 0.1),
                    format="%.3f",
                    callback=lambda s, a: self.update_config("warmup_ratio", a),
                    width=300
                )

            # Noise scheduler section
            with dpg.collapsing_header(label="Noise Scheduler", default_open=False):
                # Noise scheduler type
                dpg.add_text("Noise Scheduler Type:")

                # Get noise scheduler types
                noise_scheduler_types = []
                if HAS_MODULES:
                    for scheduler in NoiseScheduler:
                        noise_scheduler_types.append(scheduler.name)
                else:
                    noise_scheduler_types = ["ddpm", "ddim", "pndm", "lms", "euler", "euler_ancestral"]

                noise_scheduler_combo = dpg.add_combo(
                    items=noise_scheduler_types,
                    default_value=self.config.get("noise_scheduler", "ddpm"),
                    callback=lambda s, a: self.update_config("noise_scheduler", a),
                    width=300
                )

                # Prediction type
                dpg.add_text("Prediction Type:")
                prediction_type_combo = dpg.add_combo(
                    items=["epsilon", "v_prediction", "sample"],
                    default_value=self.config.get("prediction_type", "epsilon"),
                    callback=lambda s, a: self.update_config("prediction_type", a),
                    width=300
                )

                # Snr gamma
                dpg.add_text("SNR Gamma:")
                snr_gamma_input = dpg.add_input_float(
                    default_value=self.config.get("snr_gamma", 5.0),
                    format="%.1f",
                    callback=lambda s, a: self.update_config("snr_gamma", a),
                    width=300
                )

            # Training controls
            with dpg.group(horizontal=True):
                # Start training button
                self.train_button = dpg.add_button(
                    label="Start Training",
                    callback=self.on_start_training,
                    width=150,
                    height=40
                )

                # Stop training button
                self.stop_button = dpg.add_button(
                    label="Stop Training",
                    callback=self.on_stop_training,
                    width=150,
                    height=40,
                    enabled=False
                )
    
    def create_lora_tab(self, parent):
        """Create the LoRA tab"""
        with dpg.group(parent=parent):
            dpg.add_text("LoRA Settings", color=THEMES["dark_blue"]["accent"])
            dpg.add_separator()

            # Main LoRA parameters section
            with dpg.collapsing_header(label="Core LoRA Parameters", default_open=True):
                # LoRA Rank
                dpg.add_text("LoRA Rank:")
                lora_rank_input = dpg.add_input_int(
                    default_value=self.config.get("lora_rank", 32),
                    callback=lambda s, a: self.update_config("lora_rank", a),
                    width=300
                )

                # LoRA Alpha
                dpg.add_text("LoRA Alpha:")
                lora_alpha_input = dpg.add_input_float(
                    default_value=self.config.get("lora_alpha", 32.0),
                    callback=lambda s, a: self.update_config("lora_alpha", a),
                    width=300
                )

                # Dropout
                dpg.add_text("Dropout:")
                lora_dropout_input = dpg.add_slider_float(
                    default_value=self.config.get("lora_dropout", 0.0),
                    min_value=0.0,
                    max_value=0.5,
                    format="%.2f",
                    callback=lambda s, a: self.update_config("lora_dropout", a),
                    width=300
                )

                # Target modules
                dpg.add_text("Target Modules:")
                target_modules_input = dpg.add_input_text(
                    default_value=self.config.get("target_modules", ""),
                    hint="q_proj,k_proj,v_proj,out_proj",
                    callback=lambda s, a: self.update_config("target_modules", a),
                    width=500
                )

                # Target layers
                dpg.add_text("Target Layers:")
                target_layers_input = dpg.add_input_text(
                    default_value=self.config.get("target_layers", ""),
                    hint="down_blocks,mid_block,up_blocks",
                    callback=lambda s, a: self.update_config("target_layers", a),
                    width=500
                )

            # Advanced LoRA section
            with dpg.collapsing_header(label="Advanced LoRA Settings", default_open=False):
                # Network weights
                dpg.add_text("Pretrained Network Weights (optional):")
                with dpg.group(horizontal=True):
                    self.network_weights_input = dpg.add_input_text(
                        default_value=self.config.get("network_weights", ""),
                        width=500,
                        callback=lambda s, a: self.update_config("network_weights", a)
                    )
                    browse_button = dpg.add_button(label="...", width=30)

                    # Set up browse button callback
                    def browse_weights():
                        with dpg.file_dialog(
                            label="Select Network Weights File",
                            callback=lambda s, a: self._handle_weights_file_selected(a),
                            width=700,
                            height=400
                        ):
                            pass

                    dpg.set_item_callback(browse_button, browse_weights)

                # Scale option
                dpg.add_checkbox(
                    label="Scale LoRA weights",
                    default_value=self.config.get("scale_lora_weights", True),
                    callback=lambda s, a: self.update_config("scale_lora_weights", a)
                )

                # LoRA type (LyCORIS)
                dpg.add_text("LoRA Type (LyCORIS):")
                lora_type_combo = dpg.add_combo(
                    items=["lora", "loha", "lokr", "locon", "ia3", "dylora"],
                    default_value=self.config.get("lora_type", "lora"),
                    callback=lambda s, a: self.update_config("lora_type", a),
                    width=300
                )

                # LoCon / LoKr options
                dpg.add_text("LoKr Decomposition Types:")
                decomp_type_combo = dpg.add_combo(
                    items=["tucker", "svd", "full"],
                    default_value=self.config.get("decomposition_type", "tucker"),
                    callback=lambda s, a: self.update_config("decomposition_type", a),
                    width=300
                )

                # LoHa options
                dpg.add_text("LoHa Hadamard Config:")
                with dpg.group(horizontal=True):
                    use_hadamard_checkbox = dpg.add_checkbox(
                        label="Use Hadamard Product",
                        default_value=self.config.get("use_hadamard", True),
                        callback=lambda s, a: self.update_config("use_hadamard", a)
                    )

                    hadamard_dim_input = dpg.add_input_int(
                        default_value=self.config.get("hadamard_dim", 4),
                        callback=lambda s, a: self.update_config("hadamard_dim", a),
                        width=100
                    )

            # LoRA target layers section
            with dpg.collapsing_header(label="Model-Specific Settings", default_open=False):
                dpg.add_text("Target specific parts of the model:")

                # Text encoder enable/disable
                dpg.add_checkbox(
                    label="Apply to Text Encoder",
                    default_value=self.config.get("train_text_encoder", True),
                    callback=lambda s, a: self.update_config("train_text_encoder", a)
                )

                # Text encoder 2 enable/disable (for SDXL)
                dpg.add_checkbox(
                    label="Apply to Text Encoder 2 (SDXL)",
                    default_value=self.config.get("train_text_encoder_2", True),
                    callback=lambda s, a: self.update_config("train_text_encoder_2", a)
                )

                # UNet enable/disable
                dpg.add_checkbox(
                    label="Apply to UNet",
                    default_value=self.config.get("train_unet", True),
                    callback=lambda s, a: self.update_config("train_unet", a)
                )

                # Different rank for text encoder
                dpg.add_checkbox(
                    label="Use different rank for text encoder",
                    default_value=self.config.get("use_different_text_encoder_rank", False),
                    callback=lambda s, a: self.update_config("use_different_text_encoder_rank", a)
                )

                # Text encoder rank
                dpg.add_text("Text Encoder Rank:")
                text_encoder_rank_input = dpg.add_input_int(
                    default_value=self.config.get("text_encoder_rank", 4),
                    callback=lambda s, a: self.update_config("text_encoder_rank", a),
                    width=300
                )

    def _handle_weights_file_selected(self, app_data):
        """Handle weights file selection from file dialog"""
        if "file_path_name" in app_data:
            path = app_data["file_path_name"]
            dpg.set_value(self.network_weights_input, path)
            self.update_config("network_weights", path)
    
    def create_concept_tab(self, parent):
        """Create the concept tab"""
        with dpg.group(parent=parent):
            dpg.add_text("Concept Settings", color=THEMES["dark_blue"]["accent"])
            dpg.add_separator()

            # Main concept section
            with dpg.collapsing_header(label="Concept Definition", default_open=True):
                # Concept parameters
                dpg.add_text("Concept Name:")
                self.concept_name_input = dpg.add_input_text(
                    default_value=self.config.get("concept_name", "my_concept"),
                    width=300,
                    callback=lambda s, a: self.update_config("concept_name", a)
                )

                dpg.add_text("Image Directory:")
                with dpg.group(horizontal=True):
                    self.image_dir_input = dpg.add_input_text(
                        default_value=self.config.get("image_dir", ""),
                        width=500,
                        callback=lambda s, a: self.update_config("image_dir", a)
                    )
                    browse_button = dpg.add_button(label="...", width=30)

                    # Set up browse button callback
                    def browse_dir():
                        with dpg.file_dialog(
                            label="Select Image Directory",
                            directory_selector=True,
                            callback=lambda s, a: self._handle_image_dir_selected(a),
                            width=700,
                            height=400
                        ):
                            pass

                    dpg.set_item_callback(browse_button, browse_dir)

                # Trigger words section
                dpg.add_text("Trigger Words:")
                self.trigger_words_input = dpg.add_input_text(
                    default_value=self.config.get("trigger_words", ""),
                    width=500,
                    callback=lambda s, a: self.update_config("trigger_words", a)
                )

                # Instance prompt - for Dreambooth-style training
                dpg.add_text("Instance Prompt:")
                instance_prompt_input = dpg.add_input_text(
                    default_value=self.config.get("instance_prompt", "a photo of sks"),
                    width=500,
                    callback=lambda s, a: self.update_config("instance_prompt", a)
                )

                # Class prompt - for Dreambooth-style training
                dpg.add_text("Class Prompt:")
                class_prompt_input = dpg.add_input_text(
                    default_value=self.config.get("class_prompt", "a photo of"),
                    width=500,
                    callback=lambda s, a: self.update_config("class_prompt", a)
                )

            # Image preprocessing section
            with dpg.collapsing_header(label="Image Preprocessing", default_open=True):
                # Resolution
                dpg.add_text("Resolution:")
                with dpg.group(horizontal=True):
                    dpg.add_input_int(
                        default_value=self.config.get("resolution_width", 512),
                        width=145,
                        callback=lambda s, a: self.update_config("resolution_width", a)
                    )
                    dpg.add_text("x")
                    dpg.add_input_int(
                        default_value=self.config.get("resolution_height", 512),
                        width=145,
                        callback=lambda s, a: self.update_config("resolution_height", a)
                    )

                # Image preprocessing options
                preprocessing_options = [
                    ("Center crop", "center_crop", self.config.get("center_crop", True)),
                    ("Random crop", "random_crop", self.config.get("random_crop", False)),
                    ("Random flip", "random_flip", self.config.get("random_flip", True)),
                    ("Remove background", "remove_background", self.config.get("remove_background", False)),
                    ("Auto face crop", "auto_face_crop", self.config.get("auto_face_crop", False)),
                    ("Improve captions", "improve_captions", self.config.get("improve_captions", False))
                ]

                for label, key, default_value in preprocessing_options:
                    dpg.add_checkbox(
                        label=label,
                        default_value=default_value,
                        callback=lambda s, a, k=key: self.update_config(k, a)
                    )

                # Caption extension
                dpg.add_text("Caption File Extension:")
                caption_ext_combo = dpg.add_combo(
                    items=[".txt", ".caption"],
                    default_value=self.config.get("caption_extension", ".txt"),
                    callback=lambda s, a: self.update_config("caption_extension", a),
                    width=300
                )

            # Class images section (for regularization)
            with dpg.collapsing_header(label="Class/Regularization Images", default_open=False):
                # Use class/regularization images checkbox
                dpg.add_checkbox(
                    label="Use class/regularization images",
                    default_value=self.config.get("use_class_images", False),
                    callback=lambda s, a: self.update_config("use_class_images", a)
                )

                # Class images directory
                dpg.add_text("Class Images Directory:")
                with dpg.group(horizontal=True):
                    self.class_dir_input = dpg.add_input_text(
                        default_value=self.config.get("class_dir", ""),
                        width=500,
                        callback=lambda s, a: self.update_config("class_dir", a)
                    )
                    browse_button = dpg.add_button(label="...", width=30)

                    # Set up browse button callback
                    def browse_class_dir():
                        with dpg.file_dialog(
                            label="Select Class Images Directory",
                            directory_selector=True,
                            callback=lambda s, a: self._handle_class_dir_selected(a),
                            width=700,
                            height=400
                        ):
                            pass

                    dpg.set_item_callback(browse_button, browse_class_dir)

                # Number of class images
                dpg.add_text("Number of Class Images:")
                num_class_images_input = dpg.add_input_int(
                    default_value=self.config.get("num_class_images", 100),
                    callback=lambda s, a: self.update_config("num_class_images", a),
                    width=300
                )

                # Class prompt for auto-generation
                dpg.add_text("Auto-generation Class Prompt:")
                auto_class_prompt_input = dpg.add_input_text(
                    default_value=self.config.get("auto_class_prompt", "a photo of a person"),
                    width=500,
                    callback=lambda s, a: self.update_config("auto_class_prompt", a)
                )

            # Advanced concept options
            with dpg.collapsing_header(label="Advanced Options", default_open=False):
                # Concept weights
                dpg.add_text("Concept Weight:")
                concept_weight_slider = dpg.add_slider_float(
                    default_value=self.config.get("concept_weight", 1.0),
                    min_value=0.0,
                    max_value=5.0,
                    format="%.2f",
                    callback=lambda s, a: self.update_config("concept_weight", a),
                    width=300
                )

                # Prior loss weight
                dpg.add_text("Prior Loss Weight:")
                prior_weight_slider = dpg.add_slider_float(
                    default_value=self.config.get("prior_loss_weight", 1.0),
                    min_value=0.0,
                    max_value=5.0,
                    format="%.2f",
                    callback=lambda s, a: self.update_config("prior_loss_weight", a),
                    width=300
                )

                # Class prompt for prior preservation
                dpg.add_text("Prior Preservation Class Prompt:")
                prior_class_prompt_input = dpg.add_input_text(
                    default_value=self.config.get("prior_class_prompt", "a photo of a person"),
                    width=500,
                    callback=lambda s, a: self.update_config("prior_class_prompt", a)
                )

            # Concept actions
            with dpg.group(horizontal=True):
                # Generate captions button
                dpg.add_button(
                    label="Generate Captions",
                    callback=self.on_generate_captions,
                    width=150,
                    height=40
                )

                # Generate masks button
                dpg.add_button(
                    label="Generate Masks",
                    callback=self.on_generate_masks,
                    width=150,
                    height=40
                )

                # Analyze concept button
                dpg.add_button(
                    label="Analyze Concept",
                    callback=self.on_analyze_concept,
                    width=150,
                    height=40
                )

    def _handle_image_dir_selected(self, app_data):
        """Handle image directory selection"""
        if "file_path_name" in app_data:
            path = app_data["file_path_name"]
            dpg.set_value(self.image_dir_input, path)
            self.update_config("image_dir", path)

    def _handle_class_dir_selected(self, app_data):
        """Handle class directory selection"""
        if "file_path_name" in app_data:
            path = app_data["file_path_name"]
            dpg.set_value(self.class_dir_input, path)
            self.update_config("class_dir", path)

    def on_generate_captions(self):
        """Handle generate captions button"""
        dpg.set_value("status_text", "Generating captions...")
        # This would launch a captioning process in a real implementation

    def on_generate_masks(self):
        """Handle generate masks button"""
        dpg.set_value("status_text", "Generating masks...")
        # This would launch a masking process in a real implementation

    def on_analyze_concept(self):
        """Handle analyze concept button"""
        dpg.set_value("status_text", "Analyzing concept...")
        # This would analyze the concept images in a real implementation
    
    def create_sampling_tab(self, parent):
        """Create the sampling tab"""
        with dpg.group(parent=parent):
            dpg.add_text("Sampling Settings", color=THEMES["dark_blue"]["accent"])
            dpg.add_separator()

            # Main sampling parameters section
            with dpg.collapsing_header(label="Text Prompts", default_open=True):
                # Prompt
                dpg.add_text("Prompt:")
                self.prompt_input = dpg.add_input_text(
                    default_value=self.config.get("prompt", ""),
                    multiline=True,
                    width=700,
                    height=100,
                    callback=lambda s, a: self.update_config("prompt", a)
                )

                # Negative prompt
                dpg.add_text("Negative Prompt:")
                self.negative_prompt_input = dpg.add_input_text(
                    default_value=self.config.get("negative_prompt", ""),
                    multiline=True,
                    width=700,
                    height=100,
                    callback=lambda s, a: self.update_config("negative_prompt", a)
                )

            # Sampling options
            with dpg.collapsing_header(label="Sampling Parameters", default_open=True):
                # Scheduler selection
                dpg.add_text("Sampler:")
                sampler_combo = dpg.add_combo(
                    items=["ddim", "pndm", "lms", "euler", "euler_ancestral", "dpm", "dpm_solver", "heun", "dpm_solver++"],
                    default_value=self.config.get("sampler", "euler_ancestral"),
                    callback=lambda s, a: self.update_config("sampler", a),
                    width=300
                )

                # Steps
                dpg.add_text("Steps:")
                steps_slider = dpg.add_slider_int(
                    default_value=self.config.get("steps", 30),
                    min_value=1,
                    max_value=100,
                    callback=lambda s, a: self.update_config("steps", a),
                    width=300
                )

                # CFG Scale
                dpg.add_text("CFG Scale:")
                cfg_slider = dpg.add_slider_float(
                    default_value=self.config.get("cfg_scale", 7.0),
                    min_value=1.0,
                    max_value=20.0,
                    format="%.1f",
                    callback=lambda s, a: self.update_config("cfg_scale", a),
                    width=300
                )

                # Batch size
                dpg.add_text("Batch Size:")
                batch_slider = dpg.add_slider_int(
                    default_value=self.config.get("sample_batch_size", 1),
                    min_value=1,
                    max_value=16,
                    callback=lambda s, a: self.update_config("sample_batch_size", a),
                    width=300
                )

                # Number of images
                dpg.add_text("Number of Images:")
                num_images_input = dpg.add_input_int(
                    default_value=self.config.get("num_images", 4),
                    callback=lambda s, a: self.update_config("num_images", a),
                    width=300
                )

            # Advanced settings
            with dpg.collapsing_header(label="Advanced Settings", default_open=False):
                # Seed
                dpg.add_text("Seed:")
                with dpg.group(horizontal=True):
                    seed_input = dpg.add_input_int(
                        default_value=self.config.get("seed", -1),
                        callback=lambda s, a: self.update_config("seed", a),
                        width=200
                    )

                    # Random seed checkbox
                    random_seed_checkbox = dpg.add_checkbox(
                        label="Random",
                        default_value=self.config.get("seed", -1) < 0,
                        callback=lambda s, a: dpg.set_value(seed_input, -1 if a else 42)
                    )

                # Attention slicing for sampling
                dpg.add_checkbox(
                    label="Use attention slicing",
                    default_value=self.config.get("sample_attention_slicing", False),
                    callback=lambda s, a: self.update_config("sample_attention_slicing", a)
                )

                # Memory efficient attention for sampling
                dpg.add_checkbox(
                    label="Use memory efficient attention",
                    default_value=self.config.get("sample_memory_efficient_attention", True),
                    callback=lambda s, a: self.update_config("sample_memory_efficient_attention", a)
                )

                # Guidance rescale
                dpg.add_text("Guidance Rescale:")
                guidance_rescale_slider = dpg.add_slider_float(
                    default_value=self.config.get("guidance_rescale", 0.0),
                    min_value=0.0,
                    max_value=1.0,
                    format="%.2f",
                    callback=lambda s, a: self.update_config("guidance_rescale", a),
                    width=300
                )

            # Image settings
            with dpg.collapsing_header(label="Image Settings", default_open=True):
                # Resolution
                dpg.add_text("Generation Resolution:")
                with dpg.group(horizontal=True):
                    dpg.add_input_int(
                        default_value=self.config.get("width", 512),
                        width=145,
                        callback=lambda s, a: self.update_config("width", a)
                    )
                    dpg.add_text("x")
                    dpg.add_input_int(
                        default_value=self.config.get("height", 512),
                        width=145,
                        callback=lambda s, a: self.update_config("height", a)
                    )

                # Output directory
                dpg.add_text("Output Directory:")
                with dpg.group(horizontal=True):
                    self.output_dir_input = dpg.add_input_text(
                        default_value=self.config.get("output_dir", "samples"),
                        width=500,
                        callback=lambda s, a: self.update_config("output_dir", a)
                    )
                    browse_button = dpg.add_button(label="...", width=30)

                    # Set up browse button callback
                    def browse_output_dir():
                        with dpg.file_dialog(
                            label="Select Output Directory",
                            directory_selector=True,
                            callback=lambda s, a: self._handle_output_dir_selected(a),
                            width=700,
                            height=400
                        ):
                            pass

                    dpg.set_item_callback(browse_button, browse_output_dir)

                # Save images checkbox
                dpg.add_checkbox(
                    label="Save generated images",
                    default_value=self.config.get("save_images", True),
                    callback=lambda s, a: self.update_config("save_images", a)
                )

                # Image format
                dpg.add_text("Image Format:")
                format_combo = dpg.add_combo(
                    items=["png", "jpg", "webp"],
                    default_value=self.config.get("image_format", "png"),
                    callback=lambda s, a: self.update_config("image_format", a),
                    width=300
                )

            # Sample preview section
            with dpg.collapsing_header(label="Sample Preview", default_open=True):
                # Placeholder for sample image
                with dpg.group(horizontal=True):
                    for i in range(2):  # 2x2 grid of image placeholders
                        with dpg.group():
                            for j in range(2):
                                with dpg.drawlist(width=256, height=256, tag=f"sample_image_{i}_{j}"):
                                    dpg.draw_rectangle((0, 0), (256, 256), color=(30, 30, 30, 255),
                                                      fill=THEMES["dark_blue"]["background"])
                                    dpg.draw_text((100, 120), "No image", color=THEMES["dark_blue"]["text_secondary"])

            # Sample controls
            with dpg.group(horizontal=True):
                # Generate sample button
                self.sample_button = dpg.add_button(
                    label="Generate Sample",
                    callback=self.on_generate_sample,
                    width=150,
                    height=40
                )

                # Stop sampling button
                self.stop_sample_button = dpg.add_button(
                    label="Stop",
                    callback=self.on_stop_sampling,
                    width=150,
                    height=40,
                    enabled=False
                )

            # Progress indicator for sampling
            dpg.add_text("Sampling Progress:")
            self.sampling_progress_bar = dpg.add_progress_bar(default_value=0, width=500, tag="sampling_progress_bar")

    def _handle_output_dir_selected(self, app_data):
        """Handle output directory selection"""
        if "file_path_name" in app_data:
            path = app_data["file_path_name"]
            dpg.set_value(self.output_dir_input, path)
            self.update_config("output_dir", path)

    def on_generate_sample(self):
        """Handle generate sample button"""
        print("Generating sample...")
        dpg.set_value("status_text", "Generating sample...")

        # Update button states
        dpg.configure_item(self.sample_button, enabled=False)
        dpg.configure_item(self.stop_sample_button, enabled=True)

        # Start a thread to simulate sampling
        threading.Thread(target=self._sampling_worker, daemon=True).start()

    def on_stop_sampling(self):
        """Handle stop sampling button"""
        dpg.set_value("status_text", "Stopping sampling...")

        # Update button states
        dpg.configure_item(self.sample_button, enabled=True)
        dpg.configure_item(self.stop_sample_button, enabled=False)

    def _sampling_worker(self):
        """Sampling worker thread"""
        try:
            # Simulate sampling progress
            steps = self.config.get("steps", 30)
            for step in range(steps):
                # Update progress
                progress = (step + 1) / steps
                dpg.set_value("sampling_progress_bar", progress)

                # Simulate step time
                time.sleep(0.1)

                # Check if stopped
                if not dpg.get_item_configuration(self.stop_sample_button)["enabled"]:
                    break

            # Sampling complete
            dpg.set_value("status_text", "Sampling complete")

            # In a real implementation, we would load and display the generated images here

        except Exception as e:
            print(f"Error in sampling thread: {e}")
            dpg.set_value("status_text", f"Error: {e}")
        finally:
            # Reset button states
            dpg.configure_item(self.sample_button, enabled=True)
            dpg.configure_item(self.stop_sample_button, enabled=False)
    
    def update_config(self, key, value):
        """Update a config value and propagate changes"""
        self.config[key] = value
        print(f"Updated {key} = {value}")

    def on_model_change(self, sender, app_data):
        """Handle model type change"""
        print(f"Model type changed to: {app_data}")
        self.update_config("model_type", app_data)

        # Update UI elements based on model type
        if "XL" in app_data:
            # Enable SDXL-specific options
            pass
        elif "PIXART" in app_data:
            # Enable PixArt-specific options
            pass

    def on_training_method_change(self, sender, app_data):
        """Handle training method change"""
        print(f"Training method changed to: {app_data}")
        self.update_config("training_method", app_data)

    def on_start_training(self):
        """Handle start training button"""
        if self.is_training:
            return

        print("Starting training...")
        dpg.set_value("status_text", "Training started")

        # Update button states
        dpg.configure_item(self.train_button, enabled=False)
        dpg.configure_item(self.stop_button, enabled=True)

        # Start training in a separate thread
        self.is_training = True
        self.training_thread = threading.Thread(target=self._training_worker)
        self.training_thread.daemon = True
        self.training_thread.start()

    def on_stop_training(self):
        """Handle stop training button"""
        if not self.is_training:
            return

        print("Stopping training...")
        dpg.set_value("status_text", "Stopping training...")
        self.is_training = False

    def _training_worker(self):
        """Training worker thread"""
        try:
            # Simulate training progress
            total_steps = self.config.get("max_train_steps", 1000)
            for step in range(1, total_steps + 1):
                if not self.is_training:
                    break

                # Update progress
                self.progress["step"] = step
                self.progress["total_steps"] = total_steps

                # Send update to UI thread
                dpg.set_value("progress_bar", step / total_steps)
                dpg.set_value("progress_text", f"{step}/{total_steps}")

                # Simulate training step
                time.sleep(0.1)

            # Training complete or stopped
            if self.is_training:
                dpg.set_value("status_text", "Training complete")
            else:
                dpg.set_value("status_text", "Training stopped")
        except Exception as e:
            print(f"Error in training thread: {e}")
            dpg.set_value("status_text", f"Error: {e}")
        finally:
            # Reset state
            self.is_training = False
            dpg.configure_item(self.train_button, enabled=True)
            dpg.configure_item(self.stop_button, enabled=False)

    def on_load_preset(self):
        """Handle load preset button"""
        preset_name = dpg.get_value(self.preset_combo)
        if preset_name == "-- Select Preset --":
            return

        if preset_name in self.preset_paths:
            preset_path = self.preset_paths[preset_name]
            try:
                with open(preset_path, "r") as f:
                    preset_data = json.load(f)

                # Update config
                self.config.update(preset_data)

                # Update UI to reflect loaded preset
                self._update_ui_from_config()

                dpg.set_value("status_text", f"Loaded preset: {preset_name}")
            except Exception as e:
                print(f"Error loading preset {preset_name}: {e}")
                dpg.set_value("status_text", f"Error loading preset: {e}")

    def on_save_preset(self):
        """Handle save preset button"""
        # Create save dialog
        with dpg.window(label="Save Preset", modal=True, width=400, height=150, pos=(400, 300)):
            dpg.add_text("Enter preset name:")
            input_id = dpg.add_input_text(width=380)

            with dpg.group(horizontal=True):
                dpg.add_button(
                    label="Save",
                    callback=lambda: self._save_preset(dpg.get_value(input_id))
                )
                dpg.add_button(
                    label="Cancel",
                    callback=lambda: dpg.delete_item(dpg.get_current_parent_id())
                )

    def _save_preset(self, preset_name):
        """Save current config as preset with given name"""
        if not preset_name:
            return

        # Close dialog
        dpg.delete_item(dpg.get_current_parent_id())

        # Ensure presets directory exists
        if not os.path.exists(PRESETS_DIR):
            os.makedirs(PRESETS_DIR, exist_ok=True)

        # Save config to preset file
        preset_path = os.path.join(PRESETS_DIR, f"{preset_name}.json")
        try:
            with open(preset_path, "w") as f:
                json.dump(self.config, f, indent=2)

            # Update presets list
            self.load_presets()

            # Update UI
            dpg.configure_item(self.preset_combo, items=self.presets)
            dpg.set_value(self.preset_combo, preset_name)
            dpg.set_value("status_text", f"Saved preset: {preset_name}")
        except Exception as e:
            print(f"Error saving preset: {e}")
            dpg.set_value("status_text", f"Error saving preset: {e}")

    def on_load_config_button(self):
        """Handle load config button"""
        with dpg.file_dialog(
            label="Load Config",
            callback=self._handle_load_config,
            width=700,
            height=400
        ):
            dpg.add_file_extension(".json", color=(0, 255, 0, 255))

    def _handle_load_config(self, sender, app_data):
        """Handle file selection for loading config"""
        if "file_path_name" in app_data:
            config_path = app_data["file_path_name"]
            try:
                with open(config_path, "r") as f:
                    config_data = json.load(f)

                # Update config
                self.config.update(config_data)

                # Update UI to reflect loaded config
                self._update_ui_from_config()

                dpg.set_value("status_text", f"Loaded config: {os.path.basename(config_path)}")
            except Exception as e:
                print(f"Error loading config: {e}")
                dpg.set_value("status_text", f"Error loading config: {e}")

    def on_save_config_button(self):
        """Handle save config button"""
        with dpg.file_dialog(
            label="Save Config",
            callback=self._handle_save_config,
            width=700,
            height=400
        ):
            dpg.add_file_extension(".json", color=(0, 255, 0, 255))

    def _handle_save_config(self, sender, app_data):
        """Handle file selection for saving config"""
        if "file_path_name" in app_data:
            config_path = app_data["file_path_name"]

            # Add .json extension if missing
            if not config_path.lower().endswith(".json"):
                config_path += ".json"

            # Save config
            self.save_config(config_path)
            dpg.set_value("status_text", f"Saved config: {os.path.basename(config_path)}")

    def _update_ui_from_config(self):
        """Update UI elements to reflect current config"""
        # This would update all UI elements based on the config
        # For now, just update a few key elements
        if "model_type" in self.config:
            dpg.set_value(self.model_type_combo, self.config["model_type"])

        if "training_method" in self.config:
            dpg.set_value(self.training_method_combo, self.config["training_method"])

    def on_generate_sample(self):
        """Handle generate sample button"""
        print("Generating sample...")
        dpg.set_value("status_text", "Generating sample...")
    
    def create_embeddings_tab(self, parent):
        """Create the additional embeddings tab"""
        with dpg.group(parent=parent):
            dpg.add_text("Additional Text Embeddings", color=THEMES["dark_blue"]["accent"])
            dpg.add_separator()

            # Main section
            with dpg.collapsing_header(label="Additional Embeddings", default_open=True):
                # CLIP settings
                dpg.add_text("CLIP Settings", color=THEMES["dark_blue"]["text_secondary"])

                # Use OpenCLIP checkbox
                dpg.add_checkbox(
                    label="Use OpenCLIP",
                    default_value=self.config.get("use_open_clip", False),
                    callback=lambda s, a: self.update_config("use_open_clip", a)
                )

                # CLIP skip
                dpg.add_text("CLIP Skip:")
                clip_skip = dpg.add_slider_int(
                    default_value=self.config.get("clip_skip", 1),
                    min_value=1,
                    max_value=12,
                    callback=lambda s, a: self.update_config("clip_skip", a),
                    width=300
                )

                # Additional embeddings
                dpg.add_text("Additional Embeddings", color=THEMES["dark_blue"]["text_secondary"])
                dpg.add_separator()

                # T5 checkbox
                dpg.add_checkbox(
                    label="Use T5 Embeddings",
                    default_value=self.config.get("use_t5", False),
                    callback=lambda s, a: self.update_config("use_t5", a)
                )

                # T5 model selector
                dpg.add_text("T5 Model:")
                t5_model = dpg.add_combo(
                    items=["t5-small", "t5-base", "t5-large", "flan-t5-small", "flan-t5-base", "flan-t5-large"],
                    default_value=self.config.get("t5_model", "t5-base"),
                    callback=lambda s, a: self.update_config("t5_model", a),
                    width=300
                )

                # T5 weight
                dpg.add_text("T5 Weight:")
                t5_weight = dpg.add_slider_float(
                    default_value=self.config.get("t5_weight", 0.5),
                    min_value=0.0,
                    max_value=1.0,
                    format="%.2f",
                    callback=lambda s, a: self.update_config("t5_weight", a),
                    width=300
                )

                # BERT checkbox
                dpg.add_checkbox(
                    label="Use BERT Embeddings",
                    default_value=self.config.get("use_bert", False),
                    callback=lambda s, a: self.update_config("use_bert", a)
                )

                # BERT model selector
                dpg.add_text("BERT Model:")
                bert_model = dpg.add_combo(
                    items=["bert-base-uncased", "bert-large-uncased", "bert-base-cased", "bert-large-cased"],
                    default_value=self.config.get("bert_model", "bert-base-uncased"),
                    callback=lambda s, a: self.update_config("bert_model", a),
                    width=300
                )

                # BERT weight
                dpg.add_text("BERT Weight:")
                bert_weight = dpg.add_slider_float(
                    default_value=self.config.get("bert_weight", 0.5),
                    min_value=0.0,
                    max_value=1.0,
                    format="%.2f",
                    callback=lambda s, a: self.update_config("bert_weight", a),
                    width=300
                )

    def create_cloud_tab(self, parent):
        """Create the cloud tab"""
        with dpg.group(parent=parent):
            dpg.add_text("Cloud Training Settings", color=THEMES["dark_blue"]["accent"])
            dpg.add_separator()

            # Cloud provider section
            with dpg.collapsing_header(label="Cloud Provider", default_open=True):
                # Provider selector
                dpg.add_text("Cloud Provider:")
                provider_combo = dpg.add_combo(
                    items=["None", "RunPod", "AWS", "GCP", "Azure", "Custom SSH"],
                    default_value=self.config.get("cloud_provider", "None"),
                    callback=lambda s, a: self.update_config("cloud_provider", a),
                    width=300
                )

                # SSH settings
                dpg.add_text("SSH Host:")
                ssh_host = dpg.add_input_text(
                    default_value=self.config.get("ssh_host", ""),
                    callback=lambda s, a: self.update_config("ssh_host", a),
                    width=300
                )

                dpg.add_text("SSH Port:")
                ssh_port = dpg.add_input_int(
                    default_value=self.config.get("ssh_port", 22),
                    callback=lambda s, a: self.update_config("ssh_port", a),
                    width=300
                )

                dpg.add_text("SSH Username:")
                ssh_user = dpg.add_input_text(
                    default_value=self.config.get("ssh_username", ""),
                    callback=lambda s, a: self.update_config("ssh_username", a),
                    width=300
                )

                dpg.add_text("SSH Key Path:")
                with dpg.group(horizontal=True):
                    ssh_key_input = dpg.add_input_text(
                        default_value=self.config.get("ssh_key_path", ""),
                        width=500,
                        callback=lambda s, a: self.update_config("ssh_key_path", a)
                    )
                    browse_button = dpg.add_button(label="...", width=30)

                    # Set up browse button callback
                    def browse_key():
                        with dpg.file_dialog(
                            label="Select SSH Key File",
                            callback=lambda s, a: self._handle_ssh_key_selected(a),
                            width=700,
                            height=400
                        ):
                            pass

                    dpg.set_item_callback(browse_button, browse_key)

            # Sync settings
            with dpg.collapsing_header(label="Sync Settings", default_open=True):
                dpg.add_text("Remote working directory:")
                remote_dir = dpg.add_input_text(
                    default_value=self.config.get("remote_dir", "~/OneTrainer"),
                    callback=lambda s, a: self.update_config("remote_dir", a),
                    width=300
                )

                # Sync options
                dpg.add_checkbox(
                    label="Sync datasets",
                    default_value=self.config.get("sync_datasets", True),
                    callback=lambda s, a: self.update_config("sync_datasets", a)
                )

                dpg.add_checkbox(
                    label="Sync models",
                    default_value=self.config.get("sync_models", True),
                    callback=lambda s, a: self.update_config("sync_models", a)
                )

                dpg.add_checkbox(
                    label="Sync configs",
                    default_value=self.config.get("sync_configs", True),
                    callback=lambda s, a: self.update_config("sync_configs", a)
                )

            # Cloud training controls
            with dpg.group(horizontal=True):
                # Test connection button
                test_button = dpg.add_button(
                    label="Test Connection",
                    callback=self.on_test_connection,
                    width=150,
                    height=40
                )

                # Connect button
                connect_button = dpg.add_button(
                    label="Connect",
                    callback=self.on_connect,
                    width=150,
                    height=40
                )

    def _handle_ssh_key_selected(self, app_data):
        """Handle SSH key file selection"""
        if "file_path_name" in app_data:
            path = app_data["file_path_name"]
            self.update_config("ssh_key_path", path)

    def on_test_connection(self):
        """Test cloud connection"""
        dpg.set_value("status_text", "Testing connection...")

    def on_connect(self):
        """Connect to cloud provider"""
        dpg.set_value("status_text", "Connecting to cloud provider...")

    def run(self):
        """Run the UI"""
        # Configure viewport
        dpg.set_viewport_title("OneTrainer")
        dpg.set_viewport_width(1280)
        dpg.set_viewport_height(800)

        # Show viewport and start the event loop
        dpg.show_viewport()
        dpg.start_dearpygui()
        dpg.destroy_context()

# Main entry point
if __name__ == "__main__":
    try:
        ui = EnhancedOneTrainerUI()
        ui.run()
    except Exception as e:
        print(f"Error: {e}")
        traceback.print_exc()