app.py

import gradio as gr

DESCRIPTION = """
# audio-studio

Local music generation studio wrapping **ACE-Step 1.5**.

Text-to-music, reference-audio melody borrowing, and reference-audio cover generation,
with configurable artist/album/track/genre texture hints.

## Features

- **Text-to-music**: Describe what you want, get a WAV file
- **Melody borrowing**: Provide a reference audio to borrow its BPM, key, and duration
- **Cover mode**: Full sound/texture conditioning from a reference (requires ace-step backend)
- **Texture hints**: Artist, album, track, and genre hints appended to prompts
- **Two backends**: `mlx-serve` (default, ~2.9GB) or `ace-step` (~17.8GB, for cover mode)

## Quick Start

```bash
# Clone and setup
git clone https://github.com/sm3222/audio-studio.git
cd audio-studio
uv sync

# Start the backend ( mlx-serve by default)
mlx-serve serve --port 8082 --model-dir ~/.mlx-serve/models

# Generate music
uv run audio-studio "a calm piano piece" --artist "Bob Marley" --output out.wav
```

## Links

- [GitHub Repository](https://github.com/sm3222/audio-studio)
- [ACE-Step 1.5 Paper](https://arxiv.org/abs/2503.02057)
- [mlx-serve](https://github.com/lucasnewman/mlx-serve)
"""

EXAMPLE_PROMPTS = [
    "a calm piano piece with soft strings",
    "upbeat electronic dance music with synth leads",
    "acoustic folk guitar with gentle vocals",
    "heavy rock drums with distorted guitar riffs",
    "ambient atmospheric soundscape with pads",
]


def build_command(prompt, artists, albums, tracks, genres, live, instrumental, backend, reference, cover):
    """Build the CLI command from user inputs."""
    parts = ["uv run audio-studio"]

    if prompt:
        parts.append(f'"{prompt}"')

    for artist in artists.split(","):
        artist = artist.strip()
        if artist:
            parts.append(f'--artist "{artist}"')

    for album in albums.split(","):
        album = album.strip()
        if album:
            parts.append(f'--album "{album}"')

    for track in tracks.split(","):
        track = track.strip()
        if track:
            parts.append(f'--track "{track}"')

    for genre in genres.split(","):
        genre = genre.strip()
        if genre:
            parts.append(f'--genre "{genre}"')

    if live:
        parts.append("--live")

    if not instrumental:
        parts.append("--no-instrumental")

    if backend != "mlx-serve":
        parts.append(f"--backend {backend}")

    if reference:
        parts.append(f'--reference "{reference}"')

    if cover:
        parts.append("--cover")

    parts.append("--output out.wav")
    return " \\\n  ".join(parts)


with gr.Blocks(title="audio-studio") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Tab("Command Builder"):
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the music you want to generate...",
                    lines=3,
                )
                artists = gr.Textbox(
                    label="Artists (comma-separated)",
                    placeholder="e.g., Bob Marley, Lee Perry",
                )
                albums = gr.Textbox(label="Albums (comma-separated)")
                tracks = gr.Textbox(label="Tracks (comma-separated)")
                genres = gr.Textbox(label="Genres (comma-separated)")

            with gr.Column():
                live = gr.Checkbox(label="Live performance texture", value=False)
                instrumental = gr.Checkbox(label="Instrumental (no vocals)", value=True)
                backend = gr.Radio(
                    choices=["mlx-serve", "ace-step"],
                    value="mlx-serve",
                    label="Backend",
                )
                reference = gr.Textbox(
                    label="Reference audio path (optional)",
                    placeholder="/path/to/reference.wav",
                )
                cover = gr.Checkbox(
                    label="Cover mode (requires ace-step backend)",
                    value=False,
                )

        build_btn = gr.Button("Build Command", variant="primary")
        command_output = gr.Code(label="Generated Command")

        build_btn.click(
            fn=build_command,
            inputs=[prompt, artists, albums, tracks, genres, live, instrumental, backend, reference, cover],
            outputs=command_output,
        )

        gr.Examples(
            examples=[[p] for p in EXAMPLE_PROMPTS],
            inputs=prompt,
            label="Example Prompts",
        )

    with gr.Tab("Examples"):
        gr.Markdown("""
## Example Generations

```bash
# Simple text-to-music
uv run audio-studio "a calm piano piece" --output out.wav

# With artist hints
uv run audio-studio "heavy drums and bass" --artist "Bob Marley" --output out.wav

# With reference audio (melody only)
uv run audio-studio "electronic remix" --reference original.wav --output out.wav

# Full cover mode (requires ace-step backend)
uv run audio-studio "jazz cover" --reference original.wav --cover --backend ace-step --output out.wav

# Multiple hints
uv run audio-studio "ambient soundscape" \\
  --artist "Lee Perry" \\
  --genre "ambient" \\
  --live \\
  --output out.wav
```

## Texture Hints

Edit `config/hints.json` to add your own artist/album/track/genre hints:

```json
{
  "artist": {
    "bob marley": "heavy acoustic drums, analog tape echo",
    "lee perry": "heavy acoustic drums, analog tape echo"
  },
  "genre": {
    "ambient": "atmospheric pads, reverb-heavy, slow evolution"
  }
}
```
""")

    with gr.Tab("Architecture"):
        gr.Markdown("""
## Architecture

```
audio-studio/
├── src/audio_studio/
│   ├── cli.py              # CLI entrypoint
│   ├── prompt.py           # Prompt construction with hints
│   ├── analysis.py         # Reference audio analysis (librosa)
│   ├── loop.py             # Generation loop (submit → poll → download)
│   ├── acestep_client.py   # ACE-Step API client
│   └── mlx_serve_client.py # mlx-serve API client
├── config/hints.json       # Texture hint definitions
└── output/                 # Generated WAV files
```

## Backends

| Backend | Memory | Cover Mode | Notes |
|---------|--------|------------|-------|
| `mlx-serve` | ~2.9GB | No | Default, MLX-native, fast |
| `ace-step` | ~17.8GB | Yes | Original PyTorch/MPS, heavy |

## Dependencies

- Python 3.12+
- librosa (audio analysis)
- mlx-serve or ace-step API server running
- See `pyproject.toml` for full list
""")

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Soft())