-
Notifications
You must be signed in to change notification settings - Fork 1
/
waigen.py
282 lines (259 loc) · 13.7 KB
/
waigen.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
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 "Waifu harem generator"
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_seed = p.seed
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
p.seed = p.seed + 1 if p.subseed_strength == 0
return final_blend