import gradio as gr
import librosa
import librosa.display
import numpy as np
import matplotlib.pyplot as plt
import torch
from diffusers import StableDiffusionPipeline

# Load AI Art Model (Stable Diffusion)
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe.to("cpu")

def process_audio(audio_file):
    # Load audio file
    y, sr = librosa.load(audio_file, sr=22050)

    # Convert to Mel spectrogram
    mel_spectrogram = librosa.feature.melspectrogram(y=y, sr=sr)
    mel_spectrogram = librosa.power_to_db(mel_spectrogram, ref=np.max)

    # Save spectrogram image
    plt.figure(figsize=(6, 4))
    librosa.display.specshow(mel_spectrogram, sr=sr, x_axis='time', y_axis='mel')
    plt.colorbar(format='%+2.0f dB')
    plt.title("Mel Spectrogram")
    spectrogram_path = "spectrogram.png"
    plt.savefig(spectrogram_path)
    plt.close()

    # Generate AI Art (Based on Audio)
    prompt = "An abstract painting inspired by music and sound waves"
    image = pipe(prompt).images[0]
    image_path = "generated_art.png"
    image.save(image_path)

    return spectrogram_path, image_path

# Create Web Interface
interface = gr.Interface(
    fn=process_audio,
    inputs=gr.Audio(type="filepath"),
    outputs=[gr.Image(type="filepath", label="Spectrogram"), 
             gr.Image(type="filepath", label="AI-Generated Art")],
    title="Audio2Art 🎨",
    description="Upload an audio file, and this AI will generate an artwork inspired by the sound."
)

interface.launch()