From 296498c1cd3c2a5c20941cfda968afb68a3a88a6 Mon Sep 17 00:00:00 2001
From: Conner Harkness <caharkness@gmail.com>
Date: Thu, 20 Mar 2025 14:23:40 -0600
Subject: [PATCH] Added conversion to JPEG for compressed Discord attachments +
 extra options for aspect ratio control

---
 comfyui-discord.sh |  1 +
 lib/events.py      | 10 ++++++++++
 lib/images.py      | 10 ++++++++++
 lib/parser.py      | 39 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 60 insertions(+)
 create mode 100644 lib/images.py

diff --git a/comfyui-discord.sh b/comfyui-discord.sh
index f6e11b4..4969c89 100755
--- a/comfyui-discord.sh
+++ b/comfyui-discord.sh
@@ -19,6 +19,7 @@ if [[ -f ./venv/Scripts/activate ]]; then source ./venv/Scripts/activate; fi
 if [[ $INSTALLATION -eq 1 ]]
 then
     pip install discord.py
+    pip install pillow
 fi
 
 python -u bot.py
diff --git a/lib/events.py b/lib/events.py
index b10300f..ff291f4 100644
--- a/lib/events.py
+++ b/lib/events.py
@@ -13,6 +13,7 @@ from lib.helpers import *
 from lib.settings import *
 from lib.parser import *
 from lib.comfyui import *
+from lib.images import *
 
 async def on_message_or_reaction(client, obj):
     global RUN_TIME
@@ -220,6 +221,15 @@ async def on_message_or_reaction(client, obj):
             for a in attachments_buffer:
                 bytes_generated     = bytes_generated + len(a)
                 images_generated    = images_generated + 1
+                megabytes_uploading = bytes_generated / (1024 * 1024)
+
+                if megabytes_uploading > 4.5:
+                    log.write(f"PNG generated is {megabytes_uploading} MB and will be delivered as a JPEG", log.colors.fg.lightred)
+                    jpeg_bytes      = png_to_jpg(a)
+                    discord_file    = discord.File(io.BytesIO(jpeg_bytes), filename="file.jpg")
+                    discord_files.append(discord_file)
+                    continue
+
                 discord_file        = discord.File(io.BytesIO(a), filename="file.png")
                 discord_files.append(discord_file)
 
diff --git a/lib/images.py b/lib/images.py
new file mode 100644
index 0000000..c2295ee
--- /dev/null
+++ b/lib/images.py
@@ -0,0 +1,10 @@
+import io
+from PIL import Image
+
+def png_to_jpg(input_png_bytes):
+    with io.BytesIO(input_png_bytes) as png:
+        img = Image.open(png)
+        
+        with io.BytesIO() as jpg:
+            img.convert("RGB").save(jpg, format="JPEG", quality=50)
+            return jpg.getvalue()
diff --git a/lib/parser.py b/lib/parser.py
index eefde7c..0aba68e 100644
--- a/lib/parser.py
+++ b/lib/parser.py
@@ -37,10 +37,49 @@ def get_prompt_parameters(user_message, settings):
 
     params["positive"] = mutable_string.strip()
 
+
+    # Allow the user to set the square size by SSSp:
+    if params["width"] == params["height"]:
+        match = re.search(r"([1-9][0-9]{2,3})p", options)
+        if match is not None:
+            s = int(match.group(1))
+            params["width"] = s
+            params["height"] = s
+
+
+    # Allow user to specify ratio:
+    match = re.search(r"([0-9]{1,2})(x|\:)([0-9]{1,2})", options)
+    if match is not None:
+        # Use the smallest dimension for calculations:
+        size = params["width"] if params["width"] <= params["height"] else params["height"]
+
+        r_num = int(match.group(1))
+        r_den = int(match.group(3))
+
+        # Largest dimension not to exceed the size:
+        if r_num > r_den:
+            ratio = 1.0 * r_num / r_den
+            params["width"] = size
+            params["height"] = size * (1 / ratio)
+        elif r_num < r_den:
+            ratio = 1.0 * r_num / r_den
+            params["width"] = size * ratio
+            params["height"] = size
+
+
+    # Allow user to explicitly set the size by WWWxHHHp:
+    match = re.search(r"([0-9]{1,})(x)([0-9]{1,})p", options)
+    if match is not None:
+        output["width"] = int(match.group(1))
+        output["height"] = int(match.group(3))
+
+
+    # Allow all other explicit parameter overrides:
     overwrite_re = r"([A-Za-z_]{1,})=?([0-9\.]{1,})"
     for (k, v) in re.findall(overwrite_re, options):
         params[k] = float(v)
 
+
     # Process limitations, e.g. max_width and max_height:
     params_clone = copy.deepcopy(params)