comfy_fm24_newgens/gui.py

#!/usr/bin/env python3
"""
Standalone FM NewGens - GUI Interface
A user-friendly graphical interface for generating Football Manager faces
"""

import tkinter as tk
from tkinter import ttk, filedialog, messagebox, scrolledtext
import os
import sys
import json
import configparser
import threading
import queue
import time
from pathlib import Path
from PIL import Image, ImageTk
import logging

# Import the main generation logic
from comfy_fm_newgen import generate_image, get_country_name, generate_prompts_for_players, get_prompt_for_image
try:
    from lib.rtf_parser import RTF_Parser
except ImportError:
    # Fallback for when striprtf is not available
    class RTF_Parser:
        def parse_rtf(self, file_path):
            print(f"Warning: RTF parsing not available. Please install striprtf: pip install striprtf==0.0.28")
            return []

# Check if background removal is available
try:
    from lib.remove_bg import remove_bg_from_file_list, REMBG_AVAILABLE
    BG_REMOVAL_AVAILABLE = REMBG_AVAILABLE
except ImportError:
    BG_REMOVAL_AVAILABLE = False
    print("Warning: Background removal not available")
from lib.generate_xml import create_config_xml, append_to_config_xml
from lib.xml_reader import extract_from_values
from lib.general import list_profiles

class FMFaceGeneratorGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("FM Face Generator - Standalone Version")
        self.root.geometry("1000x700")
        self.root.minsize(800, 600)

        # Set environment variable to indicate GUI mode
        os.environ['FM_NEWGEN_GUI_MODE'] = 'true'

        # Initialize variables
        self.config = None
        self.app_config = None
        self.rtf = RTF_Parser()
        self.generation_thread = None
        self.stop_generation = False
        self.image_queue = queue.Queue()
        self.current_profile = None

        # Preview tracking variables
        self.generated_images = []  # List of generated image UIDs
        self.current_image_index = -1  # Current image being displayed

        # Set up the GUI
        self.setup_gui()

        # Show profile selection dialog first
        self.show_profile_selection()

        # Refresh system information
        self.refresh_system_info()

        # Start image preview update loop
        self.update_image_preview()

    def show_profile_selection(self):
        """Show profile selection dialog at startup"""
        # Load existing profiles
        profiles = self.list_profiles()

        if not profiles:
            # No profiles exist, create default one
            self.create_default_profile()
            self.current_profile = "NewGens"
            self.load_config()
            return

        # Create selection dialog
        dialog = tk.Toplevel(self.root)
        dialog.title("Select Profile")
        dialog.geometry("400x300")
        dialog.transient(self.root)
        dialog.grab_set()

        # Center the dialog
        dialog.update_idletasks()
        x = (dialog.winfo_screenwidth() - dialog.winfo_width()) // 2
        y = (dialog.winfo_screenheight() - dialog.winfo_height()) // 2
        dialog.geometry(f"+{x}+{y}")

        # Profile selection
        ttk.Label(dialog, text="Select a profile to use:").pack(pady=10)

        profile_var = tk.StringVar()
        profile_combo = ttk.Combobox(dialog, textvariable=profile_var, width=30)
        profile_combo['values'] = profiles
        profile_combo.pack(pady=10)
        profile_combo.set(profiles[0])  # Select first profile by default

        def on_select():
            selected = profile_var.get()
            if selected:
                self.current_profile = selected
                dialog.destroy()
                self.load_config()
                self.log_message(f"Selected profile: {selected}")

        def on_create_new():
            dialog.destroy()
            self.show_create_profile_dialog()

        def on_manage():
            dialog.destroy()
            self.notebook.select(1)  # Switch to profiles tab

        # Buttons
        btn_frame = ttk.Frame(dialog)
        btn_frame.pack(pady=20)

        ttk.Button(btn_frame, text="Select", command=on_select).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Create New", command=on_create_new).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Manage Profiles", command=on_manage).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Cancel", command=self.root.quit).pack(side=tk.LEFT, padx=5)

        # Wait for dialog to close
        self.root.wait_window(dialog)

    def show_create_profile_dialog(self):
        """Show dialog to create a new profile"""
        dialog = tk.Toplevel(self.root)
        dialog.title("Create New Profile")
        dialog.geometry("400x250")
        dialog.transient(self.root)
        dialog.grab_set()

        # Center the dialog
        dialog.update_idletasks()
        x = (dialog.winfo_screenwidth() - dialog.winfo_width()) // 2
        y = (dialog.winfo_screenheight() - dialog.winfo_height()) // 2
        dialog.geometry(f"+{x}+{y}")

        # Profile name
        ttk.Label(dialog, text="Profile Name:").grid(row=0, column=0, sticky=tk.W, pady=10, padx=10)
        profile_name_var = tk.StringVar()
        profile_name_entry = ttk.Entry(dialog, textvariable=profile_name_var, width=30)
        profile_name_entry.grid(row=0, column=1, sticky=tk.W, padx=10, pady=10)

        # FM Version
        ttk.Label(dialog, text="FM Version:").grid(row=1, column=0, sticky=tk.W, pady=5, padx=10)
        fm_version_var = tk.StringVar(value="2024")
        fm_version_entry = ttk.Entry(dialog, textvariable=fm_version_var, width=30)
        fm_version_entry.grid(row=1, column=1, sticky=tk.W, padx=10, pady=5)

        # RTF File
        ttk.Label(dialog, text="RTF File:").grid(row=2, column=0, sticky=tk.W, pady=5, padx=10)
        rtf_file_var = tk.StringVar()
        rtf_entry = ttk.Entry(dialog, textvariable=rtf_file_var, width=30)
        rtf_entry.grid(row=2, column=1, sticky=tk.W, padx=10, pady=5)
        ttk.Button(dialog, text="Browse...", command=lambda: self.browse_file(rtf_file_var, "RTF files", "*.rtf")).grid(row=2, column=2, padx=10, pady=5)

        # Output Directory
        ttk.Label(dialog, text="Output Directory:").grid(row=3, column=0, sticky=tk.W, pady=5, padx=10)
        output_dir_var = tk.StringVar()
        output_entry = ttk.Entry(dialog, textvariable=output_dir_var, width=30)
        output_entry.grid(row=3, column=1, sticky=tk.W, padx=10, pady=5)
        ttk.Button(dialog, text="Browse...", command=lambda: self.browse_directory(output_dir_var)).grid(row=3, column=2, padx=10, pady=5)

        def on_create():
            name = profile_name_var.get().strip()
            if not name:
                messagebox.showerror("Error", "Profile name cannot be empty")
                return

            if self.profile_exists(name):
                messagebox.showerror("Error", f"Profile '{name}' already exists")
                return

            # Create the profile
            self.create_profile(name, fm_version_var.get(), rtf_file_var.get(), output_dir_var.get())
            self.current_profile = name
            dialog.destroy()
            self.load_config()
            self.log_message(f"Created and selected profile: {name}")

        def on_cancel():
            dialog.destroy()
            self.show_profile_selection()  # Go back to profile selection

        # Buttons
        btn_frame = ttk.Frame(dialog)
        btn_frame.grid(row=4, column=0, columnspan=3, pady=20)

        ttk.Button(btn_frame, text="Create", command=on_create).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Cancel", command=on_cancel).pack(side=tk.LEFT, padx=5)

    def list_profiles(self):
        """List all available profiles"""
        if not os.path.exists('user_config.cfg'):
            return []

        config = configparser.ConfigParser()
        config.read('user_config.cfg')

        profiles = [section.split(':', 1)[1] for section in config.sections() if section.startswith('profile:')]
        return profiles

    def profile_exists(self, profile_name):
        """Check if a profile exists"""
        profiles = self.list_profiles()
        return profile_name in profiles

    def create_profile(self, name, fm_version, rtf_file, output_dir):
        """Create a new profile"""
        if not self.config:
            self.config = configparser.ConfigParser()
            self.config.add_section('models')

        section_name = f'profile:{name}'
        self.config.add_section(section_name)
        self.config.set(section_name, 'football_manager_version', fm_version)
        self.config.set(section_name, 'rtf_file', rtf_file)
        self.config.set(section_name, 'output_dir', output_dir)

        # Save to file
        with open('user_config.cfg', 'w') as configfile:
            self.config.write(configfile)

    def create_default_profile(self):
        """Create a default profile if none exists"""
        self.config = configparser.ConfigParser()
        self.config.add_section('models')
        self.config.add_section('profile:NewGens')

        self.config.set('models', 'model_name', 'SG161222/Realistic_Vision_V6.0_B1')
        self.config.set('profile:NewGens', 'football_manager_version', '2024')
        self.config.set('profile:NewGens', 'rtf_file', '')
        self.config.set('profile:NewGens', 'output_dir', './NewGens/')

        with open('user_config.cfg', 'w') as configfile:
            self.config.write(configfile)

    def browse_file(self, var, title, filetypes):
        """Browse for a file"""
        file_path = filedialog.askopenfilename(title=title, filetypes=[(title, filetypes), ("All files", "*.*")])
        if file_path:
            var.set(file_path)

    def browse_directory(self, var):
        """Browse for a directory"""
        directory = filedialog.askdirectory(title="Select Directory")
        if directory:
            var.set(directory)

    def setup_profiles_tab(self):
        """Set up the profiles management tab"""
        # Profile list
        list_frame = ttk.LabelFrame(self.profiles_frame, text="Available Profiles", padding=10)
        list_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)

        # Profile listbox
        self.profiles_listbox = tk.Listbox(list_frame, height=10)
        self.profiles_listbox.pack(fill=tk.BOTH, expand=True, pady=5)

        # Buttons
        btn_frame = ttk.Frame(list_frame)
        btn_frame.pack(fill=tk.X, pady=5)

        ttk.Button(btn_frame, text="Refresh List", command=self.refresh_profiles_list).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Select Profile", command=self.select_profile_from_list).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Create New", command=self.show_create_profile_dialog).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="Delete", command=self.delete_profile).pack(side=tk.RIGHT, padx=5)

        # Profile details
        details_frame = ttk.LabelFrame(self.profiles_frame, text="Profile Details", padding=10)
        details_frame.pack(fill=tk.X, padx=10, pady=5)

        # Current profile display
        self.current_profile_var = tk.StringVar()
        ttk.Label(details_frame, text="Current Profile:").grid(row=0, column=0, sticky=tk.W, pady=5)
        ttk.Label(details_frame, textvariable=self.current_profile_var).grid(row=0, column=1, sticky=tk.W, padx=5, pady=5)

        # Refresh the profiles list
        self.refresh_profiles_list()

    def refresh_profiles_list(self):
        """Refresh the profiles listbox"""
        self.profiles_listbox.delete(0, tk.END)
        profiles = self.list_profiles()
        for profile in profiles:
            self.profiles_listbox.insert(tk.END, profile)

        if self.current_profile:
            self.current_profile_var.set(self.current_profile)

    def refresh_system_info(self):
        """Refresh and display system information"""
        try:
            import torch

            # PyTorch version
            self.pytorch_version_var.set(torch.__version__)

            # CUDA availability
            cuda_available = torch.cuda.is_available()
            self.cuda_available_var.set("Yes" if cuda_available else "No")

            # Check CUDA availability
            if cuda_available:
                gpu_count = torch.cuda.device_count()
                current_device = torch.cuda.current_device()
                gpu_name = torch.cuda.get_device_name(current_device)
                gpu_memory = torch.cuda.get_device_properties(current_device).total_memory / 1024**3  # GB

                self.gpu_status_var.set(f"Available ({gpu_name})")
                self.gpu_memory_var.set(f"{gpu_memory:.1f} GB")

                # Check if forcing CPU
                if self.force_cpu_var.get() == "true":
                    self.gpu_status_var.set(f"Available but disabled (CPU forced)")
                    self.log_message(f"GPU available: {gpu_name} ({gpu_memory:.1f} GB) - CPU mode forced")
                else:
                    self.log_message(f"GPU available: {gpu_name} ({gpu_memory:.1f} GB)")
            else:
                self.gpu_status_var.set("Not available (using CPU)")
                self.gpu_memory_var.set("N/A")
                self.log_message("GPU not available - using CPU. Install CUDA for GPU support.")

        except ImportError:
            self.gpu_status_var.set("PyTorch not available")
            self.gpu_memory_var.set("N/A")
            self.pytorch_version_var.set("Not installed")
            self.cuda_available_var.set("N/A")
            self.log_message("PyTorch not available - cannot detect GPU status")
        except Exception as e:
            self.gpu_status_var.set("Error detecting GPU")
            self.gpu_memory_var.set("N/A")
            self.pytorch_version_var.set("Error")
            self.cuda_available_var.set("Error")
            self.log_message(f"Error detecting GPU: {str(e)}")

    def update_gpu_settings(self):
        """Update GPU settings when force CPU option changes"""
        force_cpu = self.force_cpu_var.get() == "true"
        os.environ['FM_NEWGEN_FORCE_CPU'] = 'true' if force_cpu else 'false'

        if force_cpu:
            self.log_message("GPU usage disabled - will use CPU")
        else:
            self.log_message("GPU usage enabled (if available)")
            self.refresh_system_info()  # Refresh to show actual status

    def update_gpu_memory_usage(self):
        """Update GPU memory usage display during generation"""
        try:
            import torch
            if torch.cuda.is_available() and not self.force_cpu_var.get() == "true":
                # Get current GPU memory usage
                if torch.cuda.is_available():
                    allocated = torch.cuda.memory_allocated(0) / 1024**3  # GB
                    reserved = torch.cuda.memory_reserved(0) / 1024**3   # GB
                    total = torch.cuda.get_device_properties(0).total_memory / 1024**3  # GB

                    self.gpu_usage_var.set(f"GPU Memory: {allocated:.1f}GB / {total:.1f}GB (Reserved: {reserved:.1f}GB)")
                else:
                    self.gpu_usage_var.set("GPU Memory: N/A")
            else:
                self.gpu_usage_var.set("GPU Memory: CPU Mode")
        except:
            self.gpu_usage_var.set("GPU Memory: N/A")

        # Schedule next update if generation is running
        if hasattr(self, 'generation_thread') and self.generation_thread and self.generation_thread.is_alive():
            self.root.after(2000, self.update_gpu_memory_usage)  # Update every 2 seconds during generation

    def select_profile_from_list(self):
        """Select a profile from the list"""
        selection = self.profiles_listbox.curselection()
        if selection:
            profile_name = self.profiles_listbox.get(selection[0])
            self.current_profile = profile_name
            self.load_config()
            self.log_message(f"Selected profile: {profile_name}")
            self.refresh_profiles_list()

    def delete_profile(self):
        """Delete the selected profile"""
        selection = self.profiles_listbox.curselection()
        if not selection:
            messagebox.showerror("Error", "Please select a profile to delete")
            return

        profile_name = self.profiles_listbox.get(selection[0])

        if profile_name == self.current_profile:
            messagebox.showerror("Error", "Cannot delete the currently active profile")
            return

        if messagebox.askyesno("Confirm Delete", f"Are you sure you want to delete profile '{profile_name}'?"):
            # Remove from config
            if self.config:
                self.config.remove_section(f'profile:{profile_name}')

                with open('user_config.cfg', 'w') as configfile:
                    self.config.write(configfile)

            self.log_message(f"Deleted profile: {profile_name}")
            self.refresh_profiles_list()

    def setup_gui(self):
        """Set up the main GUI components"""
        # Create main notebook (tabs)
        self.notebook = ttk.Notebook(self.root)
        self.notebook.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        # Configuration tab
        self.config_frame = ttk.Frame(self.notebook)
        self.notebook.add(self.config_frame, text="Configuration")

        # Profiles tab
        self.profiles_frame = ttk.Frame(self.notebook)
        self.notebook.add(self.profiles_frame, text="Profiles")

        # Generation tab
        self.generation_frame = ttk.Frame(self.notebook)
        self.notebook.add(self.generation_frame, text="Generation")

        # Preview tab
        self.preview_frame = ttk.Frame(self.notebook)
        self.notebook.add(self.preview_frame, text="Preview")

        # Setup each tab
        self.setup_config_tab()
        self.setup_profiles_tab()
        self.setup_generation_tab()
        self.setup_preview_tab()

        # Status bar
        self.status_var = tk.StringVar()
        self.status_var.set("Ready")
        self.status_bar = ttk.Label(self.root, textvariable=self.status_var, relief=tk.SUNKEN)
        self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)

        # Progress bar
        self.progress_var = tk.DoubleVar()
        self.progress_bar = ttk.Progressbar(self.root, variable=self.progress_var, maximum=100)
        self.progress_bar.pack(side=tk.BOTTOM, fill=tk.X)

    def setup_config_tab(self):
        """Set up the configuration tab"""
        # Model selection
        model_frame = ttk.LabelFrame(self.config_frame, text="AI Model", padding=10)
        model_frame.pack(fill=tk.X, padx=10, pady=5)

        ttk.Label(model_frame, text="Model:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.model_var = tk.StringVar()
        model_combo = ttk.Combobox(model_frame, textvariable=self.model_var, width=50)
        model_combo['values'] = [
            'SG161222/Realistic_Vision_V6.0_B1',
            'digiplay/AbsoluteReality',
            'stabilityai/stable-diffusion-2-1',
            'runwayml/stable-diffusion-v1-5'
        ]
        model_combo.grid(row=0, column=1, sticky=tk.W, padx=5, pady=5)
        model_combo.set('SG161222/Realistic_Vision_V6.0_B1')

        # Model directory
        ttk.Label(model_frame, text="Model Directory:").grid(row=1, column=0, sticky=tk.W, pady=5)
        self.model_dir_var = tk.StringVar()
        model_dir_entry = ttk.Entry(model_frame, textvariable=self.model_dir_var, width=50)
        model_dir_entry.grid(row=1, column=1, sticky=tk.W, padx=5, pady=5)
        ttk.Button(model_frame, text="Browse...", command=self.browse_model_dir).grid(row=1, column=2, padx=5, pady=5)

        # FM Configuration
        fm_frame = ttk.LabelFrame(self.config_frame, text="Football Manager", padding=10)
        fm_frame.pack(fill=tk.X, padx=10, pady=5)

        ttk.Label(fm_frame, text="FM Version:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.fm_version_var = tk.StringVar(value="2024")
        fm_version_entry = ttk.Entry(fm_frame, textvariable=self.fm_version_var, width=20)
        fm_version_entry.grid(row=0, column=1, sticky=tk.W, padx=5, pady=5)

        ttk.Label(fm_frame, text="RTF File:").grid(row=1, column=0, sticky=tk.W, pady=5)
        self.rtf_file_var = tk.StringVar()
        rtf_entry = ttk.Entry(fm_frame, textvariable=self.rtf_file_var, width=50)
        rtf_entry.grid(row=1, column=1, sticky=tk.W, padx=5, pady=5)
        ttk.Button(fm_frame, text="Browse...", command=self.browse_rtf_file).grid(row=1, column=2, padx=5, pady=5)

        ttk.Label(fm_frame, text="Output Directory:").grid(row=2, column=0, sticky=tk.W, pady=5)
        self.output_dir_var = tk.StringVar()
        output_entry = ttk.Entry(fm_frame, textvariable=self.output_dir_var, width=50)
        output_entry.grid(row=2, column=1, sticky=tk.W, padx=5, pady=5)
        ttk.Button(fm_frame, text="Browse...", command=self.browse_output_dir).grid(row=2, column=2, padx=5, pady=5)

        # System Information
        system_frame = ttk.LabelFrame(self.config_frame, text="System Information", padding=10)
        system_frame.pack(fill=tk.X, padx=10, pady=5)

        # GPU Status
        self.gpu_status_var = tk.StringVar()
        self.gpu_status_var.set("Checking GPU status...")
        ttk.Label(system_frame, text="GPU Status:").grid(row=0, column=0, sticky=tk.W, pady=2)
        ttk.Label(system_frame, textvariable=self.gpu_status_var).grid(row=0, column=1, sticky=tk.W, padx=5, pady=2)

        # GPU Memory
        self.gpu_memory_var = tk.StringVar()
        self.gpu_memory_var.set("N/A")
        ttk.Label(system_frame, text="GPU Memory:").grid(row=1, column=0, sticky=tk.W, pady=2)
        ttk.Label(system_frame, textvariable=self.gpu_memory_var).grid(row=1, column=1, sticky=tk.W, padx=5, pady=2)

        # Force CPU option
        self.force_cpu_var = tk.StringVar(value="false")
        ttk.Checkbutton(system_frame, text="Force CPU usage (disable GPU)", variable=self.force_cpu_var, onvalue="true", offvalue="false", command=self.update_gpu_settings).grid(row=2, column=0, columnspan=2, sticky=tk.W, pady=5)

        # System details
        ttk.Label(system_frame, text="PyTorch Version:").grid(row=4, column=0, sticky=tk.W, pady=2)
        self.pytorch_version_var = tk.StringVar()
        self.pytorch_version_var.set("Checking...")
        ttk.Label(system_frame, textvariable=self.pytorch_version_var).grid(row=4, column=1, sticky=tk.W, padx=5, pady=2)

        ttk.Label(system_frame, text="CUDA Available:").grid(row=5, column=0, sticky=tk.W, pady=2)
        self.cuda_available_var = tk.StringVar()
        self.cuda_available_var.set("Checking...")
        ttk.Label(system_frame, textvariable=self.cuda_available_var).grid(row=5, column=1, sticky=tk.W, padx=5, pady=2)

        # Refresh system info button
        ttk.Button(system_frame, text="Refresh System Info", command=self.refresh_system_info).grid(row=6, column=0, columnspan=2, pady=5)

        # Control buttons
        button_frame = ttk.Frame(self.config_frame)
        button_frame.pack(fill=tk.X, padx=10, pady=10)

        ttk.Button(button_frame, text="Save Configuration", command=self.save_config).pack(side=tk.LEFT, padx=5)
        ttk.Button(button_frame, text="Load Configuration", command=self.load_config).pack(side=tk.LEFT, padx=5)
        ttk.Button(button_frame, text="Reset to Defaults", command=self.reset_config).pack(side=tk.LEFT, padx=5)

    def setup_generation_tab(self):
        """Set up the generation tab"""
        # Player selection
        player_frame = ttk.LabelFrame(self.generation_frame, text="Player Selection", padding=10)
        player_frame.pack(fill=tk.X, padx=10, pady=5)

        self.player_mode_var = tk.StringVar(value="all")
        ttk.Radiobutton(player_frame, text="Process all players", variable=self.player_mode_var, value="all", command=self.update_player_controls).grid(row=0, column=0, sticky=tk.W, pady=5)
        ttk.Radiobutton(player_frame, text="Process specific player", variable=self.player_mode_var, value="specific", command=self.update_player_controls).grid(row=0, column=1, sticky=tk.W, pady=5)

        ttk.Label(player_frame, text="Player UID:").grid(row=1, column=0, sticky=tk.W, pady=5)
        self.player_uid_var = tk.StringVar()
        self.player_uid_entry = ttk.Entry(player_frame, textvariable=self.player_uid_var, width=20, state=tk.DISABLED)
        self.player_uid_entry.grid(row=1, column=1, sticky=tk.W, padx=5, pady=5)

        # Generation options
        gen_frame = ttk.LabelFrame(self.generation_frame, text="Generation Options", padding=10)
        gen_frame.pack(fill=tk.X, padx=10, pady=5)

        self.update_mode_var = tk.StringVar(value="false")
        ttk.Checkbutton(gen_frame, text="Update existing (skip already processed)", variable=self.update_mode_var, onvalue="true", offvalue="false").grid(row=0, column=0, sticky=tk.W, pady=5)

        # GPU Memory usage during generation
        self.gpu_usage_var = tk.StringVar()
        self.gpu_usage_var.set("GPU Memory: N/A")
        ttk.Label(gen_frame, textvariable=self.gpu_usage_var).grid(row=1, column=0, sticky=tk.W, pady=5)

        # Progress display
        progress_frame = ttk.LabelFrame(self.generation_frame, text="Progress", padding=10)
        progress_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)

        # Status text
        self.status_text = scrolledtext.ScrolledText(progress_frame, height=10, wrap=tk.WORD)
        self.status_text.pack(fill=tk.BOTH, expand=True, pady=5)

        # Control buttons
        control_frame = ttk.Frame(self.generation_frame)
        control_frame.pack(fill=tk.X, padx=10, pady=10)

        self.start_button = ttk.Button(control_frame, text="Start Generation", command=self.start_generation)
        self.start_button.pack(side=tk.LEFT, padx=5)

        self.stop_button = ttk.Button(control_frame, text="Stop Generation", command=self.stop_generation_thread, state=tk.DISABLED)
        self.stop_button.pack(side=tk.LEFT, padx=5)

        ttk.Button(control_frame, text="Clear Log", command=self.clear_log).pack(side=tk.RIGHT, padx=5)

    def setup_preview_tab(self):
        """Set up the preview tab"""
        # Image preview
        preview_frame = ttk.LabelFrame(self.preview_frame, text="Generated Images", padding=10)
        preview_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)

        # Image display
        self.image_label = ttk.Label(preview_frame, text="No images generated yet")
        self.image_label.pack(pady=20)

        # Image info
        info_frame = ttk.Frame(preview_frame)
        info_frame.pack(fill=tk.X, pady=10)

        ttk.Label(info_frame, text="Current Image:").grid(row=0, column=0, sticky=tk.W)
        self.current_image_var = tk.StringVar()
        ttk.Label(info_frame, textvariable=self.current_image_var).grid(row=0, column=1, sticky=tk.W, padx=5)

        ttk.Label(info_frame, text="Total Generated:").grid(row=1, column=0, sticky=tk.W)
        self.total_generated_var = tk.StringVar(value="0")
        ttk.Label(info_frame, textvariable=self.total_generated_var).grid(row=1, column=1, sticky=tk.W, padx=5)

        # Prompt display
        prompt_frame = ttk.LabelFrame(preview_frame, text="Prompt Used", padding=10)
        prompt_frame.pack(fill=tk.X, padx=10, pady=5)

        self.prompt_text = scrolledtext.ScrolledText(prompt_frame, height=4, wrap=tk.WORD)
        self.prompt_text.pack(fill=tk.X, pady=5)
        self.prompt_text.config(state=tk.DISABLED)  # Make it read-only

        # Navigation buttons
        nav_frame = ttk.Frame(preview_frame)
        nav_frame.pack(fill=tk.X, pady=10)

        ttk.Button(nav_frame, text="Previous", command=self.show_previous_image).pack(side=tk.LEFT, padx=5)
        ttk.Button(nav_frame, text="Next", command=self.show_next_image).pack(side=tk.LEFT, padx=5)
        ttk.Button(nav_frame, text="Refresh", command=self.refresh_preview).pack(side=tk.RIGHT, padx=5)

        # Scan for existing images on startup
        self.root.after(1000, self.scan_existing_images)  # Delay to allow config to load

    def browse_model_dir(self):
        """Browse for model directory"""
        directory = filedialog.askdirectory(title="Select Model Directory")
        if directory:
            self.model_dir_var.set(directory)

    def browse_rtf_file(self):
        """Browse for RTF file"""
        file_path = filedialog.askopenfilename(
            title="Select RTF File",
            filetypes=[("RTF files", "*.rtf"), ("All files", "*.*")]
        )
        if file_path:
            self.rtf_file_var.set(file_path)

    def browse_output_dir(self):
        """Browse for output directory"""
        directory = filedialog.askdirectory(title="Select Output Directory")
        if directory:
            self.output_dir_var.set(directory)

    def update_player_controls(self):
        """Update player selection controls based on mode"""
        if self.player_mode_var.get() == "specific":
            self.player_uid_entry.config(state=tk.NORMAL)
        else:
            self.player_uid_entry.config(state=tk.DISABLED)
            self.player_uid_var.set("")

    def save_config(self):
        """Save configuration to file"""
        try:
            # Update config object
            if not self.config:
                self.config = configparser.ConfigParser()
                self.config.add_section('models')

            # Use current profile or create NewGens if none selected
            profile_name = self.current_profile or "NewGens"
            profile_section = f'profile:{profile_name}'

            if not self.config.has_section(profile_section):
                self.config.add_section(profile_section)

            self.config.set('models', 'model_name', self.model_var.get())
            if self.model_dir_var.get():
                self.config.set('models', 'model_dir', self.model_dir_var.get())

            self.config.set(profile_section, 'football_manager_version', self.fm_version_var.get())
            self.config.set(profile_section, 'rtf_file', self.rtf_file_var.get())
            self.config.set(profile_section, 'output_dir', self.output_dir_var.get())

            # Save to file
            with open('user_config.cfg', 'w') as configfile:
                self.config.write(configfile)

            messagebox.showinfo("Success", f"Configuration saved to profile '{profile_name}'!")
            self.log_message(f"Configuration saved to profile '{profile_name}'")

        except Exception as e:
            messagebox.showerror("Error", f"Failed to save configuration: {str(e)}")
            self.log_message(f"Error saving configuration: {str(e)}")

    def load_config(self):
        """Load configuration from file"""
        try:
            if os.path.exists('user_config.cfg'):
                self.config = configparser.ConfigParser()
                self.config.read('user_config.cfg')

                # Load model settings
                if 'models' in self.config:
                    self.model_var.set(self.config.get('models', 'model_name', fallback='SG161222/Realistic_Vision_V6.0_B1'))
                    self.model_dir_var.set(self.config.get('models', 'model_dir', fallback=''))

                # Load current profile settings
                if self.current_profile:
                    profile_section = f'profile:{self.current_profile}'
                    if profile_section in self.config:
                        self.fm_version_var.set(self.config.get(profile_section, 'football_manager_version', fallback='2024'))
                        self.rtf_file_var.set(self.config.get(profile_section, 'rtf_file', fallback=''))
                        self.output_dir_var.set(self.config.get(profile_section, 'output_dir', fallback=''))
                        self.log_message(f"Configuration loaded for profile '{self.current_profile}'")
                    else:
                        # Profile doesn't exist, create it with defaults
                        self.create_profile(self.current_profile, '2024', '', './NewGens/')
                        self.fm_version_var.set('2024')
                        self.rtf_file_var.set('')
                        self.output_dir_var.set('./NewGens/')
                        self.log_message(f"Created new profile '{self.current_profile}' with defaults")
                else:
                    # No profile selected, use defaults
                    self.fm_version_var.set('2024')
                    self.rtf_file_var.set('')
                    self.output_dir_var.set('./NewGens/')
                    self.log_message("No profile selected, using defaults")

                # Scan for existing images after loading config
                self.root.after(500, self.scan_existing_images)
            else:
                self.log_message("No configuration file found, using defaults")

        except Exception as e:
            messagebox.showerror("Error", f"Failed to load configuration: {str(e)}")
            self.log_message(f"Error loading configuration: {str(e)}")

    def reset_config(self):
        """Reset configuration to defaults"""
        self.model_var.set('SG161222/Realistic_Vision_V6.0_B1')
        self.model_dir_var.set('')
        self.fm_version_var.set('2024')
        self.rtf_file_var.set('')
        self.output_dir_var.set('./NewGens/')

        # Update current profile with defaults
        if self.current_profile:
            profile_section = f'profile:{self.current_profile}'
            if self.config and self.config.has_section(profile_section):
                self.config.set(profile_section, 'football_manager_version', '2024')
                self.config.set(profile_section, 'rtf_file', '')
                self.config.set(profile_section, 'output_dir', './NewGens/')

                with open('user_config.cfg', 'w') as configfile:
                    self.config.write(configfile)

        self.log_message("Configuration reset to defaults")

    def log_message(self, message):
        """Add message to log"""
        self.status_text.insert(tk.END, f"[{time.strftime('%H:%M:%S')}] {message}\n")
        self.status_text.see(tk.END)

    def clear_log(self):
        """Clear the log"""
        self.status_text.delete(1.0, tk.END)

    def start_generation(self):
        """Start the generation process"""
        if not self.validate_config():
            return

        # Set environment variables for the main script
        os.environ['FM_NEWGEN_UPDATE_MODE'] = 'true' if self.update_mode_var.get() == "true" else 'false'
        os.environ['FM_NEWGEN_PROCESS_PLAYER'] = 'true' if self.player_mode_var.get() == "specific" else 'false'
        if self.player_mode_var.get() == "specific":
            os.environ['FM_NEWGEN_PLAYER_UID'] = self.player_uid_var.get()

        # Disable start button, enable stop button
        self.start_button.config(state=tk.DISABLED)
        self.stop_button.config(state=tk.NORMAL)

        # Reset progress
        self.progress_var.set(0)
        self.stop_generation = False

        # Start generation thread
        self.generation_thread = threading.Thread(target=self.generation_worker)
        self.generation_thread.daemon = True
        self.generation_thread.start()

        # Start GPU memory monitoring
        self.update_gpu_memory_usage()

    def stop_generation_thread(self):
        """Stop the generation process"""
        self.stop_generation = True
        self.stop_button.config(state=tk.DISABLED)
        self.gpu_usage_var.set("GPU Memory: N/A")  # Reset GPU usage display
        self.log_message("Stopping generation...")

    def generation_worker(self):
        """Worker thread for generation process"""
        try:
            self.log_message("Starting generation process...")

            # Load configurations
            with open("app_config.json", "r") as f:
                app_config = json.load(f)

            # Parse RTF file
            rtf_file = self.rtf.parse_rtf(self.rtf_file_var.get())
            self.log_message(f"Parsed RTF file successfully. Found {len(rtf_file)} players.")

            # Determine players to process
            update = self.update_mode_var.get() == "true"
            process_player = self.player_mode_var.get() == "specific"

            if process_player:
                player_uuid = int(self.player_uid_var.get())
                players_to_process = [p for p in rtf_file if int(p[0]) == player_uuid]
            elif update:
                # Filter out already processed players
                try:
                    output_dir = self.output_dir_var.get() if self.output_dir_var.get() else "./NewGens/"
                    values_from_config = extract_from_values(f"{output_dir}/config.xml")
                    ids_in_b = [item for item in values_from_config]
                    players_to_process = [item for item in rtf_file if item[0] not in ids_in_b]
                except FileNotFoundError:
                    players_to_process = rtf_file
            else:
                players_to_process = rtf_file

            if len(players_to_process) == 0:
                self.log_message("No players to process")
                return

            self.log_message(f"Processing {len(players_to_process)} players")

            # Generate prompts
            prompts = generate_prompts_for_players(players_to_process, app_config)

            # Generate images
            total_players = len(prompts)
            for i, prompt in enumerate(prompts):
                if self.stop_generation:
                    break

                uid = prompt.split(":")[0]
                comfy_prompt = prompt.split(":")[1]

                self.log_message(f"Generating image for player {uid} ({i+1}/{total_players})")
                generate_image(uid, comfy_prompt)

                # Update progress
                progress = (i + 1) / total_players * 100
                self.progress_var.set(progress)

                # Queue image for preview
                self.image_queue.put(uid)

            # Post-processing
            if not self.stop_generation:
                try:
                    # Remove background from images if available
                    if BG_REMOVAL_AVAILABLE:
                        try:
                            output_dir = self.output_dir_var.get() if self.output_dir_var.get() else "./NewGens/"
                            player_ids = [item[0] for item in players_to_process]

                            self.log_message(f"Starting background removal for {len(player_ids)} images...")

                            # Check if images exist before processing
                            missing_images = []
                            for uid in player_ids:
                                image_path = os.path.join(output_dir, f"{uid}.png")
                                if not os.path.exists(image_path):
                                    missing_images.append(uid)

                            if missing_images:
                                self.log_message(f"Warning: {len(missing_images)} images not found: {missing_images[:5]}...")

                            # Only process existing images
                            existing_ids = [uid for uid in player_ids if os.path.exists(os.path.join(output_dir, f"{uid}.png"))]

                            if existing_ids:
                                self.log_message(f"Processing {len(existing_ids)} images for background removal...")
                                processed_count = remove_bg_from_file_list(output_dir, existing_ids, use_gpu=False)
                                self.log_message(f"Background removal completed: {processed_count}/{len(existing_ids)} images processed successfully.")
                            else:
                                self.log_message("No images found for background removal.")

                        except Exception as e:
                            self.log_message(f"Background removal failed: {str(e)}")
                            import traceback
                            self.log_message(f"Error details: {traceback.format_exc()}")
                    else:
                        self.log_message("Background removal not available (rembg not installed). Images will have original backgrounds.")

                    # Update or create configuration XML
                    if update:
                        append_to_config_xml(
                            self.output_dir_var.get(),
                            [item[0] for item in players_to_process],
                            self.fm_version_var.get()
                        )
                    else:
                        create_config_xml(
                            self.output_dir_var.get(),
                            [item[0] for item in players_to_process],
                            self.fm_version_var.get()
                        )
                    self.log_message("Configuration XML updated")
                except Exception as e:
                    self.log_message(f"Post-processing failed: {e}")

            self.log_message("Generation complete!")

        except Exception as e:
            self.log_message(f"Generation failed: {str(e)}")
            import traceback
            self.log_message(traceback.format_exc())

        finally:
            # Re-enable start button
            self.start_button.config(state=tk.NORMAL)
            self.stop_button.config(state=tk.DISABLED)

    def validate_config(self):
        """Validate configuration before starting"""
        if not self.rtf_file_var.get():
            messagebox.showerror("Error", "Please select an RTF file")
            return False

        if not self.output_dir_var.get():
            messagebox.showerror("Error", "Please select an output directory")
            return False

        if not os.path.exists(self.rtf_file_var.get()):
            messagebox.showerror("Error", "RTF file does not exist")
            return False

        return True

    def update_image_preview(self):
        """Update image preview periodically"""
        try:
            while not self.image_queue.empty():
                uid = self.image_queue.get()
                self.show_image(uid)
        except:
            pass

        # Schedule next update
        self.root.after(1000, self.update_image_preview)

    def show_image(self, uid):
        """Show an image in the preview"""
        try:
            image_path = f"{self.output_dir_var.get()}/{uid}.png"
            if os.path.exists(image_path):
                # Add to generated images list if not already present
                if uid not in self.generated_images:
                    self.generated_images.append(uid)
                    self.update_total_count()

                # Find the index of current image
                if uid in self.generated_images:
                    self.current_image_index = self.generated_images.index(uid)

                # Load and resize image
                image = Image.open(image_path)
                image.thumbnail((300, 300), Image.Resampling.LANCZOS)

                # Convert to PhotoImage
                photo = ImageTk.PhotoImage(image)

                # Update display
                self.image_label.config(image=photo, text="")
                self.image_label.image = photo  # Keep reference
                self.current_image_var.set(f"Player {uid}")

                # Load and display the prompt for this image
                try:
                    prompt = get_prompt_for_image(uid)
                    self.prompt_text.config(state=tk.NORMAL)
                    self.prompt_text.delete(1.0, tk.END)
                    self.prompt_text.insert(tk.END, prompt)
                    self.prompt_text.config(state=tk.DISABLED)
                except Exception as e:
                    self.prompt_text.config(state=tk.NORMAL)
                    self.prompt_text.delete(1.0, tk.END)
                    self.prompt_text.insert(tk.END, f"Error loading prompt: {str(e)}")
                    self.prompt_text.config(state=tk.DISABLED)

                self.log_message(f"Preview updated: {uid}.png")

        except Exception as e:
            self.log_message(f"Failed to load image {uid}: {str(e)}")

    def show_previous_image(self):
        """Show previous image"""
        if not self.generated_images:
            self.log_message("No images available to navigate")
            return

        if self.current_image_index <= 0:
            self.log_message("Already at first image")
            return

        self.current_image_index -= 1
        uid = self.generated_images[self.current_image_index]
        self.show_image(uid)

    def show_next_image(self):
        """Show next image"""
        if not self.generated_images:
            self.log_message("No images available to navigate")
            return

        if self.current_image_index >= len(self.generated_images) - 1:
            self.log_message("Already at last image")
            return

        self.current_image_index += 1
        uid = self.generated_images[self.current_image_index]
        self.show_image(uid)

    def update_total_count(self):
        """Update the total generated count display"""
        count = len(self.generated_images)
        self.total_generated_var.set(str(count))
        self.log_message(f"Total generated images: {count}")

    def scan_existing_images(self):
        """Scan the output directory for existing images"""
        try:
            output_dir = self.output_dir_var.get()
            if not output_dir or not os.path.exists(output_dir):
                return

            # Find all PNG files in the output directory
            existing_images = []
            for file in os.listdir(output_dir):
                if file.endswith('.png') and file[:-4].isdigit():
                    uid = file[:-4]  # Remove .png extension
                    existing_images.append(uid)

            # Sort by UID (numerically)
            existing_images.sort(key=int)

            # Update our tracking list
            self.generated_images = existing_images
            self.update_total_count()

            if existing_images:
                self.log_message(f"Found {len(existing_images)} existing images")
                # Show the first image if available
                self.show_image(existing_images[0])
            else:
                self.log_message("No existing images found")

        except Exception as e:
            self.log_message(f"Error scanning existing images: {str(e)}")

    def refresh_preview(self):
        """Refresh the preview"""
        self.log_message("Refreshing preview...")
        self.scan_existing_images()
        if self.generated_images:
            self.show_image(self.generated_images[0])  # Show first image

def main():
    """Main function"""
    root = tk.Tk()
    app = FMFaceGeneratorGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main()