app.py

import streamlit as st
from gradio_client import Client, handle_file
import threading
import time
import os
import uuid
import tempfile
import shutil

# ---------------------------
# CONFIG
# ---------------------------
IMAGE_DIR = "data/images"
os.makedirs(IMAGE_DIR, exist_ok=True)

client = Client("Galaxydude2/a")
client_process = Client("SLSLK/n8")

# ---------------------------
# SESSION
# ---------------------------
if "gallery" not in st.session_state:
    st.session_state.gallery = []

if "selected_image" not in st.session_state:
    st.session_state.selected_image = None

if "page" not in st.session_state:
    st.session_state.page = "🖼️ Generieren"

if "prompt_history" not in st.session_state:
    st.session_state.prompt_history = []

# ---------------------------
# THEME
# ---------------------------
st.set_page_config(layout="wide")

st.markdown("""
<style>
.stApp { background-color: #0f0f1a; color: #e0e0ff; }
h1 { color: #ff2e82; text-align:center; }
.stButton > button {
    background-color: #ff2e82;
    color: white;
    border-radius: 10px;
}
.stButton > button:hover { background-color: #ff5ca8; }
</style>
""", unsafe_allow_html=True)

st.markdown("<h1>🎨 Omni Image Studio</h1>", unsafe_allow_html=True)

# ---------------------------
# HELPERS
# ---------------------------
def save_local_copy(src_path):
    if os.path.exists(src_path):
        new_path = os.path.join(IMAGE_DIR, f"{uuid.uuid4()}.png")
        shutil.copy(src_path, new_path)
        return new_path
    return None


def prepare_image(input_image):
    if hasattr(input_image, "getvalue"):
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        tmp.write(input_image.getvalue())
        tmp.close()
        return tmp.name
    elif isinstance(input_image, str):
        return input_image
    return None


# ---------------------------
# API FUNCTIONS
# ---------------------------
def generate_image(prompt):
    return client.predict(
        prompt=prompt,
        api_name="/generate_image"
    )


def edit_image(input_image, prompt):
    image_path = prepare_image(input_image)

    return client.predict(
        input_image=handle_file(image_path),
        edit_prompt=prompt,
        api_name="/edit_image"
    )


def process_image(input_image):
    image_path = prepare_image(input_image)

    return client_process.predict(
        image=handle_file(image_path),
        api_name="/process_image"
    )

# ---------------------------
# ASYNC + PROGRESS
# ---------------------------
def run_async(func, *args):
    result = {"done": False, "data": None}

    def task():
        result["data"] = func(*args)
        result["done"] = True

    threading.Thread(target=task).start()
    return result


def show_progress(job):
    bar = st.progress(0)
    status = st.empty()
    percent = 0

    messages = [
        "🔄 Processing...",
        "🎨 Generating pixels...",
        "🧠 AI thinking...",
        "✨ Almost done..."
    ]

    i = 0
    while not job["done"]:
        percent = min(percent + 5, 95)
        bar.progress(percent)

        status.text(messages[i % len(messages)])
        i += 1

        time.sleep(0.3)

    bar.progress(100)
    status.text("✅ Done!")

# ---------------------------
# NAV
# ---------------------------
page = st.radio(
    "Navigation",
    ["🖼️ Generieren", "✨ Bearbeiten", "🧪 Analyse", "📚 Galerie"],
    index=["🖼️ Generieren", "✨ Bearbeiten", "🧪 Analyse", "📚 Galerie"].index(st.session_state.page),
    horizontal=True
)

st.session_state.page = page

# ---------------------------
# GENERATE
# ---------------------------
if page == "🖼️ Generieren":
    prompt = st.text_area("Prompt")

    # 🧠 Prompt History
    if st.session_state.prompt_history:
        st.subheader("🧠 Prompt History")
        for i, p in enumerate(reversed(st.session_state.prompt_history[-5:])):
            if st.button(p, key=f"hist{i}"):
                prompt = p

    if st.button("Generieren"):
        if prompt:
            st.session_state.prompt_history.append(prompt)

        job = run_async(generate_image, prompt)
        show_progress(job)

        result = job["data"]

        if result:
            st.image(result)

            # 📥 Download
            with open(result, "rb") as f:
                st.download_button("📥 Download", f, file_name="generated.png")

            path = save_local_copy(result)
            if path:
                st.session_state.gallery.append({
                    "path": path,
                    "prompt": prompt
                })

# ---------------------------
# EDIT
# ---------------------------
elif page == "✨ Bearbeiten":
    if st.session_state.selected_image:
        st.image(st.session_state.selected_image)
        input_image = st.session_state.selected_image
    else:
        input_image = st.file_uploader("Upload Image")

    prompt = st.text_area("Edit Prompt")

    if st.button("Bearbeiten"):
        job = run_async(edit_image, input_image, prompt)
        show_progress(job)

        result = job["data"]

        if result:
            st.image(result)

            with open(result, "rb") as f:
                st.download_button("📥 Download", f, file_name="edited.png")

            path = save_local_copy(result)
            if path:
                st.session_state.gallery.append({
                    "path": path,
                    "prompt": prompt
                })

# ---------------------------
# ANALYSE
# ---------------------------
elif page == "🧪 Analyse":
    if st.session_state.selected_image:
        st.image(st.session_state.selected_image)
        input_image = st.session_state.selected_image
    else:
        input_image = st.file_uploader("Upload Image")

    if st.button("Analysieren"):
        job = run_async(process_image, input_image)
        show_progress(job)

        result = job["data"]

        if result:
            st.image(result)

            with open(result, "rb") as f:
                st.download_button("📥 Download", f, file_name="analysis.png")

            if os.path.exists(result):
                path = save_local_copy(result)
                st.session_state.gallery.append({
                    "path": path,
                    "prompt": "Analyse"
                })

# ---------------------------
# GALLERY
# ---------------------------
elif page == "📚 Galerie":
    if not st.session_state.gallery:
        st.info("Keine Bilder")
    else:
        cols = st.columns(3)

        for i, item in enumerate(st.session_state.gallery):
            with cols[i % 3]:
                st.image(item["path"], caption=item["prompt"])

                # 📥 Download
                with open(item["path"], "rb") as f:
                    st.download_button("📥", f, file_name=f"img_{i}.png", key=f"d{i}")

                c1, c2 = st.columns(2)

                if c1.button("✏️", key=f"e{i}"):
                    st.session_state.selected_image = item["path"]
                    st.session_state.page = "✨ Bearbeiten"
                    st.rerun()

                if c2.button("🧪", key=f"a{i}"):
                    st.session_state.selected_image = item["path"]
                    st.session_state.page = "🧪 Analyse"
                    st.rerun()

        if st.button("🗑️ Clear"):
            st.session_state.gallery = []