main.py

import spaces
import logging
import os
import pickle
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from tempfile import NamedTemporaryFile
import time
import typing as tp
import subprocess as sp
import torch
import gradio as gr
from audiocraft.data.audio_utils import f32_pcm, normalize_audio
from audiocraft.data.audio import audio_write
from audiocraft.models import JASCO
import os
from huggingface_hub import login

title = """# 🙋🏻‍♂️Welcome to 🌟Tonic's 🎼Jasco🎶AudioCraft Demo"""
description = """Facebook presents JASCO, a temporally controlled text-to-music generation model utilizing both symbolic and audio-based conditions. JASCO can generate high-quality music samples conditioned on global text descriptions along with fine-grained local controls. JASCO is based on the Flow Matching modeling paradigm together with a novel conditioning method, allowing for music generation controlled both locally (e.g., chords) and globally (text description). [run this demo locally](https://huggingface.co./spaces/Tonic/audiocraft?docker=true) or [embed this space](https://huggingface.co./spaces/Tonic/audiocraft?embed=true) or [duplicate this space](https://huggingface.co./spaces/Tonic/audiocraft?duplicate=true) to run it privately . you can also use this demo via API by clicking the link at the bottom of the page."""
join_us = """
## Join us:
🌟TeamTonic🌟 is always making cool demos! Join our active builder's 🛠️community 👻 
[![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/qdfnvSPcqP) 
On 🤗Huggingface: [MultiTransformer](https://huggingface.co./MultiTransformer) 
On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [MultiTonic](https://github.com/MultiTonic/thinking-dataset)
🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
"""

useage_instructions = """
## Overview
JASCO is a powerful text-to-music generation system that allows you to create music using text descriptions and various controls including chords, drums, and melody. This guide explains how to use each feature of the interface.

## Model Selection
Four different models are available:
1. `facebook/jasco-chords-drums-400M` - Basic model with chord and drum support (400M parameters)
2. `facebook/jasco-chords-drums-1B` - Enhanced model with chord and drum support (1B parameters)
3. `facebook/jasco-chords-drums-melody-400M` - Model with melody support (400M parameters)
4. `facebook/jasco-chords-drums-melody-1B` - Full-featured model with melody support (1B parameters)

## Input Controls

### 1. Text Description
- Enter a descriptive text about the music you want to generate
- Examples:
  - "80s pop with groovy synth bass and electric piano"
  - "Strings, woodwind, orchestral, symphony"
  - "Jazz quartet with walking bass and smooth piano"

### 2. Chord Progression
Format: `(Chord, Time), (Chord, Time), ...`
- Time is in seconds (0-10 seconds range)
- Example: `(C, 0.0), (D, 2.0), (F, 4.0), (Ab, 6.0), (Bb, 7.0), (C, 8.0)`

Supported chord types:
```python
Basic Chords: C, D, E, F, G, A, B
Minor Chords: Cm, Dm, Em, Fm, Gm, Am, Bm
Seventh Chords: C7, D7, E7, F7, G7, A7, B7
Major Seventh: Cmaj7, Dmaj7, Emaj7, Fmaj7, Gmaj7, Amaj7, Bmaj7
Minor Seventh: Cm7, Dm7, Em7, Fm7, Gm7, Am7, Bm7
Flat Chords: Ab, Bb (and their variations)
Special: N (No chord/silence)
```

### 3. Drums Input
Two options for adding drums:
1. File Upload:
   - Select "file" in Drums Input Source
   - Upload a WAV file containing drum patterns
   - Recommended length: 2-4 bars

2. Microphone Recording:
   - Select "mic" in Drums Input Source
   - Record drum patterns using your microphone
   - Keep recordings short and rhythmic

### 4. Melody Input
- Upload a melody salience matrix as a PyTorch tensor
- Format: Shape [n_melody_bins, T]
- File should be saved using `torch.save()`

### 5. Generation Parameters

#### Classifier Free Guidance (CFG) Controls:
- CFG ALL: Controls overall adherence to input conditions (default: 1.25)
  - Range: 1.0-3.0
  - Higher values = stronger conditioning

- CFG TEXT: Controls text conditioning strength (default: 2.5)
  - Range: 1.0-4.0
  - Higher values = closer match to text description

#### ODE Parameters:
- ODE Solver: Choose between 'euler' and 'dopri5'
  - euler: Faster, less accurate
  - dopri5: Slower, more accurate

- ODE Tolerance: Numerical precision (default: 1e-4)
  - Lower values = higher precision, slower generation

- Euler Steps: Number of steps for euler solver (default: 10)
  - Higher values = more accurate, slower generation

## Generation Process

1. Select a model based on your needs:
   - Use 400M models for faster generation
   - Use 1B models for higher quality
   - Choose melody-enabled models if using melody input

2. Enter your text description

3. Input chord progression:
```
Example:
(C, 0.0), (Am, 2.5), (F, 5.0), (G, 7.5)
```

4. (Optional) Add drums via file upload or microphone

5. (Optional) Upload melody matrix

6. Adjust generation parameters or use defaults

7. Click "Make Musix"

## Output
- The system generates two variations of your music
- Each generation is 10 seconds long
- Output is provided as WAV files
- You can download or play directly in the interface

## Tips for Best Results

1. Text Descriptions:
   - Be specific about instruments
   - Include genre information
   - Mention desired mood or style

2. Chord Progressions:
   - Use common progressions for better results
   - Space chords evenly
   - Include resolution points

3. Drums:
   - Use clean, clear drum patterns
   - Avoid complex patterns for better results
   - Keep volume levels consistent

4. Memory Management:
   - The interface caches models after first use
   - Switch models only when necessary
   - Clear browser cache if experiencing issues

## Example Usage

```python
# Example 1: Pop Music
Text: "Upbeat pop song with electric piano and synthesizer"
Chords: (C, 0.0), (Am, 2.5), (F, 5.0), (G, 7.5)
Model: facebook/jasco-chords-drums-400M

# Example 2: Orchestral
Text: "Epic orchestral piece with strong strings and brass"
Chords: (Cm, 0.0), (G, 3.0), (Bb, 6.0), (Cm, 8.0)
Model: facebook/jasco-chords-drums-melody-1B

# Example 3: Jazz
Text: "Smooth jazz quartet with walking bass and piano"
Chords: (Dmaj7, 0.0), (Em7, 2.5), (A7, 5.0), (Dmaj7, 7.5)
Model: facebook/jasco-chords-drums-1B
```

## Error Handling
- If generation fails, try:
  1. Simplifying chord progressions
  2. Reducing CFG values
  3. Using simpler text descriptions
  4. Checking input format compliance
  5. Refreshing the page

## Performance Considerations
- First generation may be slower due to model loading
- Subsequent generations with same model are faster
- Higher parameter models (1B) require more memory
- Melody-enabled models may be slower
"""

