import streamlit as st
import torch
from diffusers import AudioLDM2Pipeline
import scipy
import tempfile
import os

# Titel
st.set_page_config(page_title="💨 text2fart", page_icon="💀")
st.title("💨 text2fart mit Streamlit")
st.markdown("Schreib einen beliebigen Text – und hör den furzigen Klang deiner Kreativität! 😆")

# Prompt-Eingabe
prompt = st.text_input("Text eingeben:", placeholder="z. B. 'Apokalyptischer Trompetenfurz bei Sonnenuntergang'")

# Model laden (nur einmal)
@st.cache_resource
def load_model():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    # pipe = DiffusionPipeline.from_pretrained("spaceinvader/text2fart")
    pipe = AudioLDM2Pipeline.from_pretrained(
        "spaceinvader/text2fart",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        variant="fp16" if torch.cuda.is_available() else None,
    )
    return pipe.to(device)

pipe = load_model()

# Button
if st.button("💥 Furz generieren"):
    if not prompt.strip():
        st.warning("Gib bitte etwas ein. Selbst KI braucht Inspiration 😅")
    else:
        with st.spinner("Furze werden destilliert..."):
            result = pipe(prompt, num_inference_steps=20)
            audio = result["audio"][0]
            sample_rate = 16000

            # Temporäre WAV-Datei
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmpfile:
                scipy.io.wavfile.write(tmpfile.name, rate=sample_rate, data=audio)
                st.audio(tmpfile.name, format="audio/wav")
                st.success("Text erfolgreich vertont 💀💨")
                # Datei am Ende löschen (optional)
                os.unlink(tmpfile.name)