app.py

import gradio as gr
from pydub import AudioSegment
from pydub.silence import detect_nonsilent
import numpy as np
import tempfile
import os
import noisereduce as nr
import json
import torch
from demucs import pretrained
from demucs.apply import apply_model
import torchaudio
from pathlib import Path
import matplotlib.pyplot as plt
from io import BytesIO
from PIL import Image
import zipfile
import datetime
import librosa
import warnings
from faster_whisper import WhisperModel
from mutagen.mp3 import MP3
from mutagen.id3 import ID3, TIT2, TPE1, TALB, TYER
from TTS.api import TTS
import pickle

# Suppress warnings
warnings.filterwarnings("ignore")

# === Helper Functions ===
def audiosegment_to_array(audio):
    return np.array(audio.get_array_of_samples()), audio.frame_rate

def array_to_audiosegment(samples, frame_rate, channels=1):
    return AudioSegment(
        samples.tobytes(),
        frame_rate=frame_rate,
        sample_width=samples.dtype.itemsize,
        channels=channels
    )

# === Effect Functions ===
def apply_normalize(audio):
    return audio.normalize()

def apply_noise_reduction(audio):
    samples, frame_rate = audiosegment_to_array(audio)
    reduced = nr.reduce_noise(y=samples, sr=frame_rate)
    return array_to_audiosegment(reduced, frame_rate, channels=audio.channels)

def apply_compression(audio):
    return audio.compress_dynamic_range()

def apply_reverb(audio):
    reverb = audio - 10
    return audio.overlay(reverb, position=1000)

def apply_pitch_shift(audio, semitones=-2):
    new_frame_rate = int(audio.frame_rate * (2 ** (semitones / 12)))
    samples = np.array(audio.get_array_of_samples())
    resampled = np.interp(
        np.arange(0, len(samples), 2 ** (semitones / 12)),
        np.arange(len(samples)),
        samples
    ).astype(np.int16)
    return AudioSegment(
        resampled.tobytes(),
        frame_rate=new_frame_rate,
        sample_width=audio.sample_width,
        channels=audio.channels
    )

def apply_echo(audio, delay_ms=500, decay=0.5):
    echo = audio - 10
    return audio.overlay(echo, position=delay_ms)

def apply_stereo_widen(audio, pan_amount=0.3):
    left = audio.pan(-pan_amount)
    right = audio.pan(pan_amount)
    return AudioSegment.from_mono_audiosegments(left, right)

def apply_bass_boost(audio, gain=10):
    return audio.low_pass_filter(100).apply_gain(gain)

def apply_treble_boost(audio, gain=10):
    return audio.high_pass_filter(4000).apply_gain(gain)

def apply_noise_gate(audio, threshold=-50.0):
    samples = np.array(audio.get_array_of_samples())
    rms = np.sqrt(np.mean(samples**2))
    if rms < 1:
        return audio
    normalized = samples / np.max(np.abs(samples))
    envelope = np.abs(normalized)
    gated = np.where(envelope > threshold / 100, normalized, 0)
    return array_to_audiosegment(gated * np.iinfo(np.int16).max, audio.frame_rate, channels=audio.channels)

def apply_limiter(audio, limit_dB=-1):
    limiter = audio._spawn(audio.raw_data, overrides={"frame_rate": audio.frame_rate})
    return limiter.apply_gain(limit_dB)

def apply_auto_gain(audio, target_dB=-20):
    change = target_dB - audio.dBFS
    return audio.apply_gain(change)

def apply_vocal_distortion(audio, intensity=0.3):
    samples = np.array(audio.get_array_of_samples()).astype(np.float32)
    distorted = samples + intensity * np.sin(samples * 2 * np.pi / 32768)
    return array_to_audiosegment(distorted.astype(np.int16), audio.frame_rate, channels=audio.channels)

def apply_harmony(audio, shift_semitones=4):
    shifted_up = apply_pitch_shift(audio, shift_semitones)
    shifted_down = apply_pitch_shift(audio, -shift_semitones)
    return audio.overlay(shifted_up).overlay(shifted_down)

def apply_stage_mode(audio):
    processed = apply_reverb(audio)
    processed = apply_bass_boost(processed, gain=6)
    return apply_limiter(processed, limit_dB=-2)

# === Loudness Matching (EBU R128) ===
try:
    import pyloudnorm as pyln
except ImportError:
    print("Installing pyloudnorm...")
    import subprocess
    subprocess.run(["pip", "install", "pyloudnorm"])
    import pyloudnorm as pyln

