From 2193dd287ddf87320712d0450974a7ffc28a249e Mon Sep 17 00:00:00 2001
From: bennykok <itechbenny@gmail.com>
Date: Tue, 30 Jan 2024 15:17:25 +0800
Subject: [PATCH] fix: ensure logs are disabled by default

---
 custom_routes.py     |  8 +++--
 prestartup_script.py | 75 +++++++++++++++++++++++++-------------------
 2 files changed, 49 insertions(+), 34 deletions(-)

diff --git a/custom_routes.py b/custom_routes.py
index 99c4c1e..8195a61 100644
--- a/custom_routes.py
+++ b/custom_routes.py
@@ -26,6 +26,7 @@ import threading
 api = None
 api_task = None
 prompt_metadata = {}
+cd_enable_log = os.environ.get('CD_ENABLE_LOG', 'false').lower() == 'true'
 
 def post_prompt(json_data):
     prompt_server = server.PromptServer.instance
@@ -157,7 +158,9 @@ async def websocket_handler(request):
     try:
         # Send initial state to the new client
         await send("status", { 'sid': sid }, sid)
-        await send_first_time_log(sid)
+
+        if cd_enable_log:
+            await send_first_time_log(sid)
             
         async for msg in ws:
             if msg.type == aiohttp.WSMsgType.ERROR:
@@ -480,4 +483,5 @@ def run_in_new_thread(coroutine):
     t.start()
     asyncio.run_coroutine_threadsafe(coroutine, new_loop)
 
-run_in_new_thread(watch_file_changes(log_file_path, send_logs_to_websocket))
+if cd_enable_log:
+    run_in_new_thread(watch_file_changes(log_file_path, send_logs_to_websocket))
diff --git a/prestartup_script.py b/prestartup_script.py
index e2f412a..a944d6d 100644
--- a/prestartup_script.py
+++ b/prestartup_script.py
@@ -7,45 +7,56 @@ import threading
 import logging
 from logging.handlers import RotatingFileHandler
 
-handler = RotatingFileHandler('comfy-deploy.log', maxBytes=500000, backupCount=5)
+# Running with export CD_ENABLE_LOG=true; python main.py
 
-original_stdout = sys.stdout
-original_stderr = sys.stderr
+# Check for 'cd-enable-log' flag in input arguments
+# cd_enable_log = '--cd-enable-log' in sys.argv
+cd_enable_log = os.environ.get('CD_ENABLE_LOG', 'false').lower() == 'true'
 
-class StreamToLogger():
-    def __init__(self, log_level):
-        self.log_level = log_level
+def setup():
+    handler = RotatingFileHandler('comfy-deploy.log', maxBytes=500000, backupCount=5)
 
-    def write(self, buf):
-        if (self.log_level == logging.INFO):
-            original_stdout.write(buf)
-            original_stdout.flush()
-        elif (self.log_level == logging.ERROR):
-            original_stderr.write(buf)
-            original_stderr.flush()
+    original_stdout = sys.stdout
+    original_stderr = sys.stderr
 
-        for line in buf.rstrip().splitlines():
-            handler.handle(
-                logging.LogRecord(
-                    name="comfy-deploy",
-                    level=self.log_level,
-                    pathname="prestartup_script.py",
-                    lineno=1,
-                    msg=line.rstrip(),
-                    args=None,
-                    exc_info=None
+    class StreamToLogger():
+        def __init__(self, log_level):
+            self.log_level = log_level
+
+        def write(self, buf):
+            if (self.log_level == logging.INFO):
+                original_stdout.write(buf)
+                original_stdout.flush()
+            elif (self.log_level == logging.ERROR):
+                original_stderr.write(buf)
+                original_stderr.flush()
+
+            for line in buf.rstrip().splitlines():
+                handler.handle(
+                    logging.LogRecord(
+                        name="comfy-deploy",
+                        level=self.log_level,
+                        pathname="prestartup_script.py",
+                        lineno=1,
+                        msg=line.rstrip(),
+                        args=None,
+                        exc_info=None
+                    )
                 )
-            )
 
-    def flush(self):
-        if (self.log_level == logging.INFO):
-            original_stdout.flush()
-        elif (self.log_level == logging.ERROR):
-            original_stderr.flush()
+        def flush(self):
+            if (self.log_level == logging.INFO):
+                original_stdout.flush()
+            elif (self.log_level == logging.ERROR):
+                original_stderr.flush()
 
-# Redirect stdout and stderr to the logger
-sys.stdout = StreamToLogger(logging.INFO)
-sys.stderr = StreamToLogger(logging.ERROR)
+    # Redirect stdout and stderr to the logger
+    sys.stdout = StreamToLogger(logging.INFO)
+    sys.stderr = StreamToLogger(logging.ERROR)
+
+if cd_enable_log:
+   print("** Comfy Deploy logging enabled")
+   setup() 
 
 try:
     # Get the absolute path of the script's directory