From 7564dd5131ebef2b62b34ea88215f05d29bcdd60 Mon Sep 17 00:00:00 2001
From: Manuel Schmid <9307310+mashb1t@users.noreply.github.com>
Date: Sat, 23 Mar 2024 12:49:20 +0100
Subject: [PATCH] fix: load image number from preset (#2611)

* fix: add default_image_number to preset handling

* fix: use minimum image number of preset and config to prevent UI overflow
---
 modules/config.py      |  1 +
 modules/meta_parser.py | 14 +++++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/modules/config.py b/modules/config.py
index 76ffd348..6c02ca13 100644
--- a/modules/config.py
+++ b/modules/config.py
@@ -485,6 +485,7 @@ possible_preset_keys = {
     "default_scheduler": "scheduler",
     "default_overwrite_step": "steps",
     "default_performance": "performance",
+    "default_image_number": "image_number",
     "default_prompt": "prompt",
     "default_prompt_negative": "negative_prompt",
     "default_styles": "styles",
diff --git a/modules/meta_parser.py b/modules/meta_parser.py
index 10bc6896..8cd21cbc 100644
--- a/modules/meta_parser.py
+++ b/modules/meta_parser.py
@@ -27,8 +27,9 @@ def load_parameter_button_click(raw_metadata: dict | str, is_generating: bool):
         loaded_parameter_dict = json.loads(raw_metadata)
     assert isinstance(loaded_parameter_dict, dict)
 
-    results = [len(loaded_parameter_dict) > 0, 1]
+    results = [len(loaded_parameter_dict) > 0]
 
+    get_image_number('image_number', 'Image Number', loaded_parameter_dict, results)
     get_str('prompt', 'Prompt', loaded_parameter_dict, results)
     get_str('negative_prompt', 'Negative Prompt', loaded_parameter_dict, results)
     get_list('styles', 'Styles', loaded_parameter_dict, results)
@@ -92,6 +93,17 @@ def get_float(key: str, fallback: str | None, source_dict: dict, results: list,
         results.append(gr.update())
 
 
+def get_image_number(key: str, fallback: str | None, source_dict: dict, results: list, default=None):
+    try:
+        h = source_dict.get(key, source_dict.get(fallback, default))
+        assert h is not None
+        h = int(h)
+        h = min(h, modules.config.default_max_image_number)
+        results.append(h)
+    except:
+        results.append(1)
+
+
 def get_steps(key: str, fallback: str | None, source_dict: dict, results: list, default=None):
     try:
         h = source_dict.get(key, source_dict.get(fallback, default))