rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Move to async generation for Stable Horde

Browse Source
This commit is contained in:
somebody
2022-12-10 22:31:52 -06:00
parent 2410734738
commit a480a59b5f
1 changed files with 34 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
aiserver.py
View File

@@ -9695,16 +9695,43 @@ def text2img_horde(prompt: str) -> Optional[Image.Image]:
}
}
cluster_headers = {'apikey': koboldai_vars.sh_apikey if koboldai_vars.sh_apikey != '' else "0000000000",}
cluster_headers = {"apikey": koboldai_vars.sh_apikey or "0000000000"}
id_req = requests.post("https://stablehorde.net/api/v2/generate/async", json=final_submit_dict, headers=cluster_headers)
if not id_req.ok:
logger.error(f"HTTP {id_req.status_code}, expected OK-ish")
logger.error(id_req.text)
logger.error(f"Response headers: {id_req.headers}")
raise HordeException("Image seeding failed. See console for more details.")
logger.debug(final_submit_dict)
submit_req = requests.post('https://stablehorde.net/api/v2/generate/sync', json=final_submit_dict, headers=cluster_headers)
image_id = id_req.json()["id"]
if not submit_req.ok:
logger.error(submit_req.text)
return
while True:
poll_req = requests.get(f"https://stablehorde.net/api/v2/generate/check/{image_id}")
if not poll_req.ok:
logger.error(f"HTTP {poll_req.status_code}, expected OK-ish")
logger.error(poll_req.text)
logger.error(f"Response headers: {poll_req.headers}")
raise HordeException("Image polling failed. See console for more details.")
poll_j = poll_req.json()
if poll_j["finished"] > 0:
break
# This should always exist but if it doesn't 2 seems like a safe bet.
sleepy_time = int(poll_req.headers.get("retry-after", 2))
time.sleep(sleepy_time)
# Done generating, we can now fetch it.
gen_req = requests.get(f"https://stablehorde.net/api/v2/generate/status/{image_id}")
if not gen_req.ok:
logger.error(f"HTTP {gen_req.status_code}, expected OK-ish")
logger.error(gen_req.text)
logger.error(f"Response headers: {gen_req.headers}")
raise HordeException("Image fetching failed. See console for more details.")
results = gen_req.json()
results = submit_req.json()
if len(results["generations"]) > 1:
logger.warning(f"Got too many generations, discarding extras. Got {len(results['generations'])}, expected 1.")