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
import tempfile
from huggingface_hub import login
from pathlib import Path
import sys
import shutil

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
"""

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

MODEL = None
INTERRUPTING = False

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

# Set up cache directory
CACHE_DIR = os.path.join(tempfile.gettempdir(), 'audiocraft_cache')
os.makedirs(CACHE_DIR, exist_ok=True)

# Set environment variables
os.environ['AUDIOCRAFT_CACHE_DIR'] = CACHE_DIR
os.environ['TORCH_HOME'] = CACHE_DIR
os.environ['HF_HOME'] = os.path.join(CACHE_DIR, 'huggingface')
os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
os.environ['XDG_CACHE_HOME'] = CACHE_DIR

# Create necessary subdirectories
for subdir in ['models', 'cache', 'huggingface', 'drum_cache', 'transformers']:
    os.makedirs(os.path.join(CACHE_DIR, subdir), exist_ok=True)

def cleanup_cache():
    """Clean up temporary cache files while preserving chord mapping"""
    try:
        # Save the chord mapping if it exists
        mapping_file = os.path.join(CACHE_DIR, "chord_to_index_mapping.pkl")
        temp_mapping = None
        if os.path.exists(mapping_file):
            with open(mapping_file, 'rb') as f:
                temp_mapping = pickle.load(f)
        
        # Clean the cache directory
        if os.path.exists(CACHE_DIR):
            shutil.rmtree(CACHE_DIR)
        
        # Recreate necessary directories
        os.makedirs(CACHE_DIR, exist_ok=True)
        for subdir in ['models', 'cache', 'huggingface', 'drum_cache', 'transformers']:
            os.makedirs(os.path.join(CACHE_DIR, subdir), exist_ok=True)
        
        # Restore the chord mapping
        if temp_mapping is not None:
            with open(mapping_file, 'wb') as f:
                pickle.dump(temp_mapping, f)
                
    except Exception as e:
        print(f"Error cleaning cache: {e}")

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]

def patch_jasco_cache():
    """Monkey patch JASCO cache initialization"""
    from audiocraft.modules import jasco_conditioners
    
    if hasattr(jasco_conditioners, 'DrumConditioner'):
        original_init = jasco_conditioners.DrumConditioner.__init__
    elif hasattr(jasco_conditioners, 'DrumsConditioner'):
        original_init = jasco_conditioners.DrumsConditioner.__init__
    else:
        print("Warning: Could not find DrumConditioner class")
        return
    
    def new_init(self, *args, **kwargs):
        if 'cache_path' in kwargs:
            kwargs['cache_path'] = os.path.join(CACHE_DIR, 'drum_cache')
        return original_init(self, *args, **kwargs)
    
    if hasattr(jasco_conditioners, 'DrumConditioner'):
        jasco_conditioners.DrumConditioner.__init__ = new_init
    elif hasattr(jasco_conditioners, 'DrumsConditioner'):
        jasco_conditioners.DrumsConditioner.__init__ = new_init

# Apply the patch
patch_jasco_cache()

def ensure_chord_mapping():
    """Ensure chord mapping file exists in all necessary locations"""
    chord_mapping = create_default_chord_mapping()
    
    # Define all necessary paths
    paths = [
        os.path.join(CACHE_DIR, "chord_to_index_mapping.pkl"),
        os.path.join(os.path.dirname(__file__), "models", "chord_to_index_mapping.pkl"),
        "/tmp/chord_to_index_mapping.pkl",  # Additional fallback location
        os.path.join(CACHE_DIR, "models", "chord_to_index_mapping.pkl")
    ]
    
    # Create directories and save mapping to all locations
    for path in paths:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "wb") as f:
            pickle.dump(chord_mapping, f)
        print(f"Created chord mapping at: {path}")
    
    return paths[0]  # Return the primary mapping path

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

@spaces.GPU
def load_model(version='facebook/jasco-chords-drums-melody-400M'):
    global MODEL
    print("Loading model", version)
    if MODEL is None or MODEL.name != version:
        MODEL = None
        try:
            # Set up cache paths
            cache_path = os.path.join(CACHE_DIR, version.replace('/', '_'))
            os.makedirs(cache_path, exist_ok=True)
            
            # Ensure chord mapping exists and set environment
            mapping_file = ensure_chord_mapping()
            os.environ['AUDIOCRAFT_CHORD_MAPPING'] = mapping_file
            
            # Create symbolic links to ensure the mapping is found
            target_paths = [
                os.path.join(cache_path, "chord_to_index_mapping.pkl"),
                os.path.join(os.path.dirname(__file__), "chord_to_index_mapping.pkl")
            ]
            
            for target in target_paths:
                try:
                    if not os.path.exists(target):
                        os.symlink(mapping_file, target)
                except Exception as e:
                    print(f"Warning: Could not create symlink at {target}: {e}")
            
            # Set environment variables
            os.environ['AUDIOCRAFT_CACHE_DIR'] = cache_path
            os.environ['TRANSFORMERS_CACHE'] = os.path.join(cache_path, 'transformers')
            
            # Initialize JASCO with explicit paths
            kwargs = {
                'device': 'cuda',
                'chords_mapping_path': mapping_file
            }
            
            # Load the model
            MODEL = JASCO.get_pretrained(version, **kwargs)
            MODEL.name = version
            
            # Verify chord mapping is loaded
            if not hasattr(MODEL, 'chord_to_index'):
                with open(mapping_file, 'rb') as f:
                    MODEL.chord_to_index = pickle.load(f)
            
            print("Model loaded successfully")
            
        except Exception as e:
            print(f"Detailed error loading model: {str(e)}")
            raise gr.Error(f"Error loading model: {str(e)}")
        
        if MODEL is None:
            raise gr.Error("Failed to load model")
    return MODEL

# Initialize environment and mappings at startup
def initialize_environment():
    """Initialize all necessary environment variables and files"""
    # Create cache directory structure
    for subdir in ['models', 'cache', 'huggingface', 'drum_cache', 'transformers']:
        os.makedirs(os.path.join(CACHE_DIR, subdir), exist_ok=True)
    
    # Initialize chord mapping
    mapping_file = ensure_chord_mapping()
    os.environ['AUDIOCRAFT_CHORD_MAPPING'] = mapping_file
    
    # Set up environment variables
    os.environ['AUDIOCRAFT_CACHE_DIR'] = CACHE_DIR
    os.environ['TORCH_HOME'] = CACHE_DIR
    os.environ['HF_HOME'] = os.path.join(CACHE_DIR, 'huggingface')
    os.environ['TRANSFORMERS_CACHE'] = os.path.join(CACHE_DIR, 'transformers')
    os.environ['XDG_CACHE_HOME'] = CACHE_DIR
    
    print("Environment initialized successfully")
    return mapping_file

# Call initialization at startup
mapping_file = initialize_environment()

class ModelLoadingError(Exception):
    pass

def handle_model_error(func):
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            print(f"Error in {func.__name__}: {str(e)}")
            raise ModelLoadingError(f"Failed to load model: {str(e)}")
    return wrapper

@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]
    )

if __name__ == "__main__":
    try:
        # Initialize environment
        mapping_file = initialize_environment()
        print(f"Using chord mapping file: {mapping_file}")
        print(f"Chord mapping exists: {os.path.exists(mapping_file)}")
        
        # Launch the demo
        demo.queue().launch(ssr_mode=False)
    except Exception as e:
        print(f"Error launching demo: {str(e)}")
    finally:
        cleanup_cache()