|
import json |
|
import os |
|
from pathlib import Path |
|
|
|
from iopaint.schema import ( |
|
Device, |
|
InteractiveSegModel, |
|
RemoveBGModel, |
|
RealESRGANModel, |
|
ApiConfig, |
|
) |
|
|
|
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False" |
|
|
|
from datetime import datetime |
|
from json import JSONDecodeError |
|
|
|
import gradio as gr |
|
from iopaint.download import scan_models |
|
from loguru import logger |
|
|
|
from iopaint.const import * |
|
|
|
|
|
_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, |
|
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, |
|
remove_bg_model=RemoveBGModel.briaai_rmbg_1_4, |
|
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, |
|
) |
|
|
|
|
|
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, |
|
model_dir, |
|
no_half, |
|
low_mem, |
|
cpu_offload, |
|
disable_nsfw_checker, |
|
local_files_only, |
|
cpu_textencoder, |
|
device, |
|
input, |
|
output_dir, |
|
quality, |
|
enable_interactive_seg, |
|
interactive_seg_model, |
|
interactive_seg_device, |
|
enable_remove_bg, |
|
remove_bg_model, |
|
enable_anime_seg, |
|
enable_realesrgan, |
|
realesrgan_device, |
|
realesrgan_model, |
|
enable_gfpgan, |
|
gfpgan_device, |
|
enable_restoreformer, |
|
restoreformer_device, |
|
inbrowser, |
|
): |
|
config = WebConfig(**locals()) |
|
if str(config.input) == ".": |
|
config.input = None |
|
if str(config.output_dir) == ".": |
|
config.output_dir = None |
|
config.model = config.model.strip() |
|
print(config.model_dump_json(indent=4)) |
|
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: {str(_config_file.absolute())}" |
|
logger.info(msg) |
|
try: |
|
with open(_config_file, "w", encoding="utf-8") as f: |
|
f.write(config.model_dump_json(indent=4)) |
|
except Exception as e: |
|
return f"Save configure file failed: {str(e)}" |
|
return msg |
|
|
|
|
|
def change_current_model(new_model): |
|
return new_model |
|
|
|
|
|
def main(config_file: Path): |
|
global _config_file |
|
_config_file = config_file |
|
|
|
init_config = load_config(config_file) |
|
downloaded_models = [it.name for it in scan_models()] |
|
|
|
with gr.Blocks() as demo: |
|
with gr.Row(): |
|
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() |
|
|
|
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) |
|
inbrowser = gr.Checkbox(init_config.inbrowser, label=INBROWSER_HELP) |
|
|
|
with gr.Column(): |
|
model = gr.Textbox( |
|
init_config.model, |
|
label="Current Model. This is the model that will be used when the service starts. " |
|
"If the model has not been downloaded before, it will be automatically downloaded. " |
|
"You can select a model from the dropdown box below or manually enter the SD/SDXL model ID from HuggingFace, for example, runwayml/stable-diffusion-inpainting.", |
|
) |
|
with gr.Row(): |
|
recommend_model = gr.Dropdown( |
|
["lama", "mat", "migan"] + DIFFUSION_MODELS, |
|
label="Recommended Models", |
|
) |
|
downloaded_model = gr.Dropdown( |
|
downloaded_models, label="Downloaded Models" |
|
) |
|
|
|
device = gr.Radio( |
|
Device.values(), label="Device", value=init_config.device |
|
) |
|
quality = gr.Slider( |
|
value=95, |
|
label=f"Image Quality ({QUALITY_HELP})", |
|
minimum=75, |
|
maximum=100, |
|
step=1, |
|
) |
|
|
|
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( |
|
init_config.model_dir, label=f"{MODEL_DIR_HELP}" |
|
) |
|
input = gr.Textbox( |
|
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}", |
|
) |
|
|
|
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( |
|
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 |
|
) |
|
remove_bg_model = gr.Radio( |
|
RemoveBGModel.values(), |
|
label="Remove bg model", |
|
value=init_config.remove_bg_model, |
|
) |
|
with gr.Row(): |
|
enable_anime_seg = gr.Checkbox( |
|
init_config.enable_anime_seg, label=ANIMESEG_HELP |
|
) |
|
|
|
with gr.Row(): |
|
enable_realesrgan = gr.Checkbox( |
|
init_config.enable_realesrgan, label=REALESRGAN_HELP |
|
) |
|
realesrgan_device = gr.Radio( |
|
Device.values(), |
|
label="RealESRGAN Device", |
|
value=init_config.realesrgan_device, |
|
) |
|
realesrgan_model = gr.Radio( |
|
RealESRGANModel.values(), |
|
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( |
|
Device.values(), |
|
label="GFPGAN Device", |
|
value=init_config.gfpgan_device, |
|
) |
|
with gr.Row(): |
|
enable_restoreformer = gr.Checkbox( |
|
init_config.enable_restoreformer, label=RESTOREFORMER_HELP |
|
) |
|
restoreformer_device = gr.Radio( |
|
Device.values(), |
|
label="RestoreFormer Device", |
|
value=init_config.restoreformer_device, |
|
) |
|
|
|
downloaded_model.change(change_current_model, [downloaded_model], model) |
|
recommend_model.change(change_current_model, [recommend_model], model) |
|
|
|
save_btn.click( |
|
save_config, |
|
[ |
|
host, |
|
port, |
|
model, |
|
model_dir, |
|
no_half, |
|
low_mem, |
|
cpu_offload, |
|
disable_nsfw_checker, |
|
local_files_only, |
|
cpu_textencoder, |
|
device, |
|
input, |
|
output_dir, |
|
quality, |
|
enable_interactive_seg, |
|
interactive_seg_model, |
|
interactive_seg_device, |
|
enable_remove_bg, |
|
remove_bg_model, |
|
enable_anime_seg, |
|
enable_realesrgan, |
|
realesrgan_device, |
|
realesrgan_model, |
|
enable_gfpgan, |
|
gfpgan_device, |
|
enable_restoreformer, |
|
restoreformer_device, |
|
inbrowser, |
|
], |
|
message, |
|
) |
|
demo.launch(inbrowser=True, show_api=False) |
|
|