acestep/api_server.py

"""FastAPI server for ACE-Step V1.5.

Endpoints:
- POST /release_task          Create music generation task
- POST /query_result          Batch query task results
- POST /create_random_sample  Generate random music parameters via LLM
- POST /format_input          Format and enhance lyrics/caption via LLM
- GET  /v1/models             List available models
- GET  /v1/audio              Download audio file
- GET  /health                Health check

NOTE:
- In-memory queue and job store -> run uvicorn with workers=1.
"""

from __future__ import annotations

import asyncio
import glob
import json
import os
import random
import sys
import time
import traceback
import tempfile
import urllib.parse
from collections import deque
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from dataclasses import dataclass
from pathlib import Path
from threading import Lock
from typing import Any, Dict, List, Literal, Optional, Union
from uuid import uuid4
from loguru import logger

try:
    from dotenv import load_dotenv
except ImportError:  # Optional dependency
    load_dotenv = None  # type: ignore

from fastapi import FastAPI, HTTPException, Request, Depends, Header
from pydantic import BaseModel, Field
from starlette.datastructures import UploadFile as StarletteUploadFile

from acestep.handler import AceStepHandler
from acestep.llm_inference import LLMHandler
from acestep.constants import (
    DEFAULT_DIT_INSTRUCTION,
    DEFAULT_LM_INSTRUCTION,
    TASK_INSTRUCTIONS,
)
from acestep.inference import (
    GenerationParams,
    GenerationConfig,
    generate_music,
    create_sample,
    format_sample,
)
from acestep.gradio_ui.events.results_handlers import _build_generation_info
from acestep.gpu_config import (
    get_gpu_config,
    get_gpu_memory_gb,
    print_gpu_config_info,
    set_global_gpu_config,
    get_recommended_lm_model,
    is_lm_model_supported,
    GPUConfig,
    VRAM_16GB_MIN_GB,
)


# =============================================================================
# Model Auto-Download Support
# =============================================================================

# Model name to repository mapping
MODEL_REPO_MAPPING = {
    # Main unified repository (contains: acestep-v15-turbo, acestep-5Hz-lm-1.7B, Qwen3-Embedding-0.6B, vae)
    "acestep-v15-turbo": "ACE-Step/Ace-Step1.5",
    "acestep-5Hz-lm-1.7B": "ACE-Step/Ace-Step1.5",
    "vae": "ACE-Step/Ace-Step1.5",
    "Qwen3-Embedding-0.6B": "ACE-Step/Ace-Step1.5",
    # Separate model repositories
    "acestep-5Hz-lm-0.6B": "ACE-Step/acestep-5Hz-lm-0.6B",
    "acestep-5Hz-lm-4B": "ACE-Step/acestep-5Hz-lm-4B",
    "acestep-v15-base": "ACE-Step/acestep-v15-base",
    "acestep-v15-sft": "ACE-Step/acestep-v15-sft",
    "acestep-v15-turbo-shift3": "ACE-Step/acestep-v15-turbo-shift3",
}

DEFAULT_REPO_ID = "ACE-Step/Ace-Step1.5"


def _can_access_google(timeout: float = 3.0) -> bool:
    """Check if Google is accessible (to determine HuggingFace vs ModelScope)."""
    import socket
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        sock.settimeout(timeout)
        sock.connect(("www.google.com", 443))
        return True
    except (socket.timeout, socket.error, OSError):
        return False
    finally:
        sock.close()


def _download_from_huggingface(repo_id: str, local_dir: str, model_name: str) -> str:
    """Download model from HuggingFace Hub."""
    from huggingface_hub import snapshot_download

    is_unified_repo = repo_id == DEFAULT_REPO_ID or repo_id == "ACE-Step/Ace-Step1.5"

    if is_unified_repo:
        download_dir = local_dir
        print(f"[Model Download] Downloading unified repo {repo_id} to {download_dir}...")
    else:
        download_dir = os.path.join(local_dir, model_name)
        os.makedirs(download_dir, exist_ok=True)
        print(f"[Model Download] Downloading {model_name} from {repo_id} to {download_dir}...")

    snapshot_download(
        repo_id=repo_id,
        local_dir=download_dir,
        local_dir_use_symlinks=False,
    )

    return os.path.join(local_dir, model_name)


def _download_from_modelscope(repo_id: str, local_dir: str, model_name: str) -> str:
    """Download model from ModelScope."""
    from modelscope import snapshot_download

    is_unified_repo = repo_id == DEFAULT_REPO_ID or repo_id == "ACE-Step/Ace-Step1.5"

    if is_unified_repo:
        download_dir = local_dir
        print(f"[Model Download] Downloading unified repo {repo_id} from ModelScope to {download_dir}...")
    else:
        download_dir = os.path.join(local_dir, model_name)
        os.makedirs(download_dir, exist_ok=True)
        print(f"[Model Download] Downloading {model_name} from ModelScope {repo_id} to {download_dir}...")

    # ModelScope snapshot_download returns the cache path
    # Use cache_dir parameter for better compatibility across versions
    try:
        # Try with local_dir first (newer versions)
        result_path = snapshot_download(
            model_id=repo_id,
            local_dir=download_dir,
        )
        print(f"[Model Download] ModelScope download completed: {result_path}")
    except TypeError:
        # Fallback to cache_dir for older versions
        print("[Model Download] Retrying with cache_dir parameter...")
        result_path = snapshot_download(
            model_id=repo_id,
            cache_dir=download_dir,
        )
        print(f"[Model Download] ModelScope download completed: {result_path}")

    return os.path.join(local_dir, model_name)


def _ensure_model_downloaded(model_name: str, checkpoint_dir: str) -> str:
    """
    Ensure model is downloaded. Auto-detect source based on network.

    Args:
        model_name: Model directory name (e.g., "acestep-v15-turbo")
        checkpoint_dir: Target checkpoint directory

    Returns:
        Path to the model directory
    """
    model_path = os.path.join(checkpoint_dir, model_name)

    # Check if model already exists
    if os.path.exists(model_path) and os.listdir(model_path):
        print(f"[Model Download] Model {model_name} already exists at {model_path}")
        return model_path

    # Get repository ID
    repo_id = MODEL_REPO_MAPPING.get(model_name, DEFAULT_REPO_ID)

    print(f"[Model Download] Model {model_name} not found, checking network...")

    # Check for user preference
    prefer_source = os.environ.get("ACESTEP_DOWNLOAD_SOURCE", "").lower()

    # Determine download source
    if prefer_source == "huggingface":
        use_huggingface = True
        print("[Model Download] User preference: HuggingFace Hub")
    elif prefer_source == "modelscope":
        use_huggingface = False
        print("[Model Download] User preference: ModelScope")
    else:
        use_huggingface = _can_access_google()
        print(f"[Model Download] Auto-detected: {'HuggingFace Hub' if use_huggingface else 'ModelScope'}")

    if use_huggingface:
        print("[Model Download] Using HuggingFace Hub...")
        try:
            return _download_from_huggingface(repo_id, checkpoint_dir, model_name)
        except Exception as e:
            print(f"[Model Download] HuggingFace download failed: {e}")
            print("[Model Download] Falling back to ModelScope...")
            return _download_from_modelscope(repo_id, checkpoint_dir, model_name)
    else:
        print("[Model Download] Using ModelScope...")
        try:
            return _download_from_modelscope(repo_id, checkpoint_dir, model_name)
        except Exception as e:
            print(f"[Model Download] ModelScope download failed: {e}")
            print("[Model Download] Trying HuggingFace as fallback...")
            return _download_from_huggingface(repo_id, checkpoint_dir, model_name)


def _get_project_root() -> str:
    current_file = os.path.abspath(__file__)
    return os.path.dirname(os.path.dirname(current_file))


# =============================================================================
# Constants
# =============================================================================

RESULT_KEY_PREFIX = "ace_step_v1.5_"
RESULT_EXPIRE_SECONDS = 7 * 24 * 60 * 60  # 7 days
TASK_TIMEOUT_SECONDS = 3600  # 1 hour
JOB_STORE_CLEANUP_INTERVAL = 300  # 5 minutes - interval for cleaning up old jobs
JOB_STORE_MAX_AGE_SECONDS = 86400  # 24 hours - completed jobs older than this will be cleaned
STATUS_MAP = {"queued": 0, "running": 0, "succeeded": 1, "failed": 2}

LM_DEFAULT_TEMPERATURE = 0.85
LM_DEFAULT_CFG_SCALE = 2.5
LM_DEFAULT_TOP_P = 0.9


def _wrap_response(data: Any, code: int = 200, error: Optional[str] = None) -> Dict[str, Any]:
    """Wrap response data in standard format."""
    return {
        "data": data,
        "code": code,
        "error": error,
        "timestamp": int(time.time() * 1000),
        "extra": None,
    }


# =============================================================================
# Example Data for Random Sample
# =============================================================================

SIMPLE_MODE_EXAMPLES_DIR = os.path.join(_get_project_root(), "examples", "simple_mode")
CUSTOM_MODE_EXAMPLES_DIR = os.path.join(_get_project_root(), "examples", "text2music")


def _load_all_examples(sample_mode: str = "simple_mode") -> List[Dict[str, Any]]:
    """Load all example data files from the examples directory."""
    examples = []
    examples_dir = SIMPLE_MODE_EXAMPLES_DIR if sample_mode == "simple_mode" else CUSTOM_MODE_EXAMPLES_DIR
    pattern = os.path.join(examples_dir, "example_*.json")

    for filepath in glob.glob(pattern):
        try:
            with open(filepath, 'r', encoding='utf-8') as f:
                data = json.load(f)
                examples.append(data)
        except Exception as e:
            print(f"[API Server] Failed to load example file {filepath}: {e}")

    return examples


# Pre-load example data at module load time
SIMPLE_EXAMPLE_DATA: List[Dict[str, Any]] = _load_all_examples(sample_mode="simple_mode")
CUSTOM_EXAMPLE_DATA: List[Dict[str, Any]] = _load_all_examples(sample_mode="custom_mode")

# =============================================================================
# API Key Authentication
# =============================================================================

_api_key: Optional[str] = None


def set_api_key(key: Optional[str]):
    """Set the API key for authentication"""
    global _api_key
    _api_key = key


def verify_token_from_request(body: dict, authorization: Optional[str] = None) -> Optional[str]:
    """
    Verify API key from request body (ai_token) or Authorization header.
    Returns the token if valid, None if no auth required.
    """
    if _api_key is None:
        return None  # No auth required

    # Try ai_token from body first
    ai_token = body.get("ai_token") if body else None
    if ai_token:
        if ai_token == _api_key:
            return ai_token
        raise HTTPException(status_code=401, detail="Invalid ai_token")

    # Fallback to Authorization header
    if authorization:
        if authorization.startswith("Bearer "):
            token = authorization[7:]
        else:
            token = authorization
        if token == _api_key:
            return token
        raise HTTPException(status_code=401, detail="Invalid API key")

    # No token provided but auth is required
    raise HTTPException(status_code=401, detail="Missing ai_token or Authorization header")


async def verify_api_key(authorization: Optional[str] = Header(None)):
    """Verify API key from Authorization header (legacy, for non-body endpoints)"""
    if _api_key is None:
        return  # No auth required

    if not authorization:
        raise HTTPException(status_code=401, detail="Missing Authorization header")

    # Support "Bearer <key>" format
    if authorization.startswith("Bearer "):
        token = authorization[7:]
    else:
        token = authorization

    if token != _api_key:
        raise HTTPException(status_code=401, detail="Invalid API key")

# Parameter aliases for request parsing
PARAM_ALIASES = {
    "prompt": ["prompt", "caption"],
    "lyrics": ["lyrics"],
    "thinking": ["thinking"],
    "analysis_only": ["analysis_only", "analysisOnly"],
    "full_analysis_only": ["full_analysis_only", "fullAnalysisOnly"],
    "sample_mode": ["sample_mode", "sampleMode"],
    "sample_query": ["sample_query", "sampleQuery", "description", "desc"],
    "use_format": ["use_format", "useFormat", "format"],
    "model": ["model", "model_name", "modelName", "dit_model", "ditModel"],
    "key_scale": ["key_scale", "keyscale", "keyScale", "key"],
    "time_signature": ["time_signature", "timesignature", "timeSignature"],
    "audio_duration": ["audio_duration", "duration", "audioDuration", "target_duration", "targetDuration"],
    "vocal_language": ["vocal_language", "vocalLanguage", "language"],
    "bpm": ["bpm"],
    "inference_steps": ["inference_steps", "inferenceSteps"],
    "guidance_scale": ["guidance_scale", "guidanceScale"],
    "use_random_seed": ["use_random_seed", "useRandomSeed"],
    "seed": ["seed"],

    "audio_cover_strength": ["audio_cover_strength", "audioCoverStrength"],
    "reference_audio_path": ["reference_audio_path", "ref_audio_path", "referenceAudioPath", "refAudioPath"],
    "src_audio_path": ["src_audio_path", "ctx_audio_path", "sourceAudioPath", "srcAudioPath", "ctxAudioPath"],
    "task_type": ["task_type", "taskType"],
    "infer_method": ["infer_method", "inferMethod"],
    "use_tiled_decode": ["use_tiled_decode", "useTiledDecode"],
    "constrained_decoding": ["constrained_decoding", "constrainedDecoding", "constrained"],
    "constrained_decoding_debug": ["constrained_decoding_debug", "constrainedDecodingDebug"],
    "use_cot_caption": ["use_cot_caption", "cot_caption", "cot-caption"],
    "use_cot_language": ["use_cot_language", "cot_language", "cot-language"],
    "is_format_caption": ["is_format_caption", "isFormatCaption"],
    "allow_lm_batch": ["allow_lm_batch", "allowLmBatch", "parallel_thinking"],
}


