import gradio as gr
from diffusers import DiffusionPipeline
from PIL import Image
import numpy as np
from tqdm.auto import tqdm

# Descargar el modelo
pipeline = DiffusionPipeline.from_pretrained("sd-community/sdxl-flash")

# Función para generar imágenes a partir de texto con barra de progreso
def generate_image_from_text(prompt, num_steps=4):
    # Inicializar la barra de progreso
    progress_bar = tqdm(total=num_steps, desc="Generando imagen", unit="step")
    
    # Generar imágenes
    for step in range(num_steps):
        images = pipeline(prompt, num_inference_steps=step+1)
        progress_bar.update(1)
    
    progress_bar.close()
    return images[0]

# Función para generar imágenes a partir de una imagen con barra de progreso
def generate_image_from_image(image, prompt, num_steps=4):
    # Convertir la imagen a formato aceptable por el modelo (numpy array)
    image = np.array(image)
    
    # Inicializar la barra de progreso
    progress_bar = tqdm(total=num_steps, desc="Generando imagen", unit="step")
    
    # Generar imágenes
    for step in range(num_steps):
        images = pipeline(prompt, image=image, num_inference_steps=step+1)
        progress_bar.update(1)
    
    progress_bar.close()
    return images[0]

# Definir la interfaz de Gradio
text_to_image_interface = gr.Interface(
    fn=generate_image_from_text,
    inputs=["text", "number"],
    outputs="image",
    title="Texto a Imagen",
    description="Genera una imagen a partir de un texto con un número de pasos."
)

image_to_image_interface = gr.Interface(
    fn=generate_image_from_image,
    inputs=["image", "text", "number"],
    outputs="image",
    title="Imagen a Imagen",
    description="Genera una imagen a partir de una imagen y un texto con un número de pasos."
)

# Ejecutar las interfaces de Gradio
gr.Interface([text_to_image_interface, image_to_image_interface], layout="tabs").launch()