audio_engine.py

import sys, os, time, torch
from fastapi import FastAPI
from transformers import AutoProcessor, MusicgenForConditionalGeneration
from bark import SAMPLE_RATE, generate_audio, preload_models
import scipy.io.wavfile
import edge_tts
from pydub import AudioSegment
import librosa
from colorama import init, Fore
import uvicorn
import numpy as np

# --- STUDIO EFFECTS (Spotify Pedalboard) ---
from pedalboard import Pedalboard, Compressor, Reverb, Gain, HighpassFilter
from pedalboard.io import AudioFile

# --- CONFIG ---
PORT = 5555
MODEL_NAME = "facebook/musicgen-small"
OUTPUT_FOLDER = "generated_audio"

init(autoreset=True)
app = FastAPI()
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# --- LOAD MODELS ---
print(f"{Fore.CYAN}=== AUDIO ENGINE V8 (STUDIO QUALITY) ===")
device = "cpu"
if torch.cuda.is_available(): device = "cuda"

print(f"{Fore.WHITE}Loading MusicGen... ", end="")
try:
    processor = AutoProcessor.from_pretrained(MODEL_NAME)
    model = MusicgenForConditionalGeneration.from_pretrained(MODEL_NAME)
    model.to(device)
    print(f"{Fore.GREEN}DONE!")
except Exception as e: print(f"{Fore.RED}FAILED! {e}")

print(f"{Fore.WHITE}Loading Bark (High Fidelity)... ", end="")
try:
    # Patch PyTorch 2.6
    _original_load = torch.load
    def unsafe_load(*args, **kwargs):
        if 'weights_only' not in kwargs: kwargs['weights_only'] = False
        return _original_load(*args, **kwargs)
    torch.load = unsafe_load

    # CRITICAL CHANGE: We set use_small=False to get the GOOD voices
    # This uses more RAM but sounds 2x better.
    preload_models(text_use_small=False, coarse_use_small=False, fine_use_small=False)

    torch.load = _original_load
    print(f"{Fore.GREEN}DONE!")
except Exception as e:
    torch.load = _original_load
    print(f"{Fore.RED}FAILED! {e}")

# --- HELPER: STUDIO FX RACK ---
def apply_studio_fx(input_path, output_path):
    """ Adds Compressor + Reverb to raw AI vocals """
    print(f"{Fore.YELLOW}    >>> Applying Studio FX (Comp + Reverb)...")

    # Define the FX Chain
    board = Pedalboard([
        HighpassFilter(cutoff_frequency_hz=80), # Remove mud
        Compressor(threshold_db=-15, ratio=4),  # Make it punchy
        Reverb(room_size=0.4, wet_level=0.3),   # Add "Space"
        Gain(gain_db=2)                         # Boost volume
    ])

    # Process
    with AudioFile(input_path) as f:
        audio = f.read(f.frames)
        samplerate = f.samplerate

    effected = board(audio, samplerate)

    with AudioFile(output_path, 'w', samplerate, effected.shape[0]) as f:
        f.write(effected)

# --- ENDPOINTS ---

@app.post("/suno")
async def generate_suno_song(payload: dict):
    lyrics = payload.get("lyrics", "♪ La la la ♪")
    style = payload.get("style", "Pop")
    # Use "Speaker 0" (History Prompt None) or specific high quality ones
    # v2/en_speaker_6 is decent, but lets try to be robust
    voice_preset = payload.get("voice", "v2/en_speaker_6")

    print(f"{Fore.YELLOW}[SUNO V8] Creating High-Def Song: '{style}'")

    # 1. GENERATE RAW VOCALS
    print(f"{Fore.CYAN}  > Step 1: Singing (High Quality Model)...")
    raw_vocal_path = os.path.join(OUTPUT_FOLDER, "temp_raw_vox.wav")
    processed_vocal_path = os.path.join(OUTPUT_FOLDER, f"suno_vox_{int(time.time())}.wav")

    try:
        if "♪" not in lyrics: lyrics = f"♪ {lyrics} ♪"
        audio_array = generate_audio(lyrics, history_prompt=voice_preset)
        scipy.io.wavfile.write(raw_vocal_path, SAMPLE_RATE, audio_array)

        # 2. APPLY STUDIO FX (The "Good" Button)
        apply_studio_fx(raw_vocal_path, processed_vocal_path)

    except Exception as e:
        return {"status": "ERROR", "message": f"Bark Failed: {e}"}

    # 3. MEASURE DURATION
    print(f"{Fore.MAGENTA}  > Step 2: Analyzing Timing...")
    try:
        y, sr = librosa.load(processed_vocal_path, sr=32000)
        duration_sec = librosa.get_duration(y=y, sr=sr)
        target_duration = duration_sec + 2
    except: target_duration = 10

    # 4. GENERATE MUSIC
    print(f"{Fore.MAGENTA}  > Step 3: Composing Music...")
    music_path = os.path.join(OUTPUT_FOLDER, f"suno_inst_{int(time.time())}.wav")
    try:
        inputs = processor(
            text=[f"{style} song, high fidelity, studio quality"],
            padding=True,
            return_tensors="pt"
        ).to(device)
        max_tokens = int(target_duration * 50)
        audio_values = model.generate(**inputs, max_new_tokens=max_tokens)
        scipy.io.wavfile.write(music_path, rate=model.config.audio_encoder.sampling_rate, data=audio_values[0, 0].cpu().numpy())
    except Exception as e:
        return {"status": "ERROR", "message": f"MusicGen Failed: {e}"}

    # 5. MIX
    print(f"{Fore.GREEN}  > Step 4: Mixing...")
    final_path = os.path.join(OUTPUT_FOLDER, f"SUNO_SONG_{int(time.time())}.wav")
    try:
        vox = AudioSegment.from_file(processed_vocal_path)
        inst = AudioSegment.from_file(music_path)
        inst = inst - 4
        vox = vox + 1
        final = inst.overlay(vox, position=0)
        final.export(final_path, format="wav")
        print(f"{Fore.GREEN}[COMPLETE] Saved: {final_path}")
        return {"status": "SUCCESS", "file": os.path.abspath(final_path)}
    except Exception as e:
        return {"status": "ERROR", "message": f"Mix Failed: {e}"}

# --- OTHER ENDPOINTS ---
@app.post("/speak")
async def speak(payload: dict):
    # EdgeTTS is already high quality, no need to change
    text = payload.get("text")
    path = os.path.join(OUTPUT_FOLDER, f"vox_{int(time.time())}.mp3")
    communicate = edge_tts.Communicate(text, "en-US-ChristopherNeural")
    await communicate.save(path)
    return {"status": "SUCCESS", "file": os.path.abspath(path)}

@app.get("/")
def home(): return {"status": "ONLINE"}

if __name__ == "__main__":
    print(f"\n{Fore.GREEN}=== SUNO ENGINE LISTENING ON http://127.0.0.1:{PORT} ===")
    uvicorn.run(app, host="127.0.0.1", port=PORT, log_level="info")