Update acestep/resume_utils.py

acestep/resume_utils.py

@@ -1,1272 +1,325 @@
 """
-ACE-Step Inference API Module
-
-This module provides a standardized inference interface for music generation,
-designed for third-party integration. It offers both a simplified API and
-backward-compatible Gradio UI support.
+resume_utils.py
+================
+Save & Resume utility for long-running ACE-Step generations on HuggingFace
+ZeroGPU Spaces (120-second hard GPU timeout per call).
+
+DESIGN NOTE (important):
+This module is built for the LEGITIMATE resume scenario only:
+    - Same user, same Space, same account.
+    - Today's free ZeroGPU quota runs out before generation finishes.
+    - User comes back tomorrow (quota resets) and continues from where
+      they left off, using the SAME IP / SAME browser session.
+There is no IP-rotation, identity-spoofing, or quota-evasion logic here.
+All this code does is persist intermediate tensors to the Space's
+persistent disk (/data) under a task_id, so a later call (today or
+tomorrow, same account) can pick the work back up instead of starting
+from token 0 / step 0 again.
+
+Storage layout (under persistent_storage_path, e.g. /data/resume_state):
+    /data/resume_state/<task_id>/lm_tokens.pt        -> generated LM token ids (CPU tensor)
+    /data/resume_state/<task_id>/lm_meta.json        -> phase + params + timestamps
+    /data/resume_state/<task_id>/dit_latents.pt       -> latents at last checkpointed step
+    /data/resume_state/<task_id>/dit_meta.json        -> step index + DiT params + timestamps
+    /data/resume_state/<task_id>/audio_codes.json     -> final LM phase-2 output (string codes)
+
+task_id is short, human-typeable (e.g. "AB3K9F"), so a person can write it
+down and paste it into a textbox the next day.
 """
 
-import math
 import os
-import tempfile
-from typing import Optional, Union, List, Dict, Any, Tuple
-from dataclasses import dataclass, field, asdict
+import json
+import time
+import uuid
+import string
+import random
+from typing import Optional, Dict, Any
+
+import torch
 from loguru import logger
 
-from acestep.audio_utils import AudioSaver, generate_uuid_from_params
-from acestep.resume_utils import (
-    generate_task_id,
-    DEFAULT_TIMEOUT_SECONDS,
-    save_lm_audio_codes,
-    load_lm_audio_codes,
-)
-
-# HuggingFace Space environment detection
-IS_HUGGINGFACE_SPACE = os.environ.get("SPACE_ID") is not None
-
-
-@dataclass
-class GenerationParams:
-    """Configuration for music generation parameters.
-    
-    Attributes:
-        # Text Inputs
-        caption: A short text prompt describing the desired music (main prompt). < 512 characters
-        lyrics: Lyrics for the music. Use "[Instrumental]" for instrumental songs. < 4096 characters
-        instrumental: If True, generate instrumental music regardless of lyrics.
-        
-        # Music Metadata
-        bpm: BPM (beats per minute), e.g., 120. Set to None for automatic estimation. 30 ~ 300
-        keyscale: Musical key (e.g., "C Major", "Am"). Leave empty for auto-detection. A-G, #/♭, major/minor
-        timesignature: Time signature (2 for '2/4', 3 for '3/4', 4 for '4/4', 6 for '6/8'). Leave empty for auto-detection.
-        vocal_language: Language code for vocals, e.g., "en", "zh", "ja", or "unknown". see acestep/constants.py:VALID_LANGUAGES
-        duration: Target audio length in seconds. If <0 or None, model chooses automatically. 10 ~ 600
-        
-        # Generation Parameters
-        inference_steps: Number of diffusion steps (e.g., 8 for turbo, 32–100 for base model).
-        guidance_scale: CFG (classifier-free guidance) strength. Higher means following the prompt more strictly. Only support for non-turbo model.
-        seed: Integer seed for reproducibility. -1 means use random seed each time.
-        
-        # Advanced DiT Parameters
-        use_adg: Whether to use Adaptive Dual Guidance (only works for base model).
-        cfg_interval_start: Start ratio (0.0–1.0) to apply CFG.
-        cfg_interval_end: End ratio (0.0–1.0) to apply CFG.
-        shift: Timestep shift factor (default 1.0). When != 1.0, applies t = shift * t / (1 + (shift - 1) * t) to timesteps.
-        
-        # Task-Specific Parameters
-        task_type: Type of generation task. One of: "text2music", "cover", "repaint", "lego", "extract", "complete".
-        reference_audio: Path to a reference audio file for style transfer or cover tasks.
-        src_audio: Path to a source audio file for audio-to-audio tasks.
-        audio_codes: Audio semantic codes as a string (advanced use, for code-control generation).
-        repainting_start: For repaint/lego tasks: start time in seconds for region to repaint.
-        repainting_end: For repaint/lego tasks: end time in seconds for region to repaint (-1 for until end).
-        audio_cover_strength: Strength of reference audio/codes influence (range 0.0–1.0). set smaller (0.2) for style transfer tasks.
-        instruction: Optional task instruction prompt. If empty, auto-generated by system.
-        
-        # 5Hz Language Model Parameters for CoT reasoning
-        thinking: If True, enable 5Hz Language Model "Chain-of-Thought" reasoning for semantic/music metadata and codes.
-        lm_temperature: Sampling temperature for the LLM (0.0–2.0). Higher = more creative/varied results.
-        lm_cfg_scale: Classifier-free guidance scale for the LLM.
-        lm_top_k: LLM top-k sampling (0 = disabled).
-        lm_top_p: LLM top-p nucleus sampling (1.0 = disabled).
-        lm_negative_prompt: Negative prompt to use for LLM (for control).
-        use_cot_metas: Whether to let LLM generate music metadata via CoT reasoning.
-        use_cot_caption: Whether to let LLM rewrite or format the input caption via CoT reasoning.
-        use_cot_language: Whether to let LLM detect vocal language via CoT.
-    """
-    # Required Inputs
-    task_type: str = "text2music"
-    instruction: str = "Fill the audio semantic mask based on the given conditions:"
-
-    # Audio Uploads
-    reference_audio: Optional[str] = None
-    src_audio: Optional[str] = None
-
-    # LM Codes Hints
-    audio_codes: str = ""
-
-    # Text Inputs
-    caption: str = ""
-    lyrics: str = ""
-    instrumental: bool = False
-
-    # Metadata
-    vocal_language: str = "unknown"
-    bpm: Optional[int] = None
-    keyscale: str = ""
-    timesignature: str = ""
-    duration: float = -1.0
-
-    # Advanced Settings
-    inference_steps: int = 8
-    seed: int = -1
-    guidance_scale: float = 7.0
-    use_adg: bool = False
-    cfg_interval_start: float = 0.0
-    cfg_interval_end: float = 1.0
-    shift: float = 1.0
-    infer_method: str = "ode"  # "ode" or "sde" - diffusion inference method
-    # Custom timesteps (parsed from string like "0.97,0.76,0.615,0.5,0.395,0.28,0.18,0.085,0")
-    # If provided, overrides inference_steps and shift
-    timesteps: Optional[List[float]] = None
-
-    repainting_start: float = 0.0
-    repainting_end: float = -1
-    audio_cover_strength: float = 1.0
-
-    # 5Hz Language Model Parameters
-    thinking: bool = True
-    lm_temperature: float = 0.85
-    lm_cfg_scale: float = 2.0
-    lm_top_k: int = 0
-    lm_top_p: float = 0.9
-    lm_negative_prompt: str = "NO USER INPUT"
-    use_cot_metas: bool = True
-    use_cot_caption: bool = True
-    use_cot_lyrics: bool = False  # TODO: not used yet
-    use_cot_language: bool = True
-    use_constrained_decoding: bool = True
-
-    cot_bpm: Optional[int] = None
-    cot_keyscale: str = ""
-    cot_timesignature: str = ""
-    cot_duration: Optional[float] = None
-    cot_vocal_language: str = "unknown"
-    cot_caption: str = ""
-    cot_lyrics: str = ""
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert config to dictionary for JSON serialization."""
-        return asdict(self)
-
-
-@dataclass
-class GenerationConfig:
-    """Configuration for music generation.
-    
-    Attributes:
-        batch_size: Number of audio samples to generate
-        allow_lm_batch: Whether to allow batch processing in LM
-        use_random_seed: Whether to use random seed
-        seeds: Seed(s) for batch generation. Can be:
-            - None: Use random seeds (when use_random_seed=True) or params.seed (when use_random_seed=False)
-            - List[int]: List of seeds, will be padded with random seeds if fewer than batch_size
-            - int: Single seed value (will be converted to list and padded)
-        lm_batch_chunk_size: Batch chunk size for LM processing
-        constrained_decoding_debug: Whether to enable constrained decoding debug
-        audio_format: Output audio format, one of "mp3", "wav", "flac". Default: "flac"
-    """
-    batch_size: int = 2
-    allow_lm_batch: bool = False
-    use_random_seed: bool = True
-    seeds: Optional[List[int]] = None
-    lm_batch_chunk_size: int = 8
-    constrained_decoding_debug: bool = False
-    audio_format: str = "flac"  # Default to FLAC for fast saving
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert config to dictionary for JSON serialization."""
-        return asdict(self)
-
-
-@dataclass
-class GenerationResult:
-    """Result of music generation.
-    
-    Attributes:
-        # Audio Outputs
-        audios: List of audio dictionaries with paths, keys, params
-        status_message: Status message from generation
-        extra_outputs: Extra outputs from generation
-        success: Whether generation completed successfully
-        error: Error message if generation failed
-    """
 
