let extra interrogator models be configured in

This commit is contained in:
Luna 2024-08-04 17:37:17 -03:00
parent 967813b429
commit a7258a1140

View file

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