diff --git a/hakuimg/image_preprocess.py b/hakuimg/image_preprocess.py new file mode 100644 index 0000000..5d4bc81 --- /dev/null +++ b/hakuimg/image_preprocess.py @@ -0,0 +1,28 @@ +import torch +import numpy as np +from PIL import Image + + +def tensor_to_pil(image: torch.Tensor): + return Image.fromarray( + np.clip(255.0 * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) + ) + + +def pil_to_tensor(image: Image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + + +def image_preprocess(img: torch.Tensor, device: str): + use_channel_last = False + if img.ndim == 3: + img = img.unsqueeze(0) + if img.size(3) <= 4: + img = img.permute(0, 3, 1, 2) + use_channel_last = True + if img.size(1) == 4: + img = img[:, :3] + org_device = img.device + if device != "default": + img = img.to(device) + return img, use_channel_last, org_device diff --git a/hakuimg/outline_expansion.py b/hakuimg/outline_expansion.py new file mode 100644 index 0000000..c2ad095 --- /dev/null +++ b/hakuimg/outline_expansion.py @@ -0,0 +1,19 @@ +from PIL import Image +from .image_preprocess import image_preprocess, pil_to_tensor, tensor_to_pil +from pixeloe.torch.outline import outline_expansion + + +def run( + img: Image, + pixel_size: int, + thickness: int, + device: str, +): + img = pil_to_tensor(img) + img, use_channel_last, org_device = image_preprocess(img, device) + oe_image, _ = outline_expansion(img, thickness, thickness, pixel_size) + oe_image = oe_image.to(org_device) + if use_channel_last: + oe_image = oe_image.permute(0, 2, 3, 1) + oe_image = tensor_to_pil(oe_image) + return oe_image diff --git a/hakuimg/pixeloe.py b/hakuimg/pixeloe.py new file mode 100644 index 0000000..41bfcf9 --- /dev/null +++ b/hakuimg/pixeloe.py @@ -0,0 +1,37 @@ +from PIL import Image +from pixeloe.torch.pixelize import pixelize +from .image_preprocess import image_preprocess, pil_to_tensor, tensor_to_pil + + +def run( + img: Image, + pixel_size: int, + thickness: int, + num_colors: int, + mode: str, + quant_mode: str, + dither_mode: str, + device: str, + color_quant: bool, + no_post_upscale: bool, +): + img = pil_to_tensor(img) + img, use_channel_last, org_device = image_preprocess(img, device) + result, _, _ = pixelize( + img, + pixel_size, + thickness, + mode, + do_color_match=True, + do_quant=color_quant, + num_colors=num_colors, + quant_mode=quant_mode, + dither_mode=dither_mode, + no_post_upscale=no_post_upscale, + return_intermediate=True, + ) + result = result.to(org_device) + if use_channel_last: + result = result.permute(0, 2, 3, 1) + result = tensor_to_pil(result) + return result diff --git a/install.py b/install.py new file mode 100644 index 0000000..642573b --- /dev/null +++ b/install.py @@ -0,0 +1,62 @@ +import launch +import pkg_resources +from pathlib import Path +from typing import Tuple, Optional + + +repo_root = Path(__file__).parent +main_req_file = repo_root / "requirements.txt" + + +def comparable_version(version: str) -> Tuple: + return tuple(version.split(".")) + + +def get_installed_version(package: str) -> Optional[str]: + try: + return pkg_resources.get_distribution(package).version + except Exception: + return None + + +def extract_base_package(package_string: str) -> str: + base_package = package_string.split("@git")[0] + return base_package + + +def install_requirements(req_file): + with open(req_file) as file: + for package in file: + try: + package = package.strip() + if "==" in package: + package_name, package_version = package.split("==") + installed_version = get_installed_version(package_name) + if installed_version != package_version: + launch.run_pip( + f"install -U {package}", + f"a1111-sd-webui-haku-img requirement: changing {package_name} version from {installed_version} to {package_version}", + ) + elif ">=" in package: + package_name, package_version = package.split(">=") + installed_version = get_installed_version(package_name) + if not installed_version or comparable_version( + installed_version + ) < comparable_version(package_version): + launch.run_pip( + f"install -U {package}", + f"a1111-sd-webui-haku-img requirement: changing {package_name} version from {installed_version} to {package_version}", + ) + elif not launch.is_installed(extract_base_package(package)): + launch.run_pip( + f"install {package}", + f"a1111-sd-webui-haku-img requirement: {package}", + ) + except Exception as e: + print(e) + print( + f"Warning: Failed to install {package}" + ) + + +install_requirements(main_req_file) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..fed21c3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,9 @@ +torch +torchvision +pillow +numpy +opencv-python +scipy +matplotlib +blendmodes +pixeloe>=0.1.4 \ No newline at end of file diff --git a/scripts/main.py b/scripts/main.py index 75e7f68..6760850 100755 --- a/scripts/main.py +++ b/scripts/main.py @@ -18,6 +18,8 @@ from hakuimg import ( color, sketch, pixel, + pixeloe, + outline_expansion, neon, curve, chromatic, @@ -286,6 +288,60 @@ def add_tab(): pixel_btn = gr.Button("refresh", variant="primary") pixel_rst_btn = gr.Button("reset") + with gr.TabItem("PixelOE", elem_id="haku_PixelOE"): + pixeloe_pixel_size = gr.Slider( + 1, 32, 4, step=1, label="pixel size" + ) + pixeloe_thickness = gr.Slider( + 0, 6, 2, step=1, label="thickness" + ) + pixeloe_num_colors = gr.Slider( + 2, 256, 256, step=1, label="num colors" + ) + pixeloe_mode = gr.Radio( + ["contrast", "k_centroid", "lanczos", "nearest", "bilinear"], + value="contrast", + label="mode", + ) + pixeloe_quant_mode = gr.Radio( + ["kmeans", "weighted-kmeans", "repeat-kmeans"], + value="kmeans", + label="quant mode", + ) + pixeloe_dither_mode = gr.Radio( + ["ordered", "error_diffusion", "none"], + value="ordered", + label="dither mode", + ) + pixeloe_device = gr.Radio( + ["default", "cpu", "cuda", "mps"], + value="default", + label="device", + ) + pixeloe_color_quant = gr.Checkbox(label="color quant", value=False) + pixeloe_no_post_upscale = gr.Checkbox(label="no post upscale", value=False) + + with gr.Row(): + pixeloe_btn = gr.Button("refresh", variant="primary") + pixeloe_rst_btn = gr.Button("reset") + + with gr.TabItem("OutlineExpansion", elem_id="haku_OutlineExpansion"): + outline_expansion_pixel_size = gr.Slider( + 1, 32, 4, step=1, label="pixel size" + ) + outline_expansion_thickness = gr.Slider( + 1, 6, 3, step=1, label="thickness" + ) + outline_expansion_device = gr.Radio( + ["default", "cpu", "cuda", "mps"], + value="default", + label="device", + ) + + with gr.Row(): + outline_expansion_btn = gr.Button("refresh", variant="primary") + outline_expansion_rst_btn = gr.Button("reset") + with gr.TabItem("Glow", elem_id="haku_Glow"): neon_mode = gr.Radio( ["BS", "BMBL"], value="BS", label="Glow mode" @@ -534,6 +590,36 @@ def add_tab(): pixel_btn.click(pixel.run, all_p_input, image_out) pixel_rst_btn.click(lambda: [16, 8, 0, 5, "kmeans"], None, all_p_set) + # pixeloe + all_pixeloe_set = [ + pixeloe_pixel_size, + pixeloe_thickness, + pixeloe_num_colors, + pixeloe_mode, + pixeloe_quant_mode, + pixeloe_dither_mode, + pixeloe_device, + pixeloe_color_quant, + pixeloe_no_post_upscale + ] + all_pixeloe_input = [image_eff] + all_pixeloe_set + for component in all_pixeloe_set: + _release_if_possible(component, pixeloe.run, all_pixeloe_input, image_out) + pixeloe_btn.click(pixeloe.run, all_pixeloe_input, image_out) + pixeloe_rst_btn.click(lambda: [4, 2, 256, "constrast", "kmeans", "kmeans", "default", False, False], None, all_pixeloe_set) + + # outline expansion + all_outline_expansion_set = [ + outline_expansion_pixel_size, + outline_expansion_thickness, + outline_expansion_device + ] + all_outline_expansion_input = [image_eff] + all_outline_expansion_set + for component in all_outline_expansion_set: + _release_if_possible(component, outline_expansion.run, all_outline_expansion_input, image_out) + outline_expansion_btn.click(outline_expansion.run, all_outline_expansion_input, image_out) + outline_expansion_rst_btn.click(lambda: [4, 3, "default"], None, all_outline_expansion_set) + # neon all_neon_set = [ neon_blur, @@ -659,4 +745,4 @@ def on_ui_settings(): script_callbacks.on_ui_tabs(add_tab) -script_callbacks.on_ui_settings(on_ui_settings) +script_callbacks.on_ui_settings(on_ui_settings) \ No newline at end of file