| 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} | |