You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
51 lines
1.1 KiB
51 lines
1.1 KiB
# -*- coding: utf-8 -*-
|
|
import io
|
|
from os.path import abspath, dirname
|
|
|
|
import PIL
|
|
import pytest
|
|
import requests
|
|
|
|
from rudalle import get_tokenizer, get_rudalle_model, get_vae, get_realesrgan
|
|
|
|
|
|
TEST_ROOT = dirname(abspath(__file__))
|
|
|
|
|
|
@pytest.fixture(scope='module')
|
|
def realesrgan():
|
|
realesrgan = get_realesrgan('x2', device='cpu')
|
|
yield realesrgan
|
|
|
|
|
|
@pytest.fixture(scope='module')
|
|
def vae():
|
|
vae = get_vae(pretrained=False)
|
|
yield vae
|
|
|
|
|
|
@pytest.fixture(scope='module')
|
|
def dwt_vae():
|
|
vae = get_vae(pretrained=False, dwt=True)
|
|
yield vae
|
|
|
|
|
|
@pytest.fixture(scope='module')
|
|
def yttm_tokenizer():
|
|
tokenizer = get_tokenizer()
|
|
yield tokenizer
|
|
|
|
|
|
@pytest.fixture(scope='module')
|
|
def sample_image():
|
|
url = 'https://cdn.kqed.org/wp-content/uploads/sites/12/2013/12/rudolph.png'
|
|
resp = requests.get(url)
|
|
resp.raise_for_status()
|
|
image = PIL.Image.open(io.BytesIO(resp.content))
|
|
yield image
|
|
|
|
|
|
@pytest.fixture(scope='module')
|
|
def small_dalle():
|
|
model = get_rudalle_model('small', pretrained=False, fp16=False, device='cpu')
|
|
return model
|