app.py

prompt = "A peaceful mountain landscape at sunset, digital painting, high detail."
import streamlit as st
from io import BytesIO
import math
import os
from PIL import Image
import requests
import time
import dotenv

dotenv.load_dotenv()

STABILITY_KEY = os.getenv("STABILITY_KEY")


def send_generation_request(host, params):
    """Send a REST request to the Stability API and return the Response object.

    Opens any files referenced in params and ensures they are closed after the request.
    """
    headers = {
        "Accept": "image/*",
        "Authorization": f"Bearer {STABILITY_KEY}"
    }

    # Encode parameters
    files = {}
    opened_files = []
    image = params.pop("image", None)
    mask = params.pop("mask", None)
    if image is not None and image != '':
        f = open(image, 'rb')
        files["image"] = f
        opened_files.append(f)
    if mask is not None and mask != '':
        f = open(mask, 'rb')
        files["mask"] = f
        opened_files.append(f)
    if len(files) == 0:
        # keep the same form as original script
        files["none"] = ''

    # Send request
    response = requests.post(
        host,
        headers=headers,
        files=files,
        data=params
    )

    # close opened files
    for f in opened_files:
        try:
            f.close()
        except Exception:
            pass

    if not response.ok:
        raise Exception(f"HTTP {response.status_code}: {response.text}")

    return response


DEFAULT_PROMPT = (
    "A peaceful mountain landscape at sunset, digital painting, high detail."
)


def main():
    st.set_page_config(page_title="Generador de imágenes - SD3", layout="centered")
    st.title("Generador de imágenes (SD3)")
    st.write("Introduce un prompt y genera una imagen usando la API de Stability.")

    # Requirements / key check
    env_exists = os.path.exists(".env")
    if not env_exists or not STABILITY_KEY:
        st.error(
            "No se ha encontrado la variable de entorno `STABILITY_KEY`.\n\n"
            "Por favor, añade un archivo `.env` en la raíz del proyecto o exporta la variable de entorno."
        )
        st.markdown("[Abrir README (instrucciones)](README.md)")
        return

    # UI inputs (Spanish)
    prompt = st.text_area("Prompt", value=DEFAULT_PROMPT, height=120)
    col1, col2 = st.columns(2)
    with col1:
        aspect_ratio = st.selectbox("Relación de aspecto", options=["1:1", "16:9", "9:16"], index=0)
    with col2:
        output_format = st.selectbox("Formato de salida", options=["jpeg", "png"], index=0)

    seed_input = st.text_input("Semilla (opcional)", value="")

    generate_btn = st.button("Generar imagen")

    if generate_btn:
        seed = int(seed_input) if seed_input.isdigit() else math.floor(time.time())
        host = "https://api.stability.ai/v2beta/stable-image/generate/sd3"

        params = {
            "prompt": prompt,
            "aspect_ratio": aspect_ratio,
            "seed": seed,
            "output_format": output_format,
            "model": "sd3-turbo"
        }

        with st.spinner("Generando imagen, espera por favor..."):
            try:
                response = send_generation_request(host, params)

                output_image = response.content
                finish_reason = response.headers.get("finish-reason")
                seed_returned = response.headers.get("seed") or seed

                if finish_reason == 'CONTENT_FILTERED':
                    st.warning("La generación fue filtrada por el clasificador NSFW.")

                # Display image
                try:
                    img = Image.open(BytesIO(output_image))
                    st.image(img, caption=f"Imagen generada — seed: {seed_returned}")
                except Exception:
                    # If PIL cannot open, still provide raw bytes preview
                    st.write("No se pudo abrir la imagen con PIL, mostrando como bytes.")
                    st.code(output_image[:100])

                filename = f"generated_{seed_returned}.{output_format}"
                st.download_button("Descargar imagen", data=output_image, file_name=filename, mime=f"image/{output_format}")

            except Exception as e:
                st.error(f"Error al generar la imagen: {e}")


if __name__ == "__main__":
    main()