2024-02-08 09:49:54 +01:00
import json
2024-02-04 14:38:41 +01:00
import os
2024-02-08 09:49:54 +01:00
from pathlib import Path
from iopaint . schema import (
Device ,
InteractiveSegModel ,
RemoveBGModel ,
RealESRGANModel ,
ApiConfig ,
)
2024-02-04 14:38:41 +01:00
os . environ [ " GRADIO_ANALYTICS_ENABLED " ] = " False "
2023-01-20 09:52:38 +01:00
from datetime import datetime
2024-01-31 14:51:34 +01:00
from json import JSONDecodeError
2023-01-20 09:52:38 +01:00
import gradio as gr
2024-02-04 14:38:41 +01:00
from iopaint . download import scan_models
2023-01-20 09:52:38 +01:00
from loguru import logger
2024-01-05 08:19:23 +01:00
from iopaint . const import *
2023-01-20 09:52:38 +01:00
2024-01-31 14:51:34 +01:00
_config_file : Path = None
2024-02-08 09:49:54 +01:00
default_configs = dict (
host = " 127.0.0.1 " ,
port = 8080 ,
2024-02-08 10:35:55 +01:00
inbrowser = True ,
2024-02-08 09:49:54 +01:00
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 ,
2024-02-08 10:16:57 +01:00
remove_bg_model = RemoveBGModel . briaai_rmbg_1_4 ,
2024-02-08 09:49:54 +01:00
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 ,
)
2024-01-31 14:51:34 +01:00
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 :
2024-03-05 15:33:38 +01:00
print ( " Load config file failed, using default configs " )
2024-01-31 14:51:34 +01:00
return WebConfig ( * * default_configs )
else :
return WebConfig ( * * default_configs )
2023-01-20 09:52:38 +01:00
def save_config (
2024-03-05 15:33:38 +01:00
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 ,
2023-01-20 09:52:38 +01:00
) :
2024-01-31 14:51:34 +01:00
config = WebConfig ( * * locals ( ) )
if str ( config . input ) == " . " :
config . input = None
if str ( config . output_dir ) == " . " :
config . output_dir = None
2024-02-04 14:38:41 +01:00
config . model = config . model . strip ( )
2024-01-31 15:10:13 +01:00
print ( config . model_dump_json ( indent = 4 ) )
2023-01-20 09:52:38 +01:00
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 " )
2024-01-31 14:51:34 +01:00
msg = f " [ { current_time } ] Successful save config to: { str ( _config_file . absolute ( ) ) } "
2023-01-20 09:52:38 +01:00
logger . info ( msg )
try :
with open ( _config_file , " w " , encoding = " utf-8 " ) as f :
2024-01-31 14:51:34 +01:00
f . write ( config . model_dump_json ( indent = 4 ) )
2023-01-20 09:52:38 +01:00
except Exception as e :
2024-01-31 14:51:34 +01:00
return f " Save configure file failed: { str ( e ) } "
2023-01-20 09:52:38 +01:00
return msg
2024-02-04 14:38:41 +01:00
def change_current_model ( new_model ) :
return new_model
2024-01-31 14:51:34 +01:00
def main ( config_file : Path ) :
2023-01-20 09:52:38 +01:00
global _config_file
_config_file = config_file
init_config = load_config ( config_file )
2024-02-04 14:38:41 +01:00
downloaded_models = [ it . name for it in scan_models ( ) ]
2023-01-20 09:52:38 +01:00
with gr . Blocks ( ) as demo :
with gr . Row ( ) :
2024-01-31 14:51:34 +01:00
with gr . Column ( ) :
gr . Textbox ( config_file , label = " Config file " , interactive = False )
with gr . Column ( ) :
2023-01-20 09:52:38 +01:00
save_btn = gr . Button ( value = " Save configurations " )
message = gr . HTML ( )
2023-03-20 14:23:49 +01:00
2023-03-28 15:55:25 +02:00
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 )
2024-02-08 10:35:55 +01:00
inbrowser = gr . Checkbox ( init_config . inbrowser , label = INBROWSER_HELP )
2023-03-28 15:55:25 +02:00
2024-02-04 14:38:41 +01:00
with gr . Row ( ) :
recommend_model = gr . Dropdown (
[ " lama " , " mat " , " migan " ] + DIFFUSION_MODELS ,
2024-02-08 09:49:54 +01:00
label = " Recommended Models " ,
2024-02-04 14:38:41 +01:00
)
downloaded_model = gr . Dropdown (
downloaded_models , label = " Downloaded Models "
)
2024-03-05 15:33:38 +01:00
with gr . Column ( ) :
model = gr . Textbox (
init_config . model ,
label = " Current Model. Model will be automatically downloaded. "
" You can select a model in Recommended Models or Downloaded Models or manually enter the SD/SDXL model ID from HuggingFace, for example, runwayml/stable-diffusion-inpainting. " ,
)
2024-02-04 14:38:41 +01:00
2023-03-28 15:55:25 +02:00
device = gr . Radio (
2024-01-31 14:51:34 +01:00
Device . values ( ) , label = " Device " , value = init_config . device
2023-03-28 15:55:25 +02:00
)
quality = gr . Slider (
value = 95 ,
label = f " Image Quality ( { QUALITY_HELP } ) " ,
minimum = 75 ,
maximum = 100 ,
step = 1 ,
)
2024-01-31 14:51:34 +01:00
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 } "
)
2023-03-28 15:55:25 +02:00
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 " ) :
2024-01-31 14:51:34 +01:00
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 ,
)
2023-04-16 03:41:35 +02:00
with gr . Row ( ) :
enable_remove_bg = gr . Checkbox (
init_config . enable_remove_bg , label = REMOVE_BG_HELP
)
2024-02-08 09:49:54 +01:00
remove_bg_model = gr . Radio (
RemoveBGModel . values ( ) ,
label = " Remove bg model " ,
value = init_config . remove_bg_model ,
)
2023-05-09 13:07:12 +02:00
with gr . Row ( ) :
enable_anime_seg = gr . Checkbox (
init_config . enable_anime_seg , label = ANIMESEG_HELP
)
2023-04-16 03:41:35 +02:00
2023-03-28 15:55:25 +02:00
with gr . Row ( ) :
enable_realesrgan = gr . Checkbox (
init_config . enable_realesrgan , label = REALESRGAN_HELP
)
realesrgan_device = gr . Radio (
2024-01-31 14:51:34 +01:00
Device . values ( ) ,
2023-03-28 15:55:25 +02:00
label = " RealESRGAN Device " ,
value = init_config . realesrgan_device ,
)
realesrgan_model = gr . Radio (
2024-01-31 14:51:34 +01:00
RealESRGANModel . values ( ) ,
2023-03-28 15:55:25 +02:00
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 (
2024-01-31 14:51:34 +01:00
Device . values ( ) ,
2023-03-28 15:55:25 +02:00
label = " GFPGAN Device " ,
value = init_config . gfpgan_device ,
)
2023-03-30 10:07:38 +02:00
with gr . Row ( ) :
enable_restoreformer = gr . Checkbox (
init_config . enable_restoreformer , label = RESTOREFORMER_HELP
)
restoreformer_device = gr . Radio (
2024-01-31 14:51:34 +01:00
Device . values ( ) ,
2023-03-30 10:07:38 +02:00
label = " RestoreFormer Device " ,
value = init_config . restoreformer_device ,
)
2023-03-20 14:23:49 +01:00
2024-02-04 14:38:41 +01:00
downloaded_model . change ( change_current_model , [ downloaded_model ] , model )
recommend_model . change ( change_current_model , [ recommend_model ] , model )
2023-02-06 15:00:47 +01:00
save_btn . click (
save_config ,
[
host ,
port ,
model ,
2024-01-31 14:51:34 +01:00
model_dir ,
2023-02-06 15:00:47 +01:00
no_half ,
2024-01-31 14:51:34 +01:00
low_mem ,
2023-02-06 15:00:47 +01:00
cpu_offload ,
2024-01-31 14:51:34 +01:00
disable_nsfw_checker ,
2023-02-06 15:00:47 +01:00
local_files_only ,
2024-01-31 14:51:34 +01:00
cpu_textencoder ,
device ,
2023-02-06 15:00:47 +01:00
input ,
output_dir ,
2023-03-20 14:23:49 +01:00
quality ,
2023-03-28 15:55:25 +02:00
enable_interactive_seg ,
2023-04-16 03:41:35 +02:00
interactive_seg_model ,
interactive_seg_device ,
2023-03-28 15:55:25 +02:00
enable_remove_bg ,
2024-02-08 09:49:54 +01:00
remove_bg_model ,
2023-05-09 13:07:12 +02:00
enable_anime_seg ,
2023-03-28 15:55:25 +02:00
enable_realesrgan ,
realesrgan_device ,
realesrgan_model ,
enable_gfpgan ,
gfpgan_device ,
2023-03-30 10:07:38 +02:00
enable_restoreformer ,
restoreformer_device ,
2024-02-08 10:35:55 +01:00
inbrowser ,
2023-02-06 15:00:47 +01:00
] ,
message ,
)
2023-01-20 09:52:38 +01:00
demo . launch ( inbrowser = True , show_api = False )