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

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