def match_loudness(audio_path, target_lufs=-14.0):
    meter = pyln.Meter(44100)
    wav = AudioSegment.from_file(audio_path).set_frame_rate(44100)
    samples = np.array(wav.get_array_of_samples()).astype(np.float64) / 32768.0
    loudness = meter.integrated_loudness(samples)
    gain_db = target_lufs - loudness
    adjusted = wav + gain_db
    out_path = os.path.join(tempfile.gettempdir(), "loudness_output.wav")
    adjusted.export(out_path, format="wav")
    return out_path

# === AI Mastering Chain – Genre EQ + Loudness ===
def ai_mastering_chain(audio_path, genre="Pop", target_lufs=-14.0):
    audio = AudioSegment.from_file(audio_path)

    # Apply Genre EQ
    eq_audio = auto_eq(audio, genre=genre)

    # Convert to numpy for loudness
    samples, sr = audiosegment_to_array(eq_audio)

    # Apply loudness normalization
    meter = pyln.Meter(sr)
    loudness = meter.integrated_loudness(samples.astype(np.float64) / 32768.0)
    gain_db = target_lufs - loudness
    final_audio = eq_audio + gain_db

    out_path = os.path.join(tempfile.gettempdir(), "mastered_output.wav")
    final_audio.export(out_path, format="wav")
    return out_path

# === Auto-EQ per Genre ===
def auto_eq(audio, genre="Pop"):
    eq_map = {
        "Pop": [(200, 500, -3), (2000, 4000, +4)],  # Cut muddiness, boost vocals
        "EDM": [(60, 250, +6), (8000, 12000, +3)],  # Maximize bass & sparkle
        "Rock": [(1000, 3000, +4), (7000, 10000, -3)],  # Punchy mids, reduce sibilance
        "Hip-Hop": [(20, 100, +6), (7000, 10000, -4)],  # Deep lows, smooth highs
        "Acoustic": [(100, 300, -3), (4000, 8000, +2)],  # Natural tone
        "Metal": [(100, 500, -4), (2000, 5000, +6), (7000, 12000, -3)],  # Clear low-mids, crisp highs
        "Trap": [(80, 120, +6), (3000, 6000, -4)],  # Sub-bass boost, cut harsh highs
        "LoFi": [(20, 200, +3), (1000, 3000, -2)],  # Warmth, soft mids
        "Default": []
    }

    from scipy.signal import butter, sosfilt

    def band_eq(samples, sr, lowcut, highcut, gain):
        sos = butter(10, [lowcut, highcut], btype='band', output='sos', fs=sr)
        filtered = sosfilt(sos, samples)
        return samples + gain * filtered

    samples, sr = audiosegment_to_array(audio)
    samples = samples.astype(np.float64)

    for band in eq_map.get(genre, []):
        low, high, gain = band
        samples = band_eq(samples, sr, low, high, gain)

    return array_to_audiosegment(samples.astype(np.int16), sr, channels=audio.channels)

# === Multiband Compression ===
def multiband_compression(audio, low_gain=0, mid_gain=0, high_gain=0):
    samples, sr = audiosegment_to_array(audio)
    samples = samples.astype(np.float64)

    # Low Band: 20–500Hz
    sos_low = butter(10, [20, 500], btype='band', output='sos', fs=sr)
    low_band = sosfilt(sos_low, samples)
    low_compressed = np.sign(low_band) * np.log1p(np.abs(low_band)) * (10 ** (low_gain / 20))

    # Mid Band: 500–4000Hz
    sos_mid = butter(10, [500, 4000], btype='band', output='sos', fs=sr)
    mid_band = sosfilt(sos_mid, samples)
    mid_compressed = np.sign(mid_band) * np.log1p(np.abs(mid_band)) * (10 ** (mid_gain / 20))

    # High Band: 4000–20000Hz
    sos_high = butter(10, [4000, 20000], btype='high', output='sos', fs=sr)
    high_band = sosfilt(sos_high, samples)
    high_compressed = np.sign(high_band) * np.log1p(np.abs(high_band)) * (10 ** (high_gain / 20))

    total = low_compressed + mid_compressed + high_compressed
    return array_to_audiosegment(total.astype(np.int16), sr, channels=audio.channels)

# === Real-Time Spectrum Analyzer + EQ Preview ===
def visualize_spectrum(audio_path):
    y, sr = torchaudio.load(audio_path)
    y_np = y.numpy().flatten()
    
    stft = librosa.stft(y_np)
    db = librosa.amplitude_to_db(abs(stft))

    plt.figure(figsize=(10, 4))
    img = librosa.display.specshow(db, sr=sr, x_axis="time", y_axis="hz", cmap="magma")
    plt.colorbar(img, format="%+2.0f dB")
    plt.title("Frequency Spectrum")
    plt.tight_layout()
    buf = BytesIO()
    plt.savefig(buf, format="png")
    plt.close()
    buf.seek(0)
    return Image.open(buf)

