File size: 596 Bytes
2e50755 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
from diffusers import pipeline
import uuid
import os
class TextToImageTool:
def __init__(self, model_name="stabilityai/stable-diffusion-2"):
self.pipe = pipeline("text-to-image", model=model_name)
self.output_dir = "generated_images"
os.makedirs(self.output_dir, exist_ok=True)
def __call__(self, prompt: str) -> str:
"""Generates image from text and returns local file path"""
result = self.pipe(prompt)
image_path = os.path.join(self.output_dir, f"{uuid.uuid4()}.png")
result.images[0].save(image_path)
return image_path |