app.py

import os
import tempfile
import torch
import gradio as gr
from huggingface_hub import hf_hub_download, snapshot_download
import spaces

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.set_float32_matmul_precision('high')

# Download models from HuggingFace Hub on startup
def download_models():
    """Download all required model files from HuggingFace Hub."""
    cache_dir = os.environ.get("HF_HOME", os.path.expanduser("/tmp"))
    model_dir = os.path.join(cache_dir, "heartmula_models")

    if not os.path.exists(model_dir):
        os.makedirs(model_dir, exist_ok=True)

    # Download HeartMuLaGen (tokenizer and gen_config)
    print("Downloading HeartMuLaGen files...")
    for filename in ["tokenizer.json", "gen_config.json"]:
        hf_hub_download(
            repo_id="HeartMuLa/HeartMuLaGen",
            filename=filename,
            local_dir=model_dir,
        )

    # Download HeartMuLa-oss-3B
    print("Downloading HeartMuLa-oss-3B...")
    snapshot_download(
        repo_id="HeartMuLa/HeartMuLa-oss-3B-happy-new-year",
        local_dir=os.path.join(model_dir, "HeartMuLa-oss-3B"),
    )

    # Download HeartCodec-oss
    print("Downloading HeartCodec-oss...")
    snapshot_download(
        repo_id="HeartMuLa/HeartCodec-oss-20260123",
        local_dir=os.path.join(model_dir, "HeartCodec-oss"),
    )

    print("All models downloaded successfully!")
    return model_dir

from heartlib import HeartMuLaGenPipeline

model_dir = download_models()

# Determine device and dtype
if torch.cuda.is_available():
    device = torch.device("cuda")
    dtype = torch.bfloat16
else:
    device = torch.device("cpu")
    dtype = torch.float32

print(f"Loading pipeline on {device} with {dtype}...")
pipe = HeartMuLaGenPipeline.from_pretrained(
    model_dir,
    device=device,
    dtype=dtype,
    version="3B",
)
print("Pipeline loaded successfully!")


@spaces.GPU(duration=120)
def generate_music(
    lyrics: str,
    tags: str,
    max_duration_seconds: int,
    temperature: float,
    topk: int,
    cfg_scale: float,
    # progress=gr.Progress(track_tqdm=True),
):
    """Generate music from lyrics and tags."""
    if not lyrics.strip():
        raise gr.Error("Please enter some lyrics!")

    if not tags.strip():
        raise gr.Error("Please enter at least one tag!")

    # Create a temporary file for output
    with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as f:
        output_path = f.name

    max_audio_length_ms = max_duration_seconds * 1000

    with torch.no_grad():
        pipe(
            {
                "lyrics": lyrics,
                "tags": tags,
            },
            max_audio_length_ms=max_audio_length_ms,
            save_path=output_path,
            topk=topk,
            temperature=temperature,
            cfg_scale=cfg_scale,
        )

    return output_path


# Example lyrics
EXAMPLE_LYRICS = """[Intro]

[Verse]
Frustrated, degraded
Down before you're done
Rejection, depression
Can't get what you want

[Prechorus]
You ask me how I make my way
You ask me everywhere and why
You hang on every word I say
But the truth is like a lie

[Chorus]
Live to win 'til you die
'Til the light goes out in your eyes
Live to win, take it all
Just keep fighting 'til you fall

[Verse]
Obsessive, compulsive
Suffocate your mind
Confusion, delusions
Kill your dreams in time

[Prechorus]
You ask me how I took the pain
Crawled up from my lowest low
Step by step and day by day
'Til there's one last breath to go

[Chorus]
Live to win 'til you die
'Til the light goes out in your eyes
Live to win, take it all
Just keep fighting 'til you fall
Day by day, kickin' all the way
I'm not caving in
Let another round begin, live to win
Yeah, live, yeah, win

[Chorus]
Live to win 'til you die
'Til the light goes out in your eyes
Live to win, take it all
Just keep fighting 'til you fall
Day by day, kickin' all the way
I'm not caving in
Let another round begin, live to win
Live to win
Live to win
Yeah, live, yeah, win

[Outro]"""

EXAMPLE_TAGS = "rock, dark, agressive"

