app.py

import os
import tempfile
import gradio as gr
import spaces
import torch

# =============================================================================
# Both models are loaded and run LOCALLY inside the ZeroGPU allocation.
# Neither is served through HF's Inference Providers, so InferenceClient
# can't reach either one — that's why this uses heartlib / diffusers directly.
#
# requirements.txt needs:
#   gradio
#   spaces
#   torch
#   huggingface_hub
#   soundfile
#   heartlib @ git+https://github.com/HeartMuLa/heartlib.git
#   diffusers @ git+https://github.com/huggingface/diffusers.git   # AceStepPipeline isn't in a release yet
#   transformers
#   accelerate
#
# packages.txt needs:
#   ffmpeg
# =============================================================================

from huggingface_hub import snapshot_download

MODEL_HEARTMULA = "HeartMuLa"
MODEL_ACESTEP = "ACE-Step 1.5"

HEARTMULA_CKPT_DIR = os.environ.get("HEARTMULA_CKPT_DIR", "./ckpt-heartmula")
# NOTE: the diffusers-loadable checkpoint lives under a different repo than
# the raw "ACE-Step/Ace-Step1.5" release. This is the "turbo" (guidance-distilled,
# 8-step) variant — fastest for research iteration. "base"/"sft" variants exist
# too (ACE-Step/acestep-v15-base, ACE-Step/acestep-v15-sft) and use real CFG.
ACESTEP_REPO = "ACE-Step/acestep-v15-xl-turbo-diffusers"

_heartmula_pipe = None
_acestep_pipe = None


# --- HeartMuLa -----------------------------------------------------------
def _get_heartmula_pipeline():
    global _heartmula_pipe
    if _heartmula_pipe is None:
        from heartlib import HeartMuLaGenPipeline

        if not os.path.exists(os.path.join(HEARTMULA_CKPT_DIR, "gen_config.json")):
            os.makedirs(HEARTMULA_CKPT_DIR, exist_ok=True)
            snapshot_download("HeartMuLa/HeartMuLaGen", local_dir=HEARTMULA_CKPT_DIR)
            snapshot_download(
                "HeartMuLa/HeartMuLa-oss-3B-happy-new-year",
                local_dir=os.path.join(HEARTMULA_CKPT_DIR, "HeartMuLa-oss-3B"),
            )
            snapshot_download(
                "HeartMuLa/HeartCodec-oss-20260123",
                local_dir=os.path.join(HEARTMULA_CKPT_DIR, "HeartCodec-oss"),
            )

        _heartmula_pipe = HeartMuLaGenPipeline.from_pretrained(
            HEARTMULA_CKPT_DIR,
            device="cuda",
            # heartlib takes per-component dtype (matches its CLI's separate
            # --mula_dtype / --codec_dtype flags), not a single torch_dtype
            dtype={"mula": torch.bfloat16, "codec": torch.float32},
            version="3B",
            lazy_load=True,
        )
    return _heartmula_pipe


def _run_heartmula(prompt: str, lyrics: str, tags: str, duration: int):
    import soundfile as sf

    pipe = _get_heartmula_pipeline()
    # NOTE: verify current kwarg names against examples/run_music_generation.py
    # in https://github.com/HeartMuLa/heartlib — these mirror its CLI flags.
    audio, sample_rate = pipe.generate(
        lyrics=lyrics,
        tags=tags,
        max_audio_length_ms=duration * 1000,
        topk=50,
        temperature=1.0,
        cfg_scale=1.5,
    )
    tmp_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    sf.write(tmp_path, audio, sample_rate)
    return tmp_path


# --- ACE-Step 1.5 ---------------------------------------------------------
def _get_acestep_pipeline():
    global _acestep_pipe
    if _acestep_pipe is None:
        from diffusers import AceStepPipeline

        _acestep_pipe = AceStepPipeline.from_pretrained(
            ACESTEP_REPO,
            torch_dtype=torch.bfloat16,
        ).to("cuda")
    return _acestep_pipe


def _run_acestep(prompt: str, lyrics: str, tags: str, duration: int):
    import soundfile as sf

    pipe = _get_acestep_pipeline()
    output = pipe(
        prompt=prompt,
        lyrics=lyrics or "",
        audio_duration=float(duration),
        # turbo checkpoint: 8 steps is the intended default, guidance_scale
        # is ignored automatically for turbo weights (per diffusers docs)
        num_inference_steps=8,
        guidance_scale=7.0,
        shift=3.0,
        generator=torch.Generator(device="cuda").manual_seed(0),
    )
    # confirmed shape: audios[0] is (channels, samples); transpose for soundfile
    audio = output.audios[0].T.cpu().float().numpy()
    sample_rate = pipe.sample_rate

    tmp_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    sf.write(tmp_path, audio, sample_rate)
    return tmp_path


# --- Shared entry point, still the only thing wrapped in @spaces.GPU -----
@spaces.GPU(duration=180)
def generate_music(model_choice: str, prompt: str, lyrics: str, tags: str, duration: int):
    try:
        if model_choice == MODEL_HEARTMULA:
            path = _run_heartmula(prompt, lyrics, tags, duration)
            msg = "Generated with HeartMuLa-oss-3B."
        else:
            path = _run_acestep(prompt, lyrics, tags, duration)
            msg = "Generated with ACE-Step 1.5."
        return path, msg
    except Exception as e:
        return None, f"Generation error ({model_choice}): {e}"


def build_ui():
    with gr.Blocks() as demo:
        gr.Markdown(
            "# Music Generation Research (ZeroGPU)\n"
            "Compare HeartMuLa and ACE-Step 1.5, both running locally on a "
            "ZeroGPU-allocated GPU (no Inference API involved)."
        )

        with gr.Row():
            with gr.Column(scale=2):
                model_choice = gr.Radio(
                    choices=[MODEL_HEARTMULA, MODEL_ACESTEP],
                    value=MODEL_HEARTMULA,
                    label="Model",
                )
                prompt = gr.Textbox(
                    label="Prompt / style description (used by ACE-Step 1.5)",
                    placeholder="An upbeat synthwave track with driving drums",
                    lines=2,
                )
                lyrics = gr.Textbox(
                    label="Lyrics (used by both models)",
                    placeholder="[Verse]\n...\n[Chorus]\n...",
                    lines=8,
                )
                tags = gr.Textbox(
                    label="Tags — HeartMuLa only, comma-separated, no spaces",
                    placeholder="piano,happy,synthesizer,romantic",
                    visible=True,
                )
                duration = gr.Slider(5, 60, step=5, value=15, label="Duration (seconds)")
                generate_button = gr.Button("Generate audio")

            with gr.Column(scale=1):
                audio_output = gr.Audio(label="Generated audio", type="filepath")
                status = gr.Textbox(label="Status", interactive=False)

        # Only HeartMuLa uses the tags field — hide it for ACE-Step to avoid confusion
        def _toggle_tags(choice):
            return gr.update(visible=(choice == MODEL_HEARTMULA))

        model_choice.change(fn=_toggle_tags, inputs=model_choice, outputs=tags)

        generate_button.click(
            fn=generate_music,
            inputs=[model_choice, prompt, lyrics, tags, duration],
            outputs=[audio_output, status],
        )

    return demo


demo = build_ui()

if __name__ == "__main__":
    demo.launch()