astropulse/sd-palettize

Automatic1111 Extension for palettizing generated images

Palettize

Palettize uses k-means to reduce the number of colors in an image. This script is specifically designed for use with Retro Diffusion pixel art model.

Values can be changed in the "Scripts" dropdown menu.

Examples of Palettize, as well as an Aseprite compatible version can be found here.

Installation

Copy this repo's URL (https://github.com/Astropulse/sd-palettize) into the "Install from URL" section of the Extensions tab of Automatic1111's webui. You will need to restart webui to install the required dependencies.

Palette Credit

ENDESGA Adigun A. Polack Luis Miguel Maldonado PICO-8 Gabriel C. Nintendo Commodore Microsoft Atari

sd-palettize's Issues

Feature: Color Palette support

Hi @Astropulse, I tried to add to restrict the reduced colors to a color palette. But all I got is a wild mix of colors. Do you have any idea why this happens and how to improve it? Here is the result, the palette and my code:

Used Color Palette: endesga-32.txt

import modules.scripts as scripts
import hitherdither
import gradio as gr

import cv2
import numpy as np
from PIL import Image

from modules import images
from modules.processing import Processed, process_images
from modules.shared import opts, cmd_opts, state

from torch import Tensor
from torch.nn import Conv2d
from torch.nn import functional as F
from torch.nn.modules.utils import _pair
from typing import Optional

script_dir = scripts.basedir()

# Runs cv2 k_means quantization on the provided image with "k" color indexes


def palettize(img, k, d):
    image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    image = Image.fromarray(image).convert("RGB")

    if d > 0:
        if k <= 64:
            img_indexed = image.quantize(
                colors=k, method=1, kmeans=k, dither=0).convert('RGB')

            palette = [(255, 0, 0), (0, 255, 0), (0, 0, 255)] # red, green, blue
            palette = hitherdither.palette.Palette(palette)
            img_indexed = hitherdither.ordered.yliluoma.yliluomas_1_ordered_dithering(
            image, palette, order=2**d).convert('RGB')
        else:
            img_indexed = image.quantize(
                colors=k, method=1, kmeans=k, dither=0).convert('RGB')
    else:
        img_indexed = image.quantize(
            colors=k, method=1, kmeans=k, dither=0).convert('RGB')

    result = cv2.cvtColor(np.asarray(img_indexed), cv2.COLOR_RGB2BGR)
    return result

def hex_to_rgb(hex_color):
    hex_color = hex_color[-6:]
    return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))

class Script(scripts.Script):
    def title(self):
        return "Palettize"

    def show(self, is_img2img):
        return True

    def ui(self, is_img2img):
        clusters = gr.Slider(minimum=2, maximum=128, step=1,
                             label='Colors in palette', value=24)
        downscale = gr.Checkbox(
            label='Downscale before processing', value=True)
        with gr.Row():
            scale = gr.Slider(minimum=2, maximum=32, step=1,
                              label='Downscale factor', value=8)
            dither = gr.Slider(minimum=0, maximum=3, step=1,
                               label='Dithering', value=0)
        file = gr.File(label='Chosse a .txt file with hex colors on per line')

        return [downscale, scale, clusters, dither, file]
    
 
    def run(self, p, downscale, scale, clusters, dither, file):

        if dither > 0:
            if clusters <= 64:
                print(
                    f'Palettizing output to {clusters} colors with order {2**dither} dithering...')
            else:
                print('Palette too large, max colors for dithering is 64.')
                print(f'Palettizing output to {clusters} colors...')
        else:
            print(f'Palettizing output to {clusters} colors...')

        processed = process_images(p)

        generations = p.batch_size*p.n_iter

        for i in range(generations + int(generations > 1)):
            # Converts image from "Image" type to numpy array for cv2
            img = np.array(processed.images[i]).astype(np.uint8)

            if downscale:
                img = cv2.resize(img, (int(
                    img.shape[1]/scale), int(img.shape[0]/scale)), interpolation=cv2.INTER_LINEAR)

       

             # Map colors to nearest palette color
            if file is not None:
                with open(file.name, 'r', encoding='utf-8') as f:
                    content = f.read()
                    if content:
                        palette = []
                        palette = [hex_to_rgb(line.strip()) for line in content.splitlines()]
                        print(f'Loaded colors: {palette}')       
                        clusters = len(palette)
                        img = palettize(img, clusters, dither)
                        # Convert the img NumPy array to a PIL Image object
                        img = Image.fromarray(img)
                        # Quantize the image to reduce the number of colors
                        img = img.quantize(colors=len(palette))
                        # Map the colors in the quantized image to the colors in the palette
                        recolor_image = Image.new('RGB', img.size)
                        recolor_image.putdata([palette[color] for color in img.getdata()])
                        # save the recolored image
                        img = recolor_image
                        # Convert the PIL Image object to a NumPy array
                        img = np.array(img)
            else:
                img = palettize(img, clusters, dither)

                
            if downscale:
                img = cv2.resize(img, (int(
                    img.shape[1]*scale), int(img.shape[0]*scale)), interpolation=cv2.INTER_NEAREST)

            processed.images[i] = Image.fromarray(img)
            images.save_image(processed.images[i], p.outpath_samples, "palettized", processed.seed +
                            i, processed.prompt, opts.samples_format, info=processed.info, p=p)

            if generations > 1:
                grid = images.image_grid(
                    processed.images[1:generations+1], p.batch_size)
                processed.images[0] = grid

            if opts.grid_save:
                images.save_image(processed.images[0], p.outpath_grids, "palettized",
                                prompt=p.prompt, seed=processed.seed, grid=True, p=p)

        return processed

How to stop reinstalling requirement on every time I reboot sd?

I use the network in China, so I sometimes lose my connection to GitHub,Every time I start, I get a prompt like this:

This is when I am lucky, and when I am unlucky, I will not be able to install successfully and report a long error.I installed other extensions and only this one has the issue
How should I stop the repeated installation, thank you.

[Request] Semi-predefined palette?

Say, I have a 8 bit image/url input to dictate the extension to generate from, while still wanting a 32 bit image. This means the other 24 bit will be decided by the extension while 8 dictated by my image/url input.

Will this be possible to implement?

Regardless, highly appreciate the extension :)

Installation into WebUI Forge

how do I install Palettize in WebUI Forge? I'm new to working with Stable Diffusion, I use Stability Matrix for this