diff --git a/App.tsx b/App.tsx index c49eb70..0f8bdfd 100644 --- a/App.tsx +++ b/App.tsx @@ -722,6 +722,7 @@ export default function App() { lmTopP: params.lmTopP, lmNegativePrompt: params.lmNegativePrompt, lmBackend: params.lmBackend, + lmModel: params.lmModel, referenceAudioUrl: params.referenceAudioUrl, sourceAudioUrl: params.sourceAudioUrl, referenceAudioTitle: params.referenceAudioTitle, diff --git a/components/CreatePanel.tsx b/components/CreatePanel.tsx index b61041e..a041000 100644 --- a/components/CreatePanel.tsx +++ b/components/CreatePanel.tsx @@ -151,6 +151,9 @@ export const CreatePanel: React.FC = ({ const [inferenceSteps, setInferenceSteps] = useState(12); const [inferMethod, setInferMethod] = useState<'ode' | 'sde'>('ode'); const [lmBackend, setLmBackend] = useState<'pt' | 'vllm'>('pt'); + const [lmModel, setLmModel] = useState(() => { + return localStorage.getItem('ace-lmModel') || 'acestep-5Hz-lm-0.6B'; + }); const [shift, setShift] = useState(3.0); // LM Parameters (under Expert) @@ -467,6 +470,8 @@ export const CreatePanel: React.FC = ({ temperature: lmTemperature, topK: lmTopK > 0 ? lmTopK : undefined, topP: lmTopP, + lmModel: lmModel || 'acestep-5Hz-lm-0.6B', + lmBackend: lmBackend || 'pt', }, token); if (result.success) { @@ -764,6 +769,7 @@ export const CreatePanel: React.FC = ({ audioFormat, inferMethod, lmBackend, + lmModel, shift, lmTemperature, lmCfgScale, @@ -1608,6 +1614,21 @@ export const CreatePanel: React.FC = ({

PT uses less VRAM, VLLM may be faster on powerful GPUs

+ {/* LM Model */} +
+ + +

Controls the LLM used for lyrics/style enhancement. Auto-downloads if not present.

+
+ {/* Seed */}
diff --git a/server/scripts/format_sample.py b/server/scripts/format_sample.py index c9f015d..dd24b7b 100644 --- a/server/scripts/format_sample.py +++ b/server/scripts/format_sample.py @@ -17,17 +17,28 @@ sys.path.insert(0, ACESTEP_PATH) from acestep.llm_inference import LLMHandler from acestep.inference import format_sample +from pathlib import Path +from acestep.model_downloader import download_submodel # Global handler _llm_handler = None -def get_llm_handler(): +def get_llm_handler(lm_model=None, lm_backend=None): global _llm_handler if _llm_handler is None: _llm_handler = LLMHandler() - # Initialize the LLM with the 0.6B model (lighter on VRAM) checkpoint_dir = os.path.join(ACESTEP_PATH, "checkpoints") - lm_model_path = "acestep-5Hz-lm-0.6B" # Use the smaller 0.6B model + lm_model_path = lm_model or "acestep-5Hz-lm-0.6B" # Default to smallest model + backend = lm_backend or "pt" + + # Auto-download model if not present + model_dir = os.path.join(checkpoint_dir, lm_model_path) + if not os.path.exists(model_dir) or not os.listdir(model_dir): + print(f"[format_sample] Model {lm_model_path} not found, downloading...") + success, msg = download_submodel(lm_model_path, Path(checkpoint_dir)) + if not success: + raise RuntimeError(f"Failed to download model {lm_model_path}: {msg}") + print(f"[format_sample] Download complete: {msg}") if torch.cuda.is_available(): device = "cuda" elif torch.backends.mps.is_available(): @@ -38,7 +49,7 @@ def get_llm_handler(): status, success = _llm_handler.initialize( checkpoint_dir=checkpoint_dir, lm_model_path=lm_model_path, - backend="pt", # Use PyTorch backend + backend=backend, device=device, offload_to_cpu=True, ) @@ -58,9 +69,11 @@ def format_input( temperature: float = 0.85, top_k: int = 0, top_p: float = 0.9, + lm_model: str = None, + lm_backend: str = None, ): """Format caption and lyrics using the LLM.""" - handler = get_llm_handler() + handler = get_llm_handler(lm_model=lm_model, lm_backend=lm_backend) # Build user metadata for constrained decoding user_metadata = {} @@ -111,6 +124,8 @@ def main(): parser.add_argument("--temperature", type=float, default=0.85, help="LLM temperature") parser.add_argument("--top-k", type=int, default=0, help="LLM top-k sampling") parser.add_argument("--top-p", type=float, default=0.9, help="LLM top-p sampling") + parser.add_argument("--lm-model", type=str, default=None, help="LM model name (e.g. acestep-5Hz-lm-0.6B, acestep-5Hz-lm-1.7B, acestep-5Hz-lm-4B)") + parser.add_argument("--lm-backend", type=str, default=None, help="LM backend (pt or vllm)") parser.add_argument("--json", action="store_true", help="Output as JSON") args = parser.parse_args() @@ -127,6 +142,8 @@ def main(): temperature=args.temperature, top_k=args.top_k, top_p=args.top_p, + lm_model=args.lm_model, + lm_backend=args.lm_backend, ) elapsed = time.time() - start_time result["elapsed_seconds"] = elapsed diff --git a/server/src/routes/generate.ts b/server/src/routes/generate.ts index a7e05bb..de9effd 100644 --- a/server/src/routes/generate.ts +++ b/server/src/routes/generate.ts @@ -93,6 +93,7 @@ interface GenerateBody { lmTopP?: number; lmNegativePrompt?: string; lmBackend?: 'pt' | 'vllm'; + lmModel?: string; // Expert Parameters referenceAudioUrl?: string; @@ -198,6 +199,7 @@ router.post('/', authMiddleware, async (req: AuthenticatedRequest, res: Response lmTopP, lmNegativePrompt, lmBackend, + lmModel, referenceAudioUrl, sourceAudioUrl, referenceAudioTitle, @@ -264,6 +266,7 @@ router.post('/', authMiddleware, async (req: AuthenticatedRequest, res: Response lmTopP, lmNegativePrompt, lmBackend, + lmModel, referenceAudioUrl, sourceAudioUrl, referenceAudioTitle, @@ -630,7 +633,7 @@ router.get('/debug/:taskId', authMiddleware, async (req: AuthenticatedRequest, r // Format endpoint - uses LLM to enhance style/lyrics router.post('/format', authMiddleware, async (req: AuthenticatedRequest, res: Response) => { try { - const { caption, lyrics, bpm, duration, keyScale, timeSignature, temperature, topK, topP } = req.body; + const { caption, lyrics, bpm, duration, keyScale, timeSignature, temperature, topK, topP, lmModel, lmBackend } = req.body; if (!caption) { res.status(400).json({ error: 'Caption/style is required' }); @@ -660,7 +663,11 @@ router.post('/format', authMiddleware, async (req: AuthenticatedRequest, res: Re if (temperature !== undefined) args.push('--temperature', String(temperature)); if (topK && topK > 0) args.push('--top-k', String(topK)); if (topP !== undefined) args.push('--top-p', String(topP)); + if (lmModel) args.push('--lm-model', lmModel); + if (lmBackend) args.push('--lm-backend', lmBackend); + console.log(`[Format] Running: ${pythonPath} ${args.join(' ')}`); + console.log(`[Format] CWD: ${ACESTEP_DIR}`); const result = await new Promise<{ success: boolean; data?: any; error?: string }>((resolve) => { const proc = spawn(pythonPath, args, { cwd: ACESTEP_DIR, @@ -678,18 +685,29 @@ router.post('/format', authMiddleware, async (req: AuthenticatedRequest, res: Re proc.on('close', (code) => { if (code === 0 && stdout) { + // stdout may contain log lines before the JSON — extract last JSON line + const lines = stdout.trim().split('\n'); + let jsonStr = ''; + for (let i = lines.length - 1; i >= 0; i--) { + if (lines[i].startsWith('{')) { jsonStr = lines[i]; break; } + } try { - const parsed = JSON.parse(stdout); + const parsed = JSON.parse(jsonStr || stdout); resolve({ success: true, data: parsed }); } catch { + console.error('[Format] Failed to parse stdout:', stdout.slice(0, 500)); resolve({ success: false, error: 'Failed to parse format result' }); } } else { - resolve({ success: false, error: stderr || 'Format failed' }); + console.error(`[Format] Process exited with code ${code}`); + if (stdout) console.error('[Format] stdout:', stdout.slice(0, 1000)); + if (stderr) console.error('[Format] stderr:', stderr.slice(0, 1000)); + resolve({ success: false, error: stderr || stdout || `Format process exited with code ${code}` }); } }); proc.on('error', (err) => { + console.error('[Format] Spawn error:', err.message); resolve({ success: false, error: err.message }); }); }); @@ -697,10 +715,11 @@ router.post('/format', authMiddleware, async (req: AuthenticatedRequest, res: Re if (result.success && result.data) { res.json(result.data); } else { + console.error('[Format] Python error:', result.error); res.status(500).json({ success: false, error: result.error }); } } catch (error) { - console.error('Format error:', error); + console.error('[Format] Route error:', error); res.status(500).json({ error: (error as Error).message }); } }); diff --git a/server/src/services/acestep.ts b/server/src/services/acestep.ts index aed7333..fd41b63 100644 --- a/server/src/services/acestep.ts +++ b/server/src/services/acestep.ts @@ -137,6 +137,7 @@ async function submitToApi(params: GenerationParams): Promise<{ taskId: string } use_cot_language: false, // Explicitly disable CoT features that require LLM use_cot_metas: false, // Explicitly disable CoT features that require LLM lm_backend: params.lmBackend || 'pt', + lm_model_path: params.lmModel || undefined, }; if (params.duration && params.duration > 0) body.audio_duration = params.duration; @@ -388,6 +389,7 @@ export interface GenerationParams { lmTopP?: number; lmNegativePrompt?: string; lmBackend?: 'pt' | 'vllm'; + lmModel?: string; // Expert Parameters referenceAudioUrl?: string; @@ -675,6 +677,7 @@ async function processGeneration( if (params.lmTopP !== undefined) args.push('--lm-top-p', String(params.lmTopP)); if (params.lmNegativePrompt) args.push('--lm-negative-prompt', params.lmNegativePrompt); if (params.lmBackend) args.push('--lm-backend', params.lmBackend); + if (params.lmModel) args.push('--lm-model', params.lmModel); if (params.useCotMetas === false) args.push('--no-cot-metas'); if (params.useCotCaption === false) args.push('--no-cot-caption'); if (params.useCotLanguage === false) args.push('--no-cot-language'); diff --git a/services/api.ts b/services/api.ts index fb3a186..979fdf7 100644 --- a/services/api.ts +++ b/services/api.ts @@ -234,6 +234,7 @@ export interface GenerationParams { lmTopP?: number; lmNegativePrompt?: string; lmBackend?: 'pt' | 'vllm'; + lmModel?: string; // Expert Parameters referenceAudioUrl?: string; @@ -317,6 +318,8 @@ export const generateApi = { temperature?: number; topK?: number; topP?: number; + lmModel?: string; + lmBackend?: string; }, token: string): Promise<{ success: boolean; caption?: string; diff --git a/types.ts b/types.ts index 1a14e10..d19d51c 100644 --- a/types.ts +++ b/types.ts @@ -87,6 +87,7 @@ export interface GenerationParams { lmTopP: number; lmNegativePrompt: string; lmBackend?: 'pt' | 'vllm'; + lmModel?: string; // Expert Parameters referenceAudioUrl?: string;