1
0
Fork 0

chore(api): automate release tests for blend mode

This commit is contained in:
Sean Sube 2023-02-19 15:18:22 -06:00
parent 62a6f14cb5
commit c5eae688e3
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 43 additions and 14 deletions

View File

@ -1,12 +1,11 @@
import sys import sys
import traceback import traceback
from collections import Counter
from io import BytesIO from io import BytesIO
from logging import getLogger from logging import getLogger
from logging.config import dictConfig from logging.config import dictConfig
from os import environ, path from os import environ, path
from time import sleep from time import sleep
from typing import Optional from typing import List, Optional, Union
import cv2 import cv2
import numpy as np import numpy as np
@ -45,7 +44,7 @@ class TestCase:
query: str, query: str,
max_attempts: int = 20, max_attempts: int = 20,
mse_threshold: float = 0.001, mse_threshold: float = 0.001,
source: Image.Image = None, source: Union[Image.Image, List[Image.Image]] = None,
mask: Image.Image = None, mask: Image.Image = None,
) -> None: ) -> None:
self.name = name self.name = name
@ -133,7 +132,7 @@ TEST_DATA = [
), ),
source="txt2img-sd-v1-5-512-muffin", source="txt2img-sd-v1-5-512-muffin",
mask="mask-black", mask="mask-black",
mse_threshold=0.025, mse_threshold=0.010,
), ),
TestCase( TestCase(
"outpaint-horizontal-512", "outpaint-horizontal-512",
@ -143,7 +142,12 @@ TEST_DATA = [
), ),
source="txt2img-sd-v1-5-512-muffin", source="txt2img-sd-v1-5-512-muffin",
mask="mask-black", mask="mask-black",
mse_threshold=0.025, mse_threshold=0.010,
),
TestCase(
"upscale-resrgan-x2-1024-muffin",
"upscale?prompt=a+giant+pumpkin&seed=0&scheduler=ddim&upscaling=upscaling-real-esrgan-x2-plus&scale=2&outscale=2",
source="txt2img-sd-v1-5-512-muffin",
), ),
TestCase( TestCase(
"upscale-resrgan-x4-2048-muffin", "upscale-resrgan-x4-2048-muffin",
@ -151,9 +155,22 @@ TEST_DATA = [
source="txt2img-sd-v1-5-512-muffin", source="txt2img-sd-v1-5-512-muffin",
), ),
TestCase( TestCase(
"upscale-resrgan-x2-1024-muffin", "blend-512-muffin-black",
"upscale?prompt=a+giant+pumpkin&seed=0&scheduler=ddim&upscaling=upscaling-real-esrgan-x2-plus&scale=2&outscale=2", "upscale?prompt=a+giant+pumpkin&seed=0&scheduler=ddim&upscaling=upscaling-real-esrgan-x2-plus&scale=2&outscale=2",
source="txt2img-sd-v1-5-512-muffin", mask="mask-black",
source=[
"txt2img-sd-v1-5-512-muffin",
"txt2img-sd-v2-1-512-muffin",
],
),
TestCase(
"blend-512-muffin-white",
"upscale?prompt=a+giant+pumpkin&seed=0&scheduler=ddim&upscaling=upscaling-real-esrgan-x2-plus&scale=2&outscale=2",
mask="mask-white",
source=[
"txt2img-sd-v2-1-512-muffin",
"txt2img-sd-v1-5-512-muffin",
],
), ),
] ]
@ -161,13 +178,25 @@ TEST_DATA = [
def generate_image(root: str, test: TestCase) -> Optional[str]: def generate_image(root: str, test: TestCase) -> Optional[str]:
files = {} files = {}
if test.source is not None: if test.source is not None:
logger.debug("loading test source: %s", test.source) if isinstance(test.source, list):
source_path = test_path(path.join("test-refs", f"{test.source}.png")) for i in range(len(test.source)):
source_image = Image.open(source_path) source = test.source[i]
source_bytes = BytesIO() logger.debug("loading test source %s: %s", i, source)
source_image.save(source_bytes, "png") source_path = test_path(path.join("test-refs", f"{source}.png"))
source_bytes.seek(0) source_image = Image.open(source_path)
files["source"] = source_bytes source_bytes = BytesIO()
source_image.save(source_bytes, "png")
source_bytes.seek(0)
files[f"source:{i}"] = source_bytes
else:
logger.debug("loading test source: %s", test.source)
source_path = test_path(path.join("test-refs", f"{test.source}.png"))
source_image = Image.open(source_path)
source_bytes = BytesIO()
source_image.save(source_bytes, "png")
source_bytes.seek(0)
files["source"] = source_bytes
if test.mask is not None: if test.mask is not None:
logger.debug("loading test mask: %s", test.mask) logger.debug("loading test mask: %s", test.mask)