pixl8.py

from PIL import Image,features
import modules.scripts as scripts
from modules import images
from modules.processing import process_images
from modules.shared import opts
import io
import gradio as gr
import subprocess
import os
import sys

#https://github.com/mcychan/nQuantCpp licensed under the GNU General Public License v3.0

class Script(scripts.Script): 
 def __init__(self):
 self.script_dir = os.path.dirname(os.path.realpath(__name__))+os.sep+'scripts'+os.sep
 self.nQuant = 'nQuantCpp.exe'
 self.palette_algos = ['PNN', 'PNNLAB', 'NEU', 'WU', 'EAS', 'SPA', 'DIV', 'DL3', 'MMC','Median', 'Maximum Coverage', 'Octree']
 self.color_palletes = ['Automatic','NES']
 def title(self):
 return "PXL8"
 def show(self, is_img2img):
 return True

 def ui(self, is_img2img):
 with gr.Row():
 dither = gr.Checkbox(True, label="Dither")
 rescale = gr.Checkbox(True, label="Rescale to original size")
 rescale_before = gr.Checkbox(True, label="Rescale before quantize")
 with gr.Row():
 downscale = gr.Slider(minimum=1, maximum=64, step=1, value=8, label="Downscale multiplier")
 color_pal_size = gr.Slider(minimum=0, maximum=256, step=1, value=16, label="Color palette size")
 if features.check_feature("libimagequant"):
 self.palette_algos.insert(0, "libimagequant")
 palette_algo = gr.Radio(choices=self.palette_algos, value=self.palette_algos[0], label='Palette Extraction Algorithm')

 #add dropdown for color palletes
 #custom_palette = gr.Dropdown(choices=self.color_palletes, label="Custom Palette")
 
 return [downscale,rescale,rescale_before,color_pal_size,palette_algo,dither]

 def run(self, p, downscale,rescale,rescale_before,color_pal_size,palette_algo,dither): 
 def process(im):
 isDither = 1 if dither else 0
 isSmart = 1 if palette_algo in self.palette_algos[:9] else 0
 #isCustomPalette = 1 if custom_palette != "None" else 0
 temp = os.path.join(self.script_dir,"temp.png")
 out_width, out_height = im.size
 sample_width = int(out_width / downscale)
 sample_height = int(out_height / downscale)
 if rescale_before:
 work = im.resize((sample_width, sample_height), Image.NEAREST)
 else:
 work = im
 if color_pal_size != 0:
 if isSmart == 0:
 method = Image.Quantize.MEDIANCUT if palette_algo == 'Median' else Image.Quantize.MAXCOVERAGE if palette_algo == 'Maximum Coverage' else Image.Quantize.FASTOCTREE if palette_algo == 'Octree' else Image.Quantize.LIBIMAGEQUANT
 work = work.convert("RGB").quantize(colors=int(color_pal_size), method=method, dither=Image.Dither.FLOYDSTEINBERG if isDither else 0)
 else:
 work = work.convert("RGB")
 work.save(temp)
 run = subprocess.run([os.path.join(self.script_dir, self.nQuant), temp, f"/m ",str(color_pal_size), f"/a ",palette_algo, "/d ","y" if isDither==1 else "n",], stdout=subprocess.DEVNULL)
 s_t = os.path.join(self.script_dir,f"temp-{palette_algo}quant{color_pal_size}.png")
 #open s_t without holding it open
 work = Image.open(s_t).convert("RGB")
 os.remove(s_t)
 os.remove(temp)
 if rescale_before == False:
 work = work.resize((sample_width, sample_height), Image.NEAREST)
 if rescale:
 work = work.resize((out_width, out_height), Image.NEAREST)
 return work
 out = process_images(p)
 for i in range(len(out.images)):
 out.images[i] = process(out.images[i])
 images.save_image(out.images[i], p.outpath_samples, "", out.seed + i, out.prompt, opts.samples_format, info= out.info, p=p)
 return out