import streamlit as st
import requests
from PIL import Image
from io import BytesIO
import os

st.set_page_config(page_title="AI Design Generator", layout="centered")
st.title("🎨 AI Poster/Packaging Design Generator")
st.write("Enter a prompt to generate your custom design:")

prompt = st.text_input("Prompt", "modern minimalist packaging with red and white theme")

if st.button("Generate Design"):
    with st.spinner("Generating..."):
        API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2"
        headers = {"Authorization": f"Bearer {os.environ['HF_TOKEN']}"}

        payload = {"inputs": prompt}
        response = requests.post(API_URL, headers=headers, json=payload)

        if response.status_code == 200:
            image = Image.open(BytesIO(response.content))
            st.image(image, caption="Generated Design")
        else:
            st.error("Failed to generate image. Try again later.")