tools/generate.py

import torch
from diffusers import AutoPipelineForText2Image

pipe = None


def load_model():
    global pipe

    if pipe is None:
        pipe = AutoPipelineForText2Image.from_pretrained(
            "stabilityai/sd-turbo",
            torch_dtype=torch.float32
        )
        pipe.to("cpu")

    return pipe


def generate_image(prompt: str):
    model = load_model()

    result = model(
        prompt=prompt,
        num_inference_steps=20,
        guidance_scale=3.0
    )

    return result.images[0]