From 265c0f841d50eb91c4e8b3e98da7883eb50552b2 Mon Sep 17 00:00:00 2001 From: owl Date: Fri, 14 Mar 2025 21:01:18 +0700 Subject: [PATCH] add list models --- tgbot/shit/prompting.py | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/tgbot/shit/prompting.py b/tgbot/shit/prompting.py index d4b7964..521de1d 100644 --- a/tgbot/shit/prompting.py +++ b/tgbot/shit/prompting.py @@ -91,6 +91,21 @@ async def enqueue_prompt_t2i(args: dict): return await result.json() +async def list_models(): + url = f"https://{IMAGE_GENERATOR_API_HOST}/prompting/txt2img/" + async with aiohttp.ClientSession() as session: + headers = {"Authorization": + "Bearer " + IMAGE_GENERATOR_TOKEN + } + + result = await session.get( + url=url, + headers=headers, + ) + + return await result.text() + + def sdxl_image_dimension(value): try: value = int(value) @@ -136,7 +151,6 @@ def create_prompting_parser(): ) parser.add_argument("--checkpoint", type=str, - choices=["Illustrious-XL-v1.0.safetensors", "PVCStyleModelMovable_illustriousxl10.safetensors"], help="checkpoint to use", required=True) parser.add_argument("--prompt", type=str, @@ -211,7 +225,14 @@ def prompting_parse_args(args: str, parser: CustomArgParser): async def gen_image(update: Update, context: ContextTypes.DEFAULT_TYPE): - query = update.message.text.replace("prompting", "").replace("—", "--").replace("\n", "") + + query = update.message.text + + if query.startswith("list models"): + models = await list_models() + return await context.bot.send_message(chat_id=update.effective_chat.id, text=models) + + query = query.replace("prompting", "").replace("—", "--").replace("\n", "") parser = create_prompting_parser()