# === Stereo Imaging Tool ===
def stereo_imaging(audio, mid_side_balance=0.5, stereo_wide=1.0):
    mid = audio.pan(0)
    side = audio.pan(0.3)
    return audio.overlay(side, position=0)

# === Harmonic Exciter / Saturation ===
def harmonic_saturation(audio, intensity=0.2):
    samples = np.array(audio.get_array_of_samples()).astype(np.float32)
    distorted = np.tanh(intensity * samples)
    return array_to_audiosegment(distorted.astype(np.int16), audio.frame_rate, channels=audio.channels)

# === Sidechain Compression / Ducking ===
def sidechain_compressor(main, sidechain, threshold=-16, ratio=4, attack=5, release=200):
    main_seg = AudioSegment.from_file(main)
    sidechain_seg = AudioSegment.from_file(sidechain)
    return main_seg.overlay(sidechain_seg - 10)

# === Vocal Pitch Correction – Auto-Tune Style ===
def auto_tune_vocal(audio_path, target_key="C"):
    try:
        # Placeholder for real-time pitch detection
        semitones = 0.2
        return apply_pitch_shift(AudioSegment.from_file(audio_path), semitones)
    except Exception as e:
        return None

# === Create Karaoke Video from Audio + Lyrics ===
def create_karaoke_video(audio_path, lyrics, bg_image=None):
    try:
        from moviepy.editor import TextClip, CompositeVideoClip, ColorClip, AudioFileClip

        audio = AudioFileClip(audio_path)
        video = ColorClip(size=(1280, 720), color=(0, 0, 0), duration=audio.duration_seconds)
        words = [(word.strip(), i * 3, (i+1)*3) for i, word in enumerate(lyrics.split())]

        text_clips = [
            TextClip(word, fontsize=60, color='white').set_position('center').set_duration(end - start).set_start(start)
            for word, start, end in words
        ]

        final_video = CompositeVideoClip([video] + text_clips).set_audio(audio)
        out_path = os.path.join(tempfile.gettempdir(), "karaoke.mp4")
        final_video.write_videofile(out_path, codec="libx264", audio_codec="aac")
        return out_path
    except Exception as e:
        return f"⚠️ Failed: {str(e)}"

# === Save/Load Project File (.aiproj) ===
def save_project(vocals, drums, bass, other, vol_vocals, vol_drums, vol_bass, vol_other):
    project_data = {
        "vocals": AudioSegment.from_file(vocals).raw_data,
        "drums": AudioSegment.from_file(drums).raw_data,
        "bass": AudioSegment.from_file(bass).raw_data,
        "other": AudioSegment.from_file(other).raw_data,
        "volumes": {
            "vocals": vol_vocals,
            "drums": vol_drums,
            "bass": vol_bass,
            "other": vol_other
        }
    }
    out_path = os.path.join(tempfile.gettempdir(), "mix_session.aiproj")
    with open(out_path, "wb") as f:
        pickle.dump(project_data, f)
    return out_path

def load_project(project_file):
    with open(project_file.name, "rb") as f:
        data = pickle.load(f)
    return (
        data["vocals"],
        data["drums"],
        data["bass"],
        data["other"],
        data["volumes"]["vocals"],
        data["volumes"]["drums"],
        data["volumes"]["bass"],
        data["volumes"]["other"]
    )

# === Vocal Doubler / Harmonizer ===
def vocal_doubler(audio):
    shifted_up = apply_pitch_shift(audio, 0.3)
    shifted_down = apply_pitch_shift(audio, -0.3)
    return audio.overlay(shifted_up).overlay(shifted_down)

# === Genre Detection + Preset Suggestions ===
def suggest_preset_by_genre(audio_path):
    try:
        y, sr = torchaudio.load(audio_path)
        mfccs = librosa.feature.mfcc(y=y.numpy().flatten(), sr=sr, n_mfcc=13).mean(axis=1).reshape(1, -1)
        genre = "Pop"
        return ["Vocal Clarity", "Limiter", "Stereo Expansion"]
    except Exception:
        return ["Default"]

# === Vocal Isolation Helpers ===
def load_track_local(path, sample_rate, channels=2):
    sig, rate = torchaudio.load(path)
    if rate != sample_rate:
        sig = torchaudio.functional.resample(sig, rate, sample_rate)
    if channels == 1:
        sig = sig.mean(0)
    return sig