-    # Audio Outputs
-    audios: List[Dict[str, Any]] = field(default_factory=list)
-    # Generation Information
-    status_message: str = ""
-    extra_outputs: Dict[str, Any] = field(default_factory=dict)
-    # Success Status
-    success: bool = True
-    error: Optional[str] = None
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert result to dictionary for JSON serialization."""
-        return asdict(self)
-
-
-@dataclass
-class UnderstandResult:
-    """Result of music understanding from audio codes.
-    
-    Attributes:
-        # Metadata Fields
-        caption: Generated caption describing the music
-        lyrics: Generated or extracted lyrics
-        bpm: Beats per minute (None if not detected)
-        duration: Duration in seconds (None if not detected)
-        keyscale: Musical key (e.g., "C Major")
-        language: Vocal language code (e.g., "en", "zh")
-        timesignature: Time signature (e.g., "4/4")
-        
-        # Status
-        status_message: Status message from understanding
-        success: Whether understanding completed successfully
-        error: Error message if understanding failed
-    """
-    # Metadata Fields
-    caption: str = ""
-    lyrics: str = ""
-    bpm: Optional[int] = None
-    duration: Optional[float] = None
-    keyscale: str = ""
-    language: str = ""
-    timesignature: str = ""
-    
-    # Status
-    status_message: str = ""
-    success: bool = True
-    error: Optional[str] = None
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert result to dictionary for JSON serialization."""
-        return asdict(self)
-
-
-def _update_metadata_from_lm(
-    metadata: Dict[str, Any],
-    bpm: Optional[int],
-    key_scale: str,
-    time_signature: str,
-    audio_duration: Optional[float],
-    vocal_language: str,
-    caption: str,
-    lyrics: str,
-) -> Tuple[Optional[int], str, str, Optional[float]]:
-    """Update metadata fields from LM output if not provided by user."""
-
-    if bpm is None and metadata.get('bpm'):
-        bpm_value = metadata.get('bpm')
-        if bpm_value not in ["N/A", ""]:
-            try:
-                bpm = int(bpm_value)
-            except (ValueError, TypeError):
-                pass
-
-    if not key_scale and metadata.get('keyscale'):
-        key_scale_value = metadata.get('keyscale', metadata.get('key_scale', ""))
-        if key_scale_value != "N/A":
-            key_scale = key_scale_value
-
-    if not time_signature and metadata.get('timesignature'):
-        time_signature_value = metadata.get('timesignature', metadata.get('time_signature', ""))
-        if time_signature_value != "N/A":
-            time_signature = time_signature_value
-
-    if audio_duration is None or audio_duration <= 0:
-        audio_duration_value = metadata.get('duration', -1)
-        if audio_duration_value not in ["N/A", ""]:
-            try:
-                audio_duration = float(audio_duration_value)
-            except (ValueError, TypeError):
-                pass
-
-    if not vocal_language and metadata.get('vocal_language'):
-        vocal_language = metadata.get('vocal_language')
-    if not caption and metadata.get('caption'):
-        caption = metadata.get('caption')
-    if not lyrics and metadata.get('lyrics'):
-        lyrics = metadata.get('lyrics')
-    return bpm, key_scale, time_signature, audio_duration, vocal_language, caption, lyrics
-
-
-def generate_music(
-    dit_handler,
-    llm_handler,
-    params: GenerationParams,
-    config: GenerationConfig,
-    save_dir: Optional[str] = None,
-    progress=None,
-    task_id: Optional[str] = None,
-    timeout: float = DEFAULT_TIMEOUT_SECONDS,
-) -> "GenerationResult":
-    """Generate music using ACE-Step model with optional LM reasoning.
-    
-    Args:
-        dit_handler: Initialized DiT model handler (AceStepHandler instance)
-        llm_handler: Initialized LLM handler (LLMHandler instance)
-        params: Generation parameters (GenerationParams instance)
-        config: Generation configuration (GenerationConfig instance)
-        task_id: Optional tracking code for Save & Resume across ZeroGPU's
-            120-second hard timeout. Pass None on a fresh request — a new
-            task_id will be generated automatically and returned in
-            extra_outputs["task_id"] on every call (whether it finished or
-            not), so the UI can always show it. Pass the SAME task_id back
-            in to continue a previously timed-out generation (works across
-            separate calls today after quota resets, or tomorrow — it just
-            needs to be the same Space / same persistent disk).
-        timeout: Seconds of GPU time before the LM phase checkpoints itself
-            and returns early. Keep a safety margin under ZeroGPU's 120s
-            hard kill (default 110s leaves 10s for the save + cleanup).
-
-    Returns:
-        GenerationResult with generated audio files and metadata. If the
-        LM phase timed out and saved a checkpoint, success=False and
-        extra_outputs["resume_pending"]=True with extra_outputs["task_id"]
-        set — show this code to the user so they can resume.
+# ZeroGPU gives ~120s per call. We checkpoint a safety margin before the
+# hard kill so the save itself (disk I/O) has time to complete.
+DEFAULT_TIMEOUT_SECONDS = 110
+
+# How long an unfinished task is kept on disk before we consider it stale
+# and eligible for cleanup. 7 days comfortably covers "finish tomorrow".
+MAX_TASK_AGE_SECONDS = 7 * 24 * 3600
+
+
+def _resume_root(persistent_storage_path: Optional[str]) -> str:
+    """Resolve the root directory for resume state.
+
+    Falls back to a local ./resume_state directory if no persistent path
+    is configured (e.g. local dev), so the code never crashes — it just
+    won't survive a process restart in that case.
     """
-    # Always have a task_id so we can checkpoint and report it, even if
-    # the caller didn't think to ask for resume support up front.
-    if not task_id:
-        task_id = generate_task_id()
+    base = persistent_storage_path or "."
+    root = os.path.join(base, "resume_state")
+    os.makedirs(root, exist_ok=True)
+    return root
 
-    try:
-        # Phase 1: LM-based metadata and code generation (if enabled)
-        audio_code_string_to_use = params.audio_codes
-        lm_generated_metadata = None
-        lm_generated_audio_codes_list = []
-        lm_total_time_costs = {
-            "phase1_time": 0.0,
-            "phase2_time": 0.0,
-            "total_time": 0.0,
-        }
-
-        # Extract mutable copies of metadata (will be updated by LM if needed)
-        bpm = params.bpm
-        key_scale = params.keyscale
-        time_signature = params.timesignature
-        audio_duration = params.duration
-        dit_input_caption = params.caption
-        dit_input_vocal_language = params.vocal_language
-        dit_input_lyrics = params.lyrics
-        # Determine if we need to generate audio codes
-        # If user has provided audio_codes, we don't need to generate them
-        # Otherwise, check if we need audio codes (lm_dit mode) or just metas (dit mode)
-        user_provided_audio_codes = bool(params.audio_codes and str(params.audio_codes).strip())
-
-        # Determine infer_type: use "llm_dit" if we need audio codes, "dit" if only metas needed
-        # For now, we use "llm_dit" if batch mode or if user hasn't provided codes
-        # Use "dit" if user has provided codes (only need metas) or if explicitly only need metas
-        # Note: This logic can be refined based on specific requirements
-        need_audio_codes = not user_provided_audio_codes
-
-        # Determine if we should use chunk-based LM generation (always use chunks for consistency)
-        # Determine actual batch size for chunk processing
-        actual_batch_size = config.batch_size if config.batch_size is not None else 1
-
-        # Prepare seeds for batch generation
-        # Use config.seed if provided, otherwise fallback to params.seed
-        # Convert config.seed (None, int, or List[int]) to format that prepare_seeds accepts
-        seed_for_generation = ""
-        if config.seeds is not None and len(config.seeds) > 0:
-            if isinstance(config.seeds, list):
-                # Convert List[int] to comma-separated string
-                seed_for_generation = ",".join(str(s) for s in config.seeds)
-
-        # Use dit_handler.prepare_seeds to handle seed list generation and padding
-        # This will handle all the logic: padding with random seeds if needed, etc.
-        actual_seed_list, _ = dit_handler.prepare_seeds(actual_batch_size, seed_for_generation, config.use_random_seed)
-
-        # LM-based Chain-of-Thought reasoning
-        # Skip LM for cover/repaint tasks - these tasks use reference/src audio directly
-        # and don't need LM to generate audio codes
-        skip_lm_tasks = {"cover", "repaint"}
-        
-        # Determine if we should use LLM
-        # LLM is needed for:
-        # 1. thinking=True: generate audio codes via LM
-        # 2. use_cot_caption=True: enhance/generate caption via CoT
-        # 3. use_cot_language=True: detect vocal language via CoT
-        # 4. use_cot_metas=True: fill missing metadata via CoT
-        need_lm_for_cot = params.use_cot_caption or params.use_cot_language or params.use_cot_metas
-        use_lm = (params.thinking or need_lm_for_cot) and llm_handler.llm_initialized and params.task_type not in skip_lm_tasks
-        lm_status = []
-        
-        if params.task_type in skip_lm_tasks:
-            logger.info(f"Skipping LM for task_type='{params.task_type}' - using DiT directly")
-        
-        logger.info(f"[generate_music] LLM usage decision: thinking={params.thinking}, "
-                   f"use_cot_caption={params.use_cot_caption}, use_cot_language={params.use_cot_language}, "
-                   f"use_cot_metas={params.use_cot_metas}, need_lm_for_cot={need_lm_for_cot}, "
-                   f"llm_initialized={llm_handler.llm_initialized if llm_handler else False}, use_lm={use_lm}")
-        
-        if use_lm:
-            # ---- Save & Resume: skip straight to DiT if this task_id already
-            # finished its LM phase in a previous (possibly timed-out) call ----
-            saved_codes = load_lm_audio_codes(llm_handler.persistent_storage_path, task_id) if task_id else None
-            if saved_codes is not None:
-                logger.info(f"[resume] Found completed LM phase for task {task_id}, skipping straight to DiT")
-                lm_generated_metadata = saved_codes["metadata"]
-                lm_generated_audio_codes_list = [saved_codes["audio_codes"]]
-                audio_code_string_to_use = saved_codes["audio_codes"]
-                bpm, key_scale, time_signature, audio_duration, dit_input_vocal_language, dit_input_caption, dit_input_lyrics = _update_metadata_from_lm(
-                    metadata=lm_generated_metadata,
-                    bpm=bpm,
-                    key_scale=key_scale,
-                    time_signature=time_signature,
-                    audio_duration=audio_duration,
-                    vocal_language=dit_input_vocal_language,
-                    caption=dit_input_caption,
-                    lyrics=dit_input_lyrics)
-                lm_status.append(f"✅ LM phase restored from tracking code {task_id} (resumed)")
-                use_lm = False  # Skip the LM block below entirely; fall through to Phase 2 (DiT)
-
-        if use_lm:
-            # Convert sampling parameters - handle None values safely
-            top_k_value = None if not params.lm_top_k or params.lm_top_k == 0 else int(params.lm_top_k)
-            top_p_value = None if not params.lm_top_p or params.lm_top_p >= 1.0 else params.lm_top_p
-
-            # Build user_metadata from user-provided values
-            user_metadata = {}
-            if bpm is not None:
-                try:
-                    bpm_value = float(bpm)
-                    if bpm_value > 0:
-                        user_metadata['bpm'] = int(bpm_value)
-                except (ValueError, TypeError):
-                    pass
-
-            if key_scale and key_scale.strip():
-                key_scale_clean = key_scale.strip()
-                if key_scale_clean.lower() not in ["n/a", ""]:
-                    user_metadata['keyscale'] = key_scale_clean
-
-            if time_signature and time_signature.strip():
-                time_sig_clean = time_signature.strip()
-                if time_sig_clean.lower() not in ["n/a", ""]:
-                    user_metadata['timesignature'] = time_sig_clean
-
-            if audio_duration is not None:
-                try:
-                    duration_value = float(audio_duration)
-                    if duration_value > 0:
-                        user_metadata['duration'] = int(duration_value)
-                except (ValueError, TypeError):
-                    pass
-
-            user_metadata_to_pass = user_metadata if user_metadata else None
-
-            # Determine infer_type based on whether we need audio codes
-            # - "llm_dit": generates both metas and audio codes (two-phase internally)
-            # - "dit": generates only metas (single phase)
-            infer_type = "llm_dit" if need_audio_codes and params.thinking else "dit"
-
-            # Use chunk size from config, or default to batch_size if not set
-            max_inference_batch_size = int(config.lm_batch_chunk_size) if config.lm_batch_chunk_size > 0 else actual_batch_size
-            num_chunks = math.ceil(actual_batch_size / max_inference_batch_size)
-
-            # Resume (task_id-based checkpointing) is only supported for a
-            # single chunk of size 1 — i.e. generating one song at a time,
-            # which is the normal "free tier" usage pattern this is built for.
-            resume_eligible = (num_chunks == 1 and actual_batch_size == 1)
-            if task_id and not resume_eligible:
-                logger.warning(f"[resume] task_id={task_id} given but batch/chunk size > 1 — resume disabled for this call")
-
-            all_metadata_list = []
-            all_audio_codes_list = []
-
-            for chunk_idx in range(num_chunks):
-                chunk_start = chunk_idx * max_inference_batch_size
-                chunk_end = min(chunk_start + max_inference_batch_size, actual_batch_size)
-                chunk_size = chunk_end - chunk_start
-                chunk_seeds = actual_seed_list[chunk_start:chunk_end] if chunk_start < len(actual_seed_list) else None
-
-                logger.info(f"LM chunk {chunk_idx+1}/{num_chunks} (infer_type={infer_type}) "
-                            f"(size: {chunk_size}, seeds: {chunk_seeds})")
-
-                # Use the determined infer_type
-                # - "llm_dit" will internally run two phases (metas + codes)
-                # - "dit" will only run phase 1 (metas only)
-                result = llm_handler.generate_with_stop_condition(
-                    caption=params.caption or "",
-                    lyrics=params.lyrics or "",
-                    infer_type=infer_type,
-                    temperature=params.lm_temperature,
-                    cfg_scale=params.lm_cfg_scale,
-                    negative_prompt=params.lm_negative_prompt,
-                    top_k=top_k_value,
-                    top_p=top_p_value,
-                    target_duration=audio_duration,  # Pass duration to limit audio codes generation
-                    user_metadata=user_metadata_to_pass,
-                    use_cot_caption=params.use_cot_caption,
-                    use_cot_language=params.use_cot_language,
-                    use_cot_metas=params.use_cot_metas,
-                    use_constrained_decoding=params.use_constrained_decoding,
-                    constrained_decoding_debug=config.constrained_decoding_debug,
-                    batch_size=chunk_size,
-                    seeds=chunk_seeds,
-                    progress=progress,
-                    task_id=task_id if resume_eligible else None,
-                    timeout=timeout,
-                )
-
-                # Save & Resume: LM phase ran out of GPU time and checkpointed.
-                # Surface the tracking code instead of treating this as a hard failure.
-                error_val = result.get("error") or ""
-                if isinstance(error_val, str) and error_val.startswith("TIMEOUT_SAVED"):
-                    return GenerationResult(
-                        audios=[],
-                        status_message=(
-                            f"⏱️ زمان رایگان (ZeroGPU) تمام شد، اما پیشرفت ذخیره شد.\n"
-                            f"کد پیگیری شما: {task_id}\n"
-                            f"این کد را ذخیره کنید. بعداً (امروز پس از ریست سهمیه، یا فردا) با همین کد ادامه دهید."
-                        ),
-                        extra_outputs={
-                            "task_id": task_id,
-                            "resume_pending": True,
-                            "resume_phase": result.get("extra_outputs", {}).get("resume_phase"),
-                            "time_costs": result.get("extra_outputs", {}).get("time_costs", {}),
-                        },
-                        success=False,
-                        error=f"TIMEOUT_SAVED:{task_id}",
-                    )
-
-                # Check if LM generation failed
-                if not result.get("success", False):
-                    error_msg = result.get("error", "Unknown LM error")
-                    lm_status.append(f"❌ LM Error: {error_msg}")
-                    # Return early with error
-                    return GenerationResult(
-                        audios=[],
-                        status_message=f"❌ LM generation failed: {error_msg}",
-                        extra_outputs={},
-                        success=False,
-                        error=error_msg,
-                    )
-
-                # Extract metadata and audio_codes from result dict
-                if chunk_size > 1:
-                    metadata_list = result.get("metadata", [])
-                    audio_codes_list = result.get("audio_codes", [])
-                    all_metadata_list.extend(metadata_list)
-                    all_audio_codes_list.extend(audio_codes_list)
-                else:
-                    metadata = result.get("metadata", {})
-                    audio_codes = result.get("audio_codes", "")
-                    all_metadata_list.append(metadata)
-                    all_audio_codes_list.append(audio_codes)
-
-                # Collect time costs from LM extra_outputs
-                lm_extra = result.get("extra_outputs", {})
-                lm_chunk_time_costs = lm_extra.get("time_costs", {})
-                if lm_chunk_time_costs:
-                    # Accumulate time costs from all chunks
-                    for key in ["phase1_time", "phase2_time", "total_time"]:
-                        if key in lm_chunk_time_costs:
-                            lm_total_time_costs[key] += lm_chunk_time_costs[key]
-
-                    time_str = ", ".join([f"{k}: {v:.2f}s" for k, v in lm_chunk_time_costs.items()])
-                    lm_status.append(f"✅ LM chunk {chunk_idx+1}: {time_str}")
-
-            lm_generated_metadata = all_metadata_list[0] if all_metadata_list else None
-            lm_generated_audio_codes_list = all_audio_codes_list
-
-            # Set audio_code_string_to_use based on infer_type
-            if infer_type == "llm_dit":
-                # If batch mode, use list; otherwise use single string
-                if actual_batch_size > 1:
-                    audio_code_string_to_use = all_audio_codes_list
-                else:
-                    audio_code_string_to_use = all_audio_codes_list[0] if all_audio_codes_list else ""
-            else:
-                # For "dit" mode, keep user-provided codes or empty
-                audio_code_string_to_use = params.audio_codes
-
-            # Update metadata from LM if not provided by user
-            if lm_generated_metadata:
-                bpm, key_scale, time_signature, audio_duration, vocal_language, caption, lyrics = _update_metadata_from_lm(
-                    metadata=lm_generated_metadata,
-                    bpm=bpm,
-                    key_scale=key_scale,
-                    time_signature=time_signature,
-                    audio_duration=audio_duration,
-                    vocal_language=dit_input_vocal_language,
-                    caption=dit_input_caption,
-                    lyrics=dit_input_lyrics)
-                if not params.bpm:
-                    params.cot_bpm = bpm
-                if not params.keyscale:
-                    params.cot_keyscale = key_scale
-                if not params.timesignature:
-                    params.cot_timesignature = time_signature
-                if not params.duration:
-                    params.cot_duration = audio_duration
-                if not params.vocal_language:
-                    params.cot_vocal_language = vocal_language
-                if not params.caption:
-                    params.cot_caption = caption
-                if not params.lyrics:
-                    params.cot_lyrics = lyrics
-
-            # set cot caption and language if needed
-            if params.use_cot_caption:
-                dit_input_caption = lm_generated_metadata.get("caption", dit_input_caption)
-            if params.use_cot_language:
-                dit_input_vocal_language = lm_generated_metadata.get("vocal_language", dit_input_vocal_language)
-
-        # Phase 2: DiT music generation
-        # Use seed_for_generation (from config.seed or params.seed) instead of params.seed for actual generation
-        result = dit_handler.generate_music(
-            captions=dit_input_caption,
-            lyrics=dit_input_lyrics,
-            bpm=bpm,
-            key_scale=key_scale,
-            time_signature=time_signature,
-            vocal_language=dit_input_vocal_language,
-            inference_steps=params.inference_steps,
-            guidance_scale=params.guidance_scale,
-            use_random_seed=config.use_random_seed,
-            seed=seed_for_generation,  # Use config.seed (or params.seed fallback) instead of params.seed directly
-            reference_audio=params.reference_audio,
-            audio_duration=audio_duration,
-            batch_size=config.batch_size if config.batch_size is not None else 1,
-            src_audio=params.src_audio,
-            audio_code_string=audio_code_string_to_use,
-            repainting_start=params.repainting_start,
-            repainting_end=params.repainting_end,
-            instruction=params.instruction,
-            audio_cover_strength=params.audio_cover_strength,
-            task_type=params.task_type,
-            use_adg=params.use_adg,
-            cfg_interval_start=params.cfg_interval_start,
-            cfg_interval_end=params.cfg_interval_end,
-            shift=params.shift,
-            infer_method=params.infer_method,
-            timesteps=params.timesteps,
-            progress=progress,
-        )
-
-        # Check if generation failed
-        if not result.get("success", False):
-            return GenerationResult(
-                audios=[],
-                status_message=result.get("status_message", ""),
-                extra_outputs={},
-                success=False,
-                error=result.get("error"),
-            )
-
-        # Extract results from dit_handler.generate_music dict
-        dit_audios = result.get("audios", [])
-        status_message = result.get("status_message", "")
-        dit_extra_outputs = result.get("extra_outputs", {})
-
-        # Use the seed list already prepared above (from config.seed or params.seed fallback)
-        # actual_seed_list was computed earlier using dit_handler.prepare_seeds
-        seed_list = actual_seed_list
-
-        # Get base params dictionary
-        base_params_dict = params.to_dict()
-
-        # Save audio files using AudioSaver (format from config)
-        audio_format = config.audio_format if config.audio_format else "flac"
-        audio_saver = AudioSaver(default_format=audio_format)
-
-        # Use handler's temp_dir for saving files
-        if save_dir is not None:
-            os.makedirs(save_dir, exist_ok=True)
-
-        # Build audios list for GenerationResult with params and save files
-        # Audio saving and UUID generation handled here, outside of handler
-        audios = []
-        for idx, dit_audio in enumerate(dit_audios):
-            # Create a copy of params dict for this audio
-            audio_params = base_params_dict.copy()
-
-            # Update audio-specific values
-            audio_params["seed"] = seed_list[idx] if idx < len(seed_list) else None
-
-            # Add audio codes if batch mode
-            if lm_generated_audio_codes_list and idx < len(lm_generated_audio_codes_list):
-                audio_params["audio_codes"] = lm_generated_audio_codes_list[idx]
-
-            # Get audio tensor and metadata
-            audio_tensor = dit_audio.get("tensor")
-            sample_rate = dit_audio.get("sample_rate", 48000)
-
-            # Generate UUID for this audio (moved from handler)
-            batch_seed = seed_list[idx] if idx < len(seed_list) else seed_list[0] if seed_list else -1
-            audio_code_str = lm_generated_audio_codes_list[idx] if (
-                lm_generated_audio_codes_list and idx < len(lm_generated_audio_codes_list)) else audio_code_string_to_use
-            if isinstance(audio_code_str, list):
-                audio_code_str = audio_code_str[idx] if idx < len(audio_code_str) else ""
-
-            audio_key = generate_uuid_from_params(audio_params)
-
-            # Save audio file (handled outside handler)
-            audio_path = None
-            if audio_tensor is not None and save_dir is not None:
-                try:
-                    audio_file = os.path.join(save_dir, f"{audio_key}.{audio_format}")
-                    audio_path = audio_saver.save_audio(audio_tensor,
-                                                        audio_file,
-                                                        sample_rate=sample_rate,
-                                                        format=audio_format,
-                                                        channels_first=True)
-                except Exception as e:
-                    logger.error(f"[generate_music] Failed to save audio file: {e}")
-                    audio_path = ""  # Fallback to empty path
-
-            audio_dict = {
-                "path": audio_path or "",  # File path (saved here, not in handler)
-                "tensor": audio_tensor,  # Audio tensor [channels, samples], CPU, float32
-                "key": audio_key,
-                "sample_rate": sample_rate,
-                "params": audio_params,
-            }
-
-            audios.append(audio_dict)
-
-        # Merge extra_outputs: include dit_extra_outputs (latents, masks) and add LM metadata
-        extra_outputs = dit_extra_outputs.copy()
-        extra_outputs["lm_metadata"] = lm_generated_metadata
-        extra_outputs["task_id"] = task_id
-
-        # Merge time_costs from both LM and DiT into a unified dictionary
-        unified_time_costs = {}
-
-        # Add LM time costs (if LM was used)
-        if use_lm and lm_total_time_costs:
-            for key, value in lm_total_time_costs.items():
-                unified_time_costs[f"lm_{key}"] = value
-
-        # Add DiT time costs (if available)
-        dit_time_costs = dit_extra_outputs.get("time_costs", {})
-        if dit_time_costs:
-            for key, value in dit_time_costs.items():
-                unified_time_costs[f"dit_{key}"] = value
-
-        # Calculate total pipeline time
-        if unified_time_costs:
-            lm_total = unified_time_costs.get("lm_total_time", 0.0)
-            dit_total = unified_time_costs.get("dit_total_time_cost", 0.0)
-            unified_time_costs["pipeline_total_time"] = lm_total + dit_total
-
-        # Update extra_outputs with unified time_costs
-        extra_outputs["time_costs"] = unified_time_costs
-
-        if lm_status:
-            status_message = "\n".join(lm_status) + "\n" + status_message
-        else:
-            status_message = status_message
-
-        # Save & Resume: generation finished successfully end-to-end, so the
-        # checkpoint files for this task_id are no longer needed.
-        if task_id:
-            try:
-                from acestep.resume_utils import save_final_marker, cleanup_task
-                final_audio_path = audios[0]["path"] if audios else None
-                save_final_marker(llm_handler.persistent_storage_path, task_id, audio_path=final_audio_path)
-                cleanup_task(llm_handler.persistent_storage_path, task_id)
-            except Exception as e:
-                logger.warning(f"[resume] Failed to clean up completed task {task_id}: {e}")
-
-        # Create and return GenerationResult
-        return GenerationResult(
-            audios=audios,
-            status_message=status_message,
-            extra_outputs=extra_outputs,
-            success=True,
-            error=None,
-        )
 
-    except Exception as e:
-        logger.exception("Music generation failed")
-        return GenerationResult(
-            audios=[],
-            status_message=f"Error: {str(e)}",
-            extra_outputs={},
-            success=False,
-            error=str(e),
-        )
-
-
-def understand_music(
-    llm_handler,
-    audio_codes: str,
-    temperature: float = 0.85,
-    top_k: Optional[int] = None,
-    top_p: Optional[float] = None,
-    repetition_penalty: float = 1.0,
-    use_constrained_decoding: bool = True,
-    constrained_decoding_debug: bool = False,
-) -> UnderstandResult:
-    """Understand music from audio codes using the 5Hz Language Model.
-    
-    This function analyzes audio semantic codes and generates metadata about the music,
-    including caption, lyrics, BPM, duration, key scale, language, and time signature.
-    
-    If audio_codes is empty or "NO USER INPUT", the LM will generate a sample example
-    instead of analyzing existing codes.
-    
-    Note: cfg_scale and negative_prompt are not supported in understand mode.
-    
-    Args:
-        llm_handler: Initialized LLM handler (LLMHandler instance)
-        audio_codes: String of audio code tokens (e.g., "<|audio_code_123|><|audio_code_456|>...")
-                     Use empty string or "NO USER INPUT" to generate a sample example.
-        temperature: Sampling temperature for generation (0.0-2.0). Higher = more creative.
-        top_k: Top-K sampling (None or 0 = disabled)
-        top_p: Top-P (nucleus) sampling (None or 1.0 = disabled)
-        repetition_penalty: Repetition penalty (1.0 = no penalty)
-        use_constrained_decoding: Whether to use FSM-based constrained decoding for metadata
-        constrained_decoding_debug: Whether to enable debug logging for constrained decoding
-        
-    Returns:
-        UnderstandResult with parsed metadata fields and status
-        
-    Example:
-        >>> result = understand_music(llm_handler, audio_codes="<|audio_code_123|>...")
-        >>> if result.success:
-        ...     print(f"Caption: {result.caption}")
-        ...     print(f"BPM: {result.bpm}")
-        ...     print(f"Lyrics: {result.lyrics}")
+def generate_task_id() -> str:
+    """Generate a short, human-typeable tracking code, e.g. 'AB3K9F'.
+
+    Avoids ambiguous characters (0/O, 1/I/L) so users can read it back
+    off a screen and retype it the next day without errors.
     """
