fix_dearpygui_versions.py

#!/usr/bin/env python3
"""
Fix DearPyGui version issues for OneTrainer

This script checks the installed version of DearPyGui and installs a compatible
version if needed. It also applies fixes to the code to work with the installed version.
"""

import os
import sys
import importlib.util
import subprocess
import re

def get_dearpygui_version():
    """Get the installed DearPyGui version"""
    try:
        import dearpygui
        version = getattr(dearpygui, "__version__", "unknown")
        print(f"DearPyGui version: {version}")
        return version
    except ImportError:
        print("DearPyGui not installed")
        return None

def install_compatible_version():
    """Install a compatible version of DearPyGui"""
    print("Installing DearPyGui 1.8.0...")
    result = subprocess.run([sys.executable, "-m", "pip", "install", "dearpygui==1.8.0"])
    if result.returncode == 0:
        print("DearPyGui 1.8.0 installed successfully")
        return True
    else:
        print("Failed to install DearPyGui 1.8.0")
        return False

def fix_basic_components_for_v2():
    """Fix basic components for DearPyGui 2.0"""
    basic_components_path = os.path.join("dpg_ui", "components", "basic.py")
    if not os.path.exists(basic_components_path):
        print(f"Could not find {basic_components_path}")
        return False
    
    # Create backup
    backup_path = f"{basic_components_path}.bak"
    if not os.path.exists(backup_path):
        with open(basic_components_path, "r") as src:
            with open(backup_path, "w") as dst:
                dst.write(src.read())
        print(f"Created backup: {backup_path}")
    
    # Read the file
    with open(basic_components_path, "r") as f:
        content = f.read()
    
    # Fix flags
    content = re.sub(r'dialog_flags\s*=\s*dpg\.mvFileDialogFlags_\w+', 'dialog_flags = 0', content)
    
    # Remove flags
    content = re.sub(r'flags=dialog_flags,', '', content)
    
    # Write back
    with open(basic_components_path, "w") as f:
        f.write(content)
    
    print(f"Fixed {basic_components_path} for DearPyGui 2.0")
    return True

def fix_tabs_for_v2():
    """Fix tabs for DearPyGui 2.0"""
    # Fix additional_embeddings_tab.py
    tab_path = os.path.join("dpg_ui", "tabs", "additional_embeddings_tab.py")
    if os.path.exists(tab_path):
        # Create backup
        backup_path = f"{tab_path}.bak"
        if not os.path.exists(backup_path):
            with open(tab_path, "r") as src:
                with open(backup_path, "w") as dst:
                    dst.write(src.read())
            print(f"Created backup: {backup_path}")
        
        # Read the file
        with open(tab_path, "r") as f:
            content = f.read()
        
        # Replace file_entry calls with a simpler implementation
        content = re.sub(
            r'llama_path\s*=\s*Components\.file_entry\(\s*dpg\.last_item\(\),\s*default_path="",\s*width=300,\s*tooltip="[^"]+",\s*on_change=[^)]+\)',
            '''# Use a simpler approach to avoid file dialog issues
                    input_id = dpg.add_input_text(
                        default_value="",
                        width=250,
                        callback=self.on_llama_path_change
                    )
                    browse_button = dpg.add_button(
                        label="...",
                        width=30,
                        height=30
                    )
                    
                    # Set up button callback
                    def show_file_dialog():
                        with dpg.file_dialog(
                            label="Select LLaMA Model",
                            width=700,
                            height=400,
                            callback=lambda s, a: dpg.set_value(input_id, a.get("file_path_name", "")),
                            show=True
                        ):
                            pass
                            
                    dpg.set_item_callback(browse_button, show_file_dialog)''',
            content
        )
        
        # Write back
        with open(tab_path, "w") as f:
            f.write(content)
        
        print(f"Fixed {tab_path} for DearPyGui 2.0")
    
    return True

def main():
    """Main entry point"""
    print("OneTrainer DearPyGui Version Fixer")
    print("==================================")
    
    # Get current version
    version = get_dearpygui_version()
    
    # If no version found, try to install
    if not version:
        print("Installing DearPyGui...")
        if not install_compatible_version():
            print("Failed to install DearPyGui. Please install it manually.")
            return 1
    
    # If version 2.x, apply fixes for v2
    if version and version.startswith("2."):
        print(f"Detected DearPyGui {version}. Applying fixes...")
        
        # Fix basic components
        fix_basic_components_for_v2()
        
        # Fix tabs
        fix_tabs_for_v2()
        
        print("All fixes applied for DearPyGui 2.x")
        
        # Ask to install 1.8.0
        if input("Do you want to install DearPyGui 1.8.0 for better compatibility? (y/n): ").lower() == "y":
            install_compatible_version()
    
    print("Finished. You can now run the UI with: ./fixed_run_dpg_ui.sh")
    return 0

if __name__ == "__main__":
    sys.exit(main())