# Download ACE-Step / YuE weights into the local ComfyUI models tree via Hugging Face CLI.
# Example:
#   powershell -ExecutionPolicy Bypass -File scripts\download_music_models.ps1
#   powershell -File scripts\download_music_models.ps1 -ComfyRoot C:\ComfyUI -SkipYue

[CmdletBinding()]
param(
  [string]$ComfyRoot = '',
  [switch]$SkipAceStep,
  [switch]$SkipAceStep15,
  [switch]$SkipYue,
  [switch]$Force
)

$ErrorActionPreference = 'Stop'
$ProgressPreference = 'Continue'

function Write-Step([string]$Message) {
  Write-Host ""
  Write-Host "==> $Message" -ForegroundColor Cyan
}

function Write-Ok([string]$Message) {
  Write-Host "    $Message" -ForegroundColor Green
}

function Write-Skip([string]$Message) {
  Write-Host "    skip  $Message" -ForegroundColor DarkYellow
}

function Write-Warn([string]$Message) {
  Write-Host "    warn  $Message" -ForegroundColor Yellow
}

function Resolve-ComfyRoot([string]$Hint) {
  $candidates = @(
    $Hint,
    $env:COMFY_ROOT,
    $env:COMFYUI_ROOT,
    'C:\ComfyUI',
    'C:\ComfyUI_windows_portable\ComfyUI',
    'D:\ComfyUI',
    'C:\ComfyUI-MiniMax\ComfyUI',
    (Join-Path $env:USERPROFILE 'ComfyUI'),
    (Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Installs\ComfyUI (1)\ComfyUI'),
    (Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Installs\ComfyUI\ComfyUI')
  ) | Where-Object { $_ } | ForEach-Object { $_.TrimEnd('\', '/') }

  foreach ($path in $candidates) {
    if (Test-Path (Join-Path $path 'main.py')) { return (Resolve-Path $path).Path }
    if (Test-Path (Join-Path $path 'ComfyUI\main.py')) { return (Resolve-Path (Join-Path $path 'ComfyUI')).Path }
  }

  Write-Host "Could not find ComfyUI (looked for main.py)." -ForegroundColor Red
  $typed = Read-Host "Enter ComfyUI root path (or blank to abort)"
  if (-not $typed) { throw "COMFY_ROOT is required." }
  $typed = $typed.Trim().Trim('"')
  if (-not (Test-Path (Join-Path $typed 'main.py'))) {
    throw "Not a ComfyUI root (missing main.py): $typed"
  }
  return (Resolve-Path $typed).Path
}

function Get-DirBytes([string]$Path) {
  if (-not (Test-Path $Path)) { return [long]0 }
  $item = Get-Item $Path
  if (-not $item.PSIsContainer) { return [long]$item.Length }
  [long]((Get-ChildItem $Path -Recurse -File -ErrorAction SilentlyContinue | Measure-Object Length -Sum).Sum)
}

function Format-Bytes([long]$Bytes) {
  if ($Bytes -ge 1GB) { return ('{0:N2} GB' -f ($Bytes / 1GB)) }
  if ($Bytes -ge 1MB) { return ('{0:N1} MB' -f ($Bytes / 1MB)) }
  if ($Bytes -ge 1KB) { return ('{0:N0} KB' -f ($Bytes / 1KB)) }
  return "$Bytes B"
}

function Test-VerifiedTarget([string]$Path, [long]$MinBytes = 1MB) {
  if (-not (Test-Path $Path)) { return $false }
  $item = Get-Item $Path
  if (-not $item.PSIsContainer) { return $item.Length -ge $MinBytes }
  $files = Get-ChildItem $Path -Recurse -File -ErrorAction SilentlyContinue
  if (-not $files) { return $false }
  $total = [long]($files | Measure-Object Length -Sum).Sum
  return $total -ge $MinBytes
}

function Ensure-Dir([string]$Path) {
  if (-not (Test-Path $Path)) {
    New-Item -ItemType Directory -Force -Path $Path | Out-Null
  }
}

function Find-HfCli {
  foreach ($name in @('hf', 'huggingface-cli')) {
    $cmd = Get-Command $name -ErrorAction SilentlyContinue
    if ($cmd) { return $cmd.Source }
  }
  return $null
}

function Ensure-HfCli {
  $cli = Find-HfCli
  if ($cli) {
    Write-Ok "Hugging Face CLI: $cli"
    return $cli
  }

  Write-Warn "Hugging Face CLI not found (hf / huggingface-cli)."
  $answer = Read-Host "Install huggingface_hub now (Y/n)"
  if ($answer -and $answer -notmatch '^(y|yes)$') {
    throw "Hugging Face CLI is required. Install with: pip install -U huggingface_hub"
  }

  $pip = Get-Command pip -ErrorAction SilentlyContinue
  $py = Get-Command python -ErrorAction SilentlyContinue
  if ($pip) {
    Write-Step "Installing huggingface_hub via pip"
    & $pip.Source install -U huggingface_hub
  } elseif ($py) {
    Write-Step "Installing huggingface_hub via python -m pip"
    & $py.Source -m pip install -U huggingface_hub
  } else {
    throw "pip/python not on PATH. Install Python, then: pip install -U huggingface_hub"
  }

  $cli = Find-HfCli
  if (-not $cli) { throw "huggingface_hub installed but hf / huggingface-cli still not on PATH. Open a new terminal and retry." }
  Write-Ok "Hugging Face CLI: $cli"
  return $cli
}

function Invoke-HfDownload {
  param(
    [Parameter(Mandatory)] [string]$Cli,
    [Parameter(Mandatory)] [string]$Repo,
    [Parameter(Mandatory)] [string]$LocalDir,
    [string[]]$Include,
    [string]$File
  )

  Ensure-Dir $LocalDir
  $exe = Split-Path $Cli -Leaf
  $args = @()
  if ($exe -eq 'huggingface-cli.exe' -or $exe -eq 'huggingface-cli') {
    $args += @('download', $Repo)
  } else {
    $args += @('download', $Repo)
  }
  if ($File) { $args += $File }
  $args += @('--local-dir', $LocalDir)
  foreach ($pattern in $Include) {
    $args += @('--include', $pattern)
  }

  Write-Host "    hf $($args -join ' ')" -ForegroundColor DarkGray
  & $Cli @args
  if ($LASTEXITCODE -ne 0) {
    throw "hf download failed for $Repo (exit $LASTEXITCODE)"
  }
}

function Place-File([string]$Source, [string]$Dest) {
  Ensure-Dir (Split-Path $Dest -Parent)
  if ((Test-Path $Source) -and ((Resolve-Path $Source).Path -ne (Join-Path (Split-Path $Dest -Parent) (Split-Path $Dest -Leaf)))) {
    Copy-Item -Force $Source $Dest
  }
}

$root = Resolve-ComfyRoot $ComfyRoot
$cli = Ensure-HfCli

$checkpoints = Join-Path $root 'models\checkpoints'
$diffusion = Join-Path $root 'models\diffusion_models'
$textEncoders = Join-Path $root 'models\text_encoders'
$vae = Join-Path $root 'models\vae'
$ttsAce = Join-Path $root 'models\TTS\ACE-Step-v1-3.5B'
$customAce = Join-Path $root 'custom_nodes\ComfyUI-ACE-Step\models'
$yueRoot = Join-Path $root 'models\YuE'
$yueCustom = Join-Path $root 'custom_nodes\ComfyUI-YuE\models'

Write-Step "ComfyUI root"
Write-Host "    $root"
Write-Host "    checkpoints     $checkpoints"
Write-Host "    ACE-Step TTS    $ttsAce"
Write-Host "    ACE-Step custom $customAce"
Write-Host "    YuE             $yueRoot"

$before = Get-DirBytes (Join-Path $root 'models')
$downloaded = [System.Collections.Generic.List[string]]::new()
$skipped = [System.Collections.Generic.List[string]]::new()

function Pull-IfMissing {
  param(
    [string]$Label,
    [string]$Target,
    [long]$MinBytes,
    [scriptblock]$Download
  )
  if (-not $Force -and (Test-VerifiedTarget $Target $MinBytes)) {
    Write-Skip "$Label already present ($((Format-Bytes (Get-DirBytes $Target)))) — $Target"
    $skipped.Add($Label) | Out-Null
    return
  }
  Write-Step $Label
  & $Download
  if (-not (Test-VerifiedTarget $Target $MinBytes)) {
    throw "Download finished but $Target is missing or too small."
  }
  Write-Ok "$Label ready ($((Format-Bytes (Get-DirBytes $Target))))"
  $downloaded.Add($Label) | Out-Null
}

if (-not $SkipAceStep) {
  $aio = Join-Path $checkpoints 'ace_step_v1_3.5b.safetensors'
  Pull-IfMissing -Label 'ACE-Step v1 3.5B (native AIO checkpoint)' -Target $aio -MinBytes 1GB -Download {
    Ensure-Dir $checkpoints
    $tmp = Join-Path $env:TEMP 'aigen-ace-step-aio'
    Ensure-Dir $tmp
    Invoke-HfDownload -Cli $cli -Repo 'Comfy-Org/ACE-Step_ComfyUI_repackaged' -LocalDir $tmp -File 'all_in_one/ace_step_v1_3.5b.safetensors'
    $found = Get-ChildItem $tmp -Recurse -Filter 'ace_step_v1_3.5b.safetensors' | Select-Object -First 1
    if (-not $found) { throw 'AIO checkpoint not in Hugging Face download.' }
    Place-File $found.FullName $aio
  }

  Pull-IfMissing -Label 'ACE-Step v1 3.5B (split transformer / vocoder / UMT5)' -Target $ttsAce -MinBytes 2GB -Download {
    Invoke-HfDownload -Cli $cli -Repo 'ACE-Step/ACE-Step-v1-3.5B' -LocalDir $ttsAce -Include @('*.safetensors', '*.json')
  }

  if (Test-Path (Split-Path $customAce -Parent)) {
    Pull-IfMissing -Label 'ACE-Step custom-node models copy' -Target $customAce -MinBytes 2GB -Download {
      Ensure-Dir $customAce
      Copy-Item -Recurse -Force (Join-Path $ttsAce '*') $customAce
    }
  } else {
    Write-Skip "ComfyUI-ACE-Step custom node not installed — split weights stay in models\TTS"
  }
}

if (-not $SkipAceStep15) {
  $aio15 = Join-Path $checkpoints 'ace_step_1.5_turbo_aio.safetensors'
  Pull-IfMissing -Label 'ACE-Step 1.5 Turbo AIO checkpoint' -Target $aio15 -MinBytes 500MB -Download {
    Ensure-Dir $checkpoints
    $tmp = Join-Path $env:TEMP 'aigen-ace-step-15-aio'
    Ensure-Dir $tmp
    Invoke-HfDownload -Cli $cli -Repo 'Comfy-Org/ace_step_1.5_ComfyUI_files' -LocalDir $tmp -Include @('*aio*.safetensors', '*turbo_aio*')
    $found = Get-ChildItem $tmp -Recurse -Filter '*.safetensors' | Where-Object { $_.Name -match 'aio' } | Select-Object -First 1
    if (-not $found) {
      Invoke-HfDownload -Cli $cli -Repo 'Comfy-Org/ace_step_1.5_ComfyUI_files' -LocalDir $tmp
      $found = Get-ChildItem $tmp -Recurse -Filter '*.safetensors' | Where-Object { $_.Name -match 'aio|1\.5' } | Sort-Object Length -Descending | Select-Object -First 1
    }
    if (-not $found) { throw 'ACE-Step 1.5 AIO safetensors not found in repo. Download split files instead, or update the repo layout.' }
    Place-File $found.FullName $aio15
  }

  $diff15 = Join-Path $diffusion 'acestep_v1.5_turbo.safetensors'
  $enc06 = Join-Path $textEncoders 'qwen_0.6b_ace15.safetensors'
  $enc17 = Join-Path $textEncoders 'qwen_1.7b_ace15.safetensors'
  $vae15 = Join-Path $vae 'ace_1.5_vae.safetensors'
  $splitOk = (Test-VerifiedTarget $diff15 200MB) -and (Test-VerifiedTarget $enc06 50MB) -and (Test-VerifiedTarget $vae15 20MB)
  if ($Force -or -not $splitOk) {
    Write-Step 'ACE-Step 1.5 split files (diffusion / text encoder / VAE)'
    Ensure-Dir $diffusion
    Ensure-Dir $textEncoders
    Ensure-Dir $vae
    $tmp = Join-Path $env:TEMP 'aigen-ace-step-15-split'
    Ensure-Dir $tmp
    Invoke-HfDownload -Cli $cli -Repo 'Comfy-Org/ace_step_1.5_ComfyUI_files' -LocalDir $tmp -Include @('*.safetensors', '*.json')
    foreach ($pair in @(
      @{ Match = 'acestep_v1.5_turbo'; Dest = $diff15 },
      @{ Match = 'qwen_0.6b_ace15'; Dest = $enc06 },
      @{ Match = 'qwen_1.7b_ace15'; Dest = $enc17 },
      @{ Match = 'ace_1.5_vae'; Dest = $vae15 }
    )) {
      $hit = Get-ChildItem $tmp -Recurse -Filter '*.safetensors' | Where-Object { $_.Name -like "*$($pair.Match)*" } | Select-Object -First 1
      if ($hit) { Place-File $hit.FullName $pair.Dest; Write-Ok $hit.Name }
      else { Write-Warn "Did not find $($pair.Match)*.safetensors in the 1.5 repo" }
    }
    $downloaded.Add('ACE-Step 1.5 split files') | Out-Null
  } else {
    Write-Skip "ACE-Step 1.5 split files already present"
    $skipped.Add('ACE-Step 1.5 split files') | Out-Null
  }
}

if (-not $SkipYue) {
  $s1 = Join-Path $yueRoot 'YuE-s1-7B-anneal-en-cot'
  $s2 = Join-Path $yueRoot 'YuE-s2-1B-general'
  $codec = Join-Path $yueRoot 'xcodec_mini_infer'

  Pull-IfMissing -Label 'YuE stage-1 (s1 7B anneal EN CoT)' -Target $s1 -MinBytes 1GB -Download {
    Invoke-HfDownload -Cli $cli -Repo 'm-a-p/YuE-s1-7B-anneal-en-cot' -LocalDir $s1 -Include @('*.safetensors', '*.bin', '*.json', '*.model', '*.txt')
  }
  Pull-IfMissing -Label 'YuE stage-2 (s2 1B general)' -Target $s2 -MinBytes 200MB -Download {
    Invoke-HfDownload -Cli $cli -Repo 'm-a-p/YuE-s2-1B-general' -LocalDir $s2 -Include @('*.safetensors', '*.bin', '*.json', '*.model', '*.txt')
  }
  Pull-IfMissing -Label 'YuE vocoder (xcodec_mini_infer)' -Target $codec -MinBytes 20MB -Download {
    Invoke-HfDownload -Cli $cli -Repo 'm-a-p/xcodec_mini_infer' -LocalDir $codec -Include @('*.safetensors', '*.pth', '*.pt', '*.json', '*.yaml', '*.txt')
  }

  if (Test-Path (Split-Path $yueCustom -Parent)) {
    Pull-IfMissing -Label 'YuE custom-node models copy' -Target $yueCustom -MinBytes 1GB -Download {
      Ensure-Dir $yueCustom
      Copy-Item -Recurse -Force (Join-Path $yueRoot '*') $yueCustom
    }
  } else {
    Write-Skip "ComfyUI-YuE custom node not installed — YuE weights stay in models\YuE"
  }
}

$after = Get-DirBytes (Join-Path $root 'models')
$delta = [Math]::Max(0, $after - $before)

Write-Host ""
Write-Host "------------------------------------------------------------" -ForegroundColor DarkGray
Write-Host "Downloaded: $($downloaded.Count)" -ForegroundColor Green
foreach ($item in $downloaded) { Write-Host "  + $item" -ForegroundColor Green }
Write-Host "Skipped (already verified): $($skipped.Count)" -ForegroundColor DarkYellow
foreach ($item in $skipped) { Write-Host "  = $item" -ForegroundColor DarkYellow }
Write-Host "Models tree size: $(Format-Bytes $after)  (this run +$(Format-Bytes $delta))" -ForegroundColor Cyan
Write-Host "Restart ComfyUI if it was running so new checkpoints appear." -ForegroundColor DarkGray
Write-Host "------------------------------------------------------------" -ForegroundColor DarkGray