-    # Check if LLM is initialized
-    if not llm_handler.llm_initialized:
-        return UnderstandResult(
-            status_message="5Hz LM not initialized. Please initialize it first.",
-            success=False,
-            error="LLM not initialized",
-        )
-    
-    # If codes are empty, use "NO USER INPUT" to generate a sample example
-    if not audio_codes or not audio_codes.strip():
-        audio_codes = "NO USER INPUT"
-    
+    alphabet = "ABCDEFGHJKMNPQRSTUVWXYZ23456789"
+    return "".join(random.choice(alphabet) for _ in range(6))
+
+
+def _task_dir(persistent_storage_path: Optional[str], task_id: str) -> str:
+    safe_id = "".join(c for c in task_id.strip().upper() if c.isalnum())
+    if not safe_id:
+        raise ValueError("Invalid task_id")
+    d = os.path.join(_resume_root(persistent_storage_path), safe_id)
+    os.makedirs(d, exist_ok=True)
+    return d
+
+
+def task_exists(persistent_storage_path: Optional[str], task_id: str) -> bool:
     try:
-        # Call LLM understanding
-        metadata, status = llm_handler.understand_audio_from_codes(
-            audio_codes=audio_codes,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p,
-            repetition_penalty=repetition_penalty,
-            use_constrained_decoding=use_constrained_decoding,
-            constrained_decoding_debug=constrained_decoding_debug,
-        )
-        
-        # Check if LLM returned empty metadata (error case)
-        if not metadata:
-            return UnderstandResult(
-                status_message=status or "Failed to understand audio codes",
-                success=False,
-                error=status or "Empty metadata returned",
-            )
-        
-        # Extract and convert fields
-        caption = metadata.get('caption', '')
-        lyrics = metadata.get('lyrics', '')
-        keyscale = metadata.get('keyscale', '')
-        language = metadata.get('language', metadata.get('vocal_language', ''))
-        timesignature = metadata.get('timesignature', '')
-        
-        # Convert BPM to int
-        bpm = None
-        bpm_value = metadata.get('bpm')
-        if bpm_value is not None and bpm_value != 'N/A' and bpm_value != '':
-            try:
-                bpm = int(bpm_value)
-            except (ValueError, TypeError):
-                pass
-        
-        # Convert duration to float
-        duration = None
-        duration_value = metadata.get('duration')
-        if duration_value is not None and duration_value != 'N/A' and duration_value != '':
-            try:
-                duration = float(duration_value)
-            except (ValueError, TypeError):
-                pass
-        
-        # Clean up N/A values
-        if keyscale == 'N/A':
-            keyscale = ''
-        if language == 'N/A':
-            language = ''
-        if timesignature == 'N/A':
-            timesignature = ''
-        
-        return UnderstandResult(
-            caption=caption,
-            lyrics=lyrics,
-            bpm=bpm,
-            duration=duration,
-            keyscale=keyscale,
-            language=language,
-            timesignature=timesignature,
-            status_message=status,
-            success=True,
-            error=None,
-        )
-        
+        d = os.path.join(_resume_root(persistent_storage_path), task_id.strip().upper())
+        return os.path.isdir(d)
+    except Exception:
+        return False
+
+
+def get_task_status(persistent_storage_path: Optional[str], task_id: str) -> Dict[str, Any]:
+    """Return a human-readable status dict for a task_id, or {'found': False}."""
+    if not task_exists(persistent_storage_path, task_id):
+        return {"found": False}
+
+    d = _task_dir(persistent_storage_path, task_id)
+    status = {"found": True, "task_id": task_id.strip().upper(), "phase": "unknown"}
+
+    lm_meta_path = os.path.join(d, "lm_meta.json")
+    dit_meta_path = os.path.join(d, "dit_meta.json")
+    codes_path = os.path.join(d, "audio_codes.json")
+    final_audio_meta_path = os.path.join(d, "final_meta.json")
+
+    if os.path.exists(final_audio_meta_path):
+        with open(final_audio_meta_path, "r", encoding="utf-8") as f:
+            status.update(json.load(f))
+        status["phase"] = "complete"
+    elif os.path.exists(dit_meta_path):
+        with open(dit_meta_path, "r", encoding="utf-8") as f:
+            dit_meta = json.load(f)
+        status["phase"] = "dit_in_progress"
+        status["dit_step"] = dit_meta.get("current_step_idx")
+        status["dit_total_steps"] = dit_meta.get("total_steps")
+        status["saved_at"] = dit_meta.get("saved_at")
+    elif os.path.exists(codes_path):
+        status["phase"] = "lm_complete_dit_pending"
+    elif os.path.exists(lm_meta_path):
+        with open(lm_meta_path, "r", encoding="utf-8") as f:
+            lm_meta = json.load(f)
+        status["phase"] = "lm_in_progress"
+        status["lm_step"] = lm_meta.get("step")
+        status["saved_at"] = lm_meta.get("saved_at")
+
+    return status
+
+
+def cleanup_task(persistent_storage_path: Optional[str], task_id: str) -> None:
+    """Remove all state for a finished/abandoned task."""
+    import shutil
+    try:
+        d = os.path.join(_resume_root(persistent_storage_path), task_id.strip().upper())
+        if os.path.isdir(d):
+            shutil.rmtree(d)
+            logger.info(f"[resume_utils] Cleaned up task {task_id}")
     except Exception as e:
