From 5f38cf72caa6dbea7dcea3f15d7334901c1e121b Mon Sep 17 00:00:00 2001 From: db0 Date: Mon, 20 May 2024 17:58:59 +0200 Subject: [PATCH] Ensure we don't crash when QR text too large --- hordelib/horde.py | 47 +++++++++++++- tests/test_horde_inference_qr_code.py | 89 ++++++++++++++++++++++++++- 2 files changed, 132 insertions(+), 4 deletions(-) diff --git a/hordelib/horde.py b/hordelib/horde.py index f8d848bd..8535ee5a 100644 --- a/hordelib/horde.py +++ b/hordelib/horde.py @@ -23,6 +23,7 @@ from hordelib.comfy_horde import Comfy_Horde from hordelib.consts import MODEL_CATEGORY_NAMES +from hordelib.nodes.comfy_qr.qr_nodes import QRByModuleSizeSplitFunctionPatterns from hordelib.shared_model_manager import SharedModelManager from hordelib.utils.dynamicprompt import DynamicPromptParser from hordelib.utils.image_utils import ImageUtils @@ -145,7 +146,7 @@ class HordeLib: "sampler_name": {"datatype": str, "values": list(SAMPLERS_MAP.keys()), "default": "k_euler"}, "cfg_scale": {"datatype": float, "min": 1, "max": 100, "default": 8.0}, "denoising_strength": {"datatype": float, "min": 0.01, "max": 1.0, "default": 1.0}, - "control_strength": {"datatype": float, "min": 0.01, "max": 1.0, "default": 1.0}, + "control_strength": {"datatype": float, "min": 0.01, "max": 3.0, "default": 1.0}, "seed": {"datatype": int, "default": random.randint(0, sys.maxsize)}, "width": {"datatype": int, "min": 64, "max": 8192, "default": 512, "divisible": 64}, "height": {"datatype": int, "min": 64, "max": 8192, "default": 512, "divisible": 64}, @@ -942,6 +943,9 @@ def _final_pipeline_adjustments(self, payload, pipeline_data) -> tuple[dict, lis # If we have a qr code request, we check for extra texts such as the generation url if payload.get("workflow") == "qr_code": + original_width = pipeline_params.get("empty_latent_image.width", 512) + original_height = pipeline_params.get("empty_latent_image.height", 512) + pipeline_params["qr_code_split.max_image_size"] = max(original_width, original_height) for text in payload.get("extra_texts"): if text["reference"] == "qr_text": pipeline_params["qr_code_split.text"] = text["text"] @@ -958,17 +962,56 @@ def _final_pipeline_adjustments(self, payload, pipeline_data) -> tuple[dict, lis pipeline_params["qr_code_split.module_drawer"] = text["text"].capitalize() if text["reference"] == "function_layer_prompt": pipeline_params["function_layer_prompt.text"] = text["text"] + if text["reference"] == "x_offset" and text["text"].isdigit(): + pipeline_params["qr_flattened_composite.x"] = int(text["text"]) + if text["reference"] == "y_offset" and text["text"].isdigit(): + pipeline_params["qr_flattened_composite.y"] = int(text["text"]) + if text["reference"] == "qr_border" and text["text"].isdigit(): + pipeline_params["qr_code_split.border"] = int(text["text"]) if not pipeline_params.get("qr_code_split.protocol"): pipeline_params["qr_code_split.protocol"] = "None" if not pipeline_params.get("function_layer_prompt.text"): pipeline_params["function_layer_prompt.text"] = payload["prompt"] + try: + test_qr = QRByModuleSizeSplitFunctionPatterns() + _, _, _, _, _, qr_size = test_qr.generate_qr( + protocol=pipeline_params.get("qr_code_split.protocol"), + text=pipeline_params["qr_code_split.text"], + module_size=16, + max_image_size=pipeline_params["qr_code_split.max_image_size"], + fill_hexcolor="#FFFFFF", + back_hexcolor="#000000", + error_correction="High", + border=1, + module_drawer="Square", + ) + except RuntimeError as err: + logger.error(err) + pipeline_params["qr_code_split.text"] = "This QR Code is too large for this image" + test_qr = QRByModuleSizeSplitFunctionPatterns() + qr_size = 624 + if not pipeline_params.get("qr_flattened_composite.x"): + x_offset = int((original_width / 2) - qr_size / 2) + # I don't know why but through trial and error I've discovered that the QR codes + # are more legible when they're placed in an offset which is a multiple of 64 + x_offset = x_offset - (x_offset % 64) if x_offset % 64 != 0 else x_offset + pipeline_params["qr_flattened_composite.x"] = x_offset + if not pipeline_params.get("qr_flattened_composite.y"): + y_offset = int((original_height / 2) - qr_size / 2) + y_offset = y_offset - (y_offset % 64) if y_offset % 64 != 0 else y_offset + pipeline_params["qr_flattened_composite.y"] = y_offset + pipeline_params["module_layer_composite.x"] = pipeline_params["qr_flattened_composite.x"] + pipeline_params["module_layer_composite.y"] = pipeline_params["qr_flattened_composite.y"] + pipeline_params["function_layer_composite.x"] = pipeline_params["qr_flattened_composite.x"] + pipeline_params["function_layer_composite.y"] = pipeline_params["qr_flattened_composite.y"] + pipeline_params["mask_composite.x"] = pipeline_params["qr_flattened_composite.x"] + pipeline_params["mask_composite.y"] = pipeline_params["qr_flattened_composite.y"] if SharedModelManager.manager.compvis: model_details = SharedModelManager.manager.compvis.get_model_reference_info(payload["model_name"]) if model_details and model_details.get("baseline") == "stable diffusion 1": pipeline_params["controlnet_qr_model_loader.control_net_name"] = ( "control_v1p_sd15_qrcode_monster_v2.safetensors" ) - return pipeline_params, faults def _get_appropriate_pipeline(self, params): diff --git a/tests/test_horde_inference_qr_code.py b/tests/test_horde_inference_qr_code.py index 5befbea1..b57cf9b4 100644 --- a/tests/test_horde_inference_qr_code.py +++ b/tests/test_horde_inference_qr_code.py @@ -44,6 +44,14 @@ def test_qr_code_inference( "text": "https://aihorde.net", "reference": "qr_text", }, + { + "text": "256", + "reference": "x_offset", + }, + { + "text": "256", + "reference": "y_offset", + }, ], } assert hordelib_instance is not None @@ -83,6 +91,7 @@ def test_qr_code_inference_xl( "beautiful, illustration, incredible detail, 8k, abstract" "###worst quality, bad lighting, deformed, ugly, low contrast" ), + "control_strength": 1.2, "ddim_steps": 25, "n_iter": 1, "model": sdxl_refined_model_name, @@ -109,6 +118,73 @@ def test_qr_code_inference_xl( # pil_image, # ) + def test_qr_code_inference_too_large_text( + self, + shared_model_manager: type[SharedModelManager], + hordelib_instance: HordeLib, + sdxl_refined_model_name: str, + ): + data = { + "sampler_name": "k_euler", + "cfg_scale": 7.5, + "denoising_strength": 1.0, + "seed": 1111, + "height": 1024, + "width": 1024, + "karras": False, + "tiling": False, + "hires_fix": False, + "clip_skip": 1, + "prompt": "The Scream, 1893 by Edvard Munch", + "ddim_steps": 25, + "n_iter": 1, + "model": sdxl_refined_model_name, + "workflow": "qr_code", + "control_strength": 1.4, + "extra_texts": [ + { + "text": """ +Sint eos pariatur architecto repellat nihil sed distinctio aut. +Aut quae modi libero. Est ea nobis blanditiis in ut quam. +Blanditiis expedita minus tenetur dolorum. +Nisi recusandae blanditiis quia mollitia. Voluptatibus omnis non eos. +Aut voluptatum ut aspernatur minima omnis. +Quia officia est nisi exercitationem sint. +Adipisci mollitia sunt dignissimos fugiat sint magnam et sit. Cumque ipsum ullam et molestiae. +Consequatur saepe occaecati amet odio molestias odio est doloremque. +Sapiente reprehenderit qui adipisci officia delectus quas totam. +Maxime commodi rerum quod voluptas dolor ducimus. Non quibusdam ut assumenda ipsum voluptatem sit a necessitatibus. +Provident est culpa delectus nemo. Dolorem velit assumenda labore ut. +Voluptatum corporis modi id dolores necessitatibus voluptatibus at voluptate. +Rerum sed incidunt commodi quo quo. Sit sint accusantium modi eligendi molestiae maxime. + """, + "reference": "qr_text", + }, + { + "text": "48", + "reference": "x_offset", + }, + { + "text": "48", + "reference": "y_offset", + }, + ], + } + assert hordelib_instance is not None + assert shared_model_manager.manager.controlnet is not None + + pil_image = hordelib_instance.basic_inference_single_image(data).image + assert pil_image is not None + assert isinstance(pil_image, Image.Image) + + img_filename = "qr_code_too_long_text.png" + pil_image.save(f"images/{img_filename}", quality=100) + + # assert check_single_inference_image_similarity( + # f"images_expected/{img_filename}", + # pil_image, + # ) + def test_qr_code_control_strength( self, shared_model_manager: type[SharedModelManager], @@ -158,7 +234,7 @@ def test_qr_code_control_strength( # pil_image, # ) - def test_qr_code_control_size( + def test_qr_code_control_non_square( self, shared_model_manager: type[SharedModelManager], hordelib_instance: HordeLib, @@ -168,7 +244,7 @@ def test_qr_code_control_size( "sampler_name": "k_euler", "cfg_scale": 7.5, "denoising_strength": 1.0, - "seed": 1312, + "seed": 11312, "height": 1280, "width": 768, "karras": False, @@ -180,6 +256,7 @@ def test_qr_code_control_size( "beautiful, illustration, incredible detail, 8k, abstract" "###worst quality, bad lighting, deformed, ugly, low contrast" ), + "control_strength": 1, "ddim_steps": 25, "n_iter": 1, "model": sdxl_refined_model_name, @@ -249,6 +326,14 @@ def test_qr_code_control_qr_texts( "text": "Lucid Creations, ethereal brain cells", "reference": "function_layer_prompt", }, + { + "text": "48", + "reference": "x_offset", + }, + { + "text": "48", + "reference": "y_offset", + }, ], } assert hordelib_instance is not None