"""ACE-Step 1.5 turbo, fully fine-tuned on a small Japanese-song set.

The decoder weights come from a full fine-tune (every decoder parameter
updated, 1.58B of them) rather than an adapter, so they are applied by
overwriting the base decoder's parameters in place after the handler has
built the model -- there is no PEFT layer to attach.
"""

import spaces  # must precede torch: it patches torch.cuda before CUDA init

import json
import os
import tempfile

import gradio as gr
import torch
from huggingface_hub import hf_hub_download, snapshot_download

BASE_REPO = "ACE-Step/Ace-Step1.5"
FT_REPO = "swdq/acestep-v15-jpdenpa-ft"
FT_FILE = "epoch_40.pt"
LM_MODEL = "acestep-5Hz-lm-1.7B"

# The handler resolves base weights through this env var, so it has to be set
# before acestep is imported.
CHECKPOINT_DIR = snapshot_download(BASE_REPO)
os.environ["ACESTEP_CHECKPOINTS_DIR"] = CHECKPOINT_DIR

from acestep.handler import AceStepHandler  # noqa: E402
from acestep.inference import (  # noqa: E402
    GenerationConfig,
    GenerationParams,
    generate_music,
)
from acestep.llm_inference import LLMHandler  # noqa: E402

# ---------------------------------------------------------------------------
# Load once at module scope.  ZeroGPU intercepts the .to("cuda") that happens
# in here, packs the weights, and streams them in on the first @spaces.GPU
# call; loading lazily inside the handler would charge every user for it.
# ---------------------------------------------------------------------------

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

# Apply the fine-tune.  476 tensors, bf16, saved by training/scripts/train_full_ft.py.
_ft_path = hf_hub_download(FT_REPO, FT_FILE)
_blob = torch.load(_ft_path, map_location="cpu", weights_only=True)
_saved = _blob["trainable"]
_own = dict(dit.model.decoder.named_parameters())
_missing = [k for k in _saved if k not in _own]
if _missing:
    raise RuntimeError(f"{len(_missing)} fine-tuned tensors have no home in the "
                       f"decoder, e.g. {_missing[:3]}")
with torch.no_grad():
    for _k, _t in _saved.items():
        _p = _own[_k]
        _p.data.copy_(_t.to(dtype=_p.dtype, device=_p.device))
FT_META = _blob.get("meta", {})
print(f"[init] applied {len(_saved)} fine-tuned tensors "
      f"(epoch={FT_META.get('epoch')}, loss={FT_META.get('loss'):.4f})")
del _blob, _saved

llm = LLMHandler()
msg, ok = llm.initialize(
    checkpoint_dir=CHECKPOINT_DIR,
    lm_model_path=LM_MODEL,
    backend="pt",
    device="cuda",
    offload_to_cpu=False,
    dtype=None,
)
if not ok:
    raise RuntimeError(f"LM init failed: {msg}")

with open("samples/samples.json", encoding="utf-8") as f:
    SAMPLES = json.load(f)


def _duration_for(caption: str, lyrics: str, duration: float, steps: int,
                  seed: int, language: str) -> int:
    """Worst-case GPU seconds for a request.

    ZeroGPU calls this with the *same arguments as the decorated function*, not
    just the ones the estimate uses, so the signature has to mirror generate()'s
    exactly. Budget comes from the measured ~0.5 s of audio per second of
    compute at 8 steps, plus headroom for the LM pass.
    """
    return int(90 + duration * 0.6 * (steps / 8))


@spaces.GPU(duration=_duration_for)
def generate(caption: str, lyrics: str, duration: float, steps: int,
             seed: int, language: str) -> str:
    """Generate a song from a style caption and lyrics.

    Args:
        caption: Comma-separated style tags, e.g. "japanese vocaloid pop, female vocal".
        lyrics: Lyrics with section tags such as [Verse] / [Chorus].
        duration: Length of the generated audio in seconds.
        steps: Diffusion steps; 8 is the turbo default.
        seed: Random seed; the same seed and inputs reproduce the same audio.
        language: Vocal language code, e.g. "ja" or "en".

    Returns:
        Path to the generated WAV file.
    """
    if not caption.strip():
        raise gr.Error("Style caption is required.")

    # thinking stays on -- the LM still produces the audio semantic tokens.
    # What is off is the LM *overwriting the conditioning*: with the defaults
    # (use_cot_caption/use_cot_metas = True) the DiT receives an LM-authored
    # English description and an invented bpm, while this model was fine-tuned
    # on the dataset caption verbatim -- tag first, `bpm: N/A`. Leaving them on
    # drops the `jpdenpa` trigger before the DiT ever sees it.
    params = GenerationParams(
        task_type="text2music",
        thinking=True,
        use_cot_caption=False,
        use_cot_metas=False,
        caption=caption,
        lyrics=lyrics or "[Instrumental]",
        bpm=None,
        keyscale="",
        timesignature="4",
        vocal_language=language or "ja",
        duration=float(duration),
        inference_steps=int(steps),
        guidance_scale=1.0,
        seed=int(seed),
    )
    save_dir = tempfile.mkdtemp()
    res = generate_music(dit, llm, params=params,
                         config=GenerationConfig(batch_size=1, audio_format="wav"),
                         save_dir=save_dir)
    if not res.success or not res.audios:
        raise gr.Error(f"Generation failed: {res.status_message}")
    return res.audios[0]["path"]


TAG = "jpdenpa"
CSS = "footer{display:none!important}"

with gr.Blocks(title="ACE-Step 1.5 — jpdenpa full fine-tune", css=CSS) as demo:
    gr.Markdown(
        f"""
# ACE-Step 1.5 turbo — full fine-tune

Every decoder parameter (1,575,458,880 of them) was fine-tuned on 11 Japanese
songs, rather than training an adapter on top. Checkpoint: epoch
{FT_META.get('epoch', '?')}, Adafactor, lr {FT_META.get('lr', '?')},
fp32 masters over two 16 GB GPUs.

Captions were trained with the tag **`{TAG}`** in front, so keeping it
in the caption is what pulls the fine-tuned style in.
"""
    )

    with gr.Row():
        with gr.Column(scale=3):
            caption = gr.Textbox(
                label="Style caption",
                value=f"{TAG}, japanese vocaloid pop, female vocal, playful hyperpop",
                lines=2,
            )
            lyrics = gr.Textbox(
                label="Lyrics",
                value="[Verse]\nここに歌詞を入れる\n\n[Chorus]\nサビをここに",
                lines=12,
            )
        with gr.Column(scale=1):
            duration = gr.Slider(30, 240, value=120, step=10, label="Duration (s)")
            steps = gr.Slider(4, 32, value=8, step=1, label="Steps (turbo default 8)")
            seed = gr.Number(value=42, precision=0, label="Seed")
            language = gr.Dropdown(["ja", "en"], value="ja", label="Vocal language")
            go = gr.Button("Generate", variant="primary")

    out = gr.Audio(label="Output", type="filepath")
    go.click(generate, [caption, lyrics, duration, steps, seed, language], out)

    gr.Examples(
        examples=[[s["caption"], s["lyrics"]] for s in SAMPLES],
        inputs=[caption, lyrics],
        label="Prompts from the training set",
    )

    gr.Markdown("## Pre-rendered samples\n"
                "Generated locally with these same weights (seed 42, 8 steps).")
    for s in SAMPLES:
        if os.path.isfile(s["file"]):
            gr.Audio(value=s["file"], label=s["title"], type="filepath")
            gr.Markdown(f"*{s['caption']}*")

demo.queue().launch(mcp_server=True, show_error=True)