282 lines
14 KiB
Python
282 lines
14 KiB
Python
from modules.processing import Processed, StableDiffusionProcessingImg2Img, process_images, images, fix_seed
|
|
from modules.shared import opts, cmd_opts, state
|
|
from PIL import Image, ImageOps
|
|
from math import ceil
|
|
import cv2
|
|
|
|
import modules.scripts as scripts
|
|
from modules import sd_samplers
|
|
from random import randint, shuffle
|
|
import random
|
|
from skimage.util import random_noise
|
|
import gradio as gr
|
|
import numpy as np
|
|
|
|
class Script(scripts.Script):
|
|
def title(self):
|
|
return "Multi Subject Rendering"
|
|
|
|
def ui(self, is_img2img):
|
|
if is_img2img: return
|
|
txt2img_samplers_names = [s.name for s in sd_samplers.samplers]
|
|
img2img_samplers_names = [s.name for s in sd_samplers.samplers_for_img2img]
|
|
|
|
# foreground UI
|
|
with gr.Box():
|
|
foregen_prompt = gr.Textbox(label="Foreground prompt ", lines=5, max_lines=2000)
|
|
foregen_iter = gr.Slider(minimum=1, maximum=10, step=1, label='Number of foreground images ', value=5)
|
|
foregen_steps = gr.Slider(minimum=1, maximum=120, step=1, label='foreground steps ', value=24)
|
|
foregen_cfg_scale = gr.Slider(minimum=1, maximum=30, step=0.1, label='foreground cfg scale ', value=12.5)
|
|
foregen_seed_shift = gr.Slider(minimum=0, maximum=1000, step=1, label='foreground new seed+ ', value=1000)
|
|
foregen_sampler = gr.Dropdown(label="foreground sampler", choices=txt2img_samplers_names, value="DDIM")
|
|
foregen_clip = gr.Slider(minimum=0, maximum=12, step=1, label='change clip for foreground (0 = no interaction) ', value=0)
|
|
with gr.Row():
|
|
foregen_size_x = gr.Slider(minimum=64, maximum=2048, step=64, label='foreground width (64 = same as background) ', value=64)
|
|
foregen_size_y = gr.Slider(minimum=64, maximum=2048, step=64, label='foreground height (64 = same as background) ', value=64)
|
|
|
|
# blend UI
|
|
with gr.Box():
|
|
foregen_blend_prompt = gr.Textbox(label="final blend prompt", lines=2, max_lines=2000)
|
|
foregen_blend_steps = gr.Slider(minimum=1, maximum=120, step=1, label='blend steps ', value=64)
|
|
foregen_blend_cfg_scale = gr.Slider(minimum=1, maximum=30, step=0.1, label='blend cfg scale ', value=7.5)
|
|
foregen_blend_denoising_strength = gr.Slider(minimum=0.1, maximum=1, step=0.01, label='blend denoising strength ', value=0.42)
|
|
foregen_blend_sampler = gr.Dropdown(label="blend sampler", choices=img2img_samplers_names, value="DDIM")
|
|
with gr.Row():
|
|
foregen_blend_size_x = gr.Slider(minimum=64, maximum=2048, step=64, label='blend width (64 = same size as background) ', value=64)
|
|
foregen_blend_size_y = gr.Slider(minimum=64, maximum=2048, step=64, label='blend height (64 = same size as background) ', value=64)
|
|
|
|
# Misc options
|
|
foregen_x_shift = gr.Slider(minimum=0, maximum=2, step=0.01, label='Foreground distance from center multiplier ', value=1)
|
|
foregen_y_shift = gr.Slider(minimum=0, maximum=100, step=1, label='Foreground Y shift (far from center = lower) ', value=0)
|
|
foregen_treshold = gr.Slider(minimum=0, maximum=255, step=1, label='Foreground depth cut treshold ', value=92)
|
|
with gr.Row():
|
|
foregen_save_background = gr.Checkbox(label='Save background ', value=True)
|
|
foregen_save_all = gr.Checkbox(label='Save all foreground images ', value=True)
|
|
foregen_face_correction = gr.Checkbox(label='Face correction ', value=True)
|
|
foregen_random_superposition = gr.Checkbox(label='Random superposition ', value=False)
|
|
foregen_reverse_order = gr.Checkbox(label='Reverse order ', value=False)
|
|
return [foregen_prompt,
|
|
foregen_iter,
|
|
foregen_steps,
|
|
foregen_cfg_scale,
|
|
foregen_seed_shift,
|
|
foregen_sampler,
|
|
foregen_clip,
|
|
foregen_size_x,
|
|
foregen_size_y,
|
|
foregen_blend_prompt,
|
|
foregen_blend_steps,
|
|
foregen_blend_cfg_scale,
|
|
foregen_blend_denoising_strength,
|
|
foregen_blend_sampler,
|
|
foregen_blend_size_x,
|
|
foregen_blend_size_y,
|
|
foregen_x_shift,
|
|
foregen_y_shift,
|
|
foregen_treshold,
|
|
foregen_save_background,
|
|
foregen_save_all,
|
|
foregen_face_correction,
|
|
foregen_random_superposition,
|
|
foregen_reverse_order]
|
|
|
|
|
|
|
|
|
|
def run(self,p,foregen_prompt,
|
|
foregen_iter,
|
|
foregen_steps,
|
|
foregen_cfg_scale,
|
|
foregen_seed_shift,
|
|
foregen_sampler,
|
|
foregen_clip,
|
|
foregen_size_x,
|
|
foregen_size_y,
|
|
foregen_blend_prompt,
|
|
foregen_blend_steps,
|
|
foregen_blend_cfg_scale,
|
|
foregen_blend_denoising_strength,
|
|
foregen_blend_sampler,
|
|
foregen_blend_size_x,
|
|
foregen_blend_size_y,
|
|
foregen_x_shift,
|
|
foregen_y_shift,
|
|
foregen_treshold,
|
|
foregen_save_background,
|
|
foregen_save_all,
|
|
foregen_face_correction,
|
|
foregen_random_superposition,
|
|
foregen_reverse_order
|
|
):
|
|
initial_CLIP = opts.data["CLIP_stop_at_last_layers"]
|
|
import scripts.simple_depthmap as sdmg
|
|
sdmg = sdmg.SimpleDepthMapGenerator() #import midas
|
|
|
|
def cut_depth_mask(img,mask_img,foregen_treshold):
|
|
img = img.convert("RGBA")
|
|
mask_img = mask_img.convert("RGBA")
|
|
mask_datas = mask_img.getdata()
|
|
datas = img.getdata()
|
|
treshold = foregen_treshold
|
|
newData = []
|
|
for i in range(len(mask_datas)):
|
|
if mask_datas[i][0] >= foregen_treshold and mask_datas[i][1] >= foregen_treshold and mask_datas[i][2] >= foregen_treshold:
|
|
newData.append(datas[i])
|
|
else:
|
|
newData.append((255, 255, 255, 0))
|
|
mask_img.putdata(newData)
|
|
return mask_img
|
|
|
|
def paste_foreground(background,foreground,index,total_foreground,x_shift,y_shift,foregen_reverse_order):
|
|
background = background.convert("RGBA")
|
|
if not foregen_reverse_order:
|
|
index = total_foreground-index-1
|
|
image = Image.new("RGBA", background.size)
|
|
image.paste(background, (0,0), background)
|
|
alternator = -1 if index % 2 == 0 else 1
|
|
if total_foreground % 2 == 0:
|
|
foreground_shift = background.size[0]/2-foreground.size[0]/2 + background.size[0]/(total_foreground)*alternator*ceil(index/2)*x_shift - background.size[0]/(total_foreground)/2
|
|
else:
|
|
index_shift = index-(index % 2)
|
|
if index == 0:
|
|
foreground_shift = background.size[0]/2-foreground.size[0]/2
|
|
else:
|
|
foreground_shift = background.size[0]/2-foreground.size[0]/2 + background.size[0]/(total_foreground)*alternator*ceil(index/2)*x_shift
|
|
x_shift = int(foreground_shift)
|
|
y_shift = ceil(index/2)*y_shift
|
|
image.paste(foreground, (x_shift,background.size[1]-foreground.size[1]+y_shift), foreground)
|
|
return image
|
|
|
|
fix_seed(p)
|
|
p.do_not_save_samples = not foregen_save_background
|
|
|
|
foregen_size_x = p.width if foregen_size_x == 64 else foregen_size_x
|
|
foregen_size_y = p.height if foregen_size_y == 64 else foregen_size_y
|
|
foregen_blend_size_x = p.width if foregen_blend_size_x == 64 else foregen_blend_size_x
|
|
foregen_blend_size_y = p.height if foregen_blend_size_y == 64 else foregen_blend_size_y
|
|
|
|
o_sampler_name = p.sampler_name
|
|
o_prompt = p.prompt
|
|
o_cfg_scale = p.cfg_scale
|
|
o_steps = p.steps
|
|
o_do_not_save_samples = p.do_not_save_samples
|
|
o_width = p.width
|
|
o_height = p.height
|
|
o_denoising_strength = p.denoising_strength
|
|
o_firstphase_width = p.firstphase_width
|
|
o_firstphase_height = p.firstphase_height
|
|
|
|
n_iter=p.n_iter
|
|
for j in range(n_iter):
|
|
if state.interrupted:
|
|
if foregen_clip > 0:
|
|
opts.data["CLIP_stop_at_last_layers"] = initial_CLIP
|
|
break
|
|
p.n_iter=1
|
|
|
|
#background image processing
|
|
if foregen_clip > 0:
|
|
opts.data["CLIP_stop_at_last_layers"] = initial_CLIP
|
|
p.prompt = o_prompt
|
|
p.sampler_name = o_sampler_name
|
|
p.cfg_scale = o_cfg_scale
|
|
p.steps = o_steps
|
|
p.do_not_save_samples = o_do_not_save_samples
|
|
p.width = o_width
|
|
p.height = o_height
|
|
p.denoising_strength = o_denoising_strength
|
|
p.firstphase_width = o_firstphase_width
|
|
p.firstphase_height = o_firstphase_height
|
|
proc = process_images(p)
|
|
background_image = proc.images[0]
|
|
|
|
# foregrounds processing
|
|
foregen_prompts = foregen_prompt.splitlines()
|
|
foregrounds = []
|
|
if foregen_clip > 0:
|
|
opts.data["CLIP_stop_at_last_layers"] = foregen_clip
|
|
for i in range(foregen_iter):
|
|
if state.interrupted:
|
|
if foregen_clip > 0:
|
|
opts.data["CLIP_stop_at_last_layers"] = initial_CLIP
|
|
break
|
|
p.prompt = foregen_prompts[i] if len(foregen_prompts) > 1 else foregen_prompt
|
|
p.seed = p.seed + foregen_seed_shift
|
|
p.subseed = p.subseed + 1 if p.subseed_strength > 0 else p.subseed
|
|
p.cfg_scale = foregen_cfg_scale
|
|
p.steps = foregen_steps
|
|
p.do_not_save_samples = not foregen_save_all
|
|
p.sampler_name = foregen_sampler
|
|
p.width = foregen_size_x
|
|
p.height = foregen_size_y
|
|
p.denoising_strength = None
|
|
p.firstphase_width = None
|
|
p.firstphase_height = None
|
|
proc = process_images(p)
|
|
foregrounds.append(proc.images[0])
|
|
|
|
# put back clip to original settings before img2img final blend
|
|
if foregen_clip > 0:
|
|
opts.data["CLIP_stop_at_last_layers"] = initial_CLIP
|
|
|
|
#stretch background to final blend if the final blend as a specific size set
|
|
if o_width != foregen_blend_size_x or o_height != foregen_blend_size_y :
|
|
background_image = background_image.resize((foregen_blend_size_x, foregen_blend_size_y), Image.Resampling.LANCZOS)
|
|
|
|
# cut depthmaps and stick foreground on the background
|
|
random.seed(p.seed)
|
|
random_order = [k for k in range(foregen_iter)]
|
|
if foregen_random_superposition :
|
|
shuffle(random_order)
|
|
|
|
for f in range(foregen_iter):
|
|
foreground_image = foregrounds[f]
|
|
# gen depth map
|
|
foreground_image_mask = sdmg.calculate_depth_map_for_waifus(foreground_image)
|
|
# cut depth
|
|
foreground_image = cut_depth_mask(foreground_image,foreground_image_mask,foregen_treshold)
|
|
# paste foregrounds onto background
|
|
background_image = paste_foreground(background_image,foreground_image,random_order[f],foregen_iter,foregen_x_shift,foregen_y_shift,foregen_reverse_order)
|
|
|
|
# final blend
|
|
img2img_processing = StableDiffusionProcessingImg2Img(
|
|
init_images=[background_image],
|
|
resize_mode=0,
|
|
denoising_strength=foregen_blend_denoising_strength,
|
|
mask=None,
|
|
mask_blur=4,
|
|
inpainting_fill=0,
|
|
inpaint_full_res=True,
|
|
inpaint_full_res_padding=0,
|
|
inpainting_mask_invert=0,
|
|
sd_model=p.sd_model,
|
|
outpath_samples=p.outpath_samples,
|
|
outpath_grids=p.outpath_grids,
|
|
prompt=foregen_blend_prompt,
|
|
styles=p.styles,
|
|
seed=p.seed+foregen_seed_shift*(foregen_iter+1),
|
|
subseed=p.subseed,
|
|
subseed_strength=p.subseed_strength,
|
|
seed_resize_from_h=p.seed_resize_from_h,
|
|
seed_resize_from_w=p.seed_resize_from_w,
|
|
sampler_name=foregen_blend_sampler,
|
|
batch_size=p.batch_size,
|
|
n_iter=p.n_iter,
|
|
steps=foregen_blend_steps,
|
|
cfg_scale=foregen_blend_cfg_scale,
|
|
width=foregen_blend_size_x,
|
|
height=foregen_blend_size_y,
|
|
restore_faces=foregen_face_correction,
|
|
tiling=p.tiling,
|
|
do_not_save_samples=False,
|
|
do_not_save_grid=p.do_not_save_grid,
|
|
extra_generation_params=p.extra_generation_params,
|
|
overlay_images=p.overlay_images,
|
|
negative_prompt=p.negative_prompt,
|
|
eta=p.eta
|
|
)
|
|
final_blend = process_images(img2img_processing)
|
|
p.subseed = p.subseed + 1 if p.subseed_strength > 0 else p.subseed
|
|
p.seed = p.seed + 1 if p.subseed_strength == 0 else p.seed
|
|
return final_blend
|