From b3c18e0de6d398388dae1e18e822a2cff3cb0f33 Mon Sep 17 00:00:00 2001
From: owl <lonelyowl@hitler.rocks>
Date: Fri, 14 Mar 2025 16:27:57 +0700
Subject: [PATCH] add loras

---
 .gitignore              |  1 +
 tgbot/config.py         |  6 +++++-
 tgbot/shit/prompting.py | 45 +++++++++++++++++++++++++++++++++++++++--
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 5162b6b..d744e4b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -164,3 +164,4 @@ cython_debug/
 sandbox/
 
 config.ovpn
+/tesing_things.py
diff --git a/tgbot/config.py b/tgbot/config.py
index d926b82..e73d6ce 100644
--- a/tgbot/config.py
+++ b/tgbot/config.py
@@ -13,4 +13,8 @@ IMAGE_GENERATOR_TOKEN: str = os.getenv("IMAGE_GENERATOR_TOKEN")
 IMAGE_GENERATOR_API_HOST: str = os.getenv("IMAGE_GENERATOR_API_HOST")
 CDN_URL: str = os.getenv("CDN_URL")
 
-PROMPTING_USERS: list = os.getenv("PROMPTING_USERS").split(" ")
+try:
+    PROMPTING_USERS: list | None = os.getenv("PROMPTING_USERS").split(" ")
+except AttributeError:
+    PROMPTING_USERS = None
+
diff --git a/tgbot/shit/prompting.py b/tgbot/shit/prompting.py
index f30b62c..bd7e38a 100644
--- a/tgbot/shit/prompting.py
+++ b/tgbot/shit/prompting.py
@@ -99,6 +99,22 @@ def sdxl_image_dimension(value):
         raise argparse.ArgumentTypeError(f"{value} is not an integer")
 
 
+def lora(value):
+    lora_ = value.split(":")
+
+    if len(lora_) != 2:
+        raise argparse.ArgumentTypeError(f"can't parse lora '{value}', "
+                                         f"use --lora lora_name.safetensors:lora_strength")
+
+    try:
+        lora_[1] = float(lora_[1])
+
+    except ValueError:
+        argparse.ArgumentTypeError(f"lora strength {lora_[1]} is not a float")
+
+    return lora_
+
+
 def denoise(value):
     try:
         value = float(value)
@@ -156,13 +172,39 @@ def create_prompting_parser():
                         help="noise scheduler, default is normal", default="normal")
 
     parser.add_argument("--denoise", type=denoise,
-                        help="denoise level, default is 1, default=8", default=1.0)
+                        help="denoise level, default is 1", default=1.0)
+
+    parser.add_argument("--stop_at_clip_layer", type=denoise,
+                        help="clip skip, default is -1, ", default=-1)
+
+    parser.add_argument("--lora", type=lora, action="append",
+                        help="add lora to prompt, e.g. --lora lora_name.safetensors:0.8. "
+                             "Multiple loras might be added")
 
     return parser
 
 
 def prompting_parse_args(args: str, parser: CustomArgParser):
     args = parser.parse_args(shlex.split(args)).__dict__
+
+    loras = []
+
+    if args['lora']:
+
+        for lora_ in args['lora']:
+
+            loras.append(
+                {
+                    "name": lora_[0],
+                    "strength_model": lora_[1],
+                    "strength_clip": lora_[1]
+                }
+            )
+
+        del args["lora"]
+
+    args["loras"] = loras
+
     return args
 
 
@@ -210,4 +252,3 @@ async def gen_image(update: Update, context: ContextTypes.DEFAULT_TYPE):
     else:
         return await context.bot.send_message(chat_id=update.effective_chat.id,
                                               text="Some weird incomprehensible error occurred :<")
-