"""Sangeet Studio — Suno-style song generation + instant artist-voice cloning.

Chain:
  1. ACE-Step 1.5 turbo (2.4B DiT, MIT) with RagaLoRA (Indian-music LoRA) generates
     a full song with vocals (Hindi/Urdu/Punjabi/Tamil/... lyrics supported).
  2. Demucs (htdemucs) splits the song into vocals + instrumental stems.
  3. Seed-VC (SVC mode) converts the vocal stem into the user's cloned artist
     voice from a 5-25s clean reference clip.
  4. Converted vocals are remixed with the instrumental stem.

Licenses: ACE-Step MIT; RagaLoRA Apache-2.0; Seed-VC inference code is GPL-3.0
(adapted from github.com/Plachtaa/seed-vc app_svc.py) — this Space is therefore GPL-3.0.
"""

import os
import shutil
import shlex
import subprocess
import sys
import tempfile
import threading
import time

print("[startup] importing numpy/torch…", flush=True)
import numpy as np
import torch
print("[startup] torch ready — importing gradio/soundfile…", flush=True)
import gradio as gr
import soundfile as sf
print("[startup] imports done", flush=True)

# ---------------------------------------------------------------------------
# Repo setup: clone ACE-Step-1.5 and seed-vc, put them on sys.path
# ---------------------------------------------------------------------------
APP_DIR = os.path.dirname(os.path.abspath(__file__))
REPOS_DIR = os.path.join(APP_DIR, "repos")
ACE_DIR = os.path.join(REPOS_DIR, "ACE-Step-1.5")
SEEDVC_DIR = os.path.join(REPOS_DIR, "seed-vc")
CHECKPOINTS_DIR = os.path.join(ACE_DIR, "checkpoints")
os.environ["ACESTEP_CHECKPOINTS_DIR"] = CHECKPOINTS_DIR
os.makedirs(CHECKPOINTS_DIR, exist_ok=True)


def _clone(url: str, dest: str, expect_dir: str) -> None:
    """Clone a repo, then VERIFY the expected package directory actually landed.
    A proxied/failed git transfer can exit 0 with an incomplete checkout, which
    produced "No module named 'acestep'" at import time. Fall back to the
    GitHub tarball (different download route) if verification fails."""
    if not os.path.isdir(dest):
        print(f"[startup] cloning {url}…", flush=True)
        subprocess.run(
            ["git", "clone", "--depth", "1", url, dest],
            check=True, timeout=300,
        )
    if not os.path.isdir(os.path.join(dest, expect_dir)):
        print(
            f"[startup] clone incomplete ('{expect_dir}' missing) — retrying via tarball…",
            flush=True,
        )
        shutil.rmtree(dest, ignore_errors=True)
        os.makedirs(dest, exist_ok=True)
        owner_repo = url.replace("https://github.com/", "").replace(".git", "")
        cmd = (
            f"curl -fsSL https://codeload.github.com/{owner_repo}/tar.gz/HEAD "
            f"| tar -xz -C {shlex.quote(dest)} --strip-components=1"
        )
        subprocess.run(cmd, shell=True, check=True, timeout=300)
    if not os.path.isdir(os.path.join(dest, expect_dir)):
        raise RuntimeError(
            f"{dest} is missing '{expect_dir}' even after tarball fallback — "
            "network to GitHub is broken from this container."
        )
    print(f"[startup] repo ready: {dest}", flush=True)


# Repo cloning + checkpoint download happen in the background _prepare_env()
# thread (defined below) so the Gradio port always binds immediately.

# ---------------------------------------------------------------------------
# Model downloads (network + CPU only here; safe at Space startup)
# ---------------------------------------------------------------------------
_ACE_READY = threading.Event()
_ENV_READY = threading.Event()


def _download_acestep_models() -> None:
    try:
        for _p in (ACE_DIR, SEEDVC_DIR):
            if _p not in sys.path:
                sys.path.insert(0, _p)
        import importlib
        importlib.invalidate_caches()
        print(
            f"[startup] acestep package present: "
            f"{os.path.isdir(os.path.join(ACE_DIR, 'acestep'))}", flush=True,
        )
        from acestep.model_downloader import ensure_main_model, ensure_lm_model
        ok, msg = ensure_main_model(CHECKPOINTS_DIR)
        print(f"[startup] main model: {msg}", flush=True)
        ok, msg = ensure_lm_model("acestep-5Hz-lm-1.7B", CHECKPOINTS_DIR)
        print(f"[startup] LM model: {msg}", flush=True)
        _ACE_READY.set()
    except Exception as e:
        print(f"[startup] ACE-Step checkpoint download failed: {e}", flush=True)
        # Unblock waiting pipelines so they fail fast with a visible error
        _ENV_READY.set()
        _ACE_READY.set()


def _prepare_env() -> None:
    """Clone repos, extend sys.path, then download checkpoints — all in the
    background so the Gradio port binds immediately. Pipelines wait on the
    _ENV_READY / _ACE_READY events."""
    try:
        print("[startup] cloning ACE-Step and seed-vc repos…", flush=True)
        _clone("https://github.com/ACE-Step/ACE-Step-1.5.git", ACE_DIR, "acestep")
        _clone("https://github.com/Plachtaa/seed-vc.git", SEEDVC_DIR, "modules")
        for _p in (ACE_DIR, SEEDVC_DIR):
            if _p not in sys.path:
                sys.path.insert(0, _p)
        import importlib
        importlib.invalidate_caches()
        _ENV_READY.set()
        print("[startup] repos ready — downloading checkpoints…", flush=True)
        _download_acestep_models()
    except Exception as e:
        print(f"[startup] environment preparation failed: {e}", flush=True)
        _ENV_READY.set()
        _ACE_READY.set()


# Everything slow (repo clones + ~10 GB checkpoint download) happens here, in the
# background, so launch() below is reached within seconds no matter what. The
# first generation waits on the events.
threading.Thread(target=_prepare_env, daemon=True).start()

# ---------------------------------------------------------------------------
# Lazy model state
# ---------------------------------------------------------------------------
_state = {}


