dearpygui_preset_selector.py

#!/usr/bin/env python3
"""
Standalone example of improved preset selector using Dear PyGui
This module demonstrates a better preset selector implementation that handles
many presets more elegantly than a dropdown, with filtering capabilities.
"""

import os
import json
import dearpygui.dearpygui as dpg
from pathlib import Path


class PresetSelector:
    """An improved preset selector using Dear PyGui with search and filtering capabilities"""
    
    def __init__(self, presets_dir="training_presets", on_select=None):
        """
        Initialize the preset selector component
        
        Args:
            presets_dir: Directory containing the preset JSON files
            on_select: Callback function to call when a preset is selected
        """
        self.presets_dir = presets_dir
        self.on_select_callback = on_select
        
        # State variables
        self.presets = []
        self.filtered_presets = []
        self.selected_index = -1
        self.search_text = ""
        self.current_filter = "All"
        self.available_types = ["All"]
        
        # Item tags
        self.filter_combo_tag = "preset_filter_combo"
        self.search_input_tag = "preset_search_input"
        self.list_tag = "preset_list"
        self.save_input_tag = "preset_save_input"
        
        # Load available presets
        self.load_presets()
        
    def load_presets(self):
        """Load all preset files from the presets directory"""
        self.presets = []
        
        # Add default preset
        self.presets.append({
            "name": "Default Config",
            "path": Path(self.presets_dir) / "#.json",
            "is_builtin": True,
            "type": "Default"
        })
        
        # Scan directory for preset files
        if os.path.isdir(self.presets_dir):
            for filename in os.listdir(self.presets_dir):
                if filename == "#.json":
                    continue  # Skip default
                    
                if filename.endswith(".json"):
                    filepath = Path(self.presets_dir) / filename
                    name = os.path.splitext(filename)[0]
                    
                    # Determine preset type based on name
                    preset_type = "Custom"
                    if name.startswith("#"):
                        preset_type = "Built-in"
                        name = name[1:]  # Remove the # prefix for display
                        
                    # Try to identify the model type from filename
                    if "sd" in name.lower():
                        preset_type = "Stable Diffusion"
                    elif "flux" in name.lower() or "flex" in name.lower():
                        preset_type = "Flux"
                    elif "xl" in name.lower():
                        preset_type = "SDXL"
                    elif "lora" in name.lower():
                        preset_type = "LoRA"
                    elif "embedding" in name.lower():
                        preset_type = "Embedding"
                    elif "cascade" in name.lower():
                        preset_type = "Stable Cascade"
                    elif "wuerstchen" in name.lower():
                        preset_type = "Wuerstchen"
                    elif "pixart" in name.lower():
                        preset_type = "PixArt"
                    
                    # Optional: Load the file to extract more metadata
                    metadata = {}
                    try:
                        with open(filepath, 'r') as f:
                            data = json.load(f)
                            if isinstance(data, dict):
                                # Extract useful fields for filtering/display
                                if "model_type" in data:
                                    metadata["model_type"] = data["model_type"]
                                if "training_method" in data:
                                    metadata["training_method"] = data["training_method"]
                                    # Update type based on training method
                                    if data["training_method"] == "LORA":
                                        preset_type = "LoRA"
                                    elif data["training_method"] == "EMBEDDING":
                                        preset_type = "Embedding"
                    except Exception:
                        pass  # Silently continue if file can't be read
                    
                    self.presets.append({
                        "name": name,
                        "path": filepath,
                        "is_builtin": name.startswith("#"),
                        "type": preset_type,
                        "metadata": metadata
                    })
        
        # Sort presets: Default first, then built-ins, then custom presets
        self.presets.sort(key=lambda x: (
            0 if x["name"] == "Default Config" else 
            (1 if x["is_builtin"] else 2), 
            x["name"]
        ))
        
        # Update filtered list
        self.update_filtered_presets()
        
        # Build list of available preset types
        self.available_types = ["All"]
        type_set = set(preset["type"] for preset in self.presets)
        self.available_types.extend(sorted(type_set))
    
    def update_filtered_presets(self):
        """Update the filtered presets list based on search and filter"""
        search_text = self.search_text.lower()
        
        self.filtered_presets = []
        for preset in self.presets:
            # Skip if search text doesn't match
            if search_text and search_text not in preset["name"].lower():
                continue
                
            # Skip if filter doesn't match
            if self.current_filter != "All" and preset["type"] != self.current_filter:
                continue
                
            self.filtered_presets.append(preset)
        
        # Update the list widget if it exists
        if dpg.does_item_exist(self.list_tag):
            self.update_list_widget()
    
    def on_search_changed(self, sender, app_data):
        """Handle search input changes"""
        self.search_text = app_data
        self.update_filtered_presets()
    
    def on_filter_changed(self, sender, app_data):
        """Handle filter selection changes"""
        self.current_filter = app_data
        self.update_filtered_presets()
    
    def on_preset_selected(self, sender, app_data):
        """Handle preset selection in the list"""
        self.selected_index = app_data
        
        if 0 <= self.selected_index < len(self.filtered_presets):
            selected_preset = self.filtered_presets[self.selected_index]
            
            # Update save input with selected name
            dpg.set_value(self.save_input_tag, selected_preset["name"])
            
            # Call the callback if provided
            if self.on_select_callback:
                self.on_select_callback(selected_preset)
    
    def on_save(self, sender, app_data):
        """Handle saving the current configuration as a preset"""
        new_name = dpg.get_value(self.save_input_tag)
        if not new_name:
            return
            
        # This would normally save the current configuration
        print(f"Saving configuration as: {new_name}")
        
        # Mock implementation: add new preset to the list
        new_preset = {
            "name": new_name,
            "path": Path(self.presets_dir) / f"{new_name}.json",
            "is_builtin": False,
            "type": "Custom"
        }
        
        # Check if preset with this name already exists
        for i, preset in enumerate(self.presets):
            if preset["name"] == new_name:
                self.presets[i] = new_preset
                break
        else:
            # Add new preset
            self.presets.append(new_preset)
            
        # Sort and update filtered list
        self.presets.sort(key=lambda x: (
            0 if x["name"] == "Default Config" else 
            (1 if x["is_builtin"] else 2), 
            x["name"]
        ))
        self.update_filtered_presets()
    
    def update_list_widget(self):
        """Update the list widget with filtered presets"""
        # Clear the list
        dpg.delete_item(self.list_tag, children_only=True)
        
        # Add items to the list
        for i, preset in enumerate(self.filtered_presets):
            name = preset["name"]
            preset_type = preset["type"]
            
            # Format built-in presets differently
            if preset["is_builtin"]:
                dpg.add_selectable(
                    label=f"[Built-in] {name}", 
                    parent=self.list_tag, 
                    callback=lambda s, a, u=i: self.on_preset_selected(s, u),
                    color=[230, 190, 100]
                )
            else:
                dpg.add_selectable(
                    label=f"{name} ({preset_type})", 
                    parent=self.list_tag,
                    callback=lambda s, a, u=i: self.on_preset_selected(s, u)
                )
    
    def create_widget(self, parent):
        """Create the preset selector widget"""
        # Main group
        with dpg.group(parent=parent, horizontal=False):
            # Top controls
            with dpg.group(horizontal=True):
                # Filter dropdown
                dpg.add_text("Type:", color=[200, 200, 200])
                dpg.add_combo(
                    items=self.available_types,
                    default_value=self.current_filter,
                    callback=self.on_filter_changed,
                    width=150,
                    tag=self.filter_combo_tag
                )
                
                # Search box
                dpg.add_text("Search:", color=[200, 200, 200])
                dpg.add_input_text(
                    callback=self.on_search_changed,
                    width=200,
                    tag=self.search_input_tag
                )
            
            # Preset list
            with dpg.child_window(height=300, width=-1, tag=self.list_tag):
                pass  # Will be filled by update_list_widget
            
            # Bottom controls
            with dpg.group(horizontal=True):
                # Save controls
                dpg.add_text("Save as:", color=[200, 200, 200])
                dpg.add_input_text(width=200, tag=self.save_input_tag)
                dpg.add_button(label="Save", callback=self.on_save)
            
            # Initial list update
            self.update_list_widget()


