From 4ef29f95468467c569d8a68b85078a88cb038893 Mon Sep 17 00:00:00 2001
From: layerdiffusion <19834515+lllyasviel@users.noreply.github.com>
Date: Tue, 20 Aug 2024 20:40:56 -0700
Subject: [PATCH] revise space

---
 spaces.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/spaces.py b/spaces.py
index 0eb4b882..dc0d56ce 100644
--- a/spaces.py
+++ b/spaces.py
@@ -10,6 +10,7 @@ import inspect
 import functools
 import gradio.oauth
 import gradio.routes
+import contextlib
 
 from backend import memory_management
 from backend.operations import DynamicSwapInstaller
@@ -126,8 +127,11 @@ class GPUObject:
         return self
 
 
-def capture_gpu_object():
-    return GPUObject()
+def capture_gpu_object(capture=True):
+    if capture:
+        return GPUObject()
+    else:
+        return contextlib.nullcontext()
 
 
 def GPU(gpu_objects=None, manual_load=False, **kwargs):