def _load_ace(device: str):
    if not _ENV_READY.wait(timeout=600) or not _ACE_READY.wait(timeout=1800):
        raise RuntimeError(
            "Model download did not finish in time — check the Space runtime logs and restart."
        )
    if _state.get("ace_ok"):
        return _state["dit"], _state["llm"]

    for _p in (ACE_DIR, SEEDVC_DIR):
        if _p not in sys.path:
            sys.path.insert(0, _p)
    from acestep.handler import AceStepHandler
    from acestep.llm_inference import LLMHandler
    from huggingface_hub import snapshot_download

    dit = AceStepHandler()
    status, ok = dit.initialize_service(
        project_root=ACE_DIR,
        config_path="acestep-v15-turbo",
        device=device,
        offload_to_cpu=False,
    )
    if not ok:
        raise RuntimeError(f"ACE-Step DiT init failed: {status}")

    llm = LLMHandler()
    status, ok = llm.initialize(
        checkpoint_dir=CHECKPOINTS_DIR,
        lm_model_path="acestep-5Hz-lm-1.7B",
        backend="pt",
        device=device,
        offload_to_cpu=False,
    )
    if not ok:
        raise RuntimeError(f"ACE-Step LM init failed: {status}")

    lora_dir = snapshot_download("veeceey/RagaLoRA-indian-music-ace-step")
    dit.add_lora(lora_dir)
    dit.set_use_lora(True)
    dit.set_lora_scale(0.8)

    _state["dit"], _state["llm"], _state["ace_ok"] = dit, llm, True
    return dit, llm


def _load_svc(device: torch.device):
    if not _ENV_READY.wait(timeout=600):
        raise RuntimeError(
            "Environment preparation did not finish in 10 min — check the Space runtime logs."
        )
    if _state.get("svc_ok"):
        return _state["svc"]

    for _p in (ACE_DIR, SEEDVC_DIR):
        if _p not in sys.path:
            sys.path.insert(0, _p)
    import yaml
    from types import SimpleNamespace
    from hf_utils import load_custom_model_from_hf
    from modules.commons import build_model, load_checkpoint, recursive_munch

    fp16 = True
    dit_checkpoint_path, dit_config_path = load_custom_model_from_hf(
        "Plachta/Seed-VC",
        "DiT_seed_v2_uvit_whisper_base_f0_44k_bigvgan_pruned_ft_ema_v2.pth",
        "config_dit_mel_seed_uvit_whisper_base_f0_44k.yml",
    )
    config = yaml.safe_load(open(dit_config_path, "r"))
    model_params = recursive_munch(config["model_params"])
    model_params.dit_type = "DiT"
    model = build_model(model_params, stage="DiT")
    hop_length = config["preprocess_params"]["spect_params"]["hop_length"]
    sr = config["preprocess_params"]["sr"]

    model, _, _, _ = load_checkpoint(
        model, None, dit_checkpoint_path,
        load_only_params=True, ignore_modules=[], is_distributed=False,
    )
    for key in model:
        model[key].eval()
        model[key].to(device)
    model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192)

    from modules.campplus.DTDNN import CAMPPlus
    campplus_ckpt_path = load_custom_model_from_hf(
        "funasr/campplus", "campplus_cn_common.bin", config_filename=None
    )
    campplus_model = CAMPPlus(feat_dim=80, embedding_size=192)
    campplus_model.load_state_dict(torch.load(campplus_ckpt_path, map_location="cpu"))
    campplus_model.eval().to(device)

    vocoder_type = model_params.vocoder.type
    if vocoder_type == "bigvgan":
        from modules.bigvgan import bigvgan
        bigvgan_model = bigvgan.BigVGAN.from_pretrained(
            model_params.vocoder.name, use_cuda_kernel=False
        )
        bigvgan_model.remove_weight_norm()
        vocoder_fn = bigvgan_model.eval().to(device)
    else:
        raise ValueError(f"Unsupported vocoder type: {vocoder_type}")

    speech_tokenizer_type = model_params.speech_tokenizer.type
    if speech_tokenizer_type != "whisper":
        raise ValueError(f"Unsupported speech tokenizer: {speech_tokenizer_type}")
    from transformers import AutoFeatureExtractor, WhisperModel
    whisper_name = model_params.speech_tokenizer.name
    whisper_model = WhisperModel.from_pretrained(
        whisper_name, torch_dtype=torch.float16
    ).to(device)
    del whisper_model.decoder
    whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_name)

    def semantic_fn(waves_16k):
        ori_inputs = whisper_feature_extractor(
            [waves_16k.squeeze(0).cpu().numpy()],
            return_tensors="pt",
            return_attention_mask=True,
        )
        ori_input_features = whisper_model._mask_input_features(
            ori_inputs.input_features, attention_mask=ori_inputs.attention_mask
        ).to(device)
        with torch.no_grad():
            ori_outputs = whisper_model.encoder(
                ori_input_features.to(whisper_model.encoder.dtype),
                head_mask=None,
                output_attentions=False,
                output_hidden_states=False,
                return_dict=True,
            )
        S_ori = ori_outputs.last_hidden_state.to(torch.float32)
        S_ori = S_ori[:, : waves_16k.size(-1) // 320 + 1]
        return S_ori

    mel_fn_args = {
        "n_fft": config["preprocess_params"]["spect_params"]["n_fft"],
        "win_size": config["preprocess_params"]["spect_params"]["win_length"],
        "hop_size": config["preprocess_params"]["spect_params"]["hop_length"],
        "num_mels": config["preprocess_params"]["spect_params"]["n_mels"],
        "sampling_rate": sr,
        "fmin": config["preprocess_params"]["spect_params"].get("fmin", 0),
        "fmax": None
        if config["preprocess_params"]["spect_params"].get("fmax", "None") == "None"
        else 8000,
        "center": False,
    }
    from modules.audio import mel_spectrogram
    to_mel = lambda x: mel_spectrogram(x, **mel_fn_args)

    from modules.rmvpe import RMVPE
    rmvpe_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None)
    rmvpe = RMVPE(rmvpe_path, is_half=False, device=device)

    svc = SimpleNamespace(
        model=model,
        semantic_fn=semantic_fn,
        vocoder_fn=vocoder_fn,
        campplus_model=campplus_model,
        to_mel=to_mel,
        f0_fn=rmvpe.infer_from_audio,
        sr=sr,
        hop_length=hop_length,
        fp16=fp16,
        overlap_frame_len=16,
        device=device,
    )
    svc.max_context_window = sr // hop_length * 30
    svc.overlap_wave_len = svc.overlap_frame_len * hop_length

    _state["svc"], _state["svc_ok"] = svc, True
    return svc