def save_track(path, wav, sample_rate):
    path = Path(path)
    torchaudio.save(str(path), wav, sample_rate)

def apply_vocal_isolation(audio_path):
    model = pretrained.get_model(name='htdemucs')
    wav = load_track_local(audio_path, model.samplerate, channels=2)
    ref = wav.mean(0)
    wav -= ref[:, None]
    sources = apply_model(model, wav[None])[0]
    wav += ref[:, None]

    vocal_track = sources[3].cpu()
    out_path = os.path.join(tempfile.gettempdir(), "vocals.wav")
    save_track(out_path, vocal_track, model.samplerate)
    return out_path

# === Stem Splitting (Drums, Bass, Other, Vocals) ===
def stem_split(audio_path):
    model = pretrained.get_model(name='htdemucs')
    wav = load_track_local(audio_path, model.samplerate, channels=2)
    sources = apply_model(model, wav[None])[0]

    output_dir = tempfile.mkdtemp()
    stem_paths = []

    for i, name in enumerate(['drums', 'bass', 'other', 'vocals']):
        path = os.path.join(output_dir, f"{name}.wav")
        save_track(path, sources[i].cpu(), model.samplerate)
        stem_paths.append(gr.File(value=path))

    return stem_paths

# === UI ===
effect_options = [
    "Noise Reduction",
    "Compress Dynamic Range",
    "Add Reverb",
    "Pitch Shift",
    "Echo",
    "Stereo Widening",
    "Bass Boost",
    "Treble Boost",
    "Normalize",
    "Noise Gate",
    "Limiter",
    "Phaser",
    "Flanger",
    "Bitcrusher",
    "Auto Gain",
    "Vocal Distortion",
    "Harmony",
    "Stage Mode"
]

