From 5d291c28ad791a9e8b243534f334436741fbcdfe Mon Sep 17 00:00:00 2001
From: layerdiffusion <19834515+lllyasviel@users.noreply.github.com>
Date: Mon, 29 Jul 2024 11:22:20 -0600
Subject: [PATCH] implement args for new backend

---
 backend/args.py   | 7 +++++++
 backend/stream.py | 7 +------
 2 files changed, 8 insertions(+), 6 deletions(-)
 create mode 100644 backend/args.py

diff --git a/backend/args.py b/backend/args.py
new file mode 100644
index 00000000..3ef5f5af
--- /dev/null
+++ b/backend/args.py
@@ -0,0 +1,7 @@
+import argparse
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--cuda-stream", action="store_true")
+
+args = parser.parse_known_args()[0]
diff --git a/backend/stream.py b/backend/stream.py
index a231caf5..3972d0e4 100644
--- a/backend/stream.py
+++ b/backend/stream.py
@@ -1,5 +1,5 @@
 import torch
-import argparse
+from backend import args
 
 
 def stream_context():
@@ -56,11 +56,6 @@ current_stream = None
 mover_stream = None
 using_stream = False
 
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--cuda-stream", action="store_true")
-args = parser.parse_known_args()[0]
-
 if args.cuda_stream:
     current_stream = get_current_stream()
     mover_stream = get_new_stream()