start_gradio_ui.ps1

# ACE-Step Gradio Web UI Launcher (PowerShell Version)

# Set network environment variables to avoid socket error 10049 on Windows
$env:MASTER_ADDR = "127.0.0.1"
$env:MASTER_PORT = "2333"

# Disable torch compile to avoid MSVC requirement on Windows
$env:TORCH_COMPILE_DISABLE = "1"
$env:TRITON_INTERPRET = "1"

# ==================== Configuration ====================
$Port = "7860"
$ServerName = "127.0.0.1"
$Language = "ja"

# Model settings
$ConfigPath = "acestep-v15-turbo"
$LmModelPath = "acestep-5Hz-lm-1.7B"
$DownloadSource = "huggingface"
$InitService = $true
# Ensure 5Hz LM is initialized on startup
$InitLlm = $true

# Optional settings
$Share = $false
$OffloadToCpu = $true
$Backend = "pt" # Use PyTorch backend for better Windows compatibility
$EnableApi = $false
$ApiKey = $null
$AuthUsername = $null
$AuthPassword = $null

# ==================== Environment Setup ====================
$PSScriptRoot = Split-Path -Parent $MyInvocation.MyCommand.Definition
Set-Location $PSScriptRoot

Write-Host "Starting ACE-Step Gradio Web UI..." -ForegroundColor Cyan
Write-Host "Server will be available at: http://$($ServerName):$($Port)" -ForegroundColor Gray

# Check for .venv
$VenvPath = Join-Path $PSScriptRoot ".venv"
$PythonExe = Join-Path $VenvPath "Scripts\python.exe"

if (-not (Test-Path $VenvPath)) {
    Write-Host "[Setup] Virtual environment not found. Checking for uv..." -ForegroundColor Yellow

    if (Get-Command uv -ErrorAction SilentlyContinue) {
        Write-Host "Running: uv sync" -ForegroundColor Cyan
        uv sync
        if ($LASTEXITCODE -ne 0) {
            Write-Host "[Error] Failed to setup environment with uv." -ForegroundColor Red
            Read-Host "Press Enter to exit..."
            exit 1
        }
    } else {
        Write-Host "[Error] .venv not found and 'uv' is not installed." -ForegroundColor Red
        Write-Host "Please install uv or create a virtual environment manually." -ForegroundColor Yellow
        Read-Host "Press Enter to exit..."
        exit 1
    }
}

# ==================== Build Arguments ====================
$ArgList = @(
    "acestep\acestep_v15_pipeline.py",
    "--port", $Port,
    "--server-name", $ServerName,
    "--language", $Language,
    "--config_path", $ConfigPath,
    "--lm_model_path", $LmModelPath,
    "--download-source", $DownloadSource,
    "--backend", $Backend
)

if ($InitService) { $ArgList += "--init_service", "true" }
if ($InitLlm) { $ArgList += "--init_llm", "true" }
if ($Share) { $ArgList += "--share" }
if ($OffloadToCpu) { $ArgList += "--offload_to_cpu", "true" }
if ($EnableApi) { $ArgList += "--enable-api" }
if ($null -ne $ApiKey) { $ArgList += "--api-key", $ApiKey }
if ($null -ne $AuthUsername) { $ArgList += "--auth-username", $AuthUsername }
if ($null -ne $AuthPassword) { $ArgList += "--auth-password", $AuthPassword }

# ==================== Launch ====================
Write-Host "[Environment] Using virtual environment: $VenvPath" -ForegroundColor Green
Write-Host "Launching ACE-Step (Initialize 5Hz LM: $InitLlm)..." -ForegroundColor Cyan
Write-Host ""

# Execute python with the built argument list
& $PythonExe @ArgList

if ($LASTEXITCODE -ne 0) {
    Write-Host ""
    Write-Host "[Error] Application exited with code $LASTEXITCODE" -ForegroundColor Red
    Read-Host "Press Enter to exit..."
}