def _adjust_f0_semitones(f0_sequence, n_semitones):
    factor = 2 ** (n_semitones / 12)
    return f0_sequence * factor


def _crossfade(chunk1, chunk2, overlap):
    fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2
    fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2
    chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out
    return chunk2


# ---------------------------------------------------------------------------
# Seed-VC singing-voice conversion (adapted non-streaming from seed-vc app_svc.py)
# ---------------------------------------------------------------------------
@torch.no_grad()
@torch.inference_mode()
def _voice_convert(svc, source_path, target_path, diffusion_steps=30,
                   length_adjust=1.0, inference_cfg_rate=0.7,
                   auto_f0_adjust=True, pitch_shift=0):
    import librosa
    import torchaudio

    device = svc.device
    inference_module = svc.model
    mel_fn = svc.to_mel
    sr = svc.sr

    source_audio = librosa.load(source_path, sr=sr)[0]
    ref_audio = librosa.load(target_path, sr=sr)[0]

    source_audio = torch.tensor(source_audio).unsqueeze(0).float().to(device)
    ref_audio = torch.tensor(ref_audio[: sr * 25]).unsqueeze(0).float().to(device)

    ref_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000)
    converted_waves_16k = torchaudio.functional.resample(source_audio, sr, 16000)

    if converted_waves_16k.size(-1) <= 16000 * 30:
        S_alt = svc.semantic_fn(converted_waves_16k)
    else:
        overlapping_time = 5  # seconds
        S_alt_list = []
        buffer = None
        traversed_time = 0
        while traversed_time < converted_waves_16k.size(-1):
            if buffer is None:
                chunk = converted_waves_16k[:, traversed_time: traversed_time + 16000 * 30]
            else:
                chunk = torch.cat(
                    [buffer, converted_waves_16k[:, traversed_time: traversed_time + 16000 * (30 - overlapping_time)]],
                    dim=-1,
                )
            S_alt = svc.semantic_fn(chunk)
            if traversed_time == 0:
                S_alt_list.append(S_alt)
            else:
                S_alt_list.append(S_alt[:, 50 * overlapping_time:])
            buffer = chunk[:, -16000 * overlapping_time:]
            traversed_time += (
                30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time
            )
        S_alt = torch.cat(S_alt_list, dim=1)

    ori_waves_16k = torchaudio.functional.resample(ref_audio, sr, 16000)
    S_ori = svc.semantic_fn(ori_waves_16k)

    mel = mel_fn(source_audio.to(device).float())
    mel2 = mel_fn(ref_audio.to(device).float())

    target_lengths = torch.LongTensor([int(mel.size(2) * length_adjust)]).to(mel.device)
    target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device)

    feat2 = torchaudio.compliance.kaldi.fbank(
        ref_waves_16k, num_mel_bins=80, dither=0, sample_frequency=16000
    )
    feat2 = feat2 - feat2.mean(dim=0, keepdim=True)
    style2 = svc.campplus_model(feat2.unsqueeze(0))

    F0_ori = svc.f0_fn(ref_waves_16k[0], thred=0.03)
    F0_alt = svc.f0_fn(converted_waves_16k[0], thred=0.03)

    F0_ori = torch.from_numpy(F0_ori).to(device)[None]
    F0_alt = torch.from_numpy(F0_alt).to(device)[None]

    voiced_F0_ori = F0_ori[F0_ori > 1]
    voiced_F0_alt = F0_alt[F0_alt > 1]

    log_f0_alt = torch.log(F0_alt + 1e-5)
    voiced_log_f0_ori = torch.log(voiced_F0_ori + 1e-5)
    voiced_log_f0_alt = torch.log(voiced_F0_alt + 1e-5)
    median_log_f0_ori = torch.median(voiced_log_f0_ori)
    median_log_f0_alt = torch.median(voiced_log_f0_alt)

    shifted_log_f0_alt = log_f0_alt.clone()
    if auto_f0_adjust:
        shifted_log_f0_alt[F0_alt > 1] = (
            log_f0_alt[F0_alt > 1] - median_log_f0_alt + median_log_f0_ori
        )
    shifted_f0_alt = torch.exp(shifted_log_f0_alt)
    if pitch_shift != 0:
        shifted_f0_alt[F0_alt > 1] = _adjust_f0_semitones(
            shifted_f0_alt[F0_alt > 1], pitch_shift
        )

    cond, _, codes, commitment_loss, codebook_loss = inference_module.length_regulator(
        S_alt, ylens=target_lengths, n_quantizers=3, f0=shifted_f0_alt
    )
    prompt_condition, _, codes, commitment_loss, codebook_loss = inference_module.length_regulator(
        S_ori, ylens=target2_lengths, n_quantizers=3, f0=F0_ori
    )
    interpolated_shifted_f0_alt = torch.nn.functional.interpolate(
        shifted_f0_alt.unsqueeze(1), size=cond.size(1), mode="nearest"
    ).squeeze(1)
    max_source_window = svc.max_context_window - mel2.size(2)

    processed_frames = 0
    generated_wave_chunks = []
    while processed_frames < cond.size(1):
        chunk_cond = cond[:, processed_frames: processed_frames + max_source_window]
        is_last_chunk = processed_frames + max_source_window >= cond.size(1)
        cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1)
        with torch.autocast(
            device_type=device.type, dtype=torch.float16 if svc.fp16 else torch.float32
        ):
            vc_target = inference_module.cfm.inference(
                cat_condition,
                torch.LongTensor([cat_condition.size(1)]).to(mel2.device),
                mel2, style2, None, diffusion_steps,
                inference_cfg_rate=inference_cfg_rate,
            )
            vc_target = vc_target[:, :, mel2.size(-1):]
        vc_wave = svc.vocoder_fn(vc_target.float()).squeeze().cpu()
        if vc_wave.ndim == 1:
            vc_wave = vc_wave.unsqueeze(0)
        if processed_frames == 0:
            if is_last_chunk:
                generated_wave_chunks.append(vc_wave[0].cpu().numpy())
                break
            generated_wave_chunks.append(vc_wave[0, : -svc.overlap_wave_len].cpu().numpy())
            previous_chunk = vc_wave[0, -svc.overlap_wave_len:]
            processed_frames += vc_target.size(2) - svc.overlap_frame_len
        elif is_last_chunk:
            output_wave = _crossfade(
                previous_chunk.cpu().numpy(), vc_wave[0].cpu().numpy(), svc.overlap_wave_len
            )
            generated_wave_chunks.append(output_wave)
            processed_frames += vc_target.size(2) - svc.overlap_frame_len
            break
        else:
            output_wave = _crossfade(
                previous_chunk.cpu().numpy(),
                vc_wave[0, : -svc.overlap_wave_len].cpu().numpy(),
                svc.overlap_wave_len,
            )
            generated_wave_chunks.append(output_wave)
            previous_chunk = vc_wave[0, -svc.overlap_wave_len:]
            processed_frames += vc_target.size(2) - svc.overlap_frame_len

    return np.concatenate(generated_wave_chunks), sr


