File size: 1,201 Bytes
6ad0b38
 
 
 
 
 
 
 
 
e4b8c77
6ad0b38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
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")
    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
    print("下载文件")
    pipe = AutoPipelineForImage2Image.from_pretrained(
        "SimianLuo/LCM_Dreamshaper_v7",
        use_safetensors=True,
    )
    pipe.vae = AutoencoderTiny.from_pretrained(
        "madebyollin/taesd",
        torch_dtype=torch_dtype,
        use_safetensors=True,
    )
    pipe = pipe.to(device, dtype=torch_dtype)
    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:
            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()