diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index ab85e1bfb..31bc3a189 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -18,6 +18,8 @@ env: WEBUI_SRC_DIR: stable-diffusion-webui WEBUI_DEPENDENCIES_SUBDIR: repos SDXL_BASE: /share_nfs/hf_models/sd_xl_base_1.0.safetensors + REFINER_BASE: /share_nfs/onediff_ci/sd-webui/safetensors/sd_xl_refiner_1.0.safetensors + AWPAINTING_BASE: /share_nfs/onediff_ci/sd-webui/safetensors/AWPainting_v1.2.safetensors UNET_INT8: /share_nfs/hf_models/unet_int8 CONTROL_LORA_OPENPOSEXL2_RANK256: /share_nfs/hf_models/controlnet/control-lora-openposeXL2-rank256.safetensors SELENIUM_CONTAINER_NAME: selenium-test @@ -382,7 +384,10 @@ jobs: - if: matrix.test-suite == 'webui' run: docker exec ${{ env.CONTAINER_NAME }} python3 -m pytest -v -s tests/sd-webui/test_api.py - + env: + REFINER_BASE: ${{ env.REFINER_BASE }} + AWPAINTING_BASE: ${{ env.AWPAINTING_BASE }} + SDXL_BASE: ${{ env.SDXL_BASE }} - name: Show WebUI Log if: matrix.test-suite == 'webui' run: docker exec ${{ env.CONTAINER_NAME }} cat /app/${{ env.WEBUI_SRC_DIR }}/onediff_webui.log diff --git a/tests/sd-webui/test_api.py b/tests/sd-webui/test_api.py index f4cee6601..f5a9907a7 100644 --- a/tests/sd-webui/test_api.py +++ b/tests/sd-webui/test_api.py @@ -20,13 +20,14 @@ post_request_and_check, dump_image, get_threshold, + ) @pytest.fixture(scope="session", autouse=True) def change_model(): option_payload = { - "sd_model_checkpoint": "checkpoints/AWPainting_v1.2.safetensors", + "sd_model_checkpoint": "AWPainting_v1.2.safetensors", } post_request_and_check(f"{WEBUI_SERVER_URL}/{OPTIONS_API_ENDPOINT}", option_payload) @@ -64,6 +65,7 @@ def test_image_ssim(base_url, data): url = f"{base_url}/{endpoint}" generated_image = get_image_array_from_response(post_request_and_check(url, data)) target_image_path = get_target_image_filename(data) + directory, filename = os.path.split(target_image_path) target_image = np.array(Image.open(target_image_path)) ssim_value = cal_ssim(generated_image, target_image) if ssim_value < get_threshold(data): diff --git a/tests/sd-webui/utils.py b/tests/sd-webui/utils.py index 3a1bbaedd..772ed1044 100644 --- a/tests/sd-webui/utils.py +++ b/tests/sd-webui/utils.py @@ -29,6 +29,7 @@ def get_base_args() -> Dict[str, Any]: return { "prompt": "1girl", + "sd_model_checkpoint": "AWPainting_v1.2.safetensors", "negative_prompt": "", "seed": SEED, "steps": NUM_STEPS, @@ -59,11 +60,17 @@ def get_extra_args() -> List[Dict[str, Any]]: {"init_images": [get_init_image()]}, ] - return [ - quant_args, - txt2img_args, + refiner_args = [ + {}, + { + "refiner": True, + "refiner_checkpoint": "sd_xl_refiner_1.0", + "refiner_switch_at": 0.8, + }, ] + return [quant_args, txt2img_args, refiner_args] + def get_all_args() -> Iterable[Dict[str, Any]]: for extra_args in product(*get_extra_args()): @@ -77,6 +84,10 @@ def is_txt2img(data: Dict[str, Any]) -> bool: return "init_images" not in data +def is_refiner(data: Dict[str, Any]) -> bool: + return "refiner" in data + + def is_quant(data: Dict[str, Any]) -> bool: return data["script_args"][0] @@ -127,7 +138,9 @@ def get_target_image_filename(data: Dict[str, Any]) -> str: txt2img_str = "txt2img" if is_txt2img(data) else "img2img" quant_str = "-quant" if is_quant(data) else "" - return f"{parent_path}/onediff{quant_str}-{txt2img_str}-w{WIDTH}-h{HEIGHT}-seed-{SEED}-numstep-{NUM_STEPS}.png" + refiner_str = "-refiner" if is_refiner(data) else "" + + return f"{parent_path}/onediff{quant_str}{refiner_str}-{txt2img_str}-w{WIDTH}-h{HEIGHT}-seed-{SEED}-numstep-{NUM_STEPS}.png" def check_and_generate_images(): @@ -144,6 +157,7 @@ def get_data_summary(data: Dict[str, Any]) -> Dict[str, bool]: return { "is_txt2img": is_txt2img(data), "is_quant": is_quant(data), + "is_refiner": is_refiner(data), } diff --git a/tests/webui-docker-compose.yml b/tests/webui-docker-compose.yml index 6f70c713a..426eb87e3 100644 --- a/tests/webui-docker-compose.yml +++ b/tests/webui-docker-compose.yml @@ -50,6 +50,9 @@ services: - $HOME/test-container-cache-${CONTAINER_NAME}/dot-cache:/root/.cache - /share_nfs:/share_nfs:ro - /share_nfs/onediff_ci/sd-webui/images:/share_nfs/onediff_ci/sd-webui/images:rw + - ${REFINER_BASE}:/app/${WEBUI_SRC_DIR}/models/Stable-diffusion/sd_xl_refiner_1.0.safetensors:ro + - ${AWPAINTING_BASE}:/app/${WEBUI_SRC_DIR}/models/Stable-diffusion/AWPainting_v1.2.safetensors:ro + - ${SDXL_BASE}:/app/${WEBUI_SRC_DIR}/models/Stable-diffusion/sd_xl_base_1.0.safetensors:ro - ${PWD}/${WEBUI_SRC_DIR}:/app/${WEBUI_SRC_DIR} - ${PWD}/onediff_sd_webui_extensions:/app/${WEBUI_SRC_DIR}/extensions/onediff_sd_webui_extensions - $PWD:/src/onediff