From 68f62c835287f7de160c3beda3fceab3bf1ff5d8 Mon Sep 17 00:00:00 2001
From: Cragin Godley <1012191+cgodley@users.noreply.github.com>
Date: Fri, 2 Sep 2022 10:43:32 -0400
Subject: [PATCH] web: allow custom host/port, default to 127.0.0.1 for
 security reasons

---
 scripts/dream.py | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/scripts/dream.py b/scripts/dream.py
index 1535ac386c..d1db1eb11b 100755
--- a/scripts/dream.py
+++ b/scripts/dream.py
@@ -99,7 +99,7 @@ def main():
 
     cmd_parser = create_cmd_parser()
     if opt.web:
-        dream_server_loop(t2i)
+        dream_server_loop(t2i, opt.host, opt.port)
     else:
         main_loop(t2i, opt.outdir, opt.prompt_as_dir, cmd_parser, infile)
 
@@ -270,7 +270,7 @@ def get_next_command(infile=None) -> str: #command string
         print(f'#{command}')
     return command
 
-def dream_server_loop(t2i):
+def dream_server_loop(t2i, host, port):
     print('\n* --web was specified, starting web server...')
     # Change working directory to the stable-diffusion directory
     os.chdir(
@@ -279,9 +279,9 @@ def dream_server_loop(t2i):
 
     # Start server
     DreamServer.model = t2i
-    dream_server = ThreadingDreamServer(("0.0.0.0", 9090))
+    dream_server = ThreadingDreamServer((host, port))
     print("\nStarted Stable Diffusion dream server!")
-    print("Point your browser at http://localhost:9090 or use the host's DNS name or IP address.")
+    print(f"Point your browser at http://{host}:{port} or use the host's DNS name or IP address.")
 
     try:
         dream_server.serve_forever()
@@ -417,6 +417,18 @@ def create_argv_parser():
         action='store_true',
         help='Start in web server mode.',
     )
+    parser.add_argument(
+        '--host',
+        type=str,
+        default='127.0.0.1',
+        help='Web server: Host or IP to listen on'
+    )
+    parser.add_argument(
+        '--port',
+        type=int,
+        default='9090',
+        help='Web server: Port to listen on'
+    )
     parser.add_argument(
         '--weights',
         default='model',