hf_token = os.environ.get('HFTOKEN')
if hf_token:
    login(token=hf_token)

MODEL = None
MAX_BATCH_SIZE = 12
INTERRUPTING = False

os.makedirs(os.path.join(os.path.dirname(__file__), "models"), exist_ok=True)

def generate_chord_mappings():
    # Define basic chord mappings
    basic_chords = ['N', 'C', 'Dm7', 'Am', 'F', 'D', 'Ab', 'Bb'] + ['UNK']
    chord_to_index = {chord: idx for idx, chord in enumerate(basic_chords)}
    
    # Save the mapping
    mapping_path = os.path.join(os.path.dirname(__file__), "models", "chord_to_index_mapping.pkl")
    os.makedirs(os.path.dirname(mapping_path), exist_ok=True)
    
    with open(mapping_path, "wb") as f:
        pickle.dump(chord_to_index, f)
    
    return mapping_path

def create_default_chord_mapping():
    """Create a basic chord-to-index mapping with common chords"""
    basic_chords = [
        'N', 'C', 'Cm', 'C7', 'Cmaj7', 'Cm7',
        'D', 'Dm', 'D7', 'Dmaj7', 'Dm7',
        'E', 'Em', 'E7', 'Emaj7', 'Em7',
        'F', 'Fm', 'F7', 'Fmaj7', 'Fm7',
        'G', 'Gm', 'G7', 'Gmaj7', 'Gm7',
        'A', 'Am', 'A7', 'Amaj7', 'Am7',
        'B', 'Bm', 'B7', 'Bmaj7', 'Bm7',
        'Ab', 'Abm', 'Ab7', 'Abmaj7', 'Abm7',
        'Bb', 'Bbm', 'Bb7', 'Bbmaj7', 'Bbm7',
        'UNK'
    ]
    return {chord: idx for idx, chord in enumerate(basic_chords)}

def initialize_chord_mapping():
    """Initialize chord mapping file if it doesn't exist"""
    mapping_dir = os.path.join(os.path.dirname(__file__), "models")
    os.makedirs(mapping_dir, exist_ok=True)
    
    mapping_file = os.path.join(mapping_dir, "chord_to_index_mapping.pkl")
    
    if not os.path.exists(mapping_file):
        chord_to_index = create_default_chord_mapping()
        with open(mapping_file, "wb") as f:
            pickle.dump(chord_to_index, f)
    
    return mapping_file

def validate_chord(chord, chord_mapping):
    if chord not in chord_mapping:
        return 'UNK'
    return chord

mapping_file = initialize_chord_mapping()
os.environ['AUDIOCRAFT_CHORD_MAPPING'] = mapping_file

