# Standalone YuEGP only. Never install into Comfy's Python environment.
[CmdletBinding()]
param(
  [string]$Root = (Join-Path (Split-Path $PSScriptRoot -Parent) '..\YuEGP'),
  [string]$Uv = 'uv',
  [string]$Models = (Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Shared\models\yue')
)
$ErrorActionPreference = 'Stop'
$Root = [IO.Path]::GetFullPath($Root)
$revision = '2d72ff734b7a127324353c0dcd0f95ca4cc0b797'
function Run([string]$Exe, [string[]]$Arguments) {
  & $Exe @Arguments
  if ($LASTEXITCODE -ne 0) { throw "$Exe failed ($LASTEXITCODE)" }
}
if (!(Test-Path -LiteralPath $Root)) {
  Run git @('clone', 'https://github.com/deepbeepmeep/YuEGP.git', $Root)
  Run git @('-C', $Root, 'checkout', '--detach', $revision)
}
$head = & git -C $Root rev-parse HEAD
if ($head.Trim() -ne $revision) { throw "YuEGP must be checked out at $revision. No existing checkout was changed." }
Run $Uv @('venv', '--allow-existing', '--python', '3.10', (Join-Path $Root '.venv'))
$python = Join-Path $Root '.venv\Scripts\python.exe'
Run $Uv @('pip', 'install', '--python', $python, 'torch==2.7.1', 'torchaudio==2.7.1', '--index-url', 'https://download.pytorch.org/whl/cu128')
Run $Uv @('pip', 'install', '--python', $python, 'torch==2.7.1', 'torchaudio==2.7.1', 'numpy<2', 'scipy==1.10.1', 'transformers==4.48.1', 'mmgp==3.1.4-post15', 'peft==0.14.0', 'omegaconf', 'einops', 'sentencepiece', 'tqdm', 'tensorboard', 'descript-audiotools==0.7.2', 'descript-audio-codec', 'soundfile', 'psutil')
# Apply YuEGP's own upstream optimization files to its own environment only.
Copy-Item -LiteralPath (Join-Path $Root 'transformers\generation\utils.py') -Destination (Join-Path $Root '.venv\Lib\site-packages\transformers\generation\utils.py')
Copy-Item -LiteralPath (Join-Path $Root 'transformers\models\llama\modeling_llama.py') -Destination (Join-Path $Root '.venv\Lib\site-packages\transformers\models\llama\modeling_llama.py')
$codec = Join-Path $Root 'inference\xcodec_mini_infer'
if (!(Test-Path -LiteralPath $codec)) {
  $previousSmudge = $env:GIT_LFS_SKIP_SMUDGE
  try { $env:GIT_LFS_SKIP_SMUDGE = '1'; Run git @('clone', '--depth', '1', 'https://huggingface.co/m-a-p/xcodec_mini_infer', $codec) }
  finally { $env:GIT_LFS_SKIP_SMUDGE = $previousSmudge }
}
# Reuse local codec weights. Download missing assets explicitly from their HF repo.
$weights = @('final_ckpt\ckpt_00360000.pth', 'decoders\decoder_131000.pth', 'decoders\decoder_151000.pth', 'semantic_ckpts\hf_1_325000\pytorch_model.bin')
foreach ($weight in $weights) {
  $target = Join-Path $codec $weight
  if ((Test-Path -LiteralPath $target) -and (Get-Item -LiteralPath $target).Length -gt 1024) { continue }
  $source = Join-Path (Join-Path $Models 'xcodec_mini_infer') $weight
  if ((Test-Path -LiteralPath $source) -and (Get-Item -LiteralPath $source).Length -gt 1024) { Copy-Item -LiteralPath $source -Destination $target }
  else {
    Run $python @('-c', 'from huggingface_hub import hf_hub_download; import sys; hf_hub_download("m-a-p/xcodec_mini_infer",sys.argv[1],local_dir=sys.argv[2])', $weight.Replace('\','/'), $codec)
  }
}
Run $python @('-c', 'import torch,transformers; from mmgp import offload; assert torch.cuda.is_available(); print(torch.__version__,transformers.__version__,torch.cuda.get_device_name(0))')
@{ revision=$revision; modelsDir=$Models; profile=1; compile=$false } | ConvertTo-Json | Set-Content -LiteralPath (Join-Path $Root 'aigen-ready.json') -Encoding utf8
Write-Host "YuEGP ready at $Root. Set YUEGP_ROOT on the GPU host if using a different location. No Comfy environment was modified."