# Build the Gradio interface
with gr.Blocks(
    title="HeartMuLa Music Generator",
) as demo:
    gr.Markdown(
        """
        # HeartMuLa ❤️
        ## The best open-sourced music generation model in terms of lyrics controllability and music quality
        ## *Super-fast torch compiled version*
        """
    )
    gr.HTML("""
    <iframe width="100%" height="300" scrolling="no" frameborder="no" allow="autoplay; encrypted-media" src="https://w.soundcloud.com/player/?url=https%3A//api.soundcloud.com/playlists/soundcloud%253Aplaylists%253A2257612067&color=%23ff5500&auto_play=false&hide_related=false&show_comments=true&show_user=true&show_reposts=false&show_teaser=true&visual=true"></iframe><div style="font-size: 10px; color: #cccccc;line-break: anywhere;word-break: normal;overflow: hidden;white-space: nowrap;text-overflow: ellipsis; font-family: Interstate,Lucida Grande,Lucida Sans Unicode,Lucida Sans,Garuda,Verdana,Tahoma,sans-serif;font-weight: 100;"><a href="https://soundcloud.com/aleksandr-sigalov-61" title="Project Los Angeles" target="_blank" style="color: #cccccc; text-decoration: none;">Project Los Angeles</a> · <a href="https://soundcloud.com/aleksandr-sigalov-61/sets/heartmula-select" title="HeartMuLa Select" target="_blank" style="color: #cccccc; text-decoration: none;">HeartMuLa Select</a></div>
    """)
    gr.Markdown(
        """
        Generate music from lyrics and tags using [HeartMuLa](https://github.com/HeartMuLa/heartlib),
        an open-source music foundation model.

        **Instructions:**
        1. Enter your lyrics with structure tags like `[Verse]`, `[Chorus]`, `[Bridge]`, etc.
        2. Add comma-separated tags describing the music style (e.g., `piano,happy,romantic`)
        3. For best results use one tag per category
        3. Adjust generation parameters as needed
        4. Click "Generate Music" and wait for your song!

        *Note: Generation can take several minutes depending on the duration.*
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            lyrics_input = gr.Textbox(
                label="Lyrics",
                placeholder="Enter lyrics with structure tags like [Verse], [Chorus], etc.",
                lines=20,
                value=EXAMPLE_LYRICS,
            )

            tags_input = gr.Textbox(
                label="Tags",
                placeholder="piano,happy,romantic,synthesizer",
                value=EXAMPLE_TAGS,
                info="Comma-separated tags describing the music style (one tag per category)",
            )

            with gr.Accordion("Advanced Settings", open=False):
                max_duration = gr.Slider(
                    minimum=30,
                    maximum=240,
                    value=240,
                    step=10,
                    label="Max Duration (seconds)",
                    info="Maximum length of generated audio",
                )

                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=1.0,
                    step=0.1,
                    label="Temperature",
                    info="Higher = more creative, Lower = more consistent",
                )

                topk = gr.Slider(
                    minimum=1,
                    maximum=100,
                    value=50,
                    step=1,
                    label="Top-K",
                    info="Number of top tokens to sample from",
                )

                cfg_scale = gr.Slider(
                    minimum=1.0,
                    maximum=5.0,
                    value=3.0,
                    step=0.1,
                    label="CFG Scale",
                    info="Classifier-free guidance scale",
                )

            generate_btn = gr.Button("Generate Music", variant="primary", size="lg")

        with gr.Column(scale=1):
            audio_output = gr.Audio(
                label="Generated Music",
                type="filepath",
            )

            gr.Markdown(
                """
                ### Tips for Better Results
                - Use structured lyrics with section tags
                - Be specific with your style tags and use one tag per style category
                - Try different temperature values for variety
                - Shorter durations generate faster

                ### Offical Tags
                1. **GENRE** (95% - MANDATORY)
                   Examples: Pop, Rock, Electronic, Hiphop, Jazz, Classical, Techno, Trance, Ambient.
                2. **TIMBRE** (50% - Sound Texture)
                   Examples: Soft, Warm, Husky, Bright, Dark, Distorted.
                3. **GENDER** (37% - Vocal Character)
                   Examples: Male, Female.
                4. **MOOD** (32% - Emotional Vibe)
                   Examples: Happy, Sad, Energetic, Joyful, Melancholic, Relaxing, Dark.
                5. **INSTRUMENT** (25% - Dominant Sounds)
                   Examples: Piano, Synthesizer, Acoustic Guitar, Electric Guitar, Bass, Drums, Strings, Violin.
                6. **SCENE** (20% - Listening Context)
                   Examples: Dance, Workout, Dating, Study, Cinematic, Party.
                7. **REGION** (12% - Cultural Influence)
                   Examples: K-pop, Latin, Western.
                8. **TOPIC** (10% - Lyrical Theme)
                   Examples: Love, Summer, Heartbreak.
                """
            )

    generate_btn.click(
        fn=generate_music,
        inputs=[
            lyrics_input,
            tags_input,
            max_duration,
            temperature,
            topk,
            cfg_scale,
        ],
        outputs=audio_output,
    )

    gr.Markdown(
        """
        ---
        **Model:** [HeartMuLa-oss-3B](https://huggingface.co/HeartMuLa/HeartMuLa-oss-3B-happy-new-year) |
        **Paper:** [arXiv](https://arxiv.org/abs/2601.10547) |
        **Code:** [GitHub](https://github.com/HeartMuLa/heartlib)

        *Licensed under Apache 2.0*
        """
    )



demo.launch()