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

- 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

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,

components/CreatePanel.tsx

@@ -151,6 +151,9 @@ export const CreatePanel: React.FC<CreatePanelProps> = ({
   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<CreatePanelProps> = ({
         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<CreatePanelProps> = ({
         audioFormat,
         inferMethod,
         lmBackend,
+        lmModel,
         shift,
         lmTemperature,
         lmCfgScale,
@@ -1608,6 +1614,21 @@ export const CreatePanel: React.FC<CreatePanelProps> = ({
               <p className="text-[10px] text-zinc-500">PT uses less VRAM, VLLM may be faster on powerful GPUs</p>
             </div>
 
+            {/* LM Model */}
+            <div className="space-y-1.5">
+              <label className="text-xs font-medium text-zinc-600 dark:text-zinc-400">LM Model</label>
+              <select
+                value={lmModel}
+                onChange={(e) => { const v = e.target.value; setLmModel(v); localStorage.setItem('ace-lmModel', v); }}
+                className="w-full bg-zinc-50 dark:bg-black/20 border border-zinc-200 dark:border-white/10 rounded-lg px-2 py-1.5 text-xs text-zinc-900 dark:text-white focus:outline-none"
+              >
+                <option value="acestep-5Hz-lm-0.6B">0.6B (Lightest, ~0.5 GB VRAM)</option>
+                <option value="acestep-5Hz-lm-1.7B">1.7B (Balanced, ~1.5 GB VRAM)</option>
+                <option value="acestep-5Hz-lm-4B">4B (Best quality, ~4 GB VRAM)</option>
+              </select>
+              <p className="text-[10px] text-zinc-500">Controls the LLM used for lyrics/style enhancement. Auto-downloads if not present.</p>
+            </div>
+
             {/* Seed */}
             <div className="space-y-2">
               <div className="flex items-center justify-between">

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

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 });
   }
 });

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');

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;

types.ts

@@ -87,6 +87,7 @@ export interface GenerationParams {
   lmTopP: number;
   lmNegativePrompt: string;
   lmBackend?: 'pt' | 'vllm';
+  lmModel?: string;
 
   // Expert Parameters
   referenceAudioUrl?: string;