nikunjkdtechnoland
some more add more files
4b98c85
raw
history blame
10.5 kB
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)