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)] )