From 122de1af9c23d2e6d80acd3ca0e09da16e3613aa Mon Sep 17 00:00:00 2001
From: Guillermo Rodriguez <guillermo@telos.net>
Date: Tue, 17 Jan 2023 07:06:18 -0300
Subject: [PATCH] Fix redo and change a bit of telegram logic to allow pm

---
 skynet/brain.py             |  3 ++-
 skynet/frontend/telegram.py | 34 ++++++++++++++++++++--------------
 2 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/skynet/brain.py b/skynet/brain.py
index e5fe3fe..e80fc93 100644
--- a/skynet/brain.py
+++ b/skynet/brain.py
@@ -290,7 +290,8 @@ async def open_rpc_service(sock, dgpu_bus, db_pool, tls_whitelist, tls_key):
                         if prompt:
                             req = DiffusionParameters(
                                 prompt=prompt,
-                                **user_config
+                                **user_config,
+                                image=False
                             )
                             rid, img, meta = await dgpu_stream_one_img(req)
                             result = {
diff --git a/skynet/frontend/telegram.py b/skynet/frontend/telegram.py
index b0121b6..e650ae3 100644
--- a/skynet/frontend/telegram.py
+++ b/skynet/frontend/telegram.py
@@ -22,12 +22,18 @@ from . import *
 PREFIX = 'tg'
 
 
-def prepare_metainfo_caption(meta: dict) -> str:
+def prepare_metainfo_caption(tguser, meta: dict) -> str:
     prompt = meta["prompt"]
     if len(prompt) > 256:
         prompt = prompt[:256]
 
-    meta_str = f'prompt: \"{prompt}\"\n'
+    if tguser.username:
+        user = tguser.username
+    else:
+        user = f'{tguser.first_name} id: {tguser.id}'
+
+    meta_str = f'by {user}\n'
+    meta_str += f'prompt: \"{prompt}\"\n'
     meta_str += f'seed: {meta["seed"]}\n'
     meta_str += f'step: {meta["step"]}\n'
     meta_str += f'guidance: {meta["guidance"]}\n'
@@ -87,8 +93,8 @@ async def run_skynet_telegram(
         @bot.message_handler(commands=['txt2img'])
         async def send_txt2img(message):
             chat = message.chat
-            if chat.type != 'group' and chat.id != GROUP_ID:
-                return
+            # if chat.type != 'group' and chat.id != GROUP_ID:
+            #     return
 
             prompt = ' '.join(message.text.split(' ')[1:])
 
@@ -116,8 +122,8 @@ async def run_skynet_telegram(
                 img = Image.open(io.BytesIO(img_raw))
 
                 await bot.send_photo(
-                    message.chat.id,
-                    caption=prepare_metainfo_caption(result['meta']['meta']),
+                    GROUP_ID,
+                    caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']),
                     photo=img,
                     reply_to_message_id=message.id
                 )
@@ -128,8 +134,8 @@ async def run_skynet_telegram(
         @bot.message_handler(func=lambda message: True, content_types=['photo'])
         async def send_img2img(message):
             chat = message.chat
-            if chat.type != 'group' and chat.id != GROUP_ID:
-                return
+            # if chat.type != 'group' and chat.id != GROUP_ID:
+            #     return
 
             if not message.caption.startswith('/img2img'):
                 return
@@ -165,8 +171,8 @@ async def run_skynet_telegram(
                 img = Image.open(io.BytesIO(img_raw))
 
                 await bot.send_photo(
-                    message.chat.id,
-                    caption=prepare_metainfo_caption(result['meta']['meta']),
+                    GROUP_ID,
+                    caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']),
                     photo=img,
                     reply_to_message_id=message.id
                 )
@@ -177,8 +183,8 @@ async def run_skynet_telegram(
         @bot.message_handler(commands=['redo'])
         async def redo_txt2img(message):
             chat = message.chat
-            if chat.type != 'group' and chat.id != GROUP_ID:
-                return
+            # if chat.type != 'group' and chat.id != GROUP_ID:
+            #     return
 
             resp = await _rpc_call(message.from_user.id, 'redo')
 
@@ -194,8 +200,8 @@ async def run_skynet_telegram(
                 img = Image.open(io.BytesIO(img_raw))
 
                 await bot.send_photo(
-                    message.chat.id,
-                    caption=prepare_metainfo_caption(result['meta']['meta']),
+                    GROUP_ID,
+                    caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']),
                     photo=img,
                     reply_to_message_id=message.id
                 )