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()