def _parse_description_hints(description: str) -> tuple[Optional[str], bool]:
    """
    Parse a description string to extract language code and instrumental flag.
    
    This function analyzes user descriptions like "Pop rock. English" or "piano solo"
    to detect:
    - Language: Maps language names to ISO codes (e.g., "English" -> "en")
    - Instrumental: Detects patterns indicating instrumental/no-vocal music
    
    Args:
        description: User's natural language music description
        
    Returns:
        (language_code, is_instrumental) tuple:
        - language_code: ISO language code (e.g., "en", "zh") or None if not detected
        - is_instrumental: True if description indicates instrumental music
    """
    import re
    
    if not description:
        return None, False
    
    description_lower = description.lower().strip()
    
    # Language mapping: input patterns -> ISO code
    language_mapping = {
        'english': 'en', 'en': 'en',
        'chinese': 'zh', '中文': 'zh', 'zh': 'zh', 'mandarin': 'zh',
        'japanese': 'ja', '日本語': 'ja', 'ja': 'ja',
        'korean': 'ko', '한국어': 'ko', 'ko': 'ko',
        'spanish': 'es', 'español': 'es', 'es': 'es',
        'french': 'fr', 'français': 'fr', 'fr': 'fr',
        'german': 'de', 'deutsch': 'de', 'de': 'de',
        'italian': 'it', 'italiano': 'it', 'it': 'it',
        'portuguese': 'pt', 'português': 'pt', 'pt': 'pt',
        'russian': 'ru', 'русский': 'ru', 'ru': 'ru',
        'bengali': 'bn', 'bn': 'bn',
        'hindi': 'hi', 'hi': 'hi',
        'arabic': 'ar', 'ar': 'ar',
        'thai': 'th', 'th': 'th',
        'vietnamese': 'vi', 'vi': 'vi',
        'indonesian': 'id', 'id': 'id',
        'turkish': 'tr', 'tr': 'tr',
        'dutch': 'nl', 'nl': 'nl',
        'polish': 'pl', 'pl': 'pl',
    }
    
    # Detect language
    detected_language = None
    for lang_name, lang_code in language_mapping.items():
        if len(lang_name) <= 2:
            pattern = r'(?:^|\s|[.,;:!?])' + re.escape(lang_name) + r'(?:$|\s|[.,;:!?])'
        else:
            pattern = r'\b' + re.escape(lang_name) + r'\b'
        
        if re.search(pattern, description_lower):
            detected_language = lang_code
            break
    
    # Detect instrumental
    is_instrumental = False
    if 'instrumental' in description_lower:
        is_instrumental = True
    elif 'pure music' in description_lower or 'pure instrument' in description_lower:
        is_instrumental = True
    elif description_lower.endswith(' solo') or description_lower == 'solo':
        is_instrumental = True
    
    return detected_language, is_instrumental


JobStatus = Literal["queued", "running", "succeeded", "failed"]


class GenerateMusicRequest(BaseModel):
    prompt: str = Field(default="", description="Text prompt describing the music")
    lyrics: str = Field(default="", description="Lyric text")

    # New API semantics:
    # - thinking=True: use 5Hz LM to generate audio codes (lm-dit behavior)
    # - thinking=False: do not use LM to generate codes (dit behavior)
    # Regardless of thinking, if some metas are missing, server may use LM to fill them.
    thinking: bool = False
    # Sample-mode requests auto-generate caption/lyrics/metas via LM (no user prompt).
    sample_mode: bool = False
    # Description for sample mode: auto-generate caption/lyrics from description query
    sample_query: str = Field(default="", description="Query/description for sample mode (use create_sample)")
    # Whether to use format_sample() to enhance input caption/lyrics
    use_format: bool = Field(default=False, description="Use format_sample() to enhance input (default: False)")
    # Model name for multi-model support (select which DiT model to use)
    model: Optional[str] = Field(default=None, description="Model name to use (e.g., 'acestep-v15-turbo')")

    bpm: Optional[int] = None
    # Accept common client keys via manual parsing (see RequestParser).
    key_scale: str = ""
    time_signature: str = ""
    vocal_language: str = "en"
    inference_steps: int = 8
    guidance_scale: float = 7.0
    use_random_seed: bool = True
    seed: Union[int, str] = -1

    reference_audio_path: Optional[str] = None
    src_audio_path: Optional[str] = None
    audio_duration: Optional[float] = None
    batch_size: Optional[int] = None

    repainting_start: float = 0.0
    repainting_end: Optional[float] = None

    instruction: str = DEFAULT_DIT_INSTRUCTION
    audio_cover_strength: float = 1.0
    task_type: str = "text2music"
    analysis_only: bool = False
    full_analysis_only: bool = False

    use_adg: bool = False
    cfg_interval_start: float = 0.0
    cfg_interval_end: float = 1.0
    infer_method: str = "ode"  # "ode" or "sde" - diffusion inference method
    shift: float = Field(
        default=3.0,
        description="Timestep shift factor (range 1.0~5.0, default 3.0). Only effective for base models, not turbo models."
    )
    timesteps: Optional[str] = Field(
        default=None,
        description="Custom timesteps (comma-separated, e.g., '0.97,0.76,0.615,0.5,0.395,0.28,0.18,0.085,0'). Overrides inference_steps and shift."
    )

    audio_format: str = "mp3"
    use_tiled_decode: bool = True

    # 5Hz LM (server-side): used for metadata completion and (when thinking=True) codes generation.
    lm_model_path: Optional[str] = None  # e.g. "acestep-5Hz-lm-0.6B"
    lm_backend: Literal["vllm", "pt", "mlx"] = "vllm"

    constrained_decoding: bool = True
    constrained_decoding_debug: bool = False
    use_cot_caption: bool = True
    use_cot_language: bool = True
    is_format_caption: bool = False
    allow_lm_batch: bool = True

    lm_temperature: float = 0.85
    lm_cfg_scale: float = 2.5
    lm_top_k: Optional[int] = None
    lm_top_p: Optional[float] = 0.9
    lm_repetition_penalty: float = 1.0
    lm_negative_prompt: str = "NO USER INPUT"

    class Config:
        allow_population_by_field_name = True
        allow_population_by_alias = True


class CreateJobResponse(BaseModel):
    task_id: str
    status: JobStatus
    queue_position: int = 0  # 1-based best-effort position when queued
    progress_text: Optional[str] = ""
    

class JobResult(BaseModel):
    first_audio_path: Optional[str] = None
    second_audio_path: Optional[str] = None
    audio_paths: list[str] = Field(default_factory=list)

    generation_info: str = ""
    status_message: str = ""
    seed_value: str = ""

    metas: Dict[str, Any] = Field(default_factory=dict)
    bpm: Optional[int] = None
    duration: Optional[float] = None
    genres: Optional[str] = None
    keyscale: Optional[str] = None
    timesignature: Optional[str] = None
    
    # Model information
    lm_model: Optional[str] = None
    dit_model: Optional[str] = None


class JobResponse(BaseModel):
    job_id: str
    status: JobStatus
    created_at: float
    started_at: Optional[float] = None
    finished_at: Optional[float] = None

    # queue observability
    queue_position: int = 0
    eta_seconds: Optional[float] = None
    avg_job_seconds: Optional[float] = None

    result: Optional[JobResult] = None
    error: Optional[str] = None


@dataclass
class _JobRecord:
    job_id: str
    status: JobStatus
    created_at: float
    started_at: Optional[float] = None
    finished_at: Optional[float] = None
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    progress_text: str = ""
    status_text: str = ""
    env: str = "development"
    progress: float = 0.0  # 0.0 - 1.0
    stage: str = "queued"
    updated_at: Optional[float] = None
    # OpenRouter integration: synchronous wait / streaming support
    done_event: Optional[asyncio.Event] = None
    progress_queue: Optional[asyncio.Queue] = None


class _JobStore:
    def __init__(self, max_age_seconds: int = JOB_STORE_MAX_AGE_SECONDS) -> None:
        self._lock = Lock()
        self._jobs: Dict[str, _JobRecord] = {}
        self._max_age = max_age_seconds

    def create(self) -> _JobRecord:
        job_id = str(uuid4())
        now = time.time()
        rec = _JobRecord(job_id=job_id, status="queued", created_at=now, progress=0.0, stage="queued", updated_at=now)
        with self._lock:
            self._jobs[job_id] = rec
        return rec

    def create_with_id(self, job_id: str, env: str = "development") -> _JobRecord:
        """Create job record with specified ID"""
        now = time.time()
        rec = _JobRecord(
            job_id=job_id,
            status="queued",
            created_at=now,
            env=env,
            progress=0.0,
            stage="queued",
            updated_at=now,
        )
        with self._lock:
            self._jobs[job_id] = rec
        return rec

    def get(self, job_id: str) -> Optional[_JobRecord]:
        with self._lock:
            return self._jobs.get(job_id)

    def mark_running(self, job_id: str) -> None:
        with self._lock:
            rec = self._jobs[job_id]
            rec.status = "running"
            rec.started_at = time.time()
            rec.progress = max(rec.progress, 0.01)
            rec.stage = "running"
            rec.updated_at = time.time()

    def mark_succeeded(self, job_id: str, result: Dict[str, Any]) -> None:
        with self._lock:
            rec = self._jobs[job_id]
            rec.status = "succeeded"
            rec.finished_at = time.time()
            rec.result = result
            rec.error = None
            rec.progress = 1.0
            rec.stage = "succeeded"
            rec.updated_at = time.time()

    def mark_failed(self, job_id: str, error: str) -> None:
        with self._lock:
            rec = self._jobs[job_id]
            rec.status = "failed"
            rec.finished_at = time.time()
            rec.result = None
            rec.error = error
            rec.progress = rec.progress if rec.progress > 0 else 0.0
            rec.stage = "failed"
            rec.updated_at = time.time()

    def update_progress(self, job_id: str, progress: float, stage: Optional[str] = None) -> None:
        with self._lock:
            rec = self._jobs.get(job_id)
            if not rec:
                return
            rec.progress = max(0.0, min(1.0, float(progress)))
            if stage:
                rec.stage = stage
            rec.updated_at = time.time()

    def cleanup_old_jobs(self, max_age_seconds: Optional[int] = None) -> int:
        """
        Clean up completed jobs older than max_age_seconds.

        Only removes jobs with status 'succeeded' or 'failed'.
        Jobs that are 'queued' or 'running' are never removed.

        Returns the number of jobs removed.
        """
        max_age = max_age_seconds if max_age_seconds is not None else self._max_age
        now = time.time()
        removed = 0

        with self._lock:
            to_remove = []
            for job_id, rec in self._jobs.items():
                if rec.status in ("succeeded", "failed"):
                    finish_time = rec.finished_at or rec.created_at
                    age = now - finish_time
                    if age > max_age:
                        to_remove.append(job_id)

            for job_id in to_remove:
                del self._jobs[job_id]
                removed += 1

        return removed

    def get_stats(self) -> Dict[str, int]:
        """Get statistics about jobs in the store."""
        with self._lock:
            stats = {
                "total": len(self._jobs),
                "queued": 0,
                "running": 0,
                "succeeded": 0,
                "failed": 0,
            }
            for rec in self._jobs.values():
                if rec.status in stats:
                    stats[rec.status] += 1
            return stats
    
    def update_status_text(self, job_id: str, text: str) -> None:
        with self._lock:
            if job_id in self._jobs:
                self._jobs[job_id].status_text = text

    def update_progress_text(self, job_id: str, text: str) -> None:
        with self._lock:
            if job_id in self._jobs:
                self._jobs[job_id].progress_text = text