# ---------------------------------------------------------------------------
# Demucs stem split + remix
# ---------------------------------------------------------------------------
def _split_stems(song_path, out_dir):
    from demucs.separate import main as demucs_main
    demucs_main(["--two-stems", "vocals", "-n", "htdemucs", "-o", out_dir, song_path])
    base = os.path.splitext(os.path.basename(song_path))[0]
    d = os.path.join(out_dir, "htdemucs", base)
    return os.path.join(d, "vocals.wav"), os.path.join(d, "no_vocals.wav")


def _remix(instrumental_path, converted_vocals_path, out_path):
    inst, sr_i = sf.read(instrumental_path, always_2d=True)
    voc, sr_v = sf.read(converted_vocals_path, always_2d=True)
    if sr_v != sr_i:
        import torchaudio
        import librosa
        voc_mono = librosa.resample(np.mean(voc, axis=1), orig_sr=sr_v, target_sr=sr_i)
        voc = voc_mono[:, None]
    if voc.shape[1] == 1 and inst.shape[1] == 2:
        voc = np.repeat(voc, 2, axis=1)
    n = max(inst.shape[0], voc.shape[0])
    inst = np.pad(inst, ((0, n - inst.shape[0]), (0, 0)))
    voc = np.pad(voc, ((0, n - voc.shape[0]), (0, 0)))

    def _norm(x, peak):
        return x / (np.max(np.abs(x)) + 1e-9) * peak

    mix = _norm(inst, 0.9) + _norm(voc, 0.7)
    sf.write(out_path, np.clip(mix, -1.0, 1.0).astype(np.float32), sr_i)
    return out_path


# ---------------------------------------------------------------------------
# Languages (ISO 639-1 codes for ACE-Step vocal_language)
# ---------------------------------------------------------------------------
LANGS = [
    ("Auto-detect", "unknown"),
    ("English", "en"),
    ("हिन्दी Hindi", "hi"),
    ("اردو Urdu", "ur"),
    ("ਪੰਜਾਬੀ Punjabi", "pa"),
    ("தமிழ் Tamil", "ta"),
    ("తెలుగు Telugu", "te"),
    ("বাংলা Bengali", "bn"),
    ("मराठी Marathi", "mr"),
    ("ગુજરાતી Gujarati", "gu"),
    ("ಕನ್ನಡ Kannada", "kn"),
    ("മലയാളം Malayalam", "ml"),
    ("ଓଡ଼ିଆ Odia", "or"),
    ("অসমীয়া Assamese", "as"),
    ("नेपाली Nepali", "ne"),
    ("සිංහල Sinhala", "si"),
]


# ---------------------------------------------------------------------------
# Archive helper: persist finished audio to the Hub so a take can never be lost
# to an ephemeral /tmp dir or a Space restart. Defined here (module level, before
# the pipelines) so every panel — song, rap, agent — shares one code path.
# ---------------------------------------------------------------------------
OUTPUTS_REPO = os.environ.get("AGENT_OUTPUTS_REPO", "red-kit/sangeet-studio-outputs")


def _stamp():
    """UTC timestamp for archive filenames — no local-tz surprises on the box."""
    return time.strftime("%Y%m%d-%H%M%S", time.gmtime())


def _archive_output(wav_path, name):
    """Copy a finished track into OUTPUTS_REPO. Returns (link, error).

    Returns (None, None) when no HF_TOKEN is configured — the Space still works
    locally, it just cannot persist. Returns (None, <msg>) on a real failure so
    the caller can surface *why* the take was not saved instead of losing it
    silently.
    """
    token = os.environ.get("HF_TOKEN")
    if not token:
        return None, None
    try:
        from huggingface_hub import HfApi
        HfApi(token=token).upload_file(
            path_or_fileobj=wav_path,
            path_in_repo=f"variants/{name}.wav",
            repo_id=OUTPUTS_REPO,
            repo_type="dataset",
        )
        return f"https://huggingface.co/datasets/{OUTPUTS_REPO}/blob/main/variants/{name}.wav", None
    except Exception as e:
        print(f"[archive] saving {name} failed: {e}", flush=True)
        return None, str(e)


def _archive_note(wav_path, name):
    """Upload and return a short markdown suffix for the status message."""
    link, err = _archive_output(wav_path, name)
    if link:
        return f"\n\n📦 **Archived:** [{name}]({link})"
    if err:
        return (
            f"\n\n⚠️ Could not archive this take: `{err}` — download it now, "
            f"check that the `HF_TOKEN` secret can write to `{OUTPUTS_REPO}`."
        )
    return "\n\n_(Not archived — set the `HF_TOKEN` secret to auto-save takes to the Hub.)_"


