Skip to content

Commit

Permalink
Ensure we don't crash when QR text too large
Browse files Browse the repository at this point in the history
  • Loading branch information
@db0
db0 committed May 20, 2024
1 parent 3858214 commit 5f38cf7
Show file tree
Hide file tree
Showing 2 changed files with 132 additions and 4 deletions.
47 changes: 45 additions & 2 deletions hordelib/horde.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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},
Expand Down Expand Up @@ -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"]
Expand All @@ -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):
Expand Down
89 changes: 87 additions & 2 deletions tests/test_horde_inference_qr_code.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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],
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 5f38cf7

Please sign in to comment.