app.py

import os, zipfile, yaml, numpy as np, random, asyncio
from pathlib import Path
from PIL import Image
import streamlit as st
from huggingface_hub import AsyncInferenceClient
from moviepy.editor import ImageSequenceClip

try: 
    credentials = yaml.safe_load(open("config.yaml"))
except: 
    st.error("Error al cargar el archivo de configuración.")
    credentials = {"username": "", "password": ""}

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

async def generate_image(prompt, width, height, seed, model_name): 
    try: 
        return await client.text_to_image(prompt=prompt, height=height, width=width, model=model_name), int(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"
        open(prompt_file_path, "w").write(prompt_text)
        return prompt_file_path
    except Exception as e: 
        st.error(f"Error al guardar el prompt: {e}")

async def gen(prompt, width, height, model_name): 
    seed, progress_bar = PREDEFINED_SEED, st.progress(0)
    image, seed = await generate_image(prompt, width, height, seed, model_name)
    progress_bar.progress(100)
    if isinstance(image, str) and image.startswith("Error"): 
        return [image, None]
    return [str(save_image(image, seed)), str(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}")

def get_storage(): 
    files = sorted([file for file in DATA_PATH.glob("*.jpg")], key=lambda x: x.stat().st_mtime, reverse=True)
    return [str(file.resolve()) for file in files], f"Uso total: {sum([file.stat().st_size for file in files])/(1024.0 ** 3):.3f}GB"

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

def delete_all_images(): 
    try: 
        [os.remove(file) for file in DATA_PATH.glob("*.jpg") + DATA_PATH.glob("*.txt")]
        st.success("Todas las imágenes y prompts han sido borrados.")
    except Exception as e: 
        st.error(f"Error al borrar archivos: {e}")

def download_images_as_zip(): 
    zip_path = DATA_PATH / "images.zip"
    zipf = zipfile.ZipFile(zip_path, 'w')
    [zipf.write(file, arcname=file.name) for file in DATA_PATH.glob("*.jpg")]
    with open(zip_path, "rb") as zip_file: 
        st.download_button(label="Descargar imágenes en .zip", data=zip_file, file_name="images.zip", mime="application/zip")

def create_video_from_images(): 
    try: 
        image_files = sorted(DATA_PATH.glob("*.jpg"))
        image_sequence = [Image.open(image_file) for image_file in image_files]
    except: 
        st.error("No hay imágenes disponibles para crear un video.")
        return
    video_path = DATA_PATH / "output_video.mp4"
    ImageSequenceClip([np.array(img) for img in image_sequence], fps=1).write_videofile(str(video_path), codec="libx264")
    return video_path

def main():
    st.set_page_config(layout="wide")
    st.title("Generador de Imágenes Flux")
    
    if "authenticated" not in st.session_state: 
        st.session_state.authenticated = False

    if not st.session_state.authenticated:
        username, password = st.text_input("Usuario"), st.text_input("Contraseña", type="password")
        if st.button("Ingresar") and username == credentials["username"] and password == credentials["password"]: 
            st.session_state.authenticated = True
            st.success("Inicio de sesión exitoso.")
        elif st.button("Ingresar"): 
            st.error("Usuario o contraseña incorrectos.")
        return

    prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=500)
    format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9"])
    model_option = st.sidebar.selectbox("Modelo", ["enhanceaiteam/Flux-Uncensored-V2", "enhanceaiteam/Flux-uncensored"])
    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, model_option))
            image_paths, prompt_file = result[0], result[1]
            if Path(image_paths).exists(): 
                st.image(image_paths, caption="Imagen Generada")
            if prompt_file and Path(prompt_file).exists(): 
                st.write(f"Prompt utilizado: {Path(prompt_file).read_text()}")

    files, usage = get_storage()
    st.text(usage)
    cols, prompts = st.columns(6), 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}"): 
                os.remove(file)
                if prompt_file: os.remove(prompt_file)
                st.success(f"Imagen {idx+1} y su prompt fueron borrados.")

    if st.sidebar.button("Borrar todas las imágenes"): 
        delete_all_images()
    if st.sidebar.button("Descargar imágenes en .zip"): 
        download_images_as_zip()
    if st.button("Generar video con las imágenes"): 
        video_path = create_video_from_images()
        if video_path: 
            st.video(str(video_path), format="video/mp4")

if __name__ == "__main__": 
    main()