from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import io
import base64

# Cargar el modelo preentrenado de Stable Diffusion
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
pipe.to("cuda" if torch.cuda.is_available() else "cpu")  # Usar GPU si está disponible

# Función de inferencia para generar imágenes
def inference(prompt: str, num_images: int = 1):
    # Generar imágenes con el prompt
    generated_images = []
    for _ in range(num_images):
        image = pipe(prompt).images[0]
        generated_images.append(image)

    # Convertir la primera imagen a base64
    buffered = io.BytesIO()
    generated_images[0].save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')

    return {"generated_image": img_str}