src/models.py

import torch
from diffusers import DiffusionPipeline

pipe = None

def load_model():
    global pipe

    if pipe is None:
        pipe = DiffusionPipeline.from_pretrained(
            "stablediffusionapi/deliberate-v2",
            torch_dtype=torch.float16
        ).to("cuda")

    return pipe


def generate(prompt: str):
    pipe = load_model()

    image = pipe(prompt).images[0]
    return image