-        logger.exception("Music understanding failed")
-        return UnderstandResult(
-            status_message=f"Error: {str(e)}",
-            success=False,
-            error=str(e),
-        )
-
-
-@dataclass
-class CreateSampleResult:
-    """Result of creating a music sample from a natural language query.
-    
-    This is used by the "Simple Mode" / "Inspiration Mode" feature where users
-    provide a natural language description and the LLM generates a complete
-    sample with caption, lyrics, and metadata.
-    
-    Attributes:
-        # Metadata Fields
-        caption: Generated detailed music description/caption
-        lyrics: Generated lyrics (or "[Instrumental]" for instrumental music)
-        bpm: Beats per minute (None if not generated)
-        duration: Duration in seconds (None if not generated)
-        keyscale: Musical key (e.g., "C Major")
-        language: Vocal language code (e.g., "en", "zh")
-        timesignature: Time signature (e.g., "4")
-        instrumental: Whether this is an instrumental piece
-        
-        # Status
-        status_message: Status message from sample creation
-        success: Whether sample creation completed successfully
-        error: Error message if sample creation failed
-    """
-    # Metadata Fields
-    caption: str = ""
-    lyrics: str = ""
-    bpm: Optional[int] = None
-    duration: Optional[float] = None
-    keyscale: str = ""
-    language: str = ""
-    timesignature: str = ""
-    instrumental: bool = False
-    
-    # Status
-    status_message: str = ""
-    success: bool = True
-    error: Optional[str] = None
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert result to dictionary for JSON serialization."""
-        return asdict(self)
-
-
-def create_sample(
-    llm_handler,
-    query: str,
-    instrumental: bool = False,
-    vocal_language: Optional[str] = None,
-    temperature: float = 0.85,
-    top_k: Optional[int] = None,
-    top_p: Optional[float] = None,
-    repetition_penalty: float = 1.0,
-    use_constrained_decoding: bool = True,
-    constrained_decoding_debug: bool = False,
-) -> CreateSampleResult:
-    """Create a music sample from a natural language query using the 5Hz Language Model.
-    
-    This is the "Simple Mode" / "Inspiration Mode" feature that takes a user's natural
-    language description of music and generates a complete sample including:
-    - Detailed caption/description
-    - Lyrics (unless instrumental)
-    - Metadata (BPM, duration, key, language, time signature)
-    
-    Note: cfg_scale and negative_prompt are not supported in create_sample mode.
-    
-    Args:
-        llm_handler: Initialized LLM handler (LLMHandler instance)
-        query: User's natural language music description (e.g., "a soft Bengali love song")
-        instrumental: Whether to generate instrumental music (no vocals)
-        vocal_language: Allowed vocal language for constrained decoding (e.g., "en", "zh").
-                       If provided, the model will be constrained to generate lyrics in this language.
-                       If None or "unknown", no language constraint is applied.
-        temperature: Sampling temperature for generation (0.0-2.0). Higher = more creative.
-        top_k: Top-K sampling (None or 0 = disabled)
-        top_p: Top-P (nucleus) sampling (None or 1.0 = disabled)
-        repetition_penalty: Repetition penalty (1.0 = no penalty)
-        use_constrained_decoding: Whether to use FSM-based constrained decoding
-        constrained_decoding_debug: Whether to enable debug logging
-        
-    Returns:
-        CreateSampleResult with generated sample fields and status
-        
-    Example:
-        >>> result = create_sample(llm_handler, "a soft Bengali love song for a quiet evening", vocal_language="bn")
-        >>> if result.success:
-        ...     print(f"Caption: {result.caption}")
-        ...     print(f"Lyrics: {result.lyrics}")
-        ...     print(f"BPM: {result.bpm}")
+        logger.warning(f"[resume_utils] Failed to clean up task {task_id}: {e}")
+
+
+def cleanup_stale_tasks(persistent_storage_path: Optional[str], max_age_seconds: int = MAX_TASK_AGE_SECONDS) -> int:
+    """Delete tasks older than max_age_seconds. Call this occasionally (e.g. on app startup).
+
+    Returns the number of tasks removed.
     """