with gr.Blocks(title="AI Audio Studio", css="style.css") as demo:
    gr.Markdown("## 🎧 Ultimate AI Audio Studio\nUpload, edit, export — powered by AI!")

    # --- Single File Studio ---
    with gr.Tab("🎵 Single File Studio"):
        gr.Interface(
            fn=process_audio,
            inputs=[
                gr.Audio(label="Upload Audio", type="filepath"),
                gr.CheckboxGroup(choices=effect_options, label="Apply Effects in Order"),
                gr.Checkbox(label="Isolate Vocals After Effects"),
                gr.Dropdown(choices=preset_names, label="Select Preset", value=preset_names[0] if preset_names else None),
                gr.Dropdown(choices=["MP3", "WAV"], label="Export Format", value="MP3")
            ],
            outputs=[
                gr.Audio(label="Processed Audio", type="filepath"),
                gr.Image(label="Waveform Preview"),
                gr.Textbox(label="Session Log (JSON)", lines=5),
                gr.Textbox(label="Detected Genre", lines=1),
                gr.Textbox(label="Status", value="✅ Ready", lines=1)
            ],
            title="Edit One File at a Time",
            description="Apply effects, preview waveform, and get full session log.",
            flagging_mode="never",
            submit_btn="Process Audio",
            clear_btn=None
        )

    # --- AI Mastering Chain Tab ===
    with gr.Tab("🎧 AI Mastering Chain"):
        gr.Interface(
            fn=ai_mastering_chain,
            inputs=[
                gr.Audio(label="Upload Track", type="filepath"),
                gr.Dropdown(choices=["Pop", "EDM", "Rock", "Hip-Hop", "Acoustic", "Metal", "Trap", "LoFi"], label="Genre", value="Pop"),
                gr.Slider(minimum=-24, maximum=-6, value=-14, label="Target LUFS")
            ],
            outputs=gr.Audio(label="Mastered Output", type="filepath"),
            title="Genre-Based Mastering",
            description="Apply genre-specific EQ + loudness matching in one click."
        )

    # --- Multiband Compression Tab ===
    with gr.Tab("🎛 Multiband Compression"):
        gr.Interface(
            fn=multiband_compression,
            inputs=[
                gr.Audio(label="Upload Track", type="filepath"),
                gr.Slider(minimum=-12, maximum=12, value=0, label="Low Gain (20–500Hz)"),
                gr.Slider(minimum=-12, maximum=12, value=0, label="Mid Gain (500Hz–4kHz)"),
                gr.Slider(minimum=-12, maximum=12, value=0, label="High Gain (4kHz+)"),
            ],
            outputs=gr.Audio(label="EQ'd Output", type="filepath"),
            title="Adjust Frequency Bands Live",
            description="Fine-tune your sound using real-time sliders for low, mid, and high frequencies."
        )

    # --- Real-Time Spectrum Analyzer + EQ Preview ===
    with gr.Tab("📊 Real-Time Spectrum"):
        gr.Interface(
            fn=visualize_spectrum,
            inputs=gr.Audio(label="Upload Track", type="filepath"),
            outputs=gr.Image(label="Spectrum Analysis"),
            title="See the frequency breakdown of your audio"
        )

    # --- Loudness Graph Tab ===
    with gr.Tab("📈 Loudness Graph"):
        gr.Interface(
            fn=match_loudness,
            inputs=[
                gr.Audio(label="Upload Track", type="filepath"),
                gr.Slider(minimum=-24, maximum=-6, value=-14, label="Target LUFS")
            ],
            outputs=gr.Audio(label="Normalized Output", type="filepath"),
            title="Match Loudness Across Tracks",
            description="Use EBU R128 standard for consistent volume"
        )

    # --- Stereo Imaging Tool ===
    with gr.Tab("🎚 Stereo Imaging"):
        gr.Interface(
            fn=stereo_imaging,
            inputs=[
                gr.Audio(label="Upload Track", type="filepath"),
                gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Mid-Side Balance"),
                gr.Slider(minimum=0.0, maximum=2.0, value=1.0, label="Stereo Spread")
            ],
            outputs=gr.Audio(label="Imaged Output", type="filepath"),
            title="Adjust Stereo Field",
            description="Control mid-side balance and widen stereo spread."
        )

    # --- Harmonic Saturation ===
    with gr.Tab("🧬 Harmonic Saturation"):
        gr.Interface(
            fn=harmonic_saturation,
            inputs=[
                gr.Audio(label="Upload Track", type="filepath"),
                gr.Slider(minimum=0.0, maximum=1.0, value=0.2, label="Saturation Intensity")
            ],
            outputs=gr.Audio(label="Warm Output", type="filepath"),
            title="Add Analog-Style Warmth",
            description="Apply subtle distortion to enhance clarity and presence."
        )

    # --- Sidechain Compression ===
    with gr.Tab("🔁 Sidechain Compression"):
        gr.Interface(
            fn=sidechain_compressor,
            inputs=[
                gr.File(label="Main Track"),
                gr.File(label="Sidechain Track"),
                gr.Slider(minimum=-24, maximum=0, value=-16, label="Threshold (dB)"),
                gr.Number(label="Ratio", value=4),
                gr.Number(label="Attack (ms)", value=5),
                gr.Number(label="Release (ms)", value=200)
            ],
            outputs=gr.Audio(label="Ducked Output", type="filepath"),
            title="Sidechain Compression",
            description="Automatically duck background under voice or kick"
        )

    # --- Save/Load Mix Session (.aiproj) ===
    with gr.Tab("📁 Save/Load Mix Session"):
        gr.Interface(
            fn=save_project,
            inputs=[
                gr.File(label="Vocals"),
                gr.File(label="Drums"),
                gr.File(label="Bass"),
                gr.File(label="Other"),
                gr.Slider(minimum=-10, maximum=10, value=0, label="Vocals Volume"),
                gr.Slider(minimum=-10, maximum=10, value=0, label="Drums Volume"),
                gr.Slider(minimum=-10, maximum=10, value=0, label="Bass Volume"),
                gr.Slider(minimum=-10, maximum=10, value=0, label="Other Volume"),
            ],
            outputs=gr.File(label="Project File (.aiproj)"),
            title="Save Your Full Mix Session",
            description="Save stems, volumes, and settings in one file."
        )

        gr.Interface(
            fn=load_project,
            inputs=gr.File(label="Upload .aiproj File"),
            outputs=[
                gr.File(label="Vocals"),
                gr.File(label="Drums"),
                gr.File(label="Bass"),
                gr.File(label="Other"),
                gr.Slider(label="Vocals Volume"),
                gr.Slider(label="Drums Volume"),
                gr.Slider(label="Bass Volume"),
                gr.Slider(label="Other Volume")
            ],
            title="Resume Last Mix",
            description="Load saved mix session",
            allow_flagging="never"
        )

    # --- Vocal Pitch Correction (Auto-Tune) ===
    with gr.Tab("🧬 Vocal Pitch Correction"):
        gr.Interface(
            fn=auto_tune_vocal,
            inputs=[
                gr.File(label="Source Voice Clip"),
                gr.Textbox(label="Target Key", value="C", lines=1)
            ],
            outputs=gr.Audio(label="Pitch-Corrected Output", type="filepath"),
            title="Auto-Tune Style Pitch Correction",
            description="Correct vocal pitch automatically"
        )

demo.launch()