Skip to content

Instantly share code, notes, and snippets.

@SkyyySi
Last active January 11, 2023 16:42
Show Gist options
  • Save SkyyySi/01d1abbb85697f1b21ac1fdb144def01 to your computer and use it in GitHub Desktop.
Save SkyyySi/01d1abbb85697f1b21ac1fdb144def01 to your computer and use it in GitHub Desktop.
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment