import gradio as gr
from deepface import DeepFace
import numpy as np
from PIL import Image
from diffusers import StableDiffusionPipeline
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Detectar si hay GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Cargar modelo de difusión especializado en fantasía (modelo público)
pipe = StableDiffusionPipeline.from_pretrained(
    "nitrosocke/Arcane-Diffusion",  # o el que prefieras
    torch_dtype=torch.float16 if device.type == "cuda" else torch.float32
)
pipe = pipe.to(device)

# Modelo GPT-Neo para enriquecer la descripción
model_name = "EleutherAI/gpt-neo-1.3B"
model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Función para recortar el prompt a 77 tokens
def recortar_prompt(prompt, max_tokens=77):
    tokens = tokenizer(prompt, return_tensors="pt")["input_ids"]
    if len(tokens[0]) > max_tokens:
        tokens = tokens[0][:max_tokens]
        prompt = tokenizer.decode(tokens, skip_special_tokens=True)
    return prompt

# Mapeo de emociones
emociones = {
    "happy": "feliz",
    "sad": "triste",
    "angry": "enojado",
    "surprise": "sorprendido",
    "fear": "miedo",
    "disgust": "asqueroso",
    "neutral": "neutral"
}
# Analizar rostro
# Función para analizar rostro
def analizar_rostro(image):
    try:
        image_np = np.array(image)
        result = DeepFace.analyze(image_np, actions=["age", "gender", "emotion", "race"], enforce_detection=False)

        if isinstance(result, list):
            result = result[0]

        # Obtener los resultados
        gender = result.get('gender', {})
        gender_text = "una mujer" if gender.get('Woman', 0) > gender.get('Man', 0) else "un hombre"

        # Mapeo de razas
        raza = {
            "white": "de piel blanca",
            "black": "de piel negra",
            "asian": "asiático",
            "indian": "de piel india",
            "middle eastern": "de origen medio oriental",
            "latino hispanic": "latino/hispano"
        }
        color_piel = raza.get(result.get('dominant_race', '').lower(), "de piel desconocida")

        # Mapeo de emociones
        emotion = result.get('dominant_emotion', '').lower()
        emocion_text = emociones.get(emotion, "sin emoción destacada")

        # Detectar color de ojos y cabello
        color_ojos = "con ojos marrones" if emotion != "blue" else "con ojos azules"
        color_pelo = "y pelo castaño" if result.get('dominant_race', '').lower() != "blond" else "y pelo rubio"

        # Crear la descripción
        descripcion = (
            f"Una persona {color_piel}, {color_ojos} {color_pelo}, de aproximadamente {result.get('age', 'una edad desconocida')} años, "
            f"que parece estar {emocion_text}. Esta persona es {gender_text}."
        )
        return descripcion
    except Exception as e:
        return f"Error al analizar la imagen: {str(e)}"

# Generar descripción de fantasía
def generar_fantasia(descripcion_literal):
    prompt_base = (
        f"{descripcion_literal} Lleva una capa mágica que cambia de color con la luz, y su ropa está adornada con gemas brillantes. "
        "Está en un mundo fantástico con castillos flotantes, cielos con auroras mágicas y un bosque encantado donde los árboles susurran secretos. "
        "En su camino, encuentra criaturas místicas, fortalezas misteriosas y secretos antiguos que desafían su valentía. "
        "Tiene una misión que podría alterar el destino del reino. "
        "A su lado, una criatura mágica que tiene la habilidad de cambiar de forma."
    )

    inputs = tokenizer(prompt_base, return_tensors="pt").to(device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=150,  # Aumentar el número de tokens generados
        temperature=1.2,  # Aumentar la temperatura para mayor creatividad
        no_repeat_ngram_size=2,
        top_p=0.9,
        top_k=50
    )

    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if generated.startswith(prompt_base):
        generated = generated[len(prompt_base):].strip()

    return prompt_base + " " + generated

# Generar imagen con prompt mejorado y negative prompt
def generar_imagen_fantasia(descripcion_fantasia, progress=gr.Progress()):
    descripcion_fantasia_recortada = recortar_prompt(descripcion_fantasia)
    prompt_visual = (
        f"{descripcion_fantasia_recortada}. Fantasy portrait, glowing magical cloak, enchanted gems, floating castles, magical auroras in the sky, enchanted forest, cinematic lighting, fantasy art style, 8k"
    )

    descripcion_fantasia_recortada = recortar_prompt(prompt_visual)

    # Mostrar por pantalla el prompt que se utiliza para generar la imagen
    print("Prompt para generar la imagen:", descripcion_fantasia_recortada)

    # Empezar la barra de progreso
    progress(0.1)
    image = pipe(
        descripcion_fantasia_recortada,
        negative_prompt="modern clothing, blurry, low quality, photo style, watermark, nsfw, ugly, bad anatomy, disfigured, deformed, extra limbs, close up, out of frame, mutation, mutated, ugly, poorly drawn face, mutation",
        num_inference_steps=50,
        guidance_scale=7.5,
    ).images[0]

    # Completar la barra de progreso
    progress(1.0)
    return image

# Interfaz Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🌟 **Generador de Personaje de Fantasía** 🌟")
    
    with gr.Tabs():
        with gr.TabItem("1. Subir Imagen y Análisis Facial"):
            gr.Markdown("### 1. Sube una imagen para transformarte en un ser mágico")
            gr.Markdown("Puedes probar arrastrando una imagen de [https://thispersondoesnotexist.com/](https://thispersondoesnotexist.com/).")
            image_input = gr.Image(type="pil", label="Imagen de entrada")
            descripcion_output = gr.Textbox(label="Descripción literal automática (con DeepFace)", interactive=True)

            image_input.change(analizar_rostro, inputs=image_input, outputs=descripcion_output)

        with gr.TabItem("2. Generar Descripción de Fantasía"):
            gr.Markdown("### 2. Descripción de Fantasía (con EleutherAI/gpt-neo-1.3B)")
            boton_fantasia = gr.Button("✨ Generar descripción de fantasía")
            descripcion_fantasia_output = gr.Textbox(label="Descripción de fantasía", interactive=False)

            boton_fantasia.click(generar_fantasia, inputs=descripcion_output, outputs=descripcion_fantasia_output)

        with gr.TabItem("3. Generar Imagen"):
            gr.Markdown("### 3. Generar imagen de fantasía")
            boton_imagen = gr.Button("🎨 Generar imagen de fantasía (con nitrosocke/Arcane-Diffusion)")
            output_image = gr.Image(label="Imagen de fantasía generada")

            boton_imagen.click(generar_imagen_fantasia, inputs=descripcion_fantasia_output, outputs=output_image)
demo.launch()