# Standalone YuE2 only. Never install into Comfy's Python environment.
# Downloads YuE2-3B + YuE2-Vae into YUE2_ROOT\models so the UI does not hit HF on first Generate.
[CmdletBinding()]
param(
  [string]$Root = 'C:\Users\ianjm\Development\YuE2',
  [string]$Uv = 'uv',
  [string]$Models = ''
)
$ErrorActionPreference = 'Stop'
$Root = [IO.Path]::GetFullPath($Root)
if (-not $Models) { $Models = Join-Path $Root 'models' }
$Models = [IO.Path]::GetFullPath($Models)
function Run([string]$Exe, [string[]]$Arguments) {
  & $Exe @Arguments
  if ($LASTEXITCODE -ne 0) { throw "$Exe failed ($LASTEXITCODE)" }
}
Write-Host "YuE2 setup root=$Root models=$Models"
if (!(Test-Path -LiteralPath $Root)) {
  Run git @('clone', 'https://github.com/multimodal-art-projection/YuE.git', $Root)
}
Run $Uv @('venv', '--allow-existing', '--python', '3.12', (Join-Path $Root '.venv'))
$python = Join-Path $Root '.venv\Scripts\python.exe'
# Install the package first, then force CUDA torch back on (editable deps can pull CPU wheels).
Run $Uv @('pip', 'install', '--python', $python, '-e', $Root)
Run $Uv @('pip', 'install', '--python', $python, 'soundfile', 'psutil', 'huggingface_hub', 'hf_xet')
Run $Uv @('pip', 'install', '--python', $python, 'torch', 'torchaudio', '--reinstall', '--index-url', 'https://download.pytorch.org/whl/cu128')
New-Item -ItemType Directory -Force -Path $Models | Out-Null
$modelDir = Join-Path $Models 'YuE2-3B'
$vaeDir = Join-Path $Models 'YuE2-Vae'
if (!(Test-Path -LiteralPath (Join-Path $modelDir 'config.json')) -and !(Test-Path -LiteralPath (Join-Path $modelDir 'model.safetensors.index.json'))) {
  Write-Host "Downloading m-a-p/YuE2-3B → $modelDir"
  Run $python @(
    '-c',
    'from huggingface_hub import snapshot_download; import sys; snapshot_download(sys.argv[1], local_dir=sys.argv[2])',
    'm-a-p/YuE2-3B',
    $modelDir
  )
} else {
  Write-Host "Reusing existing model at $modelDir"
}
if (!(Test-Path -LiteralPath (Join-Path $vaeDir 'config.json'))) {
  Write-Host "Downloading m-a-p/YuE2-Vae → $vaeDir"
  Run $python @(
    '-c',
    'from huggingface_hub import snapshot_download; import sys; snapshot_download(sys.argv[1], local_dir=sys.argv[2])',
    'm-a-p/YuE2-Vae',
    $vaeDir
  )
} else {
  Write-Host "Reusing existing VAE at $vaeDir"
}
Run $python @('-c', 'import torch; from yue2 import YuE2Pipeline; assert torch.cuda.is_available(); print(torch.__version__, torch.cuda.get_device_name(0))')
$ready = @{
  model = $modelDir
  vae = $vaeDir
  modelId = 'm-a-p/YuE2-3B'
  vaeId = 'm-a-p/YuE2-Vae'
  modelsDir = $Models
  python = $python
  root = $Root
}
$ready | ConvertTo-Json | Set-Content -LiteralPath (Join-Path $Root 'aigen-ready.json') -Encoding utf8
Write-Host "YuE2 ready at $Root"
Write-Host "Local model=$modelDir"
Write-Host "Local vae=$vaeDir"
Write-Host "Set YUE2_ROOT=$Root YUE2_PYTHON=$python YUE2_MODEL=$modelDir YUE2_VAE=$vaeDir if not using defaults. No Comfy environment was modified."