add image quality arg

https://github.com/Sanster/lama-cleaner/issues/229
2023-03-20 21:23:49 +08:00 · 2023-03-20 21:23:49 +08:00 · 1bb25bebe6
commit 1bb25bebe6
parent 5f4c62ac18
5 changed files with 67 additions and 42 deletions
--- 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.
+"""
--- 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

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

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