# ---------------------------------------------------------------------------
# Main pipeline
# ---------------------------------------------------------------------------
def run_pipeline(
    caption, lyrics, language, duration, instrumental, thinking,
    use_lora, lora_scale, seed,
    voice_ref, vc_steps, auto_f0, pitch_shift,
    progress=gr.Progress(),
):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tmp = tempfile.mkdtemp(prefix="sangeet_")
    try:
        lyrics = (lyrics or "").strip()
        caption = (caption or "").strip()
        if not instrumental and not lyrics and not caption:
            yield "❌ Give a style prompt or lyrics first.", None, None, None, None
            return
        lyrics = lyrics[:4000]
        caption = caption[:500]
        if instrumental:
            lyrics = "[Instrumental]"

        yield "⏳ **Step 1/4** — loading ACE-Step turbo + LM + RagaLoRA…", None, None, None, None
        dit, llm = _load_ace(device)
        dit.set_use_lora(bool(use_lora))
        if use_lora:
            dit.set_lora_scale(float(lora_scale))

        yield "🎵 **Step 2/4** — generating the song…", None, None, None, None
        from acestep.inference import GenerationParams, GenerationConfig, generate_music

        params = GenerationParams(
            task_type="text2music",
            thinking=bool(thinking),
            caption=caption,
            lyrics=lyrics,
            instrumental=bool(instrumental),
            vocal_language=language,
            duration=float(duration),
            inference_steps=8,
            guidance_scale=1.0,   # turbo auto-forces CFG off; documented value
            shift=3.0,            # required timestep shift for turbo
            seed=int(seed),
        )
        config = GenerationConfig(
            batch_size=1,
            audio_format="wav",
            use_random_seed=int(seed) < 0,
        )
        result = generate_music(
            dit, llm, params=params, config=config, save_dir=os.path.join(tmp, "ace")
        )
        if not result.success:
            raise RuntimeError(result.status_message or result.error or "generation failed")
        song_path = result.audios[0]["path"]

        if voice_ref is None:
            note = _archive_note(song_path, f"{_stamp()}-song")
            yield (
                "✅ **Done** — song generated (no voice reference given, so ACE-Step's own vocals kept)." + note,
                song_path, song_path, None, None,
            )
            return

        yield "🎚️ **Step 3/4** — splitting stems (Demucs)…", None, None, None, None
        vocals_path, no_vocals_path = _split_stems(song_path, os.path.join(tmp, "stems"))

        yield "🎤 **Step 4/4** — converting vocals to the cloned artist voice (Seed-VC)…", None, None, None, None
        svc = _load_svc(device)
        converted_wave, sr = _voice_convert(
            svc, vocals_path, voice_ref,
            diffusion_steps=int(vc_steps),
            auto_f0_adjust=bool(auto_f0),
            pitch_shift=int(pitch_shift),
        )
        converted_path = os.path.join(tmp, "vocals_converted.wav")
        sf.write(converted_path, converted_wave.astype(np.float32), sr)

        final_path = _remix(no_vocals_path, converted_path, os.path.join(tmp, "final_song.wav"))
        note = _archive_note(final_path, f"{_stamp()}-song-cloned-voice")
        yield (
            "✅ **Done** — song generated, vocals swapped to the cloned artist voice, remixed." + note,
            final_path, song_path, no_vocals_path, converted_path,
        )
    except Exception as e:
        yield f"❌ **Failed:** `{e}`", None, None, None, None


# ---------------------------------------------------------------------------
# Live Rap Studio: recording -> auto beat -> polish/master
# ---------------------------------------------------------------------------
def _polish_and_mix(vocal_path, beat_path, out_path,
                    beat_gain=0.55, vocal_gain=0.9):
    """Polish the vocal (high-pass, denoise-trim, normalize, soft limiter)
    and mix it over the generated beat with a 1 s intro and fades."""
    import librosa
    from scipy.signal import butter, sosfilt

    voc, vsr = librosa.load(vocal_path, sr=44100, mono=True)
    beat, bsr = sf.read(beat_path, always_2d=True)
    if bsr != 44100:
        beat = np.stack(
            [librosa.resample(np.ascontiguousarray(beat[:, ch]).astype(np.float32),
                              orig_sr=bsr, target_sr=44100)
             for ch in range(beat.shape[1])], axis=1)
        bsr = 44100
    if beat.shape[1] == 1:
        beat = np.repeat(beat, 2, axis=1)
    if vsr != 44100:
        voc = librosa.resample(voc, orig_sr=vsr, target_sr=44100)

    # High-pass to remove rumble / handling noise, trim leading-trailing silence
    sos = butter(2, 75, btype="hp", fs=44100, output="sos")
    voc = sosfilt(sos, voc).astype(np.float32)
    voc, _ = librosa.effects.trim(voc, top_db=30)

    def _norm(x, peak):
        return (x / (np.max(np.abs(x)) + 1e-9)) * peak

    beat = _norm(beat, beat_gain)
    voc_st = _norm(np.repeat(voc[:, None], 2, axis=1), vocal_gain)

    # 1 s intro so the beat breathes before the vocal enters
    intro = int(1.0 * 44100)
    n = max(beat.shape[0], voc_st.shape[0] + intro)
    beat_p = np.pad(beat, ((0, max(0, n - beat.shape[0])), (0, 0)))
    voc_p = np.pad(voc_st, ((intro, max(0, n - voc_st.shape[0] - intro)), (0, 0)))

    mix = beat_p + voc_p
    # Soft limiter for consistent loudness, then fades
    mix = np.tanh(mix * 1.3) / np.tanh(1.3)
    fade = int(0.5 * 44100)
    env = np.ones(mix.shape[0])
    env[:fade] = np.linspace(0, 1, fade)
    env[-fade:] = np.linspace(1, 0, fade)
    mix = mix * env[:, None]
    peak = np.max(np.abs(mix)) + 1e-9
    if peak > 0.99:
        mix = mix / peak * 0.99
    sf.write(out_path, mix.astype(np.float32), 44100)
    return out_path


