update web_config

iopaint/api.py

@@ -201,8 +201,8 @@ class Api:
             enableAutoSaving=self.config.output_dir is not None,
             enableControlnet=self.model_manager.enable_controlnet,
             controlnetMethod=self.model_manager.controlnet_method,
-            disableModelSwitch=self.config.disable_model_switch,
+            disableModelSwitch=False,
-            isDesktop=self.config.gui,
+            isDesktop=False,
             samplers=self.api_samplers(),
         )
 
@@ -380,8 +380,6 @@ if __name__ == "__main__":
             disable_nsfw_checker=False,
             cpu_textencoder=False,
             device="cpu",
-            gui=False,
-            disable_model_switch=False,
             input="/Users/cwq/code/github/MI-GAN/examples/places2_512_object/images",
             output_dir="/Users/cwq/code/github/lama-cleaner/tmp",
             quality=100,

iopaint/cli.py

@@ -8,6 +8,7 @@ from typer import Option
 
 from iopaint.const import *
 from iopaint.runtime import setup_model_dir, dump_environment_info, check_device
+from iopaint.schema import InteractiveSegModel, Device, RealESRGANModel
 
 typer_app = typer.Typer(pretty_exceptions_show_locals=False, add_completion=False)
 
@@ -96,8 +97,8 @@ def start(
     port: int = Option(8080),
     model: str = Option(
         DEFAULT_MODEL,
-        help=f"Available erase models: [{', '.join(AVAILABLE_MODELS)}]. "
+        help=f"Erase models: [{', '.join(AVAILABLE_MODELS)}].\n"
-        f"You can use download command to download other SD/SDXL normal/inpainting models on huggingface",
+        f"Diffusion models: [{', '.join(DIFFUSION_MODELS)}] or any SD/SDXL normal/inpainting models on HuggingFace.",
     ),
     model_dir: Path = Option(
         DEFAULT_MODEL_DIR,
@@ -106,16 +107,13 @@ def start(
         file_okay=False,
         callback=setup_model_dir,
     ),
-    low_mem: bool = Option(
+    low_mem: bool = Option(False, help=LOW_MEM_HELP),
-        False, help="Enable attention slicing and vae tiling to save memory."
-    ),
     no_half: bool = Option(False, help=NO_HALF_HELP),
     cpu_offload: bool = Option(False, help=CPU_OFFLOAD_HELP),
     disable_nsfw_checker: bool = Option(False, help=DISABLE_NSFW_HELP),
     cpu_textencoder: bool = Option(False, help=CPU_TEXTENCODER_HELP),
     local_files_only: bool = Option(False, help=LOCAL_FILES_ONLY_HELP),
     device: Device = Option(Device.cpu),
-    disable_model_switch: bool = Option(False),
     input: Optional[Path] = Option(None, help=INPUT_HELP),
     output_dir: Optional[Path] = Option(
         None, help=OUTPUT_DIR_HELP, dir_okay=True, file_okay=False
@@ -178,8 +176,6 @@ def start(
             local_files_only=local_files_only,
             cpu_textencoder=cpu_textencoder if device == Device.cuda else False,
             device=device,
-            gui=False,
-            disable_model_switch=disable_model_switch,
             input=input,
             output_dir=output_dir,
             quality=quality,
@@ -198,3 +194,13 @@ def start(
         ),
     )
     api.launch()
+
+
+@typer_app.command(help="Start IOPaint web config page")
+def start_web_config(
+    config_file: Path = Option("config.json"),
+):
+    dump_environment_info()
+    from iopaint.web_config import main
+
+    main(config_file)

iopaint/const.py

@@ -1,7 +1,8 @@
 import json
 import os
-from enum import Enum
+from pathlib import Path
-from pydantic import BaseModel
+
+from iopaint.schema import ApiConfig, Device, InteractiveSegModel, RealESRGANModel
 
 INSTRUCT_PIX2PIX_NAME = "timbrooks/instruct-pix2pix"
 KANDINSKY22_NAME = "kandinsky-community/kandinsky-2-2-decoder-inpaint"
@@ -26,10 +27,16 @@ MPS_UNSUPPORT_MODELS = [
 
 DEFAULT_MODEL = "lama"
 AVAILABLE_MODELS = ["lama", "ldm", "zits", "mat", "fcf", "manga", "cv2", "migan"]
-
+DIFFUSION_MODELS = [
-
+    "runwayml/stable-diffusion-inpainting",
-AVAILABLE_DEVICES = ["cuda", "cpu", "mps"]
+    "Uminosachi/realisticVisionV51_v51VAE-inpainting",
-DEFAULT_DEVICE = "cuda"
+    "redstonehero/dreamshaper-inpainting",
+    "Sanster/anything-4.0-inpainting",
+    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+    "Fantasy-Studio/Paint-by-Example",
+    POWERPAINT_NAME,
+    ANYTEXT_NAME,
+]
 
 NO_HALF_HELP = """
 Using full precision(fp32) model.
@@ -40,6 +47,8 @@ CPU_OFFLOAD_HELP = """
 Offloads diffusion model's weight to CPU RAM, significantly reducing vRAM usage.
 """
 
+LOW_MEM_HELP = "Enable attention slicing and vae tiling to save memory."
+
 DISABLE_NSFW_HELP = """
 Disable NSFW checker for diffusion model.
 """
@@ -77,9 +86,10 @@ LOCAL_FILES_ONLY_HELP = """
 When loading diffusion models, using local files only, not connect to HuggingFace server.
 """
 
-DEFAULT_MODEL_DIR = os.getenv(
+DEFAULT_MODEL_DIR = os.path.abspath(
-    "XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")
+    os.getenv("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache"))
 )
+
 MODEL_DIR_HELP = f"""
 Model download directory (by setting XDG_CACHE_HOME environment variable), by default model download to {DEFAULT_MODEL_DIR}
 """
@@ -101,80 +111,40 @@ QUALITY_HELP = """
 Quality of image encoding, 0-100. Default is 95, higher quality will generate larger file size.
 """
 
-
-class Choices(str, Enum):
-    @classmethod
-    def values(cls):
-        return [member.value for member in cls]
-
-
-class RealESRGANModel(Choices):
-    realesr_general_x4v3 = "realesr-general-x4v3"
-    RealESRGAN_x4plus = "RealESRGAN_x4plus"
-    RealESRGAN_x4plus_anime_6B = "RealESRGAN_x4plus_anime_6B"
-
-
-class Device(Choices):
-    cpu = "cpu"
-    cuda = "cuda"
-    mps = "mps"
-
-
-class InteractiveSegModel(Choices):
-    vit_b = "vit_b"
-    vit_l = "vit_l"
-    vit_h = "vit_h"
-    mobile_sam = "mobile_sam"
-
-
 INTERACTIVE_SEG_HELP = "Enable interactive segmentation using Segment Anything."
-INTERACTIVE_SEG_MODEL_HELP = "Model size: vit_b < vit_l < vit_h. Bigger model size means better segmentation but slower speed."
+INTERACTIVE_SEG_MODEL_HELP = "Model size: mobile_sam < vit_b < vit_l < vit_h. Bigger model size means better segmentation but slower speed."
 REMOVE_BG_HELP = "Enable remove background. Always run on CPU"
 ANIMESEG_HELP = "Enable anime segmentation. Always run on CPU"
 REALESRGAN_HELP = "Enable realesrgan super resolution"
-GFPGAN_HELP = (
+GFPGAN_HELP = "Enable GFPGAN face restore. To also enhance background, use with --enable-realesrgan"
-    "Enable GFPGAN face restore. To enhance background, use with --enable-realesrgan"
+RESTOREFORMER_HELP = "Enable RestoreFormer face restore. To also enhance background, use with --enable-realesrgan"
-)
-RESTOREFORMER_HELP = "Enable RestoreFormer face restore. To enhance background, use with --enable-realesrgan"
 GIF_HELP = "Enable GIF plugin. Make GIF to compare original and cleaned image"
 
-
+default_configs = dict(
-class Config(BaseModel):
+    host="127.0.0.1",
-    host: str = "127.0.0.1"
+    port=8080,
-    port: int = 8080
+    model=DEFAULT_MODEL,
-    model: str = DEFAULT_MODEL
+    model_dir=DEFAULT_MODEL_DIR,
-    sd_local_model_path: str = None
+    no_half=False,
-    device: str = DEFAULT_DEVICE
+    low_mem=False,
-    gui: bool = False
+    cpu_offload=False,
-    no_gui_auto_close: bool = False
+    disable_nsfw_checker=False,
-    no_half: bool = False
+    local_files_only=False,
-    cpu_offload: bool = False
+    cpu_textencoder=False,
-    disable_nsfw: bool = False
+    device=Device.cuda,
-    sd_cpu_textencoder: bool = False
+    input=None,
-    local_files_only: bool = False
+    output_dir=None,
-    model_dir: str = DEFAULT_MODEL_DIR
+    quality=95,
-    input: str = None
+    enable_interactive_seg=False,
-    output_dir: str = None
+    interactive_seg_model=InteractiveSegModel.vit_b,
-    # plugins
+    interactive_seg_device=Device.cpu,
-    enable_interactive_seg: bool = False
+    enable_remove_bg=False,
-    interactive_seg_model: str = "vit_l"
+    enable_anime_seg=False,
-    interactive_seg_device: str = "cpu"
+    enable_realesrgan=False,
-    enable_remove_bg: bool = False
+    realesrgan_device=Device.cpu,
-    enable_anime_seg: bool = False
+    realesrgan_model=RealESRGANModel.realesr_general_x4v3,
-    enable_realesrgan: bool = False
+    enable_gfpgan=False,
-    realesrgan_device: str = "cpu"
+    gfpgan_device=Device.cpu,
-    realesrgan_model: str = RealESRGANModel.realesr_general_x4v3.value
+    enable_restoreformer=False,
-    realesrgan_no_half: bool = False
+    restoreformer_device=Device.cpu,
-    enable_gfpgan: bool = False
+)
-    gfpgan_device: str = "cpu"
-    enable_restoreformer: bool = False
-    restoreformer_device: str = "cpu"
-    enable_gif: bool = False
-
-
-def load_config(installer_config: str):
-    if os.path.exists(installer_config):
-        with open(installer_config, "r", encoding="utf-8") as f:
-            return Config(**json.load(f))
-    else:
-        return Config()

iopaint/plugins/__init__.py

@@ -8,7 +8,7 @@ from .interactive_seg import InteractiveSeg
 from .realesrgan import RealESRGANUpscaler
 from .remove_bg import RemoveBG
 from .restoreformer import RestoreFormerPlugin
-from ..const import InteractiveSegModel, Device, RealESRGANModel
+from ..schema import InteractiveSegModel, Device, RealESRGANModel
 
 
 def build_plugins(

iopaint/plugins/realesrgan.py

@@ -1,14 +1,11 @@
-from enum import Enum
-
 import cv2
 import numpy as np
 import torch
 from loguru import logger
 
-from iopaint.const import RealESRGANModel
 from iopaint.helper import download_model
 from iopaint.plugins.base_plugin import BasePlugin
-from iopaint.schema import RunPluginRequest
+from iopaint.schema import RunPluginRequest, RealESRGANModel
 
 
 class RealESRGANUpscaler(BasePlugin):

iopaint/schema.py

@@ -1,3 +1,4 @@
+import json
 import random
 from enum import Enum
 from pathlib import Path
@@ -6,7 +7,30 @@ from typing import Optional, Literal, List
 from loguru import logger
 from pydantic import BaseModel, Field, field_validator
 
-from iopaint.const import Device, InteractiveSegModel, RealESRGANModel
+
+class Choices(str, Enum):
+    @classmethod
+    def values(cls):
+        return [member.value for member in cls]
+
+
+class RealESRGANModel(Choices):
+    realesr_general_x4v3 = "realesr-general-x4v3"
+    RealESRGAN_x4plus = "RealESRGAN_x4plus"
+    RealESRGAN_x4plus_anime_6B = "RealESRGAN_x4plus_anime_6B"
+
+
+class Device(Choices):
+    cpu = "cpu"
+    cuda = "cuda"
+    mps = "mps"
+
+
+class InteractiveSegModel(Choices):
+    vit_b = "vit_b"
+    vit_l = "vit_l"
+    vit_h = "vit_h"
+    mobile_sam = "mobile_sam"
 
 
 class PluginInfo(BaseModel):
@@ -93,8 +117,6 @@ class ApiConfig(BaseModel):
     local_files_only: bool
     cpu_textencoder: bool
     device: Device
-    gui: bool
-    disable_model_switch: bool
     input: Optional[Path]
     output_dir: Optional[Path]
     quality: int

iopaint/web_config.py

@@ -1,31 +1,43 @@
-import json
-import os
 from datetime import datetime
+from json import JSONDecodeError
 
 import gradio as gr
 from loguru import logger
 
 from iopaint.const import *
 
-_config_file = None
+
+_config_file: Path = None
+
+
+class WebConfig(ApiConfig):
+    model_dir: str = DEFAULT_MODEL_DIR
+
+
+def load_config(p: Path) -> WebConfig:
+    if p.exists():
+        with open(p, "r", encoding="utf-8") as f:
+            try:
+                return WebConfig(**{**default_configs, **json.load(f)})
+            except JSONDecodeError:
+                print(f"Load config file failed, using default configs")
+                return WebConfig(**default_configs)
+    else:
+        return WebConfig(**default_configs)
 
 
 def save_config(
     host,
     port,
     model,
-    sd_local_model_path,
-    enable_controlnet,
-    controlnet_method,
-    device,
-    gui,
-    no_gui_auto_close,
-    no_half,
-    cpu_offload,
-    disable_nsfw,
-    sd_cpu_textencoder,
-    local_files_only,
     model_dir,
+    no_half,
+    low_mem,
+    cpu_offload,
+    disable_nsfw_checker,
+    local_files_only,
+    cpu_textencoder,
+    device,
     input,
     output_dir,
     quality,
@@ -41,33 +53,29 @@ def save_config(
     gfpgan_device,
     enable_restoreformer,
     restoreformer_device,
-    enable_gif,
 ):
-    config = InpaintRequest(**locals())
+    config = WebConfig(**locals())
+    if str(config.input) == ".":
+        config.input = None
+    if str(config.output_dir) == ".":
+        config.output_dir = None
+
     print(config)
     if config.input and not os.path.exists(config.input):
         return "[Error] Input file or directory does not exist"
 
     current_time = datetime.now().strftime("%H:%M:%S")
-    msg = f"[{current_time}] Successful save config to: {os.path.abspath(_config_file)}"
+    msg = f"[{current_time}] Successful save config to: {str(_config_file.absolute())}"
     logger.info(msg)
     try:
         with open(_config_file, "w", encoding="utf-8") as f:
-            json.dump(config.dict(), f, indent=4, ensure_ascii=False)
+            f.write(config.model_dump_json(indent=4))
     except Exception as e:
-        return f"Save failed: {str(e)}"
+        return f"Save configure file failed: {str(e)}"
     return msg
 
 
-def close_server(*args):
+def main(config_file: Path):
-    # TODO: make close both browser and server works
-    import os, signal
-
-    pid = os.getpid()
-    os.kill(pid, signal.SIGUSR1)
-
-
-def main(config_file: str):
     global _config_file
     _config_file = config_file
 
@@ -75,7 +83,9 @@ def main(config_file: str):
 
     with gr.Blocks() as demo:
         with gr.Row():
-            with gr.Column(scale=1):
+            with gr.Column():
+                gr.Textbox(config_file, label="Config file", interactive=False)
+            with gr.Column():
                 save_btn = gr.Button(value="Save configurations")
                 message = gr.HTML()
 
@@ -86,10 +96,12 @@ def main(config_file: str):
                     port = gr.Number(init_config.port, label="Port", precision=0)
 
                 model = gr.Radio(
-                    AVAILABLE_MODELS, label="Model", value=init_config.model
+                    AVAILABLE_MODELS + DIFFUSION_MODELS,
+                    label="Models (https://www.iopaint.com/models)",
+                    value=init_config.model,
                 )
                 device = gr.Radio(
-                    AVAILABLE_DEVICES, label="Device", value=init_config.device
+                    Device.values(), label="Device", value=init_config.device
                 )
                 quality = gr.Slider(
                     value=95,
@@ -99,8 +111,20 @@ def main(config_file: str):
                     step=1,
                 )
 
-                with gr.Column():
+                no_half = gr.Checkbox(init_config.no_half, label=f"{NO_HALF_HELP}")
-                    gui = gr.Checkbox(init_config.gui, label=f"{GUI_HELP}")
+                cpu_offload = gr.Checkbox(
+                    init_config.cpu_offload, label=f"{CPU_OFFLOAD_HELP}"
+                )
+                low_mem = gr.Checkbox(init_config.low_mem, label=f"{LOW_MEM_HELP}")
+                cpu_textencoder = gr.Checkbox(
+                    init_config.cpu_textencoder, label=f"{CPU_TEXTENCODER_HELP}"
+                )
+                disable_nsfw_checker = gr.Checkbox(
+                    init_config.disable_nsfw_checker, label=f"{DISABLE_NSFW_HELP}"
+                )
+                local_files_only = gr.Checkbox(
+                    init_config.local_files_only, label=f"{LOCAL_FILES_ONLY_HELP}"
+                )
 
                 with gr.Column():
                     model_dir = gr.Textbox(
@@ -116,16 +140,17 @@ def main(config_file: str):
                     )
 
             with gr.Tab("Plugins"):
+                with gr.Row():
                     enable_interactive_seg = gr.Checkbox(
                         init_config.enable_interactive_seg, label=INTERACTIVE_SEG_HELP
                     )
                     interactive_seg_model = gr.Radio(
-                    AVAILABLE_INTERACTIVE_SEG_MODELS,
+                        InteractiveSegModel.values(),
                         label=f"Segment Anything models. {INTERACTIVE_SEG_MODEL_HELP}",
                         value=init_config.interactive_seg_model,
                     )
                     interactive_seg_device = gr.Radio(
-                    AVAILABLE_INTERACTIVE_SEG_DEVICES,
+                        Device.values(),
                         label="Segment Anything Device",
                         value=init_config.interactive_seg_device,
                     )
@@ -143,12 +168,12 @@ def main(config_file: str):
                         init_config.enable_realesrgan, label=REALESRGAN_HELP
                     )
                     realesrgan_device = gr.Radio(
-                        REALESRGAN_AVAILABLE_DEVICES,
+                        Device.values(),
                         label="RealESRGAN Device",
                         value=init_config.realesrgan_device,
                     )
                     realesrgan_model = gr.Radio(
-                        RealESRGANModelNameList,
+                        RealESRGANModel.values(),
                         label="RealESRGAN model",
                         value=init_config.realesrgan_model,
                     )
@@ -157,7 +182,7 @@ def main(config_file: str):
                         init_config.enable_gfpgan, label=GFPGAN_HELP
                     )
                     gfpgan_device = gr.Radio(
-                        GFPGAN_AVAILABLE_DEVICES,
+                        Device.values(),
                         label="GFPGAN Device",
                         value=init_config.gfpgan_device,
                     )
@@ -166,37 +191,10 @@ def main(config_file: str):
                         init_config.enable_restoreformer, label=RESTOREFORMER_HELP
                     )
                     restoreformer_device = gr.Radio(
-                        RESTOREFORMER_AVAILABLE_DEVICES,
+                        Device.values(),
                         label="RestoreFormer Device",
                         value=init_config.restoreformer_device,
                     )
-                enable_gif = gr.Checkbox(init_config.enable_gif, label=GIF_HELP)
-
-            with gr.Tab("Diffusion Model"):
-                sd_local_model_path = gr.Textbox(
-                    init_config.sd_local_model_path, label=f"{SD_LOCAL_MODEL_HELP}"
-                )
-                enable_controlnet = gr.Checkbox(
-                    init_config.enable_controlnet, label=f"{SD_CONTROLNET_HELP}"
-                )
-                controlnet_method = gr.Radio(
-                    SD_CONTROLNET_CHOICES,
-                    label="ControlNet method",
-                    value=init_config.controlnet_method,
-                )
-                no_half = gr.Checkbox(init_config.no_half, label=f"{NO_HALF_HELP}")
-                cpu_offload = gr.Checkbox(
-                    init_config.cpu_offload, label=f"{CPU_OFFLOAD_HELP}"
-                )
-                sd_cpu_textencoder = gr.Checkbox(
-                    init_config.sd_cpu_textencoder, label=f"{CPU_TEXTENCODER_HELP}"
-                )
-                disable_nsfw = gr.Checkbox(
-                    init_config.disable_nsfw, label=f"{DISABLE_NSFW_HELP}"
-                )
-                local_files_only = gr.Checkbox(
-                    init_config.local_files_only, label=f"{LOCAL_FILES_ONLY_HELP}"
-                )
 
         save_btn.click(
             save_config,
@@ -204,18 +202,14 @@ def main(config_file: str):
                 host,
                 port,
                 model,
-                sd_local_model_path,
-                enable_controlnet,
-                controlnet_method,
-                device,
-                gui,
-                no_gui_auto_close,
-                no_half,
-                cpu_offload,
-                disable_nsfw,
-                sd_cpu_textencoder,
-                local_files_only,
                 model_dir,
+                no_half,
+                low_mem,
+                cpu_offload,
+                disable_nsfw_checker,
+                local_files_only,
+                cpu_textencoder,
+                device,
                 input,
                 output_dir,
                 quality,
@@ -231,7 +225,6 @@ def main(config_file: str):
                 gfpgan_device,
                 enable_restoreformer,
                 restoreformer_device,
-                enable_gif,
             ],
             message,
         )

requirements.txt

@@ -18,5 +18,6 @@ yacs
 piexif==1.1.3
 omegaconf
 easydict
+gradio
 
 Pillow==9.5.0 # for AnyText