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

with open("config.yaml") as f: config = yaml.safe_load(f); username = config["credentials"]["username"]; password = config["credentials"]["password"]

client = AsyncInferenceClient(); DATA_PATH = Path("./data"); DATA_PATH.mkdir(exist_ok=True); MAX_SEED = np.iinfo(np.int32).max; PREDEFINED_SEED = random.randint(0, MAX_SEED)

async def generate_image(prompt, width, height, seed):
    try:
        seed = int(seed) if seed != -1 else PREDEFINED_SEED
        return await client.text_to_image(prompt=prompt, height=height, width=width, model="enhanceaiteam/Flux-uncensored"), seed
    except Exception as e:
        return f"Error al generar imagen: {e}", None

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 prompt_file_path
    except Exception as e:
        st.error(f"Error al guardar el prompt: {e}"); return None

async def gen(prompt, width, height):
    seed = PREDEFINED_SEED; progress_bar = st.progress(0)
    image, seed = 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]
    return [save_image(image, seed), save_prompt(prompt, seed)]

def save_image(image, seed):
    try:
        image_path = DATA_PATH / f"image_{seed}.jpg"
        image.save(image_path, format="JPEG")
        return image_path
    except Exception as e:
        st.error(f"Error al guardar la imagen: {e}"); return None

def get_storage():
    files = sorted(DATA_PATH.glob("*.jpg"), 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"

def get_prompts():
    return {file.stem.replace("prompt_", ""): file for file in DATA_PATH.glob("*.txt") if file.is_file()}

def main():
    st.set_page_config(layout="wide"); st.title("Generador de Imágenes")
    prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
    width, height = (720, 1280) if st.sidebar.selectbox("Formato", ["9:16", "16:9"]) == "9:16" else (1280, 720)

    if st.sidebar.button("Generar Imagen"):
        with st.spinner("Generando imagen..."):
            result = asyncio.run(gen(prompt, width, height))
            if (image_paths := result[0]) and Path(image_paths).exists(): st.image(image_paths, caption="Imagen Generada")
            else: st.error("El archivo de imagen no existe.")
            if (prompt_file := result[1]) and Path(prompt_file).exists():
                st.write(f"Prompt utilizado: {Path(prompt_file).read_text()}")
            else: st.write("El archivo del prompt no está disponible.")

    files, usage = get_storage(); st.text(usage); cols = st.columns(6); prompts = get_prompts()
    for idx, file in enumerate(files):
        with cols[idx % 6]:
            image = Image.open(file); prompt_file = prompts.get(Path(file).stem.replace("image_", ""), None)
            st.image(image, caption=f"Imagen {idx+1}"); st.write(f"Prompt: {Path(prompt_file).read_text() if prompt_file else 'No disponible'}")
            if st.button(f"Borrar Imagen {idx+1}", key=f"delete_{idx}"):
                try: os.remove(file); os.remove(prompt_file) if prompt_file else None; 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()