From 45d29a42496e5e55ae5fc25f4ba9ce3c00aacb99 Mon Sep 17 00:00:00 2001
From: Qing <cwq1913@gmail.com>
Date: Thu, 8 Feb 2024 17:35:55 +0800
Subject: [PATCH] add inbrowser

---
 iopaint/cli.py        | 13 ++++++++++++-
 iopaint/const.py      |  2 ++
 iopaint/schema.py     |  1 +
 iopaint/web_config.py |  4 ++++
 4 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/iopaint/cli.py b/iopaint/cli.py
index 4ad06a7..88fda7b 100644
--- a/iopaint/cli.py
+++ b/iopaint/cli.py
@@ -1,3 +1,5 @@
+import webbrowser
+from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Dict, Optional
 
@@ -97,6 +99,7 @@ def run(
 def start(
     host: str = Option("127.0.0.1"),
     port: int = Option(8080),
+    inbrowser: bool = Option(False, help=INBROWSER_HELP),
     model: str = Option(
         DEFAULT_MODEL,
         help=f"Erase models: [{', '.join(AVAILABLE_MODELS)}].\n"
@@ -165,10 +168,18 @@ def start(
     from iopaint.api import Api
     from iopaint.schema import ApiConfig
 
-    app = FastAPI()
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        if inbrowser:
+            webbrowser.open(f"http://localhost:{port}", new=0, autoraise=True)
+        yield
+
+    app = FastAPI(lifespan=lifespan)
+
     api_config = ApiConfig(
         host=host,
         port=port,
+        inbrowser=inbrowser,
         model=model,
         no_half=no_half,
         low_mem=low_mem,
diff --git a/iopaint/const.py b/iopaint/const.py
index 5d5605b..148cb77 100644
--- a/iopaint/const.py
+++ b/iopaint/const.py
@@ -116,3 +116,5 @@ REALESRGAN_HELP = "Enable realesrgan super resolution"
 GFPGAN_HELP = "Enable GFPGAN face restore. To also enhance background, use with --enable-realesrgan"
 RESTOREFORMER_HELP = "Enable RestoreFormer face restore. To also enhance background, use with --enable-realesrgan"
 GIF_HELP = "Enable GIF plugin. Make GIF to compare original and cleaned image"
+
+INBROWSER_HELP = "Automatically launch IOPaint in a new tab on the default browser"
diff --git a/iopaint/schema.py b/iopaint/schema.py
index a8972ed..b7c8b44 100644
--- a/iopaint/schema.py
+++ b/iopaint/schema.py
@@ -216,6 +216,7 @@ class PowerPaintTask(str, Enum):
 class ApiConfig(BaseModel):
     host: str
     port: int
+    inbrowser: bool
     model: str
     no_half: bool
     low_mem: bool
diff --git a/iopaint/web_config.py b/iopaint/web_config.py
index 4e34c14..729fe27 100644
--- a/iopaint/web_config.py
+++ b/iopaint/web_config.py
@@ -28,6 +28,7 @@ _config_file: Path = None
 default_configs = dict(
     host="127.0.0.1",
     port=8080,
+    inbrowser=True,
     model=DEFAULT_MODEL,
     model_dir=DEFAULT_MODEL_DIR,
     no_half=False,
@@ -100,6 +101,7 @@ def save_config(
     gfpgan_device,
     enable_restoreformer,
     restoreformer_device,
+    inbrowser,
 ):
     config = WebConfig(**locals())
     if str(config.input) == ".":
@@ -146,6 +148,7 @@ def main(config_file: Path):
                 with gr.Row():
                     host = gr.Textbox(init_config.host, label="Host")
                     port = gr.Number(init_config.port, label="Port", precision=0)
+                    inbrowser = gr.Checkbox(init_config.inbrowser, label=INBROWSER_HELP)
 
                 with gr.Column():
                     model = gr.Textbox(
@@ -297,6 +300,7 @@ def main(config_file: Path):
                 gfpgan_device,
                 enable_restoreformer,
                 restoreformer_device,
+                inbrowser,
             ],
             message,
         )