joseph
/
lama_cleaner
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update web_config

This commit is contained in:
Qing 2024-01-31 21:51:34 +08:00
parent cdac68a9a9
commit 1d2d39dfa0
8 changed files with 175 additions and 188 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
iopaint/api.py
View File

@ -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,

22
iopaint/cli.py
View File

@ -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)

128
iopaint/const.py
View File

@ -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,
)

2
iopaint/plugins/__init__.py
View File

@ -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(

5
iopaint/plugins/realesrgan.py
View File

@ -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):

28
iopaint/schema.py
View File

@ -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

171
iopaint/web_config.py
View File

@ -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,
)

1
requirements.txt
View File

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