def chords_string_to_list(chords: str):
    if chords == '':
        return []
    chords = chords.replace('[', '').replace(']', '').replace(' ', '')
    chrd_times = [x.split(',') for x in chords[1:-1].split('),(')]
    
    # Load chord mapping
    mapping_path = os.path.join(os.path.dirname(__file__), "models", "chord_to_index_mapping.pkl")
    with open(mapping_path, 'rb') as f:
        chord_mapping = pickle.load(f)
    
    return [(validate_chord(x[0], chord_mapping), float(x[1])) for x in chrd_times]

# Wrap subprocess call to clean logs
_old_call = sp.call

def _call_nostderr(*args, **kwargs):
    kwargs['stderr'] = sp.DEVNULL
    kwargs['stdout'] = sp.DEVNULL
    _old_call(*args, **kwargs)

sp.call = _call_nostderr

# Preallocate process pool
pool = ProcessPoolExecutor(4)
pool.__enter__()

def interrupt():
    global INTERRUPTING
    INTERRUPTING = True

class FileCleaner:
    def __init__(self, file_lifetime: float = 3600):
        self.file_lifetime = file_lifetime
        self.files = []

    def add(self, path: tp.Union[str, Path]):
        self._cleanup()
        self.files.append((time.time(), Path(path)))

    def _cleanup(self):
        now = time.time()
        for time_added, path in list(self.files):
            if now - time_added > self.file_lifetime:
                if path.exists():
                    path.unlink()
                self.files.pop(0)
            else:
                break

file_cleaner = FileCleaner()

def chords_string_to_list(chords: str):
    if chords == '':
        return []
    chords = chords.replace('[', '').replace(']', '').replace(' ', '')
    chrd_times = [x.split(',') for x in chords[1:-1].split('),(')]
    return [(x[0], float(x[1])) for x in chrd_times]

# Create necessary directories
os.makedirs("models", exist_ok=True)

@spaces.GPU
def load_model(version='facebook/jasco-chords-drums-400M'):
    global MODEL
    print("Loading model", version)
    if MODEL is None or MODEL.name != version:
        MODEL = None
        
        # Setup model directory
        model_dir = os.path.join(os.path.dirname(__file__), "models")
        os.makedirs(model_dir, exist_ok=True)
        
        # Generate and save chord mappings
        chord_mapping_path = os.path.join(model_dir, "chord_to_index_mapping.pkl")
        if not os.path.exists(chord_mapping_path):
            chord_mapping_path = generate_chord_mappings()
        
        try:
            # Initialize JASCO with the chord mapping path
            MODEL = JASCO.get_pretrained(
                version,
                device='cuda',
                chords_mapping_path=chord_mapping_path
            )
            MODEL.name = version
        except Exception as e:
            raise gr.Error(f"Error loading model: {str(e)}")
        
        if MODEL is None:
            raise gr.Error("Failed to load model")
    
    return MODEL

@spaces.GPU
def _do_predictions(texts, chords, melody_matrix, drum_prompt, progress=False, gradio_progress=None, **gen_kwargs):
    MODEL.set_generation_params(**gen_kwargs)
    be = time.time()

    chords = chords_string_to_list(chords)

    if melody_matrix is not None:
        melody_matrix = torch.load(melody_matrix.name, weights_only=True)
        if len(melody_matrix.shape) != 2:
            raise gr.Error(f"Melody matrix should be a torch tensor of shape [n_melody_bins, T]; got: {melody_matrix.shape}")
        if melody_matrix.shape[0] > melody_matrix.shape[1]:
            melody_matrix = melody_matrix.permute(1, 0)

    if drum_prompt is None:
        preprocessed_drums_wav = None
        drums_sr = 32000
    else:
        drums_sr, drums = drum_prompt[0], f32_pcm(torch.from_numpy(drum_prompt[1])).t()
        if drums.dim() == 1:
            drums = drums[None]
        drums = normalize_audio(drums, strategy="loudness", loudness_headroom_db=16, sample_rate=drums_sr)
        preprocessed_drums_wav = drums

    try:
        outputs = MODEL.generate_music(descriptions=texts, chords=chords,
                                     drums_wav=preprocessed_drums_wav,
                                     melody_salience_matrix=melody_matrix,
                                     drums_sample_rate=drums_sr, progress=progress)
    except RuntimeError as e:
        raise gr.Error("Error while generating " + e.args[0])

    outputs = outputs.detach().cpu().float()
    out_wavs = []
    for output in outputs:
        with NamedTemporaryFile("wb", suffix=".wav", delete=False) as file:
            audio_write(
                file.name, output, MODEL.sample_rate, strategy="loudness",
                loudness_headroom_db=16, loudness_compressor=True, add_suffix=False)
            out_wavs.append(file.name)
            file_cleaner.add(file.name)
    return out_wavs

