model.py

from typing import Any


def get_pipeline():
    import torch
    from diffusers import AutoencoderTiny, AutoPipelineForImage2Image

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    pipe = AutoPipelineForImage2Image.from_pretrained(
        "SimianLuo/LCM_Dreamshaper_v7",
        use_safetensors=True,
    )
    pipe.vae = AutoencoderTiny.from_pretrained(
        "madebyollin/taesd",
        torch_dtype=torch.float16,
        use_safetensors=True,
    )
    pipe = pipe.to(device, dtype=torch.float16)
    pipe.unet.to(memory_format=torch.channels_last)
    return pipe


def get_test_pipeline():
    from PIL import Image
    from dataclasses import dataclass
    import random
    import time

    @dataclass
    class Images:
        images: list[Image.Image]

    class Pipeline:
        def __call__(self, *args: Any, **kwds: Any) -> Any:
            time.sleep(0.5)
            r = random.randint(0, 255)
            g = random.randint(0, 255)
            b = random.randint(0, 255)

            return Images(images=[Image.new("RGB", (512, 512), color=(r, g, b))])

    return Pipeline()