app.py

import os
import random
import asyncio
import numpy as np
from pathlib import Path
from PIL import Image
from insightface.app import FaceAnalysis
import streamlit as st
from huggingface_hub import InferenceClient, AsyncInferenceClient
from gradio_client import Client, handle_file
import yaml
import insightface

MAX_SEED = np.iinfo(np.int32).max
DATA_PATH = Path("./data")
DATA_PATH.mkdir(exist_ok=True)
HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN_UPSCALER")
client, llm_client = AsyncInferenceClient(), InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")

try: 
    credentials = yaml.safe_load(open("config.yaml"))
except Exception as e: 
    st.error(f"Error al cargar config: {e}")
    credentials = {"username": "", "password": ""}

def prepare_face_app():
    app = FaceAnalysis(name='buffalo_l')
    app.prepare(ctx_id=0, det_size=(640, 640))
    return app, insightface.model_zoo.get_model('onix.onnx')

app, swapper = prepare_face_app()

async def generate_image(prompt, model, w, h, scale, steps, seed):
    seed = random.randint(0, MAX_SEED) if seed == -1 else seed
    image = await client.text_to_image(prompt=prompt, height=h, width=w, guidance_scale=scale, num_inference_steps=steps, model=model)
    return image, seed if not isinstance(image, str) else (None, None)

def get_upscale_finegrain(prompt, img_path, upscale_factor):
    try:
        result = Client("finegrain/finegrain-image-enhancer", hf_token=HF_TOKEN_UPSCALER).predict(input_image=handle_file(img_path), prompt=prompt, upscale_factor=upscale_factor)
        return result[1] if isinstance(result, list) and len(result) > 1 else None
    except Exception:
        return None

async def gen(prompt, basemodel, w, h, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language):
    combined_prompt = f"{prompt} {await improve_prompt(prompt, language)}" if process_enhancer else prompt
    image, seed = await generate_image(combined_prompt, basemodel, w, h, scales, steps, seed)
    if image is None: return ["Error al generar imagen", None, combined_prompt]
    image_path = save_image(image, seed)
    prompt_file_path = save_prompt(combined_prompt, seed)
    if process_upscale:
        upscale_image_path = get_upscale_finegrain(combined_prompt, image_path, upscale_factor)
        if upscale_image_path:
            Image.open(upscale_image_path).save(DATA_PATH / f"upscale_image_{seed}.jpg", format="JPEG")
            image_path.unlink() 
            return [str(DATA_PATH / f"upscale_image_{seed}.jpg"), str(prompt_file_path)]
    return [str(image_path), str(prompt_file_path)]

async def improve_prompt(prompt, language):
    try:
        instruction_en = "With this idea, describe in English a detailed txt2img prompt in 500 characters at most, add illumination, atmosphere, cinematic elements, and characters if need it..."
        instruction_es = "Con esta idea, describe en español un prompt detallado de txt2img en un máximo de 500 caracteres, con iluminación, atmósfera, elementos cinematográficos y en su caso personajes..."
        instruction = instruction_en if language == "en" else instruction_es
        formatted_prompt = f"{prompt}: {instruction}"
        response = llm_client.text_generation(formatted_prompt, max_new_tokens=500)
        improved_text = response['generated_text'].strip() if 'generated_text' in response else response.strip()
        return improved_text[:500] if len(improved_text) > 500 else improved_text
    except Exception as e:
        return f"Error mejorando el prompt: {e}"

def save_image(image, seed):
    if image.mode == 'RGBA': image = image.convert('RGB')
    image_path = DATA_PATH / f"image_{seed}.jpg"
    image.save(image_path, format="JPEG")
    return image_path

def save_prompt(prompt_text, seed):
    prompt_file_path = DATA_PATH / f"prompt_{seed}.txt"
    open(prompt_file_path, "w").write(prompt_text)
    return prompt_file_path

def get_storage():
    files = [file for file in DATA_PATH.glob("*.jpg") if file.is_file()]
    total_size = sum([file.stat().st_size for file in files]) / (1024.0 ** 3)
    return files, f"Uso total: {total_size:.3f} GB"

def delete_image(image_path):
    try: 
        Path(image_path).unlink()
        st.success(f"Imagen {image_path} borrada.")
    except Exception as e: 
        st.error(f"Error al borrar imagen: {e}")

def delete_all_images():
    for file in DATA_PATH.glob("*.jpg"): 
        file.unlink()
    st.success("Todas las imágenes han sido borradas.")

def authenticate_user(username, password, credentials):
    return username == credentials["username"] and password == credentials["password"]