@spaces.GPU
def predict_full(model, text, chords_sym, melody_file,
                drums_file, drums_mic, drum_input_src,
                cfg_coef_all, cfg_coef_txt,
                ode_rtol, ode_atol,
                ode_solver, ode_steps,
                progress=gr.Progress()):
    global INTERRUPTING
    INTERRUPTING = False
    progress(0, desc="Loading model...")
    load_model(model)

    max_generated = 0

    def _progress(generated, to_generate):
        nonlocal max_generated
        max_generated = max(generated, max_generated)
        progress((min(max_generated, to_generate), to_generate))
        if INTERRUPTING:
            raise gr.Error("Interrupted.")
    
    MODEL.set_custom_progress_callback(_progress)

    drums = drums_mic if drum_input_src == "mic" else drums_file
    wavs = _do_predictions(
        texts=[text] * 2,
        chords=chords_sym,
        drum_prompt=drums,
        melody_matrix=melody_file,
        progress=True,
        gradio_progress=progress,
        cfg_coef_all=cfg_coef_all,
        cfg_coef_txt=cfg_coef_txt,
        ode_rtol=ode_rtol,
        ode_atol=ode_atol,
        euler=ode_solver == 'euler',
        euler_steps=ode_steps)

    return wavs

with gr.Blocks() as demo:
    gr.Markdown(title)
    with gr.Row():
        with gr.Column(): 
            with gr.Group():
                gr.Markdown(description)
        with gr.Column():
            with gr.Group():
                gr.Markdown(join_us)
    with gr.Row():
        with gr.Accordion(open=False, label="Useage Instructions"):
            gr.Markdown(useage_instructions)

    with gr.Row():
        with gr.Column():
            submit = gr.Button("🎼Make Music")
            interrupt_btn = gr.Button("❌Interrupt")

        with gr.Column():
            audio_output_0 = gr.Audio(label="🎼Jasco Stem 1", type='filepath')
            audio_output_1 = gr.Audio(label="🎼Jasco Stem 2", type='filepath')

    with gr.Row():
        with gr.Column():
            text = gr.Text(label="Input Text",
                          value="Strings, woodwind, orchestral, symphony.",
                          interactive=True)
        with gr.Column():
            model = gr.Radio([
                'facebook/jasco-chords-drums-400M',
                'facebook/jasco-chords-drums-1B',
                'facebook/jasco-chords-drums-melody-400M',
                'facebook/jasco-chords-drums-melody-1B'
            ], label="Model", value='facebook/jasco-chords-drums-melody-400M')

    gr.Markdown("### Chords Conditions")
    chords_sym = gr.Text(
        label="🎼Chord Progression",
        value="(C, 0.0), (D, 2.0), (F, 4.0), (Ab, 6.0), (Bb, 7.0), (C, 8.0)",
        interactive=True
    )

    gr.Markdown("### 🥁Drums")
    with gr.Row():
        drum_input_src = gr.Radio(["file", "mic"], value="file", label="🥁Drums Input Source")
        drums_file = gr.Audio(sources=["upload"], type="numpy", label="🥁Drums File")
        drums_mic = gr.Audio(sources=["microphone"], type="numpy", label="🥁🎙️Drums Mic")

    gr.Markdown("### 🎶Melody Conditions")
    melody_file = gr.File(label="Melody File")

    with gr.Row():
        cfg_coef_all = gr.Number(label="CFG ALL", value=1.25, step=0.25)
        cfg_coef_txt = gr.Number(label="CFG TEXT", value=2.5, step=0.25)
        ode_tol = gr.Number(label="ODE Tolerance", value=1e-4, step=1e-5)
        ode_solver = gr.Radio(['euler', 'dopri5'], label="ODE Solver", value='euler')
        ode_steps = gr.Number(label="Euler Steps", value=10, step=1)

    submit.click(
        fn=predict_full,
        inputs=[
            model, text, chords_sym, melody_file,
            drums_file, drums_mic, drum_input_src,
            cfg_coef_all, cfg_coef_txt,
            ode_tol, ode_tol, ode_solver, ode_steps
        ],
        outputs=[audio_output_0, audio_output_1]
    )
    
    interrupt_btn.click(fn=interrupt, queue=False)

    gr.Examples(
        examples=[
            [
                "80s pop with groovy synth bass and electric piano",
                "(N, 0.0), (C, 0.32), (Dm7, 3.456), (Am, 4.608), (F, 8.32), (C, 9.216)",
                None,
                None,
            ],
            [
                "Strings, woodwind, orchestral, symphony.",
                "(C, 0.0), (D, 2.0), (F, 4.0), (Ab, 6.0), (Bb, 7.0), (C, 8.0)",
                None,
                None,
            ],
        ],
        inputs=[text, chords_sym, melody_file, drums_file],
        outputs=[audio_output_0, audio_output_1]
    )

demo.queue().launch(ssr_mode=False)