recolor_image_sd_webui.py

import modules.scripts as scripts
import gradio as gr
import os

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


class Script(scripts.Script):  
	# The title of the script. This is what will be displayed in the dropdown menu.
	def title(self):
		return "Alter output colors (con/bri/sat)"


	# Determines when the script should be shown in the dropdown menu via the 
	# returned value. As an example:
	# is_img2img is True if the current tab is img2img, and False if it is txt2img.
	# Thus, return is_img2img to only show the script on the img2img tab.
	def show(self, is_img2img):
		return is_img2img


	# How the script's is displayed in the UI. See https://gradio.app/docs/#components
	# for the different UI components you can use and how to create them.
	# Most UI components can return a value, such as a boolean for a checkbox.
	# The returned values are passed to the run method as parameters.
	def ui(self, is_img2img):
		contrast   = gr.Slider(minimum=-1.0, maximum=1.0, step=0.01, value=0, label="Contrast")
		brightness = gr.Slider(minimum=-1.0, maximum=1.0, step=0.01, value=0, label="Brightness")
		saturation = gr.Slider(minimum=-1.0, maximum=1.0, step=0.01, value=0, label="Saturation") # PILlow calls this "Color"
		overwrite = gr.Checkbox(False, label="Overwrite existing files")
		return [contrast, brightness, saturation, overwrite]


	# This is where the additional processing is implemented. The parameters include
	# self, the model object "p" (a StableDiffusionProcessing class, see
	# processing.py), and the parameters returned by the ui method.
	# Custom functions can be defined here, and additional libraries can be imported 
	# to be used in processing. The return value should be a Processed object, which is
	# what is returned by the process_images method.
	def run(self, p, contrast, brightness, saturation, overwrite):
		def recolor(im, contrast, brightness, saturation):
			from PIL import Image, ImageEnhance

			raf = im

			# A value being 0 here means that no processing needs to be done, so we
			# can just skip that step.
			if contrast:
				raf = ImageEnhance.Contrast(raf).enhance(1+contrast)
			if brightness:
				raf = ImageEnhance.Brightness(raf).enhance(1+brightness)
			if saturation:
				raf = ImageEnhance.Color(raf).enhance(1+saturation)

			return raf

		# If overwrite is false, append the rotation information to the filename
		# using the "basename" parameter and save it in the same directory.
		# If overwrite is true, stop the model from saving its outputs and
		# save the rotated and flipped images instead.
		basename = ""
		if not overwrite:
			if contrast:
				basename += f"_con({contrast})"
			if brightness:
				basename += f"_bri({brightness})"
			if saturation:
				basename += f"_sat({saturation})"
		else:
			p.do_not_save_samples = True

		proc = process_images(p)

		for i, v in enumerate(proc.images):
			proc.images[i] = recolor(v, contrast, brightness, saturation)

			images.save_image(proc.images[i], p.outpath_samples, basename,
				proc.seed + i, proc.prompt, opts.samples_format, info= proc.info, p=p)

		return proc