Upload 9 files

scripts/pixelart.py

@@ -0,0 +1,148 @@
+import gradio as gr
+
+import modules.scripts as scripts
+from modules import images
+from modules.shared import opts
+
+from sd_webui_pixelart.utils import DITHER_METHODS, QUANTIZATION_METHODS, downscale_image, limit_colors, resize_image, convert_to_black_and_white, convert_to_grayscale
+
+class Script(scripts.Script):
+    def title(self):
+        return "Pixel art"
+
+
+    def show(self, is_img2img):
+        return scripts.AlwaysVisible
+
+
+    def ui(self, is_img2img):
+        quantization_methods = ['Median cut', 'Maximum coverage', 'Fast octree']
+        dither_methods = ['None', 'Floyd-Steinberg']
+
+        with gr.Accordion("Pixel art", open=False):
+            with gr.Row():
+                enabled = gr.Checkbox(label="Enable", value=False)
+
+            with gr.Column():
+                with gr.Row():
+                    downscale = gr.Slider(label="Downscale", minimum=1, maximum=32, step=2, value=8)
+                    need_rescale = gr.Checkbox(label="Rescale to original size", value=True)
+                with gr.Tabs():
+                    with gr.TabItem("Color"):
+                        enable_color_limit = gr.Checkbox(label="Enable", value=False)
+                        number_of_colors = gr.Slider(label="Palette Size", minimum=1, maximum=256, step=1, value=16)
+                        quantization_method = gr.Radio(choices=quantization_methods, value=quantization_methods[0], label='Colors quantization method')
+                        dither_method = gr.Radio(choices=dither_methods, value=dither_methods[0], label='Colors dither method')
+                        use_k_means = gr.Checkbox(label="Enable k-means for color quantization", value=True)
+                    with gr.TabItem("Grayscale"):
+                        is_grayscale = gr.Checkbox(label="Enable", value=False)
+                        number_of_shades = gr.Slider(label="Palette Size", minimum=1, maximum=256, step=1, value=16)
+                        quantization_method_grayscale = gr.Radio(choices=quantization_methods, value=quantization_methods[0], label='Colors quantization method')
+                        dither_method_grayscale = gr.Radio(choices=dither_methods, value=dither_methods[0], label='Colors dither method')
+                        use_k_means_grayscale = gr.Checkbox(label="Enable k-means for color quantization", value=True)
+                    with gr.TabItem("Black and white"):
+                        with gr.Row():
+                            is_black_and_white = gr.Checkbox(label="Enable", value=False)
+                            is_inversed_black_and_white = gr.Checkbox(label="Inverse", value=False)
+                        with gr.Row():
+                            black_and_white_threshold = gr.Slider(label="Threshold", minimum=1, maximum=256, step=1, value=128)
+                    with gr.TabItem("Custom color palette"):
+                        use_color_palette = gr.Checkbox(label="Enable", value=False)
+                        palette_image=gr.Image(label="Color palette image", type="pil")
+                        palette_colors = gr.Slider(label="Palette Size (only for complex images)", minimum=1, maximum=256, step=1, value=16)
+                        dither_method_palette = gr.Radio(choices=dither_methods, value=dither_methods[0], label='Colors dither method')
+
+        return [enabled, downscale, need_rescale, enable_color_limit, number_of_colors, quantization_method, dither_method, use_k_means, is_grayscale, number_of_shades, quantization_method_grayscale, dither_method_grayscale, use_k_means_grayscale, is_black_and_white, is_inversed_black_and_white, black_and_white_threshold, use_color_palette, palette_image, palette_colors, dither_method_palette]
+
+    def postprocess(
+        self,
+        p,
+        processed,
+        enabled,
+
+        downscale,
+        need_rescale,
+
+        enable_color_limit,
+        number_of_colors,
+        quantization_method,
+        dither_method,
+        use_k_means,
+
+        is_grayscale,
+        number_of_shades,
+        quantization_method_grayscale,
+        dither_method_grayscale,
+        use_k_means_grayscale,
+
+        is_black_and_white,
+        is_inversed_black_and_white,
+        black_and_white_threshold,
+
+        use_color_palette,
+        palette_image,
+        palette_colors,
+        dither_method_palette
+    ):
+        if not enabled:
+            return
+
+        dither = DITHER_METHODS[dither_method]
+        quantize = QUANTIZATION_METHODS[quantization_method]
+        dither_grayscale = DITHER_METHODS[dither_method_grayscale]
+        quantize_grayscale = QUANTIZATION_METHODS[quantization_method_grayscale]
+        dither_palette = DITHER_METHODS[dither_method_palette]
+
+        def process_image(original_image):
+            original_width, original_height = original_image.size
+
+            if original_image.mode != "RGB":
+                new_image = original_image.convert("RGB")
+            else:
+                new_image = original_image
+
+            new_image = downscale_image(new_image, downscale)
+
+            if use_color_palette:
+                new_image = limit_colors(
+                    image=new_image,
+                    palette=palette_image,
+                    palette_colors=palette_colors,
+                    dither=dither_palette
+                )
+
+            if is_black_and_white:
+                new_image = convert_to_black_and_white(new_image, black_and_white_threshold, is_inversed_black_and_white)
+
+            if is_grayscale:
+                new_image = convert_to_grayscale(new_image)
+                new_image = limit_colors(
+                    image=new_image,
+                    limit=int(number_of_shades),
+                    quantize=quantize_grayscale,
+                    dither=dither_grayscale,
+                    use_k_means=use_k_means_grayscale
+                )
+
+            if enable_color_limit:
+                new_image = limit_colors(
+                    image=new_image,
+                    limit=int(number_of_colors),
+                    quantize=quantize,
+                    dither=dither,
+                    use_k_means=use_k_means
+                )
+
+            if need_rescale:
+                new_image = resize_image(new_image, (original_width, original_height))
+
+            return new_image.convert('RGBA')
+
+        for i in range(len(processed.images)):
+            pixel_image = process_image(processed.images[i])
+            processed.images.append(pixel_image)
+
+            images.save_image(pixel_image, p.outpath_samples, "pixel",
+            processed.seed + i, processed.prompt, opts.samples_format, info= processed.info, p=p)
+
+        return processed

