From 16df97d731930b8f36cd0682d092eb7f8a09c4a4 Mon Sep 17 00:00:00 2001
From: Guillermo Rodriguez <guillermo@telos.net>
Date: Sun, 8 Oct 2023 10:02:15 -0300
Subject: [PATCH] Improve tg frontend ipfs results gathering parallelism

---
 skynet/dgpu/network.py               |  1 -
 skynet/frontend/telegram/__init__.py | 76 +++++++++++-----------------
 skynet/frontend/telegram/handlers.py |  1 -
 3 files changed, 30 insertions(+), 48 deletions(-)

diff --git a/skynet/dgpu/network.py b/skynet/dgpu/network.py
index 121a6e1..a029f82 100644
--- a/skynet/dgpu/network.py
+++ b/skynet/dgpu/network.py
@@ -264,7 +264,6 @@ class SkynetGPUConnector:
         ipfs_link_legacy = ipfs_link + '/image.png'
 
         async with trio.open_nursery() as n:
-
             async def get_and_set_results(link: str):
                 res = await get_ipfs_file(link, timeout=1)
                 logging.info(f'got response from {link}')
diff --git a/skynet/frontend/telegram/__init__.py b/skynet/frontend/telegram/__init__.py
index 05b8763..ea26b69 100644
--- a/skynet/frontend/telegram/__init__.py
+++ b/skynet/frontend/telegram/__init__.py
@@ -242,66 +242,50 @@ class SkynetTelegramFrontend:
         ipfs_link_legacy = ipfs_link + '/image.png'
 
         async def get_and_set_results(link: str):
-            results[link] = await get_ipfs_file(link)
+            res = await get_ipfs_file(link)
+            logging.info(f'got response from {link}')
+            if not res or res.status_code != 200:
+                logging.warning(f'couldn\'t get ipfs binary data at {link}!')
 
-        def get_image_from_resp(resp):
-            png_img = resp.raw
-            with Image.open(io.BytesIO(resp.raw)) as image:
-                w, h = image.size
+            else:
+                try:
+                    with Image.open(io.BytesIO(res.raw)) as image:
+                        w, h = image.size
 
-                if w > TG_MAX_WIDTH or h > TG_MAX_HEIGHT:
-                    logging.warning(f'result is of size {image.size}')
-                    image.thumbnail((TG_MAX_WIDTH, TG_MAX_HEIGHT))
-                    tmp_buf = io.BytesIO()
-                    image.save(tmp_buf, format='PNG')
-                    png_img = tmp_buf.getvalue()
+                        if w > TG_MAX_WIDTH or h > TG_MAX_HEIGHT:
+                            logging.warning(f'result is of size {image.size}')
+                            image.thumbnail((TG_MAX_WIDTH, TG_MAX_HEIGHT))
 
-            return png_img
+                        tmp_buf = io.BytesIO()
+                        image.save(tmp_buf, format='PNG')
+                        png_img = tmp_buf.getvalue()
+
+                        results[link] = png_img
+
+                except UnidentifiedImageError:
+                    logging.warning(f'couldn\'t get ipfs binary data at {link}!')
 
         tasks = [
             get_and_set_results(ipfs_link),
             get_and_set_results(ipfs_link_legacy)
         ]
         await asyncio.gather(*tasks)
-
         png_img = None
+        if ipfs_link_legacy in results:
+            png_img = results[ipfs_link_legacy]
 
-        resp = results[ipfs_link_legacy]
-        if not resp or resp.status_code != 200:
-            logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link_legacy}!')
+        if ipfs_link in results:
+            png_img = results[ipfs_link]
 
-        else:
-            try:
-                png_img = get_image_from_resp(resp)
-
-            except UnidentifiedImageError:
-                logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link_legacy}!')
 
         if not png_img:
-            resp = results[ipfs_link]
-            if not resp or resp.status_code != 200:
-                logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link}!')
-                await self.update_status_message(
-                    status_msg,
-                    caption,
-                    reply_markup=build_redo_menu(),
-                    parse_mode='HTML'
-                )
-                return True
-
-            else:
-                try:
-                    png_img = get_image_from_resp(resp)
-
-                except UnidentifiedImageError:
-                    logging.error(f'couldn\'t get ipfs hosted image at {ipfs_link}!')
-                    await self.update_status_message(
-                        status_msg,
-                        caption,
-                        reply_markup=build_redo_menu(),
-                        parse_mode='HTML'
-                    )
-                    return True
+            await self.update_status_message(
+                status_msg,
+                caption,
+                reply_markup=build_redo_menu(),
+                parse_mode='HTML'
+            )
+            return True
 
         logging.info(f'success! sending generated image')
         await self.bot.delete_message(
diff --git a/skynet/frontend/telegram/handlers.py b/skynet/frontend/telegram/handlers.py
index b2019a7..ef0e31f 100644
--- a/skynet/frontend/telegram/handlers.py
+++ b/skynet/frontend/telegram/handlers.py
@@ -149,7 +149,6 @@ def create_handler_context(frontend: 'SkynetTelegramFrontend'):
         user_config = {**user_row}
         del user_config['id']
 
-        breakpoint()
         if user_config['autoconf']:
             user_config = perform_auto_conf(user_config)