def run_rap_pipeline(
    rap_vocal, beat_prompt, bpm, auto_bpm, duration_extra,
    use_clone, voice_ref, vc_steps, auto_f0, pitch_shift,
    progress=gr.Progress(),
):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tmp = tempfile.mkdtemp(prefix="rap_")
    try:
        if rap_vocal is None:
            yield "❌ Record or upload a rap vocal first.", None, None, None
            return

        import librosa
        yield "⏳ **Step 1/4** — analyzing your recording…", None, None, None
        vocal_np, vocal_sr = librosa.load(rap_vocal, sr=None, mono=True)
        rec_dur = len(vocal_np) / vocal_sr
        if rec_dur < 3:
            raise RuntimeError(f"Recording too short ({rec_dur:.1f}s) — rap at least a few bars.")

        est_bpm = None
        if auto_bpm:
            try:
                tempo = librosa.beat.beat_track(y=vocal_np, sr=vocal_sr)[0]
                est_bpm = float(np.atleast_1d(tempo)[0]) or None
            except Exception:
                est_bpm = None
        use_bpm = int(round(est_bpm)) if (auto_bpm and est_bpm) else int(bpm)

        yield (
            f"🎛️ **Step 2/4** — generating the matching instrumental "
            f"({rec_dur:.0f}s vocal, {use_bpm} BPM)…",
            None, None, None,
        )
        dit, llm = _load_ace(device)
        from acestep.inference import GenerationParams, GenerationConfig, generate_music

        caption = (beat_prompt or "hard-hitting rap beat").strip()[:500]
        caption = f"{caption}, rap hip-hop instrumental backing track, no vocals"
        beat_duration = max(10.0, min(600.0, float(rec_dur) + float(duration_extra)))
        params = GenerationParams(
            task_type="text2music",
            thinking=True,
            caption=caption,
            lyrics="[Instrumental]",
            instrumental=True,
            vocal_language="unknown",
            duration=beat_duration,
            bpm=use_bpm,
            inference_steps=8,
            guidance_scale=1.0,
            shift=3.0,
            seed=-1,
        )
        config = GenerationConfig(batch_size=1, audio_format="wav", use_random_seed=True)
        result = generate_music(
            dit, llm, params=params, config=config, save_dir=os.path.join(tmp, "ace")
        )
        if not result.success:
            raise RuntimeError(result.status_message or result.error or "beat generation failed")
        beat_path = result.audios[0]["path"]
        beat_note = _archive_note(beat_path, f"rap-{_stamp()}-beat")

        voc_proc = rap_vocal
        if use_clone and voice_ref is not None:
            yield "🎤 **Step 3/4** — converting your recording to the cloned voice (Seed-VC)…", None, None, None
            svc = _load_svc(device)
            wave, sr = _voice_convert(
                svc, rap_vocal, voice_ref,
                diffusion_steps=int(vc_steps),
                auto_f0_adjust=bool(auto_f0),
                pitch_shift=int(pitch_shift),
            )
            voc_proc = os.path.join(tmp, "rap_converted.wav")
            sf.write(voc_proc, wave.astype(np.float32), sr)

        yield "✨ **Step 4/4** — polishing and mastering the mix…", None, None, None
        final_path = _polish_and_mix(voc_proc, beat_path, os.path.join(tmp, "rap_final.wav"))

        # Persist the finished master (and the beat) so the take survives a
        # page reload, a Space restart, or the /tmp dir being reclaimed.
        stamp = _stamp()
        final_note = _archive_note(final_path, f"rap-{stamp}-master")
        yield (
            f"✅ **Done** — {rec_dur:.0f}s rap over a {use_bpm} BPM generated beat, "
            "polished and mastered." + final_note,
            final_path, beat_path, voc_proc,
        )
    except Exception as e:
        yield f"❌ **Failed:** `{e}`", None, None, None


# ---------------------------------------------------------------------------
# AI Agent Studio: client brief -> LLM producer/lyricist -> variants -> Hub
# ---------------------------------------------------------------------------
AGENT_LLM_MODEL = os.environ.get("AGENT_LLM_MODEL", "Qwen/Qwen2.5-72B-Instruct")


def _agent_llm():
    token = os.environ.get("HF_TOKEN")
    if not token:
        return None
    try:
        from huggingface_hub import InferenceClient
        return InferenceClient(model=AGENT_LLM_MODEL, token=token)
    except Exception as e:
        print(f"[agent] LLM client unavailable: {e}", flush=True)
        return None


def _llm_json(client, system, user):
    import json
    import re
    text = client.chat_completion(
        messages=[{"role": "system", "content": system},
                  {"role": "user", "content": user}],
        max_tokens=1800, temperature=0.85,
    ).choices[0].message.content
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if not m:
        raise ValueError(f"LLM did not return JSON — got: {text[:300]}")
    return json.loads(m.group(0))


def _agent_plan(client, artist_style, brief, song_type, language_label, duration):
    """LLM producer: turns the client brief into a concrete production plan."""
    if client is None:
        caption = f"{artist_style or song_type}, {song_type}, polished studio production"
        lyrics = (brief or "").strip()
        if not lyrics:
            lyrics = f"[Verse]\n{song_type} night, city lights and longing\n[Chorus]\nRaise your heart, sing along"
        return {"caption": caption, "lyrics": lyrics, "bpm": None,
                "lora_scale": 0.8, "language": "unknown",
                "note": "Fallback planner — add the HF_TOKEN secret for the full LLM lyricist."}
    system = (
        "You are a world-class music producer and lyricist for South Asian music "
        "(Bollywood, qawwali, ghazal, Punjabi, Sufi, Carnatic, filmi, hip-hop). "
        "Given a client brief, return ONLY a JSON object with keys: "
        '"caption" (rich English style prompt for a music model: instruments, mood, tempo, era, vocal character), '
        '"lyrics" (full song lyrics in the requested language using [Verse]/[Chorus]/[Bridge] tags; original words, no real artist names), '
        '"bpm" (integer 60-180 or null for auto), '
        '"lora_scale" (0.0-1.0 strength of Indian-style LoRA), '
        '"language" (ISO 639-1 code for the vocals), '
        '"note" (one-line production note for the client).'
    )
    user = (
        f"Artist style (described in words — inspiration, not a voice to clone): {artist_style or 'producer\'s choice'}\n"
        f"Song type/genre: {song_type}\n"
        f"Vocal language requested: {language_label}\n"
        f"Target duration: {int(duration)} seconds\n"
        f"Client lyrics (full lyrics if provided) or keywords/themes: {brief or 'none — write from the style'}\n\n"
        "If the client provided full lyrics, polish structure and flow but keep their words; return them as 'lyrics'. "
        "If they provided keywords, write original lyrics from those themes. "
        "Match lyric length to the target duration (roughly 2 words per second for sung vocals)."
    )
    return _llm_json(client, system, user)


