app.py

from pathlib import Path
from PIL import Image
import streamlit as st
import os, random, numpy as np, yaml, time
from dataclasses import dataclass
from typing import List
from huggingface_hub import InferenceClient
import google.generativeai as genai

st.set_page_config(layout="wide")
HF_TOKEN = os.getenv("HF_TOKEN")
GEMINI_TOKEN = os.getenv("GEMINI_TOKEN")

if not HF_TOKEN:
    st.error("Error en el token de Hugging Face! 'HF_TOKEN'.")
    st.stop()

if not GEMINI_TOKEN:
    st.error("Error en el token de Gemini! 'GEMINI_TOKEN'.")
    st.stop()

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": ""}

@dataclass
class AppConfig:
    MAX_SEED: int = 1000000  
    CLEANUP_DAYS: int = 7    

MAX_SEED = AppConfig.MAX_SEED
DATA_PATH = Path("./data")
DATA_PATH.mkdir(exist_ok=True)

def get_hf_client():
    return InferenceClient(token=HF_TOKEN)

def get_gemini_client():
    genai.configure(api_key=GEMINI_TOKEN)
    return genai.GenerativeModel('gemini-1.5-flash-8b')

hf_client = get_hf_client()
gemini_client = get_gemini_client()

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

def list_saved_images(): 
    return sorted(DATA_PATH.glob("*.jpg"), key=lambda x: x.stat().st_mtime, reverse=True)

def enhance_prompt(text, client=None, use_enhancement=True):
    if not use_enhancement:
        return text[:200]
    
    if not client:
        return text[:200]
    
    try:
        enhancement_instruction = "Enhance this text description to be more suitable for text-to-image generation. Focus on vivid, descriptive language that will help an AI generate a photorealistic image. Be specific about colors, composition, lighting, and key details."
        
        response = client.generate_content(f"{enhancement_instruction}\n\nOriginal prompt: {text}")
        enhanced_prompt = response.text.strip()
        
        return enhanced_prompt[:200]
    
    except Exception as e:
        st.warning(f"Prompt enhancement error: {e}")
        return text[:200]

def save_prompt(image_name, enhanced_prompt):
    with open(DATA_PATH / "prompts.txt", "a") as f: 
        f.write(f"{image_name}: {enhanced_prompt}\n")

def generate_variations(prompt, num_variants=8, use_enhanced=True):
    instructions = [
        "Photorealistic description for txt2img prompt: ",
        "Creative, realistic text-to-image prompt: ",
        "Descriptive, true-to-life txt2img prompt: ",
        "Naturalistic scene with detailed prompt: ",
        "Realistic, elegant txt2img prompt: ",
        "Visually dynamic, hyperrealistic prompt: ",
        "Cinematic txt2img with hyperrealistic elements: ",
        "Lifelike txt2img, focusing on photorealistic depth: "
    ]
    if use_enhanced:
        prompts = [enhance_prompt(f"{instructions[i % len(instructions)]}{prompt}", client=gemini_client, use_enhancement=True) for i in range(num_variants)]
    else: 
        prompts = [enhance_prompt(prompt, use_enhancement=False) for i in range(num_variants)]
    return prompts

def generate_image(prompt, width, height, seed, model_name, client=None):
    if not client:
        st.error("No Hugging Face client available")
        return None, seed, None

    try:
        with st.spinner("Generando imagen..."):
            seed = int(seed) if seed != -1 else random.randint(0, AppConfig.MAX_SEED)
            enhanced_prompt = enhance_prompt(prompt, client=gemini_client)
            image = client.text_to_image(
                prompt=enhanced_prompt, 
                height=height, 
                width=width, 
                model=model_name, 
                seed=seed
            )
            return image, seed, enhanced_prompt
    except Exception as e:
        st.error(f"Image generation error: {e}")
        return None, seed, None

def gen(prompts, width, height, model_name, num_variants=8):
    images = []
    seeds = []
    while len(seeds) < num_variants:
        seed = random.randint(0, MAX_SEED)
        if seed not in seeds: 
            seeds.append(seed)
    
    for i in range(num_variants):
        current_prompt = prompts[i] if len(prompts) > i else prompts[-1]
        with st.spinner(f"Generando imagen {i+1}/{num_variants}"):
            image, used_seed, enhanced_prompt = generate_image(current_prompt, width, height, seeds[i], model_name, client=hf_client)
            if image:
                image_path = DATA_PATH / f"generated_image_{used_seed}.jpg"
                image.save(image_path)
                save_prompt(f"generated_image_{used_seed}.jpg", enhanced_prompt)
                images.append((str(image_path), enhanced_prompt))
                st.success(f"Imagen {i+1} generada")
    return images

def get_prompt_for_image(image_name):
    try:
        with open(DATA_PATH / "prompts.txt", "r") as f:
            for line in f:
                if line.startswith(image_name):
                    original_prompt = line.split(": ", 1)[1].strip()
                    image_path = DATA_PATH / image_name
                    
                    if image_path.exists():
                        return original_prompt
    except FileNotFoundError: 
        return "No hay prompt asociado"
    return "No hay prompt asociado"
    
def display_gallery():
    st.header("Galería de Imágenes Guardadas")
    images = list_saved_images()
    
    if images:
        cols = st.columns(4) 
        for i, image_file in enumerate(images):
            with cols[i % 4]:
                st.image(str(image_file), use_column_width=True)
                prompt = get_prompt_for_image(image_file.name)
                st.caption(prompt[:250])
                if st.button(f"Borrar", key=f"delete_{i}_{image_file}"):
                    if image_file.exists():
                        os.remove(image_file)
                        st.success("Imagen borrada")
                        st.rerun()

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.session_state['authenticated'] = True
            st.success("Autenticación exitosa.")
        else: 
            st.error("Credenciales incorrectas. Intenta de nuevo.")

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 = DATA_PATH / uploaded_image.name
        image.save(image_path)
        save_prompt(uploaded_image.name, "uploaded by user")
        st.sidebar.success(f"Imagen subida: {image_path}")

def main():
    if 'authenticated' not in st.session_state or not st.session_state['authenticated']: 
        login_form()
        return
    
    st.title("Flux +Upscale +Prompt Enhancer")
    
    if not hf_client or not gemini_client:
        st.error("No se pudo establecer conexión con los servicios. Verifique sus tokens.")
        return
    
    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", ["salomonsky/flux3"])
    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 st.sidebar.button("Generar Imágenes"):
        if prompt:
            prompts = generate_variations(prompt, num_variants=num_variants, use_enhanced=prompt_enhance)
            generated_images = gen(prompts, width, height, model_option, num_variants)
            
            st.header("Imágenes Generadas")
            cols = st.columns(4)
            for i, (image_path, image_prompt) in enumerate(generated_images):
                with cols[i % 4]:
                    st.image(image_path, use_column_width=True)
                    st.caption(image_prompt)
        else:
            st.warning("Por favor, ingrese una descripción de imagen")
        
    upload_image_to_gallery()
    display_gallery()

if __name__ == "__main__":
    main()