increase compability with jupyter

SquareFactory · Jul 28, 2022 · 436c4ce · 436c4ce
1 parent 6028900
commit 436c4ce
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/i2_client/client.py b/i2_client/client.py
@@ -13,6 +13,7 @@
 
 import archipel_utils as utils
 import msgpack
+import nest_asyncio
 import websockets
 
 log = logging.getLogger(__name__)
@@ -204,4 +205,12 @@ async def _inference(self, inputs):
             await self.__aexit__(exc_type=None, exc_value=None, traceback=None)
             return outputs
 
+        try:
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                # when executed in jupyter notebook or something
+                nest_asyncio.apply(loop)
+        except RuntimeError:
+            pass
+
         return asyncio.run(_inference(self, inputs))
diff --git a/setup.py b/setup.py
@@ -25,6 +25,7 @@
         "imutils>=0.5.4",
         "msgpack>=1.0",
         "numpy>=1.19",
+        "nest-asyncio>=1.5",
         "rich>=10.13",
         "websockets>=8.1",
         "opencv-python==4.6.0.66",