app.py

import os
import numpy as np
import random
from pathlib import Path
from PIL import Image
import streamlit as st
from huggingface_hub import AsyncInferenceClient
import asyncio
import yaml

# Configuración de cliente de inferencia
client = AsyncInferenceClient()
DATA_PATH = Path("./data")
DATA_PATH.mkdir(exist_ok=True)

# Cargar credenciales desde archivo YAML
with open("credentials.yaml", "r") as file:
    credentials = yaml.safe_load(file)

# Variables de estado
if 'logged_in' not in st.session_state:
    st.session_state.logged_in = False
if 'prompt' not in st.session_state:
    st.session_state.prompt = ""

# Función de generación de imágenes
async def generate_image(prompt, width, height, seed):
    try:
        if seed == -1:
            seed = random.randint(0, np.iinfo(np.int32).max)
        image = await client.text_to_image(
            prompt=prompt, height=height, width=width, model="enhanceaiteam/Flux-uncensored"
        )
        return image
    except Exception as e:
        return f"Error al generar imagen: {e}"

# Función para guardar prompts
def save_prompt(prompt_text, seed):
    try:
        prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
        with open(prompt_file_path, "w") as prompt_file:
            prompt_file.write(prompt_text)
        return str(prompt_file_path)
    except Exception as e:
        st.error(f"Error al guardar el prompt: {e}")
        return None

# Función para guardar imágenes
def save_image(image, seed):
    try:
        image_path = DATA_PATH / f"image_{seed}.jpg"
        image.save(image_path, format="JPEG")
        return str(image_path)
    except Exception as e:
        st.error(f"Error al guardar la imagen: {e}")
        return None

# Función para obtener archivos almacenados
def get_storage():
    files = [file for file in DATA_PATH.glob("*.jpg") if file.is_file()]
    files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
    usage = sum([file.stat().st_size for file in files])
    return [str(file.resolve()) for file in files], f"Uso total: {usage/(1024.0 ** 3):.3f}GB"

# Función principal
async def gen(prompt, width, height):
    seed = random.randint(0, np.iinfo(np.int32).max)
    progress_bar = st.progress(0)
    image = await generate_image(prompt, width, height, seed)
    progress_bar.progress(100)

    if isinstance(image, str) and image.startswith("Error"):
        progress_bar.empty()
        return [image, None]

    image_path = save_image(image, seed)
    prompt_file_path = save_prompt(prompt, seed)

    return [image_path, prompt_file_path]

def main():
    st.set_page_config(layout="wide")
    st.title("Generador de Imágenes FLUX")

    # Inicio de sesión
    if not st.session_state.logged_in:
        username = st.sidebar.text_input("Usuario")
        password = st.sidebar.text_input("Contraseña", type="password")

        if st.sidebar.button("Iniciar Sesión"):
            if username == credentials["username"] and password == credentials["password"]:
                st.session_state.logged_in = True
            else:
                st.error("Credenciales incorrectas.")
        return

    # Opciones después de iniciar sesión
    prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
    format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])

    width, height = (720, 1280) if format_option == "9:16" else (1280, 720)

    if st.sidebar.button("Generar Imagen"):
        with st.spinner("Generando imagen..."):
            result = asyncio.run(gen(prompt, width, height))
            image_paths = result[0]
            prompt_file = result[1]

        if image_paths and Path(image_paths).exists():
            st.image(image_paths, caption="Imagen Generada")
            if prompt_file and Path(prompt_file).exists():
                prompt_text = Path(prompt_file).read_text()
                st.write(f"Prompt utilizado: {prompt_text}")
        else:
            st.error("El archivo de imagen no existe.")

    files, usage = get_storage()
    st.text(usage)
    cols = st.columns(6)

    for idx, file in enumerate(files):
        with cols[idx % 6]:
            image = Image.open(file)
            prompt_text = Path(DATA_PATH / f"prompt_{file.stem.replace('image_', '')}.txt").read_text() if Path(DATA_PATH / f"prompt_{file.stem.replace('image_', '')}.txt").exists() else "No disponible"
            st.image(image, caption=f"Imagen {idx+1}")
            st.write(f"Prompt: {prompt_text}")

            if st.button(f"Borrar Imagen {idx+1}", key=f"delete_{idx}"):
                try:
                    os.remove(file)
                    os.remove(DATA_PATH / f"prompt_{file.stem.replace('image_', '')}.txt")
                    st.success(f"Imagen {idx+1} y su prompt fueron borrados.")
                except Exception as e:
                    st.error(f"Error al borrar la imagen o prompt: {e}")

if __name__ == "__main__":
    main()