File size: 833 Bytes
688f56f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
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}
|