import streamlit as st
from transformers import pipeline
from PIL import Image
import torch

# --- Config ---
st.set_page_config(page_title="Image → Prompt", page_icon="🖼️", layout="centered")

# Laad model (gecached)
@st.cache_resource
def load_model():
    return pipeline(
        "image-to-text",
        model="llava-hf/llava-v1.6-34b-hf",
        torch_dtype=torch.float16,
        device_map="auto"
    )

pipe = load_model()

# Stijlen
STIJLEN = {
    "Realistisch": ", fotorealistisch, 8k, scherp, natuurlijk licht",
    "Anime": ", anime stijl, gedetailleerde ogen, Studio Ghibli, 4k",
    "Olieverf": ", olieverf op doek, impressionistisch, Van Gogh stijl",
    "Cyberpunk": ", cyberpunk, neon lichten, regenachtige straten, Blade Runner",
    "Geen stijl": ""
}

# --- UI ---
st.title("🖼️ Image → Prompt Generator")
st.markdown("*Upload een foto, kies een stijl, krijg een perfecte AI-prompt.*")

col1, col2 = st.columns([1, 1])
with col1:
    uploaded_file = st.file_uploader("Kies een afbeelding", type=["png", "jpg", "jpeg"])
with col2:
    stijl = st.selectbox("Kies een stijl", options=list(STIJLEN.keys()))

if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Jouw afbeelding", use_column_width=True)

    with st.spinner("Prompt genereren..."):
        result = pipe(image)
        raw_prompt = result[0]["generated_text"]

        # Optimaliseer
        final_prompt = (
            raw_prompt
            .replace("a photo of", "")
            .replace("an image of", "")
            .strip()
            .capitalize()
        )
        final_prompt += STIJLEN[stijl]

    st.success("**Prompt:**")
    st.code(final_prompt, language=None)

    st.download_button(
        "📥 Download prompt",
        final_prompt,
        file_name="prompt.txt",
        mime="text/plain"
    )