# OmniAvatar-14B Setup Script for Windows
# Downloads all required models using HuggingFace CLI

Write-Host "🚀 OmniAvatar-14B Setup Script" -ForegroundColor Green
Write-Host "===============================================" -ForegroundColor Green

# Check if Python is available
try {
    $pythonVersion = python --version 2>$null
    Write-Host "✅ Python found: $pythonVersion" -ForegroundColor Green
} catch {
    Write-Host "❌ Python not found! Please install Python first." -ForegroundColor Red
    exit 1
}

# Check if pip is available
try {
    pip --version | Out-Null
    Write-Host "✅ pip is available" -ForegroundColor Green
} catch {
    Write-Host "❌ pip not found! Please ensure pip is installed." -ForegroundColor Red
    exit 1
}

# Install huggingface-cli if not available
Write-Host "📦 Checking HuggingFace CLI..." -ForegroundColor Yellow
try {
    huggingface-cli --version | Out-Null
    Write-Host "✅ HuggingFace CLI already available" -ForegroundColor Green
} catch {
    Write-Host "📦 Installing HuggingFace CLI..." -ForegroundColor Yellow
    pip install "huggingface_hub[cli]"
    if ($LASTEXITCODE -ne 0) {
        Write-Host "❌ Failed to install HuggingFace CLI" -ForegroundColor Red
        exit 1
    }
    Write-Host "✅ HuggingFace CLI installed" -ForegroundColor Green
}

# Create directories
Write-Host "📁 Creating directory structure..." -ForegroundColor Yellow
$directories = @(
    "pretrained_models",
    "pretrained_models\Wan2.1-T2V-14B",
    "pretrained_models\OmniAvatar-14B",
    "pretrained_models\wav2vec2-base-960h",
    "outputs"
)

foreach ($dir in $directories) {
    New-Item -Path $dir -ItemType Directory -Force | Out-Null
    Write-Host "✅ Created: $dir" -ForegroundColor Green
}

# Model information
$models = @(
    @{
        Name = "Wan2.1-T2V-14B"
        Repo = "Wan-AI/Wan2.1-T2V-14B"
        Description = "Base model for 14B OmniAvatar model"
        Size = "~28GB"
        LocalDir = "pretrained_models\Wan2.1-T2V-14B"
    },
    @{
        Name = "OmniAvatar-14B"
        Repo = "OmniAvatar/OmniAvatar-14B"
        Description = "LoRA and audio condition weights"
        Size = "~2GB"
        LocalDir = "pretrained_models\OmniAvatar-14B"
    },
    @{
        Name = "wav2vec2-base-960h"
        Repo = "facebook/wav2vec2-base-960h"
        Description = "Audio encoder"
        Size = "~360MB"
        LocalDir = "pretrained_models\wav2vec2-base-960h"
    }
)

Write-Host ""
Write-Host "⚠️  WARNING: This will download approximately 30GB of models!" -ForegroundColor Yellow
Write-Host "Make sure you have sufficient disk space and a stable internet connection." -ForegroundColor Yellow
Write-Host ""

$response = Read-Host "Continue with download? (y/N)"
if ($response.ToLower() -ne 'y') {
    Write-Host "❌ Download cancelled by user" -ForegroundColor Red
    exit 0
}

# Download models
foreach ($model in $models) {
    Write-Host ""
    Write-Host "📥 Downloading $($model.Name) ($($model.Size))..." -ForegroundColor Cyan
    Write-Host "📝 $($model.Description)" -ForegroundColor Gray
    
    # Check if already exists
    if ((Test-Path $model.LocalDir) -and (Get-ChildItem $model.LocalDir -Force | Measure-Object).Count -gt 0) {
        Write-Host "✅ $($model.Name) already exists, skipping..." -ForegroundColor Green
        continue
    }
    
    # Download model
    $cmd = "huggingface-cli download $($model.Repo) --local-dir $($model.LocalDir)"
    Write-Host "🚀 Running: $cmd" -ForegroundColor Gray
    
    Invoke-Expression $cmd
    
    if ($LASTEXITCODE -eq 0) {
        Write-Host "✅ $($model.Name) downloaded successfully!" -ForegroundColor Green
    } else {
        Write-Host "❌ Failed to download $($model.Name)" -ForegroundColor Red
        exit 1
    }
}

Write-Host ""
Write-Host "🎉 OmniAvatar-14B setup completed successfully!" -ForegroundColor Green
Write-Host ""
Write-Host "💡 Next steps:" -ForegroundColor Yellow
Write-Host "1. Run your app: python app.py" -ForegroundColor White
Write-Host "2. The app will now support full avatar video generation!" -ForegroundColor White
Write-Host "3. Use the Gradio interface or API endpoints" -ForegroundColor White
Write-Host ""
Write-Host "🔗 For more information visit:" -ForegroundColor Yellow
Write-Host "   https://huggingface.co/OmniAvatar/OmniAvatar-14B" -ForegroundColor Cyan