import gradio as gr
from PIL import Image
import torch
from diffusers import StableDiffusionImg2ImgPipeline, StableDiffusionInpaintPipeline
import spaces
from huggingface_hub import login
from dotenv import load_dotenv
import os

# Cargar las variables de entorno desde el archivo .env
load_dotenv()

# Obtener el token de Hugging Face del archivo .env
huggingface_token = os.getenv('HUGGINGFACE_TOKEN')

# Autenticación en Hugging Face
def huggingface_login(token):
    try:
        login(token)
        print("Autenticación exitosa con Hugging Face.")
    except Exception as e:
        print(f"Error al autenticar con Hugging Face: {e}")

# Llamar a la función de autenticación
huggingface_login(huggingface_token)

# Inicializar los pipelines globalmente
device = "cuda" if torch.cuda.is_available() else "cpu"

# Modelo para img2img
model_id_img2img = "CompVis/stable-diffusion-v1-4"
pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_img2img).to(device)

# Modelo para edición de imágenes (inpainting)
model_id_inpainting = "stabilityai/stable-diffusion-inpainting"
pipe_inpainting = StableDiffusionInpaintPipeline.from_pretrained(model_id_inpainting).to(device)

@spaces.GPU(duration=0)
def cartoonize_image(image, prompt):
    try:
        # Preprocesar imagen
        image = image.convert("RGB")
        image = image.resize((512, 512))  # Ajusta el tamaño según sea necesario

        # Generar imagen con Stable Diffusion Img2Img
        with torch.no_grad():
            result = pipe_img2img(prompt=prompt, init_image=image, strength=0.75).images[0]
        
        return result
    except Exception as e:
        return f"Error procesando la imagen: {e}"

@spaces.GPU(duration=0)
def edit_image(image, mask, prompt):
    try:
        # Preprocesar imagen
        image = image.convert("RGB")
        mask = mask.convert("L")  # Convertir la máscara a escala de grises

        # Generar la imagen editada con Stable Diffusion Inpainting
        with torch.no_grad():
            result = pipe_inpainting(prompt=prompt, image=image, mask_image=mask).images[0]
        
        return result
    except Exception as e:
        return f"Error editando la imagen: {e}"

def main():
    with gr.Blocks() as demo:
        gr.Markdown("## Conversión de Imagen a Caricatura Estilo Pixar y Edición de Imágenes")
        
        with gr.Tab("Caricaturizar"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type="pil", label="Sube tu imagen")
                    prompt = gr.Textbox(label="Prompt", value="A Pixar-style cartoon character")
                    
                    with gr.Row():
                        submit_button = gr.Button("Submit")
                    
                    output_image = gr.Image(type="pil", label="Imagen Caricaturizada")
            
            def process_image(image, prompt):
                return cartoonize_image(image, prompt)
            
            submit_button.click(
                fn=process_image,
                inputs=[input_image, prompt],
                outputs=output_image
            )
        
        with gr.Tab("Editar Imagen"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type="pil", label="Sube tu imagen a editar")
                    mask_image = gr.Image(type="pil", label="Sube una máscara (blanco para la parte a editar)")
                    prompt_edit = gr.Textbox(label="Describe los cambios", value="Cambia el color del objeto a rojo")
                    
                    with gr.Row():
                        submit_edit_button = gr.Button("Submit Edición")
                    
                    output_image_edit = gr.Image(type="pil", label="Imagen Editada")
            
            def process_edit_image(image, mask, prompt_edit):
                return edit_image(image, mask, prompt_edit)
            
            submit_edit_button.click(
                fn=process_edit_image,
                inputs=[input_image, mask_image, prompt_edit],
                outputs=output_image_edit
            )
        
    demo.launch()

if __name__ == "__main__":
    main()