kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

update web_config

Browse Source
This commit is contained in:
Qing 2023-03-28 21:55:25 +08:00
parent 14596247c1
commit e734859128
4 changed files with 151 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
lama_cleaner/const.py
View File

@ -1,4 +1,5 @@
import os
from enum import Enum
MPS_SUPPORT_MODELS = [
"instruct_pix2pix",
@ -86,3 +87,22 @@ Prevent backend auto close after the GUI window closed.
QUALITY_HELP = """
Quality of image encoding, 0-100. Default is 95, higher quality will generate larger file size.
"""
class RealESRGANModelName(str, Enum):
realesr_general_x4v3 = "realesr-general-x4v3"
RealESRGAN_x4plus = "RealESRGAN_x4plus"
RealESRGAN_x4plus_anime_6B = "RealESRGAN_x4plus_anime_6B"
RealESRGANModelNameList = [e.value for e in RealESRGANModelName]
INTERACTIVE_SEG_HELP = "Enable interactive segmentation. Always run on CPU"
REMOVE_BG_HELP = "Enable remove background. Always run on CPU"
REALESRGAN_HELP = "Enable realesrgan super resolution"
REALESRGAN_AVAILABLE_DEVICES = ["cpu", "cuda", "mps"]
GFPGAN_HELP = (
"Enable GFPGAN face restore. To enhance background, use with --enable-realesrgan"
)
GFPGAN_AVAILABLE_DEVICES = ["cpu", "cuda", "mps"]
GIF_HELP = "Enable GIF plugin. Make GIF to compare original and cleaned image"

28
lama_cleaner/parse_args.py
View File

@ -6,7 +6,6 @@ from pathlib import Path
from loguru import logger
from lama_cleaner.const import *
from lama_cleaner.plugins.realesrgan import RealESRGANModelName, RealESRGANModelNameList
from lama_cleaner.runtime import dump_environment_info
@ -80,20 +79,23 @@ def parse_args():
parser.add_argument(
"--enable-interactive-seg",
action="store_true",
help="Enable interactive segmentation. Always run on CPU",
help=INTERACTIVE_SEG_HELP,
)
parser.add_argument(
"--enable-remove-bg",
action="store_true",
help="Enable remove background. Always run on CPU",
help=REMOVE_BG_HELP,
)
parser.add_argument(
"--enable-realesrgan",
action="store_true",
help="Enable realesrgan super resolution",
help=REALESRGAN_HELP,
)
parser.add_argument(
"--realesrgan-device", default="cpu", type=str, choices=["cpu", "cuda", "mps"]
"--realesrgan-device",
default="cpu",
type=str,
choices=REALESRGAN_AVAILABLE_DEVICES,
)
parser.add_argument(
"--realesrgan-model",
@ -101,18 +103,14 @@ def parse_args():
type=str,
choices=RealESRGANModelNameList,
)
parser.add_argument(
"--enable-gfpgan",
action="store_true",
help="Enable GFPGAN face restore",
)
parser.add_argument("--enable-gfpgan", action="store_true", help=GFPGAN_HELP)
parser.add_argument(
"--gfpgan-device", default="cpu", type=str, choices=["cpu", "cuda", "mps"]
)
parser.add_argument(
"--enable-gif",
action="store_true",
help="Enable GIF plugin",
help=GIF_HELP,
)
#########
@ -200,4 +198,12 @@ def parse_args():
if not output_dir.is_dir():
parser.error(f"invalid --output-dir: {output_dir} is not a directory")
if args.enable_gfpgan:
if args.enable_realesrgan:
logger.info("Use realesrgan as GFPGAN background upscaler")
else:
logger.info(
f"GFPGAN no background upscaler, use --enable-realesrgan to enable it"
)
return args

10
lama_cleaner/plugins/realesrgan.py
View File

@ -3,19 +3,11 @@ from enum import Enum
import cv2
from loguru import logger
from lama_cleaner.const import RealESRGANModelName
from lama_cleaner.helper import download_model
from lama_cleaner.plugins.base_plugin import BasePlugin
class RealESRGANModelName(str, Enum):
realesr_general_x4v3 = "realesr-general-x4v3"
RealESRGAN_x4plus = "RealESRGAN_x4plus"
RealESRGAN_x4plus_anime_6B = "RealESRGAN_x4plus_anime_6B"
RealESRGANModelNameList = [e.value for e in RealESRGANModelName]
class RealESRGANUpscaler(BasePlugin):
name = "RealESRGAN"

84
lama_cleaner/web_config.py
View File

@ -28,6 +28,15 @@ class Config(BaseModel):
model_dir: str = DEFAULT_MODEL_DIR
input: str = None
output_dir: str = None
# plugins
enable_interactive_seg: bool = False
enable_remove_bg: bool = False
enable_realesrgan: bool = False
realesrgan_device: str = "cpu"
realesrgan_model: str = RealESRGANModelName.realesr_general_x4v3.value
enable_gfpgan: bool = False
gfpgan_device: str = "cpu"
enable_gif: bool = False
def load_config(installer_config: str):
@ -56,6 +65,14 @@ def save_config(
input,
output_dir,
quality,
enable_interactive_seg,
enable_remove_bg,
enable_realesrgan,
realesrgan_device,
realesrgan_model,
enable_gfpgan,
gfpgan_device,
enable_gif,
):
config = Config(**locals())
print(config)
@ -92,15 +109,19 @@ def main(config_file: str):
with gr.Column(scale=1):
save_btn = gr.Button(value="Save configurations")
message = gr.HTML()
# with gr.Column(scale=0, min_width=100):
# exit_btn = gr.Button(value="Close")
# exit_btn.click(close_server)
with gr.Tabs():
with gr.Tab("Common"):
with gr.Row():
host = gr.Textbox(init_config.host, label="Host")
port = gr.Number(init_config.port, label="Port", precision=0)
model = gr.Radio(AVAILABLE_MODELS, label="Model", value=init_config.model)
device = gr.Radio(AVAILABLE_DEVICES, label="Device", value=init_config.device)
model = gr.Radio(
AVAILABLE_MODELS, label="Model", value=init_config.model
)
device = gr.Radio(
AVAILABLE_DEVICES, label="Device", value=init_config.device
)
quality = gr.Slider(
value=95,
label=f"Image Quality ({QUALITY_HELP})",
@ -115,15 +136,52 @@ def main(config_file: str):
init_config.no_gui_auto_close, label=f"{NO_GUI_AUTO_CLOSE_HELP}"
)
model_dir = gr.Textbox(init_config.model_dir, label=f"{MODEL_DIR_HELP}")
with gr.Column():
model_dir = gr.Textbox(
init_config.model_dir, label=f"{MODEL_DIR_HELP}"
)
input = gr.Textbox(
init_config.input, label=f"Input file or directory. {INPUT_HELP}"
init_config.input,
label=f"Input file or directory. {INPUT_HELP}",
)
output_dir = gr.Textbox(
init_config.output_dir, label=f"Output directory. {OUTPUT_DIR_HELP}"
init_config.output_dir,
label=f"Output directory. {OUTPUT_DIR_HELP}",
)
with gr.Column():
with gr.Tab("Plugins"):
enable_interactive_seg = gr.Checkbox(
init_config.enable_interactive_seg, label=INTERACTIVE_SEG_HELP
)
enable_remove_bg = gr.Checkbox(
init_config.enable_remove_bg, label=REMOVE_BG_HELP
)
with gr.Row():
enable_realesrgan = gr.Checkbox(
init_config.enable_realesrgan, label=REALESRGAN_HELP
)
realesrgan_device = gr.Radio(
REALESRGAN_AVAILABLE_DEVICES,
label="RealESRGAN Device",
value=init_config.realesrgan_device,
)
realesrgan_model = gr.Radio(
RealESRGANModelNameList,
label="RealESRGAN model",
value=init_config.realesrgan_model,
)
with gr.Row():
enable_gfpgan = gr.Checkbox(
init_config.enable_gfpgan, label=GFPGAN_HELP
)
gfpgan_device = gr.Radio(
GFPGAN_AVAILABLE_DEVICES,
label="GFPGAN Device",
value=init_config.gfpgan_device,
)
enable_gif = gr.Checkbox(init_config.enable_gif, label=GIF_HELP)
with gr.Tab("Diffusion Model"):
sd_controlnet = gr.Checkbox(
init_config.sd_controlnet, label=f"{SD_CONTROLNET_HELP}"
)
@ -164,6 +222,14 @@ def main(config_file: str):
input,
output_dir,
quality,
enable_interactive_seg,
enable_remove_bg,
enable_realesrgan,
realesrgan_device,
realesrgan_model,
enable_gfpgan,
gfpgan_device,
enable_gif,
],
message,
)