Queue jobs until models available; fix progress tracking for UI

- Backend: jobs wait in pending_model until required DiT is installed; promote to queued when model available
- API: get_status returns pendingReason; cancel supports pending_model; POST /api/generate/retry-pending to promote and start after model download
- Frontend: SongList shows 'Waiting for model' + reason and Open Settings; poll sets generationStatus/generationPendingReason; Settings onDownloadComplete calls retryPending; cancel and pending_model handling in poll
- Progress: broaden tqdm regex so all progress lines match (INFO); add progress updater from log parser so job progress/ETA updates for UI even when callback path fails

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

api/ace_step_models.py

@@ -5,6 +5,7 @@
 """
 
 from pathlib import Path
+import shutil
 import subprocess
 import sys
 import threading
@@ -24,22 +25,22 @@ def _bundled_downloader_available() -> bool:
 
 bp = Blueprint("api_ace_step_models", __name__)
 
-# DiT variants from Tutorial (DiT Selection Summary)
+# DiT variants from Tutorial (DiT Selection Summary). size_gb: approximate for download confirmation.
 DIT_MODELS = [
-    {"id": "turbo", "label": "Turbo (default)", "description": "Best balance, 8 steps", "steps": 8, "cfg": False},
-    {"id": "turbo-shift1", "label": "Turbo shift=1", "description": "Richer details", "steps": 8, "cfg": False},
-    {"id": "turbo-shift3", "label": "Turbo shift=3", "description": "Clearer timbre", "steps": 8, "cfg": False},
-    {"id": "turbo-continuous", "label": "Turbo continuous", "description": "Flexible shift 1–5", "steps": 8, "cfg": False},
-    {"id": "sft", "label": "SFT", "description": "50 steps, CFG", "steps": 50, "cfg": True},
-    {"id": "base", "label": "Base", "description": "50 steps, CFG; lego/extract/complete", "steps": 50, "cfg": True, "exclusive_tasks": ["lego", "extract", "complete"]},
+    {"id": "turbo", "label": "Turbo (default)", "description": "Best balance, 8 steps", "steps": 8, "cfg": False, "size_gb": 8},
+    {"id": "turbo-shift1", "label": "Turbo shift=1", "description": "Richer details", "steps": 8, "cfg": False, "size_gb": 0.5},
+    {"id": "turbo-shift3", "label": "Turbo shift=3", "description": "Clearer timbre", "steps": 8, "cfg": False, "size_gb": 0.5},
+    {"id": "turbo-continuous", "label": "Turbo continuous", "description": "Flexible shift 1–5", "steps": 8, "cfg": False, "size_gb": 0.5},
+    {"id": "sft", "label": "SFT", "description": "50 steps, CFG", "steps": 50, "cfg": True, "size_gb": 8},
+    {"id": "base", "label": "Base", "description": "50 steps, CFG; lego/extract/complete", "steps": 50, "cfg": True, "exclusive_tasks": ["lego", "extract", "complete"], "size_gb": 8},
 ]
 
-# LM planner options from Tutorial
+# LM planner options from Tutorial. size_gb: approximate for download confirmation.
 LM_MODELS = [
-    {"id": "none", "label": "No LM"},
-    {"id": "0.6B", "label": "0.6B"},
-    {"id": "1.7B", "label": "1.7B (default)"},
-    {"id": "4B", "label": "4B"},
+    {"id": "none", "label": "No LM", "size_gb": 0},
+    {"id": "0.6B", "label": "0.6B", "size_gb": 2},
+    {"id": "1.7B", "label": "1.7B (default)", "size_gb": 4},
+    {"id": "4B", "label": "4B", "size_gb": 10},
 ]
 
 # ACE-Step 1.5 CLI model ids (for acestep-download --model)
@@ -333,6 +334,23 @@ def _do_download_worker(model: str, root: Path) -> None:
         _download_cancel_requested = False
 
 
+@bp.route("/models/disk-space", methods=["GET"])
+def disk_space():
+    """
+    GET /api/ace-step/models/disk-space
+    Returns free and total disk space for the models/checkpoints path (for download confirmation).
+    """
+    try:
+        root = _checkpoint_root()
+        root.mkdir(parents=True, exist_ok=True)
+        usage = shutil.disk_usage(str(root))
+        free_gb = round(usage.free / (1024 ** 3), 2)
+        total_gb = round(usage.total / (1024 ** 3), 2)
+        return jsonify({"free_gb": free_gb, "total_gb": total_gb, "path": str(root)})
+    except Exception as e:
+        return jsonify({"error": str(e), "free_gb": None, "total_gb": None, "path": ""}), 500
+
+
 @bp.route("/models/download", methods=["POST"])
 def download_model():
     """

api/generate.py

@@ -21,7 +21,7 @@ def _uppercase_track_in_instruction(instruction):
         return instruction[: m.start(2)] + m.group(2).upper() + instruction[m.end(2) :]
     return instruction
 
-from cdmf_paths import get_output_dir, get_user_data_dir, load_config
+from cdmf_paths import get_output_dir, get_user_data_dir, get_models_folder, load_config, save_config
 from cdmf_tracks import get_audio_duration, list_lora_adapters, load_track_meta, save_track_meta
 from cdmf_generation_job import GenerationCancelled
 import cdmf_state
@@ -59,6 +59,38 @@ def _is_cancel_requested(job_id: str) -> bool:
         return job_id in _cancel_requested
 
 
+def _is_model_available(dit_tag: str) -> bool:
+    """Return True if the given DiT model is installed and ready (no download needed). Used to promote pending_model jobs."""
+    if not dit_tag or not isinstance(dit_tag, str):
+        return False
+    dit = dit_tag.strip().lower()
+    DIT_15_FOLDERS = {
+        "turbo": "acestep-v15-turbo",
+        "base": "acestep-v15-base",
+        "sft": "acestep-v15-sft",
+        "turbo-shift1": "acestep-v15-turbo-shift1",
+        "turbo-shift3": "acestep-v15-turbo-shift3",
+        "turbo-continuous": "acestep-v15-turbo-continuous",
+    }
+    REQUIRED_SUBDIRS = ("music_dcae_f8c8", "music_vocoder", "ace_step_transformer", "umt5-base")
+    folder = DIT_15_FOLDERS.get(dit)
+    models_root = Path(get_models_folder()) / "checkpoints"
+    if folder:
+        candidate = models_root / folder
+        if not candidate.exists():
+            return False
+        for sub in REQUIRED_SUBDIRS:
+            if not (candidate / sub).exists():
+                return False
+        return True
+    # Legacy v1
+    try:
+        from ace_model_setup import ace_models_present
+        return ace_models_present()
+    except Exception:
+        return False
+
+
 def _refs_dir() -> Path:
     d = get_user_data_dir() / "references"
     d.mkdir(parents=True, exist_ok=True)