def main():
    """Main function to demonstrate the preset selector"""
    # Initialize DearPyGui
    dpg.create_context()
    dpg.create_viewport(title="OneTrainer Preset Selector", width=600, height=500)
    
    # Create a theme for the UI
    with dpg.theme() as global_theme:
        with dpg.theme_component(dpg.mvAll):
            # Theme similar to CustomTkinter dark-blue
            dpg.add_theme_color(dpg.mvThemeCol_WindowBg, [26, 27, 38, 255])  # Dark blue background
            dpg.add_theme_color(dpg.mvThemeCol_TitleBg, [32, 34, 48, 255])  # Title bar
            dpg.add_theme_color(dpg.mvThemeCol_TitleBgActive, [39, 41, 57, 255])  # Active title
            dpg.add_theme_color(dpg.mvThemeCol_Button, [41, 83, 154, 255])  # Buttons (blue)
            dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, [59, 113, 202, 255])  # Buttons hovered
            dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, [66, 127, 227, 255])  # Buttons active
            dpg.add_theme_color(dpg.mvThemeCol_FrameBg, [37, 38, 51, 255])  # Frame background
            dpg.add_theme_color(dpg.mvThemeCol_ChildBg, [30, 31, 45, 255])  # Frame background
            dpg.add_theme_color(dpg.mvThemeCol_CheckMark, [41, 83, 154, 255])  # Checkboxes
            dpg.add_theme_color(dpg.mvThemeCol_ScrollbarBg, [30, 31, 41, 255])  # Scrollbar bg
            dpg.add_theme_color(dpg.mvThemeCol_ScrollbarGrab, [41, 83, 154, 255])  # Scrollbar
            dpg.add_theme_color(dpg.mvThemeCol_Header, [41, 83, 154, 255])  # Headers
            
            # Add spacing to match CustomTkinter look
            dpg.add_theme_style(dpg.mvStyleVar_FramePadding, [8, 5])
            dpg.add_theme_style(dpg.mvStyleVar_ItemSpacing, [8, 5])
            dpg.add_theme_style(dpg.mvStyleVar_WindowPadding, [10, 10])
            dpg.add_theme_style(dpg.mvStyleVar_FrameRounding, 4.0)  # Rounded corners
            dpg.add_theme_style(dpg.mvStyleVar_GrabRounding, 4.0)  # Rounded corners
    
    dpg.bind_theme(global_theme)
    
    # Callback for preset selection
    def on_preset_selected(preset):
        print(f"Selected preset: {preset['name']} ({preset['path']})")
    
    # Create main window
    with dpg.window(label="OneTrainer Preset Selector", tag="main_window"):
        # Create preset selector
        preset_selector = PresetSelector(on_select=on_preset_selected)
        preset_selector.create_widget("main_window")
        
        # Add explanation
        with dpg.collapsing_header(label="About this example", default_open=True):
            dpg.add_text(
                "This is a demonstration of how Dear PyGui can improve the preset selector experience.\n"
                "The current CustomTkinter implementation uses a dropdown menu which can become\n"
                "unwieldy with many presets and often cuts off long preset names.\n\n"
                "This implementation offers several advantages:\n"
                "- Shows many presets in a scrollable list without truncation\n"
                "- Allows filtering by preset type (SD, SDXL, LoRA, etc.)\n"
                "- Provides search functionality to quickly find presets\n"
                "- Visually distinguishes built-in presets from user presets\n"
                "- Makes more efficient use of screen space\n"
                "- Conforms to the same visual style as the rest of the app"
            )
    
    # Setup and show viewport
    dpg.setup_dearpygui()
    dpg.show_viewport()
    
    # Start main loop
    while dpg.is_dearpygui_running():
        dpg.render_dearpygui_frame()
    
    # Clean up
    dpg.destroy_context()


if __name__ == "__main__":
    main()