From a7258a11401f3aaa998c521fd25e54f1eb094683 Mon Sep 17 00:00:00 2001 From: Luna Date: Sun, 4 Aug 2024 17:37:17 -0300 Subject: [PATCH] let extra interrogator models be configured in --- main.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/main.py b/main.py index 103acdf..fda6f09 100644 --- a/main.py +++ b/main.py @@ -57,7 +57,8 @@ class Interrogator: "select output_tag_string, time_taken from interrogated_posts where md5 = ? and model_name = ?", (md5_hash, self.model_id), ) - assert len(rows) == 1 # run 'fight' mode, there's missing posts + if not rows: + raise AssertionError("run fight mode first, there are missing posts..") tag_string, time_taken = rows[0][0], rows[0][1] return InterrogatorPost(self._process(tag_string.split()), time_taken) @@ -129,6 +130,7 @@ class Config: sd_webui_address: str dd_address: str dd_model_name: str + sd_webui_extras: Dict[str, str] sd_webui_models: List[str] = field(default_factory=lambda: list(DEFAULTS)) @property @@ -138,6 +140,10 @@ class Config: SDInterrogator(sd_interrogator, self.sd_webui_address) for sd_interrogator in self.sd_webui_models ] + + [ + SDInterrogator(sd_interrogator, url) + for sd_interrogator, url in self.sd_webui_extras.items() + ] + [DDInterrogator(self.dd_model_name, self.dd_address)] + [ControlInterrogator("control", None)] )