diff --git a/lama_cleaner/const.py b/lama_cleaner/const.py index 149a2c0..1aeea29 100644 --- a/lama_cleaner/const.py +++ b/lama_cleaner/const.py @@ -7,7 +7,6 @@ MPS_SUPPORT_MODELS = [ "realisticVision1.4", "sd2", "paint_by_example", - "controlnet", ] DEFAULT_MODEL = "lama" @@ -25,7 +24,6 @@ AVAILABLE_MODELS = [ "sd2", "paint_by_example", "instruct_pix2pix", - "controlnet", ] SD15_MODELS = ["sd1.5", "anything4", "realisticVision1.4"] @@ -50,7 +48,7 @@ Run Stable Diffusion text encoder model on CPU to save GPU memory. """ SD_CONTROLNET_HELP = """ -Run Stable Diffusion 1.5 inpainting model with controlNet-canny model. +Run Stable Diffusion 1.5 inpainting model with Canny ControlNet control. """ LOCAL_FILES_ONLY_HELP = """ @@ -84,3 +82,7 @@ Launch Lama Cleaner as desktop app NO_GUI_AUTO_CLOSE_HELP = """ 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. +""" diff --git a/lama_cleaner/helper.py b/lama_cleaner/helper.py index ee44206..9eb77ee 100644 --- a/lama_cleaner/helper.py +++ b/lama_cleaner/helper.py @@ -135,9 +135,9 @@ def numpy_to_bytes(image_numpy: np.ndarray, ext: str) -> bytes: return image_bytes -def pil_to_bytes(pil_img, ext: str, exif=None) -> bytes: +def pil_to_bytes(pil_img, ext: str, quality: int = 95, exif=None) -> bytes: with io.BytesIO() as output: - pil_img.save(output, format=ext, exif=exif, quality=95) + pil_img.save(output, format=ext, exif=exif, quality=quality) image_bytes = output.getvalue() return image_bytes diff --git a/lama_cleaner/parse_args.py b/lama_cleaner/parse_args.py index 6a3fb1a..c859881 100644 --- a/lama_cleaner/parse_args.py +++ b/lama_cleaner/parse_args.py @@ -37,9 +37,7 @@ def parse_args(): parser.add_argument( "--sd-cpu-textencoder", action="store_true", help=SD_CPU_TEXTENCODER_HELP ) - parser.add_argument( - "--sd-controlnet", action="store_true", help=SD_CONTROLNET_HELP - ) + parser.add_argument("--sd-controlnet", action="store_true", help=SD_CONTROLNET_HELP) parser.add_argument( "--local-files-only", action="store_true", help=LOCAL_FILES_ONLY_HELP ) @@ -70,6 +68,9 @@ def parse_args(): action="store_true", help="Disable model switch in frontend", ) + parser.add_argument( + "--quality", default=95, type=int, help=QUALITY_HELP, + ) # useless args parser.add_argument("--debug", action="store_true", help=argparse.SUPPRESS) diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py index 61a5d33..a896e64 100644 --- a/lama_cleaner/server.py +++ b/lama_cleaner/server.py @@ -93,6 +93,7 @@ is_controlnet: bool = False is_enable_file_manager: bool = False is_enable_auto_saving: bool = False is_desktop: bool = False +image_quality: int = 95 def get_image_ext(img_bytes): @@ -300,10 +301,12 @@ def process(): ext = get_image_ext(origin_image_bytes) + # fmt: off if exif is not None: - bytes_io = io.BytesIO(pil_to_bytes(Image.fromarray(res_np_img), ext, exif=exif)) + bytes_io = io.BytesIO(pil_to_bytes(Image.fromarray(res_np_img), ext, quality=image_quality, exif=exif)) else: - bytes_io = io.BytesIO(pil_to_bytes(Image.fromarray(res_np_img), ext)) + bytes_io = io.BytesIO(pil_to_bytes(Image.fromarray(res_np_img), ext, quality=image_quality)) + # fmt: on response = make_response( send_file( @@ -437,6 +440,10 @@ def main(args): global output_dir global is_enable_auto_saving global is_controlnet + global image_quality + + image_quality = args.quality + if args.sd_controlnet and args.model in SD15_MODELS: is_controlnet = True diff --git a/lama_cleaner/web_config.py b/lama_cleaner/web_config.py index 1643f27..98e33cd 100644 --- a/lama_cleaner/web_config.py +++ b/lama_cleaner/web_config.py @@ -55,6 +55,7 @@ def save_config( model_dir, input, output_dir, + quality, ): config = Config(**locals()) print(config) @@ -97,39 +98,52 @@ def main(config_file: str): with gr.Row(): host = gr.Textbox(init_config.host, label="Host") port = gr.Number(init_config.port, label="Port", precision=0) - with gr.Row(): - 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})", + minimum=75, + maximum=100, + step=1, + ) + + with gr.Column(): + gui = gr.Checkbox(init_config.gui, label=f"{GUI_HELP}") + no_gui_auto_close = gr.Checkbox( + init_config.no_gui_auto_close, label=f"{NO_GUI_AUTO_CLOSE_HELP}" ) - gui = gr.Checkbox(init_config.gui, label=f"{GUI_HELP}") - no_gui_auto_close = gr.Checkbox( - init_config.no_gui_auto_close, label=f"{NO_GUI_AUTO_CLOSE_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}") - disable_nsfw = gr.Checkbox( - init_config.disable_nsfw, label=f"{DISABLE_NSFW_HELP}" - ) - sd_controlnet = gr.Checkbox( - init_config.sd_controlnet, label=f"{SD_CONTROLNET_HELP}" - ) - sd_cpu_textencoder = gr.Checkbox( - init_config.sd_cpu_textencoder, label=f"{SD_CPU_TEXTENCODER_HELP}" - ) - enable_xformers = gr.Checkbox( - init_config.enable_xformers, label=f"{ENABLE_XFORMERS_HELP}" - ) - local_files_only = gr.Checkbox( - init_config.local_files_only, label=f"{LOCAL_FILES_ONLY_HELP}" - ) - 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}" - ) + + 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.Column(): + sd_controlnet = gr.Checkbox( + init_config.sd_controlnet, label=f"{SD_CONTROLNET_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}" + ) + disable_nsfw = gr.Checkbox( + init_config.disable_nsfw, label=f"{DISABLE_NSFW_HELP}" + ) + sd_cpu_textencoder = gr.Checkbox( + init_config.sd_cpu_textencoder, label=f"{SD_CPU_TEXTENCODER_HELP}" + ) + enable_xformers = gr.Checkbox( + init_config.enable_xformers, label=f"{ENABLE_XFORMERS_HELP}" + ) + local_files_only = gr.Checkbox( + init_config.local_files_only, label=f"{LOCAL_FILES_ONLY_HELP}" + ) + save_btn.click( save_config, [ @@ -149,6 +163,7 @@ def main(config_file: str): model_dir, input, output_dir, + quality, ], message, )