3dtours/ace-step-ui
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add LM Model selector (0.6B/1.7B/4B) with auto-download

Browse Source 
- Add lmModel parameter through full chain (types, API, routes, service)
- Add LM Model dropdown in Advanced Settings (defaults to 0.6B)
- Pass lmModel and lmBackend to format/enhance endpoint
- Update format_sample.py to accept --lm-model and --lm-backend args
- Auto-download model from HuggingFace if not present locally
- Persist model selection in localStorage
- Improve format route error logging with exit code and stdout/stderr

Fixes #9
This commit is contained in:
fspecii
2026-02-05 23:55:55 +02:00
parent dade566647
commit 6f4d50ae18
7 changed files with 74 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/scripts/format_sample.py
+22 -5
View File
@@ -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
server/src/routes/generate.ts
+23 -4
View File
@@ -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 });
}
});
server/src/services/acestep.ts
+3
View File
@@ -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');