DeepFaceLab_Linux/manage_webui.py

#!/usr/bin/env python3
import os
import sys
import time
import signal
import socket
import subprocess
import shutil

PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))
VENV_DIR = os.path.join(PROJECT_DIR, '.venv')
PID_FILE = os.path.join(PROJECT_DIR, '.webui.pid')
LOG_FILE = os.path.join(PROJECT_DIR, 'webui_server.log')
SERVER_SCRIPT = os.path.join(PROJECT_DIR, 'web_ui_server.py')

# Get correct python / pip paths in virtualenv
if sys.platform == 'win32':
    VENV_PYTHON = os.path.join(VENV_DIR, 'Scripts', 'python.exe')
    VENV_PIP = os.path.join(VENV_DIR, 'Scripts', 'pip.exe')
else:
    VENV_PYTHON = os.path.join(VENV_DIR, 'bin', 'python')
    VENV_PIP = os.path.join(VENV_DIR, 'bin', 'pip')

def is_port_in_use(port=8082):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return s.connect_ex(('127.0.0.1', port)) == 0

def check_process_running(pid):
    try:
        os.kill(pid, 0)
        return True
    except OSError:
        return False

def get_running_pid():
    if os.path.exists(PID_FILE):
        try:
            with open(PID_FILE, 'r') as f:
                pid = int(f.read().strip())
            if check_process_running(pid):
                return pid
        except ValueError:
            pass
    return None

def setup():
    print("=== Setting up virtual environment ===")
    if not os.path.exists(VENV_DIR):
        print(f"Creating virtual environment in {VENV_DIR}...")
        try:
            subprocess.run([sys.executable, '-m', 'venv', VENV_DIR], check=True)
            print("Virtual environment created successfully.")
        except subprocess.CalledProcessError as e:
            print(f"Error creating virtual environment: {e}", file=sys.stderr)
            sys.exit(1)
    else:
        print("Virtual environment already exists.")

    print("\n=== Installing dependencies ===")
    try:
        print("Upgrading pip...")
        subprocess.run([VENV_PIP, 'install', '--upgrade', 'pip'], check=True)
        print("Installing Flask...")
        subprocess.run([VENV_PIP, 'install', 'Flask'], check=True)
        print("All dependencies installed successfully.")
    except subprocess.CalledProcessError as e:
        print(f"Error installing dependencies: {e}", file=sys.stderr)
        sys.exit(1)

def start():
    pid = get_running_pid()
    if pid:
        print(f"Web UI server is already running with PID {pid}.")
        return

    if is_port_in_use(8082):
        print("Error: Port 8082 is already in use by another application.", file=sys.stderr)
        sys.exit(1)

    if not os.path.exists(VENV_PYTHON):
        print("Error: Virtual environment not found. Please run 'setup' first.", file=sys.stderr)
        sys.exit(1)

    print("Starting Web UI server in the background...")
    
    # Open log file for appending
    log_f = open(LOG_FILE, 'a')
    log_f.write(f"\n--- Server started at {time.strftime('%Y-%m-%d %H:%M:%S')} ---\n")
    log_f.flush()

    try:
        # Launch using start_new_session=True to detach from parent shell
        proc = subprocess.Popen(
            [VENV_PYTHON, SERVER_SCRIPT],
            cwd=PROJECT_DIR,
            stdout=log_f,
            stderr=log_f,
            start_new_session=True
        )
        
        # Write PID file
        with open(PID_FILE, 'w') as f:
            f.write(str(proc.pid))
            
        time.sleep(1) # Wait a bit to verify it didn't crash immediately
        if proc.poll() is None:
            print(f"Web UI server successfully started (PID: {proc.pid}).")
            print(f"Logs are being written to: {LOG_FILE}")
            print("You can access the interface at: http://localhost:8082/")
        else:
            print("Error: Server failed to start. Check log file for details.", file=sys.stderr)
            sys.exit(1)
            
    except Exception as e:
        print(f"Error launching server: {e}", file=sys.stderr)
        sys.exit(1)

def stop():
    pid = get_running_pid()
    if not pid:
        print("Web UI server is not running.")
        # Clean up stale PID file if any
        if os.path.exists(PID_FILE):
            os.remove(PID_FILE)
        return

    print(f"Stopping Web UI server (PID: {pid})...")
    try:
        os.kill(pid, signal.SIGTERM)
        # Wait up to 5 seconds for it to exit
        for _ in range(10):
            time.sleep(0.5)
            if not check_process_running(pid):
                break
        else:
            # Force kill if still running
            print("Force stopping server...")
            os.kill(pid, signal.SIGKILL)
            
        print("Server stopped.")
    except ProcessLookupError:
        print("Server was already stopped.")
    
    if os.path.exists(PID_FILE):
        os.remove(PID_FILE)

def status():
    pid = get_running_pid()
    if pid:
        print(f"Status: RUNNING (PID: {pid})")
        print("Access URLs:")
        print("  - Local: http://localhost:8082/")
        try:
            hostname = socket.gethostname()
            ip = socket.gethostbyname(hostname)
            print(f"  - Network: http://{ip}:8082/")
        except Exception:
            pass
    else:
        print("Status: STOPPED")
        
    if os.path.exists(LOG_FILE):
        print("\nLast 10 log lines:")
        try:
            with open(LOG_FILE, 'r') as f:
                lines = f.readlines()
                for line in lines[-10:]:
                    print("  " + line.strip())
        except Exception as e:
            print(f"  Failed to read logs: {e}")

def main():
    if len(sys.argv) < 2:
        print("Usage: python manage_webui.py [setup|start|stop|restart|status]")
        sys.exit(1)

    cmd = sys.argv[1].lower()

    if cmd == 'setup':
        setup()
    elif cmd == 'start':
        start()
    elif cmd == 'stop':
        stop()
    elif cmd == 'restart':
        stop()
        time.sleep(1)
        start()
    elif cmd == 'status':
        status()
    else:
        print(f"Unknown command: {cmd}", file=sys.stderr)
        print("Available commands: setup, start, stop, restart, status")
        sys.exit(1)

if __name__ == '__main__':
    main()