-    import torch
-    # Debug logging for ZeroGPU diagnosis
-    logger.info(f"[create_sample Debug] Entry: IS_HUGGINGFACE_SPACE={IS_HUGGINGFACE_SPACE}")
-    logger.info(f"[create_sample Debug] torch.cuda.is_available()={torch.cuda.is_available()}")
-    if torch.cuda.is_available():
-        logger.info(f"[create_sample Debug] torch.cuda.current_device()={torch.cuda.current_device()}")
-    logger.info(f"[create_sample Debug] llm_handler.device={llm_handler.device}, llm_handler.offload_to_cpu={llm_handler.offload_to_cpu}")
-    if llm_handler.llm is not None:
-        try:
-            logger.info(f"[create_sample Debug] Model device: {next(llm_handler.llm.parameters()).device}")
-        except Exception as e:
-            logger.info(f"[create_sample Debug] Could not get model device: {e}")
-
-    # Check if LLM is initialized
-    if not llm_handler.llm_initialized:
-        return CreateSampleResult(
-            status_message="5Hz LM not initialized. Please initialize it first.",
-            success=False,
-            error="LLM not initialized",
-        )
-    
+    import shutil
+    root = _resume_root(persistent_storage_path)
+    removed = 0
+    now = time.time()
     try:
-        # Call LLM to create sample
-        metadata, status = llm_handler.create_sample_from_query(
-            query=query,
-            instrumental=instrumental,
-            vocal_language=vocal_language,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p,
-            repetition_penalty=repetition_penalty,
-            use_constrained_decoding=use_constrained_decoding,
-            constrained_decoding_debug=constrained_decoding_debug,
-        )
-        
-        # Check if LLM returned empty metadata (error case)
-        if not metadata:
-            return CreateSampleResult(
-                status_message=status or "Failed to create sample",
-                success=False,
-                error=status or "Empty metadata returned",
-            )
-        
-        # Extract and convert fields
-        caption = metadata.get('caption', '')
-        lyrics = metadata.get('lyrics', '')
-        keyscale = metadata.get('keyscale', '')
-        language = metadata.get('language', metadata.get('vocal_language', ''))
-        timesignature = metadata.get('timesignature', '')
-        is_instrumental = metadata.get('instrumental', instrumental)
-        
-        # Convert BPM to int
-        bpm = None
-        bpm_value = metadata.get('bpm')
-        if bpm_value is not None and bpm_value != 'N/A' and bpm_value != '':
-            try:
-                bpm = int(bpm_value)
-            except (ValueError, TypeError):
-                pass
-        
-        # Convert duration to float
-        duration = None
-        duration_value = metadata.get('duration')
-        if duration_value is not None and duration_value != 'N/A' and duration_value != '':
+        for name in os.listdir(root):
+            d = os.path.join(root, name)
+            if not os.path.isdir(d):
+                continue
             try:
