Discord bot: make model entry optional for text2image

VoltaML · Jun 19, 2024 · 726ab2d · 726ab2d
1 parent 4f444c5
commit 726ab2d
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/bot/helper.py b/bot/helper.py
@@ -14,7 +14,11 @@
 
 async def find_closest_model(model: str):
     """Find the closest model to the one provided"""
+
     models, _ = await shared_bot.models.cached_loaded_models()
+
+    if model == "":
+        return models[0]
     res = difflib.get_close_matches(model, models, n=1, cutoff=0.1)
 
     if not res:

diff --git a/bot/models.py b/bot/models.py
@@ -100,8 +100,8 @@ async def load_model(self, ctx: Context, model: str) -> None:
                 response = await response.json()
 
         if status == 200:
-            await message.edit(content=f"{response['message']}: {model}")
             await shared_bot.models.refresh_loaded_models()
+            await message.edit(content=f"{response['message']}: {model}")
         else:
             await message.edit(content=f"Error: **{response.get('detail')}**")
 
@@ -121,8 +121,8 @@ async def unload_model(self, ctx: Context, model: str) -> None:
                 response = await response.json()
 
         if status == 200:
-            await message.edit(content=f"{response['message']}: {model}")
             await shared_bot.models.refresh_loaded_models()
+            await message.edit(content=f"{response['message']}: {model}")
         else:
             await message.edit(content=f"Error: {status}")
 

diff --git a/bot/txt2img.py b/bot/txt2img.py
@@ -34,7 +34,7 @@ async def dream_unsupported(
         self,
         ctx: Context,
         prompt: str,
-        model: str,
+        model: str = "",
         negative_prompt: str = "",
         guidance_scale: float = config.default_cfg,
         steps: int = config.default_steps,