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
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 choose_profile, create_or_update, process_player_or_file, get_player_input

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)

        # 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()

        # Set up the GUI
        self.setup_gui()
        self.load_config()

        # Start image preview update loop
        self.update_image_preview()

    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")

        # 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_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)

        # 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)

        # 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)

        # 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)

    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')
                self.config.add_section('profile:NewGens')

            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:NewGens', 'football_manager_version', self.fm_version_var.get())
            self.config.set('profile:NewGens', 'rtf_file', self.rtf_file_var.get())
            self.config.set('profile:NewGens', '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", "Configuration saved successfully!")
            self.log_message("Configuration saved")

        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 FM settings
                if 'profile:NewGens' in self.config:
                    self.fm_version_var.set(self.config.get('profile:NewGens', 'football_manager_version', fallback='2024'))
                    self.rtf_file_var.set(self.config.get('profile:NewGens', 'rtf_file', fallback=''))
                    self.output_dir_var.set(self.config.get('profile:NewGens', 'output_dir', fallback=''))

                self.log_message("Configuration loaded")
            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('')
        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

        # 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()

    def stop_generation_thread(self):
        """Stop the generation process"""
        self.stop_generation = True
        self.stop_button.config(state=tk.DISABLED)
        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:
                    values_from_config = extract_from_values(f"{self.output_dir_var.get()}/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:
                            remove_bg_from_file_list(self.output_dir_var.get(), [item[0] for item in players_to_process], use_gpu=False)
                            self.log_message("Background removed from images.")
                        except Exception as e:
                            self.log_message(f"Background removal failed: {e}")
                    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):
                # 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}")
                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"""
        self.log_message("Previous image functionality not implemented yet")

    def show_next_image(self):
        """Show next image"""
        self.log_message("Next image functionality not implemented yet")

    def refresh_preview(self):
        """Refresh the preview"""
        self.log_message("Refreshing preview...")
        # This would scan the output directory and show available images

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

if __name__ == "__main__":
    main()