import spaces  # Für ZeroGPU
import gradio as gr
from transformers import pipeline
import scipy.io.wavfile
import os

# Sicherstellen, dass der 'static'-Ordner existiert
os.makedirs('static', exist_ok=True)

# GPU-abhängige Funktion zur Musikgenerierung
@spaces.GPU(duration=120)  # GPU wird für diese Funktion angefordert
def generate_music(prompt, duration, diffusion_steps, cfg_scale):
    try:
        # MusicGen-Pipeline initialisieren (nur innerhalb der GPU-Session)
        synthesizer = pipeline(
            "text-to-audio",
            "facebook/musicgen-small",
            model_kwargs={"attn_implementation": "eager"}
        )
        
        # Musik generieren
        music = synthesizer(prompt, forward_params={"do_sample": True})
        
        # Datei speichern
        output_path = "static/generated_music.wav"
        scipy.io.wavfile.write(
            output_path,
            rate=music["sampling_rate"],
            data=music["audio"]
        )
        return output_path
    except Exception as e:
        return f"Error: {str(e)}"

# Benutzerdefiniertes CSS
custom_css = """
body {
    background-color: #121212;
    color: #ffffff;
    font-family: 'Arial', sans-serif;
    margin: 0;
    padding: 0;
}

#title {
    text-align: center;
    font-size: 24px;
    font-weight: bold;
    margin-bottom: 10px;
}

#description {
    text-align: center;
    font-size: 16px;
    margin-bottom: 30px;
}

button {
    background-color: #ff5722;
    color: white;
    font-weight: bold;
    border-radius: 8px;
    padding: 10px 20px;
    border: none;
    cursor: pointer;
    font-size: 16px;
}

button:hover {
    background-color: #ff784e;
}

.slider {
    accent-color: #ff5722;
}

textarea, input[type="text"] {
    background-color: #1e1e2f;
    color: white;
    border: 1px solid #444;
    padding: 10px;
    border-radius: 5px;
    font-size: 14px;
}

audio {
    border: 2px solid #444;
    border-radius: 5px;
    margin-top: 10px;
}
"""

# Gradio-Interface erstellen
with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("<h1 id='title'>Stable Audio Generator</h1>")
    gr.Markdown("<p id='description'>Generate variable-length stereo audio at 44.1kHz from text prompts using Stable Audio Open 1.0.</p>")

    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(
                label="Prompt", 
                placeholder="Describe your music (e.g., 'Relaxing piano music')."
            )
            duration = gr.Slider(
                label="Duration in Seconds", 
                minimum=10, 
                maximum=60, 
                step=1, 
                value=30
            )
            diffusion_steps = gr.Slider(
                label="Number of Diffusion Steps", 
                minimum=50, 
                maximum=200, 
                step=10, 
                value=150
            )
            cfg_scale = gr.Slider(
                label="CFG Scale", 
                minimum=1, 
                maximum=20, 
                step=1, 
                value=10
            )
            generate_button = gr.Button("Submit")
        with gr.Column(scale=1):
            output = gr.Audio(
                label="Generated Music", 
                type="filepath", 
                autoplay=True
            )
    
    with gr.Row():
        clear_button = gr.Button("Clear")
    
    generate_button.click(
        fn=generate_music, 
        inputs=[prompt, duration, diffusion_steps, cfg_scale], 
        outputs=output
    )

    clear_button.click(
        fn=lambda: None, 
        inputs=None, 
        outputs=[prompt, output]
    )

# Anwendung starten
if __name__ == "__main__":
    demo.launch()