From 054a3416f1eac7e487bfd271ce89d06bfa533105 Mon Sep 17 00:00:00 2001
From: layerdiffusion <19834515+lllyasviel@users.noreply.github.com>
Date: Mon, 19 Aug 2024 08:06:24 -0700
Subject: [PATCH] revise space logics

---
 modules_forge/forge_space.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/modules_forge/forge_space.py b/modules_forge/forge_space.py
index 2093af0c..a1943eee 100644
--- a/modules_forge/forge_space.py
+++ b/modules_forge/forge_space.py
@@ -131,6 +131,11 @@ class ForgeSpace:
     def gradio_worker(self):
         import spaces
         spaces.unload_module()
+        original_cwd = os.getcwd()
+        os.chdir(self.hf_path)
+
+        if 'models' in sys.modules:
+            del sys.modules['models']
 
         memory_management.unload_all_models()
         sys.path.insert(0, self.hf_path)
@@ -155,17 +160,21 @@ class ForgeSpace:
             server_port=port
         )
 
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+        if 'models' in sys.modules:
+            del sys.modules['models']
+
+        sys.path.remove(self.hf_path)
+        sys.path.remove(self.root_path)
+        os.chdir(original_cwd)
+
         while self.is_running:
             time.sleep(0.1)
 
         demo.close()
         self.gradio_metas = None
-
-        if module_name in sys.modules:
-            del sys.modules[module_name]
-
-        sys.path.remove(self.hf_path)
-        sys.path.remove(self.root_path)
         return