diff --git a/scripts/iib/db/update_image_data.py b/scripts/iib/db/update_image_data.py index c9b5a88..720a47f 100644 --- a/scripts/iib/db/update_image_data.py +++ b/scripts/iib/db/update_image_data.py @@ -155,10 +155,14 @@ def build_single_img_idx(conn, file_path, is_rebuild, safe_save_img_tag): meta = parsed_params.meta lora = parsed_params.extra.get("lora", []) lyco = parsed_params.extra.get("lyco", []) + if "final_width" in meta and "final_height" in meta: + size_str = str(meta["final_width"]) + " × " + str(meta["final_height"]) + else: + size_str = "Unknown Size" pos = parsed_params.pos_prompt size_tag = Tag.get_or_create( conn, - str(meta.get("Size-1", 0)) + " * " + str(meta.get("Size-2", 0)), + size_str, type="size", ) safe_save_img_tag(ImageTag(img.id, size_tag.id)) diff --git a/scripts/iib/parsers/comfyui.py b/scripts/iib/parsers/comfyui.py index 9a4d206..68c0b8b 100644 --- a/scripts/iib/parsers/comfyui.py +++ b/scripts/iib/parsers/comfyui.py @@ -22,6 +22,7 @@ class ComfyUIParser: params = None if not clz.test(img, file_path): raise Exception("The input image does not match the current parser.") + width, height = img.size try: if is_img_created_by_comfyui_with_webui_gen_info(img): info = read_sd_webui_gen_info_from_image(img, file_path) @@ -30,14 +31,20 @@ class ComfyUIParser: else: params = get_comfyui_exif_data(img) info = comfyui_exif_data_to_str(params) - except Exception as e: - logger.error('parse comfyui image failed. prompt:') - logger.error(img.info.get('prompt')) - return ImageGenerationInfo() + except Exception as e: + logger.error("parse comfyui image failed. prompt:") + logger.error(img.info.get("prompt")) + return ImageGenerationInfo( + params=ImageGenerationParams( + meta={"final_width": width, "final_height": height} + ) + ) return ImageGenerationInfo( info, ImageGenerationParams( - meta=params["meta"], pos_prompt=params["pos_prompt"], extra=params + meta=params["meta"] | {"final_width": width, "final_height": height}, + pos_prompt=params["pos_prompt"], + extra=params, ), ) diff --git a/scripts/iib/parsers/fooocus.py b/scripts/iib/parsers/fooocus.py index 7c1df19..d1b2ff6 100644 --- a/scripts/iib/parsers/fooocus.py +++ b/scripts/iib/parsers/fooocus.py @@ -60,7 +60,11 @@ class FooocusParser: return ImageGenerationInfo( info, ImageGenerationParams( - meta=params["meta"], + meta=params["meta"] + | { + "final_width": img.size[0], + "final_height": img.size[1], + }, pos_prompt=parse_generation_parameters(info)["pos_prompt"], extra={ "lora": unique_by(lora_list, lambda x: x["name"].lower()) diff --git a/scripts/iib/parsers/invoke_ai.py b/scripts/iib/parsers/invoke_ai.py index 0765ca0..534b5ad 100644 --- a/scripts/iib/parsers/invoke_ai.py +++ b/scripts/iib/parsers/invoke_ai.py @@ -50,7 +50,8 @@ class InvokeAIParser: return ImageGenerationInfo( info, ImageGenerationParams( - meta=meta_obj, + meta=meta_obj + | {"final_width": img.size[0], "final_height": img.size[1]}, pos_prompt=parse_generation_parameters(info)["pos_prompt"], # extra=params ), diff --git a/scripts/iib/parsers/novelai.py b/scripts/iib/parsers/novelai.py index 680e1cd..b675617 100644 --- a/scripts/iib/parsers/novelai.py +++ b/scripts/iib/parsers/novelai.py @@ -32,7 +32,9 @@ class NovelAIParser: return ImageGenerationInfo( info, ImageGenerationParams( - meta=params["meta"], pos_prompt=params["pos_prompt"] + meta=params["meta"] + | {"final_width": img.size[0], "final_height": img.size[1]}, + pos_prompt=params["pos_prompt"], ), ) diff --git a/scripts/iib/parsers/sd_webui.py b/scripts/iib/parsers/sd_webui.py index 88c11a5..aedb8a0 100644 --- a/scripts/iib/parsers/sd_webui.py +++ b/scripts/iib/parsers/sd_webui.py @@ -16,14 +16,21 @@ class SdWebUIParser: if not clz.test(img, file_path): raise Exception("The input image does not match the current parser.") info = read_sd_webui_gen_info_from_image(img, file_path) + width, height = img.size if not info: - return ImageGenerationInfo() + return ImageGenerationInfo( + params=ImageGenerationParams( + meta={"final_width": width, "final_height": height} + ) + ) info += ", Source Identifier: Stable Diffusion web UI" params = parse_generation_parameters(info) return ImageGenerationInfo( info, ImageGenerationParams( - meta=params["meta"], pos_prompt=params["pos_prompt"], extra=params + meta=params["meta"] | {"final_width": width, "final_height": height}, + pos_prompt=params["pos_prompt"], + extra=params, ), ) diff --git a/scripts/iib/parsers/sd_webui_stealth.py b/scripts/iib/parsers/sd_webui_stealth.py index 668b716..fdba44e 100644 --- a/scripts/iib/parsers/sd_webui_stealth.py +++ b/scripts/iib/parsers/sd_webui_stealth.py @@ -145,7 +145,10 @@ class SdWebUIStealthParser: return ImageGenerationInfo( info, ImageGenerationParams( - meta=params["meta"], pos_prompt=params["pos_prompt"], extra=params + meta=params["meta"] + | {"final_width": img.size[0], "final_height": img.size[1]}, + pos_prompt=params["pos_prompt"], + extra=params, ), ) diff --git a/scripts/iib/parsers/stable_swarm_ui.py b/scripts/iib/parsers/stable_swarm_ui.py index 8e41e04..eabe369 100644 --- a/scripts/iib/parsers/stable_swarm_ui.py +++ b/scripts/iib/parsers/stable_swarm_ui.py @@ -48,7 +48,10 @@ class StableSwarmUIParser: return ImageGenerationInfo( info, ImageGenerationParams( - meta=params["meta"], pos_prompt=params["pos_prompt"], extra=params + meta=params["meta"] + | {"final_width": img.size[0], "final_height": img.size[1]}, + pos_prompt=params["pos_prompt"], + extra=params, ), )