scripts/postprocessing_pixelart.py

@@ -0,0 +1,166 @@
+import gradio as gr
+from PIL import features
+
+from modules import scripts_postprocessing
+from modules.shared import opts
+
+from sd_webui_pixelart.utils import DITHER_METHODS, QUANTIZATION_METHODS, downscale_image, limit_colors, resize_image, convert_to_grayscale, convert_to_black_and_white
+
+
+class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
+    name = "Pixel art"
+    order = 20005
+    model = None
+
+    def ui(self):
+        quantization_methods = ['Median cut', 'Maximum coverage', 'Fast octree']
+        dither_methods = ['None', 'Floyd-Steinberg']
+
+        if features.check_feature("libimagequant"):
+            quantization_methods.insert(0, "libimagequant")
+
+        with gr.Blocks():
+            with gr.Accordion(label="Pixel art", open=False):
+                enabled = gr.Checkbox(label="Enable", value=False)
+                with gr.Row():
+                    downscale = gr.Slider(label="Downscale", minimum=1, maximum=32, step=2, value=8)
+                    need_rescale = gr.Checkbox(label="Rescale to original size", value=True)
+                with gr.Tabs():
+                    with gr.TabItem("Color"):
+                        enable_color_limit = gr.Checkbox(label="Enable", value=False)
+                        number_of_colors = gr.Slider(label="Palette Size", minimum=1, maximum=256, step=1, value=16)
+                        quantization_method = gr.Radio(choices=quantization_methods, value=quantization_methods[0], label='Colors quantization method')
+                        dither_method = gr.Radio(choices=dither_methods, value=dither_methods[0], label='Colors dither method')
+                        use_k_means = gr.Checkbox(label="Enable k-means for color quantization", value=True)
+                    with gr.TabItem("Grayscale"):
+                        is_grayscale = gr.Checkbox(label="Enable", value=False)
+                        number_of_shades = gr.Slider(label="Palette Size", minimum=1, maximum=256, step=1, value=16)
+                        quantization_method_grayscale = gr.Radio(choices=quantization_methods, value=quantization_methods[0], label='Colors quantization method')
+                        dither_method_grayscale = gr.Radio(choices=dither_methods, value=dither_methods[0], label='Colors dither method')
+                        use_k_means_grayscale = gr.Checkbox(label="Enable k-means for color quantization", value=True)
+                    with gr.TabItem("Black and white"):
+                        with gr.Row():
+                            black_and_white = gr.Checkbox(label="Enable", value=False)
+                            inversed_black_and_white = gr.Checkbox(label="Inverse", value=False)
+                        with gr.Row():
+                            black_and_white_threshold = gr.Slider(label="Threshold", minimum=1, maximum=256, step=1, value=128)
+                    with gr.TabItem("Custom color palette"):
+                        use_color_palette = gr.Checkbox(label="Enable", value=False)
+                        palette_image=gr.Image(label="Color palette image", type="pil")
+                        palette_colors = gr.Slider(label="Palette Size (only for complex images)", minimum=1, maximum=256, step=1, value=16)
+                        dither_method_palette = gr.Radio(choices=dither_methods, value=dither_methods[0], label='Colors dither method')
+
+        return {
+            "enabled": enabled,
+
+            "downscale": downscale,
+            "need_rescale": need_rescale,
+
+            "enable_color_limit": enable_color_limit,
+            "number_of_colors": number_of_colors,
+            "quantization_method": quantization_method,
+            "dither_method": dither_method,
+            "use_k_means": use_k_means,
+
+            "is_grayscale": is_grayscale,
+            "number_of_shades": number_of_shades,
+            "quantization_method_grayscale": quantization_method_grayscale,
+            "dither_method_grayscale": dither_method_grayscale,
+            "use_k_means_grayscale": use_k_means_grayscale,
+
+            "use_color_palette": use_color_palette,
+            "palette_image": palette_image,
+            "palette_colors": palette_colors,
+            "dither_method_palette": dither_method_palette,
+
+            "black_and_white": black_and_white,
+            "inversed_black_and_white": inversed_black_and_white,
+            "black_and_white_threshold": black_and_white_threshold,
+        }
+
+
+    def process(
+            self,
+            pp: scripts_postprocessing.PostprocessedImage,
+
+            enabled,
+
+            downscale,
+            need_rescale,
+
+            enable_color_limit,
+            number_of_colors,
+            quantization_method,
+            dither_method,
+            use_k_means,
+
+            is_grayscale,
+            number_of_shades,
+            quantization_method_grayscale,
+            dither_method_grayscale,
+            use_k_means_grayscale,
+
+            use_color_palette,
+            palette_image,
+            palette_colors,
+            dither_method_palette,
+
+            black_and_white,
+            inversed_black_and_white,
+            black_and_white_threshold
+        ):
+        dither = DITHER_METHODS[dither_method]
+        quantize = QUANTIZATION_METHODS[quantization_method]
+        dither_grayscale = DITHER_METHODS[dither_method_grayscale]
+        quantize_grayscale = QUANTIZATION_METHODS[quantization_method_grayscale]
+        dither_palette = DITHER_METHODS[dither_method_palette]
+
+        if not enabled:
+            return
+
+        def process_image(original_image):
+            original_width, original_height = original_image.size
+
+            if original_image.mode != "RGB":
+                new_image = original_image.convert("RGB")
+            else:
+                new_image = original_image
+
+            new_image = downscale_image(new_image, downscale)
+
+            if use_color_palette:
+                new_image = limit_colors(
+                    image=new_image,
+                    palette=palette_image,
+                    palette_colors=palette_colors,
+                    dither=dither_palette
+                )
+
+            if black_and_white:
+                new_image = convert_to_black_and_white(new_image, black_and_white_threshold, inversed_black_and_white)
+
+            if is_grayscale:
+                new_image = convert_to_grayscale(new_image)
+                new_image = limit_colors(
+                    image=new_image,
+                    limit=int(number_of_shades),
+                    quantize=quantize_grayscale,
+                    dither=dither_grayscale,
+                    use_k_means=use_k_means_grayscale
+                )
+
+            if enable_color_limit:
+                new_image = limit_colors(
+                    image=new_image,
+                    limit=int(number_of_colors),
+                    quantize=quantize,
+                    dither=dither,
+                    use_k_means=use_k_means
+                )
+
+            if need_rescale:
+                new_image = resize_image(new_image, (original_width, original_height))
+
+            return new_image.convert('RGBA')
+
+        pp.image = process_image(pp.image)

