Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
Overbill1683
Stable Diffusion Webui
提交
47f5e216
S
Stable Diffusion Webui
项目概览
Overbill1683
/
Stable Diffusion Webui
11 个月 前同步成功
通知
1767
Star
81
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
分析
仓库
DevOps
项目成员
Pages
S
Stable Diffusion Webui
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Pages
分析
分析
仓库分析
DevOps
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
提交
体验新版 GitCode,发现更多精彩内容 >>
未验证
提交
47f5e216
编写于
10月 10, 2022
作者:
J
JC-Array
提交者:
GitHub
10月 10, 2022
浏览文件
操作
浏览文件
下载
差异文件
Merge branch 'deepdanbooru_pre_process' into master
上级
aca1553b
76ef3d75
变更
4
隐藏空白更改
内联
并排
Showing
4 changed file
with
111 addition
and
22 deletion
+111
-22
modules/deepbooru.py
modules/deepbooru.py
+80
-20
modules/shared.py
modules/shared.py
+6
-0
modules/textual_inversion/preprocess.py
modules/textual_inversion/preprocess.py
+20
-2
modules/ui.py
modules/ui.py
+5
-0
未找到文件。
modules/deepbooru.py
浏览文件 @
47f5e216
import
os.path
from
concurrent.futures
import
ProcessPoolExecutor
from
multiprocessing
import
get_context
import
multiprocessing
import
time
def
get_deepbooru_tags
(
pil_image
):
"""
This method is for running only one image at a time for simple use. Used to the img2img interrogate.
"""
from
modules
import
shared
# prevents circular reference
create_deepbooru_process
(
shared
.
opts
.
deepbooru_threshold
,
shared
.
opts
.
deepbooru_sort_alpha
)
shared
.
deepbooru_process_return
[
"value"
]
=
-
1
shared
.
deepbooru_process_queue
.
put
(
pil_image
)
while
shared
.
deepbooru_process_return
[
"value"
]
==
-
1
:
time
.
sleep
(
0.2
)
tags
=
shared
.
deepbooru_process_return
[
"value"
]
release_process
()
return
tags
def
_load_tf_and_return_tags
(
pil_image
,
threshold
):
def
deepbooru_process
(
queue
,
deepbooru_process_return
,
threshold
,
alpha_sort
):
model
,
tags
=
get_deepbooru_tags_model
()
while
True
:
# while process is running, keep monitoring queue for new image
pil_image
=
queue
.
get
()
if
pil_image
==
"QUIT"
:
break
else
:
deepbooru_process_return
[
"value"
]
=
get_deepbooru_tags_from_model
(
model
,
tags
,
pil_image
,
threshold
,
alpha_sort
)
def
create_deepbooru_process
(
threshold
,
alpha_sort
):
"""
Creates deepbooru process. A queue is created to send images into the process. This enables multiple images
to be processed in a row without reloading the model or creating a new process. To return the data, a shared
dictionary is created to hold the tags created. To wait for tags to be returned, a value of -1 is assigned
to the dictionary and the method adding the image to the queue should wait for this value to be updated with
the tags.
"""
from
modules
import
shared
# prevents circular reference
shared
.
deepbooru_process_manager
=
multiprocessing
.
Manager
()
shared
.
deepbooru_process_queue
=
shared
.
deepbooru_process_manager
.
Queue
()
shared
.
deepbooru_process_return
=
shared
.
deepbooru_process_manager
.
dict
()
shared
.
deepbooru_process_return
[
"value"
]
=
-
1
shared
.
deepbooru_process
=
multiprocessing
.
Process
(
target
=
deepbooru_process
,
args
=
(
shared
.
deepbooru_process_queue
,
shared
.
deepbooru_process_return
,
threshold
,
alpha_sort
))
shared
.
deepbooru_process
.
start
()
def
release_process
():
"""
Stops the deepbooru process to return used memory
"""
from
modules
import
shared
# prevents circular reference
shared
.
deepbooru_process_queue
.
put
(
"QUIT"
)
shared
.
deepbooru_process
.
join
()
shared
.
deepbooru_process_queue
=
None
shared
.
deepbooru_process
=
None
shared
.
deepbooru_process_return
=
None
shared
.
deepbooru_process_manager
=
None
def
get_deepbooru_tags_model
():
import
deepdanbooru
as
dd
import
tensorflow
as
tf
import
numpy
as
np
this_folder
=
os
.
path
.
dirname
(
__file__
)
model_path
=
os
.
path
.
abspath
(
os
.
path
.
join
(
this_folder
,
'..'
,
'models'
,
'deepbooru'
))
if
not
os
.
path
.
exists
(
os
.
path
.
join
(
model_path
,
'project.json'
)):
# there is no point importing these every time
import
zipfile
from
basicsr.utils.download_util
import
load_file_from_url
load_file_from_url
(
r
"https://github.com/KichangKim/DeepDanbooru/releases/download/v3-20211112-sgd-e28/deepdanbooru-v3-20211112-sgd-e28.zip"
,
model_path
)
load_file_from_url
(
r
"https://github.com/KichangKim/DeepDanbooru/releases/download/v3-20211112-sgd-e28/deepdanbooru-v3-20211112-sgd-e28.zip"
,
model_path
)
with
zipfile
.
ZipFile
(
os
.
path
.
join
(
model_path
,
"deepdanbooru-v3-20211112-sgd-e28.zip"
),
"r"
)
as
zip_ref
:
zip_ref
.
extractall
(
model_path
)
os
.
remove
(
os
.
path
.
join
(
model_path
,
"deepdanbooru-v3-20211112-sgd-e28.zip"
))
...
...
@@ -24,7 +78,13 @@ def _load_tf_and_return_tags(pil_image, threshold):
model
=
dd
.
project
.
load_model_from_project
(
model_path
,
compile_model
=
True
)
return
model
,
tags
def
get_deepbooru_tags_from_model
(
model
,
tags
,
pil_image
,
threshold
,
alpha_sort
):
import
deepdanbooru
as
dd
import
tensorflow
as
tf
import
numpy
as
np
width
=
model
.
input_shape
[
2
]
height
=
model
.
input_shape
[
1
]
image
=
np
.
array
(
pil_image
)
...
...
@@ -46,28 +106,28 @@ def _load_tf_and_return_tags(pil_image, threshold):
for
i
,
tag
in
enumerate
(
tags
):
result_dict
[
tag
]
=
y
[
i
]
result_tags_out
=
[]
unsorted_tags_in_theshold
=
[]
result_tags_print
=
[]
for
tag
in
tags
:
if
result_dict
[
tag
]
>=
threshold
:
if
tag
.
startswith
(
"rating:"
):
continue
result_tags_out
.
append
(
tag
)
unsorted_tags_in_theshold
.
append
((
result_dict
[
tag
],
tag
)
)
result_tags_print
.
append
(
f
'
{
result_dict
[
tag
]
}
{
tag
}
'
)
print
(
'
\n
'
.
join
(
sorted
(
result_tags_print
,
reverse
=
True
)))
return
', '
.
join
(
result_tags_out
).
replace
(
'_'
,
' '
).
replace
(
':'
,
' '
)
# sort tags
result_tags_out
=
[]
sort_ndx
=
0
print
(
alpha_sort
)
if
alpha_sort
:
sort_ndx
=
1
def
subprocess_init_no_cuda
():
import
os
os
.
environ
[
"CUDA_VISIBLE_DEVICES"
]
=
"-1"
# sort by reverse by likelihood and normal for alpha
unsorted_tags_in_theshold
.
sort
(
key
=
lambda
y
:
y
[
sort_ndx
],
reverse
=
(
not
alpha_sort
))
for
weight
,
tag
in
unsorted_tags_in_theshold
:
result_tags_out
.
append
(
tag
)
print
(
'
\n
'
.
join
(
sorted
(
result_tags_print
,
reverse
=
True
)))
def
get_deepbooru_tags
(
pil_image
,
threshold
=
0.5
):
context
=
get_context
(
'spawn'
)
with
ProcessPoolExecutor
(
initializer
=
subprocess_init_no_cuda
,
mp_context
=
context
)
as
executor
:
f
=
executor
.
submit
(
_load_tf_and_return_tags
,
pil_image
,
threshold
,
)
ret
=
f
.
result
()
# will rethrow any exceptions
return
ret
\ No newline at end of file
return
', '
.
join
(
result_tags_out
).
replace
(
'_'
,
' '
).
replace
(
':'
,
' '
)
modules/shared.py
浏览文件 @
47f5e216
...
...
@@ -265,6 +265,12 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters"
'eta_noise_seed_delta'
:
OptionInfo
(
0
,
"Eta noise seed delta"
,
gr
.
Number
,
{
"precision"
:
0
}),
}))
if
cmd_opts
.
deepdanbooru
:
options_templates
.
update
(
options_section
((
'deepbooru-params'
,
"DeepBooru parameters"
),
{
"deepbooru_sort_alpha"
:
OptionInfo
(
True
,
"Sort Alphabetical"
,
gr
.
Checkbox
),
'deepbooru_threshold'
:
OptionInfo
(
0.5
,
"Threshold"
,
gr
.
Slider
,
{
"minimum"
:
0.0
,
"maximum"
:
1.0
,
"step"
:
0.01
}),
}))
class
Options
:
data
=
None
...
...
modules/textual_inversion/preprocess.py
浏览文件 @
47f5e216
...
...
@@ -3,11 +3,14 @@ from PIL import Image, ImageOps
import
platform
import
sys
import
tqdm
import
time
from
modules
import
shared
,
images
from
modules.shared
import
opts
,
cmd_opts
if
cmd_opts
.
deepdanbooru
:
import
modules.deepbooru
as
deepbooru
def
preprocess
(
process_src
,
process_dst
,
process_width
,
process_height
,
process_flip
,
process_split
,
process_caption
):
def
preprocess
(
process_src
,
process_dst
,
process_width
,
process_height
,
process_flip
,
process_split
,
process_caption
,
process_caption_deepbooru
=
False
):
width
=
process_width
height
=
process_height
src
=
os
.
path
.
abspath
(
process_src
)
...
...
@@ -25,10 +28,21 @@ def preprocess(process_src, process_dst, process_width, process_height, process_
if
process_caption
:
shared
.
interrogator
.
load
()
if
process_caption_deepbooru
:
deepbooru
.
create_deepbooru_process
()
def
save_pic_with_caption
(
image
,
index
):
if
process_caption
:
caption
=
"-"
+
shared
.
interrogator
.
generate_caption
(
image
)
caption
=
sanitize_caption
(
os
.
path
.
join
(
dst
,
f
"
{
index
:
05
}
-
{
subindex
[
0
]
}
"
),
caption
,
".png"
)
elif
process_caption_deepbooru
:
shared
.
deepbooru_process_return
[
"value"
]
=
-
1
shared
.
deepbooru_process_queue
.
put
(
image
)
while
shared
.
deepbooru_process_return
[
"value"
]
==
-
1
:
time
.
sleep
(
0.2
)
caption
=
"-"
+
shared
.
deepbooru_process_return
[
"value"
]
caption
=
sanitize_caption
(
os
.
path
.
join
(
dst
,
f
"
{
index
:
05
}
-
{
subindex
[
0
]
}
"
),
caption
,
".png"
)
shared
.
deepbooru_process_return
[
"value"
]
=
-
1
else
:
caption
=
filename
caption
=
os
.
path
.
splitext
(
caption
)[
0
]
...
...
@@ -80,6 +94,10 @@ def preprocess(process_src, process_dst, process_width, process_height, process_
if
process_caption
:
shared
.
interrogator
.
send_blip_to_ram
()
if
process_caption_deepbooru
:
deepbooru
.
release_process
()
def
sanitize_caption
(
base_path
,
original_caption
,
suffix
):
operating_system
=
platform
.
system
().
lower
()
if
(
operating_system
==
"windows"
):
...
...
modules/ui.py
浏览文件 @
47f5e216
...
...
@@ -1036,6 +1036,10 @@ def create_ui(wrap_gradio_gpu_call):
process_flip
=
gr
.
Checkbox
(
label
=
'Create flipped copies'
)
process_split
=
gr
.
Checkbox
(
label
=
'Split oversized images into two'
)
process_caption
=
gr
.
Checkbox
(
label
=
'Use BLIP caption as filename'
)
if
cmd_opts
.
deepdanbooru
:
process_caption_deepbooru
=
gr
.
Checkbox
(
label
=
'Use deepbooru caption as filename'
)
else
:
process_caption_deepbooru
=
gr
.
Checkbox
(
label
=
'Use deepbooru caption as filename'
,
visible
=
False
)
with
gr
.
Row
():
with
gr
.
Column
(
scale
=
3
):
...
...
@@ -1102,6 +1106,7 @@ def create_ui(wrap_gradio_gpu_call):
process_flip
,
process_split
,
process_caption
,
process_caption_deepbooru
],
outputs
=
[
ti_output
,
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录