-                duration = float(duration_value)
-            except (ValueError, TypeError):
-                pass
-        
-        # Clean up N/A values
-        if keyscale == 'N/A':
-            keyscale = ''
-        if language == 'N/A':
-            language = ''
-        if timesignature == 'N/A':
-            timesignature = ''
-        
-        return CreateSampleResult(
-            caption=caption,
-            lyrics=lyrics,
-            bpm=bpm,
-            duration=duration,
-            keyscale=keyscale,
-            language=language,
-            timesignature=timesignature,
-            instrumental=is_instrumental,
-            status_message=status,
-            success=True,
-            error=None,
-        )
-        
-    except Exception as e:
-        logger.exception("Sample creation failed")
-        return CreateSampleResult(
-            status_message=f"Error: {str(e)}",
-            success=False,
-            error=str(e),
-        )
-
-
-@dataclass
-class FormatSampleResult:
-    """Result of formatting user-provided caption and lyrics.
-    
-    This is used by the "Format" feature where users provide caption and lyrics,
-    and the LLM formats them into structured music metadata and an enhanced description.
-    
-    Attributes:
-        # Metadata Fields
-        caption: Enhanced/formatted music description/caption
-        lyrics: Formatted lyrics (may be same as input or reformatted)
-        bpm: Beats per minute (None if not detected)
-        duration: Duration in seconds (None if not detected)
-        keyscale: Musical key (e.g., "C Major")
-        language: Vocal language code (e.g., "en", "zh")
-        timesignature: Time signature (e.g., "4")
-        
-        # Status
-        status_message: Status message from formatting
-        success: Whether formatting completed successfully
-        error: Error message if formatting failed
+                mtime = os.path.getmtime(d)
+                if now - mtime > max_age_seconds:
+                    shutil.rmtree(d)
+                    removed += 1
+            except Exception:
+                continue
+    except FileNotFoundError:
+        pass
+    if removed:
+        logger.info(f"[resume_utils] Cleaned up {removed} stale task(s)")
+    return removed
+
+
+# ---------------------------------------------------------------------------
+# LM phase (token-by-token generation) save/load
+# ---------------------------------------------------------------------------
+
+def save_lm_checkpoint(
+    persistent_storage_path: Optional[str],
+    task_id: str,
+    generated_ids: torch.Tensor,
+    step: int,
+    extra: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Persist LM token generation state to disk."""
+    d = _task_dir(persistent_storage_path, task_id)
+    torch.save(generated_ids.detach().to("cpu"), os.path.join(d, "lm_tokens.pt"))
+
+    meta = {
+        "step": step,
+        "saved_at": time.time(),
+        "shape": list(generated_ids.shape),
+    }
+    if extra:
+        meta.update(extra)
+
+    with open(os.path.join(d, "lm_meta.json"), "w", encoding="utf-8") as f:
+        json.dump(meta, f, ensure_ascii=False, indent=2)
+
+    logger.info(f"[resume_utils] LM checkpoint saved for task {task_id} at step {step}")
+
+
+def load_lm_checkpoint(
+    persistent_storage_path: Optional[str],
+    task_id: str,
+    device: str = "cuda",
+) -> Optional[Dict[str, Any]]:
+    """Load previously saved LM token state, or None if not found."""
+    d = os.path.join(_resume_root(persistent_storage_path), task_id.strip().upper())
+    tokens_path = os.path.join(d, "lm_tokens.pt")
+    meta_path = os.path.join(d, "lm_meta.json")
+
+    if not (os.path.exists(tokens_path) and os.path.exists(meta_path)):
+        return None
+
+    generated_ids = torch.load(tokens_path, map_location=device)
+    with open(meta_path, "r", encoding="utf-8") as f:
+        meta = json.load(f)
+
+    logger.info(f"[resume_utils] LM checkpoint restored for task {task_id} from step {meta.get('step')}")
+    return {"generated_ids": generated_ids, "meta": meta}
+
+
+def save_lm_audio_codes(
+    persistent_storage_path: Optional[str],
+    task_id: str,
+    metadata: Dict[str, Any],
+    audio_codes: str,
+) -> None:
+    """Persist the final (completed) LM phase output: metadata + audio codes string.
+
+    Once this is saved, the LM phase is fully done — the DiT phase can be
+    resumed/started independently using this file, even in a brand-new
+    call (today or tomorrow).
     """