def run_agent_pipeline(
    artist_style, brief, song_type, agent_language, duration, n_variants,
    voice_ref, consent, vc_steps, auto_f0, pitch_shift,
    progress=gr.Progress(),
):
    import random
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tmp = tempfile.mkdtemp(prefix="agent_")
    chat = []
    outs = [None, None, None]

    def emit(text):
        chat.append({"role": "assistant", "content": text})
        return (list(chat), outs[0], outs[1], outs[2])

    try:
        if not (artist_style or "").strip() and not (brief or "").strip():
            yield emit("❌ Give me at least a style, a song type, or lyrics/keywords to work with.")
            return
        if voice_ref is not None and not consent:
            yield emit("⚠️ Voice conversion needs the consent checkbox ticked — you must own the voice or have permission to use it.")
            return
        n_variants = int(max(1, min(3, n_variants)))

        yield emit("🧠 **Producer agent:** planning the production and writing the lyrics…")
        client = _agent_llm()
        lang_label = dict(LANGS).get(agent_language, "auto")
        plan = _agent_plan(client, artist_style, brief, song_type, lang_label, duration)
        caption = str(plan.get("caption", "")).strip()[:500]
        lyrics = str(plan.get("lyrics", "")).strip()[:4000]
        bpm = plan.get("bpm")
        bpm = int(bpm) if bpm else None
        lora_scale = float(plan.get("lora_scale", 0.8))
        plan_language = str(plan.get("language") or agent_language or "unknown")
        note = str(plan.get("note", ""))

        yield emit(
            f"📋 **Production plan**\n\n"
            f"- Style: `{caption}`\n"
            f"- BPM: {bpm or 'auto'} · LoRA strength: {lora_scale:.2f} · Language: `{plan_language}`\n"
            f"- Producer note: {note}\n"
            f"- Lyricist: `{AGENT_LLM_MODEL}`" + ("" if client else " *(fallback planner — add the HF_TOKEN secret for the full LLM)*"),
        )

        dit, llm = _load_ace(device)
        dit.set_use_lora(True)
        dit.set_lora_scale(lora_scale)

        from acestep.inference import GenerationParams, GenerationConfig, generate_music

        links = []
        for i in range(n_variants):
            seed = random.randint(0, 999999)
            yield emit(f"🎛️ **Variant {i + 1}/{n_variants}** — generating (seed {seed})…")
            params = GenerationParams(
                task_type="text2music",
                thinking=True,
                caption=caption,
                lyrics=lyrics or "[Instrumental]",
                instrumental=not lyrics,
                vocal_language=plan_language,
                duration=float(duration),
                bpm=bpm,
                inference_steps=8,
                guidance_scale=1.0,
                shift=3.0,
                seed=seed,
            )
            config = GenerationConfig(batch_size=1, audio_format="wav", use_random_seed=False)
            result = generate_music(
                dit, llm, params=params, config=config, save_dir=os.path.join(tmp, f"v{i}")
            )
            if not result.success:
                raise RuntimeError(result.status_message or result.error or "generation failed")
            song_path = result.audios[0]["path"]

            if voice_ref is not None:
                yield emit(f"🎤 **Variant {i + 1}** — swapping in the reference voice…")
                vocals_path, no_vocals_path = _split_stems(song_path, os.path.join(tmp, f"stems{i}"))
                svc = _load_svc(device)
                wave, sr = _voice_convert(
                    svc, vocals_path, voice_ref,
                    diffusion_steps=int(vc_steps),
                    auto_f0_adjust=bool(auto_f0),
                    pitch_shift=int(pitch_shift),
                )
                converted = os.path.join(tmp, f"conv{i}.wav")
                sf.write(converted, wave.astype(np.float32), sr)
                song_path = _remix(no_vocals_path, converted, os.path.join(tmp, f"final{i}.wav"))

            outs[i] = song_path
            link, save_err = _archive_output(song_path, f"{_stamp()}-variant{i + 1}")
            if link:
                links.append(f"[Variant {i + 1}]({link})")
                yield emit(f"✅ **Variant {i + 1}** ready — saved to the Hub: {link}")
            elif save_err:
                yield emit(
                    f"✅ **Variant {i + 1}** ready, but archiving to the Hub failed: `{save_err}` "
                    f"— check that the HF_TOKEN secret has write access to `{OUTPUTS_REPO}`."
                )
            else:
                yield emit(f"✅ **Variant {i + 1}** ready (add the HF_TOKEN secret to also archive variants on the Hub).")

        if links:
            yield emit("📦 **All variants saved:**\n" + "\n".join(links))
        else:
            yield emit("📦 All variants delivered above.")
    except Exception as e:
        chat.append({"role": "assistant", "content": f"❌ **Failed:** `{e}`"})
        yield (list(chat), outs[0], outs[1], outs[2])


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
print("[startup] building Gradio UI…", flush=True)
demo = gr.Blocks(title="Sangeet Studio — AI Song Generation + Voice Cloning")
with demo:
    gr.Markdown(
        """# 🎵 Sangeet Studio
**Suno-style full-song generation (ACE-Step 1.5 turbo + RagaLoRA for Indian styles) with
instant artist-voice cloning (Seed-VC).**

1. Describe the style, paste lyrics (write them in the **native script** — Devanagari,
   Shahmukhi, Gurmukhi, Tamil… — transliteration degrades pronunciation), pick the language.
2. Optionally upload a **5–25 s clean reference clip** of the artist voice you want singing it.
3. Hit generate. The app produces the song, splits the stems, converts the vocals into the
   cloned voice, and remixes.
"""
    )
    with gr.Row():
        with gr.Column(scale=2):
            caption = gr.Textbox(
                label="Style prompt",
                placeholder="e.g. Coke Studio style Sufi qawwali with tabla, harmonium, deep male vocals",
            )
            lyrics = gr.Textbox(
                label="Lyrics (use section tags like [Verse] / [Chorus])",
                lines=7,
                placeholder="[Verse]\nراند کھڑکی اے...\n[Chorus]\n...",
            )
            language = gr.Dropdown(
                choices=LANGS, value="hi", label="Vocal language",
                info="Punjabi/Urdu and other long-tail languages are supported but rougher than Hindi/English.",
            )
            duration = gr.Slider(10, 240, value=60, step=5, label="Duration (seconds)")
            instrumental = gr.Checkbox(False, label="Instrumental (no vocals)")
            with gr.Accordion("Generation options", open=False):
                thinking = gr.Checkbox(True, label="LM planning (CoT) — better structure, a bit slower")
                use_lora = gr.Checkbox(True, label="Use RagaLoRA (Indian music styles)")
                lora_scale = gr.Slider(0.0, 1.0, value=0.8, step=0.05, label="RagaLoRA strength")
                seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
        with gr.Column(scale=2):
            voice_ref = gr.Audio(
                label="Artist voice reference (5–25 s clean clip, optional)",
                type="filepath",
            )
            vc_steps = gr.Slider(5, 100, value=30, step=5,
                                 label="Voice-conversion diffusion steps",
                                 info="30–50 = best quality, slower")
            auto_f0 = gr.Checkbox(True, label="Auto F0 adjust (match the clone's pitch range)")
            pitch_shift = gr.Slider(-24, 24, value=0, step=1,
                                    label="Pitch shift (semitones)", info="e.g. -5 for male↔female correction")
    with gr.Accordion("🎤 Live Rap Studio — record → auto beat → polished track", open=False):
        gr.Markdown(
            """**Record your rap with the browser mic** (or upload a dry vocal). The app estimates the
 tempo, generates a **matching rap instrumental in the same style**, keeps **your voice as the
 artist** (or converts it to a cloned artist voice), then polishes and masters the mix into a
 production-ready track. **Every finished master is auto-archived to the Hub** so it survives a
 restart or page reload — the archive link appears in the status line."""
        )
        with gr.Row():
            rap_vocal = gr.Audio(
                label="Your rap vocal — record with mic or upload",
                sources=["microphone", "upload"],
                type="filepath",
            )
            with gr.Column():
                beat_prompt = gr.Textbox(
                    label="Beat style prompt",
                    value="hard-hitting Indian hip-hop trap beat, 808 bass, tabla accents, moody",
                )
                with gr.Row():
                    bpm = gr.Slider(60, 180, value=90, step=1, label="BPM fallback",
                                    info="Used when auto-estimation is unsure")
                    auto_bpm = gr.Checkbox(True, label="Auto-estimate BPM from my recording")
                    duration_extra = gr.Slider(0, 15, value=4, step=1,
                                               label="Beat length beyond vocal (s)")
        with gr.Row():
            use_clone = gr.Checkbox(False, label="Convert my recording to a cloned artist voice")
            rap_voice_ref = gr.Audio(label="Artist voice reference (if converting)",
                                     sources=["upload"], type="filepath")
            rap_vc_steps = gr.Slider(5, 100, value=30, step=5, label="VC diffusion steps")
            rap_auto_f0 = gr.Checkbox(True, label="Auto F0 adjust")
            rap_pitch = gr.Slider(-24, 24, value=0, step=1, label="Pitch shift (semitones)")
        rap_btn = gr.Button("🎧 Turn my recording into a production-ready rap", variant="primary")
        rap_status = gr.Markdown("Ready.")
        with gr.Row():
            rap_final = gr.Audio(label="Final polished rap")
            rap_beat = gr.Audio(label="Generated instrumental")
        rap_raw = gr.Audio(label="Processed vocal (as mixed)")

    with gr.Accordion("🤖 AI Agent Studio — brief in, finished songs out", open=False):
        gr.Markdown(
            """Describe the **style in words**, give **full lyrics or keywords** + song type. The agent plans the
 production, writes/polishes lyrics with an LLM, generates **multiple variants**, converts them into your
 reference voice if provided (consented, user-supplied), and **saves every variant to the Hub**.
 Add the `HF_TOKEN` secret to this Space to enable the LLM lyricist and Hub archiving."""
        )
        with gr.Row():
            with gr.Column():
                agent_style = gr.Textbox(
                    label="Artist style (describe in words)",
                    placeholder="e.g. soulful devotional qawwali intensity, soaring male vocals, harmonium + tabla",
                )
                agent_brief = gr.Textbox(
                    label="Full lyrics OR keywords/themes",
                    lines=5,
                    placeholder="Paste full lyrics (kept + polished) — or just themes: ‘monsoon, first love, train station’",
                )
                with gr.Row():
                    agent_song_type = gr.Dropdown(
                        ["Bollywood", "Qawwali", "Ghazal", "Punjabi Bhangra", "Punjabi Pop",
                         "Sufi Rock", "Filmi", "Carnatic Fusion", "Hindustani Classical",
                         "Bhajan / Devotional", "Indian Hip-Hop / Rap", "Indie"],
                        value="Sufi Rock", label="Song type",
                    )
                    agent_language = gr.Dropdown(choices=LANGS, value="hi", label="Vocal language")
            with gr.Column():
                agent_duration = gr.Slider(10, 240, value=60, step=5, label="Duration per variant (s)")
                agent_n = gr.Slider(1, 3, value=2, step=1, label="Number of variants")
                agent_voice = gr.Audio(
                    label="Reference voice (optional — user-supplied)",
                    sources=["upload"], type="filepath",
                )
                agent_consent = gr.Checkbox(
                    False,
                    label="I confirm I have the right to use this voice (own voice or licensed/consented reference)",
                )
        with gr.Row():
            agent_vc_steps = gr.Slider(5, 100, value=30, step=5, label="VC diffusion steps")
            agent_auto_f0 = gr.Checkbox(True, label="Auto F0 adjust")
            agent_pitch = gr.Slider(-24, 24, value=0, step=1, label="Pitch shift (semitones)")
        agent_btn = gr.Button("🚀 Produce the songs", variant="primary")
        agent_chat = gr.Chatbot(label="Agent chat", height=380)
        with gr.Row():
            agent_v1 = gr.Audio(label="Variant 1")
            agent_v2 = gr.Audio(label="Variant 2")
            agent_v3 = gr.Audio(label="Variant 3")

    btn = gr.Button("🎵 Generate song (+ clone voice if reference given)", variant="primary")
    status = gr.Markdown("Ready. First generation loads the models (~2–5 min on GPU).")
    with gr.Row():
        final_audio = gr.Audio(label="Final song")
        original_audio = gr.Audio(label="Original generated song")
    with gr.Row():
        inst_audio = gr.Audio(label="Instrumental stem")
        voc_audio = gr.Audio(label="Vocals stem (voice-converted)")

    btn.click(
        run_pipeline,
        inputs=[caption, lyrics, language, duration, instrumental, thinking,
                use_lora, lora_scale, seed,
                voice_ref, vc_steps, auto_f0, pitch_shift],
        outputs=[status, final_audio, original_audio, inst_audio, voc_audio],
    )

    rap_btn.click(
        run_rap_pipeline,
        inputs=[rap_vocal, beat_prompt, bpm, auto_bpm, duration_extra,
                use_clone, rap_voice_ref, rap_vc_steps, rap_auto_f0, rap_pitch],
        outputs=[rap_status, rap_final, rap_beat, rap_raw],
    )

    agent_btn.click(
        run_agent_pipeline,
        inputs=[agent_style, agent_brief, agent_song_type, agent_language, agent_duration,
                agent_n, agent_voice, agent_consent, agent_vc_steps, agent_auto_f0, agent_pitch],
        outputs=[agent_chat, agent_v1, agent_v2, agent_v3],
    )

demo.queue(max_size=8).launch()
print("[startup] Gradio is up on port 7860", flush=True)