import torch
import gradio as gr
from diffusers import StableDiffusionPipeline

# Cargar el modelo Stable Diffusion
model_id = "stabilityai/stable-diffusion-2-1"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe.to("cuda")  # Si estás usando GPU

# Función para generar imágenes
def generate_image(prompt, style, width, height, guidance_scale, steps, seed):
    # Configurar el generador de semillas
    generator = torch.manual_seed(seed) if seed is not None else None

    # Ajustar el estilo en función del input del usuario
    if style == "Realismo":
        prompt += ", realistic, high detail, photorealistic"
    elif style == "Arte digital":
        prompt += ", digital art, sharp details, vibrant colors"
    elif style == "Ciencia Ficción":
        prompt += ", sci-fi, futuristic, high-tech"
    elif style == "Surrealismo":
        prompt += ", surreal, dream-like, abstract"

    # Generar la imagen con Stable Diffusion
    image = pipe(prompt, height=height, width=width, guidance_scale=guidance_scale, 
                 num_inference_steps=steps, generator=generator).images[0]
    return image

# Configuración de la interfaz de Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Generador de Imágenes estilo MidJourney")
    prompt = gr.Textbox(label="Descripción (Prompt)", placeholder="Describe la imagen...")
    style = gr.Dropdown(choices=["Realismo", "Arte digital", "Ciencia Ficción", "Surrealismo"], label="Estilo", value="Realismo")
    width = gr.Slider(512, 1024, step=64, label="Ancho", value=768)
    height = gr.Slider(512, 1024, step=64, label="Altura", value=768)
    guidance_scale = gr.Slider(5, 20, step=0.5, label="Nivel de guía del prompt", value=7.5)
    steps = gr.Slider(10, 100, step=5, label="Número de pasos de inferencia", value=50)
    seed = gr.Number(label="Semilla (dejar vacío para aleatorio)", value=None, optional=True)
    
    # Botón para generar la imagen
    btn = gr.Button("Generar imagen")
    
    # Salida de la imagen generada
    output = gr.Image(label="Imagen generada")

    # Vincular los inputs y outputs
    btn.click(generate_image, inputs=[prompt, style, width, height, guidance_scale, steps, seed], outputs=output)

# Lanzar la aplicación Gradio
demo.launch()