@@ -115,10 +147,45 @@ def _on_job_progress(
 register_job_progress_callback(_on_job_progress)
 
 
+def _update_job_progress_from_log(
+    percent: int, current: int, total: int, eta_seconds: float | None
+) -> None:
+    """Update current job progress from parsed tqdm log line (log handler runs in same thread as worker)."""
+    with _jobs_lock:
+        jid = cdmf_state.get_current_generation_job_id()
+        if not jid:
+            return
+        job = _jobs.get(jid)
+        if not job:
+            return
+        job["progressPercent"] = round(float(percent), 1)
+        job["progressSteps"] = f"{current}/{total}"
+        if eta_seconds is not None:
+            job["progressEta"] = round(float(eta_seconds), 1)
+
+
 def _run_generation(job_id: str) -> None:
     """Background: run generate_track_ace and update job."""
     global _generation_busy, _current_job_id
+    prev_config = None
+    config_switched = False
     try:
+        with _jobs_lock:
+            job = _jobs.get(job_id)
+            if not job or job.get("status") != "queued":
+                return
+            job_dit = job.get("dit_model") or "turbo"
+        # If required model is not installed, leave job as pending_model so it runs after user installs it.
+        if not _is_model_available(job_dit):
+            logging.info("[API generate] Job %s waiting for model %s (not installed)", job_id, job_dit)
+            with _jobs_lock:
+                job = _jobs.get(job_id)
+                if job and job.get("status") == "queued":
+                    job["status"] = "pending_model"
+                    job["pendingReason"] = (
+                        f"Model '{job_dit}' is not installed. Install it from Settings → Models to run this job."
+                    )
+            return
         with _jobs_lock:
             job = _jobs.get(job_id)
             if not job or job.get("status") != "queued":
@@ -131,8 +198,20 @@ def _run_generation(job_id: str) -> None:
             _current_job_id = job_id
 
         cdmf_state.set_current_generation_job_id(job_id)
+        cdmf_state.set_progress_updater(_update_job_progress_from_log)
         cancel_check = lambda: _is_cancel_requested(job_id)
-        from generate_ace import generate_track_ace
+        from generate_ace import generate_track_ace, clear_ace_pipeline
+
+        # Use job's dit_model (e.g. base for cover). Temporarily switch config so pipeline loads the right model.
+        with _jobs_lock:
+            j = _jobs.get(job_id)
+            job_dit = (j.get("dit_model") or "turbo") if j else "turbo"
+        prev_config = load_config() or {}
+        prev_config = dict(prev_config)
+        config_switched = job_dit != (prev_config.get("ace_step_dit_model") or "turbo")
+        if config_switched:
+            save_config({**prev_config, "ace_step_dit_model": job_dit})
+            clear_ace_pipeline()
 
         params = job.get("params") or {}
         if not isinstance(params, dict):
@@ -451,27 +530,64 @@ def _run_generation(job_id: str) -> None:
                 job["status"] = "cancelled"
                 job["error"] = "Cancelled by user"
     except Exception as e:
-        logging.exception("Generation job %s failed", job_id)
-        with _jobs_lock:
-            job = _jobs.get(job_id)
-            if job:
-                job["status"] = "failed"
-                job["error"] = str(e)
+        err_msg = str(e)
+        # Keep job queued as pending_model when model is missing so it can run after user installs it.
+        if "not installed" in err_msg.lower() or "Settings → Models" in err_msg:
+            logging.info("[API generate] Job %s waiting for model: %s", job_id, err_msg[:120])
+            with _jobs_lock:
+                job = _jobs.get(job_id)
+                if job:
+                    job["status"] = "pending_model"
+                    job["pendingReason"] = err_msg
+                    job["error"] = None
+        else:
+            logging.exception("Generation job %s failed", job_id)
+            with _jobs_lock:
+                job = _jobs.get(job_id)
+                if job:
+                    job["status"] = "failed"
+                    job["error"] = err_msg
     finally:
+        cdmf_state.set_progress_updater(None)
+        if config_switched and prev_config:
+            save_config(prev_config)
+            clear_ace_pipeline()
         cdmf_state.set_current_generation_job_id(None)
         _generation_busy = False
         with _jobs_lock:
             _current_job_id = None
             _cancel_requested.discard(job_id)
-        # Start next queued job (skips cancelled: they are no longer "queued")
+        # Promote pending_model jobs to queued when their model is now available; then start first queued job.
         with _jobs_lock:
+            for jid in _job_order:
+                j = _jobs.get(jid)
+                if j and j.get("status") == "pending_model":
+                    dit = j.get("dit_model") or "turbo"
+                    if _is_model_available(dit):
+                        j["status"] = "queued"
+                        j["pendingReason"] = None
+                        logging.info("[API generate] Job %s promoted to queued (model %s now available)", jid, dit)
             for jid in _job_order:
                 j = _jobs.get(jid)
                 if j and j.get("status") == "queued":
                     threading.Thread(target=_run_generation, args=(jid,), daemon=True).start()
                     break
 
 
+@bp.route("/model-download-status", methods=["GET"])
+def get_model_download_status():
+    """GET /api/generate/model-download-status — whether pipeline is loading (may be downloading model files)."""
+    try:
+        st = getattr(cdmf_state, "GENERATION_MODEL_LOADING", {})
+        return jsonify({
+            "in_progress": bool(st.get("in_progress")),
+            "message": st.get("message") or "Preparing model (downloading if needed)...",
+        })
+    except Exception as e:
+        logging.warning("[API generate] model-download-status failed: %s", e)
+        return jsonify({"in_progress": False, "message": ""})
+
+
 @bp.route("/lora_adapters", methods=["GET"])
 def get_lora_adapters():
     """GET /api/generate/lora_adapters — list LoRA adapters (e.g. from Training or custom_lora)."""
@@ -554,7 +670,15 @@ def _str(v):
         except (TypeError, ValueError):
             params_copy = {}
         config = load_config()
-        dit_tag = config.get("ace_step_dit_model") or params_copy.get("aceStepDitModel") or "turbo"
+        # User override from Generation tab model selector takes precedence; else auto base for cover (per docs).
+        task_for_dit = (params_copy.get("task_type") or params_copy.get("taskType") or "text2music").strip().lower()
+        explicit_dit = (params_copy.get("aceStepDitModel") or params_copy.get("ace_step_dit_model") or "").strip()
+        if explicit_dit:
+            dit_tag = explicit_dit
+        elif task_for_dit == "cover":
+            dit_tag = "base"
+        else:
+            dit_tag = config.get("ace_step_dit_model") or "turbo"
         lm_tag = config.get("ace_step_lm") or params_copy.get("aceStepLm") or "1.7B"
         with _jobs_lock:
             _jobs[job_id] = {
@@ -608,10 +732,42 @@ def get_status(job_id: str):
         "progressStage": job.get("progressStage"),
         "result": job.get("result"),
         "error": job.get("error"),
+        "pendingReason": job.get("pendingReason") if status == "pending_model" else None,
     }
     return jsonify(out)
 
 
+@bp.route("/retry-pending", methods=["POST"])
+def retry_pending():
+    """POST /api/generate/retry-pending — promote pending_model jobs to queued if model is now available, start first queued job. Call after model download completes."""
+    global _generation_busy
+    promoted = 0
+    started = None
+    with _jobs_lock:
+        for jid in _job_order:
+            j = _jobs.get(jid)
+            if j and j.get("status") == "pending_model":
+                dit = j.get("dit_model") or "turbo"
+                if _is_model_available(dit):
+                    j["status"] = "queued"
+                    j["pendingReason"] = None
+                    promoted += 1
+                    logging.info("[API generate] Job %s promoted to queued (model %s now available)", jid, dit)
+        if not _generation_busy:
+            for jid in _job_order:
+                j = _jobs.get(jid)
+                if j and j.get("status") == "queued":
+                    _generation_busy = True
+                    threading.Thread(target=_run_generation, args=(jid,), daemon=True).start()
+                    started = jid
+                    break
+    return jsonify({
+        "ok": True,
+        "promoted": promoted,
+        "startedJobId": started,
+    })
+
+
 @bp.route("/unstick", methods=["POST"])
 def unstick_queue():
     """POST /api/generate/unstick — clear stuck worker state and start the next queued job (if any)."""
@@ -641,9 +797,10 @@ def cancel_job(job_id: str):
         if not job:
             return jsonify({"error": "Job not found"}), 404
         status = job.get("status", "unknown")
-        if status == "queued":
+        if status in ("queued", "pending_model"):
             job["status"] = "cancelled"
             job["error"] = "Cancelled by user"
+            job["pendingReason"] = None
             return jsonify({"cancelled": True, "jobId": job_id, "message": "Job removed from queue."})
         if status == "running":
             _cancel_requested.add(job_id)

cdmf_pipeline_ace_step.py

@@ -1195,7 +1195,8 @@ def add_latents_noise(
                 sigma_max=sigma_max
             )
 
-        infer_steps = int(sigma_max * infer_steps)
+        # Ensure enough steps for cover/audio2audio so reference is audible (INFERENCE.md: base 32-64 recommended).
+        infer_steps = max(16, int(sigma_max * infer_steps))
         timesteps, num_inference_steps = retrieve_timesteps(
             scheduler,
             num_inference_steps=infer_steps,
@@ -1295,6 +1296,17 @@ def text2music_diffusion_process(
 
         if ref_latents is not None:
             frame_length = ref_latents.shape[-1]
+            # Cap ref length for cover/audio2audio so each diffusion step stays fast (avoids 80s+ per step on long refs)
+            max_cover_sec = float(os.environ.get("ACE_COVER_MAX_REF_SECONDS", "90"))
+            max_cover_frames = int(max_cover_sec * 44100 / 512 / 8)
+            if frame_length > max_cover_frames:
+                ref_latents = ref_latents[:, :, :, :max_cover_frames].contiguous()
+                frame_length = max_cover_frames
+                logger.info(
+                    "Capped ref_latents to %d frames (~%.0fs) for faster cover/audio2audio generation (set ACE_COVER_MAX_REF_SECONDS to override).",
+                    max_cover_frames,
+                    max_cover_sec,
+                )
 
         if len(oss_steps) > 0:
             infer_steps = max(oss_steps)
@@ -2087,6 +2099,7 @@ def __call__(
 
         ref_latents = None
         if ref_audio_input is not None and audio2audio_enable:
+            # For cover mode: ref_audio_input = source song (song to cover), per docs/ACE-Step-INFERENCE.md.
             assert ref_audio_input is not None, "ref_audio_input is required for audio2audio task"
             assert os.path.exists(
                 ref_audio_input