8bitapp.py

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
from audiocraft.models import MusicGen
from torch.cuda.amp import autocast
import warnings
import random
from bitsandbytes.nn import Linear8bitLt

# Suppress warnings for cleaner output
warnings.filterwarnings("ignore")

# Set PYTORCH_CUDA_ALLOC_CONF to manage memory fragmentation
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

# Check critical dependencies
if np.__version__ != "1.23.5":
    print(f"WARNING: NumPy version {np.__version__} is being used. Tested with numpy==1.23.5.")
if not torch.__version__.startswith(("2.1.0", "2.3.1")):
    print(f"WARNING: PyTorch version {torch.__version__} may not be compatible. Expected torch==2.1.0 or 2.3.1.")

# 1) DEVICE SETUP
device = "cuda" if torch.cuda.is_available() else "cpu"
if device != "cuda":
    print("ERROR: CUDA is required for GPU rendering. CPU rendering is disabled.")
    sys.exit(1)
print(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")

# Pre-run memory cleanup
torch.cuda.empty_cache()
gc.collect()
torch.cuda.ipc_collect()
torch.cuda.synchronize()

# 2) LOAD MUSICGEN WITH 8-BIT QUANTIZATION
try:
    print("Loading MusicGen medium model with 8-bit quantization into VRAM...")
    local_model_path = "./models/musicgen-medium"
    if not os.path.exists(local_model_path):
        print(f"ERROR: Local model path {local_model_path} does not exist.")
        print("Please download the MusicGen medium model weights and place them in the correct directory.")
        sys.exit(1)
    
    # Load MusicGen model in FP16
    musicgen_model = MusicGen.get_pretrained(local_model_path, device=device)
    
    # Apply 8-bit quantization to the language model (lm) component
    def quantize_to_8bit(model):
        # Target the lm (language model) attribute, which contains the transformer
        if not hasattr(model, 'lm'):
            raise AttributeError("MusicGen model does not have 'lm' attribute for quantization.")
        lm = model.lm
        quantized_layers = 0
        for name, module in lm.named_modules():
            if isinstance(module, torch.nn.Linear):
                # Replace with 8-bit linear layer
                parent = lm
                name_parts = name.split('.')
                for part in name_parts[:-1]:
                    parent = getattr(parent, part)
                setattr(parent, name_parts[-1], Linear8bitLt(
                    module.in_features,
                    module.out_features,
                    bias=module.bias is not None,
                    has_fp16_weights=False,
                    threshold=6.0
                ))
                quantized_layers += 1
        print(f"Quantized {quantized_layers} linear layers to 8-bit.")
        return model
    
    # Quantize the model
    musicgen_model = quantize_to_8bit(musicgen_model)
    musicgen_model.to(device)
    
    # Set generation parameters
    musicgen_model.set_generation_params(
        duration=10,  # Default chunk duration
        two_step_cfg=False  # Disable two-step CFG for stability
    )
    print("Model successfully quantized to 8-bit and loaded.")
except Exception as e:
    print(f"ERROR: Failed to load or quantize MusicGen model: {e}")
    print("Ensure model weights and bitsandbytes are installed correctly.")
    sys.exit(1)

# 3) RESOURCE MONITORING FUNCTION
def print_resource_usage(stage: str):
    """Log GPU and CPU resource usage at a given stage."""
    print(f"--- {stage} ---")
    print(f"GPU Memory Allocated: {torch.cuda.memory_allocated() / (1024**3):.2f} GB")
    print(f"GPU Memory Reserved: {torch.cuda.memory_reserved() / (1024**3):.2f} GB")
    print(f"CPU Memory Used: {psutil.virtual_memory().percent}%")
    print("---------------")

# Check available GPU memory
def check_vram_availability(required_gb=4.0):
    """Check if sufficient VRAM is available for audio generation."""
    total_vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
    allocated_vram = torch.cuda.memory_allocated() / (1024**3)
    available_vram = total_vram - allocated_vram
    if available_vram < required_gb:
        print(f"WARNING: Low VRAM available ({available_vram:.2f} GB). Reduce total_duration or chunk_duration.")
    return available_vram >= required_gb

# 4) GENRE PROMPT FUNCTIONS
def set_red_hot_chili_peppers_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Red Hot Chili Peppers-inspired funk rock prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("strong rhythmic steps" if bpm > 120 else "groovy rhythmic flow")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ", groovy basslines"
    guitar = f", {guitar_style} guitar riffs" if guitar_style != "none" else ", syncopated guitar riffs"
    return f"Instrumental funk rock{bass}{guitar}{drum}{synth}, Red Hot Chili Peppers-inspired vibe with dynamic energy and funky breakdowns, {rhythm} at {bpm} BPM."

def set_nirvana_grunge_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Nirvana-inspired grunge prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("intense rhythmic steps" if bpm > 120 else "grungy rhythmic pulse")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ", melodic basslines"
    guitar = f", {guitar_style} guitar riffs" if guitar_style != "none" else ", raw distorted guitar riffs"
    return f"Instrumental grunge{bass}{guitar}{drum}{synth}, Nirvana-inspired angst-filled sound with quiet-loud dynamics, {rhythm} at {bpm} BPM."

def set_pearl_jam_grunge_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Pearl Jam-inspired grunge prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("soulful rhythmic steps" if bpm > 120 else "driving rhythmic flow")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ", deep bass"
    guitar = f", {guitar_style} guitar leads" if guitar_style != "none" else ", soulful guitar leads"
    return f"Instrumental grunge{bass}{guitar}{drum}{synth}, Pearl Jam-inspired emotional intensity with soaring choruses, {rhythm} at {bpm} BPM."

def set_soundgarden_grunge_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Soundgarden-inspired grunge prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("heavy rhythmic steps" if bpm > 120 else "sludgy rhythmic groove")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ""
    guitar = f", {guitar_style} guitar riffs" if guitar_style != "none" else ", heavy sludgy guitar riffs"
    return f"Instrumental grunge{bass}{guitar}{drum}{synth}, Soundgarden-inspired dark, psychedelic edge, {rhythm} at {bpm} BPM."

def set_foo_fighters_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Foo Fighters-inspired alternative rock prompt."""
    styles = ["anthemic", "gritty", "melodic", "fast-paced", "driving"]
    tempos = ["upbeat", "mid-tempo", "high-energy"]
    moods = ["energetic", "introspective", "rebellious", "uplifting"]
    style = random.choice(styles)
    tempo = random.choice(tempos)
    mood = random.choice(moods)
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("powerful rhythmic steps" if bpm > 120 else "catchy rhythmic groove")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ""
    guitar = f", {guitar_style} guitar riffs" if guitar_style != "none" else f", {style} guitar riffs"
    return f"Instrumental alternative rock{bass}{guitar}{drum}{synth}, Foo Fighters-inspired {mood} vibe with powerful choruses, {rhythm} at {bpm} BPM."

def set_smashing_pumpkins_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Smashing Pumpkins-inspired alternative rock prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("dynamic rhythmic steps" if bpm > 120 else "dreamy rhythmic flow")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ""
    guitar = f", {guitar_style} guitar textures" if guitar_style != "none" else ", dreamy guitar textures"
    return f"Instrumental alternative rock{bass}{guitar}{drum}{synth}, Smashing Pumpkins-inspired blend of melancholy and aggression, {rhythm} at {bpm} BPM."

def set_radiohead_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Radiohead-inspired experimental rock prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("complex rhythmic steps" if bpm > 120 else "intricate rhythmic pulse")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ", atmospheric synths"
    bass = f", {bass_style}" if bass_style != "none" else ""
    guitar = f", {guitar_style} guitar layers" if guitar_style != "none" else ", intricate guitar layers"
    return f"Instrumental experimental rock{bass}{guitar}{drum}{synth}, Radiohead-inspired blend of introspective and innovative soundscapes, {rhythm} at {bpm} BPM."

def set_classic_rock_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Led Zeppelin-inspired classic rock prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("bluesy rhythmic steps" if bpm > 120 else "steady rhythmic groove")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ", groovy bass"
    guitar = f", {guitar_style} electric guitars" if guitar_style != "none" else ", bluesy electric guitars"
    return f"Instrumental classic rock{bass}{guitar}{drum}{synth}, Led Zeppelin-inspired raw energy with dynamic solos, {rhythm} at {bpm} BPM."

def set_alternative_rock_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Pixies-inspired alternative rock prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("quirky rhythmic steps" if bpm > 120 else "energetic rhythmic flow")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ", melodic basslines"
    guitar = f", {guitar_style} guitar riffs" if guitar_style != "none" else ", distorted guitar riffs"
    return f"Instrumental alternative rock{bass}{guitar}{drum}{synth}, Pixies-inspired quirky, energetic vibe, {rhythm} at {bpm} BPM."

def set_post_punk_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Joy Division-inspired post-punk prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("sharp rhythmic steps" if bpm > 120 else "moody rhythmic pulse")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ", driving basslines"
    guitar = f", {guitar_style} guitars" if guitar_style != "none" else ", jangly guitars"
    return f"Instrumental post-punk{bass}{guitar}{drum}{synth}, Joy Division-inspired moody, atmospheric sound, {rhythm} at {bpm} BPM."

def set_indie_rock_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate an Arctic Monkeys-inspired indie rock prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("catchy rhythmic steps" if bpm > 120 else "jangly rhythmic flow")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ""
    guitar = f", {guitar_style} guitars" if guitar_style != "none" else ", jangly guitars"
    return f"Instrumental indie rock{bass}{guitar}{drum}{synth}, Arctic Monkeys-inspired blend of catchy riffs, {rhythm} at {bpm} BPM."

def set_funk_rock_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Rage Against the Machine-inspired funk rock prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("aggressive rhythmic steps" if bpm > 120 else "funky rhythmic groove")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ""
    bass = f", {bass_style}" if bass_style != "none" else ", slap bass"
    guitar = f", {guitar_style} guitar chords" if guitar_style != "none" else ", funky guitar chords"
    return f"Instrumental funk rock{bass}{guitar}{drum}{synth}, Rage Against the Machine-inspired mix of groove and aggression, {rhythm} at {bpm} BPM."

def set_detroit_techno_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Juan Atkins-inspired Detroit techno prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("pulsing rhythmic steps" if bpm > 120 else "deep rhythmic groove")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ", crisp hi-hats"
    synth = f", {synthesizer} accents" if synthesizer != "none" else ", deep pulsing synths"
    bass = f", {bass_style}" if bass_style != "none" else ", driving basslines"
    guitar = f", {guitar_style} guitars" if guitar_style != "none" else ""
    return f"Instrumental Detroit techno{bass}{guitar}{drum}{synth}, Juan Atkins-inspired rhythmic groove, {rhythm} at {bpm} BPM."

def set_deep_house_prompt(bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style):
    """Generate a Larry Heard-inspired deep house prompt."""
    rhythm = f" with {rhythmic_steps}" if rhythmic_steps != "none" else ("soulful rhythmic steps" if bpm > 120 else "laid-back rhythmic flow")
    drum = f", {drum_beat} drums" if drum_beat != "none" else ""
    synth = f", {synthesizer} accents" if synthesizer != "none" else ", warm analog synth chords"
    bass = f", {bass_style}" if bass_style != "none" else ", deep basslines"
    guitar = f", {guitar_style} guitars" if guitar_style != "none" else ""
    return f"Instrumental deep house{bass}{guitar}{drum}{synth}, Larry Heard-inspired laid-back groove, {rhythm} at {bpm} BPM."

# 5) AUDIO PROCESSING FUNCTIONS
def apply_eq(segment):
    """Apply basic equalization to the audio segment."""
    # Low-pass filter to cut high frequencies above 8kHz
    segment = segment.low_pass_filter(8000)
    # High-pass filter to cut low frequencies below 80Hz
    segment = segment.high_pass_filter(80)
    return segment

def apply_fade(segment, fade_in_duration=1000, fade_out_duration=1000):
    """Apply fade-in and fade-out effects to the audio segment."""
    segment = segment.fade_in(fade_in_duration)
    segment = segment.fade_out(fade_out_duration)
    return segment

# 6) GENERATION & I/O FUNCTIONS
def generate_music(instrumental_prompt: str, cfg_scale: float, top_k: int, top_p: float, temperature: float, total_duration: int, chunk_duration: int, crossfade_duration: int, bpm: int, drum_beat: str, synthesizer: str, rhythmic_steps: str, bass_style: str, guitar_style: str):
    """Generate instrumental audio with 8-bit quantized model, ensuring controlled volume."""
    global musicgen_model
    # Validate input prompt
    if not instrumental_prompt.strip():
        return None, "⚠️ Please enter a valid instrumental prompt!"
    
    try:
        start_time = time.time()
        
        # Validate total_duration (restricted to radio button choices)
        if total_duration not in {30, 60, 90, 120}:
            raise ValueError(f"Invalid total_duration: {total_duration}. Must be 30, 60, 90, or 120 seconds.")
        
        # Clamp chunk_duration between 5 and 15 seconds
        chunk_duration = min(max(chunk_duration, 5), 15)
        num_chunks = max(1, total_duration // chunk_duration)
        chunk_duration = total_duration / num_chunks  # Adjust for even splitting
        overlap_duration = min(1.0, crossfade_duration / 1000.0)  # Convert ms to seconds, cap at 1s
        generation_duration = chunk_duration + overlap_duration
        
        # Validate sample rate from model
        sample_rate = getattr(musicgen_model, 'sample_rate', None)
        if not isinstance(sample_rate, int) or sample_rate <= 0:
            raise ValueError("Invalid sample_rate from musicgen_model.")
        
        audio_segments = []
        
        # Check VRAM availability (increased to 4.0GB due to 8-bit quantization)
        if not check_vram_availability(required_gb=4.0):
            return None, "⚠️ Insufficient VRAM for generation. Reduce total_duration or chunk_duration."
        
        print("Generating audio with 8-bit quantized model...")
        seed = 42  # Fixed seed for reproducibility
        torch.manual_seed(seed)
        np.random.seed(seed)
        
        # Generate audio chunks
        for i in range(num_chunks):
            chunk_prompt = instrumental_prompt
            print(f"Generating chunk {i+1}/{num_chunks} on GPU (prompt: {chunk_prompt})...")
            
            # Set model generation parameters
            musicgen_model.set_generation_params(
                duration=generation_duration,
                use_sampling=True,
                top_k=top_k,
                top_p=top_p,
                temperature=temperature,
                cfg_coef=cfg_scale
            )
            
            print_resource_usage(f"Before Chunk {i+1} Generation")
            
            # Generate audio with mixed precision
            with torch.no_grad():
                with autocast():
                    audio_chunk = musicgen_model.generate([chunk_prompt], progress=True)[0]
            
            # Move to CPU and ensure float32 dtype
            audio_chunk = audio_chunk.cpu().to(dtype=torch.float32)
            
            # Ensure audio is stereo (2 channels)
            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)
            
            # Strict validation of audio shape
            if audio_chunk.dim() != 2 or audio_chunk.shape[0] != 2:
                raise ValueError(f"Chunk {i+1}: Expected stereo audio with shape (2, samples), got {audio_chunk.shape}")
            
            # Save temporary WAV file and load with pydub
            temp_wav_path = f"temp_chunk_{i}.wav"
            torchaudio.save(temp_wav_path, audio_chunk, sample_rate, bits_per_sample=24)
            segment = AudioSegment.from_wav(temp_wav_path)
            os.remove(temp_wav_path)
            audio_segments.append(segment)
            
            # Clean up GPU memory
            torch.cuda.empty_cache()
            gc.collect()
            torch.cuda.ipc_collect()
            torch.cuda.synchronize()
            time.sleep(0.5)  # Brief pause to stabilize resources
            print_resource_usage(f"After Chunk {i+1} Generation")
        
        # Combine audio chunks
        print("Combining audio chunks...")
        if not audio_segments:
            raise ValueError("No audio segments generated.")
        
        final_segment = audio_segments[0]
        for i in range(1, len(audio_segments)):
            next_segment = audio_segments[i]
            next_segment = next_segment + 1  # Adjust amplitude if needed
            final_segment = final_segment.append(next_segment, crossfade=crossfade_duration)
        
        # Trim to exact total duration
        final_segment = final_segment[:total_duration * 1000]
        
        # Post-processing
        print("Post-processing final track...")
        final_segment = apply_eq(final_segment)
        
        # Volume adjustment: Set peak to -12 dBFS (significantly quieter than 0 dBFS)
        desired_peak = -12.0  # Target peak level in dBFS
        if not isinstance(desired_peak, (int, float)) or desired_peak >= 0:
            raise ValueError(f"Invalid desired_peak: {desired_peak}. Must be negative dBFS.")
        
        current_peak = final_segment.max_dBFS
        if current_peak is None or not isinstance(current_peak, (int, float)):
            raise ValueError("Failed to compute current peak level.")
        
        print(f"Current peak level: {current_peak:.2f} dBFS")
        gain_needed = desired_peak - current_peak
        
        # Validate gain adjustment
        if not -100 <= gain_needed <= 100:  # Reasonable range to prevent extreme adjustments
            raise ValueError(f"Gain adjustment out of bounds: {gain_needed} dB")
        
        print(f"Applying gain adjustment: {gain_needed:.2f} dB")
        final_segment = final_segment.apply_gain(gain_needed)
        
        # Verify new peak level
        new_peak = final_segment.max_dBFS
        if new_peak is None or not isinstance(new_peak, (int, float)):
            raise ValueError("Failed to compute new peak level after gain adjustment.")
        
        print(f"New peak level: {new_peak:.2f} dBFS")
        
        # Check if new peak is within 0.1 dB of desired peak
        if abs(new_peak - desired_peak) > 0.1:
            print(f"Warning: New peak {new_peak:.2f} dBFS deviates from desired {desired_peak:.2f} dBFS")
        
        final_segment = apply_fade(final_segment)
        
        # Export to MP3 (increased to 320kbps for higher quality)
        mp3_path = "output_cleaned.mp3"
        final_segment.export(
            mp3_path,
            format="mp3",
            bitrate="320k",  # Increased from 128k to 320k for better audio quality
            tags={"title": "GhostAI Instrumental", "artist": "GhostAI"}
        )
        print(f"Saved final audio to {mp3_path}")
        
        # Optional: Export to WAV for lossless quality (uncomment if desired)
        # wav_path = "output_cleaned.wav"
        # final_segment.export(wav_path, format="wav")
        # print(f"Saved lossless audio to {wav_path}")
        
        print_resource_usage("After Final Generation")
        print(f"Total Generation Time: {time.time() - start_time:.2f} seconds")
        
        return mp3_path, "✅ Done! Generated instrumental audio."
    
    except Exception as e:
        print(f"Error during generation: {e}")
        return None, f"❌ Generation failed: {e}"
    
    finally:
        # Clean up GPU memory
        torch.cuda.empty_cache()
        gc.collect()
        torch.cuda.ipc_collect()
        torch.cuda.synchronize()

# Function to clear inputs
def clear_inputs():
    """Reset all input fields to default values."""
    return "", 3.0, 250, 0.9, 1.0, 30, 10, 1000, 120, "none", "none", "none", "none", "none"

# 7) CUSTOM CSS
css = """
body { 
    background: linear-gradient(135deg, #0A0A0A 0%, #1C2526 100%); 
    color: #E0E0E0; 
    font-family: 'Orbitron', sans-serif; 
}
.header-container { 
    text-align: center; 
    padding: 10px 20px; 
    background: rgba(0, 0, 0, 0.9); 
    border-bottom: 1px solid #00FF9F; 
}
#ghost-logo { 
    font-size: 40px; 
    animation: glitch-ghost 1.5s infinite; 
}
h1 { 
    color: #A100FF; 
    font-size: 24px; 
    animation: glitch-text 2s infinite; 
}
p { 
    color: #E0E0E0; 
    font-size: 12px; 
}
.input-container, .settings-container, .output-container { 
    max-width: 1200px; 
    margin: 20px auto; 
    padding: 20px; 
    background: rgba(28, 37, 38, 0.8); 
    border-radius: 10px; 
}
.textbox { 
    background: #1A1A1A; 
    border: 1px solid #A100FF; 
    color: #E0E0E0; 
}
.genre-buttons { 
    display: flex; 
    justify-content: center; 
    flex-wrap: wrap; 
    gap: 15px; 
}
.genre-btn, button { 
    background: linear-gradient(45deg, #A100FF, #00FF9F); 
    border: none; 
    color: #0A0A0A; 
    padding: 10px 20px; 
    border-radius: 5px; 
}
.gradio-container { 
    padding: 20px; 
}
.group-container { 
    margin-bottom: 20px; 
    padding: 15px; 
    border: 1px solid #00FF9F; 
    border-radius: 8px; 
}
@keyframes glitch-ghost { 
    0% { transform: translate(0, 0); opacity: 1; }
    20% { transform: translate(-5px, 2px); opacity: 0.8; }
    100% { transform: translate(0, 0); opacity: 1; }
}
@keyframes glitch-text { 
    0% { transform: translate(0, 0); }
    20% { transform: translate(-2px, 1px); }
    100% { transform: translate(0, 0); }
}
@font-face { 
    font-family: 'Orbitron'; 
    src: url('https://fonts.gstatic.com/s/orbitron/v29/yMJRMIlzdpvBhQQL_Qq7dy0.woff2') format('woff2'); 
}
"""

# 8) BUILD WITH BLOCKS
with gr.Blocks(css=css) as demo:
    gr.Markdown("""
        <div class="header-container">
            <div id="ghost-logo">👻</div>
            <h1>GhostAI Music Generator 🎹</h1>
            <p>Summon the Sound of the Unknown</p>
        </div>
    """)
    
    with gr.Column(elem_classes="input-container"):
        gr.Markdown("### 🎸 Prompt Settings")
        instrumental_prompt = gr.Textbox(
            label="Instrumental Prompt ✍️",
            placeholder="Click a genre button or type your own instrumental prompt",
            lines=4,
            elem_classes="textbox"
        )
        with gr.Row(elem_classes="genre-buttons"):
            rhcp_btn = gr.Button("Red Hot Chili Peppers 🌶️", elem_classes="genre-btn")
            nirvana_btn = gr.Button("Nirvana Grunge 🎸", elem_classes="genre-btn")
            pearl_jam_btn = gr.Button("Pearl Jam Grunge 🦪", elem_classes="genre-btn")
            soundgarden_btn = gr.Button("Soundgarden Grunge 🌑", elem_classes="genre-btn")
            foo_fighters_btn = gr.Button("Foo Fighters 🤘", elem_classes="genre-btn")
            smashing_pumpkins_btn = gr.Button("Smashing Pumpkins 🎃", elem_classes="genre-btn")
            radiohead_btn = gr.Button("Radiohead 🧠", elem_classes="genre-btn")
            classic_rock_btn = gr.Button("Classic Rock 🎸", elem_classes="genre-btn")
            alternative_rock_btn = gr.Button("Alternative Rock 🎵", elem_classes="genre-btn")
            post_punk_btn = gr.Button("Post-Punk 🖤", elem_classes="genre-btn")
            indie_rock_btn = gr.Button("Indie Rock 🎤", elem_classes="genre-btn")
            funk_rock_btn = gr.Button("Funk 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")
    
    with gr.Column(elem_classes="settings-container"):
        gr.Markdown("### ⚙️ API Settings")
        with gr.Group(elem_classes="group-container"):
            cfg_scale = gr.Slider(
                label="CFG Scale 🎯",
                minimum=1.0,
                maximum=10.0,
                value=3.0,
                step=0.1,
                info="Controls how closely the music follows the prompt."
            )
            top_k = gr.Slider(
                label="Top-K Sampling 🔢",
                minimum=10,
                maximum=500,
                value=250,
                step=10,
                info="Limits sampling to the top k most likely tokens."
            )
            top_p = gr.Slider(
                label="Top-P Sampling 🎰",
                minimum=0.0,
                maximum=1.0,
                value=0.9,
                step=0.05,
                info="Keeps tokens with cumulative probability above p."
            )
            temperature = gr.Slider(
                label="Temperature 🔥",
                minimum=0.1,
                maximum=2.0,
                value=1.0,
                step=0.1,
                info="Controls randomness; higher values increase diversity."
            )
            total_duration = gr.Radio(
                label="Song Length ⏳ (seconds)",
                choices=[30, 60, 90, 120],
                value=30,
                info="Select the total duration of the track."
            )
            chunk_duration = gr.Slider(
                label="Chunk Duration ⏱️ (seconds)",
                minimum=5,
                maximum=15,
                value=10,
                step=1,
                info="Duration of each chunk to render (5 to 15 seconds)."
            )
            crossfade_duration = gr.Slider(
                label="Crossfade Duration 🎶 (ms)",
                minimum=100,
                maximum=2000,
                value=1000,
                step=100,
                info="Crossfade duration between chunks."
            )

        gr.Markdown("### 🎵 Musical Controls")
        with gr.Group(elem_classes="group-container"):
            bpm = gr.Slider(
                label="Tempo 🎵 (BPM)",
                minimum=60,
                maximum=180,
                value=120,
                step=1,
                info="Beats per minute to set the track's tempo."
            )
            drum_beat = gr.Dropdown(
                label="Drum Beat 🥁",
                choices=["none", "standard rock", "funk groove", "techno kick", "jazz swing"],
                value="none",
                info="Select a drum beat style to influence the rhythm."
            )
            synthesizer = gr.Dropdown(
                label="Synthesizer 🎹",
                choices=["none", "analog synth", "digital pad", "arpeggiated synth"],
                value="none",
                info="Select a synthesizer style for electronic accents."
            )
            rhythmic_steps = gr.Dropdown(
                label="Rhythmic Steps 👣",
                choices=["none", "syncopated steps", "steady steps", "complex steps"],
                value="none",
                info="Select a rhythmic step style to enhance the beat."
            )
            bass_style = gr.Dropdown(
                label="Bass Style 🎸",
                choices=["none", "slap bass", "deep bass", "melodic bass"],
                value="none",
                info="Select a bass style to shape the low end."
            )
            guitar_style = gr.Dropdown(
                label="Guitar Style 🎸",
                choices=["none", "distorted", "clean", "jangle"],
                value="none",
                info="Select a guitar style to define the riffs."
            )

        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")
        out_audio = gr.Audio(label="Generated Instrumental Track 🎵", type="filepath")
        status = gr.Textbox(label="Status 📢", interactive=False)
    
    rhcp_btn.click(set_red_hot_chili_peppers_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    nirvana_btn.click(set_nirvana_grunge_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    pearl_jam_btn.click(set_pearl_jam_grunge_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    soundgarden_btn.click(set_soundgarden_grunge_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    foo_fighters_btn.click(set_foo_fighters_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    smashing_pumpkins_btn.click(set_smashing_pumpkins_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    radiohead_btn.click(set_radiohead_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    classic_rock_btn.click(set_classic_rock_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    alternative_rock_btn.click(set_alternative_rock_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    post_punk_btn.click(set_post_punk_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    indie_rock_btn.click(set_indie_rock_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    funk_rock_btn.click(set_funk_rock_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    detroit_techno_btn.click(set_detroit_techno_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    deep_house_btn.click(set_deep_house_prompt, inputs=[bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style], outputs=instrumental_prompt)
    gen_btn.click(
        generate_music,
        inputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, chunk_duration, crossfade_duration, bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style],
        outputs=[out_audio, status]
    )
    clr_btn.click(
        clear_inputs,
        inputs=None,
        outputs=[instrumental_prompt, cfg_scale, top_k, top_p, temperature, total_duration, chunk_duration, crossfade_duration, bpm, drum_beat, synthesizer, rhythmic_steps, bass_style, guitar_style]
    )

# 9) TURN OFF OPENAPI/DOCS
app = demo.launch(
    server_name="0.0.0.0",
    server_port=9999,
    share=True,
    inbrowser=False,
    show_error=True
)
try:
    fastapi_app = demo._server.app
    fastapi_app.docs_url = None
    fastapi_app.redoc_url = None
    fastapi_app.openapi_url = None
except Exception:
    pass