import gradio as gr
from PIL import Image
from io import BytesIO
import requests
import sys

# Verificar e instalar dependencias faltantes
try:
    from diffusers import StableDiffusionPipeline
    import torch
except ImportError as e:
    print(f"Error: {e}")
    print("Por favor, instala las dependencias necesarias:")
    print("pip install torch diffusers transformers accelerate")
    sys.exit(1)

def load_models():
    """Carga los modelos con manejo de errores"""
    try:
        # Carga el modelo Stable Diffusion
        pipe_sd = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5", 
            torch_dtype=torch.float32
        )
        pipe_sd = pipe_sd.to("cpu")
        
        # URLs de ejemplo para los otros modelos
        url_dalle_mini = "https://picsum.photos/432/768"  # Imagen placeholder
        url_free_img_gen = "https://picsum.photos/432/768"  # Imagen placeholder
        
        return pipe_sd, url_dalle_mini, url_free_img_gen
        
    except Exception as e:
        print(f"Error cargando modelos: {e}")
        return None, None, None

# Cargar modelos
pipe_sd, url_dalle_mini, url_free_img_gen = load_models()

def generate_image_stable_diffusion(prompt):
    """Genera imagen con Stable Diffusion"""
    if pipe_sd is None:
        # Crear imagen placeholder si el modelo no está disponible
        return Image.new('RGB', (432, 768), color='red')
    
    try:
        image = pipe_sd(
            prompt, 
            height=768, 
            width=432,
            num_inference_steps=20
        ).images[0]
        return image
    except Exception as e:
        print(f"Error generando con Stable Diffusion: {e}")
        return Image.new('RGB', (432, 768), color='red')

def generate_image_dalle_mini(prompt):
    """Genera/simula imagen con DALL·E mini"""
    try:
        response = requests.get(url_dalle_mini, timeout=30)
        response.raise_for_status()
        image = Image.open(BytesIO(response.content))
        return image.resize((432, 768))
    except Exception as e:
        print(f"Error con DALL·E mini: {e}")
        return Image.new('RGB', (432, 768), color='blue')

def generate_image_free_img_gen(prompt):
    """Genera/simula imagen con Free Image Generation"""
    try:
        response = requests.get(url_free_img_gen, timeout=30)
        response.raise_for_status()
        image = Image.open(BytesIO(response.content))
        return image.resize((432, 768))
    except Exception as e:
        print(f"Error con Free Image Generation: {e}")
        return Image.new('RGB', (432, 768), color='green')

def generate_all_images(prompt):
    """Genera imágenes con todos los modelos"""
    if not prompt.strip():
        # Si el prompt está vacío, retornar imágenes placeholder
        placeholder = Image.new('RGB', (432, 768), color='gray')
        return placeholder, placeholder, placeholder
    
    img_sd = generate_image_stable_diffusion(prompt)
    img_dalle = generate_image_dalle_mini(prompt)
    img_free = generate_image_free_img_gen(prompt)
    
    return img_sd, img_dalle, img_free

# Crear la interfaz
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎨 BATUTO-art
    ### Generador de imágenes con múltiples modelos
    """)
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Describe la imagen que quieres generar",
                placeholder="Ej: un paisaje montañoso al atardecer...",
                lines=2
            )
            btn = gr.Button("✨ Generar imágenes", variant="primary")
    
    with gr.Row():
        img_sd = gr.Image(
            label="Stable Diffusion", 
            height=400,
            show_download_button=True
        )
        img_dalle = gr.Image(
            label="DALL·E mini", 
            height=400,
            show_download_button=True
        )
        img_free = gr.Image(
            label="Free Image Generation", 
            height=400,
            show_download_button=True
        )
    
    # Ejemplos
    gr.Examples(
        examples=[
            ["un gato astronauta en el espacio, estilo digital art"],
            ["un castillo medieval en una nube, estilo anime"],
            ["un paisaje futurista con rascacielos brillantes"]
        ],
        inputs=prompt
    )
    
    btn.click(
        fn=generate_all_images, 
        inputs=prompt, 
        outputs=[img_sd, img_dalle, img_free]
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )