From e28261663fc361af4a77af0bd330b82eea8b8990 Mon Sep 17 00:00:00 2001
From: Jeroen Oudshoorn <oudshoorn.jeroen@gmail.com>
Date: Fri, 5 Jan 2024 14:17:51 +0100
Subject: [PATCH] Update

---
 pwnagotchi/ui/web/handler.py | 4 +++-
 pwnagotchi/ui/web/server.py  | 4 ++--
 pyproject.toml               | 4 ++++
 requirements.txt             | 4 ++++
 4 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/pwnagotchi/ui/web/handler.py b/pwnagotchi/ui/web/handler.py
index f8552532..2f0af11b 100644
--- a/pwnagotchi/ui/web/handler.py
+++ b/pwnagotchi/ui/web/handler.py
@@ -8,7 +8,9 @@ from functools import wraps
 
 # https://stackoverflow.com/questions/14888799/disable-console-messages-in-flask-server
 logging.getLogger('werkzeug').setLevel(logging.ERROR)
-os.environ['WERKZEUG_RUN_MAIN'] = 'true'
+
+# Set to false, on bookworm it wouldn't load the flask server for webui
+os.environ['WERKZEUG_RUN_MAIN'] = 'false'
 
 import pwnagotchi
 import pwnagotchi.grid as grid
diff --git a/pwnagotchi/ui/web/server.py b/pwnagotchi/ui/web/server.py
index 5e96ce72..7007798f 100644
--- a/pwnagotchi/ui/web/server.py
+++ b/pwnagotchi/ui/web/server.py
@@ -11,8 +11,8 @@ from pwnagotchi.ui.web.handler import Handler
 # https://stackoverflow.com/questions/14888799/disable-console-messages-in-flask-server
 logging.getLogger('werkzeug').setLevel(logging.ERROR)
 
-# Commented this out, on bookworm it wouldnt load the flask server for webui
-os.environ['WERKZEUG_RUN_MAIN'] = 'true'
+# Set to false, on bookworm it wouldn't load the flask server for webui
+os.environ['WERKZEUG_RUN_MAIN'] = 'false'
 
 
 class Server:
diff --git a/pyproject.toml b/pyproject.toml
index dfe9d751..731a2a91 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,7 +26,11 @@ dependencies = [
     "shimmy; platform_machine!=\"armv6l\"",
     "smbus2",
     "spidev",
+    "stable_baselines3; platform_machine!=\"armv6l\"",
+    "stable_baselines3==1.8.0; platform_machine==\"armv6l\"",
     "toml",
+    "torch==2.0.1",
+    "torchvision==0.15.2",
     "tweepy",
     "websockets"
 ]
diff --git a/requirements.txt b/requirements.txt
index 609840e9..c344481f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -18,6 +18,10 @@ scapy
 shimmy; platform_machine!="armv6l"
 smbus2
 spidev
+stable_baselines3; platform_machine!="armv6l"
+stable_baselines3==1.8.0; platform_machine=="armv6l"
 toml
+torch==2.0.1
+torchvision==0.15.2
 tweepy
 websockets
\ No newline at end of file