diff --git a/iopaint/api.py b/iopaint/api.py index 651bd76..a69a7f8 100644 --- a/iopaint/api.py +++ b/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, - isDesktop=self.config.gui, + disableModelSwitch=False, + 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, diff --git a/iopaint/cli.py b/iopaint/cli.py index ef95a85..cc1f186 100644 --- a/iopaint/cli.py +++ b/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)}]. " - f"You can use download command to download other SD/SDXL normal/inpainting models on huggingface", + help=f"Erase models: [{', '.join(AVAILABLE_MODELS)}].\n" + 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( - False, help="Enable attention slicing and vae tiling to save memory." - ), + low_mem: bool = Option(False, help=LOW_MEM_HELP), 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) diff --git a/iopaint/const.py b/iopaint/const.py index f2e04fe..5cdd5d7 100644 --- a/iopaint/const.py +++ b/iopaint/const.py @@ -1,7 +1,8 @@ import json import os -from enum import Enum -from pydantic import BaseModel +from pathlib import Path + +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"] - - -AVAILABLE_DEVICES = ["cuda", "cpu", "mps"] -DEFAULT_DEVICE = "cuda" +DIFFUSION_MODELS = [ + "runwayml/stable-diffusion-inpainting", + "Uminosachi/realisticVisionV51_v51VAE-inpainting", + "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( - "XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache") +DEFAULT_MODEL_DIR = os.path.abspath( + 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 = ( - "Enable GFPGAN face restore. To enhance background, use with --enable-realesrgan" -) -RESTOREFORMER_HELP = "Enable RestoreFormer face restore. To enhance background, use with --enable-realesrgan" +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" - -class Config(BaseModel): - host: str = "127.0.0.1" - port: int = 8080 - model: str = DEFAULT_MODEL - sd_local_model_path: str = None - device: str = DEFAULT_DEVICE - gui: bool = False - no_gui_auto_close: bool = False - no_half: bool = False - cpu_offload: bool = False - disable_nsfw: bool = False - sd_cpu_textencoder: bool = False - local_files_only: bool = False - model_dir: str = DEFAULT_MODEL_DIR - input: str = None - output_dir: str = None - # plugins - enable_interactive_seg: bool = False - interactive_seg_model: str = "vit_l" - interactive_seg_device: str = "cpu" - enable_remove_bg: bool = False - enable_anime_seg: bool = False - enable_realesrgan: bool = False - realesrgan_device: str = "cpu" - realesrgan_model: str = RealESRGANModel.realesr_general_x4v3.value - realesrgan_no_half: bool = False - 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() +default_configs = dict( + host="127.0.0.1", + port=8080, + model=DEFAULT_MODEL, + model_dir=DEFAULT_MODEL_DIR, + no_half=False, + low_mem=False, + cpu_offload=False, + disable_nsfw_checker=False, + local_files_only=False, + cpu_textencoder=False, + device=Device.cuda, + input=None, + output_dir=None, + quality=95, + enable_interactive_seg=False, + interactive_seg_model=InteractiveSegModel.vit_b, + interactive_seg_device=Device.cpu, + enable_remove_bg=False, + enable_anime_seg=False, + enable_realesrgan=False, + realesrgan_device=Device.cpu, + realesrgan_model=RealESRGANModel.realesr_general_x4v3, + enable_gfpgan=False, + gfpgan_device=Device.cpu, + enable_restoreformer=False, + restoreformer_device=Device.cpu, +) diff --git a/iopaint/plugins/__init__.py b/iopaint/plugins/__init__.py index fd321cd..3e7d5cf 100644 --- a/iopaint/plugins/__init__.py +++ b/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( diff --git a/iopaint/plugins/realesrgan.py b/iopaint/plugins/realesrgan.py index 3f68303..8165fa3 100644 --- a/iopaint/plugins/realesrgan.py +++ b/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): diff --git a/iopaint/schema.py b/iopaint/schema.py index f1f82d7..6fd3980 100644 --- a/iopaint/schema.py +++ b/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 diff --git a/iopaint/web_config.py b/iopaint/web_config.py index d3da8d0..55ccd98 100644 --- a/iopaint/web_config.py +++ b/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): - # TODO: make close both browser and server works - import os, signal - - pid = os.getpid() - os.kill(pid, signal.SIGUSR1) - - -def main(config_file: str): +def main(config_file: Path): 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(): - gui = gr.Checkbox(init_config.gui, label=f"{GUI_HELP}") + 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}" + ) + 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,19 +140,20 @@ def main(config_file: str): ) with gr.Tab("Plugins"): - enable_interactive_seg = gr.Checkbox( - init_config.enable_interactive_seg, label=INTERACTIVE_SEG_HELP - ) - interactive_seg_model = gr.Radio( - AVAILABLE_INTERACTIVE_SEG_MODELS, - label=f"Segment Anything models. {INTERACTIVE_SEG_MODEL_HELP}", - value=init_config.interactive_seg_model, - ) - interactive_seg_device = gr.Radio( - AVAILABLE_INTERACTIVE_SEG_DEVICES, - label="Segment Anything Device", - value=init_config.interactive_seg_device, - ) + with gr.Row(): + enable_interactive_seg = gr.Checkbox( + init_config.enable_interactive_seg, label=INTERACTIVE_SEG_HELP + ) + interactive_seg_model = gr.Radio( + InteractiveSegModel.values(), + label=f"Segment Anything models. {INTERACTIVE_SEG_MODEL_HELP}", + value=init_config.interactive_seg_model, + ) + interactive_seg_device = gr.Radio( + Device.values(), + label="Segment Anything Device", + value=init_config.interactive_seg_device, + ) with gr.Row(): enable_remove_bg = gr.Checkbox( init_config.enable_remove_bg, label=REMOVE_BG_HELP @@ -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, ) diff --git a/requirements.txt b/requirements.txt index a636f7e..012bb90 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,5 +18,6 @@ yacs piexif==1.1.3 omegaconf easydict +gradio Pillow==9.5.0 # for AnyText