src/model.py

import torch
from diffusers import DiffusionPipeline

pipe = None

def load_model():
    print('Loading model...')
    global pipe
    print('Model loaded!')

    if pipe is None:
        print('Loading model... 2')
        pipe = DiffusionPipeline.from_pretrained(
            "stablediffusionapi/deliberate-v2",
            torch_dtype=torch.float16
        ).to("cuda")
        print('Model loaded! 2')

    return pipe


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

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