From 48c03a261fcfcebe0d141d6af4f8a07c545c6a49 Mon Sep 17 00:00:00 2001 From: lijunliang Date: Mon, 24 Jun 2024 06:27:20 +0000 Subject: [PATCH 1/4] add_src_ci_controlnet --- tests/sd-webui/test_api.py | 3 +- tests/sd-webui/utils.py | 72 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 71 insertions(+), 4 deletions(-) diff --git a/tests/sd-webui/test_api.py b/tests/sd-webui/test_api.py index 247925bd1..7a4d7cf26 100644 --- a/tests/sd-webui/test_api.py +++ b/tests/sd-webui/test_api.py @@ -22,12 +22,11 @@ get_threshold, ) -THRESHOLD = 0.97 @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) diff --git a/tests/sd-webui/utils.py b/tests/sd-webui/utils.py index 3a1bbaedd..9bd3cddf8 100644 --- a/tests/sd-webui/utils.py +++ b/tests/sd-webui/utils.py @@ -22,9 +22,32 @@ TXT2IMG_TARGET_FOLDER = "/share_nfs/onediff_ci/sd-webui/images/txt2img" SAVED_GRAPH_NAME = "saved_graph" +control_modules = ["segmentation", "canny", "depth", "openpose"] + +sd15_mapping = { + "segmentation": "control_sd15_seg", + "canny": "control_v11p_sd15_canny", + "depth": "control_v11f1p_sd15_depth", + "openpose": "control_sd15_openpose", +} + +control_mapping_imgs = { + "segmentation": "an-source.jpg", + "canny": "sk-b-src.png", + "depth": "sk-b-dep.png", + "openpose": "an-pose.png", +} + +sdxl_mapping = { + "canny": "controlnet-canny-sdxl", +} + os.makedirs(IMG2IMG_TARGET_FOLDER, exist_ok=True) os.makedirs(TXT2IMG_TARGET_FOLDER, exist_ok=True) +def get_model_img(model_name) : + img_path = str(Path(__file__).parent / control_mapping_imgs.get(model_name)) + return encode_file_to_base64(img_path) def get_base_args() -> Dict[str, Any]: return { @@ -42,6 +65,10 @@ def get_base_args() -> Dict[str, Any]: } +def get_model(module, mapping): + return mapping.get(module, "Unknown Module") + + def get_extra_args() -> List[Dict[str, Any]]: quant_args = [ { @@ -57,11 +84,42 @@ def get_extra_args() -> List[Dict[str, Any]]: txt2img_args = [ {}, {"init_images": [get_init_image()]}, + # {"init_images": ["images"]}, + ] + + + controlnet_args = [ + { + "alwayson_scripts": { + "controlnet": { + "args": [ + { + "enabled": (not txt2img_args[0]) and x, + "module": module, + "model": get_model(module, sd15_mapping), + "weight": 1.0, + "image": get_model_img(module), + "resize_mode": "Crop and Resize", + "low_vram": False, + "processor_res": 64, + "threshold_a": 64, + "threshold_b": 64, + "guidance_start": 0.0, + "guidance_end": 1.0, + "control_mode": "Balanced", + "pixel_perfect": False + } + ] + } + } + } + for x in [True, False] + for module in control_modules ] - return [ quant_args, txt2img_args, + controlnet_args ] @@ -73,6 +131,10 @@ def get_all_args() -> Iterable[Dict[str, Any]]: yield args +def get_controlnet_model(data: Dict[str, Any]) -> bool: + return data["alwayson_scripts"]["controlnet"]["args"][0]["model"] + + def is_txt2img(data: Dict[str, Any]) -> bool: return "init_images" not in data @@ -80,6 +142,9 @@ def is_txt2img(data: Dict[str, Any]) -> bool: def is_quant(data: Dict[str, Any]) -> bool: return data["script_args"][0] +def is_controlnet(data: Dict[str, Any]) -> bool: + return data["alwayson_scripts"]["controlnet"]["args"][0]["enabled"] + def encode_file_to_base64(path: str) -> str: with open(path, "rb") as file: @@ -127,7 +192,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" + controlnet_str = "-controlnet" if is_controlnet(data) else "" + controlnet_model_str = get_controlnet_model(data) if is_controlnet(data) else "" + return f"{parent_path}/onediff{quant_str}-{txt2img_str}{controlnet_str}{controlnet_model_str}-w{WIDTH}-h{HEIGHT}-seed-{SEED}-numstep-{NUM_STEPS}.png" def check_and_generate_images(): @@ -144,6 +211,7 @@ def get_data_summary(data: Dict[str, Any]) -> Dict[str, bool]: return { "is_txt2img": is_txt2img(data), "is_quant": is_quant(data), + "is_controlnet":is_controlnet(data) } From 6e2af3c9fc98d535761d2a7f8188eff9cca5cb22 Mon Sep 17 00:00:00 2001 From: lijunliang Date: Mon, 24 Jun 2024 08:14:57 +0000 Subject: [PATCH 2/4] refine_utils_logic --- tests/sd-webui/utils.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/sd-webui/utils.py b/tests/sd-webui/utils.py index 9bd3cddf8..fe0642451 100644 --- a/tests/sd-webui/utils.py +++ b/tests/sd-webui/utils.py @@ -111,7 +111,7 @@ def get_extra_args() -> List[Dict[str, Any]]: } ] } - } + } if x else {} } for x in [True, False] for module in control_modules @@ -132,7 +132,10 @@ def get_all_args() -> Iterable[Dict[str, Any]]: def get_controlnet_model(data: Dict[str, Any]) -> bool: - return data["alwayson_scripts"]["controlnet"]["args"][0]["model"] + try: + return data["alwayson_scripts"]["controlnet"]["args"][0]["model"] + except (KeyError, IndexError): + return False def is_txt2img(data: Dict[str, Any]) -> bool: @@ -143,7 +146,10 @@ def is_quant(data: Dict[str, Any]) -> bool: return data["script_args"][0] def is_controlnet(data: Dict[str, Any]) -> bool: - return data["alwayson_scripts"]["controlnet"]["args"][0]["enabled"] + try: + return data["alwayson_scripts"]["controlnet"]["args"][0]["enabled"] + except (KeyError, IndexError): + return False def encode_file_to_base64(path: str) -> str: @@ -211,7 +217,8 @@ def get_data_summary(data: Dict[str, Any]) -> Dict[str, bool]: return { "is_txt2img": is_txt2img(data), "is_quant": is_quant(data), - "is_controlnet":is_controlnet(data) + "is_controlnet":is_controlnet(data), + "controlnet_model":get_controlnet_model(data) if is_controlnet(data) else "", } From ff48b7d69256e370dcea1d81e2b93e0443109b53 Mon Sep 17 00:00:00 2001 From: lijunliang Date: Tue, 25 Jun 2024 02:01:03 +0000 Subject: [PATCH 3/4] modify path --- tests/sd-webui/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/sd-webui/utils.py b/tests/sd-webui/utils.py index fe0642451..0dc433b3b 100644 --- a/tests/sd-webui/utils.py +++ b/tests/sd-webui/utils.py @@ -46,7 +46,7 @@ os.makedirs(TXT2IMG_TARGET_FOLDER, exist_ok=True) def get_model_img(model_name) : - img_path = str(Path(__file__).parent / control_mapping_imgs.get(model_name)) + img_path = str(Path(__file__).parent / "images" /"txt2img"/control_mapping_imgs.get(model_name)) return encode_file_to_base64(img_path) def get_base_args() -> Dict[str, Any]: @@ -111,8 +111,8 @@ def get_extra_args() -> List[Dict[str, Any]]: } ] } - } if x else {} - } + } + } if x else {} for x in [True, False] for module in control_modules ] From a03872f9765ec8df475cd337e38eeeeb56acd7a1 Mon Sep 17 00:00:00 2001 From: lijunliang Date: Tue, 25 Jun 2024 02:32:59 +0000 Subject: [PATCH 4/4] add notion for function --- tests/sd-webui/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/sd-webui/utils.py b/tests/sd-webui/utils.py index 0dc433b3b..63d164c85 100644 --- a/tests/sd-webui/utils.py +++ b/tests/sd-webui/utils.py @@ -45,7 +45,7 @@ os.makedirs(IMG2IMG_TARGET_FOLDER, exist_ok=True) os.makedirs(TXT2IMG_TARGET_FOLDER, exist_ok=True) -def get_model_img(model_name) : +def get_model_img(model_name : str ) -> str : img_path = str(Path(__file__).parent / "images" /"txt2img"/control_mapping_imgs.get(model_name)) return encode_file_to_base64(img_path) @@ -65,7 +65,7 @@ def get_base_args() -> Dict[str, Any]: } -def get_model(module, mapping): +def get_model(module : str, mapping :Dict[str, str]) -> str: return mapping.get(module, "Unknown Module")