soprano-factory/gui.pyw

# gui.pyw - Soprano Reforged V3 Training Factory
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import sv_ttk
import subprocess
import threading
import sys
import queue
import os

class SopranoGUI(tk.Tk):
    def __init__(self):
        super().__init__()

        self.title("Soprano Reforged V3 - Training Factory")
        self.geometry("900x700")
        
        # Queue for thread-safe GUI updates
        self.log_queue = queue.Queue()
        self.is_running = False

        # --- Layout ---
        self.create_widgets()
        
        # --- Theme ---
        sv_ttk.set_theme("dark")
        self.refresh_display()
        
        # --- Log Loop ---
        self.after(100, self.process_logs)

    def refresh_display(self):
        self.update_idletasks()
        self.wm_attributes("-alpha", 0.99)
        self.wm_attributes("-alpha", 1.0)

    def create_widgets(self):
        main_frame = ttk.Frame(self)
        main_frame.pack(fill="both", expand=True, padx=10, pady=10)

        # 1. Tabs for Stages
        self.notebook = ttk.Notebook(main_frame)
        self.notebook.pack(fill="both", expand=True, pady=(0, 10))

        self.tab_codec = ttk.Frame(self.notebook, padding=20)
        self.tab_dataset = ttk.Frame(self.notebook, padding=20)
        self.tab_train = ttk.Frame(self.notebook, padding=20)
        self.tab_infer = ttk.Frame(self.notebook, padding=20)

        self.notebook.add(self.tab_codec, text="Stage 0: Codec")
        self.notebook.add(self.tab_dataset, text="Stage 1: Dataset")
        self.notebook.add(self.tab_train, text="Stage 2: Training")
        self.notebook.add(self.tab_infer, text="Inference")

        self.build_codec_tab()
        self.build_dataset_tab()
        self.build_train_tab()
        self.build_infer_tab()

        # 2. Console Output
        console_frame = ttk.LabelFrame(main_frame, text="V3 System Output", padding=10)
        console_frame.pack(fill="both", expand=True)
        
        self.console = tk.Text(console_frame, height=12, state="disabled", bg="#1c1c1c", fg="#f0f0f0", font=("Consolas", 10))
        self.console.pack(fill="both", expand=True, side="left")
        
        scrollbar = ttk.Scrollbar(console_frame, command=self.console.yview)
        scrollbar.pack(side="right", fill="y")
        self.console.config(yscrollcommand=scrollbar.set)

    # --- Tab Builders ---

    def build_codec_tab(self):
        f = self.tab_codec
        ttk.Label(f, text="Stage 0: Audio Codec (Float32)", font=("Segoe UI", 14, "bold")).pack(anchor="w", pady=(0, 10))
        ttk.Label(f, text="WAV Directory (e.g. ./mio_dataset/wavs/*.wav)").pack(anchor="w")
        self.codec_wav_entry = ttk.Entry(f)
        self.codec_wav_entry.pack(fill="x", pady=(5, 10))
        self.codec_wav_entry.insert(0, "./mio_dataset/wavs/*.wav")
        
        btn_frame = ttk.Frame(f)
        btn_frame.pack(fill="x", pady=10)
        ttk.Button(btn_frame, text="Browse Folder", command=lambda: self.browse_folder(self.codec_wav_entry, suffix="/*.wav")).pack(side="left", padx=(0, 10))
        ttk.Button(btn_frame, text="Start Codec Training", style="Accent.TButton", 
                   command=lambda: self.run_script("train_codec.py", ["--wav-dir", self.codec_wav_entry.get()])).pack(side="left")

    def build_dataset_tab(self):
        f = self.tab_dataset
        ttk.Label(f, text="Stage 1: Token Generation", font=("Segoe UI", 14, "bold")).pack(anchor="w", pady=(0, 10))
        ttk.Label(f, text="Dataset Input Directory").pack(anchor="w")
        self.ds_input_entry = ttk.Entry(f)
        self.ds_input_entry.pack(fill="x", pady=(5, 10))
        self.ds_input_entry.insert(0, "./mio_dataset")

        ttk.Label(f, text="Encoder Checkpoint").pack(anchor="w")
        self.ds_encoder_entry = ttk.Entry(f)
        self.ds_encoder_entry.pack(fill="x", pady=(5, 10))
        self.ds_encoder_entry.insert(0, "./weights/codec/encoder.pth")
        
        ttk.Button(f, text="Generate Dataset", style="Accent.TButton",
                   command=lambda: self.run_script("generate_dataset.py", [
                       "--input-dir", self.ds_input_entry.get(),
                       "--encoder-ckpt", self.ds_encoder_entry.get()
                   ])).pack(anchor="w", pady=10)

    def build_train_tab(self):
        f = self.tab_train
        ttk.Label(f, text="Stage 2: Joint V3 Training", font=("Segoe UI", 14, "bold")).pack(anchor="w", pady=(0, 10))
        ttk.Label(f, text="Dataset Directory (containing train.json)").pack(anchor="w")
        self.train_input_entry = ttk.Entry(f)
        self.train_input_entry.pack(fill="x", pady=(5, 10))
        self.train_input_entry.insert(0, "./mio_dataset")
        
        ttk.Button(f, text="Start Speed Run", style="Accent.TButton",
                   command=lambda: self.run_script("train.py", ["--input-dir", self.train_input_entry.get()])).pack(anchor="w")

    def build_infer_tab(self):
        f = self.tab_infer
        ttk.Label(f, text="Inference: Generate Speech", font=("Segoe UI", 14, "bold")).pack(anchor="w", pady=(0, 10))
        ttk.Label(f, text="Enter Text:").pack(anchor="w")
        self.infer_text = ttk.Entry(f)
        self.infer_text.pack(fill="x", pady=(5, 10))
        self.infer_text.insert(0, "Soprano V3 is now fully operational.")
        
        ttk.Label(f, text="Model Directory (e.g. weights/model/epoch_9)").pack(anchor="w")
        self.infer_model_entry = ttk.Entry(f)
        self.infer_model_entry.pack(fill="x", pady=(5, 10))
        self.infer_model_entry.insert(0, "./weights/model/epoch_9")
        
        ttk.Button(f, text="Generate .WAV", style="Accent.TButton",
                   command=lambda: self.run_script("inference.py", [
                       "--text", self.infer_text.get(),
                       "--model-dir", self.infer_model_entry.get(),
                       "--output", "v3_output.wav"
                   ])).pack(anchor="w", pady=10)

    # --- Subprocess Logic ---

    def browse_folder(self, entry_widget, suffix=""):
        path = filedialog.askdirectory()
        if path:
            entry_widget.delete(0, tk.END)
            entry_widget.insert(0, path + suffix)

    def log(self, message):
        self.log_queue.put(message)

    def process_logs(self):
        while not self.log_queue.empty():
            msg = self.log_queue.get()
            self.console.config(state="normal")
            self.console.insert(tk.END, msg)
            self.console.see(tk.END)
            self.console.config(state="disabled")
        self.after(100, self.process_logs)

    def run_script(self, script_name, args):
        if self.is_running:
            messagebox.showwarning("Busy", "Process already running.")
            return
        
        self.is_running = True
        self.console.config(state="normal")
        self.console.delete(1.0, tk.END)
        self.console.config(state="disabled")
        
        cmd = [sys.executable, script_name] + args
        self.log(f"RUNNING: {' '.join(cmd)}\n" + "-"*50 + "\n")
        
        thread = threading.Thread(target=self._execute_subprocess, args=(cmd,))
        thread.start()

    def _execute_subprocess(self, cmd):
        try:
            process = subprocess.Popen(
                cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, 
                text=True, bufsize=1, universal_newlines=True
            )
            for line in process.stdout:
                self.log(line)
            process.wait()
            self.log(f"\nProcess Exit Code: {process.returncode}\n")
        except Exception as e:
            self.log(f"\nError: {e}\n")
        finally:
            self.is_running = False

if __name__ == "__main__":
    app = SopranoGUI()
    app.mainloop()