60 lines
1.4 KiB
Python
60 lines
1.4 KiB
Python
from multiprocessing import Queue, Value
|
|
from os import path
|
|
from typing import List
|
|
from unittest import skipUnless
|
|
|
|
from onnx_web.params import DeviceParams, ImageParams, Size
|
|
from onnx_web.worker.context import WorkerContext
|
|
|
|
|
|
def test_needs_models(models: List[str]):
|
|
return skipUnless(
|
|
all([path.exists(model) for model in models]), "model does not exist"
|
|
)
|
|
|
|
|
|
def test_needs_onnx_models(models: List[str]):
|
|
return skipUnless(
|
|
all([path.exists(f"{model}.onnx") for model in models]), "model does not exist"
|
|
)
|
|
|
|
|
|
def test_device() -> DeviceParams:
|
|
return DeviceParams("cpu", "CPUExecutionProvider")
|
|
|
|
|
|
def test_size() -> Size:
|
|
return Size(64, 64)
|
|
|
|
|
|
def test_params() -> ImageParams:
|
|
return ImageParams("test", "txt2img", "ddim", "test", 5.0, 25, 0)
|
|
|
|
|
|
def test_worker() -> WorkerContext:
|
|
cancel = Value("L", 0)
|
|
logs = Queue()
|
|
pending = Queue()
|
|
progress = Queue()
|
|
active = Value("L", 0)
|
|
idle = Value("L", 0)
|
|
|
|
return WorkerContext(
|
|
"test",
|
|
test_device(),
|
|
cancel,
|
|
logs,
|
|
pending,
|
|
progress,
|
|
active,
|
|
idle,
|
|
3,
|
|
0.1,
|
|
)
|
|
|
|
|
|
TEST_MODEL_CORRECTION_CODEFORMER = "../models/.cache/correction-codeformer.pth"
|
|
TEST_MODEL_DIFFUSION_SD15 = "../models/stable-diffusion-onnx-v1-5"
|
|
TEST_MODEL_DIFFUSION_SD15_INPAINT = "../models/stable-diffusion-onnx-v1-inpainting"
|
|
TEST_MODEL_UPSCALING_SWINIR = "../models/.cache/upscaling-swinir.pth"
|