app.py

from pathlib import Path
from PIL import Image
import streamlit as st
import insightface
from insightface.app import FaceAnalysis
from huggingface_hub import InferenceClient
import os
import random
import numpy as np
import yaml

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

MAX_SEED = np.iinfo(np.int32).max
client = InferenceClient()
DATA_PATH = Path("./data")
DATA_PATH.mkdir(exist_ok=True)
PREDEFINED_SEED = random.randint(0, MAX_SEED)
HF_TOKEN_UPSCALER = os.environ.get("HF_TOKEN")

if not HF_TOKEN_UPSCALER:
    st.warning("HF_TOKEN no está configurado. Algunas funcionalidades pueden no funcionar.")

def get_upscale_finegrain(prompt, img_path, upscale_factor):
    try:
        upscale_client = InferenceClient("fal/AuraSR-v2", hf_token=HF_TOKEN_UPSCALER)
        result = upscale_client.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 as e:
        st.error(f"Error al mejorar la imagen: {e}")
        return None

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

def list_saved_images():
    try:
        image_files = sorted(DATA_PATH.glob("*.jpg"))
        return image_files
    except Exception as e:
        st.error(f"Error al listar imágenes guardadas: {e}")
        return []

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

app, swapper = prepare_face_app()

def sort_faces(faces):
    return sorted(faces, key=lambda x: x.bbox[0])

def get_face(faces, face_id):
    if not faces or len(faces) < face_id:
        raise ValueError("Rostro no disponible.")
    return faces[face_id - 1]

def swap_faces(source_image, source_face_index, destination_image, destination_face_index):
    faces = sort_faces(app.get(source_image))
    source_face = get_face(faces, source_face_index)
    res_faces = sort_faces(app.get(destination_image))
    if destination_face_index > len(res_faces) or destination_face_index < 1:
        raise ValueError("Índice de rostro de destino no válido.")
    res_face = get_face(res_faces, destination_face_index)
    result = swapper.get(destination_image, res_face, source_face, paste_back=True)
    return result

def generate_image(prompt, width, height, seed, model_name):
    if seed == -1:
        seed = random.randint(0, MAX_SEED)
    try:
        image = client.text_to_image(prompt=prompt, height=height, width=width, model=model_name)
        return image, seed
    except Exception as e:
        st.error(f"Error al generar imagen: {e}")
        if hasattr(e, 'response') and e.response is not None:
            st.error(f"Detalles del error: {e.response.text}") 
        return None, seed

def display_gallery():
    st.header("Galería de Imágenes Guardadas")
    images = list_saved_images()
    if images:
        cols = st.columns(8)
        for i, image_file in enumerate(images):
            with cols[i % 8]:
                st.image(str(image_file), caption=image_file.name, use_column_width=True)
                prompt = get_prompt_for_image(image_file.name)
                st.write(prompt[:100])
                if st.button(f"FaceSwap", key=f"select_{i}_{image_file.name}"):
                    st.session_state['generated_image_path'] = str(image_file)
                    st.success("Imagen seleccionada")
                if st.button(f"Borrar", key=f"delete_{i}_{image_file.name}"):
                    if image_file.exists():
                        os.remove(image_file)
                        st.success("Imagen borrada")
                        display_gallery()
                    else:
                        st.warning("La imagen no existe.")
    else:
        st.info("No hay imágenes guardadas.")

def save_prompt(prompt):
    with open(DATA_PATH / "prompts.txt", "a") as f:
        f.write(prompt + "\n")
        st.success("Prompt guardado.")

def generate_variations(prompt, num_variants, use_enhanced):
    if use_enhanced:
        instructions = [
            "With this words, create a photorealistic description for a detailed txt2img prompt in English in 200 characters maximum",
            "With this idea, write a creative, realistic, and detailed text-to-image prompt in English in 200 characters maximum",
            "With this text, generate a descriptive and True to life txt2img prompt in English in 200 characters maximum",
            "With my idea, describe a photorealistic scene with detailed illumination for a txt2img prompt in English in 200 characters maximum",
            "With this concept, give a realistic, elegant txt2img prompt in English, emphasizing photorealism in 200 characters maximum",
            "With this perspective, conform a visually dynamic and hyperrealistic txt2img prompt in English in 200 characters maximum",
            "With this inspiration, realize a cinematic txt2img prompt in English with hyperrealistic elements in 200 characters maximum",
            "With my idea, make a lifelike and txt2img prompt in English, focusing on photorealistic depth in 200 characters maximum"
        ]
        prompts = set()
        while len(prompts) < num_variants:
            instruction = random.choice(instructions)
            enhanced_prompt = f"{instruction}: {prompt}"
            prompts.add(enhanced_prompt)
        return list(prompts)
    else:
        return [prompt] * num_variants