def _env_bool(name: str, default: bool) -> bool:
    v = os.getenv(name)
    if v is None:
        return default
    return v.strip().lower() in {"1", "true", "yes", "y", "on"}




def _get_model_name(config_path: str) -> str:
    """
    Extract model name from config_path.
    
    Args:
        config_path: Path like "acestep-v15-turbo" or "/path/to/acestep-v15-turbo"
        
    Returns:
        Model name (last directory name from config_path)
    """
    if not config_path:
        return ""
    normalized = config_path.rstrip("/\\")
    return os.path.basename(normalized)


_project_env_loaded = False


def _load_project_env() -> None:
    """Load .env at most once per process to avoid epoch-boundary stalls (e.g. Windows LoRA training)."""
    global _project_env_loaded
    if _project_env_loaded or load_dotenv is None:
        return
    try:
        project_root = _get_project_root()
        env_path = os.path.join(project_root, ".env")
        if os.path.exists(env_path):
            load_dotenv(env_path, override=False)
        _project_env_loaded = True
    except Exception:
        # Optional best-effort: continue even if .env loading fails.
        pass


_load_project_env()


def _to_int(v: Any, default: Optional[int] = None) -> Optional[int]:
    if v is None:
        return default
    if isinstance(v, int):
        return v
    s = str(v).strip()
    if s == "":
        return default
    try:
        return int(s)
    except Exception:
        return default


def _to_float(v: Any, default: Optional[float] = None) -> Optional[float]:
    if v is None:
        return default
    if isinstance(v, float):
        return v
    s = str(v).strip()
    if s == "":
        return default
    try:
        return float(s)
    except Exception:
        return default


def _to_bool(v: Any, default: bool = False) -> bool:
    if v is None:
        return default
    if isinstance(v, bool):
        return v
    s = str(v).strip().lower()
    if s == "":
        return default
    return s in {"1", "true", "yes", "y", "on"}


def _map_status(status: str) -> int:
    """Map job status string to integer code."""
    return STATUS_MAP.get(status, 2)


def _parse_timesteps(s: Optional[str]) -> Optional[List[float]]:
    """Parse comma-separated timesteps string to list of floats."""
    if not s or not s.strip():
        return None
    try:
        return [float(t.strip()) for t in s.split(",") if t.strip()]
    except (ValueError, Exception):
        return None


def _is_instrumental(lyrics: str) -> bool:
    """
    Determine if the music should be instrumental based on lyrics.

    Returns True if:
    - lyrics is empty or whitespace only
    - lyrics (lowercased and trimmed) is "[inst]" or "[instrumental]"
    """
    if not lyrics:
        return True
    lyrics_clean = lyrics.strip().lower()
    if not lyrics_clean:
        return True
    return lyrics_clean in ("[inst]", "[instrumental]")


class RequestParser:
    """Parse request parameters from multiple sources with alias support."""

    def __init__(self, raw: dict):
        self._raw = dict(raw) if raw else {}
        self._param_obj = self._parse_json(self._raw.get("param_obj"))
        self._metas = self._find_metas()

    def _parse_json(self, v) -> dict:
        if isinstance(v, dict):
            return v
        if isinstance(v, str) and v.strip():
            try:
                return json.loads(v)
            except Exception:
                pass
        return {}

    def _find_metas(self) -> dict:
        for key in ("metas", "meta", "metadata", "user_metadata", "userMetadata"):
            v = self._raw.get(key)
            if v:
                return self._parse_json(v)
        return {}

    def get(self, name: str, default=None):
        """Get parameter by canonical name from all sources."""
        aliases = PARAM_ALIASES.get(name, [name])
        for source in (self._raw, self._param_obj, self._metas):
            for alias in aliases:
                v = source.get(alias)
                if v is not None:
                    return v
        return default

    def str(self, name: str, default: str = "") -> str:
        v = self.get(name)
        return str(v) if v is not None else default

    def int(self, name: str, default: Optional[int] = None) -> Optional[int]:
        return _to_int(self.get(name), default)

    def float(self, name: str, default: Optional[float] = None) -> Optional[float]:
        return _to_float(self.get(name), default)

    def bool(self, name: str, default: bool = False) -> bool:
        return _to_bool(self.get(name), default)


async def _save_upload_to_temp(upload: StarletteUploadFile, *, prefix: str) -> str:
    suffix = Path(upload.filename or "").suffix
    fd, path = tempfile.mkstemp(prefix=f"{prefix}_", suffix=suffix)
    os.close(fd)
    try:
        with open(path, "wb") as f:
            while True:
                chunk = await upload.read(1024 * 1024)
                if not chunk:
                    break
                f.write(chunk)
    except Exception:
        try:
            os.remove(path)
        except Exception:
            pass
        raise
    finally:
        try:
            await upload.close()
        except Exception:
            pass
    return path

class LogBuffer:
    def __init__(self):
        self.last_message = "Waiting"

    def write(self, message):
        msg = message.strip()
        if msg:
            self.last_message = msg

    def flush(self):
        pass

log_buffer = LogBuffer()
logger.add(lambda msg: log_buffer.write(msg), format="{time:HH:mm:ss} | {level} | {message}")

class StderrLogger:
    def __init__(self, original_stderr, buffer):
        self.original_stderr = original_stderr
        self.buffer = buffer

    def write(self, message):
        self.original_stderr.write(message) # Print to terminal
        self.buffer.write(message)          # Send to API buffer

    def flush(self):
        self.original_stderr.flush()

sys.stderr = StderrLogger(sys.stderr, log_buffer)


