mirror of https://github.com/skygpu/skynet.git
50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
import pytest
|
|
from PIL import Image
|
|
|
|
from skynet.types import ModelMode, BodyV0Params
|
|
from skynet.dgpu.compute import maybe_load_model, compute_one
|
|
|
|
|
|
@pytest.mark.parametrize('mode,model', [
|
|
(ModelMode.DIFFUSE, 'skygpu/mocker'),
|
|
(ModelMode.TXT2IMG, 'skygpu/mocker'),
|
|
(ModelMode.IMG2IMG, 'skygpu/mocker'),
|
|
(ModelMode.INPAINT, 'skygpu/mocker'),
|
|
(ModelMode.UPSCALE, 'skygpu/mocker-upscale'),
|
|
])
|
|
async def test_pipeline_mocker(inject_mockers, mode, model):
|
|
# always insert at least two inputs to make all modes pass
|
|
inputs = [
|
|
Image.new('RGB', (1, 1), color='green')
|
|
for i in range(2)
|
|
]
|
|
params = BodyV0Params(
|
|
prompt="Kronos God Realistic 4k",
|
|
model=model,
|
|
step=4,
|
|
width=1024,
|
|
height=1024,
|
|
seed=168402949,
|
|
guidance="7.5"
|
|
)
|
|
|
|
with maybe_load_model(model, mode) as model:
|
|
compute_one(model, 0, mode, params, inputs)
|
|
|
|
# disable for now (cuda)
|
|
# async def test_pipeline():
|
|
# model = 'stabilityai/stable-diffusion-xl-base-1.0'
|
|
# mode = 'txt2img'
|
|
# params = BodyV0Params(
|
|
# prompt="Kronos God Realistic 4k",
|
|
# model=model,
|
|
# step=21,
|
|
# width=1024,
|
|
# height=1024,
|
|
# seed=168402949,
|
|
# guidance="7.5"
|
|
# )
|
|
#
|
|
# with maybe_load_model(model, mode) as model:
|
|
# compute_one(model, 0, mode, params)
|