mirror of https://github.com/vladmandic/automatic
137 lines
4.8 KiB
Python
137 lines
4.8 KiB
Python
import gradio as gr
|
|
from PIL import Image
|
|
import numpy as np
|
|
from modules import shared, scripts, processing, masking
|
|
|
|
"""
|
|
Automatic Color Inpaint Script for SD.NEXT - SD & SDXL Support
|
|
|
|
Author: Artheriax
|
|
Credits: SD.NEXT team for script template
|
|
Version: v1
|
|
|
|
Contributions: A new script to automatically inpaint colors in images using Stable Diffusion, Stable Diffusion XL or Flux.
|
|
"""
|
|
|
|
## Config
|
|
|
|
# script title
|
|
supported_models = ['sd','sdxl', 'flux']
|
|
|
|
title = 'Automatic Color Inpaint'
|
|
|
|
# is script available in txt2img tab
|
|
txt2img = False
|
|
|
|
# is script available in img2img tab
|
|
img2img = True
|
|
|
|
### Script definition
|
|
|
|
class Script(scripts.Script):
|
|
def title(self):
|
|
return title
|
|
|
|
def show(self, is_img2img):
|
|
if shared.native:
|
|
return img2img if is_img2img else txt2img
|
|
return False
|
|
|
|
# Define UI for pipeline
|
|
def ui(self, _is_img2img):
|
|
with gr.Row():
|
|
gr.HTML("  ACI: Automatic Color Inpaint<br>")
|
|
with gr.Row():
|
|
color_picker = gr.ColorPicker(
|
|
label="ACI: Color to Mask",
|
|
value="#04F404", # Default to green screen green
|
|
# info="Pick the color you want to mask and inpaint."
|
|
)
|
|
tolerance_slider = gr.Slider(
|
|
minimum=0,
|
|
maximum=100,
|
|
step=1,
|
|
value=65,
|
|
label="ACI: Color tolerance",
|
|
)
|
|
denoising_slider = gr.Slider(
|
|
minimum=0.01,
|
|
maximum=1,
|
|
step=0.01,
|
|
value=0.9,
|
|
label="ACI: Denoising strength",
|
|
)
|
|
with gr.Row():
|
|
dilate_slider = gr.Slider(
|
|
minimum=0,
|
|
maximum=1,
|
|
step=0.01,
|
|
value=0.0,
|
|
label="ACI: Mask dilate",
|
|
# info="(Recommended value = 2 to remove leftovers at edges)"
|
|
)
|
|
erode_slider = gr.Slider(
|
|
minimum=0,
|
|
maximum=1,
|
|
step=0.01,
|
|
value=0,
|
|
label="ACI: Mask erode",
|
|
# info="(Recommended value = 0 for sharpness)"
|
|
)
|
|
blur_slider = gr.Slider(
|
|
minimum=0,
|
|
maximum=1,
|
|
step=0.01,
|
|
value=0.15,
|
|
label="ACI: Mask blur",
|
|
# info="(Recommended value = 0 for sharpness)"
|
|
)
|
|
return [color_picker, tolerance_slider, dilate_slider, erode_slider, blur_slider, denoising_slider]
|
|
|
|
# Run pipeline
|
|
def run(self, p: processing.StableDiffusionProcessing, *args): # pylint: disable=arguments-differ
|
|
if shared.sd_model_type not in supported_models:
|
|
shared.log.warning(f'MoD: class={shared.sd_model.__class__.__name__} model={shared.sd_model_type} required={supported_models}')
|
|
return None
|
|
if not hasattr(p, 'init_images') or p.init_images is None or len(p.init_images) == 0:
|
|
return None
|
|
color_to_mask_hex, mask_tolerance, mask_dilate, mask_erode, mask_blur, inpaint_denoising_strength = args
|
|
|
|
# Convert hex color to RGB tuple (0-255)
|
|
color_to_mask_rgb = tuple(int(color_to_mask_hex[i:i+2], 16) for i in (1, 3, 5))
|
|
|
|
shared.log.debug(f'ACI: rgb={color_to_mask_rgb} tolerance={mask_tolerance} dilate={mask_dilate} erode={mask_erode} blur={mask_blur} denoise={inpaint_denoising_strength}')
|
|
|
|
# Create Color Mask using vectorized operations
|
|
init_image = p.init_images[0].convert("RGB")
|
|
image_np = np.array(init_image)
|
|
|
|
# Calculate Euclidean distance for all pixels at once
|
|
diff = np.linalg.norm(image_np.astype(np.int16) - np.array(color_to_mask_rgb, dtype=np.int16), axis=2)
|
|
calc_tolerance = (diff.max() - diff.min()) * mask_tolerance/100
|
|
mask_np = (diff <= calc_tolerance).astype(np.uint8) * 255
|
|
|
|
mask_image = Image.fromarray(mask_np).convert("L")
|
|
|
|
# If an inpaint mask is already provided from the UI, combine it with the color mask
|
|
if p.image_mask:
|
|
combined_mask = Image.composite(
|
|
Image.new("L", mask_image.size, "white"),
|
|
p.image_mask.convert("L"),
|
|
mask_image
|
|
)
|
|
p.image_mask = combined_mask
|
|
else:
|
|
p.image_mask = mask_image
|
|
|
|
# override inpaint parameters
|
|
p.inpaint_full_res = False # always use full res
|
|
p.denoising_strength = inpaint_denoising_strength
|
|
p.mask_blur = None # do not use legacy mask blur
|
|
p.inpaint_full_res_padding = None # do not use legacy mask blur
|
|
masking.opts.mask_blur = mask_blur # new masking params triggers masking.py:run_mask
|
|
masking.opts.mask_erode = mask_erode
|
|
masking.opts.mask_dilate = mask_dilate
|
|
|
|
return None
|