Spaces:
Sleeping
Sleeping
File size: 1,176 Bytes
cf12300 b098573 cf12300 b098573 cf12300 b098573 cf12300 |
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
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()
|