def login_form(credentials):
    if 'authenticated' not in st.session_state: 
        st.session_state['authenticated'] = False
    if not st.session_state['authenticated']:
        username = st.text_input("Usuario")
        password = st.text_input("Contraseña", type='password')
        if st.button("Iniciar Sesión"):
            if authenticate_user(username, password, credentials):
                st.session_state['authenticated'] = True
                st.success("Inicio de sesión exitoso.")
            else: 
                st.error("Credenciales incorrectas.")

def upload_image():
    uploaded_file = st.sidebar.file_uploader("Sube una imagen", type=["png", "jpg", "jpeg"])
    if uploaded_file:
        image_path = DATA_PATH / uploaded_file.name
        with open(image_path, "wb") as f: 
            f.write(uploaded_file.getbuffer())
        st.sidebar.success(f"Imagen {uploaded_file.name} cargada correctamente.")
        return image_path, save_prompt("#uploadedbyuser", image_path.stem)
    return None

def gallery():
    files, usage = get_storage()
    st.sidebar.write(f"{usage}")
    if st.sidebar.button("Borrar Todas las Imágenes"): 
        delete_all_images()
    cols = st.columns(6)
    for idx, file in enumerate(files):
        with cols[idx % 6]:
            st.image(str(file))
            try:
                prompt_file_path = DATA_PATH / f"prompt_{file.stem.split('_')[-1]}.txt"
                st.write(f"Prompt: {open(prompt_file_path).read()}")
            except FileNotFoundError:
                st.write("Prompt no encontrado.")
            st.button(f"Borrar Imagen {file.name}", on_click=delete_image, args=(file,))
            if st.button(f"Swap Face en {file.name}"): 
                upload_source_and_swap(file)

def upload_source_and_swap(image_path):
    source_image = st.file_uploader("Sube la imagen source para face swap", type=["png", "jpg", "jpeg"])
    if source_image:
        source_image_path = DATA_PATH / source_image.name
        with open(source_image_path, "wb") as f: 
            f.write(source_image.getbuffer())
        st.success(f"Imagen source {source_image.name} cargada correctamente.")
        dest_face_index = st.slider("Selecciona la posición de la cara destino", 1, 4, 1)
        swapped_image_path = face_swap(image_path, source_image_path, dest_face_index)
        if swapped_image_path: 
            st.image(str(swapped_image_path), caption="Imagen con Face Swap", use_column_width=True)

def face_swap(image_path, source_image_path, dest_face_index):
    try:
        img_dest, img_src = Image.open(image_path), Image.open(source_image_path)
        faces = app.get(img_src)
        if not faces: 
            st.error("No se encontraron caras en la imagen source.")
            return None
        swapped_img = swapper.get(img_dest, faces[0], dest_index=dest_face_index - 1)  # Convertir a índice base 0
        swapped_img_path = DATA_PATH / f"swapped_{Path(image_path).stem}.jpg"
        swapped_img.save(swapped_img_path, format="JPEG")
        Path(image_path).unlink() 
        return swapped_img_path
    except Exception as e: 
        st.error(f"Error en face swap: {e}")
        return None

async def main():
    st.set_page_config(layout="wide")
    login_form(credentials)
    if not st.session_state['authenticated']:
        st.warning("Por favor, inicia sesión para acceder a la aplicación.")
        return

    prompt = st.sidebar.text_input("Descripción de la imagen", max_chars=900)
    process_enhancer, language = st.sidebar.checkbox("Mejorar Prompt", value=False), st.sidebar.selectbox("Idioma", ["en", "es"])
    basemodel, format_option, process_upscale = st.sidebar.selectbox("Modelo Base", ["black-forest-labs/FLUX.1-DEV", "black-forest-labs/FLUX.1-schnell"]), st.sidebar.selectbox("Formato", ["9:16", "16:9"]), st.sidebar.checkbox("Procesar Escalador", value=False)
    upscale_factor, scales, steps, seed = st.sidebar.selectbox("Factor de Escala", [2, 4, 8], index=0), st.sidebar.slider("Escalado", 1, 20, 10), st.sidebar.slider("Pasos", 1, 100, 20), st.sidebar.number_input("Semilla", value=-1)
    destination_face_position = st.sidebar.slider("Posición de la Cara de Destino", 1, 4, 1)
    
    w, h = (720, 1280) if format_option == "9:16" else (1280, 720)
    upload_image()

    image_path, prompt_file_path = None, None  
    if st.sidebar.button("Generar Imagen"):
        with st.spinner("Generando..."):
            image_path, prompt_file_path = await gen(prompt, basemodel, w, h, scales, steps, seed, upscale_factor, process_upscale, process_enhancer, language)
            if image_path:
                st.image(image_path, caption="Imagen Generada")
                st.write(f"Prompt: {open(prompt_file_path).read()}")
    
    if image_path:
        st.success("Imagen generada y almacenada.")
    
    gallery()

if __name__ == "__main__":
    asyncio.run(main())