def get_prompt_for_image(image_name):
    prompts = {}
    try:
        with open(DATA_PATH / "prompts.txt", "r") as f:
            for line in f:
                if line.startswith(image_name):
                    prompts[image_name] = line.split(": ", 1)[1].strip()
    except FileNotFoundError:
        return "No hay prompt asociado."
    return prompts.get(image_name, "No hay prompt asociado.")

def login_form():
    st.title("Iniciar Sesión")
    username = st.text_input("Usuario", value="admin")
    password = st.text_input("Contraseña", value="flux3x", type="password") 
    if st.button("Iniciar Sesión"):
        if authenticate_user(username, password):
            st.success("Autenticación exitosa.")
            st.session_state['authenticated'] = True
        else:
            st.error("Credenciales incorrectas. Intenta de nuevo.")

def save_image(image, filename):
    try:
        image_path = DATA_PATH / filename
        if isinstance(image, bytes):
            with open(image_path, "wb") as f:
                f.write(image)
        else:
            image.save(image_path)
        return image_path
    except Exception as e:
        st.error(f"Error al guardar la imagen: {e}")
        return None

def upload_image_to_gallery():
    uploaded_image = st.sidebar.file_uploader("Sube una imagen a la galería", type=["jpg", "jpeg", "png"])
    if uploaded_image:
        image = Image.open(uploaded_image)
        image_path = save_image(image, f"{uploaded_image.name}")
        if image_path:
            save_prompt(f"{uploaded_image.name}: uploaded by user")
            st.sidebar.success(f"Imagen subida: {image_path}")

def gen(prompts, width, height, model_name, num_variants=1):
    images = []
    try:
        for idx, prompt in enumerate(prompts[:num_variants]):
            seed = random.randint(0, MAX_SEED)
            image, seed = generate_image(prompt, width, height, seed, model_name)
            image_path = save_image(image, f"generated_image_{seed}.jpg")
            if image_path:
                save_prompt(f"generated_image_{seed}.jpg: {prompt}")
                st.success(f"Imagen {idx + 1} generada")
                images.append(str(image_path))
    except Exception as e:
        st.error(f"Error al generar imágenes: {e}")
    return images

def main():
    st.set_page_config(layout="wide")
    if 'authenticated' not in st.session_state or not st.session_state['authenticated']:
        login_form()
        return
    st.title("Flux +Upscale +Prompt Enhancer +FaceSwap")
    generated_image_path = st.session_state.get('generated_image_path')
    prompt = st.sidebar.text_area("Descripción de la imagen", height=150, max_chars=500)
    format_option = st.sidebar.selectbox("Formato", ["9:16", "16:9", "1:1"])
    model_option = st.sidebar.selectbox("Modelo", ["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-dev"])
    upscale_checkbox = st.sidebar.checkbox("Escalar imagen")
    prompt_enhance = st.sidebar.checkbox("Mejorar Prompt", True)
    num_variants = st.sidebar.slider("Número de imágenes", 1, 8, 8)
    width, height = (720, 1280) if format_option == "9:16" else (1280, 720) if format_option == "16:9" else (1280, 1280)
    
    if prompt:
        prompts = generate_variations(prompt, num_variants=num_variants, use_enhanced=prompt_enhance)
        if st.sidebar.button("Generar Imágenes"):
            images = gen(prompts, width, height, model_option, num_variants)
            
    if generated_image_path and upscale_checkbox:
        upscale_factor = st.sidebar.slider("Factor de Escalado", 1, 4, 2)
        improved_image = get_upscale_finegrain(prompt, generated_image_path, upscale_factor)
        if improved_image:
            st.image(improved_image, caption="Imagen Escalada", use_column_width=True)
            
    upload_image_to_gallery()
    display_gallery()

if __name__ == "__main__":
    main()