From 7a917228d69fd892eb3fc6602a4839ff50c39076 Mon Sep 17 00:00:00 2001
From: Won-Kyu Park <wkpark@gmail.com>
Date: Sat, 26 Oct 2024 23:32:29 +0900
Subject: [PATCH] support loading libtcmalloc_minimal.dll for windows

---
 modules/initialize_util.py | 18 ++++++++++++++++++
 webui.py                   |  4 ++++
 2 files changed, 22 insertions(+)

diff --git a/modules/initialize_util.py b/modules/initialize_util.py
index 79a72cb3aba..c89f929705b 100644
--- a/modules/initialize_util.py
+++ b/modules/initialize_util.py
@@ -213,3 +213,21 @@ def configure_cors_middleware(app):
         cors_options["allow_origin_regex"] = cmd_opts.cors_allow_origins_regex
     app.add_middleware(CORSMiddleware, **cors_options)
 
+
+def preload_malloc():
+    if os.name == 'nt':
+        import ctypes
+        # under windows we can't use LD_PRELOAD method but still we can load the libtcmalloc_minimal.dll using ctypes.cdll
+        # https://github.com/gperftools/gperftools/ (only gperftools's tcmalloc works with windows)
+        # build libtcmalloc_minimal.dll patch module and copy it into the top dir of webui
+
+        dir_path = os.path.dirname(__file__)
+        tcmallocdll = os.path.join(dir_path, "..", "libtcmalloc_minimal.dll")
+
+        if os.path.exists(tcmallocdll):
+            # load libtcmalloc_minimal.dll if it exists.
+            _tcmalloc_dll = ctypes.cdll.LoadLibrary(tcmallocdll)
+
+            return 'tcmalloc'
+
+    return None
diff --git a/webui.py b/webui.py
index 2c417168aa6..c8563cd08d7 100644
--- a/webui.py
+++ b/webui.py
@@ -10,6 +10,10 @@
 startup_timer = timer.startup_timer
 startup_timer.record("launcher")
 
+malloc = initialize_util.preload_malloc()
+if malloc is not None:
+    startup_timer.record(f"{malloc}")
+
 initialize.imports()
 
 initialize.check_versions()