sd_webui_pixelart/utils.py

@@ -0,0 +1,75 @@
+from PIL import Image
+
+# https://pillow.readthedocs.io/en/stable/reference/Image.html#dither-modes
+DITHER_METHODS = {
+    "None": Image.Dither.NONE,
+    "Floyd-Steinberg": Image.Dither.FLOYDSTEINBERG
+}
+
+#https://pillow.readthedocs.io/en/stable/reference/Image.html#quantization-methods
+QUANTIZATION_METHODS = {
+    "Median cut": Image.Quantize.MEDIANCUT,
+    "Maximum coverage": Image.Quantize.MAXCOVERAGE,
+    "Fast octree": Image.Quantize.FASTOCTREE,
+    "libimagequant": Image.Quantize.LIBIMAGEQUANT
+}
+
+
+def downscale_image(image: Image, scale: int) -> Image:
+    width, height = image.size
+    downscaled_image = image.resize((int(width / scale), int(height / scale)), Image.NEAREST)
+    return downscaled_image
+
+
+def resize_image(image: Image, size) -> Image:
+    width, height = size
+    resized_image = image.resize((width, height), Image.NEAREST)
+    return resized_image
+
+
+def limit_colors(
+        image,
+        limit: int=16,
+        palette=None,
+        palette_colors: int=256,
+        quantize: Image.Quantize=Image.Quantize.MEDIANCUT,
+        dither: Image.Dither=Image.Dither.NONE,
+        use_k_means: bool=False
+    ):
+    if use_k_means:
+        k_means_value = limit
+    else:
+        k_means_value = 0
+
+    if palette:
+        palette_image = palette
+        ppalette = palette.getcolors()
+        if ppalette:
+            color_palette = palette.quantize(colors=len(list(set(ppalette))))
+        else:
+            colors = len(palette_image.getcolors()) if palette_image.getcolors() else palette_colors
+            color_palette = palette_image.quantize(colors, kmeans=colors)
+    else:
+        # we need to get palette from image, because
+        # dither in quantize doesn't work without it
+        # https://pillow.readthedocs.io/en/stable/_modules/PIL/Image.html#Image.quantize
+        color_palette = image.quantize(colors=limit, kmeans=k_means_value, method=quantize, dither=Image.Dither.NONE)
+
+    new_image = image.quantize(palette=color_palette, dither=dither)
+
+    return new_image
+
+
+def convert_to_grayscale(image):
+    new_image = image.convert("L")
+    return new_image.convert("RGB")
+
+
+def convert_to_black_and_white(image: Image, threshold: int=128, is_inversed: bool=False):
+    if is_inversed:
+        apply_threshold = lambda x : 255 if x < threshold else 0
+    else:
+        apply_threshold = lambda x : 255 if x > threshold else 0
+
+    black_and_white_image = image.convert('L', dither=Image.Dither.NONE).point(apply_threshold, mode='1')
+    return black_and_white_image.convert("RGB")