def create_app() -> FastAPI:
    store = _JobStore()

    # API Key authentication (from environment variable)
    api_key = os.getenv("ACESTEP_API_KEY", None)
    set_api_key(api_key)

    QUEUE_MAXSIZE = int(os.getenv("ACESTEP_QUEUE_MAXSIZE", "200"))
    WORKER_COUNT = int(os.getenv("ACESTEP_QUEUE_WORKERS", "1"))  # Single GPU recommended

    INITIAL_AVG_JOB_SECONDS = float(os.getenv("ACESTEP_AVG_JOB_SECONDS", "5.0"))
    AVG_WINDOW = int(os.getenv("ACESTEP_AVG_WINDOW", "50"))

    def _path_to_audio_url(path: str) -> str:
        """Convert local file path to downloadable relative URL"""
        if not path:
            return path
        if path.startswith("http://") or path.startswith("https://"):
            return path
        encoded_path = urllib.parse.quote(path, safe="")
        return f"/v1/audio?path={encoded_path}"

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        # Clear proxy env that may affect downstream libs
        for proxy_var in ["http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"]:
            os.environ.pop(proxy_var, None)

        # Ensure compilation/temp caches do not fill up small default /tmp.
        # Triton/Inductor (and the system compiler) can create large temporary files.
        project_root = _get_project_root()
        cache_root = os.path.join(project_root, ".cache", "acestep")
        tmp_root = (os.getenv("ACESTEP_TMPDIR") or os.path.join(cache_root, "tmp")).strip()
        triton_cache_root = (os.getenv("TRITON_CACHE_DIR") or os.path.join(cache_root, "triton")).strip()
        inductor_cache_root = (os.getenv("TORCHINDUCTOR_CACHE_DIR") or os.path.join(cache_root, "torchinductor")).strip()

        for p in [cache_root, tmp_root, triton_cache_root, inductor_cache_root]:
            try:
                os.makedirs(p, exist_ok=True)
            except Exception:
                # Best-effort: do not block startup if directory creation fails.
                pass

        # Respect explicit user overrides; if ACESTEP_TMPDIR is set, it should win.
        if os.getenv("ACESTEP_TMPDIR"):
            os.environ["TMPDIR"] = tmp_root
            os.environ["TEMP"] = tmp_root
            os.environ["TMP"] = tmp_root
        else:
            os.environ.setdefault("TMPDIR", tmp_root)
            os.environ.setdefault("TEMP", tmp_root)
            os.environ.setdefault("TMP", tmp_root)

        os.environ.setdefault("TRITON_CACHE_DIR", triton_cache_root)
        os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", inductor_cache_root)

        handler = AceStepHandler()
        llm_handler = LLMHandler()
        init_lock = asyncio.Lock()
        app.state._initialized = False
        app.state._init_error = None
        app.state._init_lock = init_lock

        app.state.llm_handler = llm_handler
        app.state._llm_initialized = False
        app.state._llm_init_error = None
        app.state._llm_init_lock = Lock()
        app.state._llm_lazy_load_disabled = False  # Will be set to True if LLM skipped due to GPU config

        # Multi-model support: secondary DiT handlers
        handler2 = None
        handler3 = None
        config_path2 = os.getenv("ACESTEP_CONFIG_PATH2", "").strip()
        config_path3 = os.getenv("ACESTEP_CONFIG_PATH3", "").strip()
        
        if config_path2:
            handler2 = AceStepHandler()
        if config_path3:
            handler3 = AceStepHandler()
        
        app.state.handler2 = handler2
        app.state.handler3 = handler3
        app.state._initialized2 = False
        app.state._initialized3 = False
        app.state._config_path = os.getenv("ACESTEP_CONFIG_PATH", "acestep-v15-turbo")
        app.state._config_path2 = config_path2
        app.state._config_path3 = config_path3

        max_workers = int(os.getenv("ACESTEP_API_WORKERS", "1"))
        executor = ThreadPoolExecutor(max_workers=max_workers)

        # Queue & observability
        app.state.job_queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE)  # (job_id, req)
        app.state.pending_ids = deque()  # queued job_ids
        app.state.pending_lock = asyncio.Lock()

        # temp files per job (from multipart uploads)
        app.state.job_temp_files = {}  # job_id -> list[path]
        app.state.job_temp_files_lock = asyncio.Lock()

        # stats
        app.state.stats_lock = asyncio.Lock()
        app.state.recent_durations = deque(maxlen=AVG_WINDOW)
        app.state.avg_job_seconds = INITIAL_AVG_JOB_SECONDS

        app.state.handler = handler
        app.state.executor = executor
        app.state.job_store = store
        app.state._python_executable = sys.executable
        
        # Temporary directory for saving generated audio files
        app.state.temp_audio_dir = os.path.join(tmp_root, "api_audio")
        os.makedirs(app.state.temp_audio_dir, exist_ok=True)

        # Initialize local cache
        try:
            from acestep.local_cache import get_local_cache
            local_cache_dir = os.path.join(cache_root, "local_redis")
            app.state.local_cache = get_local_cache(local_cache_dir)
        except ImportError:
            app.state.local_cache = None

        async def _ensure_initialized() -> None:
            """Check if models are initialized (they should be loaded at startup)."""
            if getattr(app.state, "_init_error", None):
                raise RuntimeError(app.state._init_error)
            if not getattr(app.state, "_initialized", False):
                raise RuntimeError("Model not initialized")

        async def _cleanup_job_temp_files(job_id: str) -> None:
            async with app.state.job_temp_files_lock:
                paths = app.state.job_temp_files.pop(job_id, [])
            for p in paths:
                try:
                    os.remove(p)
                except Exception:
                    pass

        def _update_local_cache(job_id: str, result: Optional[Dict], status: str) -> None:
            """Update local cache with job result"""
            local_cache = getattr(app.state, 'local_cache', None)
            if not local_cache:
                return

            rec = store.get(job_id)
            env = getattr(rec, 'env', 'development') if rec else 'development'
            create_time = rec.created_at if rec else time.time()

            status_int = _map_status(status)

            if status == "succeeded" and result:
                # Check if it's a "Full Analysis" result
                if result.get("status_message") == "Full Hardware Analysis Success":
                    result_data = [result]
                else:
                    audio_paths = result.get("audio_paths", [])
                    # Final prompt/lyrics (may be modified by thinking/format)
                    final_prompt = result.get("prompt", "")
                    final_lyrics = result.get("lyrics", "")
                    # Original user input from metas
                    metas_raw = result.get("metas", {}) or {}
                    original_prompt = metas_raw.get("prompt", "")
                    original_lyrics = metas_raw.get("lyrics", "")
                    # metas contains original input + other metadata
                    metas = {
                        "bpm": metas_raw.get("bpm"),
                        "duration": metas_raw.get("duration"),
                        "genres": metas_raw.get("genres", ""),
                        "keyscale": metas_raw.get("keyscale", ""),
                        "timesignature": metas_raw.get("timesignature", ""),
                        "prompt": original_prompt,
                        "lyrics": original_lyrics,
                    }
                    # Extra fields for Discord bot
                    generation_info = result.get("generation_info", "")
                    seed_value = result.get("seed_value", "")
                    lm_model = result.get("lm_model", "")
                    dit_model = result.get("dit_model", "")

                    if audio_paths:
                        result_data = [
                            {
                                "file": p,
                                "wave": "",
                                "status": status_int,
                                "create_time": int(create_time),
                                "env": env,
                                "prompt": final_prompt,
                                "lyrics": final_lyrics,
                                "metas": metas,
                                "generation_info": generation_info,
                                "seed_value": seed_value,
                                "lm_model": lm_model,
                                "dit_model": dit_model,
                                "progress": 1.0,
                                "stage": "succeeded",
                            }
                            for p in audio_paths
                        ]
                    else:
                        result_data = [{
                            "file": "",
                            "wave": "",
                            "status": status_int,
                            "create_time": int(create_time),
                            "env": env,
                            "prompt": final_prompt,
                            "lyrics": final_lyrics,
                            "metas": metas,
                            "generation_info": generation_info,
                            "seed_value": seed_value,
                            "lm_model": lm_model,
                            "dit_model": dit_model,
                            "progress": 1.0,
                            "stage": "succeeded",
                        }]
            else:
                result_data = [{
                    "file": "",
                    "wave": "",
                    "status": status_int,
                    "create_time": int(create_time),
                    "env": env,
                    "progress": 0.0,
                    "stage": "failed" if status == "failed" else status,
                }]

            result_key = f"{RESULT_KEY_PREFIX}{job_id}"
            local_cache.set(result_key, result_data, ex=RESULT_EXPIRE_SECONDS)

        def _update_local_cache_progress(job_id: str, progress: float, stage: str) -> None:
            """Update local cache with job progress for queued/running states."""
            local_cache = getattr(app.state, 'local_cache', None)
            if not local_cache:
                return

            rec = store.get(job_id)
            env = getattr(rec, 'env', 'development') if rec else 'development'
            create_time = rec.created_at if rec else time.time()
            status_int = _map_status("running")

            result_data = [{
                "file": "",
                "wave": "",
                "status": status_int,
                "create_time": int(create_time),
                "env": env,
                "progress": float(progress),
                "stage": stage,
            }]

            result_key = f"{RESULT_KEY_PREFIX}{job_id}"
            local_cache.set(result_key, result_data, ex=RESULT_EXPIRE_SECONDS)

        async def _run_one_job(job_id: str, req: GenerateMusicRequest) -> None:
            job_store: _JobStore = app.state.job_store
            llm: LLMHandler = app.state.llm_handler
            executor: ThreadPoolExecutor = app.state.executor

            await _ensure_initialized()
            job_store.mark_running(job_id)
            _update_local_cache_progress(job_id, 0.01, "running")
            
            # Select DiT handler based on user's model choice
            # Default: use primary handler
            selected_handler: AceStepHandler = app.state.handler
            selected_model_name = _get_model_name(app.state._config_path)
            
            if req.model:
                model_matched = False
                
                # Check if it matches the second model
                if app.state.handler2 and getattr(app.state, "_initialized2", False):
                    model2_name = _get_model_name(app.state._config_path2)
                    if req.model == model2_name:
                        selected_handler = app.state.handler2
                        selected_model_name = model2_name
                        model_matched = True
                        print(f"[API Server] Job {job_id}: Using second model: {model2_name}")
                
                # Check if it matches the third model
                if not model_matched and app.state.handler3 and getattr(app.state, "_initialized3", False):
                    model3_name = _get_model_name(app.state._config_path3)
                    if req.model == model3_name:
                        selected_handler = app.state.handler3
                        selected_model_name = model3_name
                        model_matched = True
                        print(f"[API Server] Job {job_id}: Using third model: {model3_name}")
                
                if not model_matched:
                    available_models = [_get_model_name(app.state._config_path)]
                    if app.state.handler2 and getattr(app.state, "_initialized2", False):
                        available_models.append(_get_model_name(app.state._config_path2))
                    if app.state.handler3 and getattr(app.state, "_initialized3", False):
                        available_models.append(_get_model_name(app.state._config_path3))
                    print(f"[API Server] Job {job_id}: Model '{req.model}' not found in {available_models}, using primary: {selected_model_name}")
            
            # Use selected handler for generation
            h: AceStepHandler = selected_handler

            def _blocking_generate() -> Dict[str, Any]:
                """Generate music using unified inference logic from acestep.inference"""

                def _ensure_llm_ready() -> None:
                    """Ensure LLM handler is initialized when needed"""
                    with app.state._llm_init_lock:
                        initialized = getattr(app.state, "_llm_initialized", False)
                        had_error = getattr(app.state, "_llm_init_error", None)
                        if initialized or had_error is not None:
                            return
                        print("[API Server] reloading.")

                        # Check if lazy loading is disabled (GPU memory insufficient)
                        if getattr(app.state, "_llm_lazy_load_disabled", False):
                            app.state._llm_init_error = (
                                "LLM not initialized at startup. To enable LLM, set ACESTEP_INIT_LLM=true "
                                "in .env or environment variables. For this request, optional LLM features "
                                "(use_cot_caption, use_cot_language) will be auto-disabled."
                            )
                            print(f"[API Server] LLM lazy load blocked: LLM was not initialized at startup")
                            return

                        project_root = _get_project_root()
                        checkpoint_dir = os.path.join(project_root, "checkpoints")
                        lm_model_path = (req.lm_model_path or os.getenv("ACESTEP_LM_MODEL_PATH") or "acestep-5Hz-lm-0.6B").strip()
                        backend = (req.lm_backend or os.getenv("ACESTEP_LM_BACKEND") or "vllm").strip().lower()
                        if backend not in {"vllm", "pt", "mlx"}:
                            backend = "vllm"

                        # Auto-download LM model if not present
                        lm_model_name = _get_model_name(lm_model_path)
                        if lm_model_name:
                            try:
                                _ensure_model_downloaded(lm_model_name, checkpoint_dir)
                            except Exception as e:
                                print(f"[API Server] Warning: Failed to download LM model {lm_model_name}: {e}")

                        lm_device = os.getenv("ACESTEP_LM_DEVICE", os.getenv("ACESTEP_DEVICE", "auto"))
                        lm_offload = _env_bool("ACESTEP_LM_OFFLOAD_TO_CPU", False)

                        status, ok = llm.initialize(
                            checkpoint_dir=checkpoint_dir,
                            lm_model_path=lm_model_path,
                            backend=backend,
                            device=lm_device,
                            offload_to_cpu=lm_offload,
                            dtype=None,
                        )
                        if not ok:
                            app.state._llm_init_error = status
                        else:
                            app.state._llm_initialized = True

                def _normalize_metas(meta: Dict[str, Any]) -> Dict[str, Any]:
                    """Ensure a stable `metas` dict (keys always present)."""
                    meta = meta or {}
                    out: Dict[str, Any] = dict(meta)

                    # Normalize key aliases
                    if "keyscale" not in out and "key_scale" in out:
                        out["keyscale"] = out.get("key_scale")
                    if "timesignature" not in out and "time_signature" in out:
                        out["timesignature"] = out.get("time_signature")

                    # Ensure required keys exist
                    for k in ["bpm", "duration", "genres", "keyscale", "timesignature"]:
                        if out.get(k) in (None, ""):
                            out[k] = "N/A"
                    return out

                # Normalize LM sampling parameters
                lm_top_k = req.lm_top_k if req.lm_top_k and req.lm_top_k > 0 else 0
                lm_top_p = req.lm_top_p if req.lm_top_p and req.lm_top_p < 1.0 else 0.9

                # Determine if LLM is needed
                thinking = bool(req.thinking)
                sample_mode = bool(req.sample_mode)
                has_sample_query = bool(req.sample_query and req.sample_query.strip())
                use_format = bool(req.use_format)
                use_cot_caption = bool(req.use_cot_caption)
                use_cot_language = bool(req.use_cot_language)

                full_analysis_only = bool(req.full_analysis_only)

                # Unload LM for cover tasks on MPS to reduce memory; reload lazily when needed.
                if req.task_type == "cover" and h.device == "mps":
                    if getattr(app.state, "_llm_initialized", False) and getattr(llm, "llm_initialized", False):
                        try:
                            print("[API Server] unloading.")
                            llm.unload()
                            app.state._llm_initialized = False
                            app.state._llm_init_error = None
                        except Exception as e:
                            print(f"[API Server] Failed to unload LM: {e}")

                # LLM is REQUIRED for these features (fail if unavailable):
                # - thinking mode (LM generates audio codes)
                # - sample_mode (LM generates random caption/lyrics/metas)
                # - sample_query/description (LM generates from description)
                # - use_format (LM enhances caption/lyrics)
                # - full_analysis_only (LM understands audio codes)
                require_llm = thinking or sample_mode or has_sample_query or use_format or full_analysis_only

                # LLM is OPTIONAL for these features (auto-disable if unavailable):
                # - use_cot_caption or use_cot_language (LM enhances metadata)
                want_llm = use_cot_caption or use_cot_language

                # Check if LLM is available
                llm_available = True
                if require_llm or want_llm:
                    _ensure_llm_ready()
                    if getattr(app.state, "_llm_init_error", None):
                        llm_available = False

                # Fail if LLM is required but unavailable
                if require_llm and not llm_available:
                    raise RuntimeError(f"5Hz LM init failed: {app.state._llm_init_error}")

                # Auto-disable optional LLM features if unavailable
                if want_llm and not llm_available:
                    if use_cot_caption or use_cot_language:
                        print(f"[API Server] LLM unavailable, auto-disabling: use_cot_caption={use_cot_caption}->False, use_cot_language={use_cot_language}->False")
                    use_cot_caption = False
                    use_cot_language = False

                # Handle sample mode or description: generate caption/lyrics/metas via LM
                caption = req.prompt
                lyrics = req.lyrics
                bpm = req.bpm
                key_scale = req.key_scale
                time_signature = req.time_signature
                audio_duration = req.audio_duration

                # Save original user input for metas
                original_prompt = req.prompt or ""
                original_lyrics = req.lyrics or ""
                
                if sample_mode or has_sample_query:
                    # Parse description hints from sample_query (if provided)
                    sample_query = req.sample_query if has_sample_query else "NO USER INPUT"
                    parsed_language, parsed_instrumental = _parse_description_hints(sample_query)

                    # Determine vocal_language with priority:
                    # 1. User-specified vocal_language (if not default "en")
                    # 2. Language parsed from description
                    # 3. None (no constraint)
                    if req.vocal_language and req.vocal_language not in ("en", "unknown", ""):
                        sample_language = req.vocal_language
                    else:
                        sample_language = parsed_language

                    sample_result = create_sample(
                        llm_handler=llm,
                        query=sample_query,
                        instrumental=parsed_instrumental,
                        vocal_language=sample_language,
                        temperature=req.lm_temperature,
                        top_k=lm_top_k if lm_top_k > 0 else None,
                        top_p=lm_top_p if lm_top_p < 1.0 else None,
                        use_constrained_decoding=True,
                    )

                    if not sample_result.success:
                        raise RuntimeError(f"create_sample failed: {sample_result.error or sample_result.status_message}")

                    # Use generated sample data
                    caption = sample_result.caption
                    lyrics = sample_result.lyrics
                    bpm = sample_result.bpm
                    key_scale = sample_result.keyscale
                    time_signature = sample_result.timesignature
                    audio_duration = sample_result.duration

                # Apply format_sample() if use_format is True and caption/lyrics are provided
                format_has_duration = False

                if req.use_format and (caption or lyrics):
                    _ensure_llm_ready()
                    if getattr(app.state, "_llm_init_error", None):
                        raise RuntimeError(f"5Hz LM init failed (needed for format): {app.state._llm_init_error}")
                    
                    # Build user_metadata from request params (matching bot.py behavior)
                    user_metadata_for_format = {}
                    if bpm is not None:
                        user_metadata_for_format['bpm'] = bpm
                    if audio_duration is not None and float(audio_duration) > 0:
                        user_metadata_for_format['duration'] = float(audio_duration)
                    if key_scale:
                        user_metadata_for_format['keyscale'] = key_scale
                    if time_signature:
                        user_metadata_for_format['timesignature'] = time_signature
                    if req.vocal_language and req.vocal_language != "unknown":
                        user_metadata_for_format['language'] = req.vocal_language
                    
                    format_result = format_sample(
                        llm_handler=llm,
                        caption=caption,
                        lyrics=lyrics,
                        user_metadata=user_metadata_for_format if user_metadata_for_format else None,
                        temperature=req.lm_temperature,
                        top_k=lm_top_k if lm_top_k > 0 else None,
                        top_p=lm_top_p if lm_top_p < 1.0 else None,
                        use_constrained_decoding=True,
                    )
                    
                    if format_result.success:
                        # Extract all formatted data (matching bot.py behavior)
                        caption = format_result.caption or caption
                        lyrics = format_result.lyrics or lyrics
                        if format_result.duration:
                            audio_duration = format_result.duration
                            format_has_duration = True
                        if format_result.bpm:
                            bpm = format_result.bpm
                        if format_result.keyscale:
                            key_scale = format_result.keyscale
                        if format_result.timesignature:
                            time_signature = format_result.timesignature

                # Parse timesteps string to list of floats if provided
                parsed_timesteps = _parse_timesteps(req.timesteps)

                # Determine actual inference steps (timesteps override inference_steps)
                actual_inference_steps = len(parsed_timesteps) if parsed_timesteps else req.inference_steps

                # Auto-select instruction based on task_type if user didn't provide custom instruction
                # This matches gradio behavior which uses TASK_INSTRUCTIONS for each task type
                instruction_to_use = req.instruction
                if instruction_to_use == DEFAULT_DIT_INSTRUCTION and req.task_type in TASK_INSTRUCTIONS:
                    instruction_to_use = TASK_INSTRUCTIONS[req.task_type]

                # Build GenerationParams using unified interface
                # Note: thinking controls LM code generation, sample_mode only affects CoT metas
                params = GenerationParams(
                    task_type=req.task_type,
                    instruction=instruction_to_use,
                    reference_audio=req.reference_audio_path,
                    src_audio=req.src_audio_path,
                    audio_codes="",
                    caption=caption,
                    lyrics=lyrics,
                    instrumental=_is_instrumental(lyrics),
                    vocal_language=req.vocal_language,
                    bpm=bpm,
                    keyscale=key_scale,
                    timesignature=time_signature,
                    duration=audio_duration if audio_duration else -1.0,
                    inference_steps=req.inference_steps,
                    seed=req.seed,
                    guidance_scale=req.guidance_scale,
                    use_adg=req.use_adg,
                    cfg_interval_start=req.cfg_interval_start,
                    cfg_interval_end=req.cfg_interval_end,
                    shift=req.shift,
                    infer_method=req.infer_method,
                    timesteps=parsed_timesteps,
                    repainting_start=req.repainting_start,
                    repainting_end=req.repainting_end if req.repainting_end else -1,
                    audio_cover_strength=req.audio_cover_strength,
                    # LM parameters
                    thinking=thinking,  # Use LM for code generation when thinking=True
                    lm_temperature=req.lm_temperature,
                    lm_cfg_scale=req.lm_cfg_scale,
                    lm_top_k=lm_top_k,
                    lm_top_p=lm_top_p,
                    lm_negative_prompt=req.lm_negative_prompt,
                    # use_cot_metas logic:
                    # - sample_mode: metas already generated, skip Phase 1
                    # - format with duration: metas already generated, skip Phase 1
                    # - format without duration: need Phase 1 to generate duration
                    # - no format: need Phase 1 to generate all metas
                    use_cot_metas=not sample_mode and not format_has_duration,
                    use_cot_caption=use_cot_caption,  # Use local var (may be auto-disabled)
                    use_cot_language=use_cot_language,  # Use local var (may be auto-disabled)
                    use_constrained_decoding=True,
                )

                # Build GenerationConfig - default to 2 audios like gradio_ui
                batch_size = req.batch_size if req.batch_size is not None else 2
                config = GenerationConfig(
                    batch_size=batch_size,
                    allow_lm_batch=req.allow_lm_batch,
                    use_random_seed=req.use_random_seed,
                    seeds=None,  # Let unified logic handle seed generation
                    audio_format=req.audio_format,
                    constrained_decoding_debug=req.constrained_decoding_debug,
                )

                # Check LLM initialization status
                llm_is_initialized = getattr(app.state, "_llm_initialized", False)
                llm_to_pass = llm if llm_is_initialized else None

                # Progress callback for API polling
                last_progress = {"value": -1.0, "time": 0.0, "stage": ""}

                def _progress_cb(value: float, desc: str = "") -> None:
                    now = time.time()
                    try:
                        value_f = max(0.0, min(1.0, float(value)))
                    except Exception:
                        value_f = 0.0
                    stage = desc or last_progress["stage"] or "running"
                    # Throttle updates to avoid excessive cache writes
                    if (
                        value_f - last_progress["value"] >= 0.01
                        or stage != last_progress["stage"]
                        or (now - last_progress["time"]) >= 0.5
                    ):
                        last_progress["value"] = value_f
                        last_progress["time"] = now
                        last_progress["stage"] = stage
                        job_store.update_progress(job_id, value_f, stage=stage)
                        _update_local_cache_progress(job_id, value_f, stage)

                if req.full_analysis_only:
                    store.update_progress_text(job_id, "Starting Deep Analysis...")
                    # Step A: Convert source audio to semantic codes
                    # We use params.src_audio which is the server-side path
                    audio_codes = h.convert_src_audio_to_codes(params.src_audio)

                    if not audio_codes or audio_codes.startswith("❌"):
                        raise RuntimeError(f"Audio encoding failed: {audio_codes}")

                    # Step B: LLM Understanding of those specific codes
                    # This yields the deep metadata and lyrics transcription
                    metadata_dict, status_string = llm_to_pass.understand_audio_from_codes(
                        audio_codes=audio_codes,
                        temperature=0.3,
                        use_constrained_decoding=True,
                        constrained_decoding_debug=config.constrained_decoding_debug
                    )

                    if not metadata_dict:
                        raise RuntimeError(f"LLM Understanding failed: {status_string}")

                    return {
                        "status_message": "Full Hardware Analysis Success",
                        "bpm": metadata_dict.get("bpm"),
                        "keyscale": metadata_dict.get("keyscale"),
                        "timesignature": metadata_dict.get("timesignature"),
                        "duration": metadata_dict.get("duration"),
                        "genre": metadata_dict.get("genres") or metadata_dict.get("genre"),
                        "prompt": metadata_dict.get("caption", ""),
                        "lyrics": metadata_dict.get("lyrics", ""),
                        "language": metadata_dict.get("language", "unknown"),
                        "metas": metadata_dict,
                        "audio_paths": []
                    }

                if req.analysis_only:
                    lm_res = llm_to_pass.generate_with_stop_condition(
                        caption=params.caption,
                        lyrics=params.lyrics,
                        infer_type="dit",
                        temperature=req.lm_temperature,
                        top_p=req.lm_top_p,
                        use_cot_metas=True,
                        use_cot_caption=req.use_cot_caption,
                        use_cot_language=req.use_cot_language,
                        use_constrained_decoding=True
                    )

                    if not lm_res.get("success"):
                        raise RuntimeError(f"Analysis Failed: {lm_res.get('error')}")

                    metas_found = lm_res.get("metadata", {})
                    return {
                        "first_audio_path": None,
                        "audio_paths": [],
                        "raw_audio_paths": [],
                        "generation_info": "Analysis Only Mode Complete",
                        "status_message": "Success",
                        "metas": metas_found,
                        "bpm": metas_found.get("bpm"),
                        "keyscale": metas_found.get("keyscale"),
                        "duration": metas_found.get("duration"),
                        "prompt": metas_found.get("caption", params.caption),
                        "lyrics": params.lyrics,
                        "lm_model": os.getenv("ACESTEP_LM_MODEL_PATH", ""),
                        "dit_model": "None (Analysis Only)"
                    }

                # Generate music using unified interface
                sequential_runs = 1
                if req.task_type == "cover" and h.device == "mps":
                    # If user asked for multiple outputs, run sequentially on MPS to avoid OOM.
                    if config.batch_size is not None and config.batch_size > 1:
                        sequential_runs = int(config.batch_size)
                        config.batch_size = 1
                        print(f"[API Server] Job {job_id}: MPS cover sequential mode enabled (runs={sequential_runs})")

                def _progress_for_slice(start: float, end: float):
                    base = {"seen": False, "value": 0.0}
                    def _cb(value: float, desc: str = "") -> None:
                        try:
                            value_f = max(0.0, min(1.0, float(value)))
                        except Exception:
                            value_f = 0.0
                        if not base["seen"]:
                            base["seen"] = True
                            base["value"] = value_f
                        # Normalize progress to avoid initial jump (e.g., 0.51 -> 0.0)
                        if value_f <= base["value"]:
                            norm = 0.0
                        else:
                            denom = max(1e-6, 1.0 - base["value"])
                            norm = min(1.0, (value_f - base["value"]) / denom)
                        mapped = start + (end - start) * norm
                        _progress_cb(mapped, desc=desc)
                    return _cb

                aggregated_result = None
                all_audios: List[Dict[str, Any]] = []
                for run_idx in range(sequential_runs):
                    if sequential_runs > 1:
                        print(f"[API Server] Job {job_id}: Sequential cover run {run_idx + 1}/{sequential_runs}")
                    if sequential_runs > 1:
                        start = run_idx / sequential_runs
                        end = (run_idx + 1) / sequential_runs
                        progress_cb = _progress_for_slice(start, end)
                    else:
                        progress_cb = _progress_cb

                    result = generate_music(
                        dit_handler=h,
                        llm_handler=llm_to_pass,
                        params=params,
                        config=config,
                        save_dir=app.state.temp_audio_dir,
                        progress=progress_cb,
                    )
                    if not result.success:
                        raise RuntimeError(f"Music generation failed: {result.error or result.status_message}")

                    if aggregated_result is None:
                        aggregated_result = result
                    all_audios.extend(result.audios)

                # Use aggregated result with combined audios
                if aggregated_result is None:
                    raise RuntimeError("Music generation failed: no results")
                aggregated_result.audios = all_audios
                result = aggregated_result

                if not result.success:
                    raise RuntimeError(f"Music generation failed: {result.error or result.status_message}")

                # Extract results
                audio_paths = [audio["path"] for audio in result.audios if audio.get("path")]
                first_audio = audio_paths[0] if len(audio_paths) > 0 else None
                second_audio = audio_paths[1] if len(audio_paths) > 1 else None

                # Get metadata from LM or CoT results
                lm_metadata = result.extra_outputs.get("lm_metadata", {})
                metas_out = _normalize_metas(lm_metadata)
                
                # Update metas with actual values used
                if params.cot_bpm:
                    metas_out["bpm"] = params.cot_bpm
                elif bpm:
                    metas_out["bpm"] = bpm
                    
                if params.cot_duration:
                    metas_out["duration"] = params.cot_duration
                elif audio_duration:
                    metas_out["duration"] = audio_duration
                    
                if params.cot_keyscale:
                    metas_out["keyscale"] = params.cot_keyscale
                elif key_scale:
                    metas_out["keyscale"] = key_scale
                    
                if params.cot_timesignature:
                    metas_out["timesignature"] = params.cot_timesignature
                elif time_signature:
                    metas_out["timesignature"] = time_signature

                # Store original user input in metas (not the final/modified values)
                metas_out["prompt"] = original_prompt
                metas_out["lyrics"] = original_lyrics

                # Extract seed values for response (comma-separated for multiple audios)
                seed_values = []
                for audio in result.audios:
                    audio_params = audio.get("params", {})
                    seed = audio_params.get("seed")
                    if seed is not None:
                        seed_values.append(str(seed))
                seed_value = ",".join(seed_values) if seed_values else ""

                # Build generation_info using the helper function (like gradio_ui)
                time_costs = result.extra_outputs.get("time_costs", {})
                generation_info = _build_generation_info(
                    lm_metadata=lm_metadata,
                    time_costs=time_costs,
                    seed_value=seed_value,
                    inference_steps=req.inference_steps,
                    num_audios=len(result.audios),
                )

                def _none_if_na_str(v: Any) -> Optional[str]:
                    if v is None:
                        return None
                    s = str(v).strip()
                    if s in {"", "N/A"}:
                        return None
                    return s

                # Get model information
                lm_model_name = os.getenv("ACESTEP_LM_MODEL_PATH", "acestep-5Hz-lm-0.6B")
                # Use selected_model_name (set at the beginning of _run_one_job)
                dit_model_name = selected_model_name
                
                return {
                    "first_audio_path": _path_to_audio_url(first_audio) if first_audio else None,
                    "second_audio_path": _path_to_audio_url(second_audio) if second_audio else None,
                    "audio_paths": [_path_to_audio_url(p) for p in audio_paths],
                    "raw_audio_paths": list(audio_paths),
                    "generation_info": generation_info,
                    "status_message": result.status_message,
                    "seed_value": seed_value,
                    # Final prompt/lyrics (may be modified by thinking/format)
                    "prompt": caption or "",
                    "lyrics": lyrics or "",
                    # metas contains original user input + other metadata
                    "metas": metas_out,
                    "bpm": metas_out.get("bpm") if isinstance(metas_out.get("bpm"), int) else None,
                    "duration": metas_out.get("duration") if isinstance(metas_out.get("duration"), (int, float)) else None,
                    "genres": _none_if_na_str(metas_out.get("genres")),
                    "keyscale": _none_if_na_str(metas_out.get("keyscale")),
                    "timesignature": _none_if_na_str(metas_out.get("timesignature")),
                    "lm_model": lm_model_name,
                    "dit_model": dit_model_name,
                }

            t0 = time.time()
            try:
                loop = asyncio.get_running_loop()
                result = await loop.run_in_executor(executor, _blocking_generate)
                job_store.mark_succeeded(job_id, result)

                # Update local cache
                _update_local_cache(job_id, result, "succeeded")
            except Exception as e:
                error_traceback = traceback.format_exc()
                print(f"[API Server] Job {job_id} FAILED: {e}")
                print(f"[API Server] Traceback:\n{error_traceback}")
                job_store.mark_failed(job_id, error_traceback)

                # Update local cache
                _update_local_cache(job_id, None, "failed")
            finally:
                # Best-effort cache cleanup to reduce MPS memory fragmentation between jobs
                try:
                    if hasattr(h, "_empty_cache"):
                        h._empty_cache()
                    else:
                        import torch
                        if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
                            torch.mps.empty_cache()
                except Exception:
                    pass
                dt = max(0.0, time.time() - t0)
                async with app.state.stats_lock:
                    app.state.recent_durations.append(dt)
                    if app.state.recent_durations:
                        app.state.avg_job_seconds = sum(app.state.recent_durations) / len(app.state.recent_durations)

        async def _queue_worker(worker_idx: int) -> None:
            while True:
                job_id, req = await app.state.job_queue.get()
                rec = store.get(job_id)
                try:
                    async with app.state.pending_lock:
                        try:
                            app.state.pending_ids.remove(job_id)
                        except ValueError:
                            pass

                    await _run_one_job(job_id, req)

                    # Notify OpenRouter waiters after job completion
                    if rec and rec.progress_queue:
                        if rec.status == "succeeded" and rec.result:
                            await rec.progress_queue.put({"type": "result", "result": rec.result})
                        elif rec.status == "failed":
                            await rec.progress_queue.put({"type": "error", "content": rec.error or "Generation failed"})
                        await rec.progress_queue.put({"type": "done"})
                    if rec and rec.done_event:
                        rec.done_event.set()

                except Exception as exc:
                    # _run_one_job raised (e.g. _ensure_initialized failed)
                    if rec and rec.status not in ("succeeded", "failed"):
                        store.mark_failed(job_id, str(exc))
                    if rec and rec.progress_queue:
                        await rec.progress_queue.put({"type": "error", "content": str(exc)})
                        await rec.progress_queue.put({"type": "done"})
                    if rec and rec.done_event:
                        rec.done_event.set()
                finally:
                    await _cleanup_job_temp_files(job_id)
                    app.state.job_queue.task_done()

        async def _job_store_cleanup_worker() -> None:
            """Background task to periodically clean up old completed jobs."""
            while True:
                try:
                    await asyncio.sleep(JOB_STORE_CLEANUP_INTERVAL)
                    removed = store.cleanup_old_jobs()
                    if removed > 0:
                        stats = store.get_stats()
                        print(f"[API Server] Cleaned up {removed} old jobs. Current stats: {stats}")
                except asyncio.CancelledError:
                    break
                except Exception as e:
                    print(f"[API Server] Job cleanup error: {e}")

        worker_count = max(1, WORKER_COUNT)
        workers = [asyncio.create_task(_queue_worker(i)) for i in range(worker_count)]
        cleanup_task = asyncio.create_task(_job_store_cleanup_worker())
        app.state.worker_tasks = workers
        app.state.cleanup_task = cleanup_task

        # =================================================================
        # Initialize models at startup (not lazily on first request)
        # =================================================================
        print("[API Server] Initializing models at startup...")

        # Detect GPU memory and get configuration
        gpu_config = get_gpu_config()
        set_global_gpu_config(gpu_config)
        app.state.gpu_config = gpu_config

        gpu_memory_gb = gpu_config.gpu_memory_gb
        auto_offload = gpu_memory_gb > 0 and gpu_memory_gb < VRAM_16GB_MIN_GB

        # Print GPU configuration info
        print(f"\n{'='*60}")
        print("[API Server] GPU Configuration Detected:")
        print(f"{'='*60}")
        print(f"  GPU Memory: {gpu_memory_gb:.2f} GB")
        print(f"  Configuration Tier: {gpu_config.tier}")
        print(f"  Max Duration (with LM): {gpu_config.max_duration_with_lm}s")
        print(f"  Max Duration (without LM): {gpu_config.max_duration_without_lm}s")
        print(f"  Max Batch Size (with LM): {gpu_config.max_batch_size_with_lm}")
        print(f"  Max Batch Size (without LM): {gpu_config.max_batch_size_without_lm}")
        print(f"  Default LM Init: {gpu_config.init_lm_default}")
        print(f"  Available LM Models: {gpu_config.available_lm_models or 'None'}")
        print(f"{'='*60}\n")

        if auto_offload:
            print(f"[API Server] Auto-enabling CPU offload (GPU < 16GB)")
        elif gpu_memory_gb > 0:
            print(f"[API Server] CPU offload disabled by default (GPU >= 16GB)")
        else:
            print("[API Server] No GPU detected, running on CPU")

        project_root = _get_project_root()
        config_path = os.getenv("ACESTEP_CONFIG_PATH", "acestep-v15-turbo")
        device = os.getenv("ACESTEP_DEVICE", "auto")
        use_flash_attention = _env_bool("ACESTEP_USE_FLASH_ATTENTION", True)

        # Auto-determine offload settings based on GPU config if not explicitly set
        offload_to_cpu_env = os.getenv("ACESTEP_OFFLOAD_TO_CPU")
        if offload_to_cpu_env is not None:
            offload_to_cpu = _env_bool("ACESTEP_OFFLOAD_TO_CPU", False)
        else:
            offload_to_cpu = auto_offload
            if auto_offload:
                print(f"[API Server] Auto-setting offload_to_cpu=True based on GPU memory")

        offload_dit_to_cpu = _env_bool("ACESTEP_OFFLOAD_DIT_TO_CPU", False)

        # Checkpoint directory
        checkpoint_dir = os.path.join(project_root, "checkpoints")
        os.makedirs(checkpoint_dir, exist_ok=True)

        # Download and initialize primary DiT model
        dit_model_name = _get_model_name(config_path)
        if dit_model_name:
            try:
                _ensure_model_downloaded(dit_model_name, checkpoint_dir)
            except Exception as e:
                print(f"[API Server] Warning: Failed to download DiT model: {e}")

        # Download VAE model
        try:
            _ensure_model_downloaded("vae", checkpoint_dir)
        except Exception as e:
            print(f"[API Server] Warning: Failed to download VAE model: {e}")

        print(f"[API Server] Loading primary DiT model: {config_path}")
        status_msg, ok = handler.initialize_service(
            project_root=project_root,
            config_path=config_path,
            device=device,
            use_flash_attention=use_flash_attention,
            compile_model=False,
            offload_to_cpu=offload_to_cpu,
            offload_dit_to_cpu=offload_dit_to_cpu,
        )
        if not ok:
            app.state._init_error = status_msg
            print(f"[API Server] ERROR: Primary model failed to load: {status_msg}")
            raise RuntimeError(status_msg)
        app.state._initialized = True
        print(f"[API Server] Primary model loaded: {_get_model_name(config_path)}")

        # Initialize secondary model if configured
        if handler2 and config_path2:
            model2_name = _get_model_name(config_path2)
            if model2_name:
                try:
                    _ensure_model_downloaded(model2_name, checkpoint_dir)
                except Exception as e:
                    print(f"[API Server] Warning: Failed to download secondary model: {e}")

            print(f"[API Server] Loading secondary DiT model: {config_path2}")
            try:
                status_msg2, ok2 = handler2.initialize_service(
                    project_root=project_root,
                    config_path=config_path2,
                    device=device,
                    use_flash_attention=use_flash_attention,
                    compile_model=False,
                    offload_to_cpu=offload_to_cpu,
                    offload_dit_to_cpu=offload_dit_to_cpu,
                )
                app.state._initialized2 = ok2
                if ok2:
                    print(f"[API Server] Secondary model loaded: {model2_name}")
                else:
                    print(f"[API Server] Warning: Secondary model failed: {status_msg2}")
            except Exception as e:
                print(f"[API Server] Warning: Failed to initialize secondary model: {e}")
                app.state._initialized2 = False

        # Initialize third model if configured
        if handler3 and config_path3:
            model3_name = _get_model_name(config_path3)
            if model3_name:
                try:
                    _ensure_model_downloaded(model3_name, checkpoint_dir)
                except Exception as e:
                    print(f"[API Server] Warning: Failed to download third model: {e}")

            print(f"[API Server] Loading third DiT model: {config_path3}")
            try:
                status_msg3, ok3 = handler3.initialize_service(
                    project_root=project_root,
                    config_path=config_path3,
                    device=device,
                    use_flash_attention=use_flash_attention,
                    compile_model=False,
                    offload_to_cpu=offload_to_cpu,
                    offload_dit_to_cpu=offload_dit_to_cpu,
                )
                app.state._initialized3 = ok3
                if ok3:
                    print(f"[API Server] Third model loaded: {model3_name}")
                else:
                    print(f"[API Server] Warning: Third model failed: {status_msg3}")
            except Exception as e:
                print(f"[API Server] Warning: Failed to initialize third model: {e}")
                app.state._initialized3 = False

        # Initialize LLM model based on GPU configuration
        # ACESTEP_INIT_LLM controls LLM initialization:
        #   - "auto" / empty / not set: Use GPU config default (auto-detect)
        #   - "true"/"1"/"yes": Force enable LLM after GPU config is applied
        #   - "false"/"0"/"no": Force disable LLM
        #
        # Flow: GPU detection → model validation → ACESTEP_INIT_LLM override
        # This ensures GPU optimizations (offload, quantization, etc.) are always applied.
        init_llm_env = os.getenv("ACESTEP_INIT_LLM", "").strip().lower()

        # Step 1: Start with GPU auto-detection result
        init_llm = gpu_config.init_lm_default
        print(f"[API Server] GPU auto-detection: init_llm={init_llm} (VRAM: {gpu_config.gpu_memory_gb:.1f}GB, tier: {gpu_config.tier})")

        # Step 2: Apply user override if set
        if not init_llm_env or init_llm_env == "auto":
            print(f"[API Server] ACESTEP_INIT_LLM=auto, using GPU auto-detection result")
        elif init_llm_env in {"1", "true", "yes", "y", "on"}:
            if init_llm:
                print(f"[API Server] ACESTEP_INIT_LLM=true (GPU already supports LLM, no override needed)")
            else:
                init_llm = True
                print(f"[API Server] ACESTEP_INIT_LLM=true, overriding GPU auto-detection (force enable)")
        else:
            if not init_llm:
                print(f"[API Server] ACESTEP_INIT_LLM=false (GPU already disabled LLM, no override needed)")
            else:
                init_llm = False
                print(f"[API Server] ACESTEP_INIT_LLM=false, overriding GPU auto-detection (force disable)")

        if init_llm:
            print("[API Server] Loading LLM model...")

            # Auto-select LM model based on GPU config if not explicitly set
            lm_model_path_env = os.getenv("ACESTEP_LM_MODEL_PATH", "").strip()
            if lm_model_path_env:
                lm_model_path = lm_model_path_env
                print(f"[API Server] Using user-specified LM model: {lm_model_path}")
            else:
                # Get recommended LM model for this GPU tier
                recommended_lm = get_recommended_lm_model(gpu_config)
                if recommended_lm:
                    lm_model_path = recommended_lm
                    print(f"[API Server] Auto-selected LM model: {lm_model_path} based on GPU tier")
                else:
                    # No recommended model (GPU tier too low), default to smallest
                    lm_model_path = "acestep-5Hz-lm-0.6B"
                    print(f"[API Server] No recommended model for this GPU tier, using smallest: {lm_model_path}")

            # Validate LM model support (warning only, does not block)
            is_supported, warning_msg = is_lm_model_supported(lm_model_path, gpu_config)
            if not is_supported:
                print(f"[API Server] Warning: {warning_msg}")
                # Try to fall back to a supported model
                recommended_lm = get_recommended_lm_model(gpu_config)
                if recommended_lm:
                    lm_model_path = recommended_lm
                    print(f"[API Server] Falling back to supported LM model: {lm_model_path}")
                else:
                    # No supported model, but user may have forced init
                    print(f"[API Server] No GPU-validated LM model available, attempting {lm_model_path} anyway (may cause OOM)")

        if init_llm:
            lm_backend = os.getenv("ACESTEP_LM_BACKEND", "vllm").strip().lower()
            if lm_backend not in {"vllm", "pt", "mlx"}:
                lm_backend = "vllm"
            lm_device = os.getenv("ACESTEP_LM_DEVICE", device)

            # Auto-determine LM offload based on GPU config
            lm_offload_env = os.getenv("ACESTEP_LM_OFFLOAD_TO_CPU")
            if lm_offload_env is not None:
                lm_offload = _env_bool("ACESTEP_LM_OFFLOAD_TO_CPU", False)
            else:
                lm_offload = offload_to_cpu

            try:
                _ensure_model_downloaded(lm_model_path, checkpoint_dir)
            except Exception as e:
                print(f"[API Server] Warning: Failed to download LLM model: {e}")

            llm_status, llm_ok = llm_handler.initialize(
                checkpoint_dir=checkpoint_dir,
                lm_model_path=lm_model_path,
                backend=lm_backend,
                device=lm_device,
                offload_to_cpu=lm_offload,
                dtype=None,
            )
            if llm_ok:
                app.state._llm_initialized = True
                print(f"[API Server] LLM model loaded: {lm_model_path}")
            else:
                app.state._llm_init_error = llm_status
                print(f"[API Server] Warning: LLM model failed to load: {llm_status}")
        else:
            print("[API Server] Skipping LLM initialization (disabled or not supported for this GPU)")
            app.state._llm_initialized = False
            # Disable lazy loading of LLM - don't try to load it later during requests
            app.state._llm_lazy_load_disabled = True
            print("[API Server] LLM lazy loading disabled. To enable LLM:")
            print("[API Server]   - Set ACESTEP_INIT_LLM=true in .env or environment")
            print("[API Server]   - Or use --init-llm command line flag")

        print("[API Server] All models initialized successfully!")

        try:
            yield
        finally:
            cleanup_task.cancel()
            for t in workers:
                t.cancel()
            executor.shutdown(wait=False, cancel_futures=True)

    app = FastAPI(title="ACE-Step API", version="1.0", lifespan=lifespan)

    # Mount OpenRouter-compatible endpoints (/v1/chat/completions, /v1/models)
    from acestep.openrouter_adapter import create_openrouter_router
    openrouter_router = create_openrouter_router(lambda: app.state)
    app.include_router(openrouter_router)

    async def _queue_position(job_id: str) -> int:
        async with app.state.pending_lock:
            try:
                return list(app.state.pending_ids).index(job_id) + 1
            except ValueError:
                return 0

    async def _eta_seconds_for_position(pos: int) -> Optional[float]:
        if pos <= 0:
            return None
        async with app.state.stats_lock:
            avg = float(getattr(app.state, "avg_job_seconds", INITIAL_AVG_JOB_SECONDS))
        return pos * avg

    @app.post("/release_task")
    async def create_music_generate_job(request: Request, authorization: Optional[str] = Header(None)):
        content_type = (request.headers.get("content-type") or "").lower()
        temp_files: list[str] = []

        def _build_request(p: RequestParser, **kwargs) -> GenerateMusicRequest:
            """Build GenerateMusicRequest from parsed parameters."""
            return GenerateMusicRequest(
                prompt=p.str("prompt"),
                lyrics=p.str("lyrics"),
                thinking=p.bool("thinking"),
                analysis_only=p.bool("analysis_only"),
                full_analysis_only=p.bool("full_analysis_only"),
                sample_mode=p.bool("sample_mode"),
                sample_query=p.str("sample_query"),
                use_format=p.bool("use_format"),
                model=p.str("model") or None,
                bpm=p.int("bpm"),
                key_scale=p.str("key_scale"),
                time_signature=p.str("time_signature"),
                audio_duration=p.float("audio_duration"),
                vocal_language=p.str("vocal_language", "en"),
                inference_steps=p.int("inference_steps", 8),
                guidance_scale=p.float("guidance_scale", 7.0),
                use_random_seed=p.bool("use_random_seed", True),
                seed=p.int("seed", -1),
                batch_size=p.int("batch_size"),
                repainting_start=p.float("repainting_start", 0.0),
                repainting_end=p.float("repainting_end"),
                instruction=p.str("instruction", DEFAULT_DIT_INSTRUCTION),
                audio_cover_strength=p.float("audio_cover_strength", 1.0),
                reference_audio_path=p.str("reference_audio_path") or None,
                src_audio_path=p.str("src_audio_path") or None,
                task_type=p.str("task_type", "text2music"),
                use_adg=p.bool("use_adg"),
                cfg_interval_start=p.float("cfg_interval_start", 0.0),
                cfg_interval_end=p.float("cfg_interval_end", 1.0),
                infer_method=p.str("infer_method", "ode"),
                shift=p.float("shift", 3.0),
                audio_format=p.str("audio_format", "mp3"),
                use_tiled_decode=p.bool("use_tiled_decode", True),
                lm_model_path=p.str("lm_model_path") or None,
                lm_backend=p.str("lm_backend", "vllm"),
                lm_temperature=p.float("lm_temperature", LM_DEFAULT_TEMPERATURE),
                lm_cfg_scale=p.float("lm_cfg_scale", LM_DEFAULT_CFG_SCALE),
                lm_top_k=p.int("lm_top_k"),
                lm_top_p=p.float("lm_top_p", LM_DEFAULT_TOP_P),
                lm_repetition_penalty=p.float("lm_repetition_penalty", 1.0),
                lm_negative_prompt=p.str("lm_negative_prompt", "NO USER INPUT"),
                constrained_decoding=p.bool("constrained_decoding", True),
                constrained_decoding_debug=p.bool("constrained_decoding_debug"),
                use_cot_caption=p.bool("use_cot_caption", True),
                use_cot_language=p.bool("use_cot_language", True),
                is_format_caption=p.bool("is_format_caption"),
                allow_lm_batch=p.bool("allow_lm_batch", True),
                **kwargs,
            )

        if content_type.startswith("application/json"):
            body = await request.json()
            if not isinstance(body, dict):
                raise HTTPException(status_code=400, detail="JSON payload must be an object")
            verify_token_from_request(body, authorization)
            req = _build_request(RequestParser(body))

        elif content_type.endswith("+json"):
            body = await request.json()
            if not isinstance(body, dict):
                raise HTTPException(status_code=400, detail="JSON payload must be an object")
            verify_token_from_request(body, authorization)
            req = _build_request(RequestParser(body))

        elif content_type.startswith("multipart/form-data"):
            form = await request.form()
            form_dict = {k: v for k, v in form.items() if not hasattr(v, 'read')}
            verify_token_from_request(form_dict, authorization)

            # Support both naming conventions: ref_audio/reference_audio, ctx_audio/src_audio
            ref_up = form.get("ref_audio") or form.get("reference_audio")
            ctx_up = form.get("ctx_audio") or form.get("src_audio")

            reference_audio_path = None
            src_audio_path = None

            if isinstance(ref_up, StarletteUploadFile):
                reference_audio_path = await _save_upload_to_temp(ref_up, prefix="ref_audio")
                temp_files.append(reference_audio_path)
            else:
                reference_audio_path = str(form.get("ref_audio_path") or form.get("reference_audio_path") or "").strip() or None

            if isinstance(ctx_up, StarletteUploadFile):
                src_audio_path = await _save_upload_to_temp(ctx_up, prefix="ctx_audio")
                temp_files.append(src_audio_path)
            else:
                src_audio_path = str(form.get("ctx_audio_path") or form.get("src_audio_path") or "").strip() or None

            req = _build_request(
                RequestParser(dict(form)),
                reference_audio_path=reference_audio_path,
                src_audio_path=src_audio_path,
            )

        elif content_type.startswith("application/x-www-form-urlencoded"):
            form = await request.form()
            form_dict = dict(form)
            verify_token_from_request(form_dict, authorization)
            reference_audio_path = str(form.get("ref_audio_path") or form.get("reference_audio_path") or "").strip() or None
            src_audio_path = str(form.get("ctx_audio_path") or form.get("src_audio_path") or "").strip() or None
            req = _build_request(
                RequestParser(form_dict),
                reference_audio_path=reference_audio_path,
                src_audio_path=src_audio_path,
            )

        else:
            raw = await request.body()
            raw_stripped = raw.lstrip()
            # Best-effort: accept missing/incorrect Content-Type if payload is valid JSON.
            if raw_stripped.startswith(b"{") or raw_stripped.startswith(b"["):
                try:
                    body = json.loads(raw.decode("utf-8"))
                    if isinstance(body, dict):
                        verify_token_from_request(body, authorization)
                        req = _build_request(RequestParser(body))
                    else:
                        raise HTTPException(status_code=400, detail="JSON payload must be an object")
                except HTTPException:
                    raise
                except Exception:
                    raise HTTPException(
                        status_code=400,
                        detail="Invalid JSON body (hint: set 'Content-Type: application/json')",
                    )
            # Best-effort: parse key=value bodies even if Content-Type is missing.
            elif raw_stripped and b"=" in raw:
                parsed = urllib.parse.parse_qs(raw.decode("utf-8"), keep_blank_values=True)
                flat = {k: (v[0] if isinstance(v, list) and v else v) for k, v in parsed.items()}
                verify_token_from_request(flat, authorization)
                reference_audio_path = str(flat.get("ref_audio_path") or flat.get("reference_audio_path") or "").strip() or None
                src_audio_path = str(flat.get("ctx_audio_path") or flat.get("src_audio_path") or "").strip() or None
                req = _build_request(
                    RequestParser(flat),
                    reference_audio_path=reference_audio_path,
                    src_audio_path=src_audio_path,
                )
            else:
                raise HTTPException(
                    status_code=415,
                    detail=(
                        f"Unsupported Content-Type: {content_type or '(missing)'}; "
                        "use application/json, application/x-www-form-urlencoded, or multipart/form-data"
                    ),
                )

        rec = store.create()

        q: asyncio.Queue = app.state.job_queue
        if q.full():
            for p in temp_files:
                try:
                    os.remove(p)
                except Exception:
                    pass
            raise HTTPException(status_code=429, detail="Server busy: queue is full")

        if temp_files:
            async with app.state.job_temp_files_lock:
                app.state.job_temp_files[rec.job_id] = temp_files

        async with app.state.pending_lock:
            app.state.pending_ids.append(rec.job_id)
            position = len(app.state.pending_ids)

        await q.put((rec.job_id, req))
        return _wrap_response({"task_id": rec.job_id, "status": "queued", "queue_position": position})

    @app.post("/query_result")
    async def query_result(request: Request, authorization: Optional[str] = Header(None)):
        """Batch query job results"""
        content_type = (request.headers.get("content-type") or "").lower()

        if "json" in content_type:
            body = await request.json()
        else:
            form = await request.form()
            body = {k: v for k, v in form.items()}

        verify_token_from_request(body, authorization)
        task_id_list_str = body.get("task_id_list", "[]")

        # Parse task ID list
        if isinstance(task_id_list_str, list):
            task_id_list = task_id_list_str
        else:
            try:
                task_id_list = json.loads(task_id_list_str)
            except Exception:
                task_id_list = []

        local_cache = getattr(app.state, 'local_cache', None)
        data_list = []
        current_time = time.time()

        for task_id in task_id_list:
            result_key = f"{RESULT_KEY_PREFIX}{task_id}"

            # Read from local cache first
            if local_cache:
                data = local_cache.get(result_key)
                if data:
                    try:
                        data_json = json.loads(data)
                    except Exception:
                        data_json = []

                    if len(data_json) <= 0:
                        data_list.append({"task_id": task_id, "result": data, "status": 2})
                    else:
                        status = data_json[0].get("status")
                        create_time = data_json[0].get("create_time", 0)
                        if status == 0 and (current_time - create_time) > TASK_TIMEOUT_SECONDS:
                            data_list.append({"task_id": task_id, "result": data, "status": 2})
                        else:
                            data_list.append({
                                "task_id": task_id,
                                "result": data,
                                "status": int(status) if status is not None else 1,
                                "progress_text": log_buffer.last_message
                            })
                    continue

            # Fallback to job_store query
            rec = store.get(task_id)
            if rec:
                env = getattr(rec, 'env', 'development')
                create_time = rec.created_at
                status_int = _map_status(rec.status)

                if rec.result and rec.status == "succeeded":
                    # Check if it's a "Full Analysis" result
                    if rec.result.get("status_message") == "Full Hardware Analysis Success":
                         result_data = [rec.result]
                    else:
                        audio_paths = rec.result.get("audio_paths", [])
                        metas = rec.result.get("metas", {}) or {}
                        result_data = [
                            {
                                "file": p, "wave": "", "status": status_int,
                                "create_time": int(create_time), "env": env,
                                "prompt": metas.get("caption", ""),
                                "lyrics": metas.get("lyrics", ""),
                                "metas": {
                                    "bpm": metas.get("bpm"),
                                    "duration": metas.get("duration"),
                                    "genres": metas.get("genres", ""),
                                    "keyscale": metas.get("keyscale", ""),
                                    "timesignature": metas.get("timesignature", ""),
                                }
                            }
                            for p in audio_paths
                        ] if audio_paths else [{
                            "file": "", "wave": "", "status": status_int,
                            "create_time": int(create_time), "env": env,
                            "prompt": metas.get("caption", ""),
                            "lyrics": metas.get("lyrics", ""),
                            "metas": {
                                "bpm": metas.get("bpm"),
                                "duration": metas.get("duration"),
                                "genres": metas.get("genres", ""),
                                "keyscale": metas.get("keyscale", ""),
                                "timesignature": metas.get("timesignature", ""),
                            }
                        }]
                else:
                    result_data = [{
                        "file": "", "wave": "", "status": status_int,
                        "create_time": int(create_time), "env": env,
                        "prompt": "", "lyrics": "",
                        "metas": {},
                        "progress": float(rec.progress) if rec else 0.0,
                        "stage": rec.stage if rec else "queued",
                        "error": rec.error if rec.error else None,
                    }]

                current_log = log_buffer.last_message if status_int == 0 else rec.progress_text
                data_list.append({
                    "task_id": task_id,
                    "result": json.dumps(result_data, ensure_ascii=False),
                    "status": status_int,
                    "progress_text": current_log
                })
            else:
                data_list.append({"task_id": task_id, "result": "[]", "status": 0})

        return _wrap_response(data_list)

    @app.get("/health")
    async def health_check():
        """Health check endpoint for service status."""
        return _wrap_response({
            "status": "ok",
            "service": "ACE-Step API",
            "version": "1.0",
        })

    @app.get("/v1/stats")
    async def get_stats(_: None = Depends(verify_api_key)):
        """Get server statistics including job store stats."""
        job_stats = store.get_stats()
        async with app.state.stats_lock:
            avg_job_seconds = getattr(app.state, "avg_job_seconds", INITIAL_AVG_JOB_SECONDS)
        return _wrap_response({
            "jobs": job_stats,
            "queue_size": app.state.job_queue.qsize(),
            "queue_maxsize": QUEUE_MAXSIZE,
            "avg_job_seconds": avg_job_seconds,
        })

    @app.get("/v1/models")
    async def list_models(_: None = Depends(verify_api_key)):
        """List available DiT models."""
        models = []
        
        # Primary model (always available if initialized)
        if getattr(app.state, "_initialized", False):
            primary_model = _get_model_name(app.state._config_path)
            if primary_model:
                models.append({
                    "name": primary_model,
                    "is_default": True,
                })
        
        # Secondary model
        if getattr(app.state, "_initialized2", False) and app.state._config_path2:
            secondary_model = _get_model_name(app.state._config_path2)
            if secondary_model:
                models.append({
                    "name": secondary_model,
                    "is_default": False,
                })
        
        # Third model
        if getattr(app.state, "_initialized3", False) and app.state._config_path3:
            third_model = _get_model_name(app.state._config_path3)
            if third_model:
                models.append({
                    "name": third_model,
                    "is_default": False,
                })
        
        return _wrap_response({
            "models": models,
            "default_model": models[0]["name"] if models else None,
        })

    @app.post("/create_random_sample")
    async def create_random_sample_endpoint(request: Request, authorization: Optional[str] = Header(None)):
        """
        Get random sample parameters from pre-loaded example data.

        Returns a random example from the examples directory for form filling.
        """
        content_type = (request.headers.get("content-type") or "").lower()

        if "json" in content_type:
            body = await request.json()
        else:
            form = await request.form()
            body = {k: v for k, v in form.items()}

        verify_token_from_request(body, authorization)
        sample_type = body.get("sample_type", "simple_mode") or "simple_mode"

        if sample_type == "simple_mode":
            example_data = SIMPLE_EXAMPLE_DATA
        else:
            example_data = CUSTOM_EXAMPLE_DATA

        if not example_data:
            return _wrap_response(None, code=500, error="No example data available")

        random_example = random.choice(example_data)
        return _wrap_response(random_example)

    @app.post("/format_input")
    async def format_input_endpoint(request: Request, authorization: Optional[str] = Header(None)):
        """
        Format and enhance lyrics/caption via LLM.

        Takes user-provided caption and lyrics, and uses the LLM to enhance them
        with proper structure and metadata.
        """
        content_type = (request.headers.get("content-type") or "").lower()

        if "json" in content_type:
            body = await request.json()
        else:
            form = await request.form()
            body = {k: v for k, v in form.items()}

        verify_token_from_request(body, authorization)
        llm: LLMHandler = app.state.llm_handler

        # Initialize LLM if needed
        with app.state._llm_init_lock:
            if not getattr(app.state, "_llm_initialized", False):
                if getattr(app.state, "_llm_init_error", None):
                    raise HTTPException(status_code=500, detail=f"LLM init failed: {app.state._llm_init_error}")

                # Check if lazy loading is disabled
                if getattr(app.state, "_llm_lazy_load_disabled", False):
                    raise HTTPException(
                        status_code=503,
                        detail="LLM not initialized. Set ACESTEP_INIT_LLM=true in .env to enable."
                    )

                project_root = _get_project_root()
                checkpoint_dir = os.path.join(project_root, "checkpoints")
                lm_model_path = os.getenv("ACESTEP_LM_MODEL_PATH", "acestep-5Hz-lm-0.6B").strip()
                backend = os.getenv("ACESTEP_LM_BACKEND", "vllm").strip().lower()
                if backend not in {"vllm", "pt", "mlx"}:
                    backend = "vllm"

                # Auto-download LM model if not present
                lm_model_name = _get_model_name(lm_model_path)
                if lm_model_name:
                    try:
                        _ensure_model_downloaded(lm_model_name, checkpoint_dir)
                    except Exception as e:
                        print(f"[API Server] Warning: Failed to download LM model {lm_model_name}: {e}")

                lm_device = os.getenv("ACESTEP_LM_DEVICE", os.getenv("ACESTEP_DEVICE", "auto"))
                lm_offload = _env_bool("ACESTEP_LM_OFFLOAD_TO_CPU", False)

                h: AceStepHandler = app.state.handler
                status, ok = llm.initialize(
                    checkpoint_dir=checkpoint_dir,
                    lm_model_path=lm_model_path,
                    backend=backend,
                    device=lm_device,
                    offload_to_cpu=lm_offload,
                    dtype=None,
                )
                if not ok:
                    app.state._llm_init_error = status
                    raise HTTPException(status_code=500, detail=f"LLM init failed: {status}")
                app.state._llm_initialized = True

        # Parse parameters
        prompt = body.get("prompt", "") or ""
        lyrics = body.get("lyrics", "") or ""
        temperature = _to_float(body.get("temperature"), 0.85)

        # Parse param_obj if provided
        param_obj_str = body.get("param_obj", "{}")
        if isinstance(param_obj_str, dict):
            param_obj = param_obj_str
        else:
            try:
                param_obj = json.loads(param_obj_str) if param_obj_str else {}
            except json.JSONDecodeError:
                param_obj = {}

        # Extract metadata from param_obj
        duration = _to_float(param_obj.get("duration"))
        bpm = _to_int(param_obj.get("bpm"))
        key_scale = param_obj.get("key", "") or param_obj.get("key_scale", "") or ""
        time_signature = param_obj.get("time_signature", "") or body.get("time_signature", "") or ""
        language = param_obj.get("language", "") or ""

        # Build user_metadata for format_sample
        user_metadata_for_format = {}
        if bpm is not None:
            user_metadata_for_format['bpm'] = bpm
        if duration is not None and duration > 0:
            user_metadata_for_format['duration'] = int(duration)
        if key_scale:
            user_metadata_for_format['keyscale'] = key_scale
        if time_signature:
            user_metadata_for_format['timesignature'] = time_signature
        if language and language != "unknown":
            user_metadata_for_format['language'] = language

        # Call format_sample
        try:
            format_result = format_sample(
                llm_handler=llm,
                caption=prompt,
                lyrics=lyrics,
                user_metadata=user_metadata_for_format if user_metadata_for_format else None,
                temperature=temperature,
                use_constrained_decoding=True,
            )

            if not format_result.success:
                error_msg = format_result.error or format_result.status_message
                return _wrap_response(None, code=500, error=f"format_sample failed: {error_msg}")

            # Use formatted results or fallback to original
            result_caption = format_result.caption or prompt
            result_lyrics = format_result.lyrics or lyrics
            result_duration = format_result.duration or duration
            result_bpm = format_result.bpm or bpm
            result_key_scale = format_result.keyscale or key_scale
            result_time_signature = format_result.timesignature or time_signature

            return _wrap_response({
                "caption": result_caption,
                "lyrics": result_lyrics,
                "bpm": result_bpm,
                "key_scale": result_key_scale,
                "time_signature": result_time_signature,
                "duration": result_duration,
                "vocal_language": format_result.language or language or "unknown",
            })
        except Exception as e:
            return _wrap_response(None, code=500, error=f"format_sample error: {str(e)}")

    @app.get("/v1/audio")
    async def get_audio(path: str, request: Request, _: None = Depends(verify_api_key)):
        """Serve audio file by path."""
        from fastapi.responses import FileResponse

        # Security: Validate path is within allowed directory to prevent path traversal
        resolved_path = os.path.realpath(path)
        allowed_dir = os.path.realpath(request.app.state.temp_audio_dir)
        if not resolved_path.startswith(allowed_dir + os.sep) and resolved_path != allowed_dir:
            raise HTTPException(status_code=403, detail="Access denied: path outside allowed directory")
        if not os.path.exists(resolved_path):
            raise HTTPException(status_code=404, detail="Audio file not found")

        ext = os.path.splitext(resolved_path)[1].lower()
        media_types = {
            ".mp3": "audio/mpeg",
            ".wav": "audio/wav",
            ".flac": "audio/flac",
            ".ogg": "audio/ogg",
        }
        media_type = media_types.get(ext, "audio/mpeg")

        return FileResponse(resolved_path, media_type=media_type)

    return app