-    # Metadata Fields
-    caption: str = ""
-    lyrics: str = ""
-    bpm: Optional[int] = None
-    duration: Optional[float] = None
-    keyscale: str = ""
-    language: str = ""
-    timesignature: str = ""
-    
-    # Status
-    status_message: str = ""
-    success: bool = True
-    error: Optional[str] = None
-
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert result to dictionary for JSON serialization."""
-        return asdict(self)
-
-
-def format_sample(
-    llm_handler,
-    caption: str,
-    lyrics: str,
-    user_metadata: Optional[Dict[str, Any]] = None,
-    temperature: float = 0.85,
-    top_k: Optional[int] = None,
-    top_p: Optional[float] = None,
-    repetition_penalty: float = 1.0,
-    use_constrained_decoding: bool = True,
-    constrained_decoding_debug: bool = False,
-) -> FormatSampleResult:
-    """Format user-provided caption and lyrics using the 5Hz Language Model.
-    
-    This function takes user input (caption and lyrics) and generates structured
-    music metadata including an enhanced caption, BPM, duration, key, language,
-    and time signature.
-    
-    If user_metadata is provided, those values will be used to constrain the
-    decoding, ensuring the output matches user-specified values.
-    
-    Note: cfg_scale and negative_prompt are not supported in format mode.
-    
-    Args:
-        llm_handler: Initialized LLM handler (LLMHandler instance)
-        caption: User's caption/description (e.g., "Latin pop, reggaeton")
-        lyrics: User's lyrics with structure tags
-        user_metadata: Optional dict with user-provided metadata to constrain decoding.
-                      Supported keys: bpm, duration, keyscale, timesignature, language
-        temperature: Sampling temperature for generation (0.0-2.0). Higher = more creative.
-        top_k: Top-K sampling (None or 0 = disabled)
-        top_p: Top-P (nucleus) sampling (None or 1.0 = disabled)
-        repetition_penalty: Repetition penalty (1.0 = no penalty)
-        use_constrained_decoding: Whether to use FSM-based constrained decoding for metadata
-        constrained_decoding_debug: Whether to enable debug logging for constrained decoding
-        
-    Returns:
-        FormatSampleResult with formatted metadata fields and status
-        
-    Example:
-        >>> result = format_sample(llm_handler, "Latin pop, reggaeton", "[Verse 1]\\nHola mundo...")
-        >>> if result.success:
-        ...     print(f"Caption: {result.caption}")
-        ...     print(f"BPM: {result.bpm}")
-        ...     print(f"Lyrics: {result.lyrics}")
+    d = _task_dir(persistent_storage_path, task_id)
+    with open(os.path.join(d, "audio_codes.json"), "w", encoding="utf-8") as f:
+        json.dump({"metadata": metadata, "audio_codes": audio_codes, "saved_at": time.time()}, f, ensure_ascii=False, indent=2)
+    logger.info(f"[resume_utils] LM phase output (codes) saved for task {task_id}")
+
+
+def load_lm_audio_codes(persistent_storage_path: Optional[str], task_id: str) -> Optional[Dict[str, Any]]:
+    d = os.path.join(_resume_root(persistent_storage_path), task_id.strip().upper())
+    path = os.path.join(d, "audio_codes.json")
+    if not os.path.exists(path):
+        return None
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+# ---------------------------------------------------------------------------
+# DiT phase (diffusion denoise loop) save/load
+# ---------------------------------------------------------------------------
+
+def save_dit_checkpoint(
+    persistent_storage_path: Optional[str],
+    task_id: str,
+    latents: torch.Tensor,
+    current_step_idx: int,
+    total_steps: int,
+    extra: Optional[Dict[str, Any]] = None,
+) -> None:
+    """Persist DiT denoising state (latents + step index) to disk."""
+    d = _task_dir(persistent_storage_path, task_id)
+    torch.save(latents.detach().to("cpu"), os.path.join(d, "dit_latents.pt"))
+
+    meta = {
+        "current_step_idx": current_step_idx,
+        "total_steps": total_steps,
+        "saved_at": time.time(),
+        "shape": list(latents.shape),
+    }
+    if extra:
+        meta.update(extra)
+
+    with open(os.path.join(d, "dit_meta.json"), "w", encoding="utf-8") as f:
+        json.dump(meta, f, ensure_ascii=False, indent=2)
+
+    logger.info(f"[resume_utils] DiT checkpoint saved for task {task_id} at step {current_step_idx}/{total_steps}")
+
+
+def load_dit_checkpoint(
+    persistent_storage_path: Optional[str],
+    task_id: str,
+    device: str = "cuda",
+) -> Optional[Dict[str, Any]]:
+    """Load previously saved DiT latent state, or None if not found."""
+    d = os.path.join(_resume_root(persistent_storage_path), task_id.strip().upper())
+    latents_path = os.path.join(d, "dit_latents.pt")
+    meta_path = os.path.join(d, "dit_meta.json")
+
+    if not (os.path.exists(latents_path) and os.path.exists(meta_path)):
+        return None
+
+    latents = torch.load(latents_path, map_location=device)
+    with open(meta_path, "r", encoding="utf-8") as f:
+        meta = json.load(f)
+
+    logger.info(f"[resume_utils] DiT checkpoint restored for task {task_id} from step {meta.get('current_step_idx')}")
+    return {"latents": latents, "meta": meta}
+
+
+def save_final_marker(persistent_storage_path: Optional[str], task_id: str, audio_path: Optional[str] = None) -> None:
+    """Mark a task as fully complete (audio rendered). Useful for status display."""
+    d = _task_dir(persistent_storage_path, task_id)
+    with open(os.path.join(d, "final_meta.json"), "w", encoding="utf-8") as f:
+        json.dump({"completed_at": time.time(), "audio_path": audio_path}, f, ensure_ascii=False, indent=2)
+
+
+class GenerationTimer:
+    """Small helper to check elapsed time against the ZeroGPU budget.
+
+    Usage:
+        timer = GenerationTimer(timeout=110)
+        for step in ...:
+            if timer.expired():
+                # save checkpoint and return early
+                break
     """
