Various UI improvements

2026-02-04 23:15:00 -05:00
parent b75bea3860
commit 424bd3fd25
9 changed files with 262 additions and 67 deletions
@@ -9,6 +9,7 @@ import json
 import os
 import sys
 import time
+import torch

 # Get ACE-Step path from environment or use default
 ACESTEP_PATH = os.environ.get('ACESTEP_PATH', '/home/ambsd/Desktop/aceui/ACE-Step-1.5')
@@ -27,11 +28,17 @@ _llm_handler = None
 def get_handlers():
    global _handler, _llm_handler
    if _handler is None:
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif torch.backends.mps.is_available():
+            device = "mps"
+        else:
+            device = "cpu"
        _handler = AceStepHandler()
        _handler.initialize_service(
            project_root=ACESTEP_PATH,
            config_path="acestep-v15-turbo",
-            device="cuda",
+            device=device,
            offload_to_cpu=True,  # For 12GB GPU
        )
        _llm_handler = LLMHandler()  # Create but don't initialize (not enough VRAM)
@@ -96,6 +96,8 @@ interface GenerateBody {
  // Expert Parameters
  referenceAudioUrl?: string;
  sourceAudioUrl?: string;
+  referenceAudioTitle?: string;
+  sourceAudioTitle?: string;
  audioCodes?: string;
  repaintingStart?: number;
  repaintingEnd?: number;
@@ -196,6 +198,8 @@ router.post('/', authMiddleware, async (req: AuthenticatedRequest, res: Response
      lmNegativePrompt,
      referenceAudioUrl,
      sourceAudioUrl,
+      referenceAudioTitle,
+      sourceAudioTitle,
      audioCodes,
      repaintingStart,
      repaintingEnd,
@@ -259,6 +263,8 @@ router.post('/', authMiddleware, async (req: AuthenticatedRequest, res: Response
      lmNegativePrompt,
      referenceAudioUrl,
      sourceAudioUrl,
+      referenceAudioTitle,
+      sourceAudioTitle,
      audioCodes,
      repaintingStart,
      repaintingEnd,
@@ -547,6 +553,60 @@ router.get('/health', async (_req, res: Response) => {
  }
 });

+router.get('/limits', async (_req, res: Response) => {
+  try {
+    const { spawn } = await import('child_process');
+    const ACESTEP_DIR = process.env.ACESTEP_PATH || path.resolve(path.dirname(fileURLToPath(import.meta.url)), '../../../../ACE-Step-1.5');
+    const __filename = fileURLToPath(import.meta.url);
+    const __dirname = path.dirname(__filename);
+    const SCRIPTS_DIR = path.join(__dirname, '../../scripts');
+    const LIMITS_SCRIPT = path.join(SCRIPTS_DIR, 'get_limits.py');
+    const pythonPath = resolvePythonPath(ACESTEP_DIR);
+
+    const result = await new Promise<{ success: boolean; data?: any; error?: string }>((resolve) => {
+      const proc = spawn(pythonPath, [LIMITS_SCRIPT], {
+        cwd: ACESTEP_DIR,
+        env: {
+          ...process.env,
+          ACESTEP_PATH: ACESTEP_DIR,
+        },
+      });
+
+      let stdout = '';
+      let stderr = '';
+
+      proc.stdout.on('data', (data) => { stdout += data.toString(); });
+      proc.stderr.on('data', (data) => { stderr += data.toString(); });
+
+      proc.on('close', (code) => {
+        if (code === 0 && stdout) {
+          try {
+            const parsed = JSON.parse(stdout);
+            resolve({ success: true, data: parsed });
+          } catch {
+            resolve({ success: false, error: 'Failed to parse limits result' });
+          }
+        } else {
+          resolve({ success: false, error: stderr || 'Failed to read limits' });
+        }
+      });
+
+      proc.on('error', (err) => {
+        resolve({ success: false, error: err.message });
+      });
+    });
+
+    if (result.success && result.data) {
+      res.json(result.data);
+    } else {
+      res.status(500).json({ error: result.error || 'Failed to load limits' });
+    }
+  } catch (error) {
+    console.error('Limits error:', error);
+    res.status(500).json({ error: (error as Error).message });
+  }
+});
+
 router.get('/debug/:taskId', authMiddleware, async (req: AuthenticatedRequest, res: Response) => {
  try {
    const rawResponse = getJobRawResponse(req.params.taskId);
@@ -125,7 +125,6 @@ async function submitToApi(params: GenerationParams): Promise<{ taskId: string }
  const body: Record<string, unknown> = {
    prompt,
    lyrics,
-    audio_duration: params.duration ?? 60,
    batch_size: params.batchSize ?? 1,
    inference_steps: params.inferenceSteps ?? 8,
    guidance_scale: params.guidanceScale ?? 10.0,
@@ -139,6 +138,7 @@ async function submitToApi(params: GenerationParams): Promise<{ taskId: string }
    use_cot_metas: false, // Explicitly disable CoT features that require LLM
  };

+  if (params.duration && params.duration > 0) body.audio_duration = params.duration;
  if (params.bpm && params.bpm > 0) body.bpm = params.bpm;
  if (params.keyScale) body.key_scale = params.keyScale;
  if (params.timeSignature) body.time_signature = params.timeSignature;
@@ -166,20 +166,42 @@ async function submitToApi(params: GenerationParams): Promise<{ taskId: string }
  if (params.cfgIntervalStart !== undefined && params.cfgIntervalStart > 0) body.cfg_interval_start = params.cfgIntervalStart;
  if (params.cfgIntervalEnd !== undefined && params.cfgIntervalEnd < 1.0) body.cfg_interval_end = params.cfgIntervalEnd;

+  const resolveAudioPath = (audioUrl: string): string => {
+    if (audioUrl.startsWith('/audio/')) {
+      return path.join(AUDIO_DIR, audioUrl.replace('/audio/', ''));
+    }
+    if (audioUrl.startsWith('http')) {
+      try {
+        const parsed = new URL(audioUrl);
+        if (parsed.pathname.startsWith('/audio/')) {
+          return path.join(AUDIO_DIR, parsed.pathname.replace('/audio/', ''));
+        }
+      } catch {
+        // fall through
+      }
+    }
+    return audioUrl;
+  };
+
+  // Guard: cover/audio2audio requires a source or audio codes
+  if ((params.taskType === 'cover' || params.taskType === 'audio2audio') && !params.sourceAudioUrl && !params.audioCodes) {
+    throw new Error(`task_type='${params.taskType}' requires a source audio or audio codes`);
+  }
+
  // Handle reference audio - need to pass file path
  if (params.referenceAudioUrl) {
-    let refAudioPath = params.referenceAudioUrl;
-    if (refAudioPath.startsWith('/audio/')) {
-      refAudioPath = path.join(AUDIO_DIR, refAudioPath.replace('/audio/', ''));
-    }
-    body.reference_audio_path = refAudioPath;
+    body.reference_audio_path = resolveAudioPath(params.referenceAudioUrl);
  }
  if (params.sourceAudioUrl) {
-    let srcAudioPath = params.sourceAudioUrl;
-    if (srcAudioPath.startsWith('/audio/')) {
-      srcAudioPath = path.join(AUDIO_DIR, srcAudioPath.replace('/audio/', ''));
-    }
-    body.src_audio_path = srcAudioPath;
+    body.src_audio_path = resolveAudioPath(params.sourceAudioUrl);
+  }
+
+  if (params.taskType === 'cover' || params.taskType === 'audio2audio') {
+    console.log(`[ACE-Step] cover/audio2audio inputs`, {
+      reference_audio_path: body.reference_audio_path,
+      src_audio_path: body.src_audio_path,
+      has_audio_codes: Boolean(params.audioCodes),
+    });
  }

  const response = await fetch(`${ACESTEP_API}/release_task`, {
@@ -263,6 +285,20 @@ async function pollApiResult(taskId: string, maxWaitMs = 600000): Promise<ApiTas
      throw new Error(`Generation failed on API side: ${details}`);
    }

+    // Log progress while processing (if provided)
+    if (taskData.result) {
+      try {
+        const resultData = typeof taskData.result === 'string' ? JSON.parse(taskData.result) : taskData.result;
+        const item = Array.isArray(resultData) ? resultData[0] : resultData;
+        if (item && typeof item === 'object' && typeof (item as any).progress === 'number') {
+          const pct = Math.round((item as any).progress * 100);
+          console.log(`[ACE-Step] API task ${taskId} progress: ${pct}%`);
+        }
+      } catch {
+        // ignore parse failures
+      }
+    }
+
    // Still processing
    await new Promise(resolve => setTimeout(resolve, pollInterval));
  }
@@ -349,6 +385,8 @@ export interface GenerationParams {
  // Expert Parameters
  referenceAudioUrl?: string;
  sourceAudioUrl?: string;
+  referenceAudioTitle?: string;
+  sourceAudioTitle?: string;
  audioCodes?: string;
  repaintingStart?: number;
  repaintingEnd?: number;
@@ -403,6 +441,8 @@ interface ActiveJob {
  processPromise?: Promise<void>;
  rawResponse?: unknown;
  queuePosition?: number;
+  progress?: number;
+  stage?: string;
 }

 const activeJobs = new Map<string, ActiveJob>();
@@ -466,6 +506,8 @@ async function processQueue(): Promise<void> {

 // Submit generation job to queue
 export async function generateMusicViaAPI(params: GenerationParams): Promise<{ jobId: string }> {
+  // Force a fresh API availability check when starting a job
+  resetApiCache();
  const jobId = `job_${Date.now()}_${Math.random().toString(36).slice(2, 9)}`;

  const job: ActiveJob = {
@@ -573,9 +615,10 @@ async function processGeneration(
    const jobOutputDir = path.join(ACESTEP_DIR, 'output', jobId);
    await mkdir(jobOutputDir, { recursive: true });

+    const durationToSend = params.duration && params.duration > 0 ? params.duration : 60;
    const args = [
      '--prompt', prompt,
-      '--duration', String(params.duration ?? 60),
+      '--duration', String(durationToSend),
      '--batch-size', String(params.batchSize ?? 1),
      '--infer-steps', String(params.inferenceSteps ?? 8),
      '--guidance-scale', String(params.guidanceScale ?? 10.0),
@@ -701,7 +744,6 @@ function runPythonGeneration(scriptArgs: string[]): Promise<PythonResult> {
      cwd: ACESTEP_DIR,
      env: {
        ...process.env,
-        CUDA_VISIBLE_DEVICES: '0',
        ACESTEP_PATH: ACESTEP_DIR,
      },
    });
@@ -863,13 +905,16 @@ export async function getJobStatus(jobId: string): Promise<JobStatus> {

          const item = Array.isArray(resultData) ? resultData[0] : resultData;
          if (item && typeof item === 'object') {
-            const progress = typeof (item as any).progress === 'number' ? (item as any).progress : undefined;
+            const rawProgress = (item as any).progress;
+            const progress = Number.isFinite(Number(rawProgress)) ? Number(rawProgress) : undefined;
            const stage = typeof (item as any).stage === 'string' ? (item as any).stage : undefined;
+            if (progress !== undefined) job.progress = progress;
+            if (stage) job.stage = stage;
            return {
              status: job.status,
              etaSeconds: Math.max(0, 180 - elapsed),
-              progress,
-              stage,
+              progress: progress ?? job.progress,
+              stage: stage ?? job.stage,
            };
          }
        }
@@ -882,6 +927,8 @@ export async function getJobStatus(jobId: string): Promise<JobStatus> {
  return {
    status: job.status,
    etaSeconds: Math.max(0, 180 - elapsed), // 3 min estimate
+    progress: job.progress,
+    stage: job.stage,
  };
 }