New CI test for sd-webui Refiner

.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

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):

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),
     }
 
 

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