From 0f3bc95913b2fd4ae9cc9c185705bcad472265e2 Mon Sep 17 00:00:00 2001
From: henryruhs <info@henryruhs.com>
Date: Tue, 30 May 2023 13:47:44 +0200
Subject: [PATCH] max-memory support for Windows

---
 run.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/run.py b/run.py
index 5d65b7d..327bee8 100644
--- a/run.py
+++ b/run.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
-import multiprocessing
+
+import platform
 import sys
 import time
 import torch
@@ -21,7 +22,6 @@ import psutil
 import cv2
 import threading
 from PIL import Image, ImageTk
-import resource
 
 pool = None
 args = {}
@@ -44,9 +44,14 @@ if os.name == "nt":
 
 
 def limit_resources():
-    if args['max_memory'] <= 1:
+    if args['max_memory'] >= 1:
         memory = args['max_memory'] * 1024 * 1024 * 1024
-        resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))
+        if str(platform.system()).lower() == 'linux':
+            import resource
+            resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))
+        if str(platform.system()).lower() == 'windows':
+            import win32api
+            win32api.SetProcessWorkingSetSize(-1, memory, memory)
 
 
 def pre_check():