app = create_app()


def main() -> None:
    import argparse
    import uvicorn

    parser = argparse.ArgumentParser(description="ACE-Step API server")
    parser.add_argument(
        "--host",
        default=os.getenv("ACESTEP_API_HOST", "127.0.0.1"),
        help="Bind host (default from ACESTEP_API_HOST or 127.0.0.1)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=int(os.getenv("ACESTEP_API_PORT", "8001")),
        help="Bind port (default from ACESTEP_API_PORT or 8001)",
    )
    parser.add_argument(
        "--api-key",
        type=str,
        default=os.getenv("ACESTEP_API_KEY", None),
        help="API key for authentication (default from ACESTEP_API_KEY)",
    )
    parser.add_argument(
        "--download-source",
        type=str,
        choices=["huggingface", "modelscope", "auto"],
        default=os.getenv("ACESTEP_DOWNLOAD_SOURCE", "auto"),
        help="Preferred model download source: auto (default), huggingface, or modelscope",
    )
    parser.add_argument(
        "--init-llm",
        action="store_true",
        default=_env_bool("ACESTEP_INIT_LLM", False),
        help="Initialize LLM even if GPU memory is insufficient (may cause OOM). "
             "Can also be set via ACESTEP_INIT_LLM=true environment variable.",
    )
    parser.add_argument(
        "--lm-model-path",
        type=str,
        default=os.getenv("ACESTEP_LM_MODEL_PATH", ""),
        help="LM model to load (e.g., 'acestep-5Hz-lm-0.6B'). Default from ACESTEP_LM_MODEL_PATH.",
    )
    args = parser.parse_args()

    # Set API key from command line argument
    if args.api_key:
        os.environ["ACESTEP_API_KEY"] = args.api_key

    # Set download source preference
    if args.download_source and args.download_source != "auto":
        os.environ["ACESTEP_DOWNLOAD_SOURCE"] = args.download_source
        print(f"Using preferred download source: {args.download_source}")

    # Set init LLM flag
    if args.init_llm:
        os.environ["ACESTEP_INIT_LLM"] = "true"
        print("[API Server] LLM initialization enabled via --init-llm")

    # Set LM model path
    if args.lm_model_path:
        os.environ["ACESTEP_LM_MODEL_PATH"] = args.lm_model_path
        print(f"[API Server] Using LM model: {args.lm_model_path}")

    # IMPORTANT: in-memory queue/store -> workers MUST be 1
    uvicorn.run(
        "acestep.api_server:app",
        host=str(args.host),
        port=int(args.port),
        reload=False,
        workers=1,
    )

if __name__ == "__main__":
    main()