import os import torch import torchaudio import psutil import time import sys import numpy as np import gc import gradio as gr from pydub import AudioSegment import soundfile as sf import pyloudnorm as pyln from audiocraft.models import MusicGen from torch.amp import autocast import json import configparser import random import string import uvicorn from fastapi import FastAPI, HTTPException from fastapi.responses import FileResponse from pydantic import BaseModel import multiprocessing import re import datetime import warnings # ============================== # Warnings & Multiprocessing # ============================== warnings.filterwarnings("ignore", category=UserWarning) multiprocessing.set_start_method('spawn', force=True) # ============================== # CUDA / PyTorch Runtime Settings # ============================== os.environ["TORCH_NN_UTILS_LOG_LEVEL"] = "0" os.environ["CUDA_LAUNCH_BLOCKING"] = "1" os.environ["CUDA_MODULE_LOADING"] = "LAZY" os.environ["TORCH_USE_CUDA_DSA"] = "1" # Stronger allocator settings to reduce fragmentation and avoid small splits os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128,garbage_collection_threshold:0.8,expandable_segments:True" # Support a range of architectures (Turing/Ampere/Ada) os.environ["TORCH_CUDA_ARCH_LIST"] = "7.5;8.0;8.6;8.9" # Prefer TF32 on Ampere+ (perf) — also helps allocator behavior try: torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.benchmark = True except Exception: pass # ============================== # Version / Device Checks # ============================== def _parse_version_triplet(s: str): m = re.findall(r"\d+", s) m = [int(x) for x in m[:3]] while len(m) < 3: m.append(0) return tuple(m) if _parse_version_triplet(torch.__version__) < (2, 0, 0): print(f"ERROR: PyTorch {torch.__version__} incompatible. Need >=2.0.0.") sys.exit(1) device = "cuda" if torch.cuda.is_available() else "cpu" if device != "cuda": print("ERROR: CUDA required. CPU disabled.") sys.exit(1) cc_major, cc_minor = torch.cuda.get_device_capability(0) if cc_major < 7: print(f"ERROR: GPU Compute Capability {torch.cuda.get_device_capability(0)} unsupported. Need >=7.0.") sys.exit(1) gpu_name = torch.cuda.get_device_name(0) print(f"Using GPU: {gpu_name} (CUDA {torch.version.cuda}, Compute Capability {(cc_major, cc_minor)})") # Choose autocast dtype based on hardware support try: bf16_supported = torch.cuda.is_bf16_supported() except Exception: bf16_supported = False AUTOCAST_DTYPE = torch.bfloat16 if bf16_supported and cc_major >= 8 else torch.float16 # ============================== # Resource Monitoring # ============================== def print_resource_usage(stage: str): try: alloc = torch.cuda.memory_allocated() / (1024 ** 3) reserved = torch.cuda.memory_reserved() / (1024 ** 3) except Exception: alloc, reserved = 0.0, 0.0 print(f"--- {stage} ---") print(f"GPU Memory: {alloc:.2f} GB allocated, {reserved:.2f} GB reserved") print(f"CPU: {psutil.cpu_percent()}% | Memory: {psutil.virtual_memory().percent}%") print("---------------") # ============================== # Output & Metadata # ============================== output_dir = "mp3" os.makedirs(output_dir, exist_ok=True) metadata_file = os.path.join(output_dir, "songs_metadata.json") api_status = "idle" # ============================== # Prompt Variables # ============================== prompt_variables = { 'style': [ 'epic', 'gritty', 'smooth', 'lush', 'raw', 'intimate', 'driving', 'moody', 'psychedelic', 'uplifting', 'melancholic', 'aggressive', 'dreamy', 'retro', 'futuristic', 'energetic', 'brooding', 'euphoric', 'jazzy', 'cinematic', 'somber', 'triumphant', 'mystical', 'grunge', 'ethereal' ], 'key': ['C major', 'D major', 'E minor', 'F minor', 'G major', 'A minor', 'B-flat major', 'G minor', 'D minor', 'F major'], 'bpm': [80, 90, 100, 110, 120, 124, 128, 130, 140, 150, 160, 170, 180], 'time_signature': ['4/4', '3/4', '6/8'], 'guitar_style': [ 'raw distorted', 'melodic', 'fuzzy', 'crisp', 'jangly', 'clean', 'twangy', 'shimmering', 'grunge', 'bluesy', 'slide', 'wah-infused', 'chunky' ], 'bass_style': [ 'punchy', 'deep', 'groovy', 'melodic', 'throbbing', 'slappy', 'funky', 'walking', 'booming', 'resonant', 'subtle' ], 'drum_style': [ 'dynamic', 'minimal', 'hard-hitting', 'swinging', 'polyrhythmic', 'brushed', 'tight', 'loose', 'electronic', 'acoustic', 'retro', 'punchy' ], 'drum_feature': [ 'heavy snare', 'crisp cymbals', 'tight kicks', 'syncopated hits', 'rolling toms', 'ghost notes', 'blast beats' ], 'organ_style': [ 'subtle Hammond', 'swirling', 'warm Leslie', 'church', 'gritty', 'vintage', 'moody' ], 'synth_style': [ 'atmospheric', 'bright', 'eerie', 'soaring', 'chopped', 'arpeggiated', 'pulsing', 'glitchy', 'analog', 'digital', 'layered' ], 'vocal_style': [ 'chopped', 'soulful', 'haunting', 'melodic', 'harmonized', 'layered', 'ethereal', 'gruff', 'breathy' ], 'hihat_style': [ 'crisp', 'swinging', 'rapid', 'shuffling', 'open', 'tight', 'stuttered' ], 'pad_style': [ 'evolving', 'ambient', 'lush', 'dark', 'shimmering', 'warm', 'icy' ], 'kick_style': [ 'deep', 'four-on-the-floor', 'subtle', 'punchy', 'booming', 'clicky' ], 'lead_style': [ 'fluid', 'intricate', 'soaring', 'expressive', 'virtuosic', 'minimalist', 'bluesy', 'lyrical' ], 'lead_instrument': [ 'saxophone', 'trumpet', 'guitar', 'flute', 'violin', 'clarinet', 'trombone' ], 'piano_style': [ 'expressive Rhodes', 'rapid', 'smooth', 'dramatic', 'stride', 'ambient', 'classical', 'jazzy', 'sparse' ], 'keyboard_style': [ 'ornate', 'delicate', 'virtuosic', 'minimal', 'retro', 'spacey' ], 'string_style': [ 'sweeping', 'delicate', 'dramatic', 'lush', 'pizzicato', 'staccato', 'sustained' ], 'brass_style': [ 'bold', 'heroic', 'muted', 'fanfare', 'jazzy', 'smooth' ], 'woodwind_style': [ 'subtle', 'fluttering', 'melodic', 'airy', 'reedy', 'expressive' ], 'flute_style': [ 'fluttering', 'ornate', 'airy', 'breathy', 'trilling' ], 'horn_style': [ 'heroic', 'bold', 'soaring', 'mellow', 'stinging' ], 'choir_style': [ 'mystical', 'ethereal', 'dramatic', 'angelic', 'epic', 'somber' ], 'sample_style': [ 'jazzy', 'soulful', 'gritty', 'cinematic', 'vinyl', 'lo-fi', 'retro' ], 'scratch_style': [ 'crackling vinyl', 'sharp', 'rhythmic', 'chopped', 'transform' ], 'snare_style': [ 'crisp', 'booming', 'tight', 'snappy', 'rimshot', 'layered' ], 'breakdown_style': [ 'euphoric', 'stripped-down', 'intense', 'ambient', 'glitchy', 'dramatic' ], 'intro_bars': [4, 8, 16], 'verse_bars': [8, 16, 32], 'chorus_bars': [8, 16], 'bridge_bars': [4, 8, 16], 'outro_bars': [8, 16], 'build_bars': [8, 16, 32], 'drop_bars': [16, 32], 'main_bars': [16, 32], 'breakdown_bars': [8, 16], 'head_bars': [16, 32], 'solo_bars': [8, 16, 32], 'fugue_bars': [16, 32], 'coda_bars': [8, 16], 'theme_bars': [16, 32], 'development_bars': [16, 32], 'climax_bars': [8, 16], 'groove_bars': [16, 32], 'vibe': [ 'raw', 'energetic', 'melancholic', 'hypnotic', 'soulful', 'intimate', 'virtuosic', 'elegant', 'cinematic', 'gritty', 'nostalgic', 'dark', 'uplifting', 'bittersweet', 'heroic', 'dreamy', 'aggressive', 'relaxed', 'futuristic', 'retro', 'mystical', 'triumphant' ], 'production_style': [ 'lo-fi', 'warm analog', 'clean digital', 'lush', 'crisp acoustic', 'polished pop', 'grand orchestral', 'grunge', 'minimalist', 'industrial', 'vintage' ] } # ============================== # Default INI Creation # ============================== def create_default_genre_prompts_ini(ini_path): default_config = configparser.ConfigParser() default_config['Prompts'] = { 'nirvana': '{style} grunge with {guitar_style} guitar, {bass_style} bass, {drum_style} drums, {vibe} vibe in {key} at {bpm} BPM', 'classic_rock': '{style} classic rock with {guitar_style} guitar, {bass_style} bass, {drum_style} drums, {vibe} vibe in {key} at {bpm} BPM', 'detroit_techno': '{style} techno with {synth_style} synths, {kick_style} kick, {hihat_style} hi-hats, {vibe} vibe at {bpm} BPM', 'smooth_jazz': '{style} jazz with {piano_style} piano, {bass_style} bass, {drum_style} drums, {vibe} vibe in {key} at {bpm} BPM', 'alternative_rock': '{style} alternative rock with {guitar_style} guitar, {bass_style} bass, {drum_style} drums in {key} at {bpm} BPM', 'deep_house': '{style} deep house with {synth_style} synths, {kick_style} kick, {vibe} vibe at {bpm} BPM', 'bebop_jazz': '{style} bebop jazz with {piano_style} piano, {bass_style} bass, {drum_style} drums in {key} at {bpm} BPM', 'baroque_classical': '{style} baroque classical with {string_style} strings, {keyboard_style} harpsichord in {key} at {bpm} BPM', 'romantic_classical': '{style} romantic classical with {string_style} strings, {piano_style} piano in {key} at {bpm} BPM', 'boom_bap_hiphop': '{style} boom bap hip-hop with {sample_style} samples, {drum_style} drums, {scratch_style} scratches at {bpm} BPM', 'trap_hiphop': '{style} trap hip-hop with {synth_style} synths, {kick_style} kick, {snare_style} snare at {bpm} BPM', 'pop_rock': '{style} pop rock with {guitar_style} guitar, {bass_style} bass, {drum_style} drums in {key} at {bpm} BPM', 'fusion_jazz': '{style} fusion jazz with {piano_style} piano, {guitar_style} guitar, {drum_style} drums in {key} at {bpm} BPM', 'edm': '{style} EDM with {synth_style} synths, {kick_style} kick, {vibe} vibe at {bpm} BPM', 'indie_folk': '{style} indie folk with {guitar_style} guitar, {vocal_style} vocals, {drum_style} drums in {key} at {bpm} BPM', 'star_wars': '{style} epic orchestral with {brass_style} brass, {string_style} strings, {vibe} vibe in {key} at {bpm} BPM', 'star_wars_classical': '{style} classical orchestral with {string_style} strings, {horn_style} horns in {key} at {bpm} BPM', 'wutang': '{style} hip-hop with {sample_style} samples, {drum_style} drums, {scratch_style} scratches at {bpm} BPM', 'milesdavis': '{style} jazz with {lead_instrument} lead, {piano_style} piano, {bass_style} bass in {key} at {bpm} BPM' } default_config['BandNames'] = { 'nirvana': 'Nirvana, Soundgarden', 'classic_rock': 'Led Zeppelin, The Rolling Stones', 'detroit_techno': 'Underground Resistance, Jeff Mills', 'smooth_jazz': 'Pat Metheny, George Benson', 'alternative_rock': 'Radiohead, Smashing Pumpkins', 'deep_house': 'Moodymann, Theo Parrish', 'bebop_jazz': 'Charlie Parker, Dizzy Gillespie', 'baroque_classical': 'Bach, Vivaldi', 'romantic_classical': 'Chopin, Liszt', 'boom_bap_hiphop': 'A Tribe Called Quest, Pete Rock', 'trap_hiphop': 'Future, Metro Boomin', 'pop_rock': 'Coldplay, The Killers', 'fusion_jazz': 'Weather Report, Herbie Hancock', 'edm': 'Deadmau5, Skrillex', 'indie_folk': 'Fleet Foxes, Bon Iver', 'star_wars': 'John Williams', 'star_wars_classical': 'John Williams', 'wutang': 'Wu-Tang Clan', 'milesdavis': 'Miles Davis' } with open(ini_path, 'w') as f: default_config.write(f) print(f"Created default {ini_path}") # ============================== # CSS Load # ============================== css_path = "style.css" try: if not os.path.exists(css_path): print(f"ERROR: {css_path} not found. Please create style.css with the required CSS content.") sys.exit(1) with open(css_path, 'r') as f: css = f.read() except Exception as e: print(f"ERROR: Failed to read {css_path}: {e}. Please ensure style.css exists and is readable.") sys.exit(1) # ============================== # INI Load # ============================== config = configparser.ConfigParser() ini_path = "genre_prompts.ini" try: if not os.path.exists(ini_path): print(f"WARNING: {ini_path} not found. Creating default INI file.") create_default_genre_prompts_ini(ini_path) config.read(ini_path) if 'Prompts' not in config.sections() or 'BandNames' not in config.sections(): print(f"WARNING: Invalid {ini_path}. Creating default INI file.") create_default_genre_prompts_ini(ini_path) config.read(ini_path) except Exception as e: print(f"ERROR: Failed to read {ini_path}: {e}. Creating default INI file.") create_default_genre_prompts_ini(ini_path) config.read(ini_path) # ============================== # Model Load with Fallback # ============================== def load_musicgen_with_fallback(): model_paths = [ os.getenv("MUSICGEN_MODEL_PATH_LARGE", "/home/ubuntu/musicpack/models/musicgen-large"), os.getenv("MUSICGEN_MODEL_PATH_MEDIUM", "/home/ubuntu/musicpack/models/musicgen-medium"), os.getenv("MUSICGEN_MODEL_PATH_SMALL", "/home/ubuntu/musicpack/models/musicgen-small"), ] model_names = ["large", "medium", "small"] last_error = None for path, name in zip(model_paths, model_names): if not path: continue if not os.path.exists(path): print(f"NOTE: Model path not found: {path} (skipping {name})") continue try: print(f"Loading MusicGen {name} model from {path} ...") torch.cuda.empty_cache() gc.collect() with autocast('cuda', dtype=AUTOCAST_DTYPE): mdl = MusicGen.get_pretrained(path, device=device) print(f"Loaded MusicGen {name}. Sample rate: {mdl.sample_rate}Hz") return mdl, name except RuntimeError as e: last_error = e print(f"WARNING: Failed to load {name} model due to: {e}") torch.cuda.empty_cache() gc.collect() continue except Exception as e: last_error = e print(f"WARNING: Failed to load {name} model due to: {e}") torch.cuda.empty_cache() gc.collect() continue if last_error: print(f"ERROR: All model loads failed. Last error: {last_error}") raise SystemExit(1) try: musicgen_model, loaded_model_name = load_musicgen_with_fallback() # Conservative defaults; can be overridden per-call musicgen_model.set_generation_params( duration=10, use_sampling=True, top_k=50, top_p=0.0, temperature=0.8, cfg_coef=3.0, two_step_cfg=False ) sample_rate = musicgen_model.sample_rate print(f"Model active: {loaded_model_name}. Sample rate: {sample_rate}Hz") except SystemExit: sys.exit(1) # ============================== # Audio Processing Helpers # ============================== def apply_eq(segment): segment = segment.high_pass_filter(60) segment = segment.low_pass_filter(12000) segment = segment - 2.0 return segment def apply_limiter(segment, max_db=-6.0, target_lufs=-16.0): samples = np.array(segment.get_array_of_samples(), dtype=np.float32) / (2**15) if segment.channels == 2: samples = samples.reshape(-1, 2) meter = pyln.Meter(segment.frame_rate) loudness = meter.integrated_loudness(samples) normalized_samples = pyln.normalize.loudness(samples, loudness, target_lufs) if np.max(np.abs(normalized_samples)) > (10 ** (max_db / 20)): normalized_samples *= (10 ** (max_db / 20)) / np.max(np.abs(normalized_samples)) normalized_samples = (normalized_samples * (2**15)).astype(np.int16) segment = AudioSegment( normalized_samples.tobytes(), frame_rate=segment.frame_rate, sample_width=2, channels=segment.channels ) del samples, normalized_samples gc.collect() return segment def apply_fade(segment, fade_in_duration=1000, fade_out_duration=1000): segment = segment.fade_in(fade_in_duration) segment = segment.fade_out(fade_out_duration) return segment # ============================== # Names & Metadata # ============================== made_up_names = [ 'blazepulse', 'shadowrift', 'neonquest', 'thunderclash', 'stargroove', 'mysticvibe', 'ironspark', 'ghostsurge', 'velvetstorm', 'crimsonrush', 'duskblitz', 'solarflame', 'nightdrift', 'frostsaga', 'emberwave', 'coolriff', 'wildpulse', 'echoslash', 'moontide', 'skydive' ] def extract_song_keyword(prompt): if not prompt: return random.choice(made_up_names) words = re.findall(r'\b\w+\b', prompt.lower()) for word in words: if len(word) <= 15 and word.isalnum(): return word return random.choice(made_up_names) def generate_unique_title(existing_titles, genre, song_keyword, style): letters = string.ascii_uppercase numbers = string.digits max_attempts = 100 attempt = 0 while attempt < max_attempts: title_base = f"{random.choice(letters)}{random.choice(numbers)}" band_names = config['BandNames'].get(genre, "nirvana").split(',') band_name = random.choice([name.strip() for name in band_names]) existing_count = sum(1 for t in existing_titles if t.startswith(title_base) and song_keyword in t and style in t and band_name in t) if existing_count == 0: return title_base, band_name suffix = f"{random.choice(letters)}{random.choice(numbers)}".lower() title_base = f"{title_base}_{suffix}" attempt += 1 raise ValueError("Failed to generate unique title after maximum attempts") def update_metadata_storage(metadata): try: songs_metadata = [] if os.path.exists(metadata_file): with open(metadata_file, 'r') as f: songs_metadata = json.load(f) songs_metadata.append({ "title": metadata["title"], "filename": metadata["filename"], "prompt": metadata.get("prompt", ""), "duration": metadata.get("duration", 30), "volume_db": metadata.get("volume_db", -24.0), "target_lufs": metadata.get("target_lufs", -16.0), "timestamp": metadata.get("timestamp", datetime.datetime.now().strftime("%Y%m%d_%H%M%S")), "file_path": metadata.get("file_path", ""), "sample_rate": metadata.get("sample_rate", musicgen_model.sample_rate), "style": metadata.get("style", ""), "band_name": metadata.get("band_name", ""), "chunk_index": metadata.get("chunk_index", 0) }) with open(metadata_file, 'w') as f: json.dump(songs_metadata, f, indent=4) except Exception as e: print(f"ERROR: Failed to update metadata storage: {e}") def load_renders(): if not os.path.exists(metadata_file): return [], "No renders found." try: with open(metadata_file, 'r') as f: songs_metadata = json.load(f) renders = [ { "Title": entry["title"], "Filename": entry["filename"], "Prompt": entry["prompt"], "Duration (s)": entry["duration"], "Timestamp": entry["timestamp"], "Audio": entry["file_path"], "Download": f'', "Chunk": entry["chunk_index"] } for entry in songs_metadata ] return renders, "Renders loaded successfully." except Exception as e: return [], f"Error loading renders: {e}" # ============================== # Prompt Builder # ============================== def get_genre_prompt(genre): base_prompt = config['Prompts'].get(genre, "") if not base_prompt: base_prompt = "{style} grunge with {guitar_style} guitar, {bass_style} bass, {drum_style} drums, {vibe} vibe in {key} at {bpm} BPM" prompt_dict = { 'style': random.choice(prompt_variables['style']), 'key': random.choice(prompt_variables['key']), 'bpm': random.choice(prompt_variables['bpm']), 'time_signature': random.choice(prompt_variables['time_signature']), 'guitar_style': random.choice(prompt_variables['guitar_style']), 'bass_style': random.choice(prompt_variables['bass_style']), 'drum_style': random.choice(prompt_variables['drum_style']), 'drum_feature': random.choice(prompt_variables['drum_feature']), 'organ_style': random.choice(prompt_variables['organ_style']), 'synth_style': random.choice(prompt_variables['synth_style']), 'vocal_style': random.choice(prompt_variables['vocal_style']), 'hihat_style': random.choice(prompt_variables['hihat_style']), 'pad_style': random.choice(prompt_variables['pad_style']), 'kick_style': random.choice(prompt_variables['kick_style']), 'lead_style': random.choice(prompt_variables['lead_style']), 'lead_instrument': random.choice(prompt_variables['lead_instrument']), 'piano_style': random.choice(prompt_variables['piano_style']), 'keyboard_style': random.choice(prompt_variables['keyboard_style']), 'string_style': random.choice(prompt_variables['string_style']), 'brass_style': random.choice(prompt_variables['brass_style']), 'woodwind_style': random.choice(prompt_variables['woodwind_style']), 'flute_style': random.choice(prompt_variables['flute_style']), 'horn_style': random.choice(prompt_variables['horn_style']), 'choir_style': random.choice(prompt_variables['choir_style']), 'sample_style': random.choice(prompt_variables['sample_style']), 'scratch_style': random.choice(prompt_variables['scratch_style']), 'snare_style': random.choice(prompt_variables['snare_style']), 'breakdown_style': random.choice(prompt_variables['breakdown_style']), 'intro_bars': random.choice(prompt_variables['intro_bars']), 'verse_bars': random.choice(prompt_variables['verse_bars']), 'chorus_bars': random.choice(prompt_variables['chorus_bars']), 'bridge_bars': random.choice(prompt_variables['bridge_bars']), 'outro_bars': random.choice(prompt_variables['outro_bars']), 'build_bars': random.choice(prompt_variables['build_bars']), 'drop_bars': random.choice(prompt_variables['drop_bars']), 'main_bars': random.choice(prompt_variables['main_bars']), 'breakdown_bars': random.choice(prompt_variables['breakdown_bars']), 'head_bars': random.choice(prompt_variables['head_bars']), 'solo_bars': random.choice(prompt_variables['solo_bars']), 'fugue_bars': random.choice(prompt_variables['fugue_bars']), 'coda_bars': random.choice(prompt_variables['coda_bars']), 'theme_bars': random.choice(prompt_variables['theme_bars']), 'development_bars': random.choice(prompt_variables['development_bars']), 'climax_bars': random.choice(prompt_variables['climax_bars']), 'groove_bars': random.choice(prompt_variables['groove_bars']), 'vibe': random.choice(prompt_variables['vibe']), 'production_style': random.choice(prompt_variables['production_style']) } try: formatted_prompt = base_prompt.format(**prompt_dict) words = re.findall(r'\b\w+\b', formatted_prompt.lower()) val_list = [] for k, v in prompt_variables.items(): if isinstance(v, list): val_list.extend(v) if not any(word in val_list for word in words): formatted_prompt = f"{prompt_dict['style']} music with {prompt_dict['guitar_style']} guitar, {prompt_dict['bass_style']} bass, {prompt_dict['drum_style']} drums in {prompt_dict['key']} at {prompt_dict['bpm']} BPM" except KeyError: formatted_prompt = f"{prompt_dict['style']} music with {prompt_dict['guitar_style']} guitar, {prompt_dict['bass_style']} bass, {prompt_dict['drum_style']} drums in {prompt_dict['key']} at {prompt_dict['bpm']} BPM" return formatted_prompt, prompt_dict['style'] # ============================== # Adaptive Chunk Generation (OOM-safe) # ============================== def generate_chunk_oom_safe(model, text_prompt, continuation_prompt, cfg_scale, top_k, top_p, temperature, target_duration): durations_to_try = [target_duration, 20, 15, 12, 10, 8, 6, 4, 3, 2] for dur in durations_to_try: try: torch.cuda.synchronize() torch.cuda.empty_cache() model.set_generation_params( duration=dur, use_sampling=True, top_k=int(top_k), top_p=float(top_p), temperature=float(temperature), cfg_coef=float(cfg_scale), two_step_cfg=False ) with torch.no_grad(): with autocast('cuda', dtype=AUTOCAST_DTYPE): if continuation_prompt is None: # progress=False lowers overhead audio_chunk = model.generate([text_prompt], progress=False)[0] else: audio_chunk = model.generate_continuation( continuation_prompt, model.sample_rate, [text_prompt], progress=False )[0] return audio_chunk, dur except RuntimeError as e: msg = str(e).lower() if "out of memory" in msg or "cuda error" in msg: print(f"OOM at duration {dur}s — retrying with smaller chunk...") torch.cuda.empty_cache() gc.collect() continue else: raise raise RuntimeError("Failed to generate audio chunk without CUDA OOM.") # ============================== # Generation # ============================== def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p: float, temperature: float, total_duration: int, volume_db: float, genre: str = None): global musicgen_model global api_status api_status = "rendering" if not instrumental_prompt.strip() and not genre: instrumental_prompt, style = get_genre_prompt("nirvana") elif not instrumental_prompt.strip(): instrumental_prompt, style = get_genre_prompt(genre) else: words = re.findall(r'\b\w+\b', instrumental_prompt.lower()) val_list = [] for k, v in prompt_variables.items(): if isinstance(v, list): val_list.extend(v) if not any(word in val_list for word in words): instrumental_prompt, style = get_genre_prompt("nirvana") else: ek = extract_song_keyword(instrumental_prompt) style = ek if ek in prompt_variables['style'] else random.choice(prompt_variables['style']) try: start_time = time.time() base_chunk_target = 30 # target; adaptive OOM-safe will shrink if needed total_duration = max(total_duration, 30) remaining = total_duration audio_chunks = [] chunk_paths = [] continuation_prompt = None chunk_index = 0 # Titles existing_titles = [] if os.path.exists(metadata_file): with open(metadata_file, 'r') as f: songs_metadata = json.load(f) existing_titles = [entry["title"] for entry in songs_metadata] song_keyword = extract_song_keyword(instrumental_prompt) title_base, band_name = generate_unique_title(existing_titles, genre if genre else "nirvana", song_keyword, style) # Loop until we render total_duration seconds with adaptive chunks while remaining > 0: target = min(base_chunk_target, remaining) print_resource_usage(f"Before Chunk {chunk_index + 1}") try: audio_chunk, actual_dur = generate_chunk_oom_safe( musicgen_model, instrumental_prompt, continuation_prompt, cfg_scale, top_k, top_p, temperature, target ) audio_chunk = audio_chunk.cpu().to(dtype=torch.float32) if audio_chunk.dim() == 1: audio_chunk = torch.stack([audio_chunk, audio_chunk], dim=0) elif audio_chunk.dim() == 2 and audio_chunk.shape[0] == 1: audio_chunk = torch.cat([audio_chunk, audio_chunk], dim=0) elif audio_chunk.dim() == 2 and audio_chunk.shape[0] != 2: audio_chunk = audio_chunk[:1, :] audio_chunk = torch.cat([audio_chunk, audio_chunk], dim=0) elif audio_chunk.dim() > 2: audio_chunk = audio_chunk.view(2, -1) if audio_chunk.shape[0] != 2: raise ValueError(f"Expected stereo audio with shape (2, samples), got {audio_chunk.shape}") # Update continuation prompt (use up to last 2 seconds if available) samples_per_second = musicgen_model.sample_rate tail_sec = 2 tail_samples = min(int(tail_sec * samples_per_second), audio_chunk.shape[1] - 1 if audio_chunk.shape[1] > 1 else 1) if tail_samples > 0: continuation_prompt = audio_chunk[:, -tail_samples:].cpu() else: continuation_prompt = None # Save to temp wav and convert temp_wav_path = os.path.join(output_dir, f"temp_{random.randint(100, 999)}_{chunk_index}.wav") try: torchaudio.save(temp_wav_path, audio_chunk, musicgen_model.sample_rate, bits_per_sample=16) final_segment = AudioSegment.from_wav(temp_wav_path) finally: if os.path.exists(temp_wav_path): os.remove(temp_wav_path) del audio_chunk gc.collect() # Post FX print(f"Post-processing chunk {chunk_index + 1} (duration ~{actual_dur}s)...") final_segment = apply_eq(final_segment) final_segment = apply_limiter(final_segment, max_db=volume_db, target_lufs=-16.0) if chunk_index == 0: final_segment = final_segment.fade_in(1000) # if last chunk, fade out will be added after loop when combining; also safe to fade here if remaining-actual_dur==0 if remaining - actual_dur <= 0: final_segment = final_segment.fade_out(1000) # Export mp3_filename = f"{title_base.lower()}_{song_keyword}_{style}_{band_name}_chunk{chunk_index + 1}.mp3" mp3_path = os.path.join(output_dir, mp3_filename) final_segment.export( mp3_path, format="mp3", bitrate="64k", tags={"title": f"{title_base}_Chunk{chunk_index + 1}", "artist": "GhostAI"} ) print(f"Saved chunk {chunk_index + 1} to {mp3_path}") audio_chunks.append(final_segment) chunk_paths.append(mp3_path) # Metadata metadata = { "title": f"{title_base}_Chunk{chunk_index + 1}", "filename": mp3_filename, "prompt": instrumental_prompt, "duration": actual_dur, "volume_db": volume_db, "target_lufs": -16.0, "timestamp": datetime.datetime.now().strftime("%Y%m%d_%H%M%S"), "file_path": mp3_path, "sample_rate": musicgen_model.sample_rate, "style": style, "band_name": band_name, "chunk_index": chunk_index + 1 } update_metadata_storage(metadata) chunk_index += 1 remaining -= actual_dur torch.cuda.empty_cache() gc.collect() print_resource_usage(f"After Chunk {chunk_index}") except Exception as e: print(f"ERROR: Failed to process chunk {chunk_index + 1}: {e}") api_status = "idle" raise # Combine chunks if more than one if len(audio_chunks) > 1: combined_segment = audio_chunks[0] for segment in audio_chunks[1:]: combined_segment = combined_segment.append(segment, crossfade=500) combined_mp3_filename = f"{title_base.lower()}_{song_keyword}_{style}_{band_name}_combined.mp3" combined_mp3_path = os.path.join(output_dir, combined_mp3_filename) combined_segment.export( combined_mp3_path, format="mp3", bitrate="64k", tags={"title": title_base, "artist": "GhostAI"} ) print(f"Saved combined audio to {combined_mp3_path}") metadata = { "title": title_base, "filename": combined_mp3_filename, "prompt": instrumental_prompt, "duration": total_duration, "volume_db": volume_db, "target_lufs": -16.0, "timestamp": datetime.datetime.now().strftime("%Y%m%d_%H%M%S"), "file_path": combined_mp3_path, "sample_rate": musicgen_model.sample_rate, "style": style, "band_name": band_name, "chunk_index": 0 } update_metadata_storage(metadata) del combined_segment, audio_chunks gc.collect() api_status = "idle" return combined_mp3_path, "✅ Done!", False, gr.update(value=load_renders()[0]) else: # Single chunk only print(f"Saved metadata to {metadata_file}") del audio_chunks gc.collect() api_status = "idle" return chunk_paths[0], "✅ Done!", False, gr.update(value=load_renders()[0]) except Exception as e: print(f"❌ Failed: {e}") api_status = "idle" return None, f"❌ Failed: {e}", False, gr.update(value=load_renders()[0]) finally: torch.cuda.synchronize() torch.cuda.empty_cache() gc.collect() def clear_inputs(): return "", 3.0, 50, 0.0, 0.8, 30, -24.0, False def show_render_wheel(): return True def set_genre_prompt(genre: str): prompt, _ = get_genre_prompt(genre) return prompt # ============================== # Gradio UI # ============================== with gr.Blocks(css=css) as demo: gr.Markdown(""" """) with gr.Tabs(): with gr.Tab("Generate", id="generate"): with gr.Column(elem_classes="input-container"): gr.Markdown("### Instrumental Prompt") instrumental_prompt = gr.Textbox( label="Instrumental Prompt", placeholder="Select a genre or enter a custom prompt (e.g., 'coolriff grunge')", lines=4, elem_classes="textbox" ) with gr.Row(elem_classes="genre-buttons"): classic_rock_btn = gr.Button("Classic Rock", elem_classes="genre-btn") alternative_rock_btn = gr.Button("Alternative Rock", elem_classes="genre-btn") detroit_techno_btn = gr.Button("Detroit Techno", elem_classes="genre-btn") deep_house_btn = gr.Button("Deep House", elem_classes="genre-btn") smooth_jazz_btn = gr.Button("Smooth Jazz", elem_classes="genre-btn") bebop_jazz_btn = gr.Button("Bebop Jazz", elem_classes="genre-btn") baroque_classical_btn = gr.Button("Baroque Classical", elem_classes="genre-btn") romantic_classical_btn = gr.Button("Romantic Classical", elem_classes="genre-btn") boom_bap_hiphop_btn = gr.Button("Boom Bap Hip-Hop", elem_classes="genre-btn") trap_hiphop_btn = gr.Button("Trap Hip-Hop", elem_classes="genre-btn") pop_rock_btn = gr.Button("Pop Rock", elem_classes="genre-btn") fusion_jazz_btn = gr.Button("Fusion Jazz", elem_classes="genre-btn") edm_btn = gr.Button("EDM", elem_classes="genre-btn") indie_folk_btn = gr.Button("Indie Folk", elem_classes="genre-btn") star_wars_btn = gr.Button("Star Wars Epic", elem_classes="genre-btn") star_wars_classical_btn = gr.Button("Star Wars Classical", elem_classes="genre-btn") nirvana_btn = gr.Button("Nirvana", elem_classes="genre-btn") wutang_btn = gr.Button("Wu-Tang", elem_classes="genre-btn") milesdavis_btn = gr.Button("Miles Davis", elem_classes="genre-btn") with gr.Column(elem_classes="settings-container"): gr.Markdown("### Generation Settings") cfg_scale = gr.Slider( label="Guidance Scale (CFG)", minimum=1.0, maximum=10.0, value=3.0, step=0.1 ) top_k = gr.Slider( label="Top-K Sampling", minimum=10, maximum=500, value=50, step=10 ) top_p = gr.Slider( label="Top-P Sampling", minimum=0.0, maximum=1.0, value=0.0, step=0.1 ) temperature = gr.Slider( label="Temperature", minimum=0.1, maximum=2.0, value=0.8, step=0.1 ) total_duration = gr.Slider( label="Duration (seconds)", minimum=30, maximum=300, value=30, step=10 ) volume_db = gr.Slider( label="Output Volume (dBFS)", minimum=-30.0, maximum=0.0, value=-24.0, step=0.1 ) with gr.Row(elem_classes="action-buttons"): gen_btn = gr.Button("Generate Music") clr_btn = gr.Button("Clear Inputs") with gr.Column(elem_classes="output-container"): gr.Markdown("### Output") render_wheel = gr.HTML('
Generating...
', label="Rendering Status") render_state = gr.State(value=False) out_audio = gr.Audio(label="Generated Track", type="filepath", interactive=True, elem_classes="audio-container") status = gr.Textbox(label="Status", interactive=False) with gr.Tab("Renders", id="renders"): with gr.Column(elem_classes="renders-container"): gr.Markdown("### Browse Renders") renders_table = gr.DataFrame( headers=["Title", "Filename", "Prompt", "Duration (s)", "Timestamp", "Audio", "Download", "Chunk"], datatype=["str", "str", "str", "number", "str", "audio", "html", "number"], interactive=False, value=load_renders()[0], elem_classes="renders-table" ) renders_status = gr.Textbox(label="Renders Status", interactive=False, value=load_renders()[1]) # Button bindings classic_rock_btn.click(set_genre_prompt, inputs=[gr.State(value="classic_rock")], outputs=[instrumental_prompt]) alternative_rock_btn.click(set_genre_prompt, inputs=[gr.State(value="alternative_rock")], outputs=[instrumental_prompt]) detroit_techno_btn.click(set_genre_prompt, inputs=[gr.State(value="detroit_techno")], outputs=[instrumental_prompt]) deep_house_btn.click(set_genre_prompt, inputs=[gr.State(value="deep_house")], outputs=[instrumental_prompt]) smooth_jazz_btn.click(set_genre_prompt, inputs=[gr.State(value="smooth_jazz")], outputs=[instrumental_prompt]) bebop_jazz_btn.click(set_genre_prompt, inputs=[gr.State(value="bebop_jazz")], outputs=[instrumental_prompt]) baroque_classical_btn.click(set_genre_prompt, inputs=[gr.State(value="baroque_classical")], outputs=[instrumental_prompt]) romantic_classical_btn.click(set_genre_prompt, inputs=[gr.State(value="romantic_classical")], outputs=[instrumental_prompt]) boom_bap_hiphop_btn.click(set_genre_prompt, inputs=[gr.State(value="boom_bap_hiphop")], outputs=[instrumental_prompt]) trap_hiphop_btn.click(set_genre_prompt, inputs=[gr.State(value="trap_hiphop")], outputs=[instrumental_prompt]) pop_rock_btn.click(set_genre_prompt, inputs=[gr.State(value="pop_rock")], outputs=[instrumental_prompt]) fusion_jazz_btn.click(set_genre_prompt, inputs=[gr.State(value="fusion_jazz")], outputs=[instrumental_prompt]) edm_btn.click(set_genre_prompt, inputs=[gr.State(value="edm")], outputs=[instrumental_prompt]) indie_folk_btn.click(set_genre_prompt, inputs=[gr.State(value="indie_folk")], outputs=[instrumental_prompt]) star_wars_btn.click(set_genre_prompt, inputs=[gr.State(value="star_wars")], outputs=[instrumental_prompt]) star_wars_classical_btn.click(set_genre_prompt, inputs=[gr.State(value="star_wars_classical")], outputs=[instrumental_prompt]) nirvana_btn.click(set_genre_prompt, inputs=[gr.State(value="nirvana")], outputs=[instrumental_prompt]) wutang_btn.click(set_genre_prompt, inputs=[gr.State(value="wutang")], outputs=[instrumental_prompt]) milesdavis_btn.click(set_genre_prompt, inputs=[gr.State(value="milesdavis")], outputs=[instrumental_prompt]) gen_btn.click( fn=show_render_wheel, inputs=None, outputs=[render_state], ).then( fn=generate_music, inputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, volume_db, gr.State(None)], outputs=[out_audio, status, render_state, renders_table], show_progress="full" ) clr_btn.click( fn=clear_inputs, inputs=None, outputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, volume_db, render_state] ) # ============================== # FastAPI # ============================== app = FastAPI() class MusicRequest(BaseModel): prompt: str = None duration: int = 30 volume_db: float = -24.0 genre: str = None @app.get("/prompts/") async def get_prompts(): global api_status try: prompts = list(config['Prompts'].keys()) return {"status": api_status, "prompts": prompts} except Exception as e: print(f"Error fetching prompts: {e}") raise HTTPException(status_code=500, detail=f"Error fetching prompts: {e}") @app.post("/generate-music/") async def api_generate_music(request: MusicRequest): global api_status api_status = "rendering" try: instrumental_prompt = ( get_genre_prompt(request.genre)[0] if request.genre else request.prompt if request.prompt else get_genre_prompt("nirvana")[0] ) style = ( get_genre_prompt(request.genre)[1] if request.genre else extract_song_keyword(request.prompt) if request.prompt and extract_song_keyword(request.prompt) in prompt_variables['style'] else get_genre_prompt("nirvana")[1] ) if not instrumental_prompt.strip(): api_status = "idle" raise HTTPException(status_code=400, detail="Invalid prompt or genre") total_duration = max(request.duration, 30) remaining = total_duration audio_chunks = [] chunk_paths = [] continuation_prompt = None chunk_index = 0 existing_titles = [] if os.path.exists(metadata_file): with open(metadata_file, 'r') as f: songs_metadata = json.load(f) existing_titles = [entry["title"] for entry in songs_metadata] song_keyword = extract_song_keyword(request.prompt if request.prompt else instrumental_prompt) title_base, band_name = generate_unique_title(existing_titles, request.genre if request.genre else "nirvana", song_keyword, style) while remaining > 0: target = min(30, remaining) print_resource_usage(f"Before API Chunk {chunk_index + 1}") try: audio_chunk, actual_dur = generate_chunk_oom_safe( musicgen_model, instrumental_prompt, continuation_prompt, 3.0, 50, 0.0, 0.8, target ) audio_chunk = audio_chunk.cpu().to(dtype=torch.float32) if audio_chunk.dim() == 1: audio_chunk = torch.stack([audio_chunk, audio_chunk], dim=0) elif audio_chunk.dim() == 2 and audio_chunk.shape[0] == 1: audio_chunk = torch.cat([audio_chunk, audio_chunk], dim=0) elif audio_chunk.dim() == 2 and audio_chunk.shape[0] != 2: audio_chunk = audio_chunk[:1, :] audio_chunk = torch.cat([audio_chunk, audio_chunk], dim=0) elif audio_chunk.dim() > 2: audio_chunk = audio_chunk.view(2, -1) if audio_chunk.shape[0] != 2: raise ValueError(f"Expected stereo audio with shape (2, samples), got {audio_chunk.shape}") samples_per_second = musicgen_model.sample_rate tail_sec = 2 tail_samples = min(int(tail_sec * samples_per_second), audio_chunk.shape[1] - 1 if audio_chunk.shape[1] > 1 else 1) continuation_prompt = audio_chunk[:, -tail_samples:].cpu() if tail_samples > 0 else None temp_wav_path = os.path.join(output_dir, f"temp_{random.randint(100, 999)}_{chunk_index}.wav") try: torchaudio.save(temp_wav_path, audio_chunk, musicgen_model.sample_rate, bits_per_sample=16) final_segment = AudioSegment.from_wav(temp_wav_path) finally: if os.path.exists(temp_wav_path): os.remove(temp_wav_path) del audio_chunk gc.collect() final_segment = apply_eq(final_segment) final_segment = apply_limiter(final_segment, max_db=request.volume_db, target_lufs=-16.0) if chunk_index == 0: final_segment = final_segment.fade_in(1000) if remaining - actual_dur <= 0: final_segment = final_segment.fade_out(1000) mp3_filename = f"{title_base.lower()}_{song_keyword}_{style}_{band_name}_chunk{chunk_index + 1}.mp3" mp3_path = os.path.join(output_dir, mp3_filename) final_segment.export( mp3_path, format="mp3", bitrate="64k", tags={"title": f"{title_base}_Chunk{chunk_index + 1}", "artist": "GhostAI"} ) print(f"Saved API chunk {chunk_index + 1} to {mp3_path}") audio_chunks.append(final_segment) chunk_paths.append(mp3_path) metadata = { "title": f"{title_base}_Chunk{chunk_index + 1}", "filename": mp3_filename, "prompt": instrumental_prompt, "duration": actual_dur, "volume_db": request.volume_db, "target_lufs": -16.0, "timestamp": datetime.datetime.now().strftime("%Y%m%d_%H%M%S"), "file_path": mp3_path, "sample_rate": musicgen_model.sample_rate, "style": style, "band_name": band_name, "chunk_index": chunk_index + 1 } update_metadata_storage(metadata) chunk_index += 1 remaining -= actual_dur torch.cuda.empty_cache() gc.collect() print_resource_usage(f"After API Chunk {chunk_index}") except Exception as e: print(f"ERROR: Failed to process API chunk {chunk_index + 1}: {e}") api_status = "idle" raise if len(audio_chunks) > 1: combined_segment = audio_chunks[0] for segment in audio_chunks[1:]: combined_segment = combined_segment.append(segment, crossfade=500) combined_mp3_filename = f"{title_base.lower()}_{song_keyword}_{style}_{band_name}_combined.mp3" combined_mp3_path = os.path.join(output_dir, combined_mp3_filename) combined_segment.export( combined_mp3_path, format="mp3", bitrate="64k", tags={"title": title_base, "artist": "GhostAI"} ) print(f"Saved combined audio to {combined_mp3_path}") metadata = { "title": title_base, "filename": combined_mp3_filename, "prompt": instrumental_prompt, "duration": total_duration, "volume_db": request.volume_db, "target_lufs": -16.0, "timestamp": datetime.datetime.now().strftime("%Y%m%d_%H%M%S"), "file_path": combined_mp3_path, "sample_rate": musicgen_model.sample_rate, "style": style, "band_name": band_name, "chunk_index": 0 } update_metadata_storage(metadata) del combined_segment, audio_chunks gc.collect() api_status = "idle" return FileResponse(combined_mp3_path, media_type="audio/mpeg") else: print(f"Saved metadata to {metadata_file}") del audio_chunks gc.collect() api_status = "idle" return FileResponse(chunk_paths[0], media_type="audio/mpeg") except Exception as e: print(f"Error generating music: {e}") api_status = "idle" raise HTTPException(status_code=500, detail=f"Error generating music: {e}") finally: torch.cuda.synchronize() torch.cuda.empty_cache() gc.collect() @app.get("/get-song/{filename}") async def get_song(filename: str): global api_status file_path = os.path.join(output_dir, filename) if not os.path.exists(file_path): print(f"Error: Song file {filename} not found") raise HTTPException(status_code=404, detail="Song file not found") print(f"Serving file: {filename}") return FileResponse(file_path, media_type="audio/mpeg", filename=filename) @app.get("/status/") async def get_status(): global api_status return {"status": api_status} def run_fastapi(): uvicorn.run(app, host="0.0.0.0", port=8000) # ============================== # Main # ============================== if __name__ == "__main__": fastapi_process = multiprocessing.Process(target=run_fastapi) fastapi_process.start() try: demo.launch(server_name="0.0.0.0", server_port=9999, share=False, inbrowser=True, show_error=True) except Exception as e: print(f"ERROR: Failed to launch Gradio: {e}") fastapi_process.terminate() sys.exit(1) finally: fastapi_process.terminate()