Finish YuE2 host setup, weight download paths, and v1 cleanup.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-14 18:41:00 -05:00
parent e82f179688
commit 35a83648d5
5 changed files with 168 additions and 63 deletions
@@ -0,0 +1,66 @@
+# Delete named YuE v1 / YuEGP weight folders only. Never recurse-wipe Comfy models\ or custom_nodes\.
+[CmdletBinding()]
+param(
+  [string]$YueGpRoot = 'C:\Users\ianjm\Development\YuEGP',
+  [string]$YueModels = '',
+  [string]$LogPath = ''
+)
+$ErrorActionPreference = 'Continue'
+if (-not $YueModels) {
+  $YueModels = Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Shared\models\yue'
+}
+if (-not $LogPath) {
+  $LogPath = Join-Path $PSScriptRoot 'cleanup-yue-v1.log'
+}
+function Log([string]$Message) {
+  $line = "$(Get-Date -Format o) $Message"
+  Add-Content -LiteralPath $LogPath -Value $line
+  Write-Host $line
+}
+function RemoveNamed([string]$Path) {
+  if (-not $Path) { return }
+  if (!(Test-Path -LiteralPath $Path)) {
+    Log "SKIP missing $Path"
+    return
+  }
+  try {
+    Remove-Item -LiteralPath $Path -Recurse -Force
+    Log "DELETED $Path"
+  } catch {
+    Log "FAIL $Path :: $($_.Exception.Message)"
+  }
+}
+
+Log "START cleanup-yue-v1 YueGpRoot=$YueGpRoot YueModels=$YueModels"
+
+# Stage 1 / Stage 2 + codec under the shared yue models dir (named folders only)
+RemoveNamed (Join-Path $YueModels 'YuE-s1-7B-anneal-en-cot')
+RemoveNamed (Join-Path $YueModels 'YuE-s2-1B-general')
+RemoveNamed (Join-Path $YueModels 'xcodec_mini_infer')
+# Loose codec files that live beside the stage folders in this install
+foreach ($name in @('ckpt_00360000.pth', 'decoder_131000.pth', 'decoder_151000.pth')) {
+  RemoveNamed (Join-Path $YueModels $name)
+}
+
+# YuEGP checkout codec copy + job dumps
+RemoveNamed (Join-Path $YueGpRoot 'inference\xcodec_mini_infer')
+RemoveNamed (Join-Path $YueGpRoot 'aigen-jobs')
+
+# Hugging Face cache dirs that are clearly YuE v1 / old codec
+$Hf = Join-Path $env:USERPROFILE '.cache\huggingface\hub'
+RemoveNamed (Join-Path $Hf 'models--m-a-p--YuE-s1-7B-anneal-en-cot')
+RemoveNamed (Join-Path $Hf 'models--m-a-p--YuE-s2-1B-general')
+RemoveNamed (Join-Path $Hf 'models--m-a-p--xcodec_mini_infer')
+
+# Leftover Comfy YuE Stage A/B weight copies by name only (no blanket models\ purge)
+$ComfyModels = Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Shared\models'
+if (Test-Path -LiteralPath $ComfyModels) {
+  Get-ChildItem -LiteralPath $ComfyModels -Recurse -ErrorAction SilentlyContinue |
+    Where-Object { $_.Name -match '^(YUE_Stage_|YuE-s1-|YuE-s2-|yue_stage)' } |
+    ForEach-Object { RemoveNamed $_.FullName }
+} else {
+  Log "SKIP missing $ComfyModels"
+}
+
+Log "DONE cleanup-yue-v1 log=$LogPath"
+Write-Host "Left in place: ACE / H3 / Klein / Krea / LTX / upscale / LoRAs / LIBRARY_DIR / yuegp code paths."
@@ -55,7 +55,7 @@ export function createGpuProxy({ target, reservation, authorized = () => true, m
    try {
      if (req.method === 'GET' || req.method === 'HEAD') await forward()
      else {
-        if (externalBusy()) { req.resume(); return reply(res, 409, { error: { message: 'YuEGP is using the GPU.' }, code: 'GPU_BUSY' }) }
+        if (externalBusy()) { req.resume(); return reply(res, 409, { error: { message: 'A music or upscale job is using the GPU.' }, code: 'GPU_BUSY' }) }
        if (!authorized(req)) { req.resume(); return reply(res, 401, { error: 'Unauthorized' }) }
        await reservation.permit(String(req.headers['x-aigen-gpu-lease'] || ''), async () => { markWork(); await forward() })
      }
@@ -1,23 +1,66 @@
 # 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]$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'
-Run $Uv @('pip', 'install', '--python', $python, 'torch', 'torchaudio', '--index-url', 'https://download.pytorch.org/whl/cu128')
+# 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')
+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))')
-@{ model='m-a-p/YuE2-3B'; vae='m-a-p/YuE2-Vae'; python=$python } | ConvertTo-Json | Set-Content -LiteralPath (Join-Path $Root 'aigen-ready.json') -Encoding utf8
-Write-Host "YuE2 ready at $Root. Set YUE2_ROOT / YUE2_PYTHON on the GPU host if needed. Weights download from Hugging Face on first song. No Comfy environment was modified."
+$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."
@@ -70,10 +70,12 @@ export function createYue2Host({ prepare, leaseValid, spawnProcess = spawn, root
        await prepare()
        if (run.cancelled || !leaseValid(lease)) throw new Error('YuE2 start cancelled or GPU reservation expired.')
        const installation = JSON.parse(readFileSync(join(repo, 'aigen-ready.json'), 'utf8').replace(/^\uFEFF/, ''))
+        const localModel = process.env.YUE2_MODEL || installation.model || 'm-a-p/YuE2-3B'
+        const localVae = process.env.YUE2_VAE || installation.vae || 'm-a-p/YuE2-Vae'
        writeFileSync(join(dir(job.id), 'request.json'), JSON.stringify({
          ...request,
-          model: process.env.YUE2_MODEL || installation.model || 'm-a-p/YuE2-3B',
-          vae: process.env.YUE2_VAE || installation.vae || 'm-a-p/YuE2-Vae'
+          model: localModel,
+          vae: localVae
        }))
        const child = spawnProcess(executable, ['-u', worker, '--root', repo, '--request', join(dir(job.id), 'request.json')],
          { cwd: repo, windowsHide: true, shell: false, stdio: ['ignore', 'pipe', 'pipe'], env: { ...process.env, PYTHONUTF8: '1', TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD: '1' } })