-    # Check if LLM is initialized
-    if not llm_handler.llm_initialized:
-        return FormatSampleResult(
-            status_message="5Hz LM not initialized. Please initialize it first.",
-            success=False,
-            error="LLM not initialized",
-        )
-    
-    try:
-        # Call LLM formatting
-        metadata, status = llm_handler.format_sample_from_input(
-            caption=caption,
-            lyrics=lyrics,
-            user_metadata=user_metadata,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p,
-            repetition_penalty=repetition_penalty,
-            use_constrained_decoding=use_constrained_decoding,
-            constrained_decoding_debug=constrained_decoding_debug,
-        )
-        
-        # Check if LLM returned empty metadata (error case)
-        if not metadata:
-            return FormatSampleResult(
-                status_message=status or "Failed to format input",
-                success=False,
-                error=status or "Empty metadata returned",
-            )
-        
-        # Extract and convert fields
-        result_caption = metadata.get('caption', '')
-        result_lyrics = metadata.get('lyrics', lyrics)  # Fall back to input lyrics
-        keyscale = metadata.get('keyscale', '')
-        language = metadata.get('language', metadata.get('vocal_language', ''))
-        timesignature = metadata.get('timesignature', '')
-        
-        # Convert BPM to int
-        bpm = None
-        bpm_value = metadata.get('bpm')
-        if bpm_value is not None and bpm_value != 'N/A' and bpm_value != '':
-            try:
-                bpm = int(bpm_value)
-            except (ValueError, TypeError):
-                pass
-        
-        # Convert duration to float
-        duration = None
-        duration_value = metadata.get('duration')
-        if duration_value is not None and duration_value != 'N/A' and duration_value != '':
-            try:
-                duration = float(duration_value)
-            except (ValueError, TypeError):
-                pass
-        
-        # Clean up N/A values
-        if keyscale == 'N/A':
-            keyscale = ''
-        if language == 'N/A':
-            language = ''
-        if timesignature == 'N/A':
-            timesignature = ''
-        
-        return FormatSampleResult(
-            caption=result_caption,
-            lyrics=result_lyrics,
-            bpm=bpm,
-            duration=duration,
-            keyscale=keyscale,
-            language=language,
-            timesignature=timesignature,
-            status_message=status,
-            success=True,
-            error=None,
-        )
-        
-    except Exception as e:
-        logger.exception("Format sample failed")
-        return FormatSampleResult(
-            status_message=f"Error: {str(e)}",
-            success=False,
-            error=str(e),
-        )
+
+    def __init__(self, timeout: float = DEFAULT_TIMEOUT_SECONDS):
+        self.start_time = time.time()
+        self.timeout = timeout
+
+    def elapsed(self) -> float:
+        return time.time() - self.start_time
+
+    def expired(self) -> bool:
+        return self.elapsed() > self.timeout
+
+    def remaining(self) -> float:
